Add optional FixedScaleOffset dimension preservation

docs/release.md

@@ -14,6 +14,12 @@
 
 ## Unreleased
 
+### Enhancements
+
+* Add a `keep_dimensions` option to `FixedScaleOffset` for preserving multidimensional
+  array shapes during encoding and decoding.
+  By {user}`shixi-li <shixi-li>`, {issue}`852`
+
 ### Maintenance
 
 * **Migrate build system from setuptools/setup.py to meson-python.** This replaces the

src/numcodecs/fixedscaleoffset.py

@@ -20,6 +20,9 @@ class FixedScaleOffset(Codec):
         Data type to use for decoded data.
     astype : dtype, optional
         Data type to use for encoded data.
+    keep_dimensions : bool, optional
+        Preserve the input array dimensions when encoding and decoding.
+        Defaults to False for backwards compatibility.
 
     Notes
     -----
@@ -69,10 +72,11 @@ class FixedScaleOffset(Codec):
 
     codec_id = 'fixedscaleoffset'
 
-    def __init__(self, offset, scale, dtype, astype=None):
+    def __init__(self, offset, scale, dtype, astype=None, keep_dimensions=False):
         self.offset = offset
         self.scale = scale
         self.dtype = np.dtype(dtype)
+        self.keep_dimensions = keep_dimensions
         if astype is None:
             self.astype = self.dtype
         else:
@@ -84,8 +88,9 @@ def encode(self, buf):
         # normalise input
         arr = ensure_ndarray(buf).view(self.dtype)
 
-        # flatten to simplify implementation
-        arr = arr.reshape(-1, order='A')
+        if not self.keep_dimensions:
+            # flatten to simplify implementation
+            arr = arr.reshape(-1, order='A')
 
         # compute scale offset
         enc = (arr - self.offset) * self.scale
@@ -100,15 +105,27 @@ def decode(self, buf, out=None):
         # interpret buffer as numpy array
         enc = ensure_ndarray(buf).view(self.astype)
 
-        # flatten to simplify implementation
-        enc = enc.reshape(-1, order='A')
+        if not self.keep_dimensions:
+            # flatten to simplify implementation
+            enc = enc.reshape(-1, order='A')
 
         # decode scale offset
         dec = (enc / self.scale) + self.offset
 
         # convert dtype
         dec = dec.astype(self.dtype, copy=False)
 
+        if (
+            self.keep_dimensions
+            and isinstance(out, np.ndarray)
+            and dec.ndim > 1
+            and out.shape == dec.shape
+            and out.dtype == dec.dtype
+        ):
+            # Preserve logical coordinates when source and destination memory orders differ.
+            np.copyto(out, dec)
+            return out
+
         # handle output
         return ndarray_copy(dec, out)
 
@@ -120,11 +137,14 @@ def get_config(self):
             'offset': self.offset,
             'dtype': self.dtype.str,
             'astype': self.astype.str,
+            'keep_dimensions': self.keep_dimensions,
         }
 
     def __repr__(self):
         r = f'{type(self).__name__}(scale={self.scale}, offset={self.offset}, dtype={self.dtype.str!r}'
         if self.astype != self.dtype:
             r += f', astype={self.astype.str!r}'
+        if self.keep_dimensions:
+            r += ', keep_dimensions=True'
         r += ')'
         return r

tests/test_fixedscaleoffset.py

@@ -1,10 +1,12 @@
 import itertools
+from typing import Literal
 
 import numpy as np
 import pytest
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 from numcodecs.fixedscaleoffset import FixedScaleOffset
+from numcodecs.registry import get_codec
 from tests.common import (
     check_backwards_compatibility,
     check_config,
@@ -55,15 +57,89 @@ def test_encode(offset: float, scale: float, expected: list[int]):
     assert np.dtype(astype) == actual.dtype
 
 
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_keep_dimensions(order: Literal["C", "F"]):
+    shape = (8, 4, 5)
+    arr = 200.0 + np.arange(np.prod(shape), dtype="<f4").reshape(shape, order=order) * 0.001
+    codec = FixedScaleOffset(
+        offset=200.0,
+        scale=1000.0,
+        dtype="<f4",
+        astype="<u2",
+        keep_dimensions=True,
+    )
+
+    encoded = codec.encode(arr)
+    decoded = codec.decode(encoded)
+
+    assert encoded.shape == shape
+    assert decoded.shape == shape
+    assert_array_almost_equal(arr, decoded, decimal=3)
+
+
+@pytest.mark.parametrize(("input_order", "out_order"), [("C", "F"), ("F", "C")])
+def test_keep_dimensions_decode_out_preserves_logical_coordinates(
+    input_order: Literal["C", "F"], out_order: Literal["C", "F"]
+):
+    shape = (2, 3, 4)
+    arr = np.arange(np.prod(shape), dtype="<f4").reshape(shape, order=input_order)
+    codec = FixedScaleOffset(
+        offset=0,
+        scale=1,
+        dtype="<f4",
+        astype="<i4",
+        keep_dimensions=True,
+    )
+    encoded = codec.encode(arr)
+    out = np.empty(shape, dtype="<f4", order=out_order)
+
+    decoded = codec.decode(encoded, out=out)
+
+    assert decoded is out
+    assert decoded.shape == shape
+    assert decoded.dtype == arr.dtype
+    assert_array_equal(decoded, arr)
+
+
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_default_flattens_multidimensional_arrays(order: Literal["C", "F"]):
+    shape = (2, 3, 4)
+    arr = np.arange(np.prod(shape), dtype="<f8").reshape(shape, order=order)
+    codec = FixedScaleOffset(offset=0, scale=1, dtype="<f8")
+
+    encoded = codec.encode(arr)
+    decoded = codec.decode(encoded)
+
+    assert encoded.shape == (arr.size,)
+    assert decoded.shape == (arr.size,)
+
+
 def test_config():
-    codec = FixedScaleOffset(dtype='<f8', astype='<i4', scale=10, offset=100)
+    codec = FixedScaleOffset(dtype='<f8', astype='<i4', scale=10, offset=100, keep_dimensions=True)
     check_config(codec)
 
 
+def test_legacy_config():
+    config = {
+        'id': FixedScaleOffset.codec_id,
+        'dtype': '<f8',
+        'astype': '<i4',
+        'scale': 10,
+        'offset': 100,
+    }
+
+    codec = get_codec(config)
+
+    assert codec.keep_dimensions is False
+
+
 def test_repr():
     stmt = "FixedScaleOffset(scale=10, offset=100, dtype='<f8', astype='<i4')"
     check_repr(stmt)
 
+    stmt = "FixedScaleOffset(scale=10, offset=100, dtype='<f8', astype='<i4', keep_dimensions=True)"
+    check_repr(stmt)
+
 
 def test_backwards_compatibility():
     precision = [int(np.log10(codec.scale)) for codec in codecs]