Aluode/PerceptionLabPortable
0
1# Authors: The MNE-Python contributors.2# License: BSD-3-Clause3# Copyright the MNE-Python contributors.4 5import numpy as np6from scipy.signal import get_window7 8from .utils import _ensure_int, _validate_type, logger, verbose9 10###############################################################################11# Class for interpolation between adjacent points12 13 14class _Interp2:15 r"""Interpolate between two points.16 17 Parameters18 ----------19 control_points : array, shape (n_changes,)20 The control points (indices) to use.21 values : callable | array, shape (n_changes, ...)22 Callable that takes the control point and returns a list of23 arrays that must be interpolated.24 interp : str25 Can be 'zero', 'linear', 'hann', or 'cos2' (same as hann).26 27 Notes28 -----29 This will process data using overlapping windows of potentially30 different sizes to achieve a constant output value using different31 2-point interpolation schemes. For example, for linear interpolation,32 and window sizes of 6 and 17, this would look like::33 34 1 _ _35 |\ / '-. .-'36 | \ / '-. .-'37 | x |-.-|38 | / \ .-' '-.39 |/ \_.-' '-.40 0 +----|----|----|----|---41 0 5 10 15 20 2542 43 """44 45 def __init__(self, control_points, values, interp="hann", *, name="Interp2"):46 # set up interpolation47 self.control_points = np.array(control_points, int).ravel()48 if not np.array_equal(np.unique(self.control_points), self.control_points):49 raise ValueError("Control points must be sorted and unique")50 if len(self.control_points) == 0:51 raise ValueError("Must be at least one control point")52 if not (self.control_points >= 0).all():53 raise ValueError(54 f"All control points must be positive (got {self.control_points[:3]})"55 )56 if isinstance(values, np.ndarray):57 values = [values]58 if isinstance(values, list | tuple):59 for v in values:60 if not (v is None or isinstance(v, np.ndarray)):61 raise TypeError(62 'All entries in "values" must be ndarray or None, got '63 f"{type(v)}"64 )65 if v is not None and v.shape[0] != len(self.control_points):66 raise ValueError(67 "Values, if provided, must be the same length as the number of "68 f"control points ({len(self.control_points)}), got {v.shape[0]}"69 )70 use_values = values71 72 def val(pt):73 idx = np.where(control_points == pt)[0][0]74 return [v[idx] if v is not None else None for v in use_values]75 76 values = val77 self.values = values78 self.n_last = None79 self._position = 0 # start at zero80 self._left_idx = 081 self._left = self._right = self._use_interp = None82 self.name = name83 known_types = ("cos2", "linear", "zero", "hann")84 if interp not in known_types:85 raise ValueError(f'interp must be one of {known_types}, got "{interp}"')86 self._interp = interp87 88 def feed_generator(self, n_pts):89 """Feed data and get interpolators as a generator."""90 self.n_last = 091 n_pts = _ensure_int(n_pts, "n_pts")92 original_position = self._position93 stop = self._position + n_pts94 logger.debug(f" ~ {self.name} Feed {n_pts} ({self._position}-{stop})")95 used = np.zeros(n_pts, bool)96 if self._left is None: # first one97 logger.debug(f" ~ {self.name} Eval @ 0 ({self.control_points[0]})")98 self._left = self.values(self.control_points[0])99 if len(self.control_points) == 1:100 self._right = self._left101 n_used = 0102 103 # Left zero-order hold condition104 if self._position < self.control_points[self._left_idx]:105 n_use = min(self.control_points[self._left_idx] - self._position, n_pts)106 logger.debug(f" ~ {self.name} Left ZOH {n_use}")107 this_sl = slice(None, n_use)108 assert used[this_sl].size == n_use109 assert not used[this_sl].any()110 used[this_sl] = True111 yield [this_sl, self._left, None, None]112 self._position += n_use113 n_used += n_use114 self.n_last += 1115 116 # Standard interpolation condition117 stop_right_idx = np.where(self.control_points >= stop)[0]118 if len(stop_right_idx) == 0:119 stop_right_idx = [len(self.control_points) - 1]120 stop_right_idx = stop_right_idx[0]121 left_idxs = np.arange(self._left_idx, stop_right_idx)122 self.n_last += max(len(left_idxs) - 1, 0)123 for bi, left_idx in enumerate(left_idxs):124 if left_idx != self._left_idx or self._right is None:125 if self._right is not None:126 assert left_idx == self._left_idx + 1127 self._left = self._right128 self._left_idx += 1129 self._use_interp = None # need to recreate it130 eval_pt = self.control_points[self._left_idx + 1]131 logger.debug(132 f" ~ {self.name} Eval @ {self._left_idx + 1} ({eval_pt})"133 )134 self._right = self.values(eval_pt)135 assert self._right is not None136 left_point = self.control_points[self._left_idx]137 right_point = self.control_points[self._left_idx + 1]138 if self._use_interp is None:139 interp_span = right_point - left_point140 if self._interp == "zero":141 self._use_interp = None142 elif self._interp == "linear":143 self._use_interp = np.linspace(144 1.0, 0.0, interp_span, endpoint=False145 )146 else: # self._interp in ('cos2', 'hann'):147 self._use_interp = np.cos(148 np.linspace(0, np.pi / 2.0, interp_span, endpoint=False)149 )150 self._use_interp *= self._use_interp151 n_use = min(stop, right_point) - self._position152 if n_use > 0:153 logger.debug(154 f" ~ {self.name} Interp {self._interp} {n_use} "155 f"({left_point}-{right_point})"156 )157 interp_start = self._position - left_point158 assert interp_start >= 0159 if self._use_interp is None:160 this_interp = None161 else:162 this_interp = self._use_interp[interp_start : interp_start + n_use]163 assert this_interp.size == n_use164 this_sl = slice(n_used, n_used + n_use)165 assert used[this_sl].size == n_use166 assert not used[this_sl].any()167 used[this_sl] = True168 yield [this_sl, self._left, self._right, this_interp]169 self._position += n_use170 n_used += n_use171 172 # Right zero-order hold condition173 if self.control_points[self._left_idx] <= self._position:174 n_use = stop - self._position175 if n_use > 0:176 logger.debug(f" ~ {self.name} Right ZOH %s" % n_use)177 this_sl = slice(n_pts - n_use, None)178 assert not used[this_sl].any()179 used[this_sl] = True180 assert self._right is not None181 yield [this_sl, self._right, None, None]182 self._position += n_use183 n_used += n_use184 self.n_last += 1185 assert self._position == stop186 assert n_used == n_pts187 assert used.all()188 assert self._position == original_position + n_pts189 190 def feed(self, n_pts):191 """Feed data and get interpolated values."""192 # Convenience function for assembly193 out_arrays = None194 for o in self.feed_generator(n_pts):195 if out_arrays is None:196 out_arrays = [197 np.empty(v.shape + (n_pts,)) if v is not None else None198 for v in o[1]199 ]200 for ai, arr in enumerate(out_arrays):201 if arr is not None:202 if o[3] is None:203 arr[..., o[0]] = o[1][ai][..., np.newaxis]204 else:205 arr[..., o[0]] = o[1][ai][..., np.newaxis] * o[3] + o[2][ai][206 ..., np.newaxis207 ] * (1.0 - o[3])208 assert out_arrays is not None209 return out_arrays210 211 212###############################################################################213# Constant overlap-add processing class214 215 216def _check_store(store):217 _validate_type(store, (np.ndarray, list, tuple, _Storer), "store")218 if isinstance(store, np.ndarray):219 store = [store]220 if not isinstance(store, _Storer):221 if not all(isinstance(s, np.ndarray) for s in store):222 raise TypeError("All instances must be ndarrays")223 store = _Storer(*store)224 return store225 226 227class _COLA:228 r"""Constant overlap-add processing helper.229 230 Parameters231 ----------232 process : callable233 A function that takes a chunk of input data with shape234 ``(n_channels, n_samples)`` and processes it.235 store : ndarray | list of ndarray | _Storer236 The output data in which to store the results.237 n_total : int238 The total number of samples.239 n_samples : int240 The number of samples per window.241 n_overlap : int242 The overlap between windows.243 window : str244 The window to use. Default is "hann".245 tol : float246 The tolerance for COLA checking.247 248 Notes249 -----250 This will process data using overlapping windows to achieve a constant251 output value. For example, for ``n_total=27``, ``n_samples=10``,252 ``n_overlap=5`` and ``window='triang'``::253 254 1 _____ _______255 | \ /\ /\ /256 | \ / \ / \ /257 | x x x258 | / \ / \ / \259 | / \/ \/ \260 0 +----|----|----|----|----|-261 0 5 10 15 20 25262 263 This produces four windows: the first three are the requested length264 (10 samples) and the last one is longer (12 samples). The first and last265 window are asymmetric.266 """267 268 @verbose269 def __init__(270 self,271 process,272 store,273 n_total,274 n_samples,275 n_overlap,276 sfreq,277 window="hann",278 tol=1e-10,279 *,280 name="COLA",281 verbose=None,282 ):283 n_samples = _ensure_int(n_samples, "n_samples")284 n_overlap = _ensure_int(n_overlap, "n_overlap")285 n_total = _ensure_int(n_total, "n_total")286 if n_samples <= 0:287 raise ValueError(f"n_samples must be > 0, got {n_samples}")288 if n_overlap < 0:289 raise ValueError(f"n_overlap must be >= 0, got {n_overlap}")290 if n_total < 0:291 raise ValueError(f"n_total must be >= 0, got {n_total}")292 self._n_samples = int(n_samples)293 self._n_overlap = int(n_overlap)294 del n_samples, n_overlap295 if n_total < self._n_samples:296 raise ValueError(297 f"Number of samples per window ({self._n_samples}) must be at "298 f"most the total number of samples ({n_total})"299 )300 if not callable(process):301 raise TypeError(f"process must be callable, got type {type(process)}")302 self._process = process303 self._step = self._n_samples - self._n_overlap304 self._store = _check_store(store)305 self._idx = 0306 self._in_buffers = self._out_buffers = None307 self.name = name308 309 # Create our window boundaries310 window_name = window if isinstance(window, str) else "custom"311 self._window = get_window(312 window, self._n_samples, fftbins=bool((self._n_samples - 1) % 2)313 )314 self._window /= _check_cola(315 self._window, self._n_samples, self._step, window_name, tol=tol316 )317 self.starts = np.arange(0, n_total - self._n_samples + 1, self._step)318 self.stops = self.starts + self._n_samples319 delta = n_total - self.stops[-1]320 self.stops[-1] = n_total321 sfreq = float(sfreq)322 pl = "s" if len(self.starts) != 1 else ""323 logger.info(324 f" Processing {len(self.starts):4d} data chunk{pl} of (at least) "325 f"{self._n_samples / sfreq:0.1f} s with "326 f"{self._n_overlap / sfreq:0.1f} s overlap and {window_name} windowing"327 )328 del window, window_name329 if delta > 0:330 logger.info(331 f" The final {delta / sfreq} s will be lumped into the final window"332 )333 334 @property335 def _in_offset(self):336 """Compute from current processing window start and buffer len."""337 return self.starts[self._idx] + self._in_buffers[0].shape[-1]338 339 @verbose340 def feed(self, *datas, verbose=None, **kwargs):341 """Pass in a chunk of data."""342 # Append to our input buffer343 if self._in_buffers is None:344 self._in_buffers = [None] * len(datas)345 if len(datas) != len(self._in_buffers):346 raise ValueError(347 f"Got {len(datas)} array(s), needed {len(self._in_buffers)}"348 )349 current_offset = 0 # should be updated below350 for di, data in enumerate(datas):351 if not isinstance(data, np.ndarray) or data.ndim < 1:352 raise TypeError(353 f"data entry {di} must be an 2D ndarray, got {type(data)}"354 )355 if self._in_buffers[di] is None:356 # In practice, users can give large chunks, so we use357 # dynamic allocation of the in buffer. We could save some358 # memory allocation by only ever processing max_len at once,359 # but this would increase code complexity.360 self._in_buffers[di] = np.empty(data.shape[:-1] + (0,), data.dtype)361 if (362 data.shape[:-1] != self._in_buffers[di].shape[:-1]363 or self._in_buffers[di].dtype != data.dtype364 ):365 raise TypeError(366 f"data must dtype {self._in_buffers[di].dtype} and "367 f"shape[:-1]=={self._in_buffers[di].shape[:-1]}, got dtype "368 f"{data.dtype} shape[:-1]={data.shape[:-1]}"369 )370 # This gets updated on first iteration, so store it before it updates371 if di == 0:372 current_offset = self._in_offset373 logger.debug(374 f" + {self.name}[{di}] Appending "375 f"{current_offset}:{current_offset + data.shape[-1]}"376 )377 self._in_buffers[di] = np.concatenate([self._in_buffers[di], data], -1)378 if self._in_offset > self.stops[-1]:379 raise ValueError(380 f"data (shape {data.shape}) exceeded expected total buffer size ("381 f"{self._in_offset} > {self.stops[-1]})"382 )383 # Check to see if we can process the next chunk and dump outputs384 while self._idx < len(self.starts) and self._in_offset >= self.stops[self._idx]:385 start, stop = self.starts[self._idx], self.stops[self._idx]386 this_len = stop - start387 this_window = self._window.copy()388 if self._idx == len(self.starts) - 1:389 this_window = np.pad(390 self._window, (0, this_len - len(this_window)), "constant"391 )392 for offset in range(self._step, len(this_window), self._step):393 n_use = len(this_window) - offset394 this_window[offset:] += self._window[:n_use]395 if self._idx == 0:396 for offset in range(self._n_samples - self._step, 0, -self._step):397 this_window[:offset] += self._window[-offset:]398 this_proc = [in_[..., :this_len].copy() for in_ in self._in_buffers]399 logger.debug(400 f" * {self.name}[:] Processing {start}:{stop} "401 f"(e.g., {this_proc[0].flat[[0, -1]]})"402 )403 if not all(404 proc.shape[-1] == this_len == this_window.size for proc in this_proc405 ):406 raise RuntimeError("internal indexing error")407 start = self._store.idx408 stop = self._store.idx + this_len409 outs = self._process(*this_proc, start=start, stop=stop, **kwargs)410 if self._out_buffers is None:411 max_len = np.max(self.stops - self.starts)412 self._out_buffers = [413 np.zeros(o.shape[:-1] + (max_len,), o.dtype) for o in outs414 ]415 for oi, out in enumerate(outs):416 out *= this_window417 self._out_buffers[oi][..., : stop - start] += out418 self._idx += 1419 if self._idx < len(self.starts):420 next_start = self.starts[self._idx]421 else:422 next_start = self.stops[-1]423 delta = next_start - self.starts[self._idx - 1]424 logger.debug(425 f" + {self.name}[:] Shifting input and output buffers by "426 f"{delta} samples (storing {start}:{stop})"427 )428 for di in range(len(self._in_buffers)):429 self._in_buffers[di] = self._in_buffers[di][..., delta:]430 self._store(*[o[..., :delta] for o in self._out_buffers])431 for ob in self._out_buffers:432 ob[..., :-delta] = ob[..., delta:]433 ob[..., -delta:] = 0.0434 435 436def _check_cola(win, nperseg, step, window_name, tol=1e-10):437 """Check whether the Constant OverLap Add (COLA) constraint is met."""438 # adapted from SciPy439 binsums = np.sum(440 [win[ii * step : (ii + 1) * step] for ii in range(nperseg // step)], axis=0441 )442 if nperseg % step != 0:443 binsums[: nperseg % step] += win[-(nperseg % step) :]444 const = np.median(binsums)445 deviation = np.max(np.abs(binsums - const))446 if deviation > tol:447 raise ValueError(448 f"segment length {nperseg} with step {step} for {window_name} "449 "window type does not provide a constant output "450 f"({100 * deviation / const:g}% deviation)"451 )452 return const453 454 455class _Storer:456 """Store data in chunks."""457 458 def __init__(self, *outs, picks=None):459 for oi, out in enumerate(outs):460 if not isinstance(out, np.ndarray) or out.ndim < 1:461 raise TypeError(f"outs[oi] must be >= 1D ndarray, got {out}")462 self.outs = outs463 self.idx = 0464 self.picks = picks465 466 def __call__(self, *outs):467 if len(outs) != len(self.outs) or not all(468 out.shape[-1] == outs[0].shape[-1] for out in outs469 ):470 raise ValueError("Bad outs")471 idx = (Ellipsis,)472 if self.picks is not None:473 idx += (self.picks,)474 stop = self.idx + outs[0].shape[-1]475 idx += (slice(self.idx, stop),)476 for o1, o2 in zip(self.outs, outs):477 o1[idx] = o2478 self.idx = stop479