From a32fc1777f3a7802aa5aef978e1c0fc63fb3520a Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 7 Feb 2026 18:03:37 -0500 Subject: [PATCH 1/5] cosine_encoder array API --- src/ezmsg/simbiophys/cosine_encoder.py | 24 +++++--- tests/unit/test_cosine_encoder.py | 79 ++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/src/ezmsg/simbiophys/cosine_encoder.py b/src/ezmsg/simbiophys/cosine_encoder.py index 808faee..e2cd649 100644 --- a/src/ezmsg/simbiophys/cosine_encoder.py +++ b/src/ezmsg/simbiophys/cosine_encoder.py @@ -23,11 +23,12 @@ AxisArray with shape (n_samples, output_ch) containing encoded values. """ +import typing from pathlib import Path import ezmsg.core as ez import numpy as np -import numpy.typing as npt +from array_api_compat import get_namespace from ezmsg.baseproc import ( BaseStatefulTransformer, BaseTransformerUnit, @@ -82,10 +83,10 @@ class CosineEncoderState: ch_axis: Pre-built channel axis for output messages. """ - baseline: npt.NDArray[np.floating] | None = None - modulation: npt.NDArray[np.floating] | None = None - pd: npt.NDArray[np.floating] | None = None - speed_modulation: npt.NDArray[np.floating] | None = None + baseline: typing.Any = None + modulation: typing.Any = None + pd: typing.Any = None + speed_modulation: typing.Any = None ch_axis: AxisArray.CoordinateAxis | None = None @property @@ -216,9 +217,18 @@ def _reset_state(self, message: AxisArray) -> None: seed=self.settings.seed, ) + # Convert state parameters to the input array's backend + xp = get_namespace(message.data) + if xp is not np: + self.state.baseline = xp.asarray(self.state.baseline) + self.state.modulation = xp.asarray(self.state.modulation) + self.state.pd = xp.asarray(self.state.pd) + self.state.speed_modulation = xp.asarray(self.state.speed_modulation) + def _process(self, message: AxisArray) -> AxisArray: """Transform polar coordinates to encoded output.""" - polar = np.asarray(message.data, dtype=np.float64) + polar = message.data + xp = get_namespace(polar) if polar.ndim != 2 or polar.shape[1] != 2: raise ValueError(f"Expected polar coords with shape (n_samples, 2), got {polar.shape}") @@ -231,7 +241,7 @@ def _process(self, message: AxisArray) -> AxisArray: # State arrays are pre-shaped to (1, output_ch) for broadcasting output = ( self.state.baseline - + self.state.modulation * magnitude * np.cos(angle - self.state.pd) + + self.state.modulation * magnitude * xp.cos(angle - self.state.pd) + self.state.speed_modulation * magnitude ) diff --git a/tests/unit/test_cosine_encoder.py b/tests/unit/test_cosine_encoder.py index a0ff514..13954c5 100644 --- a/tests/unit/test_cosine_encoder.py +++ b/tests/unit/test_cosine_encoder.py @@ -1,5 +1,8 @@ """Unit tests for ezmsg.simbiophys.cosine_encoder module.""" +import platform +import time + import numpy as np import pytest from ezmsg.util.messages.axisarray import AxisArray @@ -10,6 +13,11 @@ CosineEncoderTransformer, ) +requires_apple_silicon = pytest.mark.skipif( + platform.machine() != "arm64" or platform.system() != "Darwin", + reason="Requires Apple Silicon for MLX", +) + class TestCosineEncoderState: """Tests for CosineEncoderState.""" @@ -230,3 +238,74 @@ def test_multiple_samples(self): msg_out = transformer(msg_in) assert msg_out.data.shape == (n_samples, 5) + + +@requires_apple_silicon +def test_cosine_encoder_mlx_benchmark(): + """Benchmark CosineEncoderTransformer: numpy vs MLX input.""" + import mlx.core as mx + + n_samples = 500 + n_chunks = 200 + output_ch = 256 + fs = 100.0 + + settings = CosineEncoderSettings( + output_ch=output_ch, + baseline=10.0, + modulation=20.0, + speed_modulation=5.0, + seed=42, + ) + + # Pre-generate chunks as numpy + rng = np.random.default_rng(42) + np_chunks = [] + for i in range(n_chunks + 1): # +1 for warmup + magnitude = np.abs(rng.standard_normal((n_samples, 1))).astype(np.float32) + angle = rng.uniform(-np.pi, np.pi, (n_samples, 1)).astype(np.float32) + polar = np.hstack([magnitude, angle]) + np_chunks.append( + AxisArray( + polar, + dims=["time", "ch"], + axes={"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * n_samples / fs)}, + ) + ) + + # MLX versions + mx_chunks = [AxisArray(data=mx.array(chunk.data), dims=chunk.dims, axes=chunk.axes) for chunk in np_chunks] + + # --- Numpy --- + xformer_np = CosineEncoderTransformer(settings) + xformer_np(np_chunks[0]) # Warmup + + t0 = time.perf_counter() + np_outputs = [xformer_np(chunk) for chunk in np_chunks[1:]] + t_numpy = time.perf_counter() - t0 + + # --- MLX --- + xformer_mx = CosineEncoderTransformer(settings) + xformer_mx(mx_chunks[0]) # Warmup + mx.eval(xformer_mx(mx_chunks[0]).data) + + t0 = time.perf_counter() + mx_outputs = [xformer_mx(chunk) for chunk in mx_chunks[1:]] + for out in mx_outputs: + mx.eval(out.data) + t_mlx = time.perf_counter() - t0 + + # Verify output is MLX array + last_mx = mx_outputs[-1] + assert isinstance(last_mx.data, mx.array), f"Expected mx.array, got {type(last_mx.data)}" + + # Correctness: compare outputs + for np_out, mx_out in zip(np_outputs, mx_outputs): + np.testing.assert_allclose(np.asarray(mx_out.data), np_out.data, rtol=5e-3, atol=1e-4) + + print( + f"\n CosineEncoder benchmark ({n_chunks} chunks, {n_samples}×2 → {output_ch} ch):" + f"\n numpy: {t_numpy:.4f}s ({t_numpy / n_chunks * 1000:.2f} ms/chunk)" + f"\n mlx: {t_mlx:.4f}s ({t_mlx / n_chunks * 1000:.2f} ms/chunk)" + f"\n ratio (mlx/numpy): {t_mlx / t_numpy:.2f}x" + ) From aea256a5a5261e081cd394057420b96d8e63ab2a Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 7 Feb 2026 18:05:49 -0500 Subject: [PATCH 2/5] dynamic_colored_noise - Array API and big speedup from vectorizing over channels (swap loop order) --- src/ezmsg/simbiophys/dynamic_colored_noise.py | 199 ++++++++---------- tests/unit/test_dynamic_colored_noise.py | 93 +++++++- 2 files changed, 179 insertions(+), 113 deletions(-) diff --git a/src/ezmsg/simbiophys/dynamic_colored_noise.py b/src/ezmsg/simbiophys/dynamic_colored_noise.py index 5728fc0..5599361 100644 --- a/src/ezmsg/simbiophys/dynamic_colored_noise.py +++ b/src/ezmsg/simbiophys/dynamic_colored_noise.py @@ -10,11 +10,10 @@ Vol. 83, No. 5, May 1995, pages 802-827. """ -from dataclasses import dataclass, field - import ezmsg.core as ez import numpy as np import numpy.typing as npt +from array_api_compat import get_namespace from ezmsg.baseproc import ( BaseStatefulTransformer, BaseTransformerUnit, @@ -36,7 +35,6 @@ def compute_kasdin_coefficients(beta: float, n_poles: int) -> npt.NDArray[np.flo - 0: white noise - 1: pink noise (1/f) - 2: brown/red noise (1/f²) - TODO: jit / vectorize over multiple values of beta n_poles: Number of IIR filter poles. More poles extend accuracy to lower frequencies but increase computation. 5 is a reasonable default. @@ -51,15 +49,27 @@ def compute_kasdin_coefficients(beta: float, n_poles: int) -> npt.NDArray[np.flo return coeffs -@dataclass -class ColoredNoiseFilterState: - """State for a single channel's colored noise filter.""" +def compute_kasdin_coefficients_batch(betas: npt.NDArray, n_poles: int) -> npt.NDArray[np.float64]: + """Compute IIR filter coefficients for multiple β values simultaneously. + + Vectorized version of :func:`compute_kasdin_coefficients` that processes + an array of beta values in one call. - delay_line: npt.NDArray[np.float64] - """Previous output samples (filter memory).""" + Args: + betas: Array of spectral exponents, shape ``(n_betas,)``. + n_poles: Number of IIR filter poles. - coeffs: npt.NDArray[np.float64] - """Current filter coefficients (exponentially smoothed).""" + Returns: + Array of shape ``(n_betas, n_poles)`` containing filter coefficients. + """ + betas = np.asarray(betas, dtype=np.float64) + a = np.ones(len(betas), dtype=np.float64) + half_betas = betas / 2.0 + coeffs = np.zeros((len(betas), n_poles), dtype=np.float64) + for k in range(1, n_poles + 1): + a *= (k - 1 - half_betas) / k + coeffs[:, k - 1] = a + return coeffs class DynamicColoredNoiseSettings(ez.Settings): @@ -88,8 +98,11 @@ class DynamicColoredNoiseSettings(ez.Settings): @processor_state class DynamicColoredNoiseState: - filter_states: list[ColoredNoiseFilterState] = field(default_factory=list) - """Per-channel filter states.""" + delay_lines: npt.NDArray[np.float64] | None = None + """Filter delay lines, shape (n_channels, n_poles).""" + + coeffs: npt.NDArray[np.float64] | None = None + """Current filter coefficients, shape (n_channels, n_poles).""" rng: np.random.Generator | None = None """Random number generator.""" @@ -153,25 +166,15 @@ def _hash_message(self, message: AxisArray) -> int: def _reset_state(self, message: AxisArray) -> None: """Initialize filter states and compute timing parameters.""" # Determine number of channels - if message.data.ndim == 1: - n_channels = 1 - else: - n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 + n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 # Initialize RNG self._state.rng = np.random.default_rng(self.settings.seed) - # Initialize filter states for each channel + # Initialize vectorized filter state initial_coeffs = compute_kasdin_coefficients(self.settings.initial_beta, self.settings.n_poles) - - self._state.filter_states = [] - for _ in range(n_channels): - self._state.filter_states.append( - ColoredNoiseFilterState( - delay_line=np.zeros(self.settings.n_poles, dtype=np.float64), - coeffs=initial_coeffs.copy(), - ) - ) + self._state.delay_lines = np.zeros((n_channels, self.settings.n_poles), dtype=np.float64) + self._state.coeffs = np.tile(initial_coeffs, (n_channels, 1)) # Compute timing parameters (only depends on sample rate, not data) time_axis = message.axes.get("time") @@ -194,6 +197,7 @@ def _reset_state(self, message: AxisArray) -> None: def _process(self, message: AxisArray) -> AxisArray: """Generate colored noise based on input β values.""" + xp = get_namespace(message.data) beta_data = np.asarray(message.data, dtype=np.float64) # Handle 1D input @@ -202,6 +206,7 @@ def _process(self, message: AxisArray) -> AxisArray: beta_data = beta_data[:, np.newaxis] n_input_samples, n_channels = beta_data.shape + n_poles = self.settings.n_poles # Get precomputed timing parameters from state output_gain = self._state.output_gain @@ -213,7 +218,6 @@ def _process(self, message: AxisArray) -> AxisArray: input_offset = time_axis.offset if time_axis is not None else 0.0 # Calculate total output samples for this chunk - # Use remainder from previous chunk for continuity total_fractional = n_input_samples * samples_per_bin + self._state.sample_remainder n_output_samples = int(total_fractional) new_remainder = total_fractional - n_output_samples @@ -221,13 +225,13 @@ def _process(self, message: AxisArray) -> AxisArray: if n_output_samples == 0: # Not enough input to produce output yet self._state.sample_remainder = total_fractional - # Return empty output empty_data = np.zeros((0, n_channels), dtype=np.float64) if was_1d: empty_data = empty_data[:, 0] + out_data = xp.asarray(empty_data) if xp is not np else empty_data return replace( message, - data=empty_data, + data=out_data, dims=message.dims, axes={ **message.axes, @@ -235,108 +239,89 @@ def _process(self, message: AxisArray) -> AxisArray: }, ) - # Generate white noise for all output samples - white_noise = self._state.rng.standard_normal((n_output_samples, n_channels)) - output = np.zeros_like(white_noise) - - # Process each channel - for ch in range(n_channels): - output[:, ch] = self._process_resampled( - white_noise[:, ch], - beta_data[:, ch], - self._state.filter_states[ch], - samples_per_bin, - self._state.sample_remainder, - alpha, - ) + # Precompute target coefficients for all bins in batch. + # For non-finite betas, we'll use current coeffs (EMA no-op). + finite_mask = np.isfinite(beta_data) # (n_input, n_channels) - # Update remainder for next chunk - self._state.sample_remainder = new_remainder - - # Apply scaling - output *= self.settings.scale + # Collect all unique finite betas and batch-compute their coefficients + all_betas_flat = beta_data.ravel() # (n_input * n_channels,) + finite_flat = finite_mask.ravel() + # Build target coefficients array: (n_input, n_channels, n_poles) + target_coeffs_all = np.zeros((n_input_samples, n_channels, n_poles), dtype=np.float64) - # Handle 1D output if input was 1D - if was_1d: - output = output[:, 0] + if np.any(finite_flat): + finite_betas = all_betas_flat[finite_flat] + finite_coeffs = compute_kasdin_coefficients_batch(finite_betas, n_poles) + # Scatter back into the full array + target_coeffs_all.reshape(-1, n_poles)[finite_flat] = finite_coeffs - return replace( - message, - data=output, - dims=["time"] if was_1d else ["time", "ch"], - axes={ - **message.axes, - "time": replace(time_axis, gain=output_gain, offset=input_offset), - }, - ) + # Generate white noise for all output samples + white_noise = self._state.rng.standard_normal((n_output_samples, n_channels)) + output = np.zeros_like(white_noise) - def _process_resampled( - self, - white_noise: npt.NDArray[np.float64], - beta_values: npt.NDArray[np.float64], - fs: ColoredNoiseFilterState, - samples_per_bin: float, - initial_remainder: float, - alpha: float, - ) -> npt.NDArray[np.float64]: - """Process with resampling - each input β defines a bin of output samples. - - Args: - white_noise: Pre-generated white noise for output samples. - beta_values: Input β values (one per input bin). - fs: Filter state for this channel. - samples_per_bin: Number of output samples per input bin (may be fractional). - initial_remainder: Fractional sample offset from previous chunk. - alpha: Exponential smoothing factor (1 - exp(-dt/tau)). - - Returns: - Colored noise output samples. - """ - n_output = len(white_noise) - n_input = len(beta_values) - output = np.zeros(n_output, dtype=np.float64) - n_poles = self.settings.n_poles + # Get mutable references to state arrays + delay_lines = self._state.delay_lines # (n_channels, n_poles) + coeffs = self._state.coeffs # (n_channels, n_poles) # Track fractional position for bin boundaries - cumulative_samples = initial_remainder + cumulative_samples = self._state.sample_remainder - for bin_idx in range(n_input): - # Calculate how many output samples this bin contributes + for bin_idx in range(n_input_samples): + # Calculate output sample range for this bin next_cumulative = cumulative_samples + samples_per_bin bin_start_out = int(cumulative_samples) - bin_end_out = min(int(next_cumulative), n_output) + bin_end_out = min(int(next_cumulative), n_output_samples) if bin_end_out <= bin_start_out: - # This bin doesn't contribute any complete samples yet cumulative_samples = next_cumulative continue - # Get target coefficients for this bin's β - beta = beta_values[bin_idx] - # Handle nan/inf beta values by using current coefficients (no update) - if np.isfinite(beta): - target_coeffs = compute_kasdin_coefficients(beta, n_poles) - else: - target_coeffs = fs.coeffs + # Get target coefficients for this bin (n_channels, n_poles) + # For non-finite channels, use current coeffs (EMA no-op) + target = target_coeffs_all[bin_idx] # (n_channels, n_poles) + non_finite_ch = ~finite_mask[bin_idx] # (n_channels,) + if np.any(non_finite_ch): + target[non_finite_ch] = coeffs[non_finite_ch] - # Generate output samples for this bin + # Generate output samples for this bin, vectorized across channels for i in range(bin_start_out, bin_end_out): - if i >= n_output: + if i >= n_output_samples: break # Exponentially smooth coefficients toward target - fs.coeffs += alpha * (target_coeffs - fs.coeffs) + coeffs += alpha * (target - coeffs) - # IIR filter: x[n] = w[n] - sum(a_k * x[n-k]) - output[i] = white_noise[i] - np.dot(fs.coeffs, fs.delay_line) + # IIR filter: y[n] = x[n] - sum(a_k * y[n-k]) + output[i] = white_noise[i] - np.einsum("cp,cp->c", coeffs, delay_lines) - # Update delay line - fs.delay_line = np.roll(fs.delay_line, 1) - fs.delay_line[0] = output[i] + # Shift delay lines (no np.roll allocation) + delay_lines[:, 1:] = delay_lines[:, :-1] + delay_lines[:, 0] = output[i] cumulative_samples = next_cumulative - return output + # Update remainder for next chunk + self._state.sample_remainder = new_remainder + + # Apply scaling + output *= self.settings.scale + + # Handle 1D output if input was 1D + if was_1d: + output = output[:, 0] + + # Convert to input's array backend + out_data = xp.asarray(output) if xp is not np else output + + return replace( + message, + data=out_data, + dims=["time"] if was_1d else ["time", "ch"], + axes={ + **message.axes, + "time": replace(time_axis, gain=output_gain, offset=input_offset), + }, + ) class DynamicColoredNoiseUnit( diff --git a/tests/unit/test_dynamic_colored_noise.py b/tests/unit/test_dynamic_colored_noise.py index 44adc82..c996adf 100644 --- a/tests/unit/test_dynamic_colored_noise.py +++ b/tests/unit/test_dynamic_colored_noise.py @@ -1,5 +1,8 @@ """Unit tests for ezmsg.simbiophys.dynamic_colored_noise module.""" +import platform +import time + import numpy as np import pytest from ezmsg.util.messages.axisarray import AxisArray @@ -10,6 +13,11 @@ compute_kasdin_coefficients, ) +requires_apple_silicon = pytest.mark.skipif( + platform.machine() != "arm64" or platform.system() != "Darwin", + reason="Requires Apple Silicon for MLX", +) + class TestComputeKasdinCoefficients: """Tests for compute_kasdin_coefficients function.""" @@ -315,11 +323,11 @@ def test_state_preservation(self): transformer(msg1) # Check state exists - assert len(transformer.state.filter_states) == 1 - assert transformer.state.filter_states[0].delay_line is not None + assert transformer.state.delay_lines is not None + assert transformer.state.delay_lines.shape == (1, 5) # Delay line should not be all zeros after processing - assert not np.allclose(transformer.state.filter_states[0].delay_line, 0.0) + assert not np.allclose(transformer.state.delay_lines[0], 0.0) def test_smoothing_tau_zero_instant_change(self): """Test that smoothing_tau=0 gives instantaneous coefficient changes.""" @@ -337,7 +345,7 @@ def test_smoothing_tau_zero_instant_change(self): expected_coeffs = compute_kasdin_coefficients(1.0, 5) # With tau=0, coefficients should exactly match target - assert np.allclose(transformer.state.filter_states[0].coeffs, expected_coeffs) + assert np.allclose(transformer.state.coeffs[0], expected_coeffs) # Now change to beta=2 beta2 = np.full((100, 1), 2.0) @@ -346,7 +354,7 @@ def test_smoothing_tau_zero_instant_change(self): expected_coeffs_2 = compute_kasdin_coefficients(2.0, 5) # Should immediately be at new target - assert np.allclose(transformer.state.filter_states[0].coeffs, expected_coeffs_2) + assert np.allclose(transformer.state.coeffs[0], expected_coeffs_2) def test_smoothing_tau_time_constant_behavior(self): """Test that smoothing_tau controls the rate of coefficient change.""" @@ -376,7 +384,7 @@ def test_smoothing_tau_time_constant_behavior(self): expected_coeffs = initial_coeffs + expected_progress * (target_coeffs - initial_coeffs) # Allow some tolerance due to discrete-time approximation - actual_coeffs = transformer.state.filter_states[0].coeffs + actual_coeffs = transformer.state.coeffs[0] assert np.allclose(actual_coeffs, expected_coeffs, rtol=0.05) @@ -591,3 +599,76 @@ def test_empty_output_accumulation(self): total_out = msg_out.data.shape[0] + msg_out2.data.shape[0] assert total_out == 1 + + +@requires_apple_silicon +def test_dynamic_colored_noise_mlx_benchmark(): + """Benchmark DynamicColoredNoiseTransformer: numpy vs MLX input.""" + import mlx.core as mx + + n_input = 50 + n_chunks = 200 + n_channels = 16 + fs = 100.0 + output_fs = 1000.0 + + settings = DynamicColoredNoiseSettings( + output_fs=output_fs, + n_poles=5, + smoothing_tau=0.01, + initial_beta=1.0, + seed=42, + ) + + # Pre-generate chunks as numpy + rng = np.random.default_rng(42) + np_chunks = [] + for i in range(n_chunks + 1): # +1 for warmup + beta = rng.uniform(0.5, 2.0, (n_input, n_channels)).astype(np.float64) + np_chunks.append( + AxisArray( + beta, + dims=["time", "ch"], + axes={"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * n_input / fs)}, + ) + ) + + # MLX versions + mx_chunks = [AxisArray(data=mx.array(chunk.data), dims=chunk.dims, axes=chunk.axes) for chunk in np_chunks] + + # --- Numpy --- + xformer_np = DynamicColoredNoiseTransformer(settings) + xformer_np(np_chunks[0]) # Warmup + + t0 = time.perf_counter() + np_outputs = [xformer_np(chunk) for chunk in np_chunks[1:]] + t_numpy = time.perf_counter() - t0 + + # --- MLX --- + xformer_mx = DynamicColoredNoiseTransformer(settings) + xformer_mx(mx_chunks[0]) # Warmup + mx.eval(xformer_mx(mx_chunks[0]).data) + + t0 = time.perf_counter() + mx_outputs = [xformer_mx(chunk) for chunk in mx_chunks[1:]] + for out in mx_outputs: + mx.eval(out.data) + t_mlx = time.perf_counter() - t0 + + # Verify output is MLX array + last_mx = mx_outputs[-1] + assert isinstance(last_mx.data, mx.array), f"Expected mx.array, got {type(last_mx.data)}" + + # Correctness: compare first chunk outputs (seeds differ due to warmup count, + # but shapes should match and values should be finite) + for np_out, mx_out in zip(np_outputs[:5], mx_outputs[:5]): + mx_data = np.asarray(mx_out.data) + assert mx_data.shape == np_out.data.shape + assert np.all(np.isfinite(mx_data)) + + print( + f"\n DynamicColoredNoise benchmark ({n_chunks} chunks, {n_input}×{n_channels} → {output_fs}Hz):" + f"\n numpy: {t_numpy:.4f}s ({t_numpy / n_chunks * 1000:.2f} ms/chunk)" + f"\n mlx: {t_mlx:.4f}s ({t_mlx / n_chunks * 1000:.2f} ms/chunk)" + f"\n ratio (mlx/numpy): {t_mlx / t_numpy:.2f}x" + ) From cc0aa56d21052047d2b416c26d0371b75879647b Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 7 Feb 2026 18:06:11 -0500 Subject: [PATCH 3/5] array-api-compat --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index a2c1086..217de36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ readme = "README.md" requires-python = ">=3.10.15" dynamic = ["version"] dependencies = [ + "array-api-compat>=1.11.0", "ezmsg>=3.6.0", "ezmsg-baseproc>=1.2.1", "ezmsg-event>=0.6.0", From 6d58705c3eff4831dbd685853b249e597d30bbd4 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 7 Feb 2026 18:29:43 -0500 Subject: [PATCH 4/5] cleanup ezmsg.simbiophys namespace --- src/ezmsg/simbiophys/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ezmsg/simbiophys/__init__.py b/src/ezmsg/simbiophys/__init__.py index b6d6e4f..8db5aa4 100644 --- a/src/ezmsg/simbiophys/__init__.py +++ b/src/ezmsg/simbiophys/__init__.py @@ -36,12 +36,12 @@ # Dynamic Colored Noise from .dynamic_colored_noise import ( - ColoredNoiseFilterState, DynamicColoredNoiseSettings, DynamicColoredNoiseState, DynamicColoredNoiseTransformer, DynamicColoredNoiseUnit, compute_kasdin_coefficients, + compute_kasdin_coefficients_batch, ) # EEG @@ -112,12 +112,12 @@ "CosineEncoderTransformer", "CosineEncoderUnit", # Dynamic Colored Noise - "ColoredNoiseFilterState", "DynamicColoredNoiseSettings", "DynamicColoredNoiseState", "DynamicColoredNoiseTransformer", "DynamicColoredNoiseUnit", "compute_kasdin_coefficients", + "compute_kasdin_coefficients_batch", # DNSS LFP "DNSSLFPProducer", "DNSSLFPSettings", From 13b80900cb2cd776d98a4035f6b276f0fbae08ba Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 7 Feb 2026 18:32:34 -0500 Subject: [PATCH 5/5] mlx as test dependency on macos --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 217de36..8c2f46d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ lint = [ test = [ "pytest>=8.0.0", "pytest-asyncio>=0.24.0", + "mlx>=0.18.0; sys_platform == 'darwin' and platform_machine == 'arm64'", ] docs = [ "sphinx>=8.0",