Add Denormalize unit

.gitignore

@@ -146,3 +146,4 @@ cython_debug/
 
 src/ezmsg/sigproc/__version__.py
 uv.lock
+*.local.json

docs/source/conf.py

@@ -72,6 +72,10 @@
     "numpy": ("https://numpy.org/doc/stable/", None),
     "scipy": ("https://scipy.org/doc/scipy/", None),
     "ezmsg": ("https://www.ezmsg.org/ezmsg/", None),
+    "ezmsg.learn": ("https://www.ezmsg.org/ezmsg-learn/", None),
+    "ezmsg.event": ("https://www.ezmsg.org/ezmsg-event/", None),
+    "ezmsg.lsl": ("https://www.ezmsg.org/ezmsg-lsl/", None),
+    "ezmsg.blackrock": ("https://www.ezmsg.org/ezmsg-blackrock/", None),
 }
 intersphinx_disabled_domains = ["std"]
 

docs/source/guides/sigproc/processors.rst

@@ -52,6 +52,13 @@ ezmsg.sigproc.decimate
     :members:
 
 
+ezmsg.sigproc.denormalize
+------------------------
+
+.. automodule:: ezmsg.sigproc.denormalize
+    :members:
+
+
 ezmsg.sigproc.downsample
 --------------------------
 

src/ezmsg/sigproc/denormalize.py

@@ -0,0 +1,86 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.sigproc.base import (
+    BaseTransformerUnit,
+    BaseStatefulTransformer,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class DenormalizeSettings(ez.Settings):
+    low_rate: float = 2.0
+    """Low end of probable rate after denormalization (Hz)."""
+
+    high_rate: float = 40.0
+    """High end of probable rate after denormalization (Hz)."""
+
+    distribution: str = "uniform"
+    """Distribution to sample rates from. Options are 'uniform', 'normal', or 'constant'."""
+
+
+@processor_state
+class DenormalizeRateState:
+    gains: npt.NDArray | None = None
+    offsets: npt.NDArray | None = None
+
+
+class DenormalizeTransformer(
+    BaseStatefulTransformer[
+        DenormalizeSettings, AxisArray, AxisArray, DenormalizeRateState
+    ]
+):
+    """
+    Scales data from a normalized distribution (mean=0, std=1) to a denormalized
+    distribution using random per-channel offsets and gains designed to keep the
+    99.9% CIs between 0 and 2x the offset.
+
+    This is useful for simulating realistic firing rates from normalized data.
+    """
+
+    def _reset_state(self, message: AxisArray) -> None:
+        ax_ix = message.get_axis_idx("ch")
+        nch = message.data.shape[ax_ix]
+        arr_size = (nch, 1) if ax_ix == 0 else (1, nch)
+        if self.settings.distribution == "uniform":
+            self.state.offsets = np.random.uniform(2.0, 40.0, size=arr_size)
+        elif self.settings.distribution == "normal":
+            self.state.offsets = np.random.normal(
+                loc=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+                scale=(self.settings.high_rate - self.settings.low_rate) / 6.0,
+                size=arr_size,
+            )
+            self.state.offsets = np.clip(
+                self.state.offsets,
+                a_min=self.settings.low_rate,
+                a_max=self.settings.high_rate,
+            )
+        elif self.settings.distribution == "constant":
+            self.state.offsets = np.full(
+                shape=arr_size,
+                fill_value=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+            )
+        else:
+            raise ValueError(f"Invalid distribution: {self.settings.distribution}")
+        # Input has std == 1
+        # Desired output has range from 0 to 2*self.state.offsets within 99.9% confidence interval
+        # For a standard normal distribution, 99.9% of data is within +/- 3.29 std devs.
+        # So, gain = offset / 3.29 to scale the std dev appropriately.
+        self.state.gains = self.state.offsets / 3.29
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        denorm = message.data * self.state.gains + self.state.offsets
+        return replace(
+            message,
+            data=np.clip(denorm, a_min=0.0, a_max=None),
+        )
+
+
+class DenormalizeRateUnit(
+    BaseTransformerUnit[
+        DenormalizeSettings, AxisArray, AxisArray, DenormalizeTransformer
+    ]
+):
+    SETTINGS = DenormalizeSettings

tests/unit/test_denormalize.py

@@ -0,0 +1,245 @@
+import copy
+
+import numpy as np
+import pytest
+from frozendict import frozendict
+from ezmsg.util.messages.axisarray import AxisArray
+
+from ezmsg.sigproc.denormalize import (
+    DenormalizeSettings,
+    DenormalizeTransformer,
+)
+
+from tests.helpers.util import assert_messages_equal
+
+
+@pytest.fixture
+def basic_input_time_ch():
+    """Create a basic input with time x ch dimensions (standard normalized data)."""
+    n_times = 100
+    n_chans = 4
+    # Normalized data with mean ~0 and std ~1
+    np.random.seed(42)
+    data = np.random.randn(n_times, n_chans)
+    return AxisArray(
+        data=data,
+        dims=["time", "ch"],
+        axes=frozendict({"time": AxisArray.TimeAxis(fs=100.0)}),
+    )
+
+
+@pytest.fixture
+def basic_input_ch_time():
+    """Create a basic input with ch x time dimensions."""
+    n_times = 100
+    n_chans = 4
+    np.random.seed(42)
+    data = np.random.randn(n_chans, n_times)
+    return AxisArray(
+        data=data,
+        dims=["ch", "time"],
+        axes=frozendict({"time": AxisArray.TimeAxis(fs=100.0)}),
+    )
+
+
+class TestDenormalizeTransformer:
+    def test_uniform_distribution(self, basic_input_time_ch):
+        """Test denormalization with uniform distribution."""
+        backup = [copy.deepcopy(basic_input_time_ch)]
+
+        xformer = DenormalizeTransformer(
+            low_rate=2.0, high_rate=40.0, distribution="uniform"
+        )
+        output = xformer(basic_input_time_ch)
+
+        # Output shape should match input shape
+        assert output.data.shape == basic_input_time_ch.data.shape
+        assert output.dims == basic_input_time_ch.dims
+
+        # All output values should be non-negative (clipped)
+        assert np.all(output.data >= 0)
+
+        # Offsets should be within the specified range
+        assert xformer.state.offsets is not None
+        assert np.all(xformer.state.offsets >= 2.0)
+        assert np.all(xformer.state.offsets <= 40.0)
+
+        # Gains should be offsets / 3.29
+        assert np.allclose(xformer.state.gains, xformer.state.offsets / 3.29)
+
+        # Verify input wasn't modified
+        assert_messages_equal([basic_input_time_ch], backup)
+
+    def test_normal_distribution(self, basic_input_time_ch):
+        """Test denormalization with normal distribution."""
+        backup = [copy.deepcopy(basic_input_time_ch)]
+
+        xformer = DenormalizeTransformer(
+            low_rate=5.0, high_rate=35.0, distribution="normal"
+        )
+        output = xformer(basic_input_time_ch)
+
+        assert output.data.shape == basic_input_time_ch.data.shape
+        assert np.all(output.data >= 0)
+
+        # Offsets should be clipped to the specified range
+        assert xformer.state.offsets is not None
+        assert np.all(xformer.state.offsets >= 5.0)
+        assert np.all(xformer.state.offsets <= 35.0)
+
+        # Gains should be offsets / 3.29
+        assert np.allclose(xformer.state.gains, xformer.state.offsets / 3.29)
+
+        assert_messages_equal([basic_input_time_ch], backup)
+
+    def test_constant_distribution(self, basic_input_time_ch):
+        """Test denormalization with constant distribution."""
+        backup = [copy.deepcopy(basic_input_time_ch)]
+
+        low_rate = 10.0
+        high_rate = 30.0
+        expected_offset = (low_rate + high_rate) / 2.0  # 20.0
+
+        xformer = DenormalizeTransformer(
+            low_rate=low_rate, high_rate=high_rate, distribution="constant"
+        )
+        output = xformer(basic_input_time_ch)
+
+        assert output.data.shape == basic_input_time_ch.data.shape
+        assert np.all(output.data >= 0)
+
+        # All offsets should be exactly the midpoint
+        assert xformer.state.offsets is not None
+        assert np.allclose(xformer.state.offsets, expected_offset)
+
+        # Gains should be offsets / 3.29
+        expected_gain = expected_offset / 3.29
+        assert np.allclose(xformer.state.gains, expected_gain)
+
+        assert_messages_equal([basic_input_time_ch], backup)
+
+    def test_invalid_distribution(self, basic_input_time_ch):
+        """Test that invalid distribution raises ValueError."""
+        xformer = DenormalizeTransformer(distribution="invalid_dist")
+
+        with pytest.raises(ValueError, match="Invalid distribution"):
+            xformer(basic_input_time_ch)
+
+    def test_ch_time_axis_order(self, basic_input_ch_time):
+        """Test denormalization with ch x time axis order."""
+        backup = [copy.deepcopy(basic_input_ch_time)]
+
+        xformer = DenormalizeTransformer(distribution="constant")
+        output = xformer(basic_input_ch_time)
+
+        assert output.data.shape == basic_input_ch_time.data.shape
+        assert output.dims == basic_input_ch_time.dims
+        assert np.all(output.data >= 0)
+
+        # When ch is axis 0, shape should be (nch, 1)
+        n_chans = basic_input_ch_time.data.shape[0]
+        assert xformer.state.offsets.shape == (n_chans, 1)
+        assert xformer.state.gains.shape == (n_chans, 1)
+
+        assert_messages_equal([basic_input_ch_time], backup)
+
+    def test_output_clipping(self):
+        """Test that negative values are clipped to zero."""
+        n_times = 50
+        n_chans = 2
+        # Create very negative input data that should result in negative output before clipping
+        data = np.full((n_times, n_chans), -10.0)  # Very negative values
+        msg_in = AxisArray(
+            data=data,
+            dims=["time", "ch"],
+            axes=frozendict({"time": AxisArray.TimeAxis(fs=100.0)}),
+        )
+
+        xformer = DenormalizeTransformer(
+            low_rate=2.0, high_rate=5.0, distribution="constant"
+        )
+        output = xformer(msg_in)
+
+        # All values should be >= 0 due to clipping
+        assert np.all(output.data >= 0)
+        # With very negative input, most/all should be clipped to 0
+        assert np.sum(output.data == 0) > 0
+
+    def test_multiple_messages_same_state(self, basic_input_time_ch):
+        """Test that state is preserved across multiple messages."""
+        xformer = DenormalizeTransformer(distribution="uniform")
+
+        # First message initializes state
+        _ = xformer(basic_input_time_ch)
+        gains_after_first = xformer.state.gains.copy()
+        offsets_after_first = xformer.state.offsets.copy()
+
+        # Second message should use same state
+        _ = xformer(basic_input_time_ch)
+
+        assert np.array_equal(xformer.state.gains, gains_after_first)
+        assert np.array_equal(xformer.state.offsets, offsets_after_first)
+
+    def test_denormalization_formula(self):
+        """Test that the denormalization formula is applied correctly."""
+        n_times = 10
+        n_chans = 2
+        data = np.ones((n_times, n_chans))  # All ones for predictable output
+        msg_in = AxisArray(
+            data=data,
+            dims=["time", "ch"],
+            axes=frozendict({"time": AxisArray.TimeAxis(fs=100.0)}),
+        )
+
+        low_rate = 20.0
+        high_rate = 20.0  # Same as low to make offset = 20.0
+        xformer = DenormalizeTransformer(
+            low_rate=low_rate, high_rate=high_rate, distribution="constant"
+        )
+        output = xformer(msg_in)
+
+        # With constant distribution and equal rates, offset = 20.0
+        # gain = offset / 3.29 = 20.0 / 3.29
+        expected_offset = 20.0
+        expected_gain = expected_offset / 3.29
+        # output = data * gain + offset = 1.0 * gain + offset
+        expected_output = 1.0 * expected_gain + expected_offset
+
+        assert np.allclose(output.data, expected_output)
+
+    def test_settings_defaults(self):
+        """Test that DenormalizeSettings has correct defaults."""
+        settings = DenormalizeSettings()
+        assert settings.low_rate == 2.0
+        assert settings.high_rate == 40.0
+        assert settings.distribution == "uniform"
+
+    def test_settings_custom(self):
+        """Test custom DenormalizeSettings values."""
+        settings = DenormalizeSettings(
+            low_rate=5.0, high_rate=50.0, distribution="normal"
+        )
+        assert settings.low_rate == 5.0
+        assert settings.high_rate == 50.0
+        assert settings.distribution == "normal"
+
+    def test_state_initialization(self, basic_input_time_ch):
+        """Test that state is properly initialized on first message."""
+        xformer = DenormalizeTransformer(distribution="constant")
+
+        # Before first message, state should have None values
+        assert xformer.state.gains is None
+        assert xformer.state.offsets is None
+
+        # After first message, state should be initialized
+        xformer(basic_input_time_ch)
+
+        assert xformer.state.gains is not None
+        assert xformer.state.offsets is not None
+
+    def test_output_not_sharing_memory(self, basic_input_time_ch):
+        """Test that output doesn't share memory with input."""
+        xformer = DenormalizeTransformer(distribution="constant")
+        output = xformer(basic_input_time_ch)
+
+        assert not np.may_share_memory(output.data, basic_input_time_ch.data)