dask-array 0.1.0__py3-none-any.whl

dask_array/tests/test_transpose_slice_pushdown.py

@@ -0,0 +1,208 @@
+"""Tests for pushing integer slices through transpose operations."""
+
+import pytest
+
+import dask_array as da
+from dask_array._test_utils import assert_eq
+
+
+def test_transpose_integer_slice_2d():
+    """x.T[0] should optimize to x[:, 0] (transpose eliminated)."""
+    x = da.ones((3, 4), chunks=2)
+
+    # Naive expression
+    result = x.T[0]
+
+    # Expected: direct slice without transpose
+    expected = x[:, 0]
+
+    # Optimized naive should match expected structure
+    assert result.expr.optimize()._name == expected.expr.optimize()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_integer_slice_scalar():
+    """x.T[0, 0] should optimize to x[0, 0] (transpose eliminated)."""
+    x = da.ones((3, 4), chunks=2)
+
+    # Naive expression
+    result = x.T[0, 0]
+
+    # Expected: direct slice without transpose
+    expected = x[0, 0]
+
+    # Optimized naive should match expected structure
+    assert result.expr.optimize()._name == expected.expr.optimize()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_integer_slice_3d():
+    """x.T[0] on 3D should optimize to x[:, :, 0].T (slice pushed through)."""
+    x = da.ones((2, 3, 4), chunks=2)
+
+    # Naive expression: x.T[0] where x.T has shape (4, 3, 2)
+    result = x.T[0]
+
+    # Expected: slice first, then transpose remaining dims
+    # x[:, :, 0] has shape (2, 3), then transpose to (3, 2)
+    expected = x[:, :, 0].T
+
+    # Optimized naive should match expected structure
+    assert result.expr.optimize()._name == expected.expr.optimize()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_mixed_slice_integer():
+    """x.T[:, 0, :] should push integer through transpose."""
+    x = da.ones((3, 4, 5), chunks=2)
+
+    # Naive expression: x.T[:, 0, :] where x.T has shape (5, 4, 3)
+    result = x.T[:, 0, :]
+
+    # Expected: x[:, 0, :] then transpose remaining (5, 3) -> (5, 3) with axes (0, 1)
+    # Actually x.T[:, 0, :] selects middle dim of x.T which is dim 1 of x
+    # So: x[:, 0, :].transpose((1, 0))
+    expected = x[:, 0, :].transpose((1, 0))
+
+    # Optimized naive should match expected structure
+    assert result.expr.optimize()._name == expected.expr.optimize()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_custom_axes_integer_slice():
+    """Integer slice with custom transpose axes."""
+    x = da.ones((2, 3, 4), chunks=2)
+
+    # Naive: x.transpose((2, 0, 1))[0] - shape (4, 2, 3) -> slice dim 0 -> (2, 3)
+    result = x.transpose((2, 0, 1))[0]
+
+    # Expected: x[:, :, 0] then transpose remaining with reduced axes
+    # Original axes (2, 0, 1): out[0]=in[2], out[1]=in[0], out[2]=in[1]
+    # Slice out[0] (which is in[2]) -> remaining: out[1]=in[0], out[2]=in[1]
+    # After renumbering: axes (0, 1) = identity, so no transpose needed
+    expected = x[:, :, 0]
+
+    # Optimized naive should match expected structure
+    assert result.expr.optimize()._name == expected.expr.optimize()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_slice_task_count():
+    """Verify task count reduction when slice pushes through transpose."""
+    from dask_array._collection import Array
+
+    x = da.ones((4, 6), chunks=2)
+
+    # Without optimization: slice(transpose(ones)) has transpose layer
+    result = x.T[0]
+    unopt_graph = dict(result.__dask_graph__())
+
+    # With optimization: transpose is eliminated, becomes slice(ones)
+    optimized = result.expr.optimize()
+    opt_result = Array(optimized)
+    opt_graph = dict(opt_result.__dask_graph__())
+
+    # Optimized should have fewer tasks (no transpose layer)
+    assert len(opt_graph) < len(unopt_graph), (
+        f"Optimized graph should be smaller: {len(opt_graph)} vs {len(unopt_graph)}"
+    )
+
+    # Specifically: unoptimized has ones(6) + transpose(6) + getitem(2) = 14
+    # Optimized has ones(6) + getitem(2) = 8 (transpose eliminated)
+    assert len(opt_graph) == 8, f"Expected 8 tasks (6 ones + 2 getitem), got {len(opt_graph)}"
+    assert len(unopt_graph) == 14, f"Expected 14 unoptimized tasks, got {len(unopt_graph)}"
+
+
+# --- Transpose through Elemwise Tests ---
+
+
+def test_transpose_pushes_through_elemwise_add():
+    """(x + y).T should optimize to x.T + y.T."""
+    x = da.ones((10, 5), chunks=5)
+    y = da.ones((10, 5), chunks=5)
+
+    result = (x + y).T
+    expected = x.T + y.T
+
+    # Simplified should match expected structure
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+
+    # Verify correctness
+    assert_eq(result, expected)
+
+
+def test_transpose_pushes_through_elemwise_mul():
+    """(x * y).T should optimize to x.T * y.T."""
+    x = da.ones((6, 4), chunks=2)
+    y = da.ones((6, 4), chunks=2)
+
+    result = (x * y).T
+    expected = x.T * y.T
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, expected)
+
+
+def test_transpose_through_elemwise_broadcasting_no_pushdown():
+    """Transpose doesn't push through elemwise when broadcasting (different ndim)."""
+    from dask_array.manipulation._transpose import Transpose
+
+    x = da.ones((6, 4), chunks=2)
+    y = da.ones((4,), chunks=2)  # broadcasts along axis 1
+
+    result = (x + y).T  # (6, 4) + (4,) -> (6, 4), then .T -> (4, 6)
+
+    # We don't push through broadcasting cases, so outer op is still Transpose
+    opt = result.expr.simplify()
+    assert isinstance(opt, Transpose)
+
+    # But result is still correct
+    import numpy as np
+
+    x_np = np.ones((6, 4))
+    y_np = np.ones((4,))
+    assert_eq(result, (x_np + y_np).T)
+
+
+def test_transpose_pushes_through_elemwise_scalar():
+    """Transpose through elemwise with scalar."""
+    x = da.ones((5, 3), chunks=2)
+
+    result = (x + 1).T
+    expected = x.T + 1
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, expected)
+
+
+def test_transpose_pushes_through_unary_elemwise():
+    """Transpose through unary elemwise (e.g. negative)."""
+    x = da.ones((4, 6), chunks=2)
+
+    result = (-x).T
+    expected = -(x.T)
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, expected)
+
+
+def test_transpose_custom_axes_through_elemwise():
+    """Custom transpose axes through elemwise."""
+    x = da.ones((2, 3, 4), chunks=2)
+    y = da.ones((2, 3, 4), chunks=2)
+
+    result = (x + y).transpose((2, 0, 1))
+    expected = x.transpose((2, 0, 1)) + y.transpose((2, 0, 1))
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, expected)

dask_array/tests/test_visualize.py

@@ -0,0 +1,94 @@
+"""Tests for expression visualization."""
+
+import dask_array as da
+from dask_array._visualize import expr_table
+
+
+def test_expr_table_contains_shapes():
+    """Test that expr_table output contains array shapes."""
+    x = da.ones((10, 100), chunks=(5, 50))
+    y = x.sum()
+
+    table = expr_table(y.expr)
+    text = str(table)
+
+    # Should contain the input shape
+    assert "(10, 100)" in text
+    # Should contain scalar output shape
+    assert "()" in text
+
+
+def test_expr_table_contains_bytes():
+    """Test that expr_table output contains byte sizes."""
+    x = da.ones((10, 100), chunks=(5, 50))
+    y = x.sum()
+
+    table = expr_table(y.expr)
+    text = str(table)
+
+    # 1000 float64 elements = 8000 bytes = 7.8 kiB
+    assert "kiB" in text
+    # Scalar output = 8 bytes
+    assert "8 B" in text
+
+
+def test_expr_table_contains_operation_names():
+    """Test that expr_table shows operation names."""
+    x = da.ones((10, 10), chunks=5)
+    y = x + 1
+    z = y.sum()
+
+    table = expr_table(z.expr)
+    text = str(table)
+
+    assert "Reduction" in text
+    assert "Ones" in text
+
+
+def test_expr_table_styling_emphasis():
+    """Test that large arrays are bold and small arrays are dim."""
+    x = da.ones((100, 100), chunks=50)
+    y = x.sum()  # Reduces to scalar
+
+    table = expr_table(y.expr)
+    rich_table = table._table
+
+    operation_cells = rich_table.columns[0]._cells
+    shape_cells = rich_table.columns[1]._cells
+    bytes_cells = rich_table.columns[2]._cells
+
+    # The scalar reduction row is small relative to the source and should dim
+    # non-operation data. Operation names should remain emphasized.
+    assert "bold" in operation_cells[0].spans[0].style
+    assert shape_cells[0].style == "dim"
+    assert bytes_cells[0].style == "dim"
+
+    # The large source row should stay visually emphasized.
+    assert any("bold" in span.style for span in operation_cells[1].spans)
+    assert shape_cells[1].style is None
+    assert bytes_cells[1].style is None
+
+
+def test_expr_table_html_output():
+    """Test that HTML output is generated for Jupyter."""
+    x = da.ones((10, 10), chunks=5)
+    y = x.sum()
+
+    table = expr_table(y.expr)
+    html = table._repr_html_()
+
+    assert "<pre>" in html
+    assert "Reduction" in html
+    assert "Ones" in html
+
+
+def test_expr_repr_html():
+    """Test that ArrayExpr._repr_html_ works for Jupyter display."""
+    x = da.ones((10, 10), chunks=5)
+    y = x.sum()
+
+    # The expression should have _repr_html_ for Jupyter
+    html = y.expr._repr_html_()
+
+    assert "<pre>" in html
+    assert "Reduction" in html

dask_array/tests/test_xarray.py

@@ -0,0 +1,193 @@
+"""Tests for xarray ChunkManager integration.
+
+This module tests the DaskArrayExprManager which registers as the "dask"
+chunk manager, replacing xarray's built-in DaskManager. This allows it to
+handle dask_array.Array types.
+"""
+
+import numpy as np
+import pytest
+
+xr = pytest.importorskip("xarray")
+
+import dask_array as da
+from dask_array._xarray import DaskArrayExprManager
+
+
+class TestDaskArrayExprManager:
+    """Tests for the DaskArrayExprManager class."""
+
+    def test_init(self):
+        manager = DaskArrayExprManager()
+        assert manager.array_cls is da.Array
+        assert manager.available is True
+
+    def test_is_chunked_array(self):
+        manager = DaskArrayExprManager()
+        arr = da.ones((10, 10), chunks=(5, 5))
+        assert manager.is_chunked_array(arr)
+        assert not manager.is_chunked_array(np.ones((10, 10)))
+
+    def test_is_chunked_array_legacy_dask(self):
+        """Test that manager rejects legacy dask.array.Array."""
+        import dask.array as legacy_da
+
+        manager = DaskArrayExprManager()
+        arr = legacy_da.ones((10, 10), chunks=(5, 5))
+        assert not manager.is_chunked_array(arr)
+
+    def test_chunks(self):
+        manager = DaskArrayExprManager()
+        arr = da.ones((10, 10), chunks=(5, 5))
+        assert manager.chunks(arr) == ((5, 5), (5, 5))
+
+    def test_normalize_chunks(self):
+        manager = DaskArrayExprManager()
+        result = manager.normalize_chunks((5, 5), shape=(10, 10))
+        assert result == ((5, 5), (5, 5))
+
+    def test_from_array(self):
+        manager = DaskArrayExprManager()
+        x = np.arange(100).reshape(10, 10)
+        arr = manager.from_array(x, chunks=(5, 5))
+        assert isinstance(arr, da.Array)
+        assert arr.chunks == ((5, 5), (5, 5))
+
+    def test_from_array_lazy_indexing_adapter_uses_numpy_meta(self):
+        from xarray.core.indexing import (
+            ImplicitToExplicitIndexingAdapter,
+            LazilyIndexedArray,
+            NumpyIndexingAdapter,
+            OuterIndexer,
+        )
+
+        manager = DaskArrayExprManager()
+        base = NumpyIndexingAdapter(np.ones((2, 3)))
+        lazy = LazilyIndexedArray(base, OuterIndexer((slice(None), slice(None))))
+        adapter = ImplicitToExplicitIndexingAdapter(lazy, OuterIndexer)
+
+        arr = manager.from_array(adapter, chunks=(1, 3))
+        out = arr * 1.0
+
+        assert isinstance(arr.expr._meta, np.ndarray)
+        assert isinstance(out.expr._meta, np.ndarray)
+        assert out.expr._meta.shape == (0, 0)
+
+    def test_compute(self):
+        manager = DaskArrayExprManager()
+        arr = da.ones((10, 10), chunks=(5, 5))
+        (result,) = manager.compute(arr)
+        np.testing.assert_array_equal(result, np.ones((10, 10)))
+
+    def test_array_api(self):
+        manager = DaskArrayExprManager()
+        api = manager.array_api
+        assert hasattr(api, "ones")
+        assert hasattr(api, "zeros")
+        assert hasattr(api, "full_like")
+
+
+class TestXarrayIntegration:
+    """Tests for xarray integration with DaskArrayExprManager."""
+
+    def test_manager_discoverable(self):
+        """Test that the manager is discoverable via xarray as 'dask'."""
+        from xarray.namedarray.parallelcompat import list_chunkmanagers
+
+        managers = list_chunkmanagers()
+        assert "dask" in managers
+        # Our manager should be the one registered (replaces built-in)
+        assert isinstance(managers["dask"], DaskArrayExprManager)
+
+    def test_get_chunked_array_type_selects_manager_once(self):
+        """Test xarray sees dask_array.Array through one chunk manager."""
+        from xarray.namedarray.parallelcompat import get_chunked_array_type
+
+        arr = da.ones((10, 10), chunks=(5, 5))
+
+        assert isinstance(get_chunked_array_type(arr), DaskArrayExprManager)
+
+    def test_dask_new_collection_roundtrip(self):
+        """Test Dask can rebuild dask_array.Array from its expression."""
+        from dask._collections import new_collection
+
+        arr = da.arange(6, chunks=(3,)) + 1
+
+        rebuilt = new_collection(arr.expr)
+
+        assert isinstance(rebuilt, da.Array)
+        np.testing.assert_array_equal(rebuilt.compute(), np.arange(6) + 1)
+
+    def test_dataarray_from_dask_array(self):
+        """Test creating a DataArray from a dask_array.Array."""
+        arr = da.ones((10, 20), chunks=(5, 10))
+        da_xr = xr.DataArray(arr, dims=["x", "y"])
+
+        assert da_xr.shape == (10, 20)
+        assert da_xr.chunks == ((5, 5), (10, 10))
+
+    def test_dataarray_compute(self):
+        """Test computing a DataArray backed by dask_array.Array."""
+        arr = da.arange(100, chunks=25).reshape(10, 10)
+        da_xr = xr.DataArray(arr, dims=["x", "y"])
+
+        result = da_xr.compute()
+        expected = np.arange(100).reshape(10, 10)
+        np.testing.assert_array_equal(result.values, expected)
+
+    def test_dataarray_operations(self):
+        """Test that DataArray operations work with dask_array.Array."""
+        arr = da.ones((10, 20), chunks=(5, 10))
+        da_xr = xr.DataArray(arr, dims=["x", "y"])
+
+        # Test arithmetic
+        result = (da_xr + 1).compute()
+        np.testing.assert_array_equal(result.values, np.full((10, 20), 2.0))
+
+        # Test reduction
+        result = da_xr.sum().compute()
+        assert result.values == 200.0
+
+    def test_dataarray_rechunk(self):
+        """Test rechunking a DataArray."""
+        arr = da.ones((10, 20), chunks=(5, 10))
+        da_xr = xr.DataArray(arr, dims=["x", "y"])
+
+        rechunked = da_xr.chunk({"x": 10, "y": 5})
+        assert rechunked.chunks == ((10,), (5, 5, 5, 5))
+
+    def test_dataset_from_dask_arrays(self):
+        """Test creating a Dataset from dask_array.Arrays."""
+        arr1 = da.ones((10, 20), chunks=(5, 10))
+        arr2 = da.zeros((10, 20), chunks=(5, 10))
+
+        ds = xr.Dataset(
+            {
+                "var1": (["x", "y"], arr1),
+                "var2": (["x", "y"], arr2),
+            }
+        )
+
+        assert ds["var1"].shape == (10, 20)
+        assert ds["var2"].shape == (10, 20)
+
+        result = ds.compute()
+        np.testing.assert_array_equal(result["var1"].values, np.ones((10, 20)))
+        np.testing.assert_array_equal(result["var2"].values, np.zeros((10, 20)))
+
+    def test_apply_ufunc(self):
+        """Test xr.apply_ufunc with dask_array.Array."""
+        arr = da.ones((10, 20), chunks=(5, 10))
+        da_xr = xr.DataArray(arr, dims=["x", "y"])
+
+        result = xr.apply_ufunc(
+            lambda x: x * 2,
+            da_xr,
+            dask="parallelized",
+            output_dtypes=[float],
+        )
+
+        assert isinstance(result.data, da.Array)
+
+        computed = result.compute()
+        np.testing.assert_array_equal(computed.values, np.full((10, 20), 2.0))

dask_array-0.1.0.dist-info/METADATA

@@ -0,0 +1,48 @@
+Metadata-Version: 2.4
+Name: dask-array
+Version: 0.1.0
+Summary: Parallel arrays with task graphs
+Project-URL: Homepage, https://github.com/mrocklin/dask-array
+Project-URL: Source, https://github.com/mrocklin/dask-array
+Project-URL: Issues, https://github.com/mrocklin/dask-array/issues
+License-Expression: BSD-3-Clause
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: dask>=2024.1.0
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: numpy!=2.4.0,>=2.0.0
+Requires-Dist: rich
+Requires-Dist: toolz>=0.8.2
+Provides-Extra: complete
+Requires-Dist: h5py; extra == 'complete'
+Requires-Dist: scipy; extra == 'complete'
+Requires-Dist: tiledb; extra == 'complete'
+Requires-Dist: zarr>=2.0; extra == 'complete'
+Provides-Extra: sparse
+Requires-Dist: numba>=0.60; (python_version < '3.13') and extra == 'sparse'
+Requires-Dist: sparse; (python_version < '3.13') and extra == 'sparse'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: pytest-cov; extra == 'test'
+Requires-Dist: pytest-xdist; extra == 'test'
+Requires-Dist: xarray; extra == 'test'
+Description-Content-Type: text/markdown
+
+# dask-array
+
+Expression-based dask array implementation.
+
+## Installation
+
+```bash
+pip install dask-array
+```
+
+## Usage
+
+```python
+import dask_array as da
+
+x = da.ones((1000, 1000), chunks=(100, 100))
+result = x.sum().compute()
+```