Forward masks through CodecStack to mask-aware codecs

src/numcodecs_combinators/stack.py

@@ -24,6 +24,15 @@
 from .abc import CodecCombinatorMixin
 
 
+def _is_mask_aware(codec: Codec) -> bool:
+    """Check for the mask-aware codec protocol of
+    [`numcodecs-mask`](https://numcodecs-mask.readthedocs.io) structurally, so
+    that `numcodecs-combinators` does not need to depend on it."""
+    return callable(getattr(codec, "encode_masked", None)) and callable(
+        getattr(codec, "decode_masked", None)
+    )
+
+
 class CodecStack(Codec, CodecCombinatorMixin, tuple[Codec]):
     """
     A stack of codecs, which makes up a combined codec.
@@ -136,6 +145,94 @@ def decode(self, buf: Buffer, out: Optional[Buffer] = None) -> Buffer:
             )
         return numcodecs.compat.ndarray_copy(decoded, out)  # type: ignore
 
+    def encode_masked(
+        self, buf: Buffer, mask: np.ndarray[tuple[int, ...], np.dtype[np.bool]]
+    ) -> Buffer:
+        """Encode the data in `buf`, ignoring the values where `mask` is
+        [`True`][True].
+
+        The stack forwards the mask to every codec in the stack that
+        implements the mask-aware codec protocol of
+        [`numcodecs-mask`](https://numcodecs-mask.readthedocs.io), i.e. that
+        provides `encode_masked` and `decode_masked` methods; all other codecs
+        encode as usual. All mask-aware codecs in the stack must encode data
+        of the same shape as the mask, e.g. shape-preserving filters may
+        precede them but reshaping codecs must not.
+
+        Parameters
+        ----------
+        buf : Buffer
+            Data to be encoded. May be any object supporting the new-style
+            buffer protocol. The values at masked positions are unspecified.
+        mask : np.ndarray[tuple[int, ...], np.dtype[np.bool]]
+            The [boolean][numpy.bool] mask, of the same shape as the data, of
+            the values that do not need to be preserved.
+
+        Returns
+        -------
+        enc : Buffer
+            Encoded data. May be any object supporting the new-style buffer
+            protocol.
+        """
+
+        if len(self) == 0:
+            return buf
+
+        encoded = buf
+        for codec in self:
+            encoded_ndarray = np.asarray(
+                numcodecs.compat.ensure_contiguous_ndarray_like(encoded, flatten=False)
+            )
+            if _is_mask_aware(codec):
+                encoded = codec.encode_masked(encoded_ndarray, mask)  # type: ignore
+            else:
+                encoded = codec.encode(encoded_ndarray)
+        return encoded
+
+    def decode_masked(
+        self,
+        buf: Buffer,
+        mask: np.ndarray[tuple[int, ...], np.dtype[np.bool]],
+        out: Optional[Buffer] = None,
+    ) -> Buffer:
+        """Decode the data in `buf`, which was encoded with the same `mask`.
+
+        The mask is forwarded to every codec in the stack that implements the
+        mask-aware codec protocol of
+        [`numcodecs-mask`](https://numcodecs-mask.readthedocs.io), matching
+        [`encode_masked`][numcodecs_combinators.stack.CodecStack.encode_masked].
+
+        Parameters
+        ----------
+        buf : Buffer
+            Encoded data. May be any object supporting the new-style buffer
+            protocol.
+        mask : np.ndarray[tuple[int, ...], np.dtype[np.bool]]
+            The [boolean][numpy.bool] mask, of the same shape as the decoded
+            data, that was passed to
+            [`encode_masked`][numcodecs_combinators.stack.CodecStack.encode_masked].
+        out : Buffer, optional
+            Writeable buffer to store decoded data. N.B. if provided, this
+            buffer must be exactly the right size to store the decoded data.
+
+        Returns
+        -------
+        dec : Buffer
+            Decoded data. May be any object supporting the new-style buffer
+            protocol. The values at masked positions are unspecified.
+        """
+
+        decoded = buf
+        for codec in reversed(self):
+            decoded_ndarray = numcodecs.compat.ensure_contiguous_ndarray_like(
+                decoded, flatten=False
+            )
+            if _is_mask_aware(codec):
+                decoded = codec.decode_masked(decoded_ndarray, mask, out=None)  # type: ignore
+            else:
+                decoded = codec.decode(decoded_ndarray, out=None)
+        return numcodecs.compat.ndarray_copy(decoded, out)  # type: ignore
+
     def encode_decode(self, buf: Buffer) -> Buffer:
         """
         Encode, then decode the data in `buf`.

tests/test_stack.py

@@ -114,3 +114,81 @@ def test_map():
             )
         )
     )
+
+
+class SparseCodec(Codec):
+    """Test codec implementing the mask-aware protocol of numcodecs-mask:
+    it only stores the unmasked values and records the masks it was given."""
+
+    codec_id = "sparse-test"
+
+    def __init__(self):
+        self.masks: list[np.ndarray] = []
+
+    def encode(self, buf):
+        return self.encode_masked(buf, np.zeros(np.shape(buf), dtype=np.bool))
+
+    def decode(self, buf, out=None):
+        return self.decode_masked(buf, np.zeros(np.shape(out), dtype=np.bool), out)
+
+    def encode_masked(self, buf, mask):
+        a = numcodecs.compat.ensure_ndarray(buf)
+        assert mask.shape == a.shape and mask.dtype == np.bool
+        self.masks.append(np.copy(mask))
+        return a[~mask].tobytes()
+
+    def decode_masked(self, buf, mask, out=None):
+        assert mask.dtype == np.bool
+        decoded = np.zeros(mask.shape, dtype=np.float64)
+        decoded[~mask] = np.frombuffer(
+            numcodecs.compat.ensure_bytes(buf), dtype=np.float64
+        )
+        return numcodecs.compat.ndarray_copy(decoded, out)
+
+    def get_config(self) -> dict:
+        return dict(id=type(self).codec_id)
+
+
+class ClampFilterCodec(Codec):
+    """Test filter codec that clamps on encode and passes through on decode,
+    preserving the shape of the data."""
+
+    codec_id = "clamp-test"
+
+    def encode(self, buf):
+        return np.minimum(numcodecs.compat.ensure_ndarray(buf), 10.0)
+
+    def decode(self, buf, out=None):
+        return numcodecs.compat.ndarray_copy(numcodecs.compat.ensure_ndarray(buf), out)
+
+    def get_config(self) -> dict:
+        return dict(id=type(self).codec_id)
+
+
+def test_encode_decode_masked():
+    data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    mask = np.array([[False, True, False], [True, False, False]])
+
+    # the mask is forwarded through mask-unaware, shape-preserving filters to
+    # mask-aware codecs
+    inner = SparseCodec()
+    stack = CodecStack(ClampFilterCodec(), inner)
+
+    encoded = stack.encode_masked(data, mask)
+    assert len(inner.masks) == 1
+    np.testing.assert_array_equal(inner.masks[0], mask)
+    # the mask-aware codec only stored the four unmasked values
+    assert len(numcodecs.compat.ensure_bytes(encoded)) == 4 * 8
+
+    decoded = np.asarray(stack.decode_masked(encoded, mask)).reshape(data.shape)
+    np.testing.assert_array_equal(decoded[~mask], data[~mask])
+
+    # an empty stack and a stack of only mask-unaware codecs still roundtrip
+    np.testing.assert_array_equal(CodecStack().encode_masked(data, mask), data)
+    stack = CodecStack(dict(id="zlib", level=1))
+    decoded = np.frombuffer(
+        numcodecs.compat.ensure_bytes(
+            stack.decode_masked(stack.encode_masked(data, mask), mask)
+        )
+    ).reshape(data.shape)
+    np.testing.assert_array_equal(decoded, data)