From 23b3ff31f9dfb2aaa8daf9b5e7de28edb82e6d3b Mon Sep 17 00:00:00 2001 From: Sara Faghih-Naini Date: Mon, 5 Oct 2026 16:47:07 +0200 Subject: [PATCH] Forward masks through CodecStack to codecs that implement the mask-aware protocol of numcodecs-mask --- src/numcodecs_combinators/stack.py | 97 ++++++++++++++++++++++++++++++ tests/test_stack.py | 78 ++++++++++++++++++++++++ 2 files changed, 175 insertions(+) diff --git a/src/numcodecs_combinators/stack.py b/src/numcodecs_combinators/stack.py index 7fe8566..4ea690e 100644 --- a/src/numcodecs_combinators/stack.py +++ b/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`. diff --git a/tests/test_stack.py b/tests/test_stack.py index a1b9fc5..c646990 100644 --- a/tests/test_stack.py +++ b/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)