Declare microvolts as the unit of the velocity systems' output

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([
     [

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

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(

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),
         )

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,12 +107,15 @@ 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 (
             (self.INPUT_SIGNAL, self.RATE_ENCODER.INPUT_SIGNAL),
             (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),
         )

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)

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"