diff --git a/.gitignore b/.gitignore index b78272ee..b2b010e7 100644 --- a/.gitignore +++ b/.gitignore @@ -146,3 +146,4 @@ cython_debug/ src/ezmsg/sigproc/__version__.py uv.lock +*.local.json diff --git a/docs/source/conf.py b/docs/source/conf.py index b2b43cdd..a81eddad 100644 --- a/docs/source/conf.py +++ b/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"] diff --git a/docs/source/guides/sigproc/processors.rst b/docs/source/guides/sigproc/processors.rst index dff5ca35..8662d8ea 100644 --- a/docs/source/guides/sigproc/processors.rst +++ b/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 -------------------------- diff --git a/src/ezmsg/sigproc/denormalize.py b/src/ezmsg/sigproc/denormalize.py new file mode 100644 index 00000000..e3743e20 --- /dev/null +++ b/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 diff --git a/tests/unit/test_denormalize.py b/tests/unit/test_denormalize.py new file mode 100644 index 00000000..37d99a7a --- /dev/null +++ b/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)