#!/usr/bin/env python """Artifact handling — ICA blink removal, ICLabel auto-labelling, autoreject epoch repair. A flag-driven pipeline so the same code powers batch analysis, the interactive review UI (clean_ui.py), and the live de-blink in cap_gui. Every step is optional and reports what it did, so you can compare "raw vs cleaned" and see exactly what each module removed. Steps (in order): notch → band-pass → bad-channel detect+interpolate → CAR → ICA(remove eye/muscle/…) then, at the epoch stage: autoreject (repair/drop bad epochs) Constraints for THIS cap: no dedicated EOG/EMG electrodes, so eye components are found either by ICLabel (a trained classifier) or by correlation with the frontal channels FP1/FP2 used as an EOG proxy. python src/analysis/artifacts.py recordings/xxx.npz # full clean + report python src/analysis/artifacts.py --synth-blinks # demo: inject+remove blinks """ from __future__ import annotations import argparse import sys from dataclasses import dataclass, asdict from pathlib import Path import numpy as np HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE.parent)) # src/ sys.path.insert(0, str(HERE)) # src/analysis/ import load as loadmod # noqa: E402 from load import detect_bad_channels # noqa: E402 EOG_PROXY = ("FP1", "FP2") ICLABEL_KEEP = {"brain", "other"} # remove everything else (over prob) @dataclass class CleanFlags: notch: float = 50.0 l_freq: float = 1.0 h_freq: float = 40.0 interp: bool = True # bad-channel detect + spherical interpolation car: bool = True # common average reference ica: bool = True # ICA artifact removal ica_method: str = "iclabel" # "iclabel" | "eog" (eog = FP1/FP2 correlation proxy) ica_prob: float = 0.80 # ICLabel confidence needed to drop a component # ------------------------------------------------------------------- ICA core def fit_ica(raw, n_components=None, seed=42): """Fit ICA suitable for ICLabel: extended-infomax on a 1 Hz-highpassed copy.""" import mne fit_raw = raw.copy().filter(1.0, None, verbose="ERROR") # highpass 1 Hz for stable ICA picks = mne.pick_types(fit_raw.info, eeg=True, exclude="bads") if n_components is None: n_components = min(20, len(picks) - 1) ica = mne.preprocessing.ICA(n_components=n_components, method="infomax", fit_params=dict(extended=True), max_iter="auto", random_state=seed, verbose="ERROR") ica.fit(fit_raw, picks=picks, verbose="ERROR") return ica def label_ica(ica, raw, method="iclabel", eog_ch=EOG_PROXY, prob=0.80): """Return (labels, probs, exclude). ICLabel classifies every component; the EOG proxy only flags eye components by correlating with FP1/FP2.""" n = ica.n_components_ if method == "iclabel": from mne_icalabel import label_components raw_car = raw.copy().set_eeg_reference("average", verbose="ERROR") res = label_components(raw_car, ica, method="iclabel") labels = list(res["labels"]) probs = np.asarray(res["y_pred_proba"]).ravel() exclude = [i for i in range(n) if labels[i] not in ICLABEL_KEEP and probs[i] >= prob] return labels, probs, exclude # EOG proxy: correlate components with the frontal channels exclude, scoremap = [], np.zeros(n) for ch in eog_ch: if ch not in raw.ch_names: continue idx, scores = ica.find_bads_eog(raw, ch_name=ch, verbose="ERROR") exclude += idx scoremap = np.maximum(scoremap, np.abs(np.asarray(scores)[:n])) exclude = sorted(set(exclude)) labels = ["eye (EOG proxy)" if i in exclude else "kept" for i in range(n)] return labels, scoremap, exclude # ------------------------------------------------------------- raw-level clean def preprocess(raw, flags: CleanFlags, verbose=False): """Apply the flag-driven pipeline. Returns (clean_raw, report).""" import mne report = {"interpolated": [], "ica": None} raw = raw.copy() if flags.notch: raw.notch_filter(flags.notch, verbose="ERROR") if flags.l_freq or flags.h_freq: raw.filter(flags.l_freq or None, flags.h_freq or None, verbose="ERROR") bad = detect_bad_channels(raw) # always FLAG (so CAR/ICA exclude them) … raw.info["bads"] = bad report["bad"] = bad if flags.interp and bad and raw.get_montage() is not None: raw.interpolate_bads(reset_bads=True, verbose="ERROR") # … interpolate only if asked report["interpolated"] = bad if flags.car: raw.set_eeg_reference("average", verbose="ERROR") if flags.ica: try: ica = fit_ica(raw) labels, probs, exclude = label_ica(ica, raw, flags.ica_method, prob=flags.ica_prob) ica.exclude = exclude ica.apply(raw, verbose="ERROR") report["ica"] = dict(method=flags.ica_method, n_components=ica.n_components_, labels=labels, probs=[float(p) for p in np.ravel(probs)], exclude=exclude, removed=len(exclude)) report["_ica_obj"] = ica except Exception as e: report["ica"] = dict(error=f"{type(e).__name__}: {e}") if verbose: print(" interpolated:", report["interpolated"]) print(" ICA:", {k: v for k, v in (report["ica"] or {}).items() if k != "probs"}) return raw, report # --------------------------------------------------------- epoch-level (autoreject) def clean_epochs(epochs, seed=11, verbose=False): """autoreject: cross-validated per-channel thresholds → repair or drop bad epochs. Returns (epochs_clean, reject_log). Epoch-based & offline — not a live operation.""" from autoreject import AutoReject ar = AutoReject(random_state=seed, n_jobs=1, verbose=False) clean = ar.fit_transform(epochs, return_log=False) log = ar.get_reject_log(epochs) if verbose: print(f" autoreject: {len(epochs)} → {len(clean)} epochs " f"({int(log.bad_epochs.sum())} dropped)") return clean, log # --------------------------------------------------- live de-blink linear operator def build_deblink_operator(ica, info): """Precompute the 32×32 sensor→sensor cleaning matrix M (and bias b) for a fitted ICA, so the live scope can de-blink a chunk with one matmul: x_clean ≈ M @ x + b. Derived by probing the (affine) ica.apply with impulses — exact & version-independent.""" import mne nch = info["nchan"] probe = np.hstack([np.zeros((nch, 1)), np.eye(nch)]) * 1e-6 # col0 = zero (bias) r = mne.io.RawArray(probe, info, verbose="ERROR") ica.apply(r, verbose="ERROR") out = r.get_data() # (nch, nch+1) volts b = out[:, 0] M = (out[:, 1:] - b[:, None]) / 1e-6 # unit response (dimensionless) return M.astype(np.float32), (b / 1e-6).astype(np.float32) class LiveDeblink: """Fit ICA+labels once on a calibration buffer, then apply M @ chunk online.""" def __init__(self, M, b, info, report): self.M, self.b, self.info, self.report = M, b, info, report @classmethod def calibrate(cls, buf_uv, fs, ch_names, method="iclabel", prob=0.80): import mne info = mne.create_info(list(ch_names), fs, "eeg") info.set_montage(mne.channels.make_standard_montage("standard_1020"), match_case=False, on_missing="ignore", verbose="ERROR") raw = mne.io.RawArray(buf_uv * 1e-6, info, verbose="ERROR") # A dead/railed channel (e.g. an electrode pinned at +full scale) would otherwise # poison the MEAN reference (~fullscale/n_ch on every channel) and make the data # rank-deficient once ICA centres it. Exclude it from both. raw.info["bads"] = detect_bad_channels(raw) raw.set_eeg_reference("average", verbose="ERROR") ica = fit_ica(raw) labels, probs, exclude = label_ica(ica, raw, method, prob=prob) ica.exclude = exclude M, b = build_deblink_operator(ica, info) rep = dict(method=method, n_components=ica.n_components_, labels=labels, exclude=exclude, removed=len(exclude)) return cls(M, b, info, rep) def apply(self, chunk_uv): # (nch, m) µV -> cleaned return self.M @ chunk_uv + self.b[:, None] # --------------------------------------------------------------------- CLI / demo def _synth_blinks_demo(): """Generate a synthetic MI recording WITH blink artifacts, clean it, and report that ICA removed the blinks while the C3/C4 mu-ERD survived.""" sys.path.insert(0, str(HERE)) from erd_ers import synth_mi_recording, _band_timecourse, _tfr, MU, IMAGERY path = synth_mi_recording(reps=12, out=HERE.parents[1] / "recordings" / "synth_mi_blinks.npz", blinks=True) raw, events, event_id = loadmod.read_recording(path) flags = CleanFlags(ica=True, ica_method="iclabel") clean, report = preprocess(raw, flags, verbose=True) fp1 = raw.ch_names.index("FP1") print(f" FP1 std: raw {raw.get_data()[fp1].std()*1e6:.1f} µV → clean " f"{clean.get_data()[fp1].std()*1e6:.1f} µV (blink power should drop)") return path def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("path", nargs="?") ap.add_argument("--synth-blinks", action="store_true", help="inject+remove blinks demo") ap.add_argument("--method", default="iclabel", choices=["iclabel", "eog"]) ap.add_argument("--epochs", action="store_true", help="also run autoreject on epochs") args = ap.parse_args() if args.synth_blinks: _synth_blinks_demo(); return if not args.path: ap.error("give a recording path or use --synth-blinks") raw, events, event_id = loadmod.read_recording(args.path) clean, report = preprocess(raw, CleanFlags(ica_method=args.method), verbose=True) if args.epochs and len(events): import mne ep = mne.Epochs(clean, events, event_id, tmin=loadmod.DEFAULT_TMIN, tmax=loadmod.DEFAULT_TMAX, baseline=None, preload=True, verbose="ERROR") clean_epochs(ep, verbose=True) if __name__ == "__main__": main()