ezmsg-sigproc 2.13.1__tar.gz → 2.15.0__tar.gz

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.13.1
+Version: 2.15.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT
 License-File: LICENSE
 Requires-Python: >=3.10.15
 Requires-Dist: array-api-compat>=1.11.1
-Requires-Dist: ezmsg-baseproc>=1.1.0
+Requires-Dist: ezmsg-baseproc>=1.3.0
 Requires-Dist: ezmsg>=3.6.0
 Requires-Dist: numba>=0.61.0
 Requires-Dist: numpy>=1.26.0

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.10.15"
 dynamic = ["version"]
 dependencies = [
   "array-api-compat>=1.11.1",
-  "ezmsg-baseproc>=1.1.0",
+  "ezmsg-baseproc>=1.3.0",
   "ezmsg>=3.6.0",
   "numba>=0.61.0",
   "numpy>=1.26.0",

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/src/ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.13.1'
-__version_tuple__ = version_tuple = (2, 13, 1)
+__version__ = version = '2.15.0'
+__version_tuple__ = version_tuple = (2, 15, 0)
 
 __commit_id__ = commit_id = None

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/src/ezmsg/sigproc/affinetransform.py

@@ -23,6 +23,8 @@ from ezmsg.baseproc import (
 from ezmsg.util.messages.axisarray import AxisArray, AxisBase
 from ezmsg.util.messages.util import replace
 
+from ezmsg.sigproc.util.array import array_device, is_float_dtype, xp_asarray, xp_create
+
 
 def _find_block_diagonal_clusters(weights: np.ndarray) -> list[tuple[np.ndarray, np.ndarray]] | None:
     """Detect block-diagonal structure in a weight matrix.
@@ -245,13 +247,25 @@ class AffineTransformTransformer(
 
             self._state.new_axis = replace(message.axes[axis], data=np.array(new_labels))
 
-        # Convert to match message.data namespace for efficient operations in _process
+        # Convert to match message.data namespace and device for _process.
+        # Weights are stored as numpy float64 after cluster detection; some
+        # devices (e.g. MPS) don't support float64, so we downcast weight
+        # arrays to the message's dtype when the message is floating-point.
         xp = get_namespace(message.data)
+        dev = array_device(message.data)
+        msg_dt = message.data.dtype
+        # Downcast weights dtype only for float message data (avoids casting
+        # float weights to integer when message data happens to be int).
+        w_dt = msg_dt if is_float_dtype(xp, msg_dt) else None
         if self._state.weights is not None:
-            self._state.weights = xp.asarray(self._state.weights)
+            self._state.weights = xp_asarray(xp, self._state.weights, dtype=w_dt, device=dev)
         if self._state.clusters is not None:
             self._state.clusters = [
-                (xp.asarray(in_idx), xp.asarray(out_idx), xp.asarray(sub_w))
+                (
+                    xp_asarray(xp, in_idx, device=dev),
+                    xp_asarray(xp, out_idx, device=dev),
+                    xp_asarray(xp, sub_w, dtype=w_dt, device=dev),
+                )
                 for in_idx, out_idx, sub_w in self._state.clusters
             ]
 
@@ -345,7 +359,7 @@ class AffineTransformTransformer(
 
         # Pre-allocate output (omitted channels stay zero)
         out_shape = data.shape[:-1] + (self._state.n_out,)
-        result = xp.zeros(out_shape, dtype=data.dtype)
+        result = xp_create(xp.zeros, out_shape, dtype=data.dtype, device=array_device(data))
 
         for in_idx, out_idx, sub_weights in self._state.clusters:
             chunk = xp.take(data, in_idx, axis=data.ndim - 1)
@@ -371,7 +385,10 @@ class AffineTransformTransformer(
                 # The weights are stacked A|B where A is the transform and B is a single row
                 #  in the equation y = Ax + B. This supports NeuroKey's weights matrices.
                 sample_shape = data.shape[:axis_idx] + (1,) + data.shape[axis_idx + 1 :]
-                data = xp.concat((data, xp.ones(sample_shape, dtype=data.dtype)), axis=axis_idx)
+                data = xp.concat(
+                    (data, xp_create(xp.ones, sample_shape, dtype=data.dtype, device=array_device(data))),
+                    axis=axis_idx,
+                )
 
             if axis_idx in [-1, len(message.dims) - 1]:
                 data = xp.matmul(data, self._state.weights)

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/src/ezmsg/sigproc/rollingscaler.py

@@ -12,8 +12,6 @@ from ezmsg.baseproc import (
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ezmsg.sigproc.sampler import SampleMessage
-
 
 class RollingScalerSettings(ez.Settings):
     axis: str = "time"
@@ -168,8 +166,8 @@ class RollingScalerProcessor(BaseAdaptiveTransformer[RollingScalerSettings, Axis
 
         self._state.samples.append((n_b, mean_b, M2_b))
 
-    def partial_fit(self, message: SampleMessage) -> None:
-        x = message.sample.data
+    def partial_fit(self, message: AxisArray) -> None:
+        x = message.data
         self._add_batch_stats(x)
 
     def _process(self, message: AxisArray) -> AxisArray:

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/src/ezmsg/sigproc/sampler.py

@@ -21,7 +21,7 @@ from ezmsg.util.messages.util import replace
 
 from .util.axisarray_buffer import HybridAxisArrayBuffer
 from .util.buffer import UpdateStrategy
-from .util.message import SampleMessage, SampleTriggerMessage
+from .util.message import SampleTriggerMessage
 from .util.profile import profile_subpub
 
 
@@ -75,17 +75,7 @@ class SamplerState:
 
 
 class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, AxisArray, SamplerState]):
-    def __call__(self, message: AxisArray | SampleTriggerMessage) -> list[SampleMessage]:
-        # TODO: Currently we have a single entry point that accepts both
-        #  data and trigger messages and we choose a code path based on
-        #  the message type. However, in the future we will likely replace
-        #  SampleTriggerMessage with an agumented form of AxisArray,
-        #  leveraging its attrs field, which makes this a bit harder.
-        #  We should probably force callers of this object to explicitly
-        #  call `push_trigger` for trigger messages. This will also
-        #  simplify typing somewhat because `push_trigger` should not
-        #  return anything yet we currently have it returning an empty
-        #  list just to be compatible with __call__.
+    def __call__(self, message: AxisArray | SampleTriggerMessage) -> list[AxisArray]:
         if isinstance(message, AxisArray):
             return super().__call__(message)
         else:
@@ -109,7 +99,7 @@ class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, Axi
             self._state.triggers = deque()
         self._state.triggers.clear()
 
-    def _process(self, message: AxisArray) -> list[SampleMessage]:
+    def _process(self, message: AxisArray) -> list[AxisArray]:
         self._state.buffer.write(message)
 
         # How much data in the buffer?
@@ -119,7 +109,7 @@ class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, Axi
         )
 
         # Process in reverse order so that we can remove triggers safely as we iterate.
-        msgs_out: list[SampleMessage] = []
+        msgs_out: list[AxisArray] = []
         for trig_ix in range(len(self._state.triggers) - 1, -1, -1):
             trig = self._state.triggers[trig_ix]
             if trig.period is None:
@@ -152,13 +142,13 @@ class SamplerTransformer(BaseStatefulTransformer[SamplerSettings, AxisArray, Axi
             # Note: buffer will trim itself as needed based on buffer_dur.
 
             # Prepare output and drop trigger
-            msgs_out.append(SampleMessage(trigger=copy.copy(trig), sample=buff_axarr))
+            msgs_out.append(replace(buff_axarr, attrs={**buff_axarr.attrs, "trigger": copy.copy(trig)}))
             del self._state.triggers[trig_ix]
 
         msgs_out.reverse()  # in-place
         return msgs_out
 
-    def push_trigger(self, message: SampleTriggerMessage) -> list[SampleMessage]:
+    def push_trigger(self, message: SampleTriggerMessage) -> list[AxisArray]:
         # Input is a trigger message that we will use to sample the buffer.
 
         if self._state.buffer is None:
@@ -198,7 +188,7 @@ class Sampler(BaseTransformerUnit[SamplerSettings, AxisArray, AxisArray, Sampler
     SETTINGS = SamplerSettings
 
     INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
-    OUTPUT_SIGNAL = ez.OutputStream(SampleMessage)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
     @ez.subscriber(INPUT_TRIGGER)
     async def on_trigger(self, msg: SampleTriggerMessage) -> None:
@@ -228,7 +218,8 @@ def sampler(
 
     Returns:
         A generator that expects `.send` either an :obj:`AxisArray` containing streaming data messages,
-        or a :obj:`SampleTriggerMessage` containing a trigger, and yields the list of :obj:`SampleMessage` s.
+        or a :obj:`SampleTriggerMessage` containing a trigger, and yields the list of :obj:`AxisArray` s
+        with trigger info stored in ``attrs["trigger"]``.
     """
     return SamplerTransformer(
         settings=SamplerSettings(

ezmsg_sigproc-2.15.0/src/ezmsg/sigproc/util/array.py

@@ -0,0 +1,59 @@
+"""Portable helpers for Array API interoperability.
+
+These utilities smooth over differences between Array API libraries
+(NumPy, PyTorch, MLX, CuPy, etc.) — in particular around ``device``
+placement and ``dtype`` introspection, which are not uniformly supported.
+"""
+
+import numpy as np
+
+
+def array_device(x):
+    """Return the device of an array, or ``None`` for device-less libraries."""
+    try:
+        from array_api_compat import device
+
+        return device(x)
+    except (AttributeError, TypeError):
+        return None
+
+
+def xp_asarray(xp, obj, *, dtype=None, device=None):
+    """Portable ``xp.asarray`` that omits unsupported kwargs.
+
+    Some Array API libraries (e.g. MLX) don't accept a ``device`` keyword.
+    This helper builds the kwargs dict dynamically so that only supported
+    arguments are forwarded.
+    """
+    kwargs = {}
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    if device is not None:
+        kwargs["device"] = device
+    return xp.asarray(obj, **kwargs)
+
+
+def xp_create(fn, *args, dtype=None, device=None, **extra):
+    """Call a creation function (``zeros``, ``ones``, ``eye``) portably.
+
+    Omits ``device`` if it is ``None`` (for libraries that don't support it).
+    """
+    kwargs = dict(extra)
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    if device is not None:
+        kwargs["device"] = device
+    return fn(*args, **kwargs)
+
+
+def is_float_dtype(xp, dtype) -> bool:
+    """Check whether *dtype* is a real floating-point type, portably."""
+    try:
+        return xp.isdtype(dtype, "real floating")
+    except AttributeError:
+        pass
+    # Fallback for libraries without isdtype (e.g. MLX).
+    try:
+        return xp.issubdtype(dtype, xp.floating)
+    except (AttributeError, TypeError):
+        return np.issubdtype(np.dtype(dtype), np.floating)

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/tests/integration/ezmsg/test_sampler_system.py

@@ -90,8 +90,10 @@ def test_sampler_system(test_name: str | None = None):
     os.remove(test_filename)
     ez.logger.info(f"Analyzing recording of {len(messages)} messages...")
     assert len(messages) >= n_msgs
-    assert all([_.sample.data.shape == (int(freq * sample_dur), 1) for _ in messages])
+    assert all([_.data.shape == (int(freq * sample_dur), 1) for _ in messages])
     # Test the sample window slice vs the trigger timestamps
-    latencies = [_.sample.axes["time"].offset - (_.trigger.timestamp + _.trigger.period[0]) for _ in messages]
+    latencies = [
+        _.axes["time"].offset - (_.attrs["trigger"].timestamp + _.attrs["trigger"].period[0]) for _ in messages
+    ]
     assert all([0 <= _ < 1 / freq for _ in latencies])
     # Given that the input is a pure sinusoid, we could test that the signal has expected characteristics.

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/tests/unit/test_rollingscaler.py

@@ -6,7 +6,6 @@ from ezmsg.sigproc.rollingscaler import (
     RollingScalerProcessor,
     RollingScalerSettings,
 )
-from ezmsg.sigproc.sampler import SampleMessage
 
 
 def _axisarray_from_ndarray(x: np.ndarray, fs: float = 100.0, t0: float = 0.0) -> AxisArray:
@@ -47,7 +46,7 @@ def test_partial_fit_updates_and_process_uses_snapshot():
             [5.0, 3.0],
         ]
     )
-    proc.partial_fit(SampleMessage(trigger=None, sample=_axisarray_from_ndarray(epoch)))
+    proc.partial_fit(_axisarray_from_ndarray(epoch))
 
     x = np.array([[1.0, 3.0], [3.0, 3.0], [5.0, 3.0]])
     msg = _axisarray_from_ndarray(x)
@@ -68,7 +67,7 @@ def test_rolling_window_drops_oldest():
 
     def add_epoch(vals):
         arr = np.array(vals, dtype=float).reshape(-1, 1)
-        proc.partial_fit(SampleMessage(trigger=None, sample=_axisarray_from_ndarray(arr)))
+        proc.partial_fit(_axisarray_from_ndarray(arr))
 
     add_epoch([1, 1, 1, 1])
     add_epoch([3, 3])

{ezmsg_sigproc-2.13.1 → ezmsg_sigproc-2.15.0}/tests/unit/test_sampler.py

@@ -70,9 +70,9 @@ def test_sampler():
 
     assert len(samples) == n_trigs
     # Check sample data size. Note: sampler puts the time axis first.
-    assert all([_.sample.data.shape == (int(fs * period_dur), n_chans) for _ in samples])
+    assert all([_.data.shape == (int(fs * period_dur), n_chans) for _ in samples])
     # Compare the sample window slice against the trigger timestamps
-    latencies = [_.sample.axes["time"].offset - (_.trigger.timestamp + _.trigger.period[0]) for _ in samples]
+    latencies = [_.axes["time"].offset - (_.attrs["trigger"].timestamp + _.attrs["trigger"].period[0]) for _ in samples]
     assert all([0 <= _ < 1 / fs for _ in latencies])
     # Check the sample trigger value matches the trigger input.
-    assert all([_.trigger.value == ["Start", "Stop"][ix % 2] for ix, _ in enumerate(samples)])
+    assert all([_.attrs["trigger"].value == ["Start", "Stop"][ix % 2] for ix, _ in enumerate(samples)])