dask-array 0.1.0__py3-none-any.whl

dask_array/fft.py

@@ -0,0 +1,320 @@
+"""FFT operations for array-expr."""
+
+from __future__ import annotations
+
+import inspect
+import warnings
+from collections.abc import Sequence
+
+import numpy as np
+
+try:
+    import scipy
+    import scipy.fftpack
+except ImportError:
+    scipy = None
+
+from dask_array._collection import asarray, concatenate
+from dask_array._map_blocks import map_blocks
+from dask_array.creation import arange
+from dask_array._numpy_compat import NUMPY_GE_200
+from dask.utils import derived_from, skip_doctest
+
+chunk_error = (
+    "Dask array only supports taking an FFT along an axis that \n"
+    "has a single chunk. An FFT operation was tried on axis %s \n"
+    "which has chunks %s. To change the array's chunks use "
+    "dask.Array.rechunk."
+)
+
+fft_preamble = """
+    Wrapping of %s
+
+    The axis along which the FFT is applied must have only one chunk. To change
+    the array's chunking use dask.Array.rechunk.
+
+    The %s docstring follows below:
+
+    """
+
+
+def _fft_out_chunks(a, s, axes):
+    """For computing the output chunks of [i]fft*"""
+    if s is None:
+        return a.chunks
+    chunks = list(a.chunks)
+    for i, axis in enumerate(axes):
+        chunks[axis] = (s[i],)
+    return tuple(chunks)
+
+
+def _rfft_out_chunks(a, s, axes):
+    """For computing the output chunks of rfft*"""
+    if s is None:
+        s = [a.chunks[axis][0] for axis in axes]
+    s = list(s)
+    s[-1] = s[-1] // 2 + 1
+    chunks = list(a.chunks)
+    for i, axis in enumerate(axes):
+        chunks[axis] = (s[i],)
+    return tuple(chunks)
+
+
+def _irfft_out_chunks(a, s, axes):
+    """For computing the output chunks of irfft*"""
+    if s is None:
+        s = [a.chunks[axis][0] for axis in axes]
+        s[-1] = 2 * (s[-1] - 1)
+    chunks = list(a.chunks)
+    for i, axis in enumerate(axes):
+        chunks[axis] = (s[i],)
+    return tuple(chunks)
+
+
+def _hfft_out_chunks(a, s, axes):
+    assert len(axes) == 1
+
+    axis = axes[0]
+
+    if s is None:
+        s = [2 * (a.chunks[axis][0] - 1)]
+
+    n = s[0]
+
+    chunks = list(a.chunks)
+    chunks[axis] = (n,)
+    return tuple(chunks)
+
+
+def _ihfft_out_chunks(a, s, axes):
+    assert len(axes) == 1
+
+    axis = axes[0]
+
+    if s is None:
+        s = [a.chunks[axis][0]]
+    else:
+        assert len(s) == 1
+
+    n = s[0]
+
+    chunks = list(a.chunks)
+    if n % 2 == 0:
+        m = (n // 2) + 1
+    else:
+        m = (n + 1) // 2
+    chunks[axis] = (m,)
+    return tuple(chunks)
+
+
+_out_chunk_fns = {
+    "fft": _fft_out_chunks,
+    "ifft": _fft_out_chunks,
+    "rfft": _rfft_out_chunks,
+    "irfft": _irfft_out_chunks,
+    "hfft": _hfft_out_chunks,
+    "ihfft": _ihfft_out_chunks,
+}
+
+
+def fft_wrap(fft_func, kind=None, dtype=None, allow_fftpack=False):
+    """Wrap 1D, 2D, and ND real and complex FFT functions
+
+    Takes a function that behaves like ``numpy.fft`` functions and
+    a specified kind to match it to that are named after the functions
+    in the ``numpy.fft`` API.
+
+    Supported kinds include:
+
+        * fft
+        * fft2
+        * fftn
+        * ifft
+        * ifft2
+        * ifftn
+        * rfft
+        * rfft2
+        * rfftn
+        * irfft
+        * irfft2
+        * irfftn
+        * hfft
+        * ihfft
+
+    Examples
+    --------
+    >>> import dask_array.fft as dff
+    >>> parallel_fft = dff.fft_wrap(np.fft.fft)
+    >>> parallel_ifft = dff.fft_wrap(np.fft.ifft)
+    """
+    if scipy is not None:
+        if fft_func.__module__.startswith("scipy.fftpack"):
+            if not allow_fftpack:
+                warnings.warn(
+                    f"Function {fft_func.__name__} from `scipy.fftpack` does not "
+                    "match NumPy's API and is considered legacy. Please use "
+                    "`scipy.fft` instead. To suppress this warning and allow usage"
+                    ", set `allow_fftpack=True`. Support for `scipy.fftpack` will "
+                    "be deprecated in future releases.",
+                    FutureWarning,
+                )
+            # If allow_fftpack is True, we proceed but we skip passing the norm
+            # argument.
+
+    if kind is None:
+        kind = fft_func.__name__
+    try:
+        out_chunk_fn = _out_chunk_fns[kind.rstrip("2n")]
+    except KeyError:
+        raise ValueError(f"Given unknown `kind` {kind}.")
+
+    def func(a, s=None, axes=None, norm=None):
+        a = asarray(a)
+        if axes is None:
+            if kind.endswith("2"):
+                axes = (-2, -1)
+            elif kind.endswith("n"):
+                if s is None:
+                    axes = tuple(range(a.ndim))
+                else:
+                    if NUMPY_GE_200:
+                        # Match deprecation in numpy
+                        warnings.warn(
+                            "DeprecationWarning: `axes` should not be `None` "
+                            "if `s` is not `None` (Deprecated in NumPy 2.0)",
+                            DeprecationWarning,
+                        )
+                    axes = tuple(range(len(s)))
+            else:
+                axes = (-1,)
+        elif len(set(axes)) < len(axes):
+            raise ValueError("Duplicate axes not allowed.")
+
+        _dtype = dtype
+        if _dtype is None:
+            sample = np.ones(a.ndim * (8,), dtype=a.dtype)
+            try:
+                _dtype = fft_func(sample, axes=axes, norm=norm).dtype
+            except TypeError:
+                _dtype = fft_func(sample).dtype
+
+        for each_axis in axes:
+            if len(a.chunks[each_axis]) != 1:
+                raise ValueError(chunk_error % (each_axis, a.chunks[each_axis]))
+
+        chunks = out_chunk_fn(a, s, axes)
+
+        args = (s, axes, norm)
+        if kind.endswith("fft"):
+            axis = None if axes is None else axes[0]
+            n = None if s is None else s[0]
+            args = (n, axis, norm)
+
+        return map_blocks(fft_func, a, *args, dtype=_dtype, chunks=chunks)
+
+    if kind.endswith("fft"):
+        _func = func
+
+        def func(a, n=None, axis=None, norm=None):
+            s = None
+            if n is not None:
+                s = (n,)
+
+            axes = None
+            if axis is not None:
+                axes = (axis,)
+
+            return _func(a, s, axes, norm)
+
+    func_mod = inspect.getmodule(fft_func)
+    func_name = fft_func.__name__
+    func_fullname = f"{func_mod.__name__}.{func_name}"
+    if fft_func.__doc__ is not None:
+        func.__doc__ = fft_preamble % (2 * (func_fullname,))
+        func.__doc__ += fft_func.__doc__
+        func.__doc__ = skip_doctest(func.__doc__)
+    func.__name__ = func_name
+    return func
+
+
+fft = fft_wrap(np.fft.fft)
+fft2 = fft_wrap(np.fft.fft2)
+fftn = fft_wrap(np.fft.fftn)
+ifft = fft_wrap(np.fft.ifft)
+ifft2 = fft_wrap(np.fft.ifft2)
+ifftn = fft_wrap(np.fft.ifftn)
+rfft = fft_wrap(np.fft.rfft)
+rfft2 = fft_wrap(np.fft.rfft2)
+rfftn = fft_wrap(np.fft.rfftn)
+irfft = fft_wrap(np.fft.irfft)
+irfft2 = fft_wrap(np.fft.irfft2)
+irfftn = fft_wrap(np.fft.irfftn)
+hfft = fft_wrap(np.fft.hfft)
+ihfft = fft_wrap(np.fft.ihfft)
+
+
+def _fftfreq_block(i, n, d):
+    r = i.copy()
+    r[i >= (n + 1) // 2] -= n
+    r /= n * d
+    return r
+
+
+@derived_from(np.fft)
+def fftfreq(n, d=1.0, chunks="auto"):
+    n = int(n)
+    d = float(d)
+
+    r = arange(n, dtype=float, chunks=chunks)
+
+    return map_blocks(_fftfreq_block, r, dtype=float, n=n, d=d)
+
+
+@derived_from(np.fft)
+def rfftfreq(n, d=1.0, chunks="auto"):
+    n = int(n)
+    d = float(d)
+
+    r = arange(n // 2 + 1, dtype=float, chunks=chunks)
+    r = r / (n * d)
+
+    return r
+
+
+def _fftshift_helper(x, axes=None, inverse=False):
+    if axes is None:
+        axes = list(range(x.ndim))
+    elif not isinstance(axes, Sequence):
+        axes = (axes,)
+
+    y = x
+    for i in axes:
+        n = y.shape[i]
+        n_2 = (n + int(inverse is False)) // 2
+
+        l = y.ndim * [slice(None)]
+        l[i] = slice(None, n_2)
+        l = tuple(l)
+
+        r = y.ndim * [slice(None)]
+        r[i] = slice(n_2, None)
+        r = tuple(r)
+
+        y = concatenate([y[r], y[l]], axis=i)
+
+        if len(x.chunks[i]) == 1:
+            y = y.rechunk({i: x.chunks[i]})
+
+    return y
+
+
+@derived_from(np.fft)
+def fftshift(x, axes=None):
+    x = asarray(x)
+    return _fftshift_helper(x, axes=axes, inverse=False)
+
+
+@derived_from(np.fft)
+def ifftshift(x, axes=None):
+    x = asarray(x)
+    return _fftshift_helper(x, axes=axes, inverse=True)

dask_array/io/__init__.py

@@ -0,0 +1,39 @@
+"""IO functions for array-expr."""
+
+from __future__ import annotations
+
+from dask_array.io._base import IO
+from dask_array.io._from_array import FromArray
+from dask_array.io._from_delayed import FromDelayed, from_delayed
+from dask_array.io._from_graph import FromGraph
+from dask_array.io._from_npy_stack import FromNpyStack, from_npy_stack
+from dask_array.io._store import (
+    get_scheduler_lock,
+    load_chunk,
+    load_store_chunk,
+    store,
+    to_hdf5,
+)
+from dask_array.io._tiledb import from_tiledb, to_tiledb
+from dask_array.io._to_npy_stack import to_npy_stack
+from dask_array.io._zarr import from_zarr, to_zarr
+
+__all__ = [
+    "IO",
+    "FromArray",
+    "FromDelayed",
+    "FromGraph",
+    "FromNpyStack",
+    "from_delayed",
+    "from_npy_stack",
+    "from_tiledb",
+    "from_zarr",
+    "get_scheduler_lock",
+    "load_chunk",
+    "load_store_chunk",
+    "store",
+    "to_hdf5",
+    "to_npy_stack",
+    "to_tiledb",
+    "to_zarr",
+]

dask_array/io/_base.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from dask_array._expr import ArrayExpr
+
+
+class IO(ArrayExpr):
+    # Whether rechunk can be pushed into this IO expression by modifying its chunks.
+    # False by default since many IO expressions have chunks that affect computation
+    # (e.g., Random generates different values with different chunks).
+    _can_rechunk_pushdown = False

dask_array/io/_from_array.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+import functools
+from itertools import product
+
+import numpy as np
+
+from dask_array.io._base import IO
+from dask_array._core_utils import (
+    getter,
+    getter_nofancy,
+    graph_from_arraylike,
+    normalize_chunks,
+    slices_from_chunks,
+)
+from dask_array._utils import meta_from_array
+from dask.utils import SerializableLock
+
+
+class FromArray(IO):
+    _parameters = [
+        "array",
+        "_chunks",
+        "lock",
+        "getitem",
+        "inline_array",
+        "meta",
+        "asarray",
+        "fancy",
+        "_name_override",
+        "_region",  # Slice region for pushdown (tuple of slices or None)
+    ]
+    _defaults = {
+        "_chunks": "auto",
+        "getitem": None,
+        "inline_array": False,
+        "meta": None,
+        "asarray": None,
+        "fancy": True,
+        "lock": False,
+        "_name_override": None,
+        "_region": None,
+    }
+    # FromArray reads static data, so rechunk can be pushed in safely
+    _can_rechunk_pushdown = True
+    # Slicing can be pushed into FromArray by slicing the source array
+    _slice_pushdown = True
+
+    @functools.cached_property
+    def _name(self):
+        # _name_override is a prefix, deterministic token is always appended
+        prefix = self.operand("_name_override") or "fromarray"
+        return f"{prefix}-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _effective_shape(self):
+        """Shape after applying region slice."""
+        region = self.operand("_region")
+        if region is None:
+            return self.array.shape
+        # Compute shape from region slices
+        return tuple(len(range(*slc.indices(dim_size))) for slc, dim_size in zip(region, self.array.shape))
+
+    @functools.cached_property
+    def chunks(self):
+        # Normalize chunks lazily - keeps repr compact with user-provided chunks
+        # Pass previous_chunks from underlying array (h5py, zarr) for alignment
+        previous_chunks = getattr(self.array, "chunks", None)
+        # Handle zarr 3.x shards attribute for write alignment
+        if hasattr(self.array, "shards") and self.array.shards is not None and self.operand("_chunks") == "auto":
+            previous_chunks = self.array.shards
+        return normalize_chunks(
+            self.operand("_chunks"),
+            self._effective_shape,
+            dtype=self.array.dtype,
+            previous_chunks=previous_chunks,
+        )
+
+    @functools.cached_property
+    def _meta(self):
+        if self.operand("meta") is not None:
+            return meta_from_array(self.operand("meta"), ndim=len(self._effective_shape), dtype=self.array.dtype)
+        return meta_from_array(self.array, dtype=getattr(self.array, "dtype", None))
+
+    @functools.cached_property
+    def asarray_arg(self):
+        if self.operand("asarray") is None:
+            return not hasattr(self.array, "__array_function__")
+        else:
+            return self.operand("asarray")
+
+    def _layer(self):
+        lock = self.operand("lock")
+        region = self.operand("_region")
+        # Note: lock=True is already normalized to SerializableLock() in from_array()
+
+        is_ndarray = type(self.array) in (np.ndarray, np.ma.core.MaskedArray)
+        is_single_block = all(len(c) == 1 for c in self.chunks)
+
+        # Get slices for chunks (based on effective shape after region)
+        slices = slices_from_chunks(self.chunks)
+
+        # If region is set, offset all slices by the region start
+        if region is not None:
+            region_starts = tuple(slc.indices(dim_size)[0] for slc, dim_size in zip(region, self.array.shape))
+            slices = [
+                tuple(slice(s.start + offset, s.stop + offset, s.step) for s, offset in zip(slc, region_starts))
+                for slc in slices
+            ]
+
+        # Always use the getter for h5py etc. Not using isinstance(x, np.ndarray)
+        # because np.matrix is a subclass of np.ndarray.
+        if is_ndarray and not is_single_block and not lock:
+            # eagerly slice numpy arrays to prevent memory blowup
+            # GH5367, GH5601
+            keys = product([self._name], *(range(len(bds)) for bds in self.chunks))
+            values = [self.array[slc] for slc in slices]
+            dsk = dict(zip(keys, values))
+        elif is_ndarray and is_single_block and not lock:
+            # Single block - slice with region (or full array) and copy
+            if region is not None:
+                dsk = {(self._name,) + (0,) * self.array.ndim: self.array[region].copy()}
+            else:
+                dsk = {(self._name,) + (0,) * self.array.ndim: self.array.copy()}
+        else:
+            getitem = self.operand("getitem")
+            if getitem is None:
+                if self.operand("fancy"):
+                    getitem = getter
+                else:
+                    getitem = getter_nofancy
+
+            # For non-numpy arrays with region, we need custom graph generation
+            # to apply the offset slices
+            if region is not None:
+                keys = list(product([self._name], *(range(len(bds)) for bds in self.chunks)))
+                if self.inline_array:
+                    dsk = {k: (getitem, self.array, slc, self.asarray_arg, lock) for k, slc in zip(keys, slices)}
+                else:
+                    # Put array in graph once, reference by key
+                    arr_key = ("array-" + self._name,)
+                    dsk = {arr_key: self.array}
+                    dsk.update({k: (getitem, arr_key, slc, self.asarray_arg, lock) for k, slc in zip(keys, slices)})
+            else:
+                dsk = graph_from_arraylike(
+                    self.array,
+                    chunks=self.chunks,
+                    shape=self.array.shape,
+                    name=self._name,
+                    lock=lock,
+                    getitem=getitem,
+                    asarray=self.asarray_arg,
+                    inline_array=self.inline_array,
+                    dtype=self.array.dtype,
+                )
+        return dict(dsk)  # this comes as a legacy HLG for now
+
+    def __str__(self):
+        return "FromArray(...)"
+
+    def __dask_tokenize__(self):
+        from dask.tokenize import _tokenize_deterministic
+
+        # Handle non-serializable locks by using their id()
+        # Locks are identity-based objects, so using id() is semantically correct
+        lock = self.operand("lock")
+        if lock and not isinstance(lock, (bool, SerializableLock)):
+            lock_token = ("lock-id", id(lock))
+        else:
+            lock_token = lock
+
+        operands = [lock_token if p == "lock" else self.operand(p) for p in self._parameters]
+        return _tokenize_deterministic(type(self), *operands)
+
+    def _simplify_up(self, parent, dependents):
+        """Allow slice operations to push into FromArray."""
+        from dask_array.slicing import SliceSlicesIntegers
+
+        if isinstance(parent, SliceSlicesIntegers):
+            return self._accept_slice(parent)
+        return None
+
+    def _accept_slice(self, slice_expr):
+        """Accept a slice by setting region (deferred slice).
+
+        Pushes the slice into the FromArray expression by recording it as a region,
+        which is then applied during layer generation.
+        """
+        from numbers import Integral
+
+        from dask_array.slicing._basic import (
+            SliceSlicesIntegers,
+            _compose_slices,
+            _compute_sliced_chunks,
+        )
+
+        index = slice_expr.index
+
+        # Only handle slices and integers (no None/newaxis, no fancy indexing)
+        if any(idx is None for idx in index):
+            return None
+        if any(not isinstance(idx, (slice, Integral)) for idx in index):
+            return None
+        # Don't push non-unit step slices - _layer doesn't handle them correctly
+        if any(isinstance(idx, slice) and idx.step is not None and idx.step != 1 for idx in index):
+            return None
+
+        source = self.array
+        old_chunks = self.chunks  # Use normalized chunks property
+        old_region = self.operand("_region")
+
+        # Pad index to full dimensions
+        full_index = index + (slice(None),) * (source.ndim - len(index))
+
+        # Check if any integers are present - they need special handling
+        has_integers = any(isinstance(idx, Integral) for idx in full_index)
+
+        # Convert integers to 1-element slices for region calculation
+        region_index = tuple(slice(idx, idx + 1) if isinstance(idx, Integral) else idx for idx in full_index)
+
+        # Compute new region by combining with existing region
+        if old_region is not None:
+            # Compose slices: new slice is relative to old region
+            new_region = tuple(
+                _compose_slices(old_slc, new_slc, dim_size)
+                for old_slc, new_slc, dim_size in zip(old_region, region_index, source.shape)
+            )
+        else:
+            new_region = region_index
+
+        # Compute new chunks - use same chunk sizes but clipped to new shape
+        new_chunks = tuple(
+            _compute_sliced_chunks(dim_chunks, slc, dim_size)
+            for dim_chunks, slc, dim_size in zip(old_chunks, region_index, self._effective_shape)
+        )
+
+        # Create new FromArray with region (deferred slice)
+        new_io = FromArray(
+            source,  # Keep original source, don't slice
+            new_chunks,
+            lock=self.operand("lock"),
+            getitem=self.operand("getitem"),
+            inline_array=self.inline_array,
+            meta=self.operand("meta"),
+            asarray=self.operand("asarray"),
+            fancy=self.operand("fancy"),
+            _name_override=self.operand("_name_override"),
+            _region=new_region,
+        )
+
+        # If integers were present, apply them to extract elements
+        if has_integers:
+            # Build index with 0s for integer dims (they're now size-1)
+            extract_index = tuple(0 if isinstance(idx, Integral) else slice(None) for idx in full_index)
+            return SliceSlicesIntegers(new_io, extract_index, False)
+
+        return new_io