8000 ENH: add block() function. · numpy/numpy@3d4c5d1 · GitHub
[go: up one dir, main page]
Skip to content

Commit 3d4c5d1

Browse files
sottesotte
committed
ENH: add block() function.
`block` is similar to Matlab's square bracket stacking functionality for block matices. `block` is an addition to the current stacking functions vstack, hstack, stack.
1 parent fd1d7de commit 3d4c5d1

File tree

2 files changed

+161
-7
lines changed

2 files changed

+161
-7
lines changed

‎numpy/core/shape_base.py

Lines changed: 88 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
from __future__ import division, absolute_import, print_function
22

3-
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'vstack', 'hstack',
4-
'stack']
3+
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
4+
'stack', 'vstack']
5+
56

67
from . import numeric as _nx
7-
from .numeric import asanyarray, newaxis
8+
from .numeric import array, asanyarray, newaxis
89

910
def atleast_1d(*arys):
1011
"""
@@ -202,6 +203,7 @@ def vstack(tup):
202203
dstack : Stack arrays in sequence depth wise (along third dimension).
203204
concatenate : Join a sequence of arrays along an existing axis.
204205
vsplit : Split array into a list of multiple sub-arrays vertically.
206+
block : Create block arrays.
205207
206208
Notes
207209
-----
@@ -253,6 +255,7 @@ def hstack(tup):
253255
dstack : Stack arrays in sequence depth wise (along third axis).
254256
concatenate : Join a sequence of arrays along an existing axis.
255257
hsplit : Split array along second axis.
258+
block : Create block arrays.
256259
257260
Notes
258261
-----
@@ -279,6 +282,7 @@ def hstack(tup):
279282
else:
280283
return _nx.concatenate(arrs, 1)
281284

285+
282286
def stack(arrays, axis=0):
283287
"""
284288
Join a sequence of arrays along a new axis.
@@ -305,6 +309,7 @@ def stack(arrays, axis=0):
305309
--------
306310
concatenate : Join a sequence of arrays along an existing axis.
307311
split : Split array into a list of multiple sub-arrays of equal size.
312+
block : Create block arrays.
308313
309314
Examples
310315
--------
@@ -348,3 +353,83 @@ def stack(arrays, axis=0):
348353
sl = (slice(None),) * axis + (_nx.newaxis,)
349354
expanded_arrays = [arr[sl] for arr in arrays]
350355
return _nx.concatenate(expanded_arrays, axis=axis)
356+
357+
358+
def block(tup_tup):
359+
"""
360+
Create block arrays similar to Matlab's "square bracket stacking":
361+
362+
[A A; B B]
363+
364+
You can create a block array with the same notation you use for
365+
`np.array`.
366+
367+
Parameters
368+
----------
369+
tup_tup : sequence of sequence of ndarrays
370+
1-D arrays are treated as row vectors.
371+
372+
Returns
373+
-------
374+
stacked : ndarray
375+
The 2-D array formed by stacking the given arrays.
376+
377+
See Also
378+
--------
379+
stack : Stack arrays in sequence along a new dimension.
380+
hstack : Stack arrays in sequence horizontally (column wise).
381+
vstack : Stack arrays in sequence vertically (row wise).
382+
dstack : Stack arrays in sequence depth wise (along third dimension).
383+
concatenate : Join a sequence of arrays together.
384+
vsplit : Split array into a list of multiple sub-arrays vertically.
385+
386+
Examples
387+
--------
388+
Stacking in a row:
389+
>>> A = np.array([[1, 2, 3]])
390+
>>> B = np.array([[2, 3, 4]])
391+
>>> block([A, B])
392+
array([[1, 2, 3, 2, 3, 4]])
393+
394+
>>> A = np.array([[1, 1], [1, 1]])
395+
>>> B = 2 * A
396+
>>> block([A, B])
397+
array([[1, 1, 2, 2],
398+
[1, 1, 2, 2]])
399+
400+
>>> # the tuple notation also works
401+
>>> block((A, B))
402+
array([[1, 1, 2, 2],
403+
[1, 1, 2, 2]])
404+
405+
>>> # block matrix with arbitrary shaped elements
406+
>>> One = np.array([[1, 1, 1]])
407+
>>> Two = np.array([[2, 2, 2]])
408+
>>> Three = np.array([[3, 3, 3, 3, 3, 3]])
409+
>>> four = np.array([4, 4, 4, 4, 4, 4])
410+
>>> five = np.array([5])
411+
>>> six = np.array([6, 6, 6, 6, 6])
412+
>>> Zeros = np.zeros((2, 6), dtype=int)
413+
>>> block([[One, Two],
414+
... [Three],
415+
... [four],
416+
... [five, six],
417+
... [Zeros]])
418+
array([[1, 1, 1, 2, 2, 2],
419+
[3, 3, 3, 3, 3, 3],
420+
[4, 4, 4, 4, 4, 4],
421+
[5, 6, 6, 6, 6, 6],
422+
[0, 0, 0, 0, 0, 0],
423+
[0, 0, 0, 0, 0, 0]])
424+
425+
>>> # 1-D vectors are treated as row arrays
426+
>>> a = np.array([1, 1])
427+
>>> b = np.array([2, 2])
428+
>>> block([a, b])
429+
array([[1, 1, 2, 2]])
430+
"""
431+
if isinstance(tup_tup[0], list) or isinstance(tup_tup[0], tuple):
432+
result = vstack([hstack(row) for row in tup_tup])
433+
else:
434+
result = hstack(tup_tup)
435+
return atleast_2d(result)

‎numpy/core/tests/test_shape_base.py

Lines changed: 73 additions & 4 deletions
384
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
from __future__ import division, absolute_import, print_function
22

3+
import warnings
34
import numpy as np
4-
from numpy.compat import long
55
from numpy.core import (array, arange, atleast_1d, atleast_2d, atleast_3d,
6-
vstack, hstack, newaxis, concatenate, stack)
7-
from numpy.testing import (TestCase, assert_, assert_raises, assert_array_equal,
8-
assert_equal, run_module_suite, assert_raises_regex)
6+
block, vstack, hstack, newaxis, concatenate, stack)
7+
from numpy.testing import (TestCase, assert_, assert_raises,
8+
assert_array_equal, assert_equal, run_module_suite,
9+
assert_raises_regex, assert_almost_equal)
10+
11+
from numpy.compat import long
912

1013
class TestAtleast1d(TestCase):
1114
def test_0D_array(self):
@@ -315,5 +318,71 @@ def test_stack():
315318
stack, [m, m])
316319

317320

321+
class TestBlock(TestCase):
322+
def test_block_row_wise(self):
323+
A = np.ones((2, 2))
324+
B = 2 * A
325+
assert_almost_equal(block([A, B]),
326+
np.array([[1, 1, 2, 2],
327+
[1, 1, 2, 2]]))
328+
# tuple notation
329+
assert_almost_equal(block((A, B)),
330+
np.array([[1, 1, 2, 2],
331+
[1, 1, 2, 2]]))
332+
333+
def test_block_column_wise(self):
334+
A = np.ones((2, 2))
335+
B = 2 * A
336+
assert_almost_equal(block([[A], [B]]),
337+
np.array([[1, 1],
338+
[1, 1],
339+
[2, 2],
340+
[2, 2]]))
341+
# tuple notation with only one element per tuple does not make much
342+
# sense. test it anyway just to make sure
343+
assert_almost_equal(block(((A, ), (B, ))),
344+
np.array([[1, 1],
345+
[1, 1],
346+
[2, 2],
347+
[2, 2]]))
348+
349+
def test_block_complex(self):
350+
# # # a bit more complex
351+
One = np.array([[1, 1, 1]])
352+
Two = np.array([[2, 2, 2]])
353+
Three = np.array([[3, 3, 3, 3, 3, 3]])
354+
four = np.array([4, 4, 4, 4, 4, 4])
355+
five = np.array([5])
356+
six = np.array([6, 6, 6, 6, 6])
357+
Zeros = np.zeros((2, 6))
358+
result = block([[One, Two],
359+
[Three],
360+
[four],
361+
[five, six],
362+
[Zeros]])
363+
expected = np.array([[1, 1, 1, 2, 2, 2],
364+
[3, 3, 3, 3, 3, 3],
365+
[4, 4, 4, 4, 4, 4],
366+
[5, 6, 6, 6, 6, 6],
367+
[0, 0, 0, 0, 0, 0],
368+
[0, 0, 0, 0, 0, 0]])
369+
assert_almost_equal(result, expected)
370+
371+
def test_block_with_1d_arrays(self):
372+
# # # 1-D vectors are treated as row arrays
373+
a = np.array([1, 2, 3])
374+
b = np.array([2, 3, 4])
375+
assert_almost_equal(block([a, b]),
376+
np.array([[1, 2, 3, 2, 3, 4]]))
377+
assert_almost_equal(block([[a], [b]]),
378+
np.array([[1, 2, 3],
379+
[2, 3, 4]]))
380+
a = np.array([1, 2, 3])
381+
b = np.array([2, 3, 4])
382+
assert_almost_equal(block([[a, b], [a, b]]),
383+
np.array([[1, 2, 3, 2, 3, 4],
+
[1, 2, 3, 2, 3, 4]]))
385+
386+
318387
if __name__ == "__main__":
319388
run_module_suite()

0 commit comments

Comments
 (0)
0