← Walkthroughmapa/preprocessing/frontend.py
Source snapshot · source SHA-256 1bfc67715492
1"""The released checkpoint's offline voltage-to-bands recipe.23Main path: prepare_recording() -> PreparedRecording.window().45Input: a full voltage recording plus contact labels/regions.6Output: cached STFT magnitudes and statistics; window() selects normalized clips.7Next: models/stem.py turns each clip's frequency bins into learned tokens.89Process a whole session, then slice windows on its 32 Hz clock. Raw arrays are10channels × samples. Supply cleaned input and any explicit contact exclusions. No task labels are used.11"""1213from __future__ import annotations1415import json16import math17from dataclasses import dataclass18from importlib.metadata import version19from pathlib import Path2021import numpy as np22import torch23from scipy.signal import resample_poly2425from mapa.data.anatomy import V14_DKT_REGION_LABELS26from mapa.data.normalize import SessionRobustZNormalizer27from mapa.models.sidecar import SensorSidecar, build_sidecar28from mapa.models.stem import INPUT_CLIP_Z29from mapa.preprocessing.stft import _single_stft_raw_view_chunked3031SAMPLE_RATE = 204832FRAME_RATE = 3233# name, FFT length, first bin, last bin (inclusive). Hz = bin * 2048 / FFT length.34# Bin centers (inclusive): slow 2-14 Hz by 2; mid 16-56 by 8; fast 64-160 by 16.35# Window lengths: 500/125/62.5 ms; all three use hop=64 samples (31.25 ms).36BANDS = (("slow", 1024, 1, 7), ("mid", 256, 2, 7), ("fast", 128, 4, 10))373839def region_ids_from_names(names):40 """Map exact hemisphere-qualified DKT names; None explicitly selects ID 74.4142 Unknown strings raise, so misspellings cannot silently become unassigned.43 Anatomical localization/atlas assignment must be performed upstream.44 """45 lookup = {name: i for i, name in enumerate(V14_DKT_REGION_LABELS)}46 return torch.tensor([74 if name is None else lookup[name] for name in names], dtype=torch.long)474849def _filter(data, labels):50 import mne5152 raw = mne.io.RawArray(53 np.asarray(data, dtype=np.float64),54 mne.create_info(list(labels), SAMPLE_RATE, "seeg"),55 verbose=False,56 )57 freqs = np.arange(60.0, min(SAMPLE_RATE / 2, 301), 60.0)58 raw.notch_filter(freqs, phase="zero", filter_length="auto", verbose=False)59 raw.filter(0.5, None, verbose=False)60 return raw.get_data()616263def _reference(data, sidecar):64 data = np.array(data, copy=True)65 for array in range(sidecar.n_arrays):66 ids = np.flatnonzero(sidecar.array_id.numpy() == array)67 data[ids] -= data[ids].mean(axis=0, keepdims=True)68 return data697071def preprocessing_config():72 """Return a fresh, JSON-serializable description of the released frontend."""73 return {74 "recipe": "mapa-v0.1.0",75 "sample_rate_hz": SAMPLE_RATE,76 "frame_rate_hz": FRAME_RATE,77 "resample": "scipy.signal.resample_poly; default Kaiser window; input/output float32",78 "model_reference": "mean of surviving selected contacts per shaft, before model filters",79 "model_notches_hz": [60, 120, 180, 240, 300],80 "highpass_hz": 0.5,81 "filter": "MNE default FIR; zero phase; reflect_limited padding",82 "contact_exclusions": "supplied by caller",83 "artifact_cleaning": "external contact detection and window rejection",84 "model_input_clip_z": list(INPUT_CLIP_Z),85 "bands": [86 {"name": n, "n_fft": fft, "hop": 64, "bins_inclusive": [lo, hi]}87 for n, fft, lo, hi in BANDS88 ],89 "stft": {90 "window": "periodic Hann",91 "center": True,92 "padding": "reflect",93 "normalized": False,94 "representation": "magnitude; no log",95 },96 "normalization": {97 "fit": "entire session, separately per contact and frequency bin",98 "median": "torch lower median",99 "mad_scale": 1.4826,100 "sigma_floor": 1e-6,101 "constant_bins": "zero",102 },103 "window_start": "round(seconds * 32), ties to even",104 "voltage_scaling": "none",105 }106107108@dataclass109class PreparedRecording:110 """Raw STFT caches plus frozen statistics; normalized clips are created on demand."""111112 bands: tuple[torch.Tensor, ...]113 normalizers: tuple[SessionRobustZNormalizer, ...]114 sidecar: SensorSidecar115 input_indices: np.ndarray116 duration_s: float117 metadata: dict118119 def window(self, start_s, duration_s=1.0, *, device="cpu"):120 """Return three (1, contacts, bins, frames) normalized tensors.121122 Reject windows outside the recording. No artifact detection or rejection123 is performed. Window selection is the caller's responsibility.124 """125 if not math.isfinite(start_s) or start_s < 0:126 raise ValueError("start_s must be finite and nonnegative")127 if not math.isfinite(duration_s) or duration_s <= 0:128 raise ValueError("duration_s must be finite and positive")129 frames = round(duration_s * FRAME_RATE)130 if frames < 8 or frames % 8 or not math.isclose(frames / FRAME_RATE, duration_s):131 raise ValueError("duration must be a positive multiple of 0.25 seconds")132 start = round(start_s * FRAME_RATE)133 snapped = start / FRAME_RATE134 if start_s + duration_s > self.duration_s or snapped + duration_s > self.duration_s:135 raise ValueError("window extends beyond the recording")136 # Extract a clip from the cached STFT frames, not another FFT window.137 # Normalize with the full-session median/sigma fitted in prepare_recording():138 # z = (clip - median) / sigma; bins with sigma < 1e-6 become zero. No refit.139 # (C,F_b,session_frames) -> (C,F_b,frames) -> (1,C,F_b,frames).140 clip_bands = []141 for band, norm in zip(self.bands, self.normalizers):142 clip = band[..., start : start + frames] # select cached STFT frames143 normalized = norm.transform(clip) # use stored session statistics144 batched = normalized.unsqueeze(0) # add the one-clip batch axis145 clip_bands.append(batched.to(device))146 return clip_bands147148 def save(self, output_dir, *, subject_id, trial_id):149 """Write raw band caches/statistics compatible with the evaluation loader.150151 Use a new root per recording. Existing directories are never overwritten.152 No raw voltage is written. Configuration/selection/statistics accompany the caches.153 """154 root = Path(output_dir)155 root.mkdir(parents=True, exist_ok=False)156 key = json.dumps({"subject_id": int(subject_id), "trial_id": int(trial_id)})157 for (name, fft, lo, hi), band, norm in zip(BANDS, self.bands, self.normalizers):158 leaf = root / name159 leaf.mkdir()160 np.save(leaf / "session.npy", band.numpy())161 np.savez(162 leaf / "session.stats.npz", median=norm.median.numpy(), sigma=norm.sigma.numpy()163 )164 meta = {165 "key": key,166 "ch_names": list(self.sidecar.labels),167 "total_frames": band.shape[-1],168 "sample_rate": SAMPLE_RATE,169 "band_hop": 64,170 "band_nperseg": fft,171 "f_bins": hi - lo + 1,172 "n_channels": band.shape[0],173 "dtype": "float32",174 }175 (leaf / "session.json").write_text(json.dumps(meta, indent=2) + "\n")176 (root / "spans").mkdir()177 (root / "spans/session.json").write_text(178 json.dumps(179 {180 "subject_id": int(subject_id),181 "trial_id": int(trial_id),182 "bad_windows_s": [],183 },184 indent=2,185 )186 + "\n"187 )188 (root / "preprocessing.json").write_text(json.dumps(self.metadata, indent=2) + "\n")189190191def prepare_recording(192 voltage,193 sample_rate,194 labels,195 region_ids,196 *,197 bad_labels=(),198 keep_labels=None,199 chunk_frames=960,200):201 """Prepare a full unreferenced session using the released offline recipe.202203 Preserve acquisition units (no implicit conversion to MNE volts). The original204 Brain Treebank voltage scale is retained by the research loader; do not assume205 it was converted to volts. Fixed numerical floors make arbitrary rescaling206 non-equivalent for near-constant channels. See the preprocessing guide.207208 bad_labels supplies contact exclusions determined upstream.209 This function does not detect bad contacts. keep_labels restricts the output210 montage before referencing.211 This is important for reproducing Neuroprobe Lite's contact budget.212 """213 if not isinstance(chunk_frames, int) or chunk_frames <= 0:214 raise ValueError("chunk_frames must be a positive integer")215 if not math.isfinite(sample_rate) or sample_rate < 2 * 160 or int(sample_rate) != sample_rate:216 raise ValueError("sample_rate must be an integer >= 320 Hz; paper rates were 1024/2048 Hz")217 x = np.asarray(voltage, dtype=np.float32)218 labels = tuple(labels)219 ids = np.asarray(region_ids)220 if x.ndim != 2 or x.shape[0] != len(labels) or ids.shape != (len(labels),):221 raise ValueError("voltage, labels and region_ids must have matching contact axes")222 if not len(labels) or len(set(labels)) != len(labels):223 raise ValueError("provide unique contact labels")224 if not np.issubdtype(ids.dtype, np.integer) or np.any((ids < 0) | (ids > 74)):225 raise ValueError("region_ids must be integer DKT indices 0..74")226 if x.shape[1] / sample_rate < 5:227 raise ValueError("supply a full session (at least 5 seconds), not an isolated trial")228 bad = set(bad_labels)229 keep = set(labels) if keep_labels is None else set(keep_labels)230 if (bad | keep) - set(labels):231 raise ValueError("bad_labels/keep_labels contain labels absent from the recording")232 # 1. Select contacts; keep the voltage rows and sidecar metadata aligned.233 original_indices = np.array([234 i for i, label in enumerate(labels) if label not in bad and label in keep235 ], dtype=np.int64)236 if not original_indices.size:237 raise ValueError("no contacts remain after supplied exclusions and montage selection")238 x = x[original_indices]239 if not np.isfinite(x).all():240 raise ValueError("selected contacts must have finite voltage")241 sidecar = build_sidecar([labels[i] for i in original_indices], region_id=ids[original_indices])242 if any(int((sidecar.array_id == i).sum()) < 2 for i in range(sidecar.n_arrays)):243 raise ValueError("a surviving shaft has fewer than two contacts for mean referencing")244 # 2. Resample to 2048 Hz when needed.245 factor = math.gcd(SAMPLE_RATE, int(sample_rate))246 if sample_rate != SAMPLE_RATE:247 x = resample_poly(x, SAMPLE_RATE // factor, int(sample_rate) // factor, axis=1).astype(248 np.float32249 )250 # 3. Subtract each shaft's mean, then apply harmonic notches and a high-pass filter.251 model_voltage = _filter(_reference(x.astype(np.float64), sidecar), sidecar.labels).astype(252 np.float32253 )254 waveform = torch.from_numpy(model_voltage)255 bands, normalizers = [], []256 # 4. Voltage (C,samples) -> magnitude STFT (C,F_b,session_frames), no log.257 # Slow: 2-14 Hz, 7 bins, 500 ms window; mid: 16-56 Hz, 6 bins, 125 ms;258 # fast: 64-160 Hz, 7 bins, 62.5 ms. Ranges are bin centers, inclusive.259 # All bands use a 31.25 ms hop (32 frames/s); the stem subsamples later.260 for _, fft, lo, hi in BANDS:261 band = _single_stft_raw_view_chunked(262 waveform,263 sample_rate=SAMPLE_RATE,264 nperseg=fft,265 hop_length=64,266 k0=lo,267 k1=hi,268 log_eps=1e-6,269 chunk_frames=chunk_frames,270 )271 # 5. Fit once over session time, separately for each contact/frequency bin:272 # median and sigma = 1.4826 * median(abs(band - median)), both (C,F_b,1).273 # Keep raw magnitudes + these statistics; window() applies normalization.274 norm = SessionRobustZNormalizer(sigma_floor=1e-6).fit(band)275 bands.append(band)276 normalizers.append(norm)277 metadata = {278 "recipe": preprocessing_config(),279 "input_sample_rate_hz": sample_rate,280 "input_units": "unchanged",281 "input_labels": list(labels),282 "output_labels": list(sidecar.labels),283 "input_indices": original_indices.tolist(),284 "region_ids": sidecar.region_id.tolist(),285 "known_exclusions": sorted(bad),286 "versions": {name: version(name) for name in ("numpy", "scipy", "torch", "mne")},287 }288 return PreparedRecording(289 tuple(bands),290 tuple(normalizers),291 sidecar,292 original_indices,293 x.shape[1] / SAMPLE_RATE,294 metadata,295 )