dask-array 0.1.0__py3-none-any.whl

dask_array/io/_from_delayed.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from dask_array.io._base import IO
+from dask_array._utils import meta_from_array
+
+if TYPE_CHECKING:
+    pass
+
+
+class FromDelayed(IO):
+    """Expression for creating an array from a delayed value."""
+
+    _parameters = ["value", "shape", "dtype", "_meta", "_name_prefix"]
+    _defaults = {"dtype": None, "_meta": None, "_name_prefix": None}
+
+    @functools.cached_property
+    def _meta(self):
+        meta = self.operand("_meta")
+        dtype = self.operand("dtype")
+        shape = self.operand("shape")
+        if meta is not None:
+            if dtype is None:
+                dtype = getattr(meta, "dtype", None)
+            return meta_from_array(meta, dtype=dtype)
+        if dtype is not None:
+            return np.empty((0,) * len(shape), dtype=dtype)
+        return np.empty((0,) * len(shape))
+
+    @functools.cached_property
+    def chunks(self):
+        return tuple((d,) for d in self.operand("shape"))
+
+    @functools.cached_property
+    def _name(self):
+        prefix = self.operand("_name_prefix")
+        if prefix:
+            return prefix
+        return "from-delayed-" + self.deterministic_token
+
+    def _layer(self):
+        from dask._task_spec import Alias
+        from dask.base import is_dask_collection
+
+        value = self.operand("value")
+        shape = self.operand("shape")
+        key = (self._name,) + (0,) * len(shape)
+        task = Alias(key=key, target=value.key)
+        result = {key: task}
+        # Include the delayed value's graph
+        if is_dask_collection(value):
+            result.update(value.__dask_graph__())
+        return result
+
+
+def from_delayed(value, shape, dtype=None, meta=None, name=None):
+    """Create a dask array from a dask delayed value
+
+    This routine is useful for constructing dask arrays in an ad-hoc fashion
+    using dask delayed, particularly when combined with stack and concatenate.
+
+    The dask array will consist of a single chunk.
+
+    Examples
+    --------
+    >>> import dask
+    >>> import dask_array as da
+    >>> import numpy as np
+    >>> value = dask.delayed(np.ones)(5)
+    >>> array = da.from_delayed(value, (5,), dtype=float)
+    >>> array
+    dask.array<from-value, shape=(5,), dtype=float64, chunksize=(5,), chunktype=numpy.ndarray>
+    >>> array.compute()
+    array([1., 1., 1., 1., 1.])
+    """
+    from dask_array._new_collection import new_collection
+    from dask.delayed import Delayed, delayed
+
+    # Convert to Delayed if it has a key but isn't a Delayed
+    if not isinstance(value, Delayed) and hasattr(value, "key"):
+        value = delayed(value)
+
+    return new_collection(
+        FromDelayed(
+            value=value,
+            shape=shape,
+            dtype=dtype,
+            _meta=meta,
+            _name_prefix=name,
+        )
+    )

dask_array/io/_from_graph.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import functools
+
+from dask import istask
+from dask_array._expr import ArrayExpr
+
+
+class FromGraph(ArrayExpr):
+    _parameters = ["layer", "_meta", "chunks", "keys", "name_prefix", "_dependencies"]
+    _defaults = {"_dependencies": ()}
+
+    @functools.cached_property
+    def _meta(self):
+        return self.operand("_meta")
+
+    @functools.cached_property
+    def chunks(self):
+        return self.operand("chunks")
+
+    @functools.cached_property
+    def _name(self):
+        return self.operand("name_prefix") + "-" + self.deterministic_token
+
+    def dependencies(self):
+        return list(self.operand("_dependencies"))
+
+    def _layer(self):
+        layer = self.operand("layer")
+        our_keys = set(self.operand("keys"))
+        is_hlg = hasattr(layer, "layers")
+
+        # Persist case: layer is a dict of computed values with potentially
+        # different keys (optimization can change key names). Just rename.
+        if not is_hlg:
+            layer_keys = {k for k in layer if isinstance(k, tuple)}
+            if layer_keys and not (layer_keys & our_keys):
+                return {(self._name, *k[1:]) if isinstance(k, tuple) else k: v for k, v in layer.items()}
+
+        # HLG case (e.g., from BlockView): contains tasks and dependencies.
+        # Rename output keys and preserve dependency structure.
+        dsk = dict(layer)
+        result = {}
+        for k, v in dsk.items():
+            if k in our_keys:
+                new_key = (self._name, *k[1:])
+                if istask(v):
+                    result[new_key] = k  # Alias to original
+                    result[k] = v  # Keep original task
+                else:
+                    result[new_key] = v  # Simple rename
+            else:
+                result[k] = v  # Dependency - keep as-is
+        return result

dask_array/io/_from_npy_stack.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import functools
+import os
+import pickle
+from itertools import product
+
+import numpy as np
+
+from dask_array.io._base import IO
+
+
+class FromNpyStack(IO):
+    """Expression for loading an array from a stack of .npy files."""
+
+    _parameters = ["dirname", "mmap_mode"]
+    _defaults = {"mmap_mode": "r"}
+
+    @functools.cached_property
+    def _info(self):
+        """Load and cache the info file."""
+        dirname = self.operand("dirname")
+        with open(os.path.join(dirname, "info"), "rb") as f:
+            return pickle.load(f)
+
+    @functools.cached_property
+    def _meta(self):
+        info = self._info
+        return np.empty((0,) * len(info["chunks"]), dtype=info["dtype"])
+
+    @functools.cached_property
+    def chunks(self):
+        return self._info["chunks"]
+
+    @functools.cached_property
+    def _name(self):
+        return "from-npy-stack-" + self.deterministic_token
+
+    def _layer(self):
+        dirname = self.operand("dirname")
+        mmap_mode = self.operand("mmap_mode")
+        info = self._info
+        chunks = info["chunks"]
+        axis = info["axis"]
+
+        keys = list(product([self._name], *[range(len(c)) for c in chunks]))
+        values = [(np.load, os.path.join(dirname, f"{i}.npy"), mmap_mode) for i in range(len(chunks[axis]))]
+        return dict(zip(keys, values))
+
+
+def from_npy_stack(dirname, mmap_mode="r"):
+    """Load dask array from stack of npy files
+
+    Parameters
+    ----------
+    dirname: string
+        Directory of .npy files
+    mmap_mode: (None or 'r')
+        Read data in memory map mode
+
+    See Also
+    --------
+    to_npy_stack
+    """
+    from dask_array._new_collection import new_collection
+
+    return new_collection(FromNpyStack(dirname=dirname, mmap_mode=mmap_mode))

dask_array/io/_store.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+from collections.abc import Collection
+from threading import Lock
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.typing import ArrayLike
+
+    from dask.delayed import Delayed
+
+from dask.base import named_schedulers
+from dask.utils import SerializableLock
+
+from dask_array._utils import is_arraylike
+from dask_array.slicing._utils import fuse_slice
+
+
+def get_scheduler_lock(collection, scheduler):
+    """Get an appropriate lock for the given collection and scheduler."""
+    if scheduler is None:
+        scheduler = collection.__dask_scheduler__
+    actual_get = named_schedulers.get(scheduler, scheduler)
+    # Only use locks for non-distributed schedulers
+    if actual_get is named_schedulers.get("synchronous", None):
+        return False
+    return SerializableLock()
+
+
+def load_store_chunk(
+    x: Any,
+    out: Any,
+    index: slice | None,
+    region: slice | None,
+    lock: Any,
+    return_stored: bool,
+    load_stored: bool,
+) -> Any:
+    """
+    A function inserted in a Dask graph for storing a chunk.
+
+    Parameters
+    ----------
+    x: array-like
+        An array (potentially a NumPy one)
+    out: array-like
+        Where to store results.
+    index: slice-like
+        Where to store result from ``x`` in ``out``.
+    lock: Lock-like or False
+        Lock to use before writing to ``out``.
+    return_stored: bool
+        Whether to return ``out``.
+    load_stored: bool
+        Whether to return the array stored in ``out``.
+        Ignored if ``return_stored`` is not ``True``.
+
+    Returns
+    -------
+
+    If return_stored=True and load_stored=False
+        out
+    If return_stored=True and load_stored=True
+        out[index]
+    If return_stored=False and compute=False
+        None
+
+    Examples
+    --------
+
+    >>> a = np.ones((5, 6))
+    >>> b = np.empty(a.shape)
+    >>> load_store_chunk(a, b, (slice(None), slice(None)), None, False, False, False)
+    """
+    if region:
+        # Equivalent to `out[region][index]`
+        if index:
+            index = fuse_slice(region, index)
+        else:
+            index = region
+    if lock:
+        lock.acquire()
+    try:
+        if x is not None and x.size != 0:
+            if is_arraylike(x):
+                out[index] = x
+            else:
+                out[index] = np.asanyarray(x)
+
+        if return_stored and load_stored:
+            return out[index]
+        elif return_stored and not load_stored:
+            return out
+        else:
+            return None
+    finally:
+        if lock:
+            lock.release()
+
+
+A = TypeVar("A", bound="ArrayLike")
+
+
+def load_chunk(out: A, index: slice, lock: Any, region: slice | None) -> A:
+    """Load a chunk from an array-like object.
+
+    This is used for loading stored chunks back into dask arrays.
+    """
+    return load_store_chunk(
+        None,
+        out=out,
+        region=region,
+        index=index,
+        lock=lock,
+        return_stored=True,
+        load_stored=True,
+    )
+
+
+def store(
+    sources,
+    targets,
+    lock: bool | Lock = True,
+    regions: tuple[slice, ...] | Collection[tuple[slice, ...]] | None = None,
+    compute: bool = True,
+    return_stored: bool = False,
+    load_stored: bool | None = None,
+    **kwargs,
+):
+    """Store dask arrays in array-like objects, overwrite data in target
+
+    This stores dask arrays into object that supports numpy-style setitem
+    indexing.  It stores values chunk by chunk so that it does not have to
+    fill up memory.  For best performance you can align the block size of
+    the storage target with the block size of your array.
+
+    If your data fits in memory then you may prefer calling
+    ``np.array(myarray)`` instead.
+
+    Parameters
+    ----------
+
+    sources: Array or collection of Arrays
+    targets: array-like or Delayed or collection of array-likes and/or Delayeds
+        These should support setitem syntax ``target[10:20] = ...``.
+        If sources is a single item, targets must be a single item; if sources is a
+        collection of arrays, targets must be a matching collection.
+    lock: boolean or threading.Lock, optional
+        Whether or not to lock the data stores while storing.
+        Pass True (lock each file individually), False (don't lock) or a
+        particular :class:`threading.Lock` object to be shared among all writes.
+    regions: tuple of slices or collection of tuples of slices, optional
+        Each ``region`` tuple in ``regions`` should be such that
+        ``target[region].shape = source.shape``
+        for the corresponding source and target in sources and targets,
+        respectively. If this is a tuple, the contents will be assumed to be
+        slices, so do not provide a tuple of tuples.
+    compute: boolean, optional
+        If true compute immediately; return :class:`dask.delayed.Delayed` otherwise.
+    return_stored: boolean, optional
+        Optionally return the stored result (default False).
+    load_stored: boolean, optional
+        Optionally return the stored result, loaded in to memory (default None).
+        If None, ``load_stored`` is True if ``return_stored`` is True and
+        ``compute`` is False. *This is an advanced option.*
+        When False, store will return the appropriate ``target`` for each chunk that is stored.
+        Directly computing this result is not what you want.
+        Instead, you can use the returned ``target`` to execute followup operations to the store.
+    kwargs:
+        Parameters passed to compute/persist (only used if compute=True)
+
+    Returns
+    -------
+
+    If return_stored=True
+        tuple of Arrays
+    If return_stored=False and compute=True
+        None
+    If return_stored=False and compute=False
+        Delayed
+
+    Examples
+    --------
+
+    >>> import h5py  # doctest: +SKIP
+    >>> f = h5py.File('myfile.hdf5', mode='a')  # doctest: +SKIP
+    >>> dset = f.create_dataset('/data', shape=x.shape,
+    ...                                  chunks=x.chunks,
+    ...                                  dtype='f8')  # doctest: +SKIP
+
+    >>> store(x, dset)  # doctest: +SKIP
+
+    Alternatively store many arrays at the same time
+
+    >>> store([x, y, z], [dset1, dset2, dset3])  # doctest: +SKIP
+    """
+    from dask.base import persist
+    from dask.layers import ArraySliceDep
+
+    from dask_array._collection import Array
+    from dask_array._map_blocks import map_blocks
+
+    if isinstance(sources, Array):
+        sources = [sources]
+        targets = [targets]
+    targets = cast("Collection[ArrayLike | Delayed]", targets)
+
+    if any(not isinstance(s, Array) for s in sources):
+        raise ValueError("All sources must be dask array objects")
+
+    if len(sources) != len(targets):
+        raise ValueError(f"Different number of sources [{len(sources)}] and targets [{len(targets)}]")
+
+    if isinstance(regions, tuple) or regions is None:
+        regions_list = [regions] * len(sources)
+    else:
+        regions_list = list(regions)
+        if len(sources) != len(regions_list):
+            raise ValueError(
+                f"Different number of sources [{len(sources)}] and "
+                f"targets [{len(targets)}] than regions [{len(regions_list)}]"
+            )
+    del regions
+
+    if load_stored is None:
+        load_stored = return_stored and not compute
+
+    if lock is True:
+        lock = get_scheduler_lock(Array, kwargs.get("scheduler"))
+
+    arrays = []
+    for s, t, r in zip(sources, targets, regions_list):
+        slices = ArraySliceDep(s.chunks)
+        arrays.append(
+            map_blocks(
+                load_store_chunk,
+                s,
+                t,
+                slices,
+                region=r,
+                lock=lock,
+                return_stored=return_stored,
+                load_stored=load_stored,
+                name="store-map",
+                meta=s._meta,
+            )
+        )
+
+    if compute:
+        if not return_stored:
+            import dask
+
+            dask.compute(arrays, **kwargs)
+            return None
+        else:
+            stored_persisted = persist(*arrays, **kwargs)
+            arrays = []
+            for s, r in zip(stored_persisted, regions_list):
+                slices = ArraySliceDep(s.chunks)
+                arrays.append(
+                    map_blocks(
+                        load_chunk,
+                        s,
+                        slices,
+                        lock=lock,
+                        region=r,
+                        name="load-stored",
+                        meta=s._meta,
+                    )
+                )
+    if len(arrays) == 1:
+        return arrays[0]
+    return tuple(arrays)
+
+
+def to_hdf5(filename, *args, chunks=True, **kwargs):
+    """Store arrays in HDF5 file
+
+    This saves several dask arrays into several datapaths in an HDF5 file.
+    It creates the necessary datasets and handles clean file opening/closing.
+
+    Parameters
+    ----------
+    chunks: tuple or ``True``
+        Chunk shape, or ``True`` to pass the chunks from the dask array.
+        Defaults to ``True``.
+
+    Examples
+    --------
+
+    >>> da.to_hdf5('myfile.hdf5', '/x', x)  # doctest: +SKIP
+
+    or
+
+    >>> da.to_hdf5('myfile.hdf5', {'/x': x, '/y': y})  # doctest: +SKIP
+
+    Optionally provide arguments as though to ``h5py.File.create_dataset``
+
+    >>> da.to_hdf5('myfile.hdf5', '/x', x, compression='lzf', shuffle=True)  # doctest: +SKIP
+
+    >>> da.to_hdf5('myfile.hdf5', '/x', x, chunks=(10,20,30))  # doctest: +SKIP
+
+    This can also be used as a method on a single Array
+
+    >>> x.to_hdf5('myfile.hdf5', '/x')  # doctest: +SKIP
+
+    See Also
+    --------
+    da.store
+    h5py.File.create_dataset
+    """
+    from dask_array._collection import Array
+
+    if len(args) == 1 and isinstance(args[0], dict):
+        data = args[0]
+    elif len(args) == 2 and isinstance(args[0], str) and isinstance(args[1], Array):
+        data = {args[0]: args[1]}
+    else:
+        raise ValueError("Please provide {'/data/path': array} dictionary")
+
+    import h5py
+
+    with h5py.File(filename, mode="a") as f:
+        dsets = [
+            f.require_dataset(
+                dp,
+                shape=x.shape,
+                dtype=x.dtype,
+                chunks=tuple(c[0] for c in x.chunks) if chunks is True else chunks,
+                **kwargs,
+            )
+            for dp, x in data.items()
+        ]
+        store(list(data.values()), dsets)

dask_array/io/_tiledb.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from dask_array.io._zarr import _check_regular_chunks
+from dask_array.core._conversion import from_array
+
+
+def _tiledb_to_chunks(tiledb_array):
+    schema = tiledb_array.schema
+    return list(schema.domain.dim(i).tile for i in range(schema.ndim))
+
+
+def from_tiledb(uri, attribute=None, chunks=None, storage_options=None, **kwargs):
+    """Load array from the TileDB storage format
+
+    See https://docs.tiledb.io for more information about TileDB.
+
+    Parameters
+    ----------
+    uri: TileDB array or str
+        Location to save the data
+    attribute: str or None
+        Attribute selection (single-attribute view on multi-attribute array)
+
+
+    Returns
+    -------
+
+    A Dask Array
+
+    Examples
+    --------
+
+    >>> import tempfile, tiledb
+    >>> import dask_array as da, numpy as np
+    >>> uri = tempfile.NamedTemporaryFile().name
+    >>> _ = tiledb.from_numpy(uri, np.arange(0,9).reshape(3,3))  # create a tiledb array
+    >>> tdb_ar = da.from_tiledb(uri)  # read back the array
+    >>> tdb_ar.shape
+    (3, 3)
+    >>> tdb_ar.mean().compute()
+    4.0
+    """
+    import tiledb
+
+    tiledb_config = storage_options or dict()
+    key = tiledb_config.pop("key", None)
+
+    if isinstance(uri, tiledb.Array):
+        tdb = uri
+    else:
+        tdb = tiledb.open(uri, attr=attribute, config=tiledb_config, key=key)
+
+    if tdb.schema.sparse:
+        raise ValueError("Sparse TileDB arrays are not supported")
+
+    if not attribute:
+        if tdb.schema.nattr > 1:
+            raise TypeError("keyword 'attribute' must be providedwhen loading a multi-attribute TileDB array")
+        else:
+            attribute = tdb.schema.attr(0).name
+
+    if tdb.iswritable:
+        raise ValueError("TileDB array must be open for reading")
+
+    chunks = chunks or _tiledb_to_chunks(tdb)
+
+    assert len(chunks) == tdb.schema.ndim
+
+    return from_array(tdb, chunks, name=f"tiledb-{uri}")
+
+
+def to_tiledb(
+    darray,
+    uri,
+    compute=True,
+    return_stored=False,
+    storage_options=None,
+    key=None,
+    **kwargs,
+):
+    """Save array to the TileDB storage format
+
+    Save 'array' using the TileDB storage manager, to any TileDB-supported URI,
+    including local disk, S3, or HDFS.
+
+    See https://docs.tiledb.io for more information about TileDB.
+
+    Parameters
+    ----------
+
+    darray: dask.array
+        A dask array to write.
+    uri:
+        Any supported TileDB storage location.
+    storage_options: dict
+        Dict containing any configuration options for the TileDB backend.
+        see https://docs.tiledb.io/en/stable/tutorials/config.html
+    compute, return_stored: see ``store()``
+    key: str or None
+        Encryption key
+
+    Returns
+    -------
+
+    None
+        Unless ``return_stored`` is set to ``True`` (``False`` by default)
+
+    Notes
+    -----
+
+    TileDB only supports regularly-chunked arrays.
+    TileDB `tile extents`_ correspond to form 2 of the dask
+    `chunk specification`_, and the conversion is
+    done automatically for supported arrays.
+
+    Examples
+    --------
+
+    >>> import dask_array as da, tempfile
+    >>> uri = tempfile.NamedTemporaryFile().name
+    >>> data = da.random.random(5,5)
+    >>> da.to_tiledb(data, uri)
+    >>> import tiledb
+    >>> tdb_ar = tiledb.open(uri)
+    >>> all(tdb_ar == data)
+    True
+
+    .. _chunk specification: https://docs.tiledb.io/en/stable/tutorials/tiling-dense.html
+    .. _tile extents: http://docs.dask.org/en/latest/array-chunks.html
+    """
+    import tiledb
+
+    tiledb_config = storage_options or dict()
+    # encryption key, if any
+    key = key or tiledb_config.pop("key", None)
+
+    if not _check_regular_chunks(darray.chunks):
+        raise ValueError(
+            "Attempt to save array to TileDB with irregular chunking, please call `arr.rechunk(...)` first."
+        )
+
+    if isinstance(uri, str):
+        chunks = [c[0] for c in darray.chunks]
+        # create a suitable, empty, writable TileDB array
+        tdb = tiledb.empty_like(uri, darray, tile=chunks, config=tiledb_config, key=key, **kwargs)
+    elif isinstance(uri, tiledb.Array):
+        tdb = uri
+        # sanity checks
+        if not ((darray.dtype == tdb.dtype) and (darray.ndim == tdb.ndim)):
+            raise ValueError("Target TileDB array layout is not compatible with source array")
+    else:
+        raise ValueError(
+            "'uri' must be string pointing to supported TileDB store location or an open, writable TileDB array."
+        )
+
+    if not (tdb.isopen and tdb.iswritable):
+        raise ValueError("Target TileDB array is not open and writable.")
+
+    return darray.store(tdb, lock=False, compute=compute, return_stored=return_stored)