fix negative value handling

datajoint · Mar 5, 2024 · da169c8 · da169c8
1 parent 5e72bf3
commit da169c8
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 29 deletions.
diff --git a/examples/workbook.ipynb b/examples/workbook.ipynb
diff --git a/src/poisson_numcodecs/codec.py b/src/poisson_numcodecs/codec.py
@@ -18,13 +18,15 @@ def make_anscombe_lookup(
     :param input_max: the maximum value in the input
     :param beta: the grayscale quantization step expressed in units of noise std dev
     """
-    xx = (np.r_[:input_max + 1] - zero_level) / sensitivity
+    xx = (np.r_[:input_max + 1] - zero_level) / sensitivity  # input expressed in photon rates
     zero_slope = 1 / beta / np.sqrt(3/8)  # slope for negative values
-    offset = -zero_level * zero_slope
-    lookup_table = np.floor(0.75 + offset +
+    offset = zero_level * zero_slope / sensitivity
+    lookup_table = np.round(offset +
           (xx < 0) * (xx * zero_slope) +
           (xx >= 0) * (2.0 / beta * (np.sqrt(np.maximum(0, xx) + 3/8) - np.sqrt(3/8))))
-    return lookup_table.astype(output_type)
+    lookup = lookup_table.astype(output_type)
+    assert np.diff(lookup_table).min() >= 0, "non-monotonic lookup generated"
+    return lookup
 
 
 def make_inverse_lookup(lookup_table: np.ndarray, output_type='int16') -> np.ndarray:
@@ -69,7 +71,8 @@ def __init__(self,
     def encode(self, buf: np.ndarray) -> np.ndarray:
         lookup_table = make_anscombe_lookup(
             self.photon_sensitivity, 
-            output_type=self.encoded_dtype
+            output_type=self.encoded_dtype,
+            zero_level=self.zero_level,
         )
         encoded = lookup(buf, lookup_table)
         shape = [encoded.ndim] + list(encoded.shape)
@@ -79,7 +82,8 @@ def encode(self, buf: np.ndarray) -> np.ndarray:
     def decode(self, buf: bytes, out=None) -> np.ndarray:
         lookup_table = make_anscombe_lookup(
             self.photon_sensitivity, 
-            output_type=self.encoded_dtype
+            output_type=self.encoded_dtype,
+            zero_level=self.zero_level,
         )
         inverse_table = make_inverse_lookup(
             lookup_table,

diff --git a/tests/test_poisson_codec.py b/tests/test_poisson_codec.py
@@ -1,3 +1,4 @@
+from numpy._typing._generic_alias import NDArray
 from poisson_numcodecs import PoissonCodec
 import numpy as np
 import pytest
@@ -11,28 +12,30 @@ def make_poisson_ramp_signals(shape=(10, 1, 1), min_rate=1, max_rate=5, dtype="i
     output_array = np.zeros(shape, dtype=dtype)
     for x_ind in range(x):
         for y_ind in range(y):
-            output_array[x_ind, y_ind, :] = np.random.poisson(
+            output_array[x_ind, y_ind, :] = sensitivity * np.random.poisson(
                 np.arange(min_rate, max_rate, (max_rate-min_rate)/times))
-    return output_array
+    return output_array.astype('int16')
+
+sensitivity = 100.0
 
 @pytest.fixture
 def test_data(dtype="int16"):
-    test2d = make_poisson_ramp_signals(shape=(50, 1, 1), min_rate=1, max_rate=5, dtype=dtype)
+    test2d =  make_poisson_ramp_signals(shape=(50, 1, 1), min_rate=1, max_rate=5, dtype=dtype)
     test2d_long = make_poisson_ramp_signals(shape=(1, 50, 1), min_rate=1, max_rate=5, dtype=dtype)
     return [test2d, test2d_long]
 
 def test_poisson_encode_decode(test_data):
     poisson_codec = PoissonCodec(
         zero_level=0,
-        photon_sensitivity=1.0,
+        photon_sensitivity=sensitivity,
         encoded_dtype='uint8', 
         decoded_dtype='int16'
     )
 
     for example_data in test_data:
         encoded = poisson_codec.encode(example_data)
         decoded = poisson_codec.decode(encoded)
-        np.testing.assert_allclose(decoded, example_data, atol=1e-3)
+        assert np.abs(decoded - example_data).max() < sensitivity / 2
 
 if __name__ == '__main__':
     list_data = test_data("int16")