ENH: Add a block function for creating stacked block arrays.

Add a block function to the current stacking functions vstack, hstack,
stack. It is similar to Matlab's square bracket stacking functionality
for block matrices.

4 files changed, 244 insertions, 7 deletions

diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
index 4b8a27503..a509675ce 100644
--- a/doc/release/1.13.0-notes.rst
+++ b/doc/release/1.13.0-notes.rst
@@ -160,6 +160,11 @@ being iterated over.
 For consistency with ``ndarray`` and ``broadcast``, ``d.ndim`` is a shorthand
 for ``len(d.shape)``.
 
+``np.block`` function for creating blocked arrays
+-------------------------------------------------
+Add a new ``block`` function to the current stacking functions ``vstack``,
+``hstack``, and ``stack``. It is similar to Matlab's square bracket
+notation for creating block matrices.
 
 Improvements
 ============
diff --git a/doc/source/reference/routines.array-manipulation.rst b/doc/source/reference/routines.array-manipulation.rst
index b9cf6f448..cc93d1029 100644
--- a/doc/source/reference/routines.array-manipulation.rst
+++ b/doc/source/reference/routines.array-manipulation.rst
@@ -72,6 +72,7 @@ Joining arrays
    dstack
    hstack
    vstack
+   block
 
 Splitting arrays
 ================
diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
index d7fb0dfb3..639df04b3 100644
--- a/numpy/core/shape_base.py
+++ b/numpy/core/shape_base.py
@@ -1,10 +1,11 @@
 from __future__ import division, absolute_import, print_function
 
-__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'vstack', 'hstack',
-           'stack']
+__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
+           'stack', 'vstack']
+
 
 from . import numeric as _nx
-from .numeric import asanyarray, newaxis
+from .numeric import array, asanyarray, newaxis
 from .multiarray import normalize_axis_index
 
 def atleast_1d(*arys):
@@ -207,6 +208,7 @@ def vstack(tup):
     dstack : Stack arrays in sequence depth wise (along third dimension).
     concatenate : Join a sequence of arrays along an existing axis.
     vsplit : Split array into a list of multiple sub-arrays vertically.
+    block : Assemble arrays from blocks.
 
     Notes
     -----
@@ -262,6 +264,7 @@ def hstack(tup):
     dstack : Stack arrays in sequence depth wise (along third axis).
     concatenate : Join a sequence of arrays along an existing axis.
     hsplit : Split array along second axis.
+    block : Create block arrays.
 
     Notes
     -----
@@ -289,6 +292,7 @@ def hstack(tup):
     else:
         return _nx.concatenate(arrs, 1)
 
+
 def stack(arrays, axis=0):
     """
     Join a sequence of arrays along a new axis.
@@ -315,6 +319,7 @@ def stack(arrays, axis=0):
     --------
     concatenate : Join a sequence of arrays along an existing axis.
     split : Split array into a list of multiple sub-arrays of equal size.
+    block : Create block arrays.
 
     Examples
     --------
@@ -354,3 +359,105 @@ def stack(arrays, axis=0):
     sl = (slice(None),) * axis + (_nx.newaxis,)
     expanded_arrays = [arr[sl] for arr in arrays]
     return _nx.concatenate(expanded_arrays, axis=axis)
+
+
+def block(*arrays):
+    """
+    Create a block array consisting of other arrays.
+
+    You can create a 2-D blocked array with the same notation you use for
+    `np.array`.
+
+    Parameters
+    ----------
+    arrays : sequence of sequence of ndarrays
+        1-D arrays are treated as row vectors.
+
+    Returns
+    -------
+    blocked : ndarray
+        The 2-D array assembled from the given blocks.
+
+    See Also
+    --------
+    stack : Stack arrays in sequence along a new dimension.
+    hstack : Stack arrays in sequence horizontally (column wise).
+    vstack : Stack arrays in sequence vertically (row wise).
+    dstack : Stack arrays in sequence depth wise (along third dimension).
+    concatenate : Join a sequence of arrays together.
+    vsplit : Split array into a list of multiple sub-arrays vertically.
+
+    Notes
+    -----
+    ``block`` is similar to Matlab's "square bracket stacking": ``[A A; B B]``
+
+    Examples
+    --------
+    Stacking in a row:
+    >>> A = np.array([[1, 2, 3]])
+    >>> B = np.array([[2, 3, 4]])
+    >>> block([A, B])
+    array([[1, 2, 3, 2, 3, 4]])
+
+    Stacking in a column:
+    >>> A = np.array([[1, 2, 3]])
+    >>> B = np.array([[2, 3, 4]])
+    >>> block(A, B)
+    array([[1, 2, 3],
+           [2, 3, 4]])
+
+    1-D vectors are treated as row arrays
+    >>> a = np.array([1, 1])
+    >>> b = np.array([2, 2])
+    >>> block([a, b])
+    array([[1, 1, 2, 2]])
+
+    >>> a = np.array([1, 1])
+    >>> b = np.array([2, 2])
+    >>> block(a, b)
+    array([[1, 1],
+           [2, 2]])
+
+    The tuple notation also works:
+    >>> A = np.ones((2, 2))
+    >>> B = 2 * A
+    >>> block((A, B))
+    array([[1, 1, 2, 2],
+           [1, 1, 2, 2]])
+
+    Block array with arbitrary shaped elements
+    >>> One = np.array([[1, 1, 1]])
+    >>> Two = np.array([[2, 2, 2]])
+    >>> Three = np.array([[3, 3, 3, 3, 3, 3]])
+    >>> four = np.array([4, 4, 4, 4, 4, 4])
+    >>> five = np.array([5])
+    >>> six = np.array([6, 6, 6, 6, 6])
+    >>> Zeros = np.zeros((2, 6), dtype=int)
+    >>> block([One, Two],
+    ...        Three,
+    ...        four,
+    ...        [five, six],
+    ...        Zeros)
+    array([[1, 1, 1, 2, 2, 2],
+           [3, 3, 3, 3, 3, 3],
+           [4, 4, 4, 4, 4, 4],
+           [5, 6, 6, 6, 6, 6],
+           [0, 0, 0, 0, 0, 0],
+           [0, 0, 0, 0, 0, 0]])
+
+
+    """
+    if len(arrays) < 1:
+        raise TypeError("need at least one array to create a block array")
+
+    result = []
+    for row in arrays:
+        if isinstance(row, (list, tuple)):
+            result.append(hstack(row))
+        else:
+            result.append(row)
+
+    if len(result) > 1:
+        return vstack(result)
+    else:
+        return atleast_2d(result[0])
diff --git a/numpy/core/tests/test_shape_base.py b/numpy/core/tests/test_shape_base.py
index a42b08e48..6557d456e 100644
--- a/numpy/core/tests/test_shape_base.py
+++ b/numpy/core/tests/test_shape_base.py
@@ -1,11 +1,14 @@
 from __future__ import division, absolute_import, print_function
 
+import warnings
 import numpy as np
-from numpy.compat import long
 from numpy.core import (array, arange, atleast_1d, atleast_2d, atleast_3d,
-                        vstack, hstack, newaxis, concatenate, stack)
-from numpy.testing import (TestCase, assert_, assert_raises, assert_array_equal,
-                           assert_equal, run_module_suite, assert_raises_regex)
+                        block, vstack, hstack, newaxis, concatenate, stack)
+from numpy.testing import (TestCase, assert_, assert_raises,
+                           assert_array_equal, assert_equal, run_module_suite,
+                           assert_raises_regex, assert_almost_equal)
+
+from numpy.compat import long
 
 class TestAtleast1d(TestCase):
     def test_0D_array(self):
@@ -330,5 +333,126 @@ def test_stack():
                         stack, [m, m])
 
 
+class TestBlock(TestCase):
+    def test_block_simple_row_wise(self):
+        A = np.ones((2, 2))
+        B = 2 * A
+        desired = np.array([[1, 1, 2, 2],
+                            [1, 1, 2, 2]])
+        result = block([A, B])
+        assert_equal(desired, result)
+        # with tuples
+        result = block((A, B))
+        assert_almost_equal(desired, result)
+
+    def test_block_simple_column_wise(self):
+        A = np.ones((2, 2))
+        B = 2 * A
+        expected = np.array([[1, 1],
+                             [1, 1],
+                             [2, 2],
+                             [2, 2]])
+        result = block(A, B)
+        assert_almost_equal(expected, result)
+
+    def test_block_needless_brackts(self):
+        A = np.ones((2, 2))
+        B = 2 * A
+        expected = np.array([[1, 1],
+                             [1, 1],
+                             [2, 2],
+                             [2, 2]])
+        result = block([A], [B])  # here are the needless brackets
+        assert_almost_equal(expected, result)
+
+    def test_block_with_1d_arrays_row_wise(self):
+        # # # 1-D vectors are treated as row arrays
+        a = np.array([1, 2, 3])
+        b = np.array([2, 3, 4])
+        expected = np.array([[1, 2, 3, 2, 3, 4]])
+        result = block([a, b])
+        assert_almost_equal(expected, result)
+
+    def test_block_with_1d_arrays_multiple_rows(self):
+        a = np.array([1, 2, 3])
+        b = np.array([2, 3, 4])
+        expected = np.array([[1, 2, 3, 2, 3, 4],
+                             [1, 2, 3, 2, 3, 4]])
+        result = block([a, b], [a, b])
+        assert_almost_equal(expected, result)
+
+    def test_block_with_1d_arrays_column_wise(self):
+        # # # 1-D vectors are treated as row arrays
+        a = np.array([1, 2, 3])
+        b = np.array([2, 3, 4])
+        expected = np.array([[1, 2, 3],
+                             [2, 3, 4]])
+        result = block(a, b)
+        assert_almost_equal(expected, result)
+
+    def test_block_mixed_1d_and_2d(self):
+        A = np.ones((2, 2))
+        B = np.array([2, 2])
+        result = block(A, B)
+        expected = np.array([[1, 1],
+                             [1, 1],
+                             [2, 2]])
+        assert_almost_equal(expected, result)
+
+    def test_block_complex(self):
+        # # # a bit more complex
+        One = np.array([[1, 1, 1]])
+        Two = np.array([[2, 2, 2]])
+        Three = np.array([[3, 3, 3, 3, 3, 3]])
+        four = np.array([4, 4, 4, 4, 4, 4])
+        five = np.array([5])
+        six = np.array([6, 6, 6, 6, 6])
+        Zeros = np.zeros((2, 6))
+
+        result = block([One, Two],
+                       Three,
+                       four,
+                       [five, six],
+                       Zeros)
+        expected = np.array([[1, 1, 1, 2, 2, 2],
+                             [3, 3, 3, 3, 3, 3],
+                             [4, 4, 4, 4, 4, 4],
+                             [5, 6, 6, 6, 6, 6],
+                             [0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 0, 0]])
+        assert_almost_equal(result, expected)
+
+        # additional [] around rows should not have any influence
+        result = block([One, Two],
+                       Three,
+                       [four],
+                       [five, six],
+                       Zeros)
+        assert_almost_equal(result, expected)
+
+        result = block([One, Two],
+                       [Three],
+                       [four],
+                       [five, six],
+                       Zeros)
+        assert_almost_equal(result, expected)
+
+        result = block([One, Two],
+                       [Three],
+                       [four],
+                       [five, six],
+                       [Zeros])
+        assert_almost_equal(result, expected)
+
+    def test_block_with_mismatched_shape(self):
+        a = np.array([0, 0])
+        b = np.eye(2)
+        assert_raises(ValueError, np.block, (a, b))
+        assert_raises(ValueError, np.block, (b, a))
+
+    def test_not_list_or_tuple_as_input(self):
+        assert_raises(TypeError, np.block)
+
+
 if __name__ == "__main__":
     run_module_suite()