From 11cd8fa16bc3130121519baf7d630baaa3bd177f Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Tue, 6 Oct 2026 22:20:24 -0400 Subject: [PATCH] Declare microvolts as the unit of the velocity systems' output Velocity2Spike and Velocity2LFP now stamp attrs["unit"] = "microvolts" on every output message. The amplitude is set by the DNSS spike templates (microvolts) and the LFP scale is chosen to match, but nothing declared it, so consumers such as an LSL outlet could not report a unit. VelocityEncoder inherits the unit through Add and LineNoise without an extra stage. The stamping stage is private to the system package. --- src/ezmsg/simbiophys/dnss/wfs.py | 4 ++ src/ezmsg/simbiophys/system/_signal_unit.py | 29 ++++++++ .../simbiophys/system/velocity2ecephys.py | 8 ++- src/ezmsg/simbiophys/system/velocity2lfp.py | 14 +++- src/ezmsg/simbiophys/system/velocity2spike.py | 11 ++- .../integration/test_velocity_encoder_unit.py | 68 +++++++++++++++++++ tests/unit/test_signal_unit.py | 28 ++++++++ 7 files changed, 153 insertions(+), 9 deletions(-) create mode 100644 src/ezmsg/simbiophys/system/_signal_unit.py create mode 100644 tests/integration/test_velocity_encoder_unit.py create mode 100644 tests/unit/test_signal_unit.py diff --git a/src/ezmsg/simbiophys/dnss/wfs.py b/src/ezmsg/simbiophys/dnss/wfs.py index 59e84bf..efa646c 100644 --- a/src/ezmsg/simbiophys/dnss/wfs.py +++ b/src/ezmsg/simbiophys/dnss/wfs.py @@ -2,6 +2,10 @@ import numpy as np +WF_UNIT = "microvolts" +"""Unit of the values in :data:`wf_orig`, written as an AxisArray's ``unit`` attr +by systems whose output amplitude is set by these waveforms.""" + # fmt: off wf_orig = np.array([ [ diff --git a/src/ezmsg/simbiophys/system/_signal_unit.py b/src/ezmsg/simbiophys/system/_signal_unit.py new file mode 100644 index 0000000..b8c3981 --- /dev/null +++ b/src/ezmsg/simbiophys/system/_signal_unit.py @@ -0,0 +1,29 @@ +"""Declare the physical unit of a simulated signal on its messages. + +The systems in this package manufacture their amplitudes in a fixed unit (the +DNSS spike templates are microvolts and the LFP scale is chosen to match), but +none of the generators along the way know it. This stage writes it to the +message's ``unit`` attr at the end of each system so consumers -- e.g. an LSL +outlet, which publishes attrs as stream metadata -- can report it. +""" + +import ezmsg.core as ez +from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit +from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.messages.util import replace + + +class SignalUnitSettings(ez.Settings): + unit: str + """Value written to the output's ``unit`` attr.""" + + +class SignalUnitTransformer(BaseTransformer[SignalUnitSettings, AxisArray, AxisArray]): + def _process(self, message: AxisArray) -> AxisArray: + if message.attrs.get("unit") == self.settings.unit: + return message + return replace(message, attrs={**message.attrs, "unit": self.settings.unit}) + + +class SignalUnit(BaseTransformerUnit[SignalUnitSettings, AxisArray, AxisArray, SignalUnitTransformer]): + SETTINGS = SignalUnitSettings diff --git a/src/ezmsg/simbiophys/system/velocity2ecephys.py b/src/ezmsg/simbiophys/system/velocity2ecephys.py index dc3263b..479effc 100644 --- a/src/ezmsg/simbiophys/system/velocity2ecephys.py +++ b/src/ezmsg/simbiophys/system/velocity2ecephys.py @@ -58,7 +58,7 @@ class VelocityEncoderSettings(ez.Settings): """Number of cosine-encoded LFP sources.""" drift_scale: float = 4.0 - """Amplitude of always-on slow 1/f drift added to each LFP source before + """Amplitude (microvolts) of always-on slow 1/f drift added to each LFP source before mixing (a shared low-frequency field drift across channels). Velocity- independent, so baseline wander is present even at rest. Set to 0 to disable.""" @@ -67,7 +67,7 @@ class VelocityEncoderSettings(ez.Settings): line noise (pass-through). Added to every channel of the final signal.""" line_noise_amp: float = 10.0 - """Amplitude of the line-noise sinusoid (same units as the output signal).""" + """Amplitude of the line-noise sinusoid in microvolts (the output's unit).""" line_noise_drift_rate: float = 0.002 """Line-noise frequency drift rate in Hz per second (recording-clock drift).""" @@ -95,7 +95,9 @@ class VelocityEncoder(ez.Collection): Output: AxisArray with shape (M, output_ch) containing combined spike and LFP - signals at output_fs sampling rate. + signals at output_fs sampling rate, in microvolts (``attrs["unit"]``). + Both branches declare the unit and ``Add`` and line noise keep it, so + no stage here stamps it again. Example: >>> encoder = VelocityEncoder(VelocityEncoderSettings( diff --git a/src/ezmsg/simbiophys/system/velocity2lfp.py b/src/ezmsg/simbiophys/system/velocity2lfp.py index 74d93a6..77458c4 100644 --- a/src/ezmsg/simbiophys/system/velocity2lfp.py +++ b/src/ezmsg/simbiophys/system/velocity2lfp.py @@ -32,7 +32,9 @@ from ..baseline_drift import BaselineDriftSettings, BaselineDriftUnit from ..cosine_encoder import CosineEncoderSettings, CosineEncoderUnit +from ..dnss.wfs import WF_UNIT from ..dynamic_colored_noise import DynamicColoredNoiseSettings, DynamicColoredNoiseUnit +from ._signal_unit import SignalUnit, SignalUnitSettings def _make_mixing_weights(n_in: int, output_ch: int, seed: int) -> np.ndarray: @@ -71,7 +73,7 @@ class Velocity2LFPSettings(ez.Settings): max_velocity: float = 315.0 drift_scale: float = 4.0 - """Amplitude of always-on slow 1/f drift added to each LFP source before + """Amplitude (microvolts) of always-on slow 1/f drift added to each LFP source before mixing, so it appears as a shared low-frequency field drift across output channels. Velocity-independent, so baseline wander is present even at rest. Set to 0 to disable. Roughly ~20% of the per-channel LFP std at defaults.""" @@ -109,7 +111,10 @@ class Velocity2LFP(ez.Collection): Output: AxisArray with shape (M, output_ch) containing LFP-like colored noise - at output_fs sampling rate. + at output_fs sampling rate, in microvolts (``attrs["unit"]``). The noise + scale is chosen to sit under the spike templates of + :obj:`~ezmsg.simbiophys.system.velocity2spike.Velocity2Spike`, so the + two branches share a unit and can be summed. """ SETTINGS = Velocity2LFPSettings @@ -121,6 +126,7 @@ class Velocity2LFP(ez.Collection): PINK_NOISE = DynamicColoredNoiseUnit() BASELINE_DRIFT = BaselineDriftUnit() # Always-on slow 1/f wander per source MIX_NOISE = AffineTransform() # Project n_lfp_sources to output_ch sensors + UNIT = SignalUnit() OUTPUT_SIGNAL = ez.OutputTopic(AxisArray) def configure(self) -> None: @@ -181,6 +187,7 @@ def configure(self) -> None: _make_mixing_weights, output_ch=self.SETTINGS.output_ch, seed=self.SETTINGS.seed ) self.MIX_NOISE.apply_settings(AffineTransformSettings(weights=make_mixing_weights, axis="ch")) + self.UNIT.apply_settings(SignalUnitSettings(unit=WF_UNIT)) def network(self) -> ez.NetworkDefinition: return ( @@ -189,5 +196,6 @@ def network(self) -> ez.NetworkDefinition: (self.CLIP_BETA.OUTPUT_SIGNAL, self.PINK_NOISE.INPUT_SIGNAL), (self.PINK_NOISE.OUTPUT_SIGNAL, self.BASELINE_DRIFT.INPUT_SIGNAL), (self.BASELINE_DRIFT.OUTPUT_SIGNAL, self.MIX_NOISE.INPUT_SIGNAL), - (self.MIX_NOISE.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), + (self.MIX_NOISE.OUTPUT_SIGNAL, self.UNIT.INPUT_SIGNAL), + (self.UNIT.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), ) diff --git a/src/ezmsg/simbiophys/system/velocity2spike.py b/src/ezmsg/simbiophys/system/velocity2spike.py index b428aec..967c130 100644 --- a/src/ezmsg/simbiophys/system/velocity2spike.py +++ b/src/ezmsg/simbiophys/system/velocity2spike.py @@ -26,7 +26,8 @@ from ezmsg.util.messages.axisarray import AxisArray from ..cosine_encoder import CosineEncoderSettings, CosineEncoderUnit -from ..dnss.wfs import wf_orig +from ..dnss.wfs import WF_UNIT, wf_orig +from ._signal_unit import SignalUnit, SignalUnitSettings class Velocity2SpikeSettings(ez.Settings): @@ -71,7 +72,7 @@ class Velocity2Spike(ez.Collection): Output: AxisArray with shape (M, output_ch) containing spike waveforms at - output_fs sampling rate. + output_fs sampling rate, in microvolts (``attrs["unit"]``). """ SETTINGS = Velocity2SpikeSettings @@ -82,6 +83,7 @@ class Velocity2Spike(ez.Collection): CLIP_RATE = Clip() SPIKE_EVENT = PoissonEventUnit() WAVEFORMS = SparseKernelInserterUnit() + UNIT = SignalUnit() OUTPUT_SIGNAL = ez.OutputTopic(AxisArray) def configure(self) -> None: @@ -105,6 +107,8 @@ def configure(self) -> None: kernel=MultiKernel({i + 1: ArrayKernel(wf.astype(np.float32)) for i, wf in enumerate(wf_orig)}), ) ) + # The inserted templates set the amplitude, so they set the unit. + self.UNIT.apply_settings(SignalUnitSettings(unit=WF_UNIT)) def network(self) -> ez.NetworkDefinition: return ( @@ -112,5 +116,6 @@ def network(self) -> ez.NetworkDefinition: (self.RATE_ENCODER.OUTPUT_SIGNAL, self.CLIP_RATE.INPUT_SIGNAL), (self.CLIP_RATE.OUTPUT_SIGNAL, self.SPIKE_EVENT.INPUT_SIGNAL), (self.SPIKE_EVENT.OUTPUT_SIGNAL, self.WAVEFORMS.INPUT_SIGNAL), - (self.WAVEFORMS.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), + (self.WAVEFORMS.OUTPUT_SIGNAL, self.UNIT.INPUT_SIGNAL), + (self.UNIT.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), ) diff --git a/tests/integration/test_velocity_encoder_unit.py b/tests/integration/test_velocity_encoder_unit.py new file mode 100644 index 0000000..44c0148 --- /dev/null +++ b/tests/integration/test_velocity_encoder_unit.py @@ -0,0 +1,68 @@ +"""The velocity systems declare their output unit on every message.""" + +import os +from dataclasses import field + +import ezmsg.core as ez +from ezmsg.baseproc import Clock, ClockSettings +from ezmsg.sigproc.diff import DiffSettings, DiffUnit +from ezmsg.util.messagecodec import message_log +from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings +from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.terminate import TerminateOnTotal, TerminateOnTotalSettings + +from ezmsg.simbiophys.oscillator import SpiralGenerator, SpiralGeneratorSettings +from ezmsg.simbiophys.system.velocity2ecephys import VelocityEncoder, VelocityEncoderSettings +from tests.helpers.util import get_test_fn + +CURSOR_FS = 50.0 + + +class EncoderUnitTestSettings(ez.Settings): + encoder_settings: VelocityEncoderSettings + log_settings: MessageLoggerSettings + term_settings: TerminateOnTotalSettings = field(default_factory=TerminateOnTotalSettings) + + +class EncoderUnitTest(ez.Collection): + SETTINGS = EncoderUnitTestSettings + + CLOCK = Clock() + SPIRAL = SpiralGenerator() + DIFF = DiffUnit() + ENCODER = VelocityEncoder() + SINK = MessageLogger() + TERM = TerminateOnTotal() + + def configure(self) -> None: + self.CLOCK.apply_settings(ClockSettings(dispatch_rate=CURSOR_FS)) + self.SPIRAL.apply_settings(SpiralGeneratorSettings(fs=CURSOR_FS, r_mean=150.0, r_amp=150.0)) + self.DIFF.apply_settings(DiffSettings(axis="time", scale_by_fs=True)) + self.ENCODER.apply_settings(self.SETTINGS.encoder_settings) + self.SINK.apply_settings(self.SETTINGS.log_settings) + self.TERM.apply_settings(self.SETTINGS.term_settings) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.CLOCK.OUTPUT_SIGNAL, self.SPIRAL.INPUT_CLOCK), + (self.SPIRAL.OUTPUT_SIGNAL, self.DIFF.INPUT_SIGNAL), + (self.DIFF.OUTPUT_SIGNAL, self.ENCODER.INPUT_SIGNAL), + (self.ENCODER.OUTPUT_SIGNAL, self.SINK.INPUT_MESSAGE), + (self.SINK.OUTPUT_MESSAGE, self.TERM.INPUT_MESSAGE), + ) + + +def test_velocity_encoder_output_is_microvolts(test_name: str | None = None): + test_filename = get_test_fn(test_name) + settings = EncoderUnitTestSettings( + # Line noise on, so the unit is shown to survive every stage after Add. + encoder_settings=VelocityEncoderSettings(output_fs=3_000.0, output_ch=8, line_noise_freq=60.0), + log_settings=MessageLoggerSettings(output=test_filename), + term_settings=TerminateOnTotalSettings(total=20), + ) + ez.run(SYSTEM=EncoderUnitTest(settings)) + + messages: list[AxisArray] = list(message_log(test_filename)) + os.remove(test_filename) + assert messages + assert all(msg.attrs.get("unit") == "microvolts" for msg in messages) diff --git a/tests/unit/test_signal_unit.py b/tests/unit/test_signal_unit.py new file mode 100644 index 0000000..3c85e36 --- /dev/null +++ b/tests/unit/test_signal_unit.py @@ -0,0 +1,28 @@ +"""Unit tests for ezmsg.simbiophys.system._signal_unit.""" + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray + +from ezmsg.simbiophys.system._signal_unit import SignalUnitSettings, SignalUnitTransformer + + +def _msg(attrs: dict) -> AxisArray: + return AxisArray( + np.zeros((10, 2)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=1000.0)}, + attrs=attrs, + ) + + +def test_sets_unit_and_keeps_other_attrs(): + msg = _msg({"lsl_source_id": "abc"}) + out = SignalUnitTransformer(SignalUnitSettings(unit="microvolts"))(msg) + assert out.attrs == {"lsl_source_id": "abc", "unit": "microvolts"} + assert "unit" not in msg.attrs + assert out.data is msg.data + + +def test_replaces_a_different_unit(): + out = SignalUnitTransformer(SignalUnitSettings(unit="microvolts"))(_msg({"unit": "volts"})) + assert out.attrs["unit"] == "microvolts"