ComDAS 0.1.0__py3-none-any.whl

comdas/__init__.py

@@ -0,0 +1,35 @@
+"""
+ComDAS: compressed DAS patches for DASCore.
+
+ComDAS patches work identically to normal ``dascore.Patch`` objects,
+but the data they hold is stored compressed in memory
+(:class:`~comdas.arrays.duck_array.DuckArray`) and
+decompressed on access. Compression is controlled via
+the :class:`~comdas.codecs.base.Codec` interface. Both the in-memory
+array and the on-disk ``COMDAS`` container format
+(:mod:`comdas.io.h5_container`) are codec-agnostic, working with any
+registered codec by name rather than hardcoding a specific one.
+"""
+
+from comdas.codecs.base import Codec, CompressedPayload
+from comdas.codecs.svd import SVDCodec, SVDPayload
+from comdas.codecs.wavelet import WaveletCodec, WaveletPayload
+from comdas.arrays.duck_array import DuckArray
+
+__all__ = [
+    "Codec",
+    "CompressedPayload",
+    "SVDCodec",
+    "SVDPayload",
+    "WaveletCodec",
+    "WaveletPayload",
+    "DuckArray",
+]
+
+try:
+    from comdas.utils import compress_patch  # requires dascore installed
+    from comdas.io import write_compressed  # noqa: F401 -- also self-registers the COMDAS FiberIO plugin
+
+    __all__ += ["compress_patch", "write_compressed"]
+except ImportError:  # pragma: no cover
+    pass

comdas/arrays/__init__.py

@@ -0,0 +1,9 @@
+"""
+Implements a NumPy-like array (`duck array
+<https://docs.xarray.dev/en/stable/user-guide/duckarrays.html>`_) to hold
+compressed Patch data.
+"""
+
+from comdas.arrays.duck_array import DuckArray
+
+__all__ = ["DuckArray"]

comdas/arrays/duck_array.py

@@ -0,0 +1,323 @@
+"""
+Implements a NumPy-like `duck array
+<https://docs.xarray.dev/en/stable/user-guide/duckarrays.html>`_ for compressed
+Patch data.
+"""
+
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+
+from comdas.codecs.base import Codec, CompressedPayload
+
+
+class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
+    """
+    A NumPy-like duck array backed by a compressed payload and codec defining
+    how the data is compressed and decompressed.
+
+    Reads are performed via ``codec.decode_partial``. Writes go into a
+    sparse overlay dict rather than immediately forcing a complete recompression
+    of the underlying compressed payload. :meth:`consolidate` can be called to
+    merge the overlay into the compressed payload. :meth:`consolidate` will be
+    called automatically when the number of pending writes exceeds the codec's
+    specified maximum number of pending writes.
+
+    If the specified codec implements partial writes efficiently, the codec's
+    partial writes function will always be used to update the compressed
+    payload without accumulation.
+
+    :param payload: The compressed payload backing this duck array.
+    :type payload: CompressedPayload
+    :param codec: The codec defining how the data is compressed and decompressed.
+    :type codec: Codec
+    :param max_pending_writes: Maximum number of pending writes before automatic
+        consolidation.
+    :type max_pending_writes: int or None
+    """
+
+    def __init__(
+        self,
+        payload: CompressedPayload,
+        codec: Codec,
+        *,
+        max_pending_writes: int | None = None,
+    ):
+        self._payload = payload
+        self._codec = codec
+        self._overlay: dict[tuple[int, ...], object] = {}
+        self._max_pending_writes = (
+            max_pending_writes
+            if max_pending_writes is not None
+            else codec.max_pending_writes
+        )
+
+    # Array basics
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """
+        Get the shape of the decompressed/original array.
+
+        :rtype: tuple[int, ...]
+        """
+        return self._payload.original_shape
+
+    @property
+    def dtype(self) -> np.dtype:
+        """
+        Get the data type of the decompressed/original array.
+
+        :rtype: numpy.dtype
+        """
+        return self._payload.dtype
+
+    @property
+    def ndim(self) -> int:
+        """
+        Get the number of dimensions of the decompressed/original array.
+
+        :rtype: int
+        """
+        return len(self.shape)
+
+    def __len__(self) -> int:
+        return self.shape[0]
+
+    @property
+    def payload(self) -> CompressedPayload:
+        """
+        Get the current compressed payload backing this array.
+
+        :rtype: CompressedPayload
+        """
+        return self._payload
+
+    @property
+    def codec(self) -> Codec:
+        """
+        Get the codec used to encode/decode :attr:`payload`.
+
+        :rtype: Codec
+        """
+        return self._codec
+
+    @property
+    def overlay_size(self) -> int:
+        """
+        Get the number of elements currently held in the write overlay.
+
+        :rtype: int
+        """
+        return len(self._overlay)
+
+    def __array__(self, dtype=None) -> np.ndarray:
+        """
+        Fully materialize this array as a dense NumPy array.
+
+        :param dtype: Optional dtype to cast the result to.
+        :type dtype: numpy.dtype or None
+        :returns: The decompressed array, with any pending overlay
+            writes applied.
+        :rtype: numpy.ndarray
+        """
+        out = self._codec.decode(self._payload).copy()
+        for idx, val in self._overlay.items():
+            out[idx] = val
+        if dtype is not None:
+            out = out.astype(dtype, copy=False)
+        return out
+
+    def __array_function__(self, func, types, args, kwargs):
+        """
+        Apply a numpy function to this array. Decompresses the
+        array before applying the NumPy function.
+
+        NOTE: Some codecs support efficient operations directly
+        on the compressed representation. We should provide a way
+        to implement these in the future.
+
+        NOTE: Applying a chain of functions will decompress the array
+        at each step. We should optimize chained operations to avoid
+        this.
+
+        :param func: The NumPy function being dispatched.
+        :param types: Types of the arguments involved.
+        :param args: Positional arguments to ``func``.
+        :param kwargs: Keyword arguments to ``func``.
+        :returns: The result of calling ``func`` on materialized arrays.
+        """
+        args = [np.asarray(a) if isinstance(a, DuckArray) else a for a in args]
+        return func(*args, **kwargs)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        """
+        Materializes the array then applies the ufunc.
+
+        NOTE: We should explore per-codec optimizations.
+
+        :param ufunc: The NumPy ufunc being dispatched.
+        :param method: The ufunc method being called (e.g. ``"__call__"``).
+        :param inputs: Positional inputs to the ufunc.
+        :param kwargs: Keyword arguments to the ufunc.
+        :returns: The result of calling the ufunc on materialized arrays.
+        """
+        inputs = [
+            np.asarray(i) if isinstance(i, DuckArray) else i for i in inputs
+        ]
+        return getattr(ufunc, method)(*inputs, **kwargs)
+
+    def __repr__(self) -> str:
+        ratio = self._codec.compression_ratio(self._payload)
+        return (
+            f"DuckArray(shape={self.shape}, dtype={self.dtype}, "
+            f"codec={type(self._codec).__name__}, ratio={ratio:.1f}x, "
+            f"overlay={self.overlay_size})"
+        )
+
+    def _key_to_index_lists(self, key) -> list[list[int]]:
+        """
+        Normalize a NumPy-style key into per-axis integer index lists.
+
+        :param key: A NumPy-style index/slice key.
+        :returns: One list of concrete integer indices per axis.
+        :rtype: list[list[int]]
+        """
+        if not isinstance(key, tuple):
+            key = (key,)
+        key = key + (slice(None),) * (self.ndim - len(key))
+        axes: list[list[int]] = []
+        for k, dim in zip(key, self.shape):
+            if isinstance(k, (int, np.integer)):
+                axes.append([int(k)])
+            elif isinstance(k, slice):
+                axes.append(list(range(*k.indices(dim))))
+            else:
+                axes.append(np.atleast_1d(np.asarray(k)).tolist())
+        return axes
+
+    @staticmethod
+    def _normalize_key(key, ndim: int) -> tuple:
+        """
+        Pad ``key`` with trailing full slices to length ``ndim``.
+
+        :param key: A NumPy-style index/slice key.
+        :param ndim: Number of dimensions to pad to.
+        :type ndim: int
+        :returns: The normalized, tuple-form key.
+        :rtype: tuple
+        """
+        if not isinstance(key, tuple):
+            key = (key,)
+        return key + (slice(None),) * (ndim - len(key))
+
+    def __getitem__(self, key):
+        """
+        Read values from the array.
+
+        :param key: A NumPy-style index/slice key.
+        :returns: The requested subset of the array.
+        :rtype: numpy.ndarray or scalar
+        """
+        if not self._overlay:
+            return self._codec.decode_partial(self._payload, key)
+
+        axes = self._key_to_index_lists(key)
+        axes_sets = [set(a) for a in axes]
+        touched = [
+            (idx, val)
+            for idx, val in self._overlay.items()
+            if all(idx[d] in axes_sets[d] for d in range(self.ndim))
+        ]
+        base = np.array(
+            self._codec.decode_partial(self._payload, key), copy=True
+        )
+        if not touched:
+            return base
+
+        norm_key = self._normalize_key(key, self.ndim)
+        for idx, val in touched:
+            local = tuple(
+                axes[d].index(idx[d])
+                for d, k in enumerate(norm_key)
+                if not isinstance(k, (int, np.integer))
+            )
+            base[local] = val
+        return base
+
+    def __setitem__(self, key, value) -> None:
+        """
+        Write a value. Attempts to not decompress the entire array if possible.
+
+        Three strategies are tried, in order:
+
+        1. If the codec implements :meth:`~comdas.codecs.base.Codec.partial_write`,
+           it's used directly. In this case, the payload is patched in place.
+           A codec may raise :class:`NotImplementedError` for a particular ``key``
+           it can't handle (e.g. it only supports single elements), in
+           this case it goes to strategy 2 for that call.
+        2. Otherwise, small writes land in a sparse overlay dict,
+           applied on read and folded in by :meth:`consolidate`.
+        3. A write that touches more than
+           :attr:`~comdas.codecs.base.Codec.max_pending_writes`
+           elements or would push the overlay size past it auto-consolidates
+           instead.
+
+        :param key: A NumPy-style index/slice key.
+        :param value: The value(s) to write; broadcast against the
+            shape implied by ``key``, matching normal NumPy assignment.
+        """
+        if self._codec.supports_partial_write():
+            try:
+                self._payload = self._codec.partial_write(
+                    self._payload, key, value
+                )
+                return
+            except NotImplementedError:
+                pass
+
+        axes = self._key_to_index_lists(key)
+        size = 1
+        for a in axes:
+            size *= len(a)
+
+        if (
+            size > self._max_pending_writes
+            or (self.overlay_size + size) > self._max_pending_writes
+        ):
+            dense = self.__array__()
+            dense[key] = value
+            self.consolidate(dense=dense)
+            return
+
+        indices = list(itertools.product(*axes))
+        values = np.broadcast_to(value, tuple(len(a) for a in axes)).ravel()
+        for idx, val in zip(indices, values):
+            self._overlay[idx] = val
+
+    def consolidate(
+        self,
+        *,
+        dense: np.ndarray | None = None,
+        **encode_kwargs,
+    ) -> None:
+        """
+        Fold the write overlay into a freshly-compressed payload.
+
+        By default, re-encodes using whatever compression level the
+        current payload already used.
+
+        :param dense: Precomputed dense array to encode, if the
+            caller already has one.
+        :type dense: numpy.ndarray or None
+        :param encode_kwargs: Additional codec-specific kwargs
+            forwarded to ``codec.encode``.
+        """
+        if dense is None:
+            dense = self.__array__()
+
+        encode_kwargs = self._payload.compression_params
+
+        self._payload = self._codec.encode(dense, **encode_kwargs)
+        self._overlay.clear()

comdas/codecs/__init__.py

@@ -0,0 +1,14 @@
+"""Compression codec interfaces and implementations."""
+
+from comdas.codecs.base import Codec, CompressedPayload
+from comdas.codecs.svd import SVDCodec, SVDPayload
+from comdas.codecs.wavelet import WaveletCodec, WaveletPayload
+
+__all__ = [
+    "Codec",
+    "CompressedPayload",
+    "SVDCodec",
+    "SVDPayload",
+    "WaveletCodec",
+    "WaveletPayload",
+]

comdas/codecs/base.py

@@ -0,0 +1,264 @@
+"""
+Base classes for compression codecs.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class CompressedPayload:
+    """
+    Base container for a codec's compressed payload (representation of an array).
+
+    Codecs should subclass this to hold the arrays/parameters the method requires.
+
+    :ivar original_shape: The shape of the array before compression.
+    :vartype original_shape: tuple[int, ...]
+    :ivar dtype: The data type of the array before compression.
+    :vartype dtype: numpy.dtype
+    """
+
+    original_shape: tuple[int, ...]
+    dtype: np.dtype
+
+    @property
+    def compression_params(self) -> dict:
+        """
+        The parameters used for compression.
+
+        The default implementation returns an empty dict. Codecs with
+        parameters relevant to re-encoding should override this.
+
+        :rtype: dict
+        """
+        return {}
+
+
+class Codec(ABC):
+    """
+    Base class for all codecs.
+
+    A codec is responsible for compressing and decompressing arrays. It knows
+    how to turn a dense NumPy array into a :class:`CompressedPayload` and back,
+    and how to (de)serialize that payload to/from an HDF5 file.
+
+    Subclasses are auto-registered by name on class creation. This allows
+    the reader to automatically recognize and use the appropriate codec by name,
+    without having to import every codec.
+
+    Most codecs do not support random access writes and algorithms can be
+    inefficient at performing many small writes, so writes to the payload are
+    buffered until max_pending_writes is reached. Once this value is exceeded,
+    the entire array will be decompressed, the writes applied, and the
+    array recompressed. By default, this value is set to 10,000.
+
+    If a codec supports random access writes efficiently, it should
+    implement the :meth:`partial_write` method. If this method is
+    implemented, it will always be immediately used for applying writes,
+    bypassing max_pending_writes entirely.
+
+    :cvar name: Unique, human-readable name for the codec.
+    :vartype name: str
+    :cvar version: Codec format version, bumped if the codec's payload
+        structure changes in a way that is not backward-compatible.
+    :vartype version: str
+    :cvar max_pending_writes: Maximum number of pending writes before automatic consolidation.
+        This is only used by codecs that do not implement :meth:`partial_write`.
+    :vartype max_pending_writes: int
+    """
+
+    name: str = ""
+    version: str = "1"
+    max_pending_writes: int = 10_000
+
+    _registry: dict[str, type[Codec]] = {}
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        """
+        Register ``cls`` in the registry under ``cls.name``.
+
+        :param kwargs: Forwarded to :meth:`object.__init_subclass__`.
+        :raises ValueError: If a concrete subclass doesn't set a
+            non-empty ``name``.
+        """
+        super().__init_subclass__(**kwargs)
+        if getattr(cls, "__abstractmethods__", None):
+            return
+        if not cls.name:
+            raise ValueError(
+                f"{cls.__name__} must define a `name` for codec registration."
+            )
+        Codec._registry[cls.name] = cls
+
+    @classmethod
+    def get_registered(cls, name: str) -> type[Codec]:
+        """
+        Look up a codec class by name.
+
+        :param name: The codec's ``name`` attribute.
+        :type name: str
+        :returns: The codec class registered under ``name``.
+        :rtype: type[Codec]
+        :raises KeyError: If no codec is registered under ``name``.
+        """
+        try:
+            return cls._registry[name]
+        except KeyError as err:
+            raise KeyError(
+                f"No codec registered under name {name!r}. Known codecs: {sorted(cls._registry)}"
+            ) from err
+
+    @abstractmethod
+    def encode(self, array: np.ndarray, **kwargs) -> CompressedPayload:
+        """
+        Compress ``array`` into a :class:`CompressedPayload`.
+
+        :param array: The dense array to compress.
+        :type array: numpy.ndarray
+        :param kwargs: Codec-specific compression parameters.
+        :returns: The compressed payload.
+        :rtype: CompressedPayload
+        """
+
+    @abstractmethod
+    def decode(self, payload: CompressedPayload) -> np.ndarray:
+        """
+        Fully decompress the payload.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: CompressedPayload
+        :returns: The decompressed array.
+        :rtype: numpy.ndarray
+        """
+
+    def decode_partial(self, payload: CompressedPayload, key) -> np.ndarray:
+        """
+        Partially decompress the payload, reconstructing only the elements
+        addressed by ``key``.
+
+        The default implementation uses a full decode. Codecs that support
+        partial decompression or random access read/write should overwrite
+        this method.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: CompressedPayload
+        :param key: A NumPy-style index/slice key.
+        :returns: The requested subset of the reconstructed array.
+        :rtype: numpy.ndarray
+        """
+        return self.decode(payload)[key]
+
+    def partial_write(
+        self, payload: CompressedPayload, key, value
+    ) -> CompressedPayload:
+        """
+        Write ``value`` into the compressed representation at ``key``,
+        without a full decode -> update -> re-encode round trip.
+
+        This is an optional method, most codecs do not support random access
+        writes. The default implementation always raises a
+        :class:`NotImplementedError`, which
+        :class:`~comdas.arrays.duck_array.DuckArray`
+        interprets as "this codec or write attempt doesn't
+        support true partial writes" and falls back to its overlay +
+        :attr:`default_overlay_flush_threshold` strategy instead.
+
+        A codec should implement this for the cases it can handle
+        cheaply and raise :class:`NotImplementedError` for cases it can't.
+        The fallback applies per call, not just
+        per codec. See :meth:`supports_partial_write`.
+
+        Implementations must apply exactly the same assignment
+        semantics as ``dense_array[key] = value`` would on the fully
+        decompressed array, including NumPy's broadcasting rules.
+
+        :param payload: The payload before the write.
+        :type payload: CompressedPayload
+        :param key: A NumPy-style index/slice key.
+        :param value: The value(s) to write, broadcast against the
+            shape implied by ``key`` exactly as plain NumPy assignment
+            would.
+        :returns: The payload reflecting the write.
+        :rtype: CompressedPayload
+        :raises NotImplementedError: Always, unless overridden.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support partial_write()."
+        )
+
+    def supports_partial_write(self) -> bool:
+        """
+        Whether this codec has overridden :meth:`partial_write`.
+
+        :returns: True if :meth:`partial_write` is overridden.
+        :rtype: bool
+        """
+        return type(self).partial_write is not Codec.partial_write
+
+    @abstractmethod
+    def compressed_size_bytes(self, payload: CompressedPayload) -> int:
+        """
+        Get the total size, in bytes, of everything stored in ``payload``.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: CompressedPayload
+        :returns: Total size in bytes of the payload's stored arrays.
+        :rtype: int
+        """
+
+    def compression_ratio(self, payload: CompressedPayload) -> float:
+        """
+        Helper to determine the compression ratio of the payload.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: CompressedPayload
+        :returns: Original bytes divided by compressed bytes.
+        :rtype: float
+        """
+        original_bytes = (
+            int(np.prod(payload.original_shape))
+            * np.dtype(payload.dtype).itemsize
+        )
+        return original_bytes / self.compressed_size_bytes(payload)
+
+    @abstractmethod
+    def payload_to_group(self, payload: CompressedPayload, group) -> None:
+        """
+        Write this codec's own arrays/params into an open HDF5 group.
+
+        Only codec-specific data belongs here. The general ``COMDAS``
+        container already writes ``original_shape``, ``dtype``,
+        ``codec_name``, and ``codec_version`` at the container level
+        before calling this. Don't duplicate those fields.
+
+        :param payload: The payload to serialize.
+        :type payload: CompressedPayload
+        :param group: An open, writable ``h5py.Group`` dedicated to
+            this one patch.
+        :type group: h5py.Group
+        """
+
+    @abstractmethod
+    def payload_from_group(
+        self, group, *, original_shape: tuple[int, ...], dtype: np.dtype
+    ) -> CompressedPayload:
+        """
+        Reconstruct this codec's payload from an HDF5 group.
+
+        :param group: An open, readable ``h5py.Group`` previously
+            written by :meth:`payload_to_group`.
+        :type group: h5py.Group
+        :param original_shape: The array shape, as already read from
+            the container's own (codec-independent) attributes.
+        :type original_shape: tuple[int, ...]
+        :param dtype: The array dtype, as already read from the
+            container's own (codec-independent) attributes.
+        :type dtype: numpy.dtype
+        :returns: The reconstructed payload.
+        :rtype: CompressedPayload
+        """