| 593 | |
| 594 | |
| 595 | def test_stack(): |
| 596 | a, b, c = ( |
| 597 | Array( |
| 598 | graph_from_arraylike(object(), chunks=(2, 3), shape=(4, 6), name=name), |
| 599 | name, |
| 600 | chunks=(2, 3), |
| 601 | dtype="f8", |
| 602 | shape=(4, 6), |
| 603 | ) |
| 604 | for name in "ABC" |
| 605 | ) |
| 606 | |
| 607 | s = stack([a, b, c], axis=0) |
| 608 | |
| 609 | colon = slice(None, None, None) |
| 610 | |
| 611 | assert s.shape == (3, 4, 6) |
| 612 | assert s.chunks == ((1, 1, 1), (2, 2), (3, 3)) |
| 613 | assert s.chunksize == (1, 2, 3) |
| 614 | assert s.dask[(s.name, 0, 1, 0)] == (getitem, ("A", 1, 0), (None, colon, colon)) |
| 615 | assert s.dask[(s.name, 2, 1, 0)] == (getitem, ("C", 1, 0), (None, colon, colon)) |
| 616 | assert same_keys(s, stack([a, b, c], axis=0)) |
| 617 | |
| 618 | s2 = stack([a, b, c], axis=1) |
| 619 | assert s2.shape == (4, 3, 6) |
| 620 | assert s2.chunks == ((2, 2), (1, 1, 1), (3, 3)) |
| 621 | assert s2.chunksize == (2, 1, 3) |
| 622 | assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ("B", 0, 0), (colon, None, colon)) |
| 623 | assert s2.dask[(s2.name, 1, 1, 0)] == (getitem, ("B", 1, 0), (colon, None, colon)) |
| 624 | assert same_keys(s2, stack([a, b, c], axis=1)) |
| 625 | |
| 626 | s2 = stack([a, b, c], axis=2) |
| 627 | assert s2.shape == (4, 6, 3) |
| 628 | assert s2.chunks == ((2, 2), (3, 3), (1, 1, 1)) |
| 629 | assert s2.chunksize == (2, 3, 1) |
| 630 | assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ("A", 0, 1), (colon, colon, None)) |
| 631 | assert s2.dask[(s2.name, 1, 1, 2)] == (getitem, ("C", 1, 1), (colon, colon, None)) |
| 632 | assert same_keys(s2, stack([a, b, c], axis=2)) |
| 633 | |
| 634 | pytest.raises(ValueError, lambda: stack([])) |
| 635 | pytest.raises(ValueError, lambda: stack([a, b, c], axis=3)) |
| 636 | |
| 637 | assert set(b.dask.keys()).issubset(s2.dask.keys()) |
| 638 | |
| 639 | assert stack([a, b, c], axis=-1).chunks == stack([a, b, c], axis=2).chunks |
| 640 | |
| 641 | |
| 642 | def test_stack_zero_size(): |