dask-array 0.1.0__py3-none-any.whl

dask_array/manipulation/_transpose.py

@@ -0,0 +1,309 @@
+"""Transpose operations: transpose, swapaxes, moveaxis, rollaxis."""
+
+from __future__ import annotations
+
+import functools
+
+import numpy as np
+
+from dask._task_spec import Task, TaskRef
+from dask_array._blockwise import Blockwise
+from dask_array._utils import meta_from_array
+
+
+class Transpose(Blockwise):
+    _parameters = ["array", "axes"]
+    func = staticmethod(np.transpose)
+    align_arrays = False
+    adjust_chunks = None
+    concatenate = None
+    token = "transpose"
+
+    @property
+    def new_axes(self):
+        return {}
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def _meta_provided(self):
+        return self.array._meta
+
+    @functools.cached_property
+    def _meta(self):
+        meta = self.array._meta
+        if meta is None or getattr(meta, "ndim", None) != len(self.axes):
+            return meta_from_array(None, ndim=len(self.axes), dtype=self.array.dtype)
+        return np.transpose(meta, self.axes)
+
+    @property
+    def dtype(self):
+        return self.array.dtype
+
+    @property
+    def out_ind(self):
+        return self.axes
+
+    @property
+    def kwargs(self):
+        return {"axes": self.axes}
+
+    @property
+    def args(self):
+        return (self.array, tuple(range(self.array.ndim)))
+
+    @functools.cached_property
+    def _inverse_axes(self):
+        """Inverse permutation of axes."""
+        inv = [0] * len(self.axes)
+        for i, a in enumerate(self.axes):
+            inv[a] = i
+        return tuple(inv)
+
+    def _task(self, key, block_id: tuple[int, ...]) -> Task:
+        """Generate task for a specific output block."""
+        # Map output block_id to input block_id using inverse permutation
+        # For axes=(1,0), output block (i,j) needs input block (j,i)
+        input_block_id = self._input_block_id(self.array, block_id)
+        return Task(key, self.func, TaskRef((self.array._name, *input_block_id)), **self.kwargs)
+
+    def _input_block_id(self, dep, block_id: tuple[int, ...]) -> tuple[int, ...]:
+        """Map output block_id to input block_id using inverse permutation."""
+        return tuple(block_id[self._inverse_axes[d]] for d in range(len(block_id)))
+
+    def _simplify_down(self):
+        # Transpose(Transpose(x)) -> single Transpose with composed axes
+        if isinstance(self.array, Transpose):
+            axes = tuple(self.array.axes[i] for i in self.axes)
+            return Transpose(self.array.array, axes)
+        # Identity transpose -> return the array
+        if self.axes == tuple(range(self.array.ndim)):
+            return self.array
+        # Transpose(Elemwise(x, y)) -> Elemwise(Transpose(x), Transpose(y))
+        from dask_array._blockwise import Elemwise
+
+        if isinstance(self.array, Elemwise):
+            return self._pushdown_through_elemwise()
+
+    def _pushdown_through_elemwise(self):
+        """Push transpose through elemwise by transposing each input."""
+        from dask_array._blockwise import Elemwise
+        from dask_array._core_utils import is_scalar_for_elemwise
+
+        elemwise = self.array
+        axes = self.axes
+        out_ndim = len(axes)
+
+        # Only push through if all array inputs have the same ndim as output
+        # Broadcasting cases require index transformations we don't handle
+        for arg in elemwise.elemwise_args:
+            if is_scalar_for_elemwise(arg):
+                continue
+            if arg.ndim != out_ndim:
+                return None
+
+        # Check where/out as well
+        if hasattr(elemwise.where, "ndim") and elemwise.where.ndim != out_ndim:
+            return None
+        if hasattr(elemwise.out, "ndim") and elemwise.out.ndim != out_ndim:
+            return None
+
+        # Transpose each array input
+        new_args = [arg if is_scalar_for_elemwise(arg) else Transpose(arg, axes) for arg in elemwise.elemwise_args]
+
+        # Transpose where/out if they are arrays
+        new_where = elemwise.where
+        if hasattr(new_where, "ndim"):
+            new_where = Transpose(new_where, axes)
+
+        new_out = elemwise.out
+        if hasattr(new_out, "ndim"):
+            new_out = Transpose(new_out, axes)
+
+        return Elemwise(
+            elemwise.op,
+            elemwise.operand("dtype"),
+            elemwise.operand("name"),
+            new_where,
+            new_out,
+            elemwise.operand("_user_kwargs"),
+            *new_args,
+        )
+
+    def _simplify_up(self, parent, dependents):
+        """Allow slice and shuffle operations to push through Transpose."""
+        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_shuffle(self, shuffle_expr):
+        """Accept a shuffle being pushed through Transpose.
+
+        Maps shuffle axis through transpose to get input axis.
+        """
+        axes = self.axes
+        shuffle_axis = shuffle_expr.axis
+
+        # Map shuffle axis through transpose: axes[i] tells us which input axis
+        # becomes output axis i. So to shuffle output axis `shuffle_axis`, we need
+        # to shuffle input axis `axes[shuffle_axis]`.
+        input_axis = axes[shuffle_axis]
+
+        from dask_array._shuffle import Shuffle
+
+        shuffled_input = Shuffle(self.array, shuffle_expr.indexer, input_axis, shuffle_expr.operand("name"))
+        return Transpose(shuffled_input, axes)
+
+    def _accept_slice(self, slice_expr):
+        """Accept a slice being pushed through Transpose.
+
+        Maps output slice indices through transpose axes to get input slice.
+        """
+        from numbers import Integral
+
+        from dask_array._new_collection import new_collection
+
+        axes = self.axes
+        index = slice_expr.index
+
+        # Don't handle None/newaxis (adds dimensions)
+        if any(idx is None for idx in index):
+            return None
+
+        # Pad index to full length
+        full_index = index + (slice(None),) * (self.ndim - len(index))
+
+        # Map output slice through transpose axes to get input slice
+        # axes[i] tells us which input axis becomes output axis i
+        # So output axis i gets slice full_index[i], which should go to input axis axes[i]
+        input_index = [slice(None)] * len(axes)
+        for out_axis, in_axis in enumerate(axes):
+            input_index[in_axis] = full_index[out_axis]
+
+        sliced_input = new_collection(self.array)[tuple(input_index)]
+
+        # Check if any dimensions were removed by integer indexing
+        has_integers = any(isinstance(idx, Integral) for idx in full_index)
+
+        if not has_integers:
+            # No dimension changes - just apply original transpose
+            return Transpose(sliced_input.expr, axes)
+
+        # Integer indices remove dimensions - compute new axes for remaining dims
+        # Track which input dimensions remain (those not indexed by integers)
+        remaining_input_dims = [
+            in_axis for out_axis, in_axis in enumerate(axes) if not isinstance(full_index[out_axis], Integral)
+        ]
+
+        if len(remaining_input_dims) <= 1:
+            # 0 or 1 dimension left - no transpose needed
+            return sliced_input.expr
+
+        # Map old input dim indices to new (post-slice) indices
+        # After slicing, input dims are renumbered 0, 1, 2, ...
+        sorted_remaining = sorted(remaining_input_dims)
+        dim_map = {old: new for new, old in enumerate(sorted_remaining)}
+
+        # Build new axes: for each remaining output dim, what's the new input dim?
+        new_axes = tuple(dim_map[in_dim] for in_dim in remaining_input_dims)
+
+        # Check if it's an identity transpose
+        if new_axes == tuple(range(len(new_axes))):
+            return sliced_input.expr
+
+        return Transpose(sliced_input.expr, new_axes)
+
+
+def transpose(a, axes=None):
+    """Reverse or permute the axes of an array.
+
+    See Also
+    --------
+    numpy.transpose
+    """
+    from dask_array.core import asanyarray
+
+    a = asanyarray(a)
+    if axes is not None:
+        return a.transpose(axes)
+    return a.transpose()
+
+
+def swapaxes(a, axis1, axis2):
+    """Interchange two axes of an array.
+
+    See Also
+    --------
+    numpy.swapaxes
+    """
+    from dask_array.core import asanyarray
+
+    a = asanyarray(a)
+    if axis1 == axis2:
+        return a
+    if axis1 < 0:
+        axis1 = axis1 + a.ndim
+    if axis2 < 0:
+        axis2 = axis2 + a.ndim
+    ind = list(range(a.ndim))
+    ind[axis1], ind[axis2] = ind[axis2], ind[axis1]
+    return transpose(a, ind)
+
+
+def moveaxis(a, source, destination):
+    """Move axes of an array to new positions.
+
+    See Also
+    --------
+    numpy.moveaxis
+    """
+    from dask_array.core import asanyarray
+    from dask_array._numpy_compat import normalize_axis_tuple
+
+    a = asanyarray(a)
+    source = normalize_axis_tuple(source, a.ndim, "source")
+    destination = normalize_axis_tuple(destination, a.ndim, "destination")
+    if len(source) != len(destination):
+        raise ValueError("`source` and `destination` arguments must have the same number of elements")
+
+    order = [n for n in range(a.ndim) if n not in source]
+
+    for dest, src in sorted(zip(destination, source)):
+        order.insert(dest, src)
+
+    return transpose(a, order)
+
+
+def rollaxis(a, axis, start=0):
+    """Roll the specified axis backwards, until it lies in a given position.
+
+    See Also
+    --------
+    numpy.rollaxis
+    """
+    from dask_array.core import asanyarray
+    from dask_array._numpy_compat import normalize_axis_index
+
+    a = asanyarray(a)
+    n = a.ndim
+    axis = normalize_axis_index(axis, n)
+    if start < 0:
+        start += n
+    msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
+    if not (0 <= start < n + 1):
+        raise ValueError(msg % ("start", -n, "start", n + 1, start))
+    if axis < start:
+        start -= 1
+    if axis == start:
+        return a[...]
+    axes = list(range(0, n))
+    axes.remove(axis)
+    axes.insert(start, axis)
+    return transpose(a, axes)

dask_array/random/__init__.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from threading import Lock
+
+from ._generator import Generator, default_rng
+from ._random_state import RandomState
+
+
+# Lazy RNG-state machinery
+#
+# Many of the RandomState methods are exported as functions in da.random for
+# backward compatibility reasons. Their usage is discouraged.
+# Use da.random.default_rng() to get a Generator based rng and use its
+# methods instead.
+
+_cached_states: dict[str, RandomState] = {}
+_cached_states_lock = Lock()
+
+
+def _make_api(attr):
+    def wrapper(*args, **kwargs):
+        from dask_array._backends_array import array_creation_dispatch
+
+        key = array_creation_dispatch.backend
+        with _cached_states_lock:
+            try:
+                state = _cached_states[key]
+            except KeyError:
+                _cached_states[key] = state = RandomState()
+        return getattr(state, attr)(*args, **kwargs)
+
+    wrapper.__name__ = getattr(RandomState, attr).__name__
+    wrapper.__doc__ = getattr(RandomState, attr).__doc__
+    return wrapper
+
+
+# RandomState only
+
+seed = _make_api("seed")
+
+beta = _make_api("beta")
+binomial = _make_api("binomial")
+chisquare = _make_api("chisquare")
+choice = _make_api("choice")
+exponential = _make_api("exponential")
+f = _make_api("f")
+gamma = _make_api("gamma")
+geometric = _make_api("geometric")
+gumbel = _make_api("gumbel")
+hypergeometric = _make_api("hypergeometric")
+laplace = _make_api("laplace")
+logistic = _make_api("logistic")
+lognormal = _make_api("lognormal")
+logseries = _make_api("logseries")
+multinomial = _make_api("multinomial")
+negative_binomial = _make_api("negative_binomial")
+noncentral_chisquare = _make_api("noncentral_chisquare")
+noncentral_f = _make_api("noncentral_f")
+normal = _make_api("normal")
+pareto = _make_api("pareto")
+permutation = _make_api("permutation")
+poisson = _make_api("poisson")
+power = _make_api("power")
+random_sample = _make_api("random_sample")
+random = _make_api("random_sample")
+randint = _make_api("randint")
+random_integers = _make_api("random_integers")
+rayleigh = _make_api("rayleigh")
+standard_cauchy = _make_api("standard_cauchy")
+standard_exponential = _make_api("standard_exponential")
+standard_gamma = _make_api("standard_gamma")
+standard_normal = _make_api("standard_normal")
+standard_t = _make_api("standard_t")
+triangular = _make_api("triangular")
+uniform = _make_api("uniform")
+vonmises = _make_api("vonmises")
+wald = _make_api("wald")
+weibull = _make_api("weibull")
+zipf = _make_api("zipf")
+
+__all__ = [
+    "Generator",
+    "RandomState",
+    "default_rng",
+    "seed",
+    "beta",
+    "binomial",
+    "chisquare",
+    "choice",
+    "exponential",
+    "f",
+    "gamma",
+    "geometric",
+    "gumbel",
+    "hypergeometric",
+    "laplace",
+    "logistic",
+    "lognormal",
+    "logseries",
+    "multinomial",
+    "negative_binomial",
+    "noncentral_chisquare",
+    "noncentral_f",
+    "normal",
+    "pareto",
+    "permutation",
+    "poisson",
+    "power",
+    "random_sample",
+    "random",
+    "randint",
+    "random_integers",
+    "rayleigh",
+    "standard_cauchy",
+    "standard_exponential",
+    "standard_gamma",
+    "standard_normal",
+    "standard_t",
+    "triangular",
+    "uniform",
+    "vonmises",
+    "wald",
+    "weibull",
+    "zipf",
+]

dask_array/random/_choice.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from itertools import product
+from numbers import Integral
+
+import numpy as np
+
+from dask._task_spec import TaskRef
+from dask_array._collection import Array
+from dask_array.core._conversion import asarray
+from dask_array.io import IO
+from dask_array._core_utils import normalize_chunks
+from dask_array._utils import asarray_safe
+from dask_array._backends_array import array_creation_dispatch
+from dask.utils import cached_property, random_state_data
+
+from ._expr import _spawn_bitgens
+from ._generator import Generator
+from ._random_state import RandomState
+
+
+def _choice_rng(state_data, a, size, replace, p, axis, shuffle):
+    from ._expr import _rng_from_bitgen
+
+    state = _rng_from_bitgen(state_data)
+    return state.choice(a, size=size, replace=replace, p=p, axis=axis, shuffle=shuffle)
+
+
+def _choice_rs(state_data, a, size, replace, p):
+    state = array_creation_dispatch.RandomState(state_data)
+    return state.choice(a, size=size, replace=replace, p=p)
+
+
+def _choice_validate_params(state, a, size, replace, p, axis, chunks):
+    """Validate and normalize parameters for choice.
+
+    Returns expressions for array/p (or int/None) so they participate in lowering.
+    """
+    # Normalize and validate `a`
+    if isinstance(a, Integral):
+        if isinstance(state, Generator):
+            if state._backend_name == "cupy":
+                raise NotImplementedError("`choice` not supported for cupy-backed `Generator`.")
+            meta = state._backend.random.default_rng().choice(1, size=(), p=None)
+        elif isinstance(state, RandomState):
+            # On windows the output dtype differs if p is provided or
+            # # absent, see https://github.com/numpy/numpy/issues/9867
+            dummy_p = state._backend.array([1]) if p is not None else p
+            meta = state._backend.random.RandomState().choice(1, size=(), p=dummy_p)
+        else:
+            raise ValueError("Unknown generator class")
+        len_a = a
+        if a < 0:
+            raise ValueError("a must be greater than 0")
+        a_expr = None  # No expression for int a
+    else:
+        a = asarray(a)
+        a = a.rechunk(a.shape)
+        meta = a._meta
+        if a.ndim != 1:
+            raise ValueError("a must be one dimensional")
+        len_a = len(a)
+        a_expr = a.expr  # Store expression so it gets lowered
+
+    # Normalize and validate `p`
+    p_expr = None
+    if p is not None:
+        if not isinstance(p, Array):
+            # If p is not a dask array, first check the sum is close
+            # to 1 before converting.
+            p = asarray_safe(p, like=p)
+            if not np.isclose(p.sum(), 1, rtol=1e-7, atol=0):
+                raise ValueError("probabilities do not sum to 1")
+            p = asarray(p)
+        else:
+            p = p.rechunk(p.shape)
+
+        if p.ndim != 1:
+            raise ValueError("p must be one dimensional")
+        if len(p) != len_a:
+            raise ValueError("a and p must have the same size")
+
+        p_expr = p.expr  # Store expression so it gets lowered
+
+    if size is None:
+        size = ()
+    elif not isinstance(size, (tuple, list)):
+        size = (size,)
+
+    if axis != 0:
+        raise ValueError("axis must be 0 since a is one dimensional")
+
+    chunks = normalize_chunks(chunks, size, dtype=np.float64)
+    if not replace and len(chunks[0]) > 1:
+        err_msg = "replace=False is not currently supported for dask.array.choice with multi-chunk output arrays"
+        raise NotImplementedError(err_msg)
+
+    # For int a, return the int value; for array a, return None (use a_expr)
+    a_val = a if isinstance(a, Integral) else None
+    return a_val, a_expr, size, replace, p_expr, axis, chunks, meta
+
+
+class RandomChoice(IO):
+    _parameters = [
+        "a_val",  # int value of a (or None if a is an array)
+        "a_expr",  # expression for a (or None if a is an int)
+        "chunks",
+        "_meta",
+        "_state",
+        "replace",
+        "p_expr",  # expression for p (or None)
+        "axis",
+        "shuffle",
+    ]
+    _defaults = {"axis": None, "shuffle": None}
+    _funcname = "da.random.choice-"
+
+    @cached_property
+    def chunks(self):
+        return self.operand("chunks")
+
+    @cached_property
+    def sizes(self):
+        return list(product(*self.chunks))
+
+    @cached_property
+    def state_data(self):
+        return random_state_data(len(self.sizes), self._state)
+
+    @cached_property
+    def _meta(self):
+        return self.operand("_meta")
+
+    # No custom dependencies() needed - base class finds Expr operands automatically
+    # (a_expr and p_expr are included when they're expressions, excluded when None)
+
+    @property
+    def _a_arg(self):
+        """Value to pass to choice: int or TaskRef to single-chunk array."""
+        if self.a_val is not None:
+            return self.a_val
+        return TaskRef((self.a_expr._name, 0))
+
+    @property
+    def _p_arg(self):
+        """Value to pass to choice: None or TaskRef to single-chunk array."""
+        if self.p_expr is None:
+            return None
+        return TaskRef((self.p_expr._name, 0))
+
+    def _layer(self) -> dict:
+        keys = product([self._name], *[range(len(bd)) for bd in self.chunks])
+        return {
+            k: (_choice_rs, state, self._a_arg, size, self.replace, self._p_arg)
+            for k, state, size in zip(keys, self.state_data, self.sizes)
+        }
+
+
+class RandomChoiceGenerator(RandomChoice):
+    # Keep axis and shuffle as required parameters (no defaults)
+    _defaults = {}
+
+    @cached_property
+    def state_data(self):
+        return _spawn_bitgens(self._state, len(self.sizes))
+
+    def _layer(self) -> dict:
+        keys = product([self._name], *[range(len(bd)) for bd in self.chunks])
+        return {
+            k: (
+                _choice_rng,
+                bitgen,
+                self._a_arg,
+                size,
+                self.replace,
+                self._p_arg,
+                self.axis,
+                self.shuffle,
+            )
+            for k, bitgen, size in zip(keys, self.state_data, self.sizes)
+        }