diff --git a/mkdocs.yml b/mkdocs.yml index 01e101c..5fb07b9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -55,6 +55,7 @@ plugins: - https://docs.python.org/3.12/objects.inv - https://numpy.org/doc/2.2/objects.inv - https://numcodecs.readthedocs.io/en/v0.16.0/objects.inv + - https://numcodecs-combinators.readthedocs.io/en/v0.2.15/objects.inv watch: - src diff --git a/pyproject.toml b/pyproject.toml index 369c7f1..f18e492 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,9 @@ readme = "README.md" license = "MPL-2.0" requires-python = ">=3.10" dependencies = [ + "leb128~=1.0.8", "numcodecs>=0.13.0,<0.17", + "numcodecs-combinators~=0.2.15", "numpy~=2.0", "typing-extensions~=4.6", ] @@ -20,6 +22,7 @@ dev = ["mypy~=1.14", "pytest~=8.3"] [project.entry-points."numcodecs.codecs"] "replace.filter" = "numcodecs_replace:ReplaceFilterCodec" +"mask.meta" = "numcodecs_replace:MaskMetaCodec" [tool.setuptools.packages.find] where = ["src"] diff --git a/src/numcodecs_replace/__init__.py b/src/numcodecs_replace/__init__.py index 8f6510e..6b95099 100644 --- a/src/numcodecs_replace/__init__.py +++ b/src/numcodecs_replace/__init__.py @@ -2,15 +2,20 @@ Value replacement codecs for the [`numcodecs`][numcodecs] buffer compression API. """ -__all__ = ["ReplaceFilterCodec", "Replacement"] +__all__ = ["ReplaceFilterCodec", "MaskMetaCodec", "Replacement"] +from collections.abc import Callable from enum import Enum, auto +from functools import reduce +from io import BytesIO from typing import Literal, TypeVar +import leb128 import numcodecs.compat import numcodecs.registry import numpy as np from numcodecs.abc import Codec +from numcodecs_combinators.abc import CodecCombinatorMixin from typing_extensions import ( Buffer, # MSPV 3.12 assert_never, # MSPV 3.11 @@ -103,7 +108,7 @@ class ReplaceFilterCodec(Codec): Mapping from values to be replaced to the replacement values. """ - __slots__ = "_replacements" + __slots__: tuple[str, ...] = ("_replacements",) _replacements: dict[ int | float, int | float | Replacement, @@ -217,3 +222,282 @@ def __repr__(self) -> str: numcodecs.registry.register_codec(ReplaceFilterCodec) + + +class MaskMetaCodec(Codec, CodecCombinatorMixin): + """ + Meta-codec that masks a value during encoding and restores it during decoding. + + The data is encoded with the `codec`, the [boolean][numpy.bool] mask of + where values need to be restored is encoded with the `bitmap_codec`. + + Multiple [`MaskMetaCodec`][.]s can be nested to mask several values. + The masked value can be replaced with a fill value before further encoding + by including a [`ReplaceFilterCodec`][..ReplaceFilterCodec] in the `codec`, + e.g. by stacking using the + [`numcodecs-combinators`](https://numcodecs-combinators.readthedocs.io) + package. + + Parameters + ---------- + mask : int | float + The value to be masked. + codec : dict | Codec + The configuration or instantiated codec that encodes the data. + bitmap_codec : dict | Codec + The configuration or instantiated codec that encodes the + [boolean][numpy.bool] mask of where values need to be restored. + + For instance, the [`numcodecs.PackBits`][numcodecs.packbits.PackBits] + codec can be used to pack the mask into a byte array. + """ + + __slots__: tuple[str, ...] = ("_mask", "_codec", "_bitmap_codec") + _mask: int | float + _codec: Codec + _bitmap_codec: Codec + + codec_id: str = "mask.meta" # type: ignore + + def __init__( + self, + *, + mask: int | float, + codec: dict | Codec, + bitmap_codec: dict | Codec, + ) -> None: + self._mask = mask + self._codec = ( + codec if isinstance(codec, Codec) else numcodecs.registry.get_codec(codec) + ) + self._bitmap_codec = ( + bitmap_codec + if isinstance(bitmap_codec, Codec) + else numcodecs.registry.get_codec(bitmap_codec) + ) + + def encode(self, buf: Buffer) -> Buffer: + """Encode the data in `buf`. + + Parameters + ---------- + buf : Buffer + Data to be encoded. May be any object supporting the new-style + buffer protocol. + + Returns + ------- + enc : Buffer + Encoded data. May be any object supporting the new-style buffer + protocol. + """ + + a = np.copy(numcodecs.compat.ensure_ndarray(buf)) + dtype, shape = a.dtype, a.shape + + if isinstance(self._mask, int) or not np.isnan(self._mask): + is_masked = a == self._mask + else: + is_masked = np.isnan(a) + + # message: dtype shape encoded-dtype encoded-shape [padding] encoded + # bitmap-dtype bitmap-shape [padding] bitmap + message: list[bytes | bytearray] = [] + + message.append(leb128.u.encode(len(dtype.str))) + message.append(dtype.str.encode("ascii")) + + message.append(leb128.u.encode(len(shape))) + for s in shape: + message.append(leb128.u.encode(s)) + + encoded = self._codec.encode(a) + encoded = numcodecs.compat.ensure_ndarray(encoded) + + message.append(leb128.u.encode(len(encoded.dtype.str))) + message.append(encoded.dtype.str.encode("ascii")) + + message.append(leb128.u.encode(encoded.ndim)) + for s in encoded.shape: + message.append(leb128.u.encode(s)) + + # insert padding to align with encoded itemsize + message.append( + b"\0" + * ( + encoded.dtype.itemsize + - (sum(len(m) for m in message) % encoded.itemsize) + ) + ) + + # ensure that the encoded values are encoded in little endian binary + message.append(encoded.astype(encoded.dtype.newbyteorder("<")).tobytes()) + + bitmap = self._bitmap_codec.encode(is_masked) + bitmap = numcodecs.compat.ensure_ndarray(bitmap) + + message.append(leb128.u.encode(len(bitmap.dtype.str))) + message.append(bitmap.dtype.str.encode("ascii")) + + message.append(leb128.u.encode(bitmap.ndim)) + for s in bitmap.shape: + message.append(leb128.u.encode(s)) + + # insert padding to align with bitmap itemsize + message.append( + b"\0" + * (bitmap.dtype.itemsize - (sum(len(m) for m in message) % bitmap.itemsize)) + ) + + # ensure that the bitmap values are encoded in little endian binary + message.append(bitmap.astype(bitmap.dtype.newbyteorder("<")).tobytes()) + + return b"".join(message) + + def decode(self, buf: Buffer, out: None | Buffer = None) -> Buffer: + """ + Decode the data in `buf`. + + Parameters + ---------- + buf : Buffer + Encoded data. May be any object supporting the new-style buffer + protocol. + 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. + """ + + b = numcodecs.compat.ensure_bytes(buf) + b_io = BytesIO(b) + + # message: dtype shape encoded-dtype encoded-shape [padding] encoded + # bitmap-dtype bitmap-shape [padding] bitmap + dtype = np.dtype(b_io.read(leb128.u.decode_reader(b_io)[0]).decode("ascii")) + shape = tuple( + leb128.u.decode_reader(b_io)[0] + for _ in range(leb128.u.decode_reader(b_io)[0]) + ) + + encoded_dtype = np.dtype( + b_io.read(leb128.u.decode_reader(b_io)[0]).decode("ascii") + ) + encoded_shape = tuple( + leb128.u.decode_reader(b_io)[0] + for _ in range(leb128.u.decode_reader(b_io)[0]) + ) + encoded_size = reduce(lambda a, b: a * b, encoded_shape, 1) + + # remove padding to align with encoded itemsize + b_io.read(encoded_dtype.itemsize - (b_io.tell() % encoded_dtype.itemsize)) + + encoded = ( + np.frombuffer( + b_io.read(encoded_size * encoded_dtype.itemsize), + dtype=encoded_dtype.newbyteorder("<"), + count=encoded_size, + ) + .astype(encoded_dtype) + .reshape(encoded_shape) + ) + + decoded = np.empty(shape, dtype=dtype) + self._codec.decode(encoded, out=decoded) + + bitmap_dtype = np.dtype( + b_io.read(leb128.u.decode_reader(b_io)[0]).decode("ascii") + ) + bitmap_shape = tuple( + leb128.u.decode_reader(b_io)[0] + for _ in range(leb128.u.decode_reader(b_io)[0]) + ) + bitmap_size = reduce(lambda a, b: a * b, bitmap_shape, 1) + + # remove padding to align with bitmap itemsize + b_io.read(bitmap_dtype.itemsize - (b_io.tell() % bitmap_dtype.itemsize)) + + bitmap = ( + np.frombuffer( + b_io.read(bitmap_size * bitmap_dtype.itemsize), + dtype=bitmap_dtype.newbyteorder("<"), + count=bitmap_size, + ) + .astype(bitmap_dtype) + .reshape(bitmap_shape) + ) + + is_masked = np.empty(shape, dtype=np.bool) + self._bitmap_codec.decode(bitmap, out=is_masked) + + decoded[is_masked] = self._mask + + return numcodecs.compat.ndarray_copy(decoded, out) # type: ignore + + def get_config(self) -> dict: + """ + Returns the configuration of this mask meta-codec. + + [`numcodecs.registry.get_codec(config)`][numcodecs.registry.get_codec] + can be used to reconstruct this codec from the returned config. + + Returns + ------- + config : dict + Configuration of this mask meta-codec. + """ + + return dict( + id=type(self).codec_id, + mask=self._mask, + codec=self._codec.get_config(), + bitmap_codec=self._bitmap_codec.get_config(), + ) + + def __repr__(self) -> str: + return f"{type(self).__name__}(mask={self._mask!r}, codec={self._codec!r}, bitmap_codec={self._bitmap_codec!r})" + + def map(self, mapper: Callable[[Codec], Codec]) -> "MaskMetaCodec": + """ + Apply the `mapper` to this mask meta-codec. + + In the returned [`MaskMetaCodec`][..], the `codec` and + `bitmap_codec` are replaced by their mapped codecs. + + The `mapper` should recursively apply itself to any inner codecs that + also implement the + [`CodecCombinatorMixin`][numcodecs_combinators.abc.CodecCombinatorMixin] + mixin. + + To automatically handle the recursive application as a caller, you can + use + ```python + numcodecs_combinators.map_codec(codec, mapper) + ``` + instead. + + Parameters + ---------- + mapper : Callable[[Codec], Codec] + The callable that should be applied to the wrapped `codec` and + `bitmap_codec` to map over this mask meta-codec. + + Returns + ------- + mapped : MaskMetaCodec + The mapped mask meta-codec. + """ + + return MaskMetaCodec( + mask=self._mask, + codec=mapper(self._codec), + bitmap_codec=mapper(self._bitmap_codec), + ) + + +numcodecs.registry.register_codec(MaskMetaCodec) diff --git a/tests/test_replace.py b/tests/test_replace.py index 1e838ea..0b5b88b 100644 --- a/tests/test_replace.py +++ b/tests/test_replace.py @@ -6,15 +6,43 @@ import pytest -def test_from_config(): - codec = numcodecs.registry.get_codec(dict(id="replace.filter", replacements={})) +def test_filter_from_config(): + config = dict(id="replace.filter", replacements={}) + + codec = numcodecs.registry.get_codec(config) assert codec.__class__.__name__ == "ReplaceFilterCodec" assert codec.__class__.__module__ == "numcodecs_replace" assert repr(codec) == "ReplaceFilterCodec(replacements={})" + assert json.dumps(codec.get_config(), sort_keys=True) == json.dumps( + config, sort_keys=True + ) + + +def test_meta_from_config(): + config = dict( + id="mask.meta", + mask=np.nan, + codec=dict(id="zlib", level=1), + bitmap_codec=dict(id="packbits"), + ) + + codec = numcodecs.registry.get_codec(config) + assert codec.__class__.__name__ == "MaskMetaCodec" + assert codec.__class__.__module__ == "numcodecs_replace" + + assert ( + repr(codec) + == "MaskMetaCodec(mask=nan, codec=Zlib(level=1), bitmap_codec=PackBits())" + ) + + assert json.dumps(codec.get_config(), sort_keys=True) == json.dumps( + config, sort_keys=True + ) + -def check_roundtrip(data: np.ndarray): +def check_filter_roundtrip(data: np.ndarray): config = dict( id="replace.filter", replacements={ @@ -89,20 +117,20 @@ def check_roundtrip(data: np.ndarray): @np.errstate(invalid="ignore") -def test_roundtrip(): - check_roundtrip(np.zeros(tuple())) +def test_filter_roundtrip(): + check_filter_roundtrip(np.zeros(tuple())) with pytest.warns(RuntimeWarning, match="empty slice"): - check_roundtrip(np.zeros((0,))) - check_roundtrip(np.arange(1000).reshape(10, 10, 10)) - check_roundtrip(np.array([4.2, -2.4, np.nan, -np.nan, 0.0, -0.0])) - check_roundtrip(np.array([np.inf, -np.inf, np.nan, -np.nan, 0.0, -0.0])) - check_roundtrip( + check_filter_roundtrip(np.zeros((0,))) + check_filter_roundtrip(np.arange(1000).reshape(10, 10, 10)) + check_filter_roundtrip(np.array([4.2, -2.4, np.nan, -np.nan, 0.0, -0.0])) + check_filter_roundtrip(np.array([np.inf, -np.inf, np.nan, -np.nan, 0.0, -0.0])) + check_filter_roundtrip( np.array( [np.inf, -np.inf, np.nan, -np.nan, 0.0, -0.0], dtype=np.dtype(np.float64).newbyteorder("<"), ) ) - check_roundtrip( + check_filter_roundtrip( np.array( [np.inf, -np.inf, np.nan, -np.nan, 0.0, -0.0], dtype=np.dtype(np.float64).newbyteorder(">"),