ezmsg-sigproc 1.8.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. ezmsg/sigproc/__version__.py +2 -2
  2. ezmsg/sigproc/activation.py +36 -39
  3. ezmsg/sigproc/adaptive_lattice_notch.py +231 -0
  4. ezmsg/sigproc/affinetransform.py +169 -163
  5. ezmsg/sigproc/aggregate.py +119 -104
  6. ezmsg/sigproc/bandpower.py +58 -52
  7. ezmsg/sigproc/base.py +1242 -0
  8. ezmsg/sigproc/butterworthfilter.py +37 -33
  9. ezmsg/sigproc/cheby.py +29 -17
  10. ezmsg/sigproc/combfilter.py +163 -0
  11. ezmsg/sigproc/decimate.py +19 -10
  12. ezmsg/sigproc/detrend.py +29 -0
  13. ezmsg/sigproc/diff.py +81 -0
  14. ezmsg/sigproc/downsample.py +78 -78
  15. ezmsg/sigproc/ewma.py +197 -0
  16. ezmsg/sigproc/extract_axis.py +41 -0
  17. ezmsg/sigproc/filter.py +257 -141
  18. ezmsg/sigproc/filterbank.py +247 -199
  19. ezmsg/sigproc/math/abs.py +17 -22
  20. ezmsg/sigproc/math/clip.py +24 -24
  21. ezmsg/sigproc/math/difference.py +34 -30
  22. ezmsg/sigproc/math/invert.py +13 -25
  23. ezmsg/sigproc/math/log.py +28 -33
  24. ezmsg/sigproc/math/scale.py +18 -26
  25. ezmsg/sigproc/quantize.py +71 -0
  26. ezmsg/sigproc/resample.py +298 -0
  27. ezmsg/sigproc/sampler.py +241 -259
  28. ezmsg/sigproc/scaler.py +55 -218
  29. ezmsg/sigproc/signalinjector.py +52 -43
  30. ezmsg/sigproc/slicer.py +81 -89
  31. ezmsg/sigproc/spectrogram.py +77 -75
  32. ezmsg/sigproc/spectrum.py +203 -168
  33. ezmsg/sigproc/synth.py +546 -393
  34. ezmsg/sigproc/transpose.py +131 -0
  35. ezmsg/sigproc/util/asio.py +156 -0
  36. ezmsg/sigproc/util/message.py +31 -0
  37. ezmsg/sigproc/util/profile.py +55 -12
  38. ezmsg/sigproc/util/typeresolution.py +83 -0
  39. ezmsg/sigproc/wavelets.py +154 -153
  40. ezmsg/sigproc/window.py +269 -211
  41. {ezmsg_sigproc-1.8.1.dist-info → ezmsg_sigproc-2.0.0.dist-info}/METADATA +2 -1
  42. ezmsg_sigproc-2.0.0.dist-info/RECORD +51 -0
  43. ezmsg_sigproc-1.8.1.dist-info/RECORD +0 -39
  44. {ezmsg_sigproc-1.8.1.dist-info → ezmsg_sigproc-2.0.0.dist-info}/WHEEL +0 -0
  45. {ezmsg_sigproc-1.8.1.dist-info → ezmsg_sigproc-2.0.0.dist-info}/licenses/LICENSE.txt +0 -0
ezmsg/sigproc/window.py CHANGED
@@ -3,19 +3,23 @@ import traceback
3
3
  import typing
4
4
 
5
5
  import ezmsg.core as ez
6
- import numpy as np
7
6
  import numpy.typing as npt
8
7
  import sparse
8
+ from array_api_compat import is_pydata_sparse_namespace, get_namespace
9
9
  from ezmsg.util.messages.axisarray import (
10
10
  AxisArray,
11
11
  slice_along_axis,
12
12
  sliding_win_oneaxis,
13
13
  replace,
14
14
  )
15
- from ezmsg.util.generator import consumer
16
15
 
17
- from .base import GenAxisArray
16
+ from .base import (
17
+ BaseStatefulTransformer,
18
+ BaseTransformerUnit,
19
+ processor_state,
20
+ )
18
21
  from .util.sparse import sliding_win_oneaxis as sparse_sliding_win_oneaxis
22
+ from .util.profile import profile_subpub
19
23
 
20
24
 
21
25
  class Anchor(enum.Enum):
@@ -24,142 +28,174 @@ class Anchor(enum.Enum):
24
28
  MIDDLE = "middle"
25
29
 
26
30
 
27
- @consumer
28
- def windowing(
29
- axis: str | None = None,
30
- newaxis: str = "win",
31
- window_dur: float | None = None,
32
- window_shift: float | None = None,
33
- zero_pad_until: str = "input",
34
- anchor: str | Anchor = Anchor.BEGINNING,
35
- ) -> typing.Generator[AxisArray, AxisArray, None]:
36
- """
37
- Apply a sliding window along the specified axis to input streaming data.
38
- The `windowing` method is perhaps the most useful and versatile method in ezmsg.sigproc, but its parameterization
39
- can be difficult. Please read the argument descriptions carefully.
40
-
41
- Args:
42
- axis: The axis along which to segment windows.
43
- If None, defaults to the first dimension of the first seen AxisArray.
44
- Note: The windowed axis must be an AxisArray.LinearAxis, not an AxisArray.CoordinateAxis.
45
- newaxis: New axis on which windows are delimited, immediately
46
- preceding the target windowed axis. The data length along newaxis may be 0 if
47
- this most recent push did not provide enough data for a new window.
48
- If window_shift is None then the newaxis length will always be 1.
49
- window_dur: The duration of the window in seconds.
50
- If None, the function acts as a passthrough and all other parameters are ignored.
51
- window_shift: The shift of the window in seconds.
52
- If None (default), windowing operates in "1:1 mode", where each input yields exactly one most-recent window.
53
- zero_pad_until: Determines how the function initializes the buffer.
54
- Can be one of "input" (default), "full", "shift", or "none". If `window_shift` is None then this field is
55
- ignored and "input" is always used.
56
-
57
- - "input" (default) initializes the buffer with the input then prepends with zeros to the window size.
58
- The first input will always yield at least one output.
59
- - "shift" fills the buffer until `window_shift`.
60
- No outputs will be yielded until at least `window_shift` data has been seen.
61
- - "none" does not pad the buffer. No outputs will be yielded until at least `window_dur` data has been seen.
62
- anchor: Determines the entry in `axis` that gets assigned `0`, which references the
63
- value in `newaxis`. Can be of class :obj:`Anchor` or a string representation of an :obj:`Anchor`.
64
-
65
- Returns:
66
- A primed generator that accepts an :obj:`AxisArray` via `.send(axis_array)`
67
- and yields an :obj:`AxisArray` with the data payload containing a windowed version of the input data.
68
- """
69
- # Check arguments
70
- if newaxis is None:
71
- ez.logger.warning("`newaxis` must not be None. Setting to 'win'.")
72
- newaxis = "win"
73
- if window_shift is None and zero_pad_until != "input":
74
- ez.logger.warning(
75
- "`zero_pad_until` must be 'input' if `window_shift` is None. "
76
- f"Ignoring received argument value: {zero_pad_until}"
77
- )
78
- zero_pad_until = "input"
79
- elif window_shift is not None and zero_pad_until == "input":
80
- ez.logger.warning(
81
- "windowing is non-deterministic with `zero_pad_until='input'` as it depends on the size "
82
- "of the first input. We recommend using 'shift' when `window_shift` is float-valued."
83
- )
84
- try:
85
- anchor = Anchor(anchor)
86
- except ValueError:
87
- raise ValueError(f"Invalid anchor: {anchor}. Valid anchor are: {', '.join([e.value for e in Anchor])}")
31
+ class WindowSettings(ez.Settings):
32
+ axis: str | None = None
33
+ newaxis: str | None = None # new axis for output. No new axes if None
34
+ window_dur: float | None = None # Sec. passthrough if None
35
+ window_shift: float | None = None # Sec. Use "1:1 mode" if None
36
+ zero_pad_until: str = "full" # "full", "shift", "input", "none"
37
+ anchor: str | Anchor = Anchor.BEGINNING
88
38
 
89
- msg_out = AxisArray(np.array([]), dims=[""])
90
39
 
91
- # State variables
40
+ @processor_state
41
+ class WindowState:
92
42
  buffer: npt.NDArray | sparse.SparseArray | None = None
43
+
93
44
  window_samples: int | None = None
45
+
94
46
  window_shift_samples: int | None = None
95
- # Number of incoming samples to ignore. Only relevant when shift > window.:
47
+
96
48
  shift_deficit: int = 0
97
- b_1to1 = window_shift is None
98
- newaxis_warned: bool = b_1to1
49
+ """ Number of incoming samples to ignore. Only relevant when shift > window."""
50
+
51
+ newaxis_warned: bool = False
52
+
99
53
  out_newaxis: AxisArray.LinearAxis | None = None
54
+
100
55
  out_dims: list[str] | None = None
101
56
 
102
- check_inputs = {"samp_shape": None, "fs": None, "key": None}
103
- concat_fun = np.concatenate
104
- sliding_win_fun = sliding_win_oneaxis
105
57
 
106
- while True:
107
- msg_in: AxisArray = yield msg_out
58
+ class WindowTransformer(
59
+ BaseStatefulTransformer[WindowSettings, AxisArray, AxisArray, WindowState]
60
+ ):
61
+ """
62
+ Apply a sliding window along the specified axis to input streaming data.
63
+ The `windowing` method is perhaps the most useful and versatile method in ezmsg.sigproc, but its parameterization
64
+ can be difficult. Please read the argument descriptions carefully.
65
+ """
108
66
 
109
- if window_dur is None:
110
- msg_out = msg_in
111
- continue
67
+ def __init__(self, *args, **kwargs) -> None:
68
+ """
69
+
70
+ Args:
71
+ axis: The axis along which to segment windows.
72
+ If None, defaults to the first dimension of the first seen AxisArray.
73
+ Note: The windowed axis must be an AxisArray.LinearAxis, not an AxisArray.CoordinateAxis.
74
+ newaxis: New axis on which windows are delimited, immediately
75
+ preceding the target windowed axis. The data length along newaxis may be 0 if
76
+ this most recent push did not provide enough data for a new window.
77
+ If window_shift is None then the newaxis length will always be 1.
78
+ window_dur: The duration of the window in seconds.
79
+ If None, the function acts as a passthrough and all other parameters are ignored.
80
+ window_shift: The shift of the window in seconds.
81
+ If None (default), windowing operates in "1:1 mode",
82
+ where each input yields exactly one most-recent window.
83
+ zero_pad_until: Determines how the function initializes the buffer.
84
+ Can be one of "input" (default), "full", "shift", or "none".
85
+ If `window_shift` is None then this field is ignored and "input" is always used.
86
+
87
+ - "input" (default) initializes the buffer with the input then prepends with zeros to the window size.
88
+ The first input will always yield at least one output.
89
+ - "shift" fills the buffer until `window_shift`.
90
+ No outputs will be yielded until at least `window_shift` data has been seen.
91
+ - "none" does not pad the buffer. No outputs will be yielded until
92
+ at least `window_dur` data has been seen.
93
+ anchor: Determines the entry in `axis` that gets assigned `0`, which references the
94
+ value in `newaxis`. Can be of class :obj:`Anchor` or a string representation of an :obj:`Anchor`.
95
+ """
96
+ super().__init__(*args, **kwargs)
97
+
98
+ # Sanity-check settings
99
+ # if self.settings.newaxis is None:
100
+ # ez.logger.warning("`newaxis=None` will be replaced with `newaxis='win'`.")
101
+ # object.__setattr__(self.settings, "newaxis", "win")
102
+ if (
103
+ self.settings.window_shift is None
104
+ and self.settings.zero_pad_until != "input"
105
+ ):
106
+ ez.logger.warning(
107
+ "`zero_pad_until` must be 'input' if `window_shift` is None. "
108
+ f"Ignoring received argument value: {self.settings.zero_pad_until}"
109
+ )
110
+ object.__setattr__(self.settings, "zero_pad_until", "input")
111
+ elif (
112
+ self.settings.window_shift is not None
113
+ and self.settings.zero_pad_until == "input"
114
+ ):
115
+ ez.logger.warning(
116
+ "windowing is non-deterministic with `zero_pad_until='input'` as it depends on the size "
117
+ "of the first input. We recommend using `zero_pad_until='shift'` when `window_shift` is float-valued."
118
+ )
119
+ try:
120
+ object.__setattr__(self.settings, "anchor", Anchor(self.settings.anchor))
121
+ except ValueError:
122
+ raise ValueError(
123
+ f"Invalid anchor: {self.settings.anchor}. Valid anchor are: {', '.join([e.value for e in Anchor])}"
124
+ )
112
125
 
113
- axis = axis or msg_in.dims[0]
114
- axis_idx = msg_in.get_axis_idx(axis)
115
- axis_info = msg_in.get_axis(axis)
126
+ def _hash_message(self, message: AxisArray) -> int:
127
+ axis = self.settings.axis or message.dims[0]
128
+ axis_idx = message.get_axis_idx(axis)
129
+ axis_info = message.get_axis(axis)
116
130
  fs = 1.0 / axis_info.gain
131
+ samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
132
+
133
+ return hash(samp_shape + (fs, message.key))
117
134
 
118
- if not newaxis_warned and newaxis in msg_in.dims:
135
+ def _reset_state(self, message: AxisArray) -> None:
136
+ _newaxis = self.settings.newaxis or "win"
137
+ if not self._state.newaxis_warned and _newaxis in message.dims:
119
138
  ez.logger.warning(
120
- f"newaxis {newaxis} present in input dims. Using {newaxis}_win instead"
139
+ f"newaxis {_newaxis} present in input dims. Using {_newaxis}_win instead"
121
140
  )
122
- newaxis_warned = True
123
- newaxis = f"{newaxis}_win"
124
-
125
- samp_shape = msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1 :]
126
-
127
- # If buffer unset or input stats changed, create a new buffer
128
- b_reset = buffer is None
129
- b_reset = b_reset or samp_shape != check_inputs["samp_shape"]
130
- b_reset = b_reset or fs != check_inputs["fs"]
131
- b_reset = b_reset or msg_in.key != check_inputs["key"]
132
- if b_reset:
133
- # Update check variables
134
- check_inputs["samp_shape"] = samp_shape
135
- check_inputs["fs"] = fs
136
- check_inputs["key"] = msg_in.key
137
-
138
- if isinstance(msg_in.data, sparse.SparseArray):
139
- concat_fun = sparse.concatenate
140
- sliding_win_fun = sparse_sliding_win_oneaxis
141
- zero_fun = sparse.zeros
142
- else:
143
- concat_fun = np.concatenate
144
- sliding_win_fun = sliding_win_oneaxis # Requires updated signature in ezmsg dev.
145
- zero_fun = np.zeros
146
-
147
- window_samples = int(window_dur * fs)
148
- if not b_1to1:
149
- window_shift_samples = int(window_shift * fs)
150
- if zero_pad_until == "none":
151
- req_samples = window_samples
152
- elif zero_pad_until == "shift" and not b_1to1:
153
- req_samples = window_shift_samples
154
- else: # i.e. zero_pad_until == "input"
155
- req_samples = msg_in.data.shape[axis_idx]
156
- n_zero = max(0, window_samples - req_samples)
157
- init_buffer_shape = (
158
- msg_in.data.shape[:axis_idx]
159
- + (n_zero,)
160
- + msg_in.data.shape[axis_idx + 1 :]
141
+ self._state.newaxis_warned = True
142
+ self.settings.newaxis = f"{_newaxis}_win"
143
+
144
+ axis = self.settings.axis or message.dims[0]
145
+ axis_idx = message.get_axis_idx(axis)
146
+ axis_info = message.get_axis(axis)
147
+ fs = 1.0 / axis_info.gain
148
+
149
+ xp = get_namespace(message.data)
150
+
151
+ self._state.window_samples = int(self.settings.window_dur * fs)
152
+ if self.settings.window_shift is not None:
153
+ # If window_shift is None, we are in "1:1 mode" and window_shift_samples is not used.
154
+ self._state.window_shift_samples = int(self.settings.window_shift * fs)
155
+ if self.settings.zero_pad_until == "none":
156
+ req_samples = self._state.window_samples
157
+ elif (
158
+ self.settings.zero_pad_until == "shift"
159
+ and self.settings.window_shift is not None
160
+ ):
161
+ req_samples = self._state.window_shift_samples
162
+ else: # i.e. zero_pad_until == "input"
163
+ req_samples = message.data.shape[axis_idx]
164
+ n_zero = max(0, self._state.window_samples - req_samples)
165
+ init_buffer_shape = (
166
+ message.data.shape[:axis_idx]
167
+ + (n_zero,)
168
+ + message.data.shape[axis_idx + 1 :]
169
+ )
170
+ self._state.buffer = xp.zeros(init_buffer_shape, dtype=message.data.dtype)
171
+
172
+ # Prepare reusable parts of output
173
+ if self._state.out_newaxis is None:
174
+ self._state.out_dims = (
175
+ list(message.dims[:axis_idx])
176
+ + [_newaxis]
177
+ + list(message.dims[axis_idx:])
178
+ )
179
+ self._state.out_newaxis = replace(
180
+ axis_info,
181
+ gain=0.0
182
+ if self.settings.window_shift is None
183
+ else axis_info.gain * self._state.window_shift_samples,
184
+ offset=0.0, # offset modified per-msg below
161
185
  )
162
- buffer = zero_fun(init_buffer_shape, dtype=msg_in.data.dtype)
186
+
187
+ def __call__(self, message: AxisArray) -> AxisArray:
188
+ if self.settings.window_dur is None:
189
+ # Shortcut for no windowing
190
+ return message
191
+ return super().__call__(message)
192
+
193
+ def _process(self, message: AxisArray) -> AxisArray:
194
+ axis = self.settings.axis or message.dims[0]
195
+ axis_idx = message.get_axis_idx(axis)
196
+ axis_info = message.get_axis(axis)
197
+
198
+ xp = get_namespace(message.data)
163
199
 
164
200
  # Add new data to buffer.
165
201
  # Currently, we concatenate the new time samples and clip the output.
@@ -168,155 +204,177 @@ def windowing(
168
204
  # is generally faster than np.roll and slicing anyway, but this could still
169
205
  # be a performance bottleneck for large memory arrays.
170
206
  # A circular buffer might be faster.
171
- buffer = concat_fun((buffer, msg_in.data), axis=axis_idx)
207
+ self._state.buffer = xp.concatenate(
208
+ (self._state.buffer, message.data), axis=axis_idx
209
+ )
172
210
 
173
211
  # Create a vector of buffer timestamps to track axis `offset` in output(s)
174
- buffer_tvec = np.arange(buffer.shape[axis_idx]).astype(float)
212
+ buffer_tvec = xp.asarray(range(self._state.buffer.shape[axis_idx]), dtype=float)
213
+
175
214
  # Adjust so first _new_ sample at index 0.
176
- buffer_tvec -= buffer_tvec[-msg_in.data.shape[axis_idx]]
215
+ buffer_tvec -= buffer_tvec[-message.data.shape[axis_idx]]
177
216
  # Convert form indices to 'units' (probably seconds).
178
217
  buffer_tvec *= axis_info.gain
179
218
  buffer_tvec += axis_info.offset
180
219
 
181
- if not b_1to1 and shift_deficit > 0:
182
- n_skip = min(buffer.shape[axis_idx], shift_deficit)
220
+ if self.settings.window_shift is not None and self._state.shift_deficit > 0:
221
+ n_skip = min(self._state.buffer.shape[axis_idx], self._state.shift_deficit)
183
222
  if n_skip > 0:
184
- buffer = slice_along_axis(buffer, slice(n_skip, None), axis_idx)
223
+ self._state.buffer = slice_along_axis(
224
+ self._state.buffer, slice(n_skip, None), axis_idx
225
+ )
185
226
  buffer_tvec = buffer_tvec[n_skip:]
186
- shift_deficit -= n_skip
187
-
188
- # Prepare reusable parts of output
189
- if out_newaxis is None:
190
- out_dims = msg_in.dims[:axis_idx] + [newaxis] + msg_in.dims[axis_idx:]
191
- out_newaxis = replace(
192
- axis_info,
193
- gain=0.0 if b_1to1 else axis_info.gain * window_shift_samples,
194
- offset=0.0, # offset modified per-msg below
195
- )
227
+ self._state.shift_deficit -= n_skip
196
228
 
197
229
  # Generate outputs.
198
230
  # Preliminary copy of axes without the axes that we are modifying.
199
- out_axes = {k: v for k, v in msg_in.axes.items() if k not in [newaxis, axis]}
231
+ _newaxis = self.settings.newaxis or "win"
232
+ out_axes = {k: v for k, v in message.axes.items() if k not in [_newaxis, axis]}
200
233
 
201
234
  # Update targeted (windowed) axis so that its offset is relative to the new axis
202
- if anchor == Anchor.BEGINNING:
235
+ if self.settings.anchor == Anchor.BEGINNING:
203
236
  out_axes[axis] = replace(axis_info, offset=0.0)
204
- elif anchor == Anchor.END:
205
- out_axes[axis] = replace(axis_info, offset=-window_dur)
206
- elif anchor == Anchor.MIDDLE:
207
- out_axes[axis] = replace(axis_info, offset=-window_dur / 2)
237
+ elif self.settings.anchor == Anchor.END:
238
+ out_axes[axis] = replace(axis_info, offset=-self.settings.window_dur)
239
+ elif self.settings.anchor == Anchor.MIDDLE:
240
+ out_axes[axis] = replace(axis_info, offset=-self.settings.window_dur / 2)
208
241
 
209
242
  # How we update .data and .axes[newaxis] depends on the windowing mode.
210
- if b_1to1:
243
+ if self.settings.window_shift is None:
211
244
  # one-to-one mode -- Each send yields exactly one window containing only the most recent samples.
212
- buffer = slice_along_axis(buffer, slice(-window_samples, None), axis_idx)
213
- out_dat = buffer.reshape(
214
- buffer.shape[:axis_idx] + (1,) + buffer.shape[axis_idx:]
245
+ self._state.buffer = slice_along_axis(
246
+ self._state.buffer, slice(-self._state.window_samples, None), axis_idx
215
247
  )
216
- win_offset = buffer_tvec[-window_samples]
217
- elif buffer.shape[axis_idx] >= window_samples:
248
+ out_dat = self._state.buffer.reshape(
249
+ self._state.buffer.shape[:axis_idx]
250
+ + (1,)
251
+ + self._state.buffer.shape[axis_idx:]
252
+ )
253
+ win_offset = buffer_tvec[-self._state.window_samples]
254
+ elif self._state.buffer.shape[axis_idx] >= self._state.window_samples:
218
255
  # Deterministic window shifts.
219
- out_dat = sliding_win_fun(buffer, window_samples, axis_idx, step=window_shift_samples)
220
- offset_view = sliding_win_oneaxis(buffer_tvec, window_samples, 0)[
221
- ::window_shift_samples
256
+ sliding_win_fun = (
257
+ sparse_sliding_win_oneaxis
258
+ if is_pydata_sparse_namespace(xp)
259
+ else sliding_win_oneaxis
260
+ )
261
+ out_dat = sliding_win_fun(
262
+ self._state.buffer,
263
+ self._state.window_samples,
264
+ axis_idx,
265
+ step=self._state.window_shift_samples,
266
+ )
267
+ offset_view = sliding_win_fun(buffer_tvec, self._state.window_samples, 0)[
268
+ :: self._state.window_shift_samples
222
269
  ]
223
270
  win_offset = offset_view[0, 0]
224
271
 
225
272
  # Drop expired beginning of buffer and update shift_deficit
226
- multi_shift = window_shift_samples * out_dat.shape[axis_idx]
227
- shift_deficit = max(0, multi_shift - buffer.shape[axis_idx])
228
- buffer = slice_along_axis(buffer, slice(multi_shift, None), axis_idx)
273
+ multi_shift = self._state.window_shift_samples * out_dat.shape[axis_idx]
274
+ self._state.shift_deficit = max(
275
+ 0, multi_shift - self._state.buffer.shape[axis_idx]
276
+ )
277
+ self._state.buffer = slice_along_axis(
278
+ self._state.buffer, slice(multi_shift, None), axis_idx
279
+ )
229
280
  else:
230
281
  # Not enough data to make a new window. Return empty data.
231
282
  empty_data_shape = (
232
- msg_in.data.shape[:axis_idx]
233
- + (0, window_samples)
234
- + msg_in.data.shape[axis_idx + 1 :]
283
+ message.data.shape[:axis_idx]
284
+ + (0, self._state.window_samples)
285
+ + message.data.shape[axis_idx + 1 :]
235
286
  )
236
- out_dat = np.zeros(empty_data_shape, dtype=msg_in.data.dtype)
287
+ out_dat = xp.zeros(empty_data_shape, dtype=message.data.dtype)
237
288
  # out_newaxis will have first timestamp in input... but mostly meaningless because output is size-zero.
238
289
  win_offset = axis_info.offset
239
290
 
240
- if anchor == Anchor.END:
241
- win_offset += window_dur
242
- elif anchor == Anchor.MIDDLE:
243
- win_offset += window_dur / 2
244
- out_newaxis = replace(out_newaxis, offset=win_offset)
291
+ if self.settings.anchor == Anchor.END:
292
+ win_offset += self.settings.window_dur
293
+ elif self.settings.anchor == Anchor.MIDDLE:
294
+ win_offset += self.settings.window_dur / 2
295
+ self._state.out_newaxis = replace(self._state.out_newaxis, offset=win_offset)
245
296
 
246
297
  msg_out = replace(
247
- msg_in, data=out_dat, dims=out_dims, axes={**out_axes, newaxis: out_newaxis}
298
+ message,
299
+ data=out_dat,
300
+ dims=self._state.out_dims,
301
+ axes={**out_axes, _newaxis: self._state.out_newaxis},
248
302
  )
303
+ return msg_out
249
304
 
250
305
 
251
- class WindowSettings(ez.Settings):
252
- axis: str | None = None
253
- newaxis: str | None = None # new axis for output. No new axes if None
254
- window_dur: float | None = None # Sec. passthrough if None
255
- window_shift: float | None = None # Sec. Use "1:1 mode" if None
256
- zero_pad_until: str = "full" # "full", "shift", "input", "none"
257
- anchor: str | Anchor = Anchor.BEGINNING
258
-
259
-
260
- class WindowState(ez.State):
261
- cur_settings: WindowSettings
262
- gen: typing.Generator
263
-
264
-
265
- class Window(GenAxisArray):
266
- """:obj:`Unit` for :obj:`bandpower`."""
267
-
306
+ class Window(
307
+ BaseTransformerUnit[WindowSettings, AxisArray, AxisArray, WindowTransformer]
308
+ ):
268
309
  SETTINGS = WindowSettings
269
-
270
310
  INPUT_SIGNAL = ez.InputStream(AxisArray)
271
311
  OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
272
312
 
273
- def construct_generator(self):
274
- self.STATE.gen = windowing(
275
- axis=self.SETTINGS.axis,
276
- newaxis=self.SETTINGS.newaxis,
277
- window_dur=self.SETTINGS.window_dur,
278
- window_shift=self.SETTINGS.window_shift,
279
- zero_pad_until=self.SETTINGS.zero_pad_until,
280
- anchor=self.SETTINGS.anchor,
281
- )
282
-
283
313
  @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
284
314
  @ez.publisher(OUTPUT_SIGNAL)
285
- async def on_signal(self, msg: AxisArray) -> typing.AsyncGenerator:
315
+ @profile_subpub(trace_oldest=False)
316
+ async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
317
+ """
318
+ override superclass on_signal so we can opt to yield once or multiple times after dropping the win axis.
319
+ """
320
+ # TODO: The transfomer overwrites settings.newaxis from None to "win",
321
+ # then we no longer know if the user wants to trim out the newaxis from the unit.
322
+ xp = get_namespace(message.data)
286
323
  try:
287
- out_msg = self.STATE.gen.send(msg)
288
- if out_msg.data.size > 0:
324
+ ret = self.processor(message)
325
+ if ret.data.size > 0:
289
326
  if (
290
327
  self.SETTINGS.newaxis is not None
291
328
  or self.SETTINGS.window_dur is None
292
329
  ):
293
330
  # Multi-win mode or pass-through mode.
294
- yield self.OUTPUT_SIGNAL, out_msg
331
+ yield self.OUTPUT_SIGNAL, ret
295
332
  else:
296
333
  # We need to split out_msg into multiple yields, dropping newaxis.
297
- axis_idx = out_msg.get_axis_idx("win")
298
- win_axis = out_msg.axes["win"]
299
- offsets = win_axis.value(np.arange(out_msg.data.shape[axis_idx]))
300
- for msg_ix in range(out_msg.data.shape[axis_idx]):
334
+ axis_idx = ret.get_axis_idx("win")
335
+ win_axis = ret.axes["win"]
336
+ offsets = win_axis.value(
337
+ xp.asarray(range(ret.data.shape[axis_idx]))
338
+ )
339
+ for msg_ix in range(ret.data.shape[axis_idx]):
301
340
  # Need to drop 'win' and replace self.SETTINGS.axis from axes.
302
341
  _out_axes = {
303
342
  **{
304
343
  k: v
305
- for k, v in out_msg.axes.items()
344
+ for k, v in ret.axes.items()
306
345
  if k not in ["win", self.SETTINGS.axis]
307
346
  },
308
347
  self.SETTINGS.axis: replace(
309
- out_msg.axes[self.SETTINGS.axis], offset=offsets[msg_ix]
348
+ ret.axes[self.SETTINGS.axis], offset=offsets[msg_ix]
310
349
  ),
311
350
  }
312
- _out_msg = replace(
313
- out_msg,
314
- data=slice_along_axis(out_msg.data, msg_ix, axis_idx),
315
- dims=out_msg.dims[:axis_idx] + out_msg.dims[axis_idx + 1 :],
351
+ _ret = replace(
352
+ ret,
353
+ data=slice_along_axis(ret.data, msg_ix, axis_idx),
354
+ dims=ret.dims[:axis_idx] + ret.dims[axis_idx + 1 :],
316
355
  axes=_out_axes,
317
356
  )
318
- yield self.OUTPUT_SIGNAL, _out_msg
319
- except (StopIteration, GeneratorExit):
320
- ez.logger.debug(f"Window closed in {self.address}")
357
+ yield self.OUTPUT_SIGNAL, _ret
358
+
321
359
  except Exception:
322
360
  ez.logger.info(traceback.format_exc())
361
+
362
+
363
+ def windowing(
364
+ axis: str | None = None,
365
+ newaxis: str | None = None,
366
+ window_dur: float | None = None,
367
+ window_shift: float | None = None,
368
+ zero_pad_until: str = "full",
369
+ anchor: str | Anchor = Anchor.BEGINNING,
370
+ ) -> WindowTransformer:
371
+ return WindowTransformer(
372
+ WindowSettings(
373
+ axis=axis,
374
+ newaxis=newaxis,
375
+ window_dur=window_dur,
376
+ window_shift=window_shift,
377
+ zero_pad_until=zero_pad_until,
378
+ anchor=anchor,
379
+ )
380
+ )
@@ -1,11 +1,12 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ezmsg-sigproc
3
- Version: 1.8.1
3
+ Version: 2.0.0
4
4
  Summary: Timeseries signal processing implementations in ezmsg
5
5
  Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>
6
6
  License-Expression: MIT
7
7
  License-File: LICENSE.txt
8
8
  Requires-Python: >=3.10.15
9
+ Requires-Dist: array-api-compat>=1.11.1
9
10
  Requires-Dist: ezmsg>=3.6.0
10
11
  Requires-Dist: numba>=0.61.0
11
12
  Requires-Dist: numpy>=1.26.0
@@ -0,0 +1,51 @@
1
+ ezmsg/sigproc/__init__.py,sha256=8K4IcOA3-pfzadoM6s2Sfg5460KlJUocGgyTJTJl96U,52
2
+ ezmsg/sigproc/__version__.py,sha256=2thmcF9DS_Zp1zHI3N0kjBeMAuCP9mdDGnL0clqQpS8,511
3
+ ezmsg/sigproc/activation.py,sha256=qWAhpbFBxSoqbGy4P9JKE5LY-5v8rQI1U81OvNxBG2Y,2820
4
+ ezmsg/sigproc/adaptive_lattice_notch.py,sha256=3M65PrZpdgBlQtE7Ph4Gu2ISIyWw4j8Xxhm5PpSkLFw,9102
5
+ ezmsg/sigproc/affinetransform.py,sha256=WU495KoDKZfHPS3Dumh65rgf639koNlfDIx_torIByg,8662
6
+ ezmsg/sigproc/aggregate.py,sha256=sdVzSXDg9BUNT-ljyvrWLeoZtRTlfisP0OxEchbgyMM,6111
7
+ ezmsg/sigproc/bandpower.py,sha256=N9pDz1X6ZTNP6VpCcfoXhj32j_9KpMaMuVYimoS6Jpc,2083
8
+ ezmsg/sigproc/base.py,sha256=PQr03O2P1v9LzcSR0GJLvPpBCLtnmGaz76gUeXphcH4,48753
9
+ ezmsg/sigproc/butterworthfilter.py,sha256=7ZP4CRsXBt3-5dzyUjD45vc0J3Fhpm4CLrk-ps28jhc,5305
10
+ ezmsg/sigproc/cheby.py,sha256=-aSauAwxJmmSSiRaw5qGY9rvYFOmk1bZlS4gGrS0jls,3737
11
+ ezmsg/sigproc/combfilter.py,sha256=5UCfzGESpS5LSx6rxZv8_n25ZUvOOmws-mM_gpTZNhU,4777
12
+ ezmsg/sigproc/decimate.py,sha256=Lz46fBllWagu17QeQzgklm6GWCV-zPysiydiby2IElU,2347
13
+ ezmsg/sigproc/detrend.py,sha256=7bpjFKdk2b6FdVn2GEtMbWtCuk7ToeiYKEBHVbN4Gd0,903
14
+ ezmsg/sigproc/diff.py,sha256=P5BBjR7KdaCL9aD3GG09cmC7a-3cxDeEUw4nKdQ1HY8,2895
15
+ ezmsg/sigproc/downsample.py,sha256=0X6EwPZ_XTwA2-nx5w-2HmMZUEDFuGAYF5EmPSuuVj8,3721
16
+ ezmsg/sigproc/ewma.py,sha256=W_VS2MxiO1J7z2XS6rtnLnCEXxdRPQbMKtZduBwqTEQ,6369
17
+ ezmsg/sigproc/ewmfilter.py,sha256=EPlocRdKORj575VV1YUzcNsVcq-pYgdEJ7_m9WfpVnY,4795
18
+ ezmsg/sigproc/extract_axis.py,sha256=Gl8Hl_Ho2pPzchPjfseVHVRAqxj6eOvUQZlzfYRA7eI,1603
19
+ ezmsg/sigproc/filter.py,sha256=i5adfND0NATrk2RewkWQ0C3RKRGiElr5AIB2eZE4Dr8,11225
20
+ ezmsg/sigproc/filterbank.py,sha256=pJzv_G6chgWa1ARmRjMAMgt9eEGnA-ZbMSge4EWrcYY,13633
21
+ ezmsg/sigproc/messages.py,sha256=y_twVPK7TxRj8ajmuSuBuxwvLTgyv9OF7Y7v9bw1tfs,926
22
+ ezmsg/sigproc/quantize.py,sha256=VzaqE6PatibEjkk7XrGO-ubAXYurAed9FYOn4bcQZQk,2193
23
+ ezmsg/sigproc/resample.py,sha256=XQzEbUq44qTx5tXX2QXd14hkMb7C3LXT3CqbC161X1M,11600
24
+ ezmsg/sigproc/sampler.py,sha256=qrw-7US3mqrGS7lOio7P_za0MSPgBhSxinIrMf1P3Os,11026
25
+ ezmsg/sigproc/scaler.py,sha256=fCLHvCNUSgv0XChf8iS9s5uHCSCVjCasM2TCvyG5BwQ,4111
26
+ ezmsg/sigproc/signalinjector.py,sha256=hGC837JyDLtAGrfsdMwzEoOqWXiwP7r7sGlUC9nahTY,2948
27
+ ezmsg/sigproc/slicer.py,sha256=QKiq8wOTXf3kwWSCiZEGn9rA9HaM_q6PbXvtfpgjTXw,5417
28
+ ezmsg/sigproc/spectral.py,sha256=_2qO6as4Nesmc9V1WW2EXNMH5pPz8aVTEcIPOi4-g2o,322
29
+ ezmsg/sigproc/spectrogram.py,sha256=TXbn8oXhBHCV-Ut8HRRUXJBYB-LbUxzTHpSgR0UnqtM,2941
30
+ ezmsg/sigproc/spectrum.py,sha256=xTSP8QFCG9M3NHveFkcks_wI-RzD7kM_fR1dmaLtiEQ,9737
31
+ ezmsg/sigproc/synth.py,sha256=DdE9yEXGrDRb745cOgKNpY2frI5uM2VHmCsaZO-UkBk,24547
32
+ ezmsg/sigproc/transpose.py,sha256=AIwz8X2AS7Tf1_aidND132uDuB04M4f3-0iRYq0ysC8,4530
33
+ ezmsg/sigproc/wavelets.py,sha256=g9nYRF4oVov2uLC0tfzPOLjaQah_HhM0ckhQ4m23mms,7507
34
+ ezmsg/sigproc/window.py,sha256=VAFqZsHu-J3hfBnbbUk9d7MIsbbPIgrqg3OSz7uhl_o,16242
35
+ ezmsg/sigproc/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
36
+ ezmsg/sigproc/math/abs.py,sha256=9E0A-p_Qa1SVzqbr1sesjgpu6-XGUZkcRJVK5jcNc0U,685
37
+ ezmsg/sigproc/math/clip.py,sha256=bBGfy45CKsUIhrznYZdpgUa0Lz7GsOuFl87IeL0qNA8,1057
38
+ ezmsg/sigproc/math/difference.py,sha256=Mr7lA_JZ8ix-FK_4s7zfA7IAJgj7SkWonDNZ5k2DZg4,2321
39
+ ezmsg/sigproc/math/invert.py,sha256=OTiT5gXs088-tvZe8piqxLoInYf1SMY4Ph62cu6DCGM,609
40
+ ezmsg/sigproc/math/log.py,sha256=bx0om3Qi3ZShExEZ-IH5Xrg3XFjNEmjVygWlXWjyrv8,1479
41
+ ezmsg/sigproc/math/scale.py,sha256=kMQRPYnm1o_9lC1EtIkoZOWaAWOWWbeT4ri1q7Hs7Fc,898
42
+ ezmsg/sigproc/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
43
+ ezmsg/sigproc/util/asio.py,sha256=PQew73hB1oRmp7pfTqx-c4uo1zqgjxvZcTZCROQrEP4,5270
44
+ ezmsg/sigproc/util/message.py,sha256=l_b1b6bXX8N6VF9RbUELzsHs73cKkDURBdIr0lt3CY0,909
45
+ ezmsg/sigproc/util/profile.py,sha256=KNJ_QkKelQHNEp2C8MhqzdhYydMNULc_NQq3ccMfzIk,5775
46
+ ezmsg/sigproc/util/sparse.py,sha256=8Ke0jh3jRPi_TwIdLTwLdojQiaqPs6QV-Edqpx81VoI,1036
47
+ ezmsg/sigproc/util/typeresolution.py,sha256=5R7xmG-F4CkdqQ5aoQnqM-htQb-VwAJl58jJgxtClys,3146
48
+ ezmsg_sigproc-2.0.0.dist-info/METADATA,sha256=6QpdeS-5kMO5DogCy7LpB6ho0sQN5rnpfRwF7TwgXg4,2479
49
+ ezmsg_sigproc-2.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
50
+ ezmsg_sigproc-2.0.0.dist-info/licenses/LICENSE.txt,sha256=seu0tKhhAMPCUgc1XpXGGaCxY1YaYvFJwqFuQZAl2go,1100
51
+ ezmsg_sigproc-2.0.0.dist-info/RECORD,,