dask-array 0.1.0__py3-none-any.whl

dask_array/linalg/_utils.py

@@ -0,0 +1,74 @@
+"""Shared utilities for linear algebra operations."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def _cumsum_blocks(it):
+    """Yield cumulative (start, end) pairs for block sizes."""
+    total = 0
+    for x in it:
+        total_previous = total
+        total += x
+        yield (total_previous, total)
+
+
+def _nanmin(m, n):
+    """Return min(m, n), handling NaN values.
+
+    If either value is NaN, return the other value if it's not NaN,
+    otherwise return NaN.
+    """
+    k_0 = min([m, n])
+    k_1 = m if np.isnan(n) else n
+    return k_1 if np.isnan(k_0) else k_0
+
+
+def _has_uncertain_chunks(chunks):
+    """Check if any chunk sizes are uncertain (NaN)."""
+    return any(np.isnan(c) for cs in chunks for c in cs)
+
+
+def _cumsum_part(last, new):
+    """Compute (start, end) from previous cumsum and new value."""
+    return (last[1], last[1] + new)
+
+
+def _get_block_size(block):
+    """Get min dimension from a block (used for R block sizes)."""
+    return min(block.shape)
+
+
+def _make_slice(cumsum_pair, n):
+    """Create slice tuple from cumsum pair and column count."""
+    return (slice(cumsum_pair[0], cumsum_pair[1]), slice(0, n))
+
+
+def _getitem_with_slice(arr, slc):
+    """Apply slice to array."""
+    return arr[slc]
+
+
+def _get_n(arr):
+    """Get the number of columns from an array."""
+    return arr.shape[1]
+
+
+def _solve_triangular_lower(a, b):
+    """Solve triangular system with lower triangular matrix."""
+    from dask_array._utils import solve_triangular_safe
+
+    return solve_triangular_safe(a, b, lower=True)
+
+
+def _solve_triangular_upper(a, b):
+    """Solve triangular system with upper triangular matrix."""
+    from dask_array._utils import solve_triangular_safe
+
+    return solve_triangular_safe(a, b, lower=False)
+
+
+def _transpose(x):
+    """Transpose a matrix (used for P_inv and U.T)."""
+    return x.T

dask_array/manipulation/__init__.py

@@ -0,0 +1,45 @@
+"""Array manipulation functions: flip, transpose, reshape, expand_dims, etc."""
+
+# Import from module files
+from dask_array.manipulation._expand import (
+    atleast_1d,
+    atleast_2d,
+    atleast_3d,
+    expand_dims,
+)
+from dask_array.manipulation._flip import flip, fliplr, flipud, rot90
+from dask_array.manipulation._roll import roll
+from dask_array.manipulation._transpose import (
+    moveaxis,
+    rollaxis,
+    swapaxes,
+    transpose,
+)
+
+
+def __getattr__(name):
+    """Lazy import of reshape and ravel to avoid circular imports."""
+    if name in ("reshape", "ravel"):
+        from dask_array import _collection
+
+        return getattr(_collection, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "flip",
+    "flipud",
+    "fliplr",
+    "rot90",
+    "swapaxes",
+    "moveaxis",
+    "rollaxis",
+    "transpose",
+    "expand_dims",
+    "atleast_1d",
+    "atleast_2d",
+    "atleast_3d",
+    "roll",
+    "reshape",
+    "ravel",
+]

dask_array/manipulation/_expand.py

@@ -0,0 +1,321 @@
+"""Expand operations: expand_dims, atleast_1d, atleast_2d, atleast_3d."""
+
+from __future__ import annotations
+
+import functools
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task, TaskRef
+from dask_array._expr import ArrayExpr
+from dask_array._chunk import getitem
+
+
+class ExpandDims(ArrayExpr):
+    """Dimension expansion expression.
+
+    Adds new axes of size 1 at specified positions using numpy indexing with None.
+    This is more efficient than reshape for dimension expansion and integrates
+    better with slice pushdown optimizations.
+
+    Parameters
+    ----------
+    array : ArrayExpr
+        The input array expression.
+    axes : tuple of int
+        Positions where new axes should be inserted (in output coordinates).
+    """
+
+    _parameters = ["array", "axes"]
+
+    @functools.cached_property
+    def _meta(self):
+        meta = self.array._meta
+        for ax in sorted(self.axes):
+            meta = np.expand_dims(meta, axis=ax)
+        return meta
+
+    @functools.cached_property
+    def chunks(self):
+        chunks = list(self.array.chunks)
+        for ax in sorted(self.axes):
+            chunks.insert(ax, (1,))
+        return tuple(chunks)
+
+    @functools.cached_property
+    def _name(self):
+        return f"expand-dims-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _indexer(self):
+        """Build indexer tuple with None at expansion axes."""
+        out_ndim = self.array.ndim + len(self.axes)
+        return tuple(None if i in self.axes else slice(None) for i in range(out_ndim))
+
+    def _layer(self) -> dict:
+        indexer = self._indexer
+        axes = sorted(self.axes)
+        input_name = self.array._name
+
+        dsk = {}
+        for block_id in np.ndindex(self.array.numblocks):
+            in_key = (input_name,) + block_id
+            # Insert 0 at expansion axes: ('x', 2, 3) -> ('expand', 0, 2, 0, 3)
+            out_block_id = list(block_id)
+            for ax in axes:
+                out_block_id.insert(ax, 0)
+            out_key = (self._name,) + tuple(out_block_id)
+            dsk[out_key] = Task(out_key, getitem, TaskRef(in_key), indexer)
+
+        return dsk
+
+    def _simplify_up(self, parent, dependents):
+        """Allow slice and shuffle operations to push through ExpandDims."""
+        from dask_array._shuffle import Shuffle
+        from dask_array.slicing import SliceSlicesIntegers
+
+        if isinstance(parent, SliceSlicesIntegers):
+            return self._accept_slice(parent)
+        if isinstance(parent, Shuffle):
+            return self._accept_shuffle(parent)
+        return None
+
+    def _accept_slice(self, slice_expr):
+        """Accept a slice being pushed through ExpandDims.
+
+        Maps output slice indices to input, removing expanded axes.
+        """
+        from numbers import Integral
+
+        axes = set(self.axes)
+        index = slice_expr.index
+        out_ndim = self.ndim
+
+        # Pad index to full output length
+        full_index = index + (slice(None),) * (out_ndim - len(index))
+
+        # Build input slice by removing expanded axes
+        input_index = []
+        dims_removed = 0
+        for i, idx in enumerate(full_index):
+            if i in axes:
+                # This is an expanded axis (size 1)
+                if isinstance(idx, Integral):
+                    # Integer index on size-1 axis must be 0 or -1
+                    if idx not in (0, -1):
+                        return None  # Out of bounds
+                    dims_removed += 1
+                elif idx == slice(None):
+                    pass  # Keep this expansion axis
+                elif isinstance(idx, slice):
+                    # Slicing a size-1 axis - normalize
+                    start, stop, step = idx.indices(1)
+                    if stop <= start:
+                        return None  # Empty result
+                else:
+                    return None  # Can't handle this
+            else:
+                input_index.append(idx)
+
+        # Slice the input
+        if all(idx == slice(None) for idx in input_index):
+            sliced_input = self.array
+        else:
+            sliced_input = new_collection(self.array)[tuple(input_index)].expr
+
+        # Compute new axes positions after slicing
+        # Axes that had integer indexing are removed
+        new_axes = []
+        removed_count = 0
+        for i, idx in enumerate(full_index):
+            if i in axes:
+                if isinstance(idx, Integral):
+                    removed_count += 1
+                else:
+                    # Adjust for removed dimensions before this position
+                    new_axes.append(i - removed_count)
+
+        if not new_axes:
+            # All expansion axes were removed by integer indexing
+            return sliced_input
+
+        return ExpandDims(sliced_input, tuple(new_axes))
+
+    def _accept_shuffle(self, shuffle_expr):
+        """Accept a shuffle being pushed through ExpandDims.
+
+        Maps shuffle axis through expansion to input axis.
+        """
+        from dask_array._shuffle import Shuffle
+
+        axes = set(self.axes)
+        shuffle_axis = shuffle_expr.axis
+
+        # Can't shuffle on an expanded axis (size 1, nothing to shuffle)
+        if shuffle_axis in axes:
+            return None
+
+        # Map output axis to input axis (subtract number of expansion axes before it)
+        input_axis = shuffle_axis - sum(1 for ax in axes if ax < shuffle_axis)
+
+        shuffled_input = Shuffle(
+            self.array,
+            shuffle_expr.indexer,
+            input_axis,
+            shuffle_expr.operand("name"),
+        )
+        return ExpandDims(shuffled_input, self.axes)
+
+
+def expand_dims(a, axis):
+    """Expand the shape of an array.
+
+    Insert a new axis that will appear at the axis position in the expanded
+    array shape.
+
+    Parameters
+    ----------
+    a : array_like
+        Input array.
+    axis : int or tuple of ints
+        Position in the expanded axes where the new axis (or axes) is placed.
+
+    Returns
+    -------
+    result : Array
+        Array with the number of dimensions increased.
+
+    See Also
+    --------
+    numpy.expand_dims
+    """
+    from dask_array._utils import validate_axis
+
+    if axis is None:
+        raise TypeError("axis must be an integer, not None")
+
+    if type(axis) not in (tuple, list):
+        axis = (axis,)
+
+    out_ndim = len(axis) + a.ndim
+    axis = validate_axis(axis, out_ndim)
+
+    return new_collection(ExpandDims(a.expr, tuple(sorted(axis))))
+
+
+def atleast_1d(*arys):
+    """Convert inputs to arrays with at least one dimension.
+
+    Parameters
+    ----------
+    arys : array_like
+        One or more array-like sequences. Non-array inputs are converted
+        to arrays. Arrays that already have one or more dimensions are
+        preserved.
+
+    Returns
+    -------
+    ret : Array or tuple of Arrays
+        An array, or tuple of arrays, each with a.ndim >= 1.
+
+    See Also
+    --------
+    numpy.atleast_1d
+    """
+    from dask_array.core import asanyarray
+    from dask_array._numpy_compat import NUMPY_GE_200
+
+    new_arys = []
+    for x in arys:
+        x = asanyarray(x)
+        if x.ndim == 0:
+            x = x[None]
+        new_arys.append(x)
+
+    if len(new_arys) == 1:
+        return new_arys[0]
+    else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
+        return new_arys
+
+
+def atleast_2d(*arys):
+    """View inputs as arrays with at least two dimensions.
+
+    Parameters
+    ----------
+    arys : array_like
+        One or more array-like sequences. Non-array inputs are converted
+        to arrays. Arrays that already have two or more dimensions are
+        preserved.
+
+    Returns
+    -------
+    ret : Array or tuple of Arrays
+        An array, or tuple of arrays, each with a.ndim >= 2.
+
+    See Also
+    --------
+    numpy.atleast_2d
+    """
+    from dask_array.core import asanyarray
+    from dask_array._numpy_compat import NUMPY_GE_200
+
+    new_arys = []
+    for x in arys:
+        x = asanyarray(x)
+        if x.ndim == 0:
+            x = x[None, None]
+        elif x.ndim == 1:
+            x = x[None, :]
+        new_arys.append(x)
+
+    if len(new_arys) == 1:
+        return new_arys[0]
+    else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
+        return new_arys
+
+
+def atleast_3d(*arys):
+    """View inputs as arrays with at least three dimensions.
+
+    Parameters
+    ----------
+    arys : array_like
+        One or more array-like sequences. Non-array inputs are converted
+        to arrays. Arrays that already have three or more dimensions are
+        preserved.
+
+    Returns
+    -------
+    ret : Array or tuple of Arrays
+        An array, or tuple of arrays, each with a.ndim >= 3.
+
+    See Also
+    --------
+    numpy.atleast_3d
+    """
+    from dask_array.core import asanyarray
+    from dask_array._numpy_compat import NUMPY_GE_200
+
+    new_arys = []
+    for x in arys:
+        x = asanyarray(x)
+        if x.ndim == 0:
+            x = x[None, None, None]
+        elif x.ndim == 1:
+            x = x[None, :, None]
+        elif x.ndim == 2:
+            x = x[:, :, None]
+        new_arys.append(x)
+
+    if len(new_arys) == 1:
+        return new_arys[0]
+    else:
+        if NUMPY_GE_200:
+            new_arys = tuple(new_arys)
+        return new_arys

dask_array/manipulation/_flip.py

@@ -0,0 +1,92 @@
+"""Flip operations: flip, flipud, fliplr, rot90."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+
+
+def flip(m, axis=None):
+    """Reverse element order along axis.
+
+    See Also
+    --------
+    numpy.flip
+    """
+    from dask_array.core import asanyarray
+
+    m = asanyarray(m)
+
+    sl = m.ndim * [slice(None)]
+    if axis is None:
+        axis = range(m.ndim)
+    if not isinstance(axis, Iterable):
+        axis = (axis,)
+    try:
+        for ax in axis:
+            sl[ax] = slice(None, None, -1)
+    except IndexError as e:
+        raise ValueError(f"`axis` of {axis} invalid for {m.ndim}-D array") from e
+    sl = tuple(sl)
+
+    return m[sl]
+
+
+def flipud(m):
+    """Flip array in the up/down direction.
+
+    See Also
+    --------
+    numpy.flipud
+    """
+    return flip(m, 0)
+
+
+def fliplr(m):
+    """Flip array in the left/right direction.
+
+    See Also
+    --------
+    numpy.fliplr
+    """
+    return flip(m, 1)
+
+
+def rot90(m, k=1, axes=(0, 1)):
+    """Rotate an array by 90 degrees in the plane specified by axes.
+
+    See Also
+    --------
+    numpy.rot90
+    """
+    from dask_array.core import asanyarray
+    from dask_array.manipulation._transpose import transpose
+
+    axes = tuple(axes)
+    if len(axes) != 2:
+        raise ValueError("len(axes) must be 2.")
+
+    m = asanyarray(m)
+
+    if axes[0] == axes[1] or np.absolute(axes[0] - axes[1]) == m.ndim:
+        raise ValueError("Axes must be different.")
+
+    if axes[0] >= m.ndim or axes[0] < -m.ndim or axes[1] >= m.ndim or axes[1] < -m.ndim:
+        raise ValueError(f"Axes={axes} out of range for array of ndim={m.ndim}.")
+
+    k %= 4
+
+    if k == 0:
+        return m[:]
+    if k == 2:
+        return flip(flip(m, axes[0]), axes[1])
+
+    axes_list = list(range(0, m.ndim))
+    (axes_list[axes[0]], axes_list[axes[1]]) = (axes_list[axes[1]], axes_list[axes[0]])
+
+    if k == 1:
+        return transpose(flip(m, axes[1]), axes_list)
+    else:
+        # k == 3
+        return flip(transpose(m, axes_list), axes[1])

dask_array/manipulation/_roll.py

@@ -0,0 +1,78 @@
+"""Roll operation."""
+
+from __future__ import annotations
+
+from numbers import Integral
+
+
+def roll(array, shift, axis=None):
+    """Roll array elements along a given axis.
+
+    Elements that roll beyond the last position are re-introduced at the first.
+
+    Parameters
+    ----------
+    array : array_like
+        Input array.
+    shift : int or tuple of ints
+        The number of places by which elements are shifted.
+    axis : int or tuple of ints, optional
+        Axis or axes along which elements are shifted. By default, the
+        array is flattened before shifting, after which the original shape
+        is restored.
+
+    Returns
+    -------
+    result : Array
+        Array with the same shape as array.
+
+    See Also
+    --------
+    numpy.roll
+    """
+    # Import here to avoid circular imports
+    from dask_array._collection import concatenate, ravel
+
+    result = array
+
+    if axis is None:
+        result = ravel(result)
+
+        if not isinstance(shift, Integral):
+            raise TypeError("Expect `shift` to be an instance of Integral when `axis` is None.")
+
+        shift = (shift,)
+        axis = (0,)
+    else:
+        try:
+            len(shift)
+        except TypeError:
+            shift = (shift,)
+        try:
+            len(axis)
+        except TypeError:
+            axis = (axis,)
+
+    if len(shift) != len(axis):
+        raise ValueError("Must have the same number of shifts as axes.")
+
+    for i, s in zip(axis, shift):
+        shape = result.shape[i]
+        s = 0 if shape == 0 else -s % shape
+
+        sl1 = result.ndim * [slice(None)]
+        sl2 = result.ndim * [slice(None)]
+
+        sl1[i] = slice(s, None)
+        sl2[i] = slice(None, s)
+
+        sl1 = tuple(sl1)
+        sl2 = tuple(sl2)
+
+        result = concatenate([result[sl1], result[sl2]], axis=i)
+
+    result = result.reshape(array.shape)
+    # Ensure that the output is always a new array object
+    result = result.copy() if result is array else result
+
+    return result