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"