ComDAS 0.1.0__py3-none-any.whl

comdas/codecs/svd.py

@@ -0,0 +1,258 @@
+"""
+Truncated Singular Value Decomposition (SVD) codec for 2D arrays.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from warnings import warn
+
+import numpy as np
+
+from comdas.codecs.base import Codec, CompressedPayload
+
+
+@dataclass
+class SVDPayload(CompressedPayload):
+    """
+    Truncated-SVD representation of a 2D array:
+    ``A ~= (U * s) @ Vt``.
+
+    :ivar U: Left singular vectors, shape ``(m, k)``.
+    :vartype U: numpy.ndarray
+    :ivar s: Singular values, shape ``(k,)``.
+    :vartype s: numpy.ndarray
+    :ivar Vt: Right singular vectors (transposed), shape ``(k, n)``.
+    :vartype Vt: numpy.ndarray
+    """
+
+    U: np.ndarray
+    s: np.ndarray
+    Vt: np.ndarray
+
+    @property
+    def rank(self) -> int:
+        """
+        The truncation rank ``k`` used on this payload.
+
+        :rtype: int
+        """
+        return self.s.shape[0]
+
+    @property
+    def compression_params(self) -> dict:
+        """
+        The compression parameters used for this payload.
+
+        :rtype: dict
+        """
+        return {
+            "rank": self.rank,
+        }
+
+
+class SVDCodec(Codec):
+    """
+    Truncated-Singular Value Decomposition (SVD) codec for 2D arrays.
+
+    This codec compresses a 2D array by computing its truncated singular value
+    decomposition (SVD) and storing the left singular vectors, singular
+    values, and right singular vectors (transposed) as a :class:`SVDPayload`.
+
+    Truncation rank can be chosen either as a fixed integer ``rank`` or
+    as a retained energy fraction ``energy`` in ``(0,1]``. Exactly one option
+    must be specified, either as a codec-level default or per-call parameter.
+
+    SVD does not implement partial writes efficiently, so the full array is
+    re-encoded every 10,000 pending writes (or on demand).
+
+    :param rank: Default truncation rank ``k`` to use if not specified per-call.
+    :type rank: int | None
+    :param energy: Default retained energy fraction to use if not specified per-call.
+    :type energy: float | None
+    :raises ValueError: If both of ``rank`` and ``energy`` are specified.
+    """
+
+    name = "SVD"
+    version = "1"
+
+    # Writes are expensive with SVD
+    max_pending_writes = 10_000
+
+    def __init__(self, rank: int | None = None, energy: float | None = None):
+        if rank is not None and energy is not None:
+            raise ValueError(
+                "Provide at most one of rank / energy to act as codec default."
+            )
+        self.default_rank = rank
+        self.default_energy = energy
+
+    def encode(
+        self,
+        array: np.ndarray,
+        *,
+        rank: int | None = None,
+        energy: float | None = None,
+        **kwargs,
+    ) -> SVDPayload:
+        """
+        Compress a 2D array via truncated SVD.
+
+        :param array: The dense 2D array to compress.
+        :type array: numpy.ndarray
+        :param rank: Fixed truncation rank. Overrides
+            ``self.default_rank`` for this call.
+        :type rank: int or None
+        :param energy: Retained energy fraction in ``(0, 1]``.
+            Overrides ``self.default_energy`` for this call.
+        :type energy: float or None
+        :param kwargs: Accepted and ignored for compatibility.
+        :returns: The truncated-SVD payload.
+        :rtype: SVDPayload
+        :raises ValueError: If ``array`` is not 2D, or if both
+            ``rank`` and ``energy`` are specified.
+        """
+        array = np.asarray(array)
+        if array.ndim != 2:
+            raise ValueError(
+                f"SVDCodec only supports 2D arrays, got shape {array.shape}"
+            )
+
+        if rank is not None and energy is not None:
+            raise ValueError(
+                "Specify only one of `rank` or `energy` per call or use the codec default."
+            )
+        if rank is None and energy is None:
+            rank = self.default_rank
+            energy = self.default_energy
+        if (rank is None) == (energy is None):
+            raise ValueError(
+                "Exactly one of `rank` or `energy` must be specified, "
+                "either here or as a codec default."
+            )
+
+        U, s, Vt = np.linalg.svd(array, full_matrices=False)
+
+        if rank is None:
+            cumulative = np.cumsum(s**2)
+            cumulative /= cumulative[-1]
+            rank = int(np.searchsorted(cumulative, energy) + 1)
+            rank = min(rank, s.shape[0])
+        rank = max(1, min(rank, s.shape[0]))
+
+        return SVDPayload(
+            original_shape=array.shape,
+            dtype=array.dtype,
+            U=np.ascontiguousarray(U[:, :rank]),
+            s=np.ascontiguousarray(s[:rank]),
+            Vt=np.ascontiguousarray(Vt[:rank, :]),
+        )
+
+    def decode(self, payload: SVDPayload) -> np.ndarray:
+        """
+        Fully reconstruct the dense array as ``(U * s) @ Vt``.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: SVDPayload
+        :returns: The reconstructed dense array.
+        :rtype: numpy.ndarray
+        """
+        out = (payload.U * payload.s) @ payload.Vt
+        return out.astype(payload.dtype, copy=False)
+
+    def decode_partial(self, payload: SVDPayload, key) -> np.ndarray:
+        """
+        Reconstruct only the requested rows/cols of ``payload``.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: SVDPayload
+        :param key: An int/slice index in the NumPy-style.
+        :returns: The requested subset of the reconstructed array.
+        :rtype: numpy.ndarray
+        """
+        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))
+        ):
+            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]
+
+        row_is_int = isinstance(row_key, (int, np.integer))
+        col_is_int = isinstance(col_key, (int, np.integer))
+
+        U_sub = payload.U[row_key, :]
+        Vt_sub = payload.Vt[:, col_key]
+        if row_is_int:
+            U_sub = U_sub[np.newaxis, :]
+        if col_is_int:
+            Vt_sub = Vt_sub[:, np.newaxis]
+
+        out = (U_sub * payload.s) @ Vt_sub
+        out = out.astype(payload.dtype, copy=False)
+
+        if row_is_int:
+            out = out[0]
+        if col_is_int:
+            out = out[..., 0]
+        return out
+
+    def compressed_size_bytes(self, payload: SVDPayload) -> int:
+        """
+        Get the total number of bytes used by ``U``, ``s``, and ``Vt``.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: SVDPayload
+        :returns: Combined byte size of the three stored arrays.
+        :rtype: int
+        """
+        return payload.U.nbytes + payload.s.nbytes + payload.Vt.nbytes
+
+    def payload_to_group(self, payload: SVDPayload, group) -> None:
+        """
+        Write ``U``, ``s``, ``Vt``, and ``rank`` into an open ``h5py.Group``.
+
+        :param payload: The payload to serialize.
+        :type payload: SVDPayload
+        :param group: An open, writable ``h5py.Group`` dedicated to
+            this one patch.
+        :type group: h5py.Group
+        """
+        group.create_dataset("U", data=payload.U)
+        group.create_dataset("s", data=payload.s)
+        group.create_dataset("Vt", data=payload.Vt)
+        group.attrs["rank"] = payload.rank
+
+    def payload_from_group(
+        self, group, *, original_shape: tuple[int, ...], dtype: np.dtype
+    ) -> SVDPayload:
+        """
+        Construct an :class:`SVDPayload` 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 read from the
+            container's own attributes.
+        :type original_shape: tuple[int, ...]
+        :param dtype: The array dtype, as read from the container's
+            own attributes.
+        :type dtype: numpy.dtype
+        :returns: The reconstructed payload.
+        :rtype: SVDPayload
+        """
+        return SVDPayload(
+            original_shape=original_shape,
+            dtype=dtype,
+            U=group["U"][:],
+            s=group["s"][:],
+            Vt=group["Vt"][:],
+        )

comdas/codecs/wavelet.py

@@ -0,0 +1,273 @@
+"""
+Wavelet compression via PyWavelets.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import pywt
+
+from comdas.codecs.base import Codec, CompressedPayload
+
+
+@dataclass
+class WaveletPayload(CompressedPayload):
+    """
+    Thresholded wavelet-coefficient representation of an array.
+
+    Retained coefficients are stored as ``(indices, values)`` pairs.
+
+    :ivar wavelet: Name of the PyWavelets wavelet.
+    :vartype wavelet: str
+    :ivar mode: PyWavelets signal extension mode.
+    :vartype mode: str
+    :ivar indices: Flat indices of retained coefficients.
+    :vartype indices: numpy.ndarray
+    :ivar values: Values of retained coefficients.
+    :vartype values: numpy.ndarray
+    """
+
+    wavelet: str
+    mode: str
+    indices: np.ndarray
+    values: np.ndarray
+
+    @property
+    def coeff_shapes(self) -> list:
+        """
+        Get the coefficient shapes of the decomposition before thresholding.
+
+        :rtype: list
+        """
+        return pywt.wavedecn_shapes(
+            self.original_shape, self.wavelet, mode=self.mode
+        )
+
+    @property
+    def n_coeffs(self) -> int:
+        """
+        Total number of coefficients in the full decomposition without thresholding.
+
+        :rtype: int
+        """
+        shapes = self.coeff_shapes
+        total = int(np.prod(shapes[0]))
+        for details in shapes[1:]:
+            total += sum(int(np.prod(s)) for s in details.values())
+        return total
+
+    @property
+    def compression_params(self) -> dict:
+        """
+        The compression parameters used for this payload.
+
+        :rtype: dict
+        """
+        return {
+            "wavelet": self.wavelet,
+            "mode": self.mode,
+            "keep_fraction": self.values.size / self.n_coeffs,
+        }
+
+
+class WaveletCodec(Codec):
+    """
+    Wavelet-thresholding codec.
+
+    The array is decomposed with a multilevel n-dimensional discrete wavelet
+    transform (:func:`pywt.wavedecn`) at the maximum useful level for the
+    array shape, and only the largest-magnitude coefficients are retained.
+
+    Coefficients to keep can be chosen either as a fraction ``keep_fraction`` in
+    ``(0, 1]`` of all coefficients, or as an absolute magnitude
+    ``threshold`` where coefficients with ``|c| > threshold`` are kept. Exactly
+    one option must be specified, either as a codec-level default or
+    per-call parameter.
+
+    :param wavelet: PyWavelets wavelet name. (default: ``"db4"``)
+    :type wavelet: str
+    :param mode: PyWavelets signal extension mode. (default: ``"symmetric"``)
+    :type mode: str
+    :param keep_fraction: Default fraction of coefficients to retain.
+    :type keep_fraction: float | None
+    :param threshold: Default absolute coefficient magnitude threshold.
+    :type threshold: float | None
+    :raises ValueError: If both of ``keep_fraction`` and ``threshold`` are specified.
+    """
+
+    name = "WAVELET"
+    version = "1"
+
+    def __init__(
+        self,
+        wavelet: str = "db4",
+        mode: str = "symmetric",
+        keep_fraction: float | None = None,
+        threshold: float | None = None,
+    ):
+        if keep_fraction is not None and threshold is not None:
+            raise ValueError(
+                "Provide at most one of keep_fraction / threshold to act as codec default."
+            )
+        self.default_wavelet = wavelet
+        self.default_mode = mode
+        self.default_keep_fraction = keep_fraction
+        self.default_threshold = threshold
+
+    def encode(
+        self,
+        array: np.ndarray,
+        *,
+        wavelet: str | None = None,
+        mode: str | None = None,
+        keep_fraction: float | None = None,
+        threshold: float | None = None,
+        **kwargs,
+    ) -> WaveletPayload:
+        """
+        Compress an array via wavelet coefficient thresholding.
+
+        :param array: The dense array to compress.
+        :type array: numpy.ndarray
+        :param wavelet: Overrides ``self.default_wavelet`` for this call.
+        :type wavelet: str or None
+        :param mode: Overrides ``self.default_mode`` for this call.
+        :type mode: str or None
+        :param keep_fraction: Fraction of coefficients to retain, in ``(0, 1]``.
+            Overrides ``self.default_keep_fraction`` for this call.
+        :type keep_fraction: float or None
+        :param threshold: Absolute coefficient magnitude threshold.
+            Overrides ``self.default_threshold`` for this call.
+        :type threshold: float or None
+        :param kwargs: Accepted and ignored for compatibility.
+        :returns: The sparse wavelet payload.
+        :rtype: WaveletPayload
+        :raises ValueError: If both or neither of ``keep_fraction`` and
+            ``threshold`` are specified, or either is out of range.
+        """
+        array = np.asarray(array)
+        wavelet = wavelet if wavelet is not None else self.default_wavelet
+        mode = mode if mode is not None else self.default_mode
+
+        if keep_fraction is not None and threshold is not None:
+            raise ValueError(
+                "Specify only one of `keep_fraction` or `threshold` per call or use the codec default."
+            )
+        if keep_fraction is None and threshold is None:
+            keep_fraction = self.default_keep_fraction
+            threshold = self.default_threshold
+        if (keep_fraction is None) == (threshold is None):
+            raise ValueError(
+                "Exactly one of `keep_fraction` or `threshold` must be specified, "
+                "either here or as a codec default."
+            )
+        if keep_fraction is not None and not 0 < keep_fraction <= 1:
+            raise ValueError(
+                f"`keep_fraction` must be in (0, 1], got {keep_fraction}"
+            )
+        if threshold is not None and threshold < 0:
+            raise ValueError(f"`threshold` must be >= 0, got {threshold}")
+
+        coeffs = pywt.wavedecn(array, wavelet, mode=mode)
+        flat = pywt.ravel_coeffs(coeffs)[0]
+        magnitude = np.abs(flat)
+
+        if keep_fraction is not None:
+            k = max(1, min(flat.size, int(round(keep_fraction * flat.size))))
+            indices = np.argpartition(magnitude, -k)[-k:]
+            indices.sort()
+        else:
+            indices = np.flatnonzero(magnitude > threshold)
+
+        index_dtype = (
+            np.uint32 if flat.size <= np.iinfo(np.uint32).max else np.uint64
+        )
+        return WaveletPayload(
+            original_shape=array.shape,
+            dtype=array.dtype,
+            wavelet=wavelet,
+            mode=mode,
+            indices=indices.astype(index_dtype),
+            values=np.ascontiguousarray(flat[indices]),
+        )
+
+    def decode(self, payload: WaveletPayload) -> np.ndarray:
+        """
+        Reconstruct the dense array via the inverse wavelet transform.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: WaveletPayload
+        :returns: The reconstructed dense array.
+        :rtype: numpy.ndarray
+        """
+        shapes = payload.coeff_shapes
+        template = [np.zeros(shapes[0])] + [
+            {key: np.zeros(shape) for key, shape in details.items()}
+            for details in shapes[1:]
+        ]
+        _, slices, coeff_shapes = pywt.ravel_coeffs(template)
+        flat = np.zeros(payload.n_coeffs, dtype=payload.values.dtype)
+        flat[payload.indices] = payload.values
+
+        coeffs = pywt.unravel_coeffs(
+            flat, slices, coeff_shapes, output_format="wavedecn"
+        )
+        out = pywt.waverecn(coeffs, payload.wavelet, mode=payload.mode)
+        out = out[tuple(slice(0, n) for n in payload.original_shape)]
+        return out.astype(payload.dtype, copy=False)
+
+    def compressed_size_bytes(self, payload: WaveletPayload) -> int:
+        """
+        Get payload size.
+
+        :param payload: A payload previously produced by :meth:`encode`.
+        :type payload: WaveletPayload
+        :returns: Combined byte size of the stored arrays.
+        :rtype: int
+        """
+        return payload.indices.nbytes + payload.values.nbytes
+
+    def payload_to_group(self, payload: WaveletPayload, group) -> None:
+        """
+        Write ``indices``, ``values``, and transform parameters into an
+        open ``h5py.Group``.
+
+        :param payload: The payload to serialize.
+        :type payload: WaveletPayload
+        :param group: An open, writable ``h5py.Group`` dedicated to
+            this one patch.
+        :type group: h5py.Group
+        """
+        group.create_dataset("indices", data=payload.indices)
+        group.create_dataset("values", data=payload.values)
+        group.attrs["wavelet"] = payload.wavelet
+        group.attrs["mode"] = payload.mode
+
+    def payload_from_group(
+        self, group, *, original_shape: tuple[int, ...], dtype: np.dtype
+    ) -> WaveletPayload:
+        """
+        Construct a :class:`WaveletPayload` 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 read from the
+            container's own attributes.
+        :type original_shape: tuple[int, ...]
+        :param dtype: The array dtype, as read from the container's
+            own attributes.
+        :type dtype: numpy.dtype
+        :returns: The reconstructed payload.
+        :rtype: WaveletPayload
+        """
+        return WaveletPayload(
+            original_shape=original_shape,
+            dtype=dtype,
+            wavelet=str(group.attrs["wavelet"]),
+            mode=str(group.attrs["mode"]),
+            indices=group["indices"][:],
+            values=group["values"][:],
+        )

comdas/io/__init__.py

@@ -0,0 +1,18 @@
+"""
+I/O support for ComDAS patches.
+
+Importing this subpackage registers the ``COMDAS`` HDF5 container
+format (:class:`~comdas.io.h5_container.ComdasV1`) with DASCore's
+FiberIO manager, so ``dc.spool(path)`` / ``dc.read(...)`` / ``dc.write(...)``
+work with it directly. :func:`~comdas.io.h5_container.write_compressed`
+is a small convenience wrapper for the write side.
+"""
+
+from comdas.io.h5_container import (
+    ComdasV1,
+    FORMAT_NAME,
+    FORMAT_VERSION,
+    write_compressed,
+)
+
+__all__ = ["ComdasV1", "FORMAT_NAME", "FORMAT_VERSION", "write_compressed"]