|
1 | 1 | from __future__ import division, absolute_import, print_function
|
2 | 2 |
|
3 |
| -__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'vstack', 'hstack'] |
| 3 | +__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'vstack', 'hstack', |
| 4 | + 'stack'] |
4 | 5 |
|
5 | 6 | from . import numeric as _nx
|
6 | 7 | from .numeric import array, asanyarray, newaxis
|
@@ -275,3 +276,87 @@ def hstack(tup):
|
275 | 276 | return _nx.concatenate(arrs, 0)
|
276 | 277 | else:
|
277 | 278 | return _nx.concatenate(arrs, 1)
|
| 279 | + |
| 280 | + |
| 281 | +def stack(tup_tup): |
| 282 | + """ |
| 283 | + Stack arrays similar to matlab's "square bracket stacking". |
| 284 | +
|
| 285 | + [A A; B B] |
| 286 | +
|
| 287 | + Parameters |
| 288 | + ---------- |
| 289 | + tup_tup : sequence of sequence of ndarrays |
| 290 | + 1-D arrays are treated as row vectors. |
| 291 | +
|
| 292 | + Returns |
| 293 | + ------- |
| 294 | + stacked : ndarray |
| 295 | + The 2-D array formed by stacking the given arrays. |
| 296 | +
|
| 297 | + See Also |
| 298 | + -------- |
| 299 | + hstack : Stack arrays in sequence horizontally (column wise). |
| 300 | + vstack : Stack arrays in sequence vertically (row wise). |
| 301 | + dstack : Stack arrays in sequence depth wise (along third dimension). |
| 302 | + concatenate : Join a sequence of arrays together. |
| 303 | + vsplit : Split array into a list of multiple sub-arrays vertically. |
| 304 | +
|
| 305 | + Examples |
| 306 | + -------- |
| 307 | + >>> A = np.array([[1, 2, 3]]) |
| 308 | + >>> B = np.array([[2, 3, 4]]) |
| 309 | + >>> stack([A, B]) |
| 310 | + array([[1, 2, 3, 2, 3, 4]]) |
| 311 | +
|
| 312 | + >>> A = np.array([[1, 2, 3]]).T |
| 313 | + >>> B = np.array([[2, 3, 4]]).T |
| 314 | + >>> stack([A, B]) |
| 315 | + array([[1, 2], |
| 316 | + [2, 3], |
| 317 | + [3, 4]]) |
| 318 | +
|
| 319 | + >>> A = np.array([[1, 2, 3]]) |
| 320 | + >>> B = np.array([[2, 3, 4]]) |
| 321 | + >>> stack([[A, A], [B, B]]) |
| 322 | + array([[1, 2, 3, 1, 2, 3], |
| 323 | + [2, 3, 4, 2, 3, 4]]) |
| 324 | +
|
| 325 | + >>> A = np.array([[1, 2, 3]]) |
| 326 | + >>> B = np.array([[2, 3, 4]]) |
| 327 | + >>> stack([[A], [B]]) |
| 328 | + array([[1, 2, 3], |
| 329 | + [2, 3, 4]]) |
| 330 | +
|
| 331 | + >>> # 1-D vectors are treated as row arrays |
| 332 | + >>> a = np.array([1, 2, 3]) |
| 333 | + >>> b = np.array([2, 3, 4]) |
| 334 | + >>> stack([a, b]) |
| 335 | + array([[1, 2, 3, 2, 3, 4]]) |
| 336 | +
|
| 337 | + >>> # 1-D vectors are treated as row arrays |
| 338 | + >>> a = np.array([1, 2, 3]) |
| 339 | + >>> b = np.array([2, 3, 4]) |
| 340 | + >>> stack([[a, b], [a, b]]) |
| 341 | + array([[1, 2, 3, 2, 3, 4], |
| 342 | + [1, 2, 3, 2, 3, 4]]) |
| 343 | +
|
| 344 | + >>> # a bit more complex |
| 345 | + >>> A = np.array([[1, 2, 3]]) |
| 346 | + >>> B = np.array([[2, 3, 4, 5, 6, 7]]) |
| 347 | + >>> c = np.array([8, 3, 4, 5, 6, 7]) |
| 348 | + >>> d1 = np.array([9]) |
| 349 | + >>> d2 = np.array([8, 7, 6, 5, 4]) |
| 350 | + >>> stack([[A, A], [B], [c], [d1, d2]]) |
| 351 | + array([[1, 2, 3, 1, 2, 3], |
| 352 | + [2, 3, 4, 5, 6, 7], |
| 353 | + [8, 3, 4, 5, 6, 7], |
| 354 | + [9, 8, 7, 6, 5, 4]]) |
| 355 | +
|
| 356 | +
|
| 357 | + """ |
| 358 | + if isinstance(tup_tup[0], list) or isinstance(tup_tup[0], tuple): |
| 359 | + result = vstack([hstack([a for a in row]) for row in tup_tup]) |
| 360 | + else: |
| 361 | + result = hstack([a for a in tup_tup]) |
| 362 | + return atleast_2d(result) |
0 commit comments