ComDAS 0.1.0__tar.gz → 0.2.0__tar.gz

{comdas-0.1.0 → comdas-0.2.0}/.github/workflows/publish.yml

@@ -3,6 +3,7 @@ name: Upload Python Package
 on:
   release:
     types: [published]
+  workflow_dispatch:
 
 permissions:
   contents: read

{comdas-0.1.0 → comdas-0.2.0}/.github/workflows/testing.yml

@@ -9,6 +9,7 @@ on:
             - ".github/workflows/testing.yml"
     pull_request:
         branches: [main]
+    workflow_dispatch:
 
 permissions:
     contents: read

{comdas-0.1.0 → comdas-0.2.0}/ComDAS.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ComDAS
-Version: 0.1.0
+Version: 0.2.0
 Summary: DASCORE companion package that adds support for compressed DAS data
 Author-email: Byron-Selvage <bselvage@mines.edu>
 License: LGPL-3.0-or-later
@@ -44,8 +44,8 @@ With ComDAS you can:
 
 | Codec | Main setting | Reference |
 | --- | --- | --- |
-| [SVD](https://byron-selvage.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
-| [Wavelet](https://byron-selvage.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
+| [SVD](https://dasdae.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
+| [Wavelet](https://dasdae.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
 
 ## Installation
 
@@ -57,12 +57,12 @@ pip install comdas
 
 ## Documentation
 
-Full documentation is available at [byron-selvage.github.io/ComDAS](https://byron-selvage.github.io/ComDAS/), including:
+Full documentation is available at [dasdae.github.io/ComDAS](https://dasdae.github.io/ComDAS/), including:
 
-- [Quick start](https://byron-selvage.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
-- [Codecs](https://byron-selvage.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
-- [Examples](https://byron-selvage.github.io/ComDAS/examples/compression_comparison.html)
-- [API reference](https://byron-selvage.github.io/ComDAS/api/index.html)
+- [Quick start](https://dasdae.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
+- [Codecs](https://dasdae.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
+- [Examples](https://dasdae.github.io/ComDAS/examples/compression_comparison.html)
+- [API reference](https://dasdae.github.io/ComDAS/api/index.html)
 
 ## License
 

{comdas-0.1.0 → comdas-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ComDAS
-Version: 0.1.0
+Version: 0.2.0
 Summary: DASCORE companion package that adds support for compressed DAS data
 Author-email: Byron-Selvage <bselvage@mines.edu>
 License: LGPL-3.0-or-later
@@ -44,8 +44,8 @@ With ComDAS you can:
 
 | Codec | Main setting | Reference |
 | --- | --- | --- |
-| [SVD](https://byron-selvage.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
-| [Wavelet](https://byron-selvage.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
+| [SVD](https://dasdae.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
+| [Wavelet](https://dasdae.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
 
 ## Installation
 
@@ -57,12 +57,12 @@ pip install comdas
 
 ## Documentation
 
-Full documentation is available at [byron-selvage.github.io/ComDAS](https://byron-selvage.github.io/ComDAS/), including:
+Full documentation is available at [dasdae.github.io/ComDAS](https://dasdae.github.io/ComDAS/), including:
 
-- [Quick start](https://byron-selvage.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
-- [Codecs](https://byron-selvage.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
-- [Examples](https://byron-selvage.github.io/ComDAS/examples/compression_comparison.html)
-- [API reference](https://byron-selvage.github.io/ComDAS/api/index.html)
+- [Quick start](https://dasdae.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
+- [Codecs](https://dasdae.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
+- [Examples](https://dasdae.github.io/ComDAS/examples/compression_comparison.html)
+- [API reference](https://dasdae.github.io/ComDAS/api/index.html)
 
 ## License
 

{comdas-0.1.0 → comdas-0.2.0}/README.md

@@ -19,8 +19,8 @@ With ComDAS you can:
 
 | Codec | Main setting | Reference |
 | --- | --- | --- |
-| [SVD](https://byron-selvage.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
-| [Wavelet](https://byron-selvage.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
+| [SVD](https://dasdae.github.io/ComDAS/codecs/svd.html) | `rank` or retained `energy` | [NumPy SVD](https://numpy.org/doc/1.22/reference/generated/numpy.linalg.svd.html) |
+| [Wavelet](https://dasdae.github.io/ComDAS/codecs/wavelet.html) | `keep_fraction` or `threshold` | [PyWavelets](https://doi.org/10.21105/joss.01237); [DeVore et al. (1992)](https://doi.org/10.1109/18.119733) |
 
 ## Installation
 
@@ -32,12 +32,12 @@ pip install comdas
 
 ## Documentation
 
-Full documentation is available at [byron-selvage.github.io/ComDAS](https://byron-selvage.github.io/ComDAS/), including:
+Full documentation is available at [dasdae.github.io/ComDAS](https://dasdae.github.io/ComDAS/), including:
 
-- [Quick start](https://byron-selvage.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
-- [Codecs](https://byron-selvage.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
-- [Examples](https://byron-selvage.github.io/ComDAS/examples/compression_comparison.html)
-- [API reference](https://byron-selvage.github.io/ComDAS/api/index.html)
+- [Quick start](https://dasdae.github.io/ComDAS/quickstart.html): compress a patch, save a compressed file, and read it back
+- [Codecs](https://dasdae.github.io/ComDAS/codecs/svd.html): configuration options and examples for each codec
+- [Examples](https://dasdae.github.io/ComDAS/examples/compression_comparison.html)
+- [API reference](https://dasdae.github.io/ComDAS/api/index.html)
 
 ## License
 

{comdas-0.1.0 → comdas-0.2.0}/comdas/arrays/duck_array.py

@@ -6,13 +6,80 @@ Patch data.
 
 from __future__ import annotations
 
-import itertools
-
 import numpy as np
 
 from comdas.codecs.base import Codec, CompressedPayload
 
 
+def _materialize(obj):
+    """
+    Replace every :class:`DuckArray` in ``obj`` with a dense NumPy array,
+    recursing into lists, tuples, and dicts (e.g. the list passed to
+    ``np.concatenate``).
+
+    :param obj: An argument (or container of arguments) to a NumPy call.
+    :returns: ``obj`` with DuckArrays materialized.
+    """
+    if isinstance(obj, DuckArray):
+        return np.asarray(obj)
+    if type(obj) is list:
+        return [_materialize(o) for o in obj]
+    if type(obj) is tuple:
+        return tuple(_materialize(o) for o in obj)
+    if type(obj) is dict:
+        return {k: _materialize(v) for k, v in obj.items()}
+    return obj
+
+
+def _densify_out(kwargs: dict) -> tuple[dict, list]:
+    """
+    Swap any DuckArray in an ``out=`` kwarg for a dense temporary buffer.
+
+    The buffer starts as a copy of the DuckArray's current values, so
+    elements a ``where=`` mask leaves untouched are preserved.
+
+    :param kwargs: Keyword arguments to a NumPy function/ufunc.
+    :returns: ``(new_kwargs, pairs)`` where ``pairs`` is a list of
+        ``(duck_array, dense_buffer)`` to pass to :func:`_write_back_out`.
+    """
+    out = kwargs.get("out")
+    if out is None:
+        return kwargs, []
+    is_tuple = isinstance(out, tuple)
+    pairs = []
+    new_out = []
+    for o in out if is_tuple else (out,):
+        if isinstance(o, DuckArray):
+            dense = np.asarray(o)
+            pairs.append((o, dense))
+            new_out.append(dense)
+        else:
+            new_out.append(o)
+    new_kwargs = dict(kwargs)
+    new_kwargs["out"] = tuple(new_out) if is_tuple else new_out[0]
+    return new_kwargs, pairs
+
+
+def _write_back_out(result, pairs: list):
+    """
+    Fold dense ``out`` buffers back into their DuckArrays (re-encoding
+    them) and return those DuckArrays in place of the buffers.
+
+    :param result: The value returned by the NumPy call.
+    :param pairs: As returned by :func:`_densify_out`.
+    :returns: ``result``, with dense buffers replaced by their DuckArrays.
+    """
+    if not pairs:
+        return result
+    swap = {}
+    for duck, dense in pairs:
+        duck.consolidate(dense=dense)
+        swap[id(dense)] = duck
+    if isinstance(result, tuple):
+        return tuple(swap.get(id(r), r) for r in result)
+    return swap.get(id(result), result)
+
+
 class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
     """
     A NumPy-like duck array backed by a compressed payload and codec defining
@@ -131,8 +198,13 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
 
     def __array_function__(self, func, types, args, kwargs):
         """
-        Apply a numpy function to this array. Decompresses the
-        array before applying the NumPy function.
+        Apply a numpy function to this array. Decompresses any DuckArrays
+        among the arguments (including inside lists/tuples/dicts, e.g. for
+        ``np.concatenate``) before applying the NumPy function.
+
+        If a DuckArray is given as ``out=``, the result is computed into a
+        dense buffer and then re-encoded into that DuckArray (see
+        :meth:`consolidate`), which is returned in place of the buffer.
 
         NOTE: Some codecs support efficient operations directly
         on the compressed representation. We should provide a way
@@ -148,13 +220,19 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
         :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)
+        kwargs, out_pairs = _densify_out(kwargs)
+        result = func(*_materialize(args), **_materialize(kwargs))
+        return _write_back_out(result, out_pairs)
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         """
         Materializes the array then applies the ufunc.
 
+        DuckArrays given as ``out=`` (including via in-place operators such
+        as ``+=``) or as the target of ``ufunc.at`` are updated by computing
+        into a dense buffer and re-encoding it (see :meth:`consolidate`).
+        Re-encoding with a lossy codec can introduce reconstruction error.
+
         NOTE: We should explore per-codec optimizations.
 
         :param ufunc: The NumPy ufunc being dispatched.
@@ -163,10 +241,20 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
         :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)
+        kwargs, out_pairs = _densify_out(kwargs)
+
+        # ufunc.at modifies its first input in place.
+        if method == "at" and isinstance(inputs[0], DuckArray):
+            target = inputs[0]
+            dense = np.asarray(target)
+            ufunc.at(dense, *_materialize(inputs[1:]), **_materialize(kwargs))
+            target.consolidate(dense=dense)
+            return None
+
+        result = getattr(ufunc, method)(
+            *_materialize(inputs), **_materialize(kwargs)
+        )
+        return _write_back_out(result, out_pairs)
 
     def __repr__(self) -> str:
         ratio = self._codec.compression_ratio(self._payload)
@@ -176,41 +264,25 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
             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]]
+    def _selected_positions(self, key) -> tuple[np.ndarray, ...]:
         """
-        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:
+        Find the physical position of every element selected by ``key``.
+
+        NumPy itself applies ``key`` (to zero-copy broadcast views of
+        per-axis index grids), so the result follows NumPy's indexing
+        semantics: negative indices, bounds checks, slices,
+        broadcast/paired advanced indices, ellipsis, newaxis, and boolean
+        masks.
+
+        :param key: A NumPy-style index key.
+        :returns: One integer array per axis, each with the shape of
+            ``dense[key]``, giving the coordinates of each
+            selected element.
+        :rtype: tuple[numpy.ndarray, ...]
+        :raises IndexError: If ``key`` is invalid for this array's shape.
         """
-        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))
+        grids = np.indices(self.shape, sparse=True)
+        return tuple(np.broadcast_to(g, self.shape)[key] for g in grids)
 
     def __getitem__(self, key):
         """
@@ -223,28 +295,24 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
         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))
-        ]
+        positions = self._selected_positions(key)
         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
+
+        # Match pending overlay writes to the selected elements by flat
+        # physical index.
+        selected = np.ravel_multi_index(positions, self.shape)
+        overlay_idx = np.ravel_multi_index(
+            tuple(np.array(list(self._overlay)).T), self.shape
+        )
+        overlay_vals = np.array(list(self._overlay.values()), dtype=self.dtype)
+        hit = np.isin(selected, overlay_idx)
+        if hit.any():
+            order = np.argsort(overlay_idx)
+            loc = np.searchsorted(overlay_idx, selected[hit], sorter=order)
+            base[hit] = overlay_vals[order[loc]]
+        return base[()] if base.ndim == 0 else base
 
     def __setitem__(self, key, value) -> None:
         """
@@ -277,10 +345,9 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
             except NotImplementedError:
                 pass
 
-        axes = self._key_to_index_lists(key)
-        size = 1
-        for a in axes:
-            size *= len(a)
+        positions = self._selected_positions(key)
+        target_shape = positions[0].shape
+        size = positions[0].size
 
         if (
             size > self._max_pending_writes
@@ -291,10 +358,21 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
             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
+        values = np.asarray(value).astype(self.dtype, copy=False)
+        # Like NumPy assignment, allow extra leading length-1 dims on value.
+        while values.ndim > len(target_shape) and values.shape[0] == 1:
+            values = values[0]
+        try:
+            values = np.broadcast_to(values, target_shape)
+        except ValueError:
+            raise ValueError(
+                f"could not broadcast input array from shape "
+                f"{np.shape(value)} into shape {target_shape}"
+            ) from None
+
+        flat_positions = [p.ravel() for p in positions]
+        for idx, val in zip(zip(*flat_positions), values.ravel()):
+            self._overlay[tuple(int(i) for i in idx)] = val
 
     def consolidate(
         self,
@@ -305,19 +383,22 @@ class DuckArray(np.lib.mixins.NDArrayOperatorsMixin):
         """
         Fold the write overlay into a freshly-compressed payload.
 
-        By default, re-encodes using whatever compression level the
-        current payload already used.
+        By default, re-encodes using whatever compression parameters the
+        current payload already used. Any ``encode_kwargs`` given override
+        those (see :meth:`~comdas.codecs.base.Codec.merge_compression_params`).
 
         :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``.
+        :param encode_kwargs: Codec-specific kwargs forwarded to
+            ``codec.encode``, overriding the current payload's parameters.
         """
         if dense is None:
             dense = self.__array__()
 
-        encode_kwargs = self._payload.compression_params
+        params = self._codec.merge_compression_params(
+            self._payload.compression_params, encode_kwargs
+        )
 
-        self._payload = self._codec.encode(dense, **encode_kwargs)
+        self._payload = self._codec.encode(dense, **params)
         self._overlay.clear()

{comdas-0.1.0 → comdas-0.2.0}/comdas/codecs/base.py

@@ -70,11 +70,17 @@ class Codec(ABC):
     :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
+    :cvar exclusive_params: Names of mutually exclusive ``encode`` parameters
+        that select the compression mode (e.g. ``("rank", "energy")``). When
+        a caller overrides one of them, the others are dropped from the
+        payload's stored parameters. See :meth:`merge_compression_params`.
+    :vartype exclusive_params: tuple[str, ...]
     """
 
     name: str = ""
     version: str = "1"
     max_pending_writes: int = 10_000
+    exclusive_params: tuple[str, ...] = ()
 
     _registry: dict[str, type[Codec]] = {}
 
@@ -211,20 +217,55 @@ class Codec(ABC):
         :rtype: int
         """
 
+    def merge_compression_params(self, stored: dict, overrides: dict) -> dict:
+        """
+        Combine a payload's stored compression parameters with explicit
+        overrides, for re-encoding (e.g. in
+        :meth:`~comdas.arrays.duck_array.DuckArray.consolidate`).
+
+        Overrides win. If an override selects one of the mutually exclusive
+        :attr:`exclusive_params`, the other exclusive params are dropped
+        from ``stored`` so the old mode doesn't conflict with the new one.
+        Conflicts *within* ``overrides`` are left for :meth:`encode` to
+        reject.
+
+        :param stored: The payload's ``compression_params``.
+        :type stored: dict
+        :param overrides: Explicit keyword arguments from the caller.
+        :type overrides: dict
+        :returns: The parameters to pass to :meth:`encode`.
+        :rtype: dict
+        """
+        merged = dict(stored)
+        if any(k in self.exclusive_params for k in overrides):
+            for k in self.exclusive_params:
+                merged.pop(k, None)
+        merged.update(overrides)
+        return merged
+
     def compression_ratio(self, payload: CompressedPayload) -> float:
         """
-        Helper to determine the compression ratio of the payload.
+        Helper to determine the compression ratio of the payload's stored
+        arrays.
+
+        This only accounts for the payload's arrays, not metadata or HDF5
+        container overhead, so it is not the on-disk file compression ratio.
 
         :param payload: A payload previously produced by :meth:`encode`.
         :type payload: CompressedPayload
-        :returns: Original bytes divided by compressed bytes.
+        :returns: Original bytes divided by compressed bytes. ``inf`` if the
+            payload stores no bytes for a non-empty array, and ``nan`` if the
+            original array is empty too.
         :rtype: float
         """
         original_bytes = (
             int(np.prod(payload.original_shape))
             * np.dtype(payload.dtype).itemsize
         )
-        return original_bytes / self.compressed_size_bytes(payload)
+        compressed_bytes = self.compressed_size_bytes(payload)
+        if compressed_bytes == 0:
+            return float("inf") if original_bytes > 0 else float("nan")
+        return original_bytes / compressed_bytes
 
     @abstractmethod
     def payload_to_group(self, payload: CompressedPayload, group) -> None:

{comdas-0.1.0 → comdas-0.2.0}/comdas/codecs/svd.py

@@ -78,6 +78,7 @@ class SVDCodec(Codec):
 
     # Writes are expensive with SVD
     max_pending_writes = 10_000
+    exclusive_params = ("rank", "energy")
 
     def __init__(self, rank: int | None = None, energy: float | None = None):
         if rank is not None and energy is not None:
@@ -166,26 +167,34 @@ class SVDCodec(Codec):
 
         :param payload: A payload previously produced by :meth:`encode`.
         :type payload: SVDPayload
-        :param key: An int/slice index in the NumPy-style.
+        :param key: A NumPy-style index. Up to two ints/slices take the
+            fast path; anything else falls back to a full decode.
         :returns: The requested subset of the reconstructed array.
         :rtype: numpy.ndarray
+        :raises IndexError: If ``key`` is invalid for a 2D array (e.g. too
+            many indices or out of bounds), as with NumPy.
         """
         if not isinstance(key, tuple):
             key = (key,)
-        key = key + (slice(None),) * (2 - len(key))
-        row_key, col_key = key[0], key[1]
-
-        # If using complex indexing (anything more than a simple int or slice)
-        # Fall back to full decode
-        if not (
-            isinstance(row_key, (int, np.integer, slice))
-            and isinstance(col_key, (int, np.integer, slice))
-        ):
+
+        # Validate the key against the array shape using NumPy's rules.
+        np.broadcast_to(np.empty((), dtype=bool), payload.original_shape)[key]
+
+        def _is_simple(k) -> bool:
+            return isinstance(k, (int, np.integer, slice)) and not isinstance(
+                k, (bool, np.bool_)
+            )
+
+        # If using complex indexing (anything more than up to two simple
+        # ints/slices, e.g. arrays, ellipsis, newaxis) fall back to full decode
+        if len(key) > 2 or not all(_is_simple(k) for k in key):
             warn(
                 "Only simple int/slice indexing is supported for partial decoding. Fully decoding entire array and then applying the index."
             )
             return self.decode(payload)[key]
 
+        key = key + (slice(None),) * (2 - len(key))
+        row_key, col_key = key
         row_is_int = isinstance(row_key, (int, np.integer))
         col_is_int = isinstance(col_key, (int, np.integer))
 

{comdas-0.1.0 → comdas-0.2.0}/comdas/codecs/wavelet.py

@@ -27,12 +27,16 @@ class WaveletPayload(CompressedPayload):
     :vartype indices: numpy.ndarray
     :ivar values: Values of retained coefficients.
     :vartype values: numpy.ndarray
+    :ivar threshold: The absolute threshold used to select coefficients, if
+        encoded in threshold mode. ``None`` if encoded with ``keep_fraction``.
+    :vartype threshold: float or None
     """
 
     wavelet: str
     mode: str
     indices: np.ndarray
     values: np.ndarray
+    threshold: float | None = None
 
     @property
     def coeff_shapes(self) -> list:
@@ -63,13 +67,18 @@ class WaveletPayload(CompressedPayload):
         """
         The compression parameters used for this payload.
 
+        Threshold-encoded payloads report their ``threshold``, so re-encoding
+        applies the same rule (a derived ``keep_fraction`` could be 0, which
+        isn't a valid input). Otherwise the retained ``keep_fraction``.
+
         :rtype: dict
         """
-        return {
-            "wavelet": self.wavelet,
-            "mode": self.mode,
-            "keep_fraction": self.values.size / self.n_coeffs,
-        }
+        params = {"wavelet": self.wavelet, "mode": self.mode}
+        if self.threshold is not None:
+            params["threshold"] = self.threshold
+        else:
+            params["keep_fraction"] = self.values.size / self.n_coeffs
+        return params
 
 
 class WaveletCodec(Codec):
@@ -98,6 +107,7 @@ class WaveletCodec(Codec):
     """
 
     name = "WAVELET"
+    exclusive_params = ("keep_fraction", "threshold")
     version = "1"
 
     def __init__(
@@ -191,6 +201,7 @@ class WaveletCodec(Codec):
             mode=mode,
             indices=indices.astype(index_dtype),
             values=np.ascontiguousarray(flat[indices]),
+            threshold=float(threshold) if threshold is not None else None,
         )
 
     def decode(self, payload: WaveletPayload) -> np.ndarray:
@@ -244,6 +255,8 @@ class WaveletCodec(Codec):
         group.create_dataset("values", data=payload.values)
         group.attrs["wavelet"] = payload.wavelet
         group.attrs["mode"] = payload.mode
+        if payload.threshold is not None:
+            group.attrs["threshold"] = payload.threshold
 
     def payload_from_group(
         self, group, *, original_shape: tuple[int, ...], dtype: np.dtype
@@ -270,4 +283,9 @@ class WaveletCodec(Codec):
             mode=str(group.attrs["mode"]),
             indices=group["indices"][:],
             values=group["values"][:],
+            threshold=(
+                float(group.attrs["threshold"])
+                if "threshold" in group.attrs
+                else None
+            ),
         )