diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index a23efb3d2..592aed487 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -241,6 +241,7 @@ Available Datasets datasets/pyhealth.datasets.COVID19CXRDataset datasets/pyhealth.datasets.ChestXray14Dataset datasets/pyhealth.datasets.PhysioNetDeIDDataset + datasets/pyhealth.datasets.EEGBCIDataset datasets/pyhealth.datasets.TUABDataset datasets/pyhealth.datasets.TUEVDataset datasets/pyhealth.datasets.ClinVarDataset diff --git a/docs/api/datasets/pyhealth.datasets.EEGBCIDataset.rst b/docs/api/datasets/pyhealth.datasets.EEGBCIDataset.rst new file mode 100644 index 000000000..8f4d427e9 --- /dev/null +++ b/docs/api/datasets/pyhealth.datasets.EEGBCIDataset.rst @@ -0,0 +1,7 @@ +pyhealth.datasets.EEGBCIDataset +================================ + +.. autoclass:: pyhealth.datasets.EEGBCIDataset + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 8724176a8..c7910e626 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -223,6 +223,7 @@ Available Tasks Sleep Staging Sleep Staging (SleepEDF) Temple University EEG Tasks + EEGBCI Tasks Sleep Staging v2 Benchmark EHRShot ChestX-ray14 Binary Classification diff --git a/docs/api/tasks/pyhealth.tasks.eegbci.rst b/docs/api/tasks/pyhealth.tasks.eegbci.rst new file mode 100644 index 000000000..b2682057f --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.eegbci.rst @@ -0,0 +1,7 @@ +pyhealth.tasks.eegbci +===================== + +.. automodule:: pyhealth.tasks.eegbci + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/eeg/eegbci/README.md b/examples/eeg/eegbci/README.md new file mode 100644 index 000000000..fadb008c6 --- /dev/null +++ b/examples/eeg/eegbci/README.md @@ -0,0 +1,56 @@ +# EEGBCI Pattern Discovery + +This example uses `EEGBCIDataset` and `EEGBCIPatternDiscovery` to create +2-second EEGBCI windows with task labels, Welch bandpower features, and cautious +frequency-profile interpretations. + +The interpretations are exploratory signal metadata. They are not clinical +diagnoses and do not prove a subject's cognition. + +Run a tiny real-data example: + +```bash +python examples/eeg/eegbci/eegbci_pattern_discovery.py \ + --subjects 1 \ + --runs 3 \ + --max-windows 20 \ + --download +``` + +Outputs are written to `outputs/eegbci_pattern_discovery/` by default: + +- `eegbci_pattern_windows.csv` +- `eegbci_pattern_summary.md` + +The CSV has one row per emitted 2-second window. Key columns include subject/run +metadata, `event_code`, decoded `task_label`, raw EEGBCI numeric label +(`eegbci_label` / `label`), PyHealth model-local label (`model_label`), +absolute window timing, band powers, relative band powers, `dominant_band`, +frequency ratios, and `interpretation`. + +The moment-report columns add analysis-grade fields: + +- `analysis_version` +- `state_hypothesis`, `state_confidence`, and `evidence_score` +- `evidence_summary` +- `rest_reference_scope` and rest-normalized relative band deltas +- `task_state_relation`, `task_state_rationale`, and `task_state_confidence` +- `is_low_confidence`, `is_possible_artifact`, and `is_mixed_or_ambiguous` + +The `interpretation` column is report-level text derived from these moment-report +fields. Legacy task-level fields such as `brain_state_hypothesis`, `confidence`, +and `quality_flags` are intentionally not written to the CSV. + +The Markdown report summarizes state counts, task-label/state agreement, +rest-normalized bandpower deltas, confidence and quality flags, representative +windows, limitations, and next checks. These labels are signal-pattern +summaries from short EEG windows, not clinical findings or evidence of a +subject's cognition. + +Implementation details are tracked in +`docs/eeg_pattern_discovery/moment_report_implementation_plan.md`. + +`--root` points to the local EEGBCI data directory. With `--download`, MNE +downloads any missing EDF files under that root. PyHealth task caches are stored +under the configured PyHealth cache directory and are keyed by the requested +subject/run selection. diff --git a/examples/eeg/eegbci/eegbci_pattern_discovery.py b/examples/eeg/eegbci/eegbci_pattern_discovery.py new file mode 100644 index 000000000..bd86142b5 --- /dev/null +++ b/examples/eeg/eegbci/eegbci_pattern_discovery.py @@ -0,0 +1,659 @@ +from __future__ import annotations + +import argparse +import sys +from collections import Counter +from pathlib import Path + +import pandas as pd + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from pyhealth.datasets import EEGBCIDataset +from pyhealth.tasks import EEGBCIPatternDiscovery + + +ANALYSIS_VERSION = "eegbci_pattern_moment_report_v1" +REPORT_BANDS = ("delta", "theta", "alpha", "beta", "gamma") +STATE_CONFIDENCE_RANK = {"low": 0, "medium": 1, "high": 2} + + +def scalar_value(value): + if hasattr(value, "item"): + return value.item() + return value + + +def parse_int_list(value: str) -> list[int]: + items: list[int] = [] + for raw_part in value.split(","): + part = raw_part.strip() + if not part: + raise ValueError("Empty value in integer list") + if "-" in part: + start_text, end_text = part.split("-", 1) + start = int(start_text.strip()) + end = int(end_text.strip()) + if start > end: + raise ValueError("Range start must be <= range end") + items.extend(range(start, end + 1)) + else: + items.append(int(part)) + return items + + +def sample_to_row(sample: dict) -> dict: + bandpower = sample["bandpower"] + model_label = scalar_value(sample["label"]) + eegbci_label = scalar_value(sample.get("eegbci_label", model_label)) + return { + "patient_id": sample["patient_id"], + "record_id": sample["record_id"], + "subject_id": sample["subject_id"], + "run": sample["run"], + "run_type": sample["run_type"], + "trial_id": sample["trial_id"], + "event_code": sample["event_code"], + "task_label": sample["task_label"], + "label_family": sample["label_family"], + "label": eegbci_label, + "eegbci_label": eegbci_label, + "model_label": model_label, + "start_time": sample["start_time"], + "end_time": sample["end_time"], + "dominant_band": bandpower["dominant_band"], + "alpha_beta_ratio": bandpower["alpha_beta_ratio"], + "theta_beta_ratio": bandpower["theta_beta_ratio"], + **{key: value for key, value in bandpower.items() if key.endswith("_power")}, + **{key: value for key, value in bandpower.items() if key.endswith("_relative")}, + } + + +def _mean_band_values(rows: list[dict]) -> dict: + means = {} + for band in REPORT_BANDS: + key = f"{band}_relative" + values = [float(row[key]) for row in rows if row.get(key) not in ("", None)] + if values: + means[key] = sum(values) / len(values) + return means + + +def build_rest_baselines(rows: list[dict]) -> dict: + rest_rows = [row for row in rows if row.get("task_label") == "rest"] + same_subject_run = {} + same_subject_all_runs = {} + + subject_run_keys = sorted({(row["subject_id"], row["run"]) for row in rest_rows}) + for key in subject_run_keys: + subject_id, run = key + grouped = [ + row + for row in rest_rows + if row["subject_id"] == subject_id and row["run"] == run + ] + same_subject_run[key] = _mean_band_values(grouped) + + subject_keys = sorted({row["subject_id"] for row in rest_rows}) + for subject_id in subject_keys: + grouped = [row for row in rest_rows if row["subject_id"] == subject_id] + same_subject_all_runs[subject_id] = _mean_band_values(grouped) + + return { + "same_subject_run": same_subject_run, + "same_subject_all_runs": same_subject_all_runs, + "global_rest": _mean_band_values(rest_rows) if rest_rows else None, + } + + +def _baseline_for_row(row: dict, baselines: dict) -> tuple[str, dict | None]: + subject_run_key = (row["subject_id"], row["run"]) + if subject_run_key in baselines["same_subject_run"]: + return "same_subject_run", baselines["same_subject_run"][subject_run_key] + if row["subject_id"] in baselines["same_subject_all_runs"]: + return "same_subject_all_runs", baselines["same_subject_all_runs"][row["subject_id"]] + if baselines["global_rest"]: + return "global_rest", baselines["global_rest"] + return "unavailable", None + + +def _clip01(value: float) -> float: + return max(0.0, min(1.0, value)) + + +def derive_state_hypothesis(row: dict) -> dict: + delta = float(row.get("delta_relative", 0.0) or 0.0) + theta = float(row.get("theta_relative", 0.0) or 0.0) + alpha = float(row.get("alpha_relative", 0.0) or 0.0) + beta = float(row.get("beta_relative", 0.0) or 0.0) + gamma = float(row.get("gamma_relative", 0.0) or 0.0) + alpha_beta = float(row.get("alpha_beta_ratio", 0.0) or 0.0) + theta_beta = float(row.get("theta_beta_ratio", 0.0) or 0.0) + + scores = { + "idle_alpha_profile": _clip01((alpha - 0.25) + min(alpha_beta / 8.0, 0.40)), + "sensorimotor_engagement_profile": _clip01( + (beta - 0.20) + + max(gamma - 0.12, 0.0) + + max(0.0, 1.5 - alpha_beta) / 6.0 + ), + "slow_wave_dominant_pattern": _clip01( + (delta + theta) - 0.45 + min(theta_beta / 8.0, 0.20) + ), + "possible_artifact_profile": _clip01( + (gamma - 0.22) * 2.0 + max(delta - 0.50, 0.0) + ), + } + ordered = sorted(scores.items(), key=lambda item: item[1], reverse=True) + winner, winning_score = ordered[0] + runner_up = ordered[1][1] + margin = winning_score - runner_up + + if winning_score < 0.20 or margin < 0.08: + state = "mixed_ambiguous_profile" + evidence_score = round(max(winning_score, 0.10), 3) + confidence = "low" + else: + state = winner + evidence_score = round(winning_score, 3) + if winning_score >= 0.65 and margin >= 0.20: + confidence = "high" + elif winning_score >= 0.35 and margin >= 0.12: + confidence = "medium" + else: + confidence = "low" + + return { + "state_hypothesis": state, + "state_confidence": confidence, + "evidence_score": evidence_score, + "evidence_summary": ( + f"delta={delta:.3f}; theta={theta:.3f}; alpha={alpha:.3f}; " + f"beta={beta:.3f}; gamma={gamma:.3f}; alpha_beta={alpha_beta:.3f}; " + f"margin={margin:.3f}" + ), + } + + +def derive_task_state_relation(row: dict) -> dict: + label_family = row.get("label_family", "") + task_label = row.get("task_label", "") + state = row.get("state_hypothesis", "") + + if state == "possible_artifact_profile": + relation = "not_applicable" + confidence = "medium" + rationale = ( + "Artifact-like frequency evidence is flagged for inspection instead of " + "task-label comparison." + ) + elif state == "mixed_ambiguous_profile": + relation = "ambiguous" + confidence = "low" + rationale = ( + "No frequency-profile state won clearly enough to compare strongly with " + "the task label." + ) + elif task_label == "rest" and state == "idle_alpha_profile": + relation = "supports_label" + confidence = "medium" + rationale = "The idle-like alpha profile is consistent with a rest-labeled EEGBCI window." + elif label_family == "motor_execution" and state == "sensorimotor_engagement_profile": + relation = "supports_label" + confidence = "medium" + rationale = ( + "The motor-engaged frequency profile is consistent with an " + "execution-labeled window." + ) + elif label_family == "motor_imagery" and state == "sensorimotor_engagement_profile": + relation = "adds_detail" + confidence = "medium" + rationale = ( + "The motor-engaged frequency profile adds signal detail to an " + "imagery-labeled window." + ) + elif label_family in {"motor_execution", "motor_imagery"} and state == "idle_alpha_profile": + relation = "disagrees" + confidence = "medium" + rationale = "The idle-like alpha profile does not align with a motor-labeled EEGBCI window." + elif state == "slow_wave_dominant_pattern": + relation = "adds_detail" + confidence = "low" + rationale = "The slow-wave dominant pattern adds frequency detail but is not a direct task match." + else: + relation = "ambiguous" + confidence = "low" + rationale = ( + "The task label and frequency-profile state do not have a stronger " + "deterministic mapping." + ) + + return { + "task_state_relation": relation, + "task_state_rationale": rationale, + "task_state_confidence": confidence, + } + + +def derive_quality_columns(row: dict) -> dict: + flags = str(row.get("quality_flags", "")) + state = row.get("state_hypothesis", "") + confidence = row.get("state_confidence", row.get("confidence", "")) + return { + "is_low_confidence": confidence == "low", + "is_possible_artifact": state == "possible_artifact_profile" + or "artifact" in flags + or "high_gamma" in flags, + "is_mixed_or_ambiguous": state == "mixed_ambiguous_profile" + or "ambiguous" in flags, + } + + +def derive_moment_interpretation(row: dict) -> str: + state = row.get("state_hypothesis", "missing") + confidence = row.get("state_confidence", "missing") + evidence = row.get("evidence_score", "") + relation = row.get("task_state_relation", "missing") + task = row.get("task_label", "missing") + dominant = row.get("dominant_band", "missing") + scope = row.get("rest_reference_scope", "missing") + return ( + f"The segment is consistent with `{state}` based on a `{dominant}`-dominant " + f"frequency profile ({confidence} confidence, evidence {evidence}). " + f"The task label is `{task}`, the task/state relation is `{relation}`, " + f"and the rest reference is `{scope}`." + ) + + +BASE_OUTPUT_COLUMNS = ( + "patient_id", + "record_id", + "subject_id", + "run", + "run_type", + "trial_id", + "event_code", + "task_label", + "label_family", + "label", + "eegbci_label", + "model_label", + "start_time", + "end_time", + "dominant_band", + "alpha_beta_ratio", + "theta_beta_ratio", + "interpretation", + "delta_power", + "theta_power", + "alpha_power", + "beta_power", + "gamma_power", + "delta_relative", + "theta_relative", + "alpha_relative", + "beta_relative", + "gamma_relative", +) + +MOMENT_REPORT_COLUMNS = ( + "analysis_version", + "state_hypothesis", + "state_confidence", + "evidence_score", + "evidence_summary", + "rest_reference_scope", + "rest_delta_relative_delta", + "rest_theta_relative_delta", + "rest_alpha_relative_delta", + "rest_beta_relative_delta", + "rest_gamma_relative_delta", + "task_state_relation", + "task_state_rationale", + "task_state_confidence", + "is_low_confidence", + "is_possible_artifact", + "is_mixed_or_ambiguous", +) + +OUTPUT_COLUMNS = BASE_OUTPUT_COLUMNS + MOMENT_REPORT_COLUMNS + + +def annotate_moment_rows(rows: list[dict], baselines: dict) -> list[dict]: + annotated = [] + for row in rows: + next_row = dict(row) + scope, baseline = _baseline_for_row(next_row, baselines) + next_row["analysis_version"] = ANALYSIS_VERSION + next_row["rest_reference_scope"] = scope + + for band in REPORT_BANDS: + source_key = f"{band}_relative" + delta_key = f"rest_{band}_relative_delta" + if baseline and source_key in baseline and next_row.get(source_key) not in ("", None): + next_row[delta_key] = round( + float(next_row[source_key]) - float(baseline[source_key]), 6 + ) + else: + next_row[delta_key] = "" + + next_row.update(derive_state_hypothesis(next_row)) + next_row.update(derive_task_state_relation(next_row)) + next_row["interpretation"] = derive_moment_interpretation(next_row) + next_row.update(derive_quality_columns(next_row)) + annotated.append(next_row) + return annotated + + +def _stable_row_key(row: dict) -> tuple: + return ( + row.get("subject_id", 0), + row.get("run", 0), + float(row.get("start_time", 0.0) or 0.0), + ) + + +def _strongest_row(rows: list[dict]) -> dict | None: + if not rows: + return None + return sorted( + rows, + key=lambda row: ( + -float(row.get("evidence_score", 0.0) or 0.0), + -STATE_CONFIDENCE_RANK.get(row.get("state_confidence", "low"), 0), + *_stable_row_key(row), + ), + )[0] + + +def select_representative_windows(rows: list[dict]) -> dict: + definitions = { + "strongest_idle_like": "idle_alpha_profile", + "strongest_motor_engaged": "sensorimotor_engagement_profile", + "strongest_slow_wave": "slow_wave_dominant_pattern", + "strongest_artifact_like": "possible_artifact_profile", + } + cards = {} + absent = [] + + for card_name, state in definitions.items(): + candidate = _strongest_row( + [row for row in rows if row.get("state_hypothesis") == state] + ) + if candidate is None: + absent.append(card_name) + else: + cards[card_name] = candidate + + ambiguous = [ + row for row in rows if row.get("state_hypothesis") == "mixed_ambiguous_profile" + ] + if ambiguous: + cards["most_ambiguous"] = sorted( + ambiguous, + key=lambda row: ( + float(row.get("evidence_score", 0.0) or 0.0), + -STATE_CONFIDENCE_RANK.get(row.get("state_confidence", "low"), 0), + *_stable_row_key(row), + ), + )[0] + else: + absent.append("most_ambiguous") + + disagreement = _strongest_row( + [row for row in rows if row.get("task_state_relation") == "disagrees"] + ) + if disagreement is None: + absent.append("strongest_task_state_disagreement") + else: + cards["strongest_task_state_disagreement"] = disagreement + + return {"cards": cards, "absent": absent} + + +def _format_count_lines(counter: Counter) -> list[str]: + if not counter: + return ["- None"] + return [f"- {label}: {count}" for label, count in counter.most_common()] + + +def _format_card(row: dict) -> list[str]: + bands = ", ".join( + f"{band}={float(row.get(f'{band}_relative', 0.0) or 0.0):.3f}" + for band in REPORT_BANDS + ) + deltas = ", ".join( + f"{band}={row.get(f'rest_{band}_relative_delta', '')}" + for band in REPORT_BANDS + ) + return [ + f"- Subject {row.get('subject_id')} run {row.get('run')} trial {row.get('trial_id')}", + f" - Task: {row.get('task_label')} from {row.get('start_time')}s to {row.get('end_time')}s", + ( + f" - State: {row.get('state_hypothesis')} " + f"({row.get('state_confidence')}, evidence {row.get('evidence_score')})" + ), + f" - Dominant band: {row.get('dominant_band')}; relative bands: {bands}", + f" - Rest deltas: {deltas}; scope: {row.get('rest_reference_scope')}", + ( + f" - Task relation: {row.get('task_state_relation')} " + f"({row.get('task_state_confidence')})" + ), + ( + f" - Flags: low_confidence={row.get('is_low_confidence')}, " + f"possible_artifact={row.get('is_possible_artifact')}, " + f"mixed_or_ambiguous={row.get('is_mixed_or_ambiguous')}" + ), + f" - Rationale: {row.get('task_state_rationale')}", + ] + + +def render_summary(rows: list[dict], config: dict) -> str: + state_counts = Counter(row.get("state_hypothesis", "missing") for row in rows) + task_counts = Counter(row.get("task_label", "missing") for row in rows) + confidence_counts = Counter(row.get("state_confidence", "missing") for row in rows) + relation_counts = Counter(row.get("task_state_relation", "missing") for row in rows) + unavailable_rest = sum( + row.get("rest_reference_scope") == "unavailable" for row in rows + ) + low_confidence = sum(bool(row.get("is_low_confidence")) for row in rows) + artifacts = sum(bool(row.get("is_possible_artifact")) for row in rows) + ambiguous = sum(bool(row.get("is_mixed_or_ambiguous")) for row in rows) + representatives = select_representative_windows(rows) + + executive = [] + if not rows: + executive.append("No windows were produced for the requested configuration.") + else: + top_state, top_state_count = state_counts.most_common(1)[0] + executive.append( + f"Processed {len(rows)} windows. Most common state: `{top_state}` " + f"({top_state_count}/{len(rows)})." + ) + if low_confidence == len(rows): + executive.append("Every window is low confidence.") + if len(state_counts) == 1: + executive.append( + "Every window maps to the same state; broaden coverage or review thresholds." + ) + if unavailable_rest == len(rows): + executive.append("No rest baseline was available for the emitted rows.") + if config.get("output_was_capped"): + executive.append("Output was capped by `--max-windows`.") + + lines = [ + "# EEGBCI Pattern Discovery Moment Report", + "", + f"Analysis version: `{ANALYSIS_VERSION}`", + "", + "## Executive Result", + "", + *[f"- {item}" for item in executive], + "", + "## Run Configuration", + "", + f"- Subjects: {config.get('subjects')}", + f"- Runs: {config.get('runs')}", + f"- Max windows: {config.get('max_windows')}", + f"- Baseline source rows: {config.get('baseline_row_count')}", + "", + "## Window Coverage", + "", + f"- Output windows: {len(rows)}", + f"- Task labels: {dict(task_counts)}", + "", + "## Moment-State Summary", + "", + *_format_count_lines(state_counts), + "", + "## Task Label x State Matrix", + "", + ] + + matrix = Counter( + (row.get("task_label", "missing"), row.get("state_hypothesis", "missing")) + for row in rows + ) + if matrix: + for (task_label, state), count in sorted(matrix.items()): + lines.append(f"- {task_label} x {state}: {count}") + else: + lines.append("- None") + + lines.extend( + [ + "", + "## Rest-Normalized Bandpower Summary", + "", + f"- Rows with unavailable rest baseline: {unavailable_rest}", + ] + ) + for band in REPORT_BANDS: + key = f"rest_{band}_relative_delta" + values = [float(row.get(key)) for row in rows if row.get(key) not in ("", None)] + if values: + lines.append(f"- {band}: mean delta {sum(values) / len(values):.3f}") + else: + lines.append(f"- {band}: unavailable") + + lines.extend( + [ + "", + "## Confidence and Quality Audit", + "", + f"- State confidence: {dict(confidence_counts)}", + f"- Task-state relations: {dict(relation_counts)}", + f"- Low-confidence rows: {low_confidence}", + f"- Possible artifact rows: {artifacts}", + f"- Mixed or ambiguous rows: {ambiguous}", + "", + "## Representative Windows", + "", + ] + ) + if representatives["cards"]: + for card_name, row in representatives["cards"].items(): + lines.append(f"### {card_name.replace('_', ' ').title()}") + lines.extend(_format_card(row)) + lines.append("") + else: + lines.append("- None") + if representatives["absent"]: + lines.append( + f"- Absent representative classes: {', '.join(representatives['absent'])}" + ) + + lines.extend( + [ + "", + "## Limitations", + "", + ( + "- These labels are signal-pattern summaries from short EEG windows. " + "They are not clinical findings and should not be read as evidence " + "of a subject's cognition." + ), + ] + ) + if unavailable_rest: + lines.append("- No rest baseline was available for at least one emitted row.") + if config.get("output_was_capped"): + lines.append( + "- The output was capped, so the artifact may not represent all requested windows." + ) + + lines.extend( + [ + "", + "## Next Checks", + "", + "- Run with broader subjects/runs to verify that state diversity improves.", + "- Inspect possible artifact rows before drawing conclusions from state counts.", + "- Compare rest-normalized deltas against the raw relative band shares.", + ] + ) + return "\n".join(lines).rstrip() + "\n" + + +def write_summary(rows: list[dict], path: Path, config: dict) -> None: + path.write_text(render_summary(rows, config), encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--root", default="~/.cache/pyhealth/eegbci") + parser.add_argument("--subjects", default="1,2,3") + parser.add_argument("--runs", default="3-14") + parser.add_argument("--output-dir", default="outputs/eegbci_pattern_discovery") + parser.add_argument("--max-windows", type=int, default=None) + parser.add_argument("--download", action="store_true") + args = parser.parse_args() + + output_dir = Path(args.output_dir).expanduser() + output_dir.mkdir(parents=True, exist_ok=True) + + requested_subjects = parse_int_list(args.subjects) + requested_runs = parse_int_list(args.runs) + dataset = EEGBCIDataset( + root=str(Path(args.root).expanduser()), + subjects=requested_subjects, + runs=requested_runs, + download=args.download, + ) + sample_dataset = dataset.set_task(EEGBCIPatternDiscovery(compute_stft=False)) + + all_rows = [sample_to_row(sample) for sample in sample_dataset] + baseline_row_count = sum(row.get("task_label") == "rest" for row in all_rows) + baselines = build_rest_baselines(all_rows) + annotated_rows = annotate_moment_rows(all_rows, baselines) + output_rows = ( + annotated_rows[: args.max_windows] + if args.max_windows is not None + else annotated_rows + ) + output_was_capped = ( + args.max_windows is not None and len(annotated_rows) > len(output_rows) + ) + + csv_path = output_dir / "eegbci_pattern_windows.csv" + summary_path = output_dir / "eegbci_pattern_summary.md" + pd.DataFrame(output_rows, columns=OUTPUT_COLUMNS).to_csv(csv_path, index=False) + write_summary( + output_rows, + summary_path, + { + "subjects": getattr(dataset, "subjects", requested_subjects), + "runs": getattr(dataset, "runs", requested_runs), + "max_windows": args.max_windows, + "baseline_row_count": baseline_row_count, + "output_was_capped": output_was_capped, + }, + ) + print(f"Wrote {csv_path}") + print(f"Wrote {summary_path}") + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index c29955e7d..3b7d6213a 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -82,6 +82,7 @@ def __init__(self, *args, **kwargs): split_by_visit, split_by_visit_conformal, ) +from .eegbci import EEGBCIDataset from .tuab import TUABDataset from .tuev import TUEVDataset from .utils import ( diff --git a/pyhealth/datasets/configs/eegbci.yaml b/pyhealth/datasets/configs/eegbci.yaml new file mode 100644 index 000000000..edb0e5b7a --- /dev/null +++ b/pyhealth/datasets/configs/eegbci.yaml @@ -0,0 +1,13 @@ +version: "1.0.0" +tables: + records: + file_path: "eegbci-pyhealth.csv" + patient_id: "patient_id" + timestamp: null + attributes: + - "record_id" + - "subject_id" + - "run" + - "run_type" + - "signal_file" + - "source" diff --git a/pyhealth/datasets/eegbci.py b/pyhealth/datasets/eegbci.py new file mode 100644 index 000000000..8e83022ed --- /dev/null +++ b/pyhealth/datasets/eegbci.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import hashlib +import json +import logging +from pathlib import Path +from typing import Optional + +import mne +import pandas as pd + +from .base_dataset import BaseDataset +from pyhealth.tasks.eegbci import EEGBCIPatternDiscovery, run_type_for_run + +logger = logging.getLogger(__name__) + +EEGBCI_METADATA_COLUMNS = { + "patient_id", + "record_id", + "subject_id", + "run", + "run_type", + "signal_file", + "source", +} + + +class EEGBCIDataset(BaseDataset): + """PhysioNet EEG Motor Movement/Imagery metadata dataset.""" + + def __init__( + self, + root: str, + dataset_name: Optional[str] = None, + config_path: Optional[str] = None, + subjects: Optional[list[int]] = None, + runs: Optional[list[int]] = None, + download: bool = False, + **kwargs, + ) -> None: + if config_path is None: + config_path = Path(__file__).parent / "configs" / "eegbci.yaml" + self.root = root + self.subjects = self._normalize_selection( + list(subjects) if subjects is not None else [1, 2, 3] + ) + self.runs = self._normalize_selection( + list(runs) if runs is not None else list(range(3, 15)) + ) + self.download = download + self.selection_key = self._build_selection_key() + self.metadata_file_name = self._metadata_file_name() + self.prepare_metadata() + dataset_name = dataset_name or "eegbci" + super().__init__( + root=root, + tables=["records"], + dataset_name=f"{dataset_name}_{self.selection_key}", + config_path=config_path, + **kwargs, + ) + if self.config is not None: + self.config.tables["records"].file_path = self.metadata_file_name + + @staticmethod + def _normalize_selection(values: list[int]) -> list[int]: + return sorted({int(value) for value in values}) + + def _build_selection_key(self) -> str: + payload = { + "subjects": [int(subject) for subject in self.subjects], + "runs": [int(run) for run in self.runs], + } + digest = hashlib.sha1( + json.dumps(payload, sort_keys=True).encode("utf-8") + ).hexdigest()[:10] + subject_part = "-".join(f"{int(subject):03d}" for subject in self.subjects) + run_part = "-".join(f"{int(run):02d}" for run in self.runs) + return f"s{subject_part}_r{run_part}_{digest}" + + def _metadata_file_name(self) -> str: + return f"eegbci-pyhealth-{self.selection_key}.csv" + + def _find_local_edf(self, subject: int, run: int) -> Path | None: + root = Path(self.root) + filename = f"S{subject:03d}R{run:02d}.edf" + canonical_path = ( + root / "files" / "eegmmidb" / "1.0.0" / f"S{subject:03d}" / filename + ) + if canonical_path.exists(): + return canonical_path + matches = sorted(root.rglob(filename)) + return matches[0] if matches else None + + def _requested_pairs(self) -> list[tuple[int, int]]: + return sorted( + (int(subject), int(run)) + for subject in self.subjects + for run in self.runs + ) + + def _metadata_matches_request(self, csv_path: Path) -> bool: + try: + df = pd.read_csv(csv_path) + except Exception: + return False + if not EEGBCI_METADATA_COLUMNS.issubset(df.columns): + return False + pairs = sorted((int(row.subject_id), int(row.run)) for row in df.itertuples()) + return pairs == self._requested_pairs() + + def prepare_metadata(self) -> None: + root = Path(self.root) + csv_path = root / self.metadata_file_name + if csv_path.exists() and self._metadata_matches_request(csv_path): + return + + rows: list[dict] = [] + for subject in self.subjects: + paths_by_run: dict[int, Path] = {} + if self.download: + downloaded = mne.datasets.eegbci.load_data( + subject, self.runs, path=str(root), update_path=False + ) + for path in downloaded: + p = Path(path) + for run in self.runs: + if p.name == f"S{subject:03d}R{run:02d}.edf": + paths_by_run[run] = p + for run in self.runs: + signal_file = paths_by_run.get(run) or self._find_local_edf(subject, run) + if signal_file is None: + raise FileNotFoundError( + f"Missing EEGBCI EDF for subject {subject}, run {run}. " + "Pass download=True to fetch it with MNE." + ) + rows.append( + { + "patient_id": f"S{subject:03d}", + "record_id": f"R{run:02d}", + "subject_id": int(subject), + "run": int(run), + "run_type": run_type_for_run(run), + "signal_file": str(signal_file), + "source": "physionet_eegbci", + } + ) + + df = pd.DataFrame(rows) + df.sort_values(["subject_id", "run"], inplace=True) + df.reset_index(drop=True, inplace=True) + csv_path.parent.mkdir(parents=True, exist_ok=True) + df.to_csv(csv_path, index=False) + logger.info("Wrote EEGBCI metadata to %s", csv_path) + + @property + def default_task(self) -> EEGBCIPatternDiscovery: + return EEGBCIPatternDiscovery() diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index cc95ef94e..1a0db720f 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -69,6 +69,7 @@ EEGEventsTUEV, EEGAbnormalTUAB ) +from .eegbci import EEGBCIPatternDiscovery, EEGMotorImageryEEGBCI from .variant_classification import ( MutationPathogenicityPrediction, VariantClassificationClinVar, diff --git a/pyhealth/tasks/eegbci.py b/pyhealth/tasks/eegbci.py new file mode 100644 index 000000000..b834e0ce7 --- /dev/null +++ b/pyhealth/tasks/eegbci.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +import mne +import numpy as np +import torch + +from pyhealth.tasks import BaseTask + +EEGBCI_RUN_TYPES = { + 3: "motor_execution_left_right", + 4: "motor_imagery_left_right", + 5: "motor_execution_fists_feet", + 6: "motor_imagery_fists_feet", + 7: "motor_execution_left_right", + 8: "motor_imagery_left_right", + 9: "motor_execution_fists_feet", + 10: "motor_imagery_fists_feet", + 11: "motor_execution_left_right", + 12: "motor_imagery_left_right", + 13: "motor_execution_fists_feet", + 14: "motor_imagery_fists_feet", +} + +EEGBCI_LABELS = { + "rest": 0, + "execute_left_fist": 1, + "execute_right_fist": 2, + "imagine_left_fist": 3, + "imagine_right_fist": 4, + "execute_both_fists": 5, + "execute_both_feet": 6, + "imagine_both_fists": 7, + "imagine_both_feet": 8, +} + + +def run_type_for_run(run: int) -> str: + try: + return EEGBCI_RUN_TYPES[int(run)] + except KeyError as exc: + raise ValueError(f"Unsupported EEGBCI run: {run}") from exc + + +def label_family_for_run(run: int) -> str: + run_type = run_type_for_run(run) + if "execution" in run_type: + return "motor_execution" + if "imagery" in run_type: + return "motor_imagery" + return "baseline" + + +def task_label_for_event(run: int, event_code: str) -> str: + code = str(event_code).strip() + if code == "T0": + return "rest" + run_type = run_type_for_run(run) + mapping = { + "motor_execution_left_right": { + "T1": "execute_left_fist", + "T2": "execute_right_fist", + }, + "motor_imagery_left_right": { + "T1": "imagine_left_fist", + "T2": "imagine_right_fist", + }, + "motor_execution_fists_feet": { + "T1": "execute_both_fists", + "T2": "execute_both_feet", + }, + "motor_imagery_fists_feet": { + "T1": "imagine_both_fists", + "T2": "imagine_both_feet", + }, + } + try: + return mapping[run_type][code] + except KeyError as exc: + raise ValueError(f"Unsupported EEGBCI event {event_code!r} for run {run}") from exc + + +def numeric_label_for_task(task_label: str) -> int: + try: + return EEGBCI_LABELS[task_label] + except KeyError as exc: + raise ValueError(f"Unsupported EEGBCI task label: {task_label}") from exc + + +EEGBCI_COMPAT_CHANNELS = ( + "FC5", + "FC3", + "FC1", + "FC2", + "FC4", + "FC6", + "C5", + "C3", + "C1", + "C2", + "C4", + "C6", + "CP5", + "CP3", + "CP4", + "CP6", +) + + +def normalize_eegbci_channel_name(name: str) -> str: + clean = name.upper().replace(".", "").replace("EEG ", "").replace("-REF", "") + aliases = { + "T9": "FT9", + "T10": "FT10", + } + return aliases.get(clean, clean) + + +def select_eegbci_channels( + data: np.ndarray, + ch_names: List[str], + channel_mode: str = "compat16", +) -> Tuple[np.ndarray, List[str]]: + if channel_mode == "all": + return data, list(ch_names) + if channel_mode != "compat16": + raise ValueError("channel_mode must be one of {'compat16', 'all'}") + + normalized_to_index = { + normalize_eegbci_channel_name(name): idx for idx, name in enumerate(ch_names) + } + missing = [ch for ch in EEGBCI_COMPAT_CHANNELS if ch not in normalized_to_index] + if missing: + raise ValueError(f"Missing EEGBCI channels for compat16 mode: {missing}") + indices = [normalized_to_index[ch] for ch in EEGBCI_COMPAT_CHANNELS] + return data[indices], list(EEGBCI_COMPAT_CHANNELS) + + +def normalize_signal(signal: np.ndarray, mode: str | None) -> np.ndarray: + if mode is None: + return signal + if mode == "95th_percentile": + scale = np.quantile( + np.abs(signal), q=0.95, axis=-1, method="linear", keepdims=True + ) + return signal / (scale + 1e-8) + if mode == "div_by_100": + return signal / 100.0 + raise ValueError("normalization must be one of {None, '95th_percentile', 'div_by_100'}") + + +BANDS = { + "delta": (0.5, 4.0), + "theta": (4.0, 8.0), + "alpha": (8.0, 13.0), + "beta": (13.0, 30.0), + "gamma": (30.0, 45.0), +} + + +def compute_band_powers(data: np.ndarray, sfreq: float) -> Dict[str, float | str]: + from scipy.signal import welch + + if data.ndim != 2: + raise ValueError("data must have shape (channels, time)") + nperseg = min(data.shape[-1], int(sfreq * 2)) + freqs, psd = welch(data, fs=sfreq, nperseg=nperseg, axis=-1) + mean_psd = psd.mean(axis=0) + + features: Dict[str, float | str] = {} + total_power = 0.0 + band_values: Dict[str, float] = {} + for band, (low, high) in BANDS.items(): + mask = (freqs >= low) & (freqs < high) + value = float(np.trapezoid(mean_psd[mask], freqs[mask])) if np.any(mask) else 0.0 + features[f"{band}_power"] = value + band_values[band] = value + total_power += value + + denom = total_power + 1e-12 + for band, value in band_values.items(): + features[f"{band}_relative"] = float(value / denom) + + features["dominant_band"] = max(band_values, key=band_values.get) + features["alpha_beta_ratio"] = float( + band_values["alpha"] / (band_values["beta"] + 1e-12) + ) + features["theta_beta_ratio"] = float( + band_values["theta"] / (band_values["beta"] + 1e-12) + ) + return features + + +def interpret_band_profile(features: Dict[str, float | str]) -> Dict[str, str]: + dominant = str(features["dominant_band"]) + alpha_rel = float(features.get("alpha_relative", 0.0)) + beta_rel = float(features.get("beta_relative", 0.0)) + theta_rel = float(features.get("theta_relative", 0.0)) + gamma_rel = float(features.get("gamma_relative", 0.0)) + alpha_beta = float(features.get("alpha_beta_ratio", 0.0)) + theta_beta = float(features.get("theta_beta_ratio", 0.0)) + + quality_flags: List[str] = [] + hypothesis = "mixed_frequency_profile" + confidence = "low" + + if dominant == "alpha" and alpha_rel >= 0.45 and alpha_beta >= 2.0: + hypothesis = "relaxed_or_idle" + confidence = "medium" + elif dominant == "beta" and beta_rel >= 0.35: + hypothesis = "active_sensorimotor_processing" + confidence = "medium" + elif dominant == "theta" and theta_rel >= 0.35 and theta_beta >= 1.5: + hypothesis = "slow_wave_or_drowsy_pattern" + confidence = "medium" + elif dominant == "gamma" and gamma_rel >= 0.30: + hypothesis = "high_frequency_or_artifact_pattern" + confidence = "low" + quality_flags.append("possible_muscle_artifact") + + if confidence == "low": + quality_flags.append("low_confidence") + + return { + "brain_state_hypothesis": hypothesis, + "confidence": confidence, + "quality_flags": ";".join(quality_flags) if quality_flags else "none", + "interpretation": ( + f"The segment is consistent with {hypothesis} based on a " + f"{dominant}-dominant frequency profile." + ), + } + + +def iter_annotation_windows( + raw: mne.io.BaseRaw, + run: int, + window_size: float = 2.0, +) -> List[Dict[str, Any]]: + sfreq = float(raw.info["sfreq"]) + window_samples = int(round(window_size * sfreq)) + windows: List[Dict[str, Any]] = [] + for idx, annotation in enumerate(raw.annotations): + event_code = str(annotation["description"]) + if event_code not in {"T0", "T1", "T2"}: + continue + start_sample = int( + raw.time_as_index([float(annotation["onset"])], use_rounding=True)[0] + ) + duration_samples = int(round(float(annotation["duration"]) * sfreq)) + n_full_windows = duration_samples // window_samples + for window_idx in range(n_full_windows): + s0 = start_sample + window_idx * window_samples + s1 = s0 + window_samples + task_label = task_label_for_event(run, event_code) + windows.append( + { + "trial_id": f"ann{idx:04d}_win{window_idx:03d}", + "event_code": event_code, + "task_label": task_label, + "label_family": label_family_for_run(run), + "label": numeric_label_for_task(task_label), + "start_time": s0 / sfreq, + "end_time": s1 / sfreq, + "start_sample": s0, + "end_sample": s1, + } + ) + return windows + + +class EEGMotorImageryEEGBCI(BaseTask): + task_name: str = "EEGBCI_motor_imagery" + input_schema: Dict[str, str] = {"signal": "tensor", "stft": "tensor"} + output_schema: Dict[str, str] = {"label": "multiclass"} + + def __init__( + self, + window_size: float = 2.0, + resample_rate: float | None = 200, + bandpass_filter: Tuple[float, float] | None = (0.5, 45.0), + channel_mode: str = "compat16", + normalization: str | None = "95th_percentile", + compute_stft: bool = True, + ) -> None: + super().__init__() + self.window_size = window_size + self.resample_rate = resample_rate + self.bandpass_filter = bandpass_filter + self.channel_mode = channel_mode + self.normalization = normalization + self.compute_stft = compute_stft + if not compute_stft: + self.input_schema = {"signal": "tensor"} + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + return self._base_samples_from_patient(patient) + + def read_raw(self, signal_file: str) -> mne.io.BaseRaw: + raw = mne.io.read_raw_edf(signal_file, preload=True, verbose="error") + raw.pick_types(eeg=True, stim=False, exclude=[]) + if self.bandpass_filter is not None: + raw.filter( + l_freq=self.bandpass_filter[0], + h_freq=self.bandpass_filter[1], + verbose="error", + ) + if self.resample_rate is not None: + raw.resample(self.resample_rate, n_jobs=1, verbose="error") + return raw + + def _base_samples_from_patient(self, patient: Any) -> List[Dict[str, Any]]: + samples: List[Dict[str, Any]] = [] + for event in patient.get_events("records"): + raw = self.read_raw(event.signal_file) + data = raw.get_data(units="uV") + selected, selected_names = select_eegbci_channels( + data, raw.ch_names, self.channel_mode + ) + selected = normalize_signal(selected, self.normalization) + sfreq = float(raw.info["sfreq"]) + for idx, window in enumerate( + iter_annotation_windows(raw, int(event.run), self.window_size) + ): + signal_np = selected[:, window["start_sample"] : window["end_sample"]] + if signal_np.shape[-1] != int(round(self.window_size * sfreq)): + continue + signal = torch.FloatTensor(signal_np) + sample = { + "patient_id": patient.patient_id, + "record_id": event.record_id, + "subject_id": int(event.subject_id), + "run": int(event.run), + "run_type": event.run_type, + "signal_file": event.signal_file, + "trial_id": f"{patient.patient_id}_{event.record_id}_{idx:04d}", + "event_code": window["event_code"], + "task_label": window["task_label"], + "label_family": window["label_family"], + "label": int(window["label"]), + "eegbci_label": int(window["label"]), + "signal": signal, + "channel_names": selected_names, + "start_time": window["start_time"], + "end_time": window["end_time"], + "sample_rate": sfreq, + } + if self.compute_stft: + from pyhealth.models.tfm_tokenizer import get_stft_torch + + sample["stft"] = get_stft_torch( + signal.unsqueeze(0), resampling_rate=int(round(sfreq)) + ).squeeze(0) + samples.append(sample) + raw.close() + return samples + + +class EEGBCIPatternDiscovery(EEGMotorImageryEEGBCI): + task_name: str = "EEGBCI_pattern_discovery" + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + samples = self._base_samples_from_patient(patient) + for sample in samples: + features = compute_band_powers( + sample["signal"].detach().cpu().numpy(), + float(sample["sample_rate"]), + ) + interpretation = interpret_band_profile(features) + sample["bandpower"] = features + sample.update(interpretation) + return samples diff --git a/tests/core/test_eegbci.py b/tests/core/test_eegbci.py new file mode 100644 index 000000000..275e8f212 --- /dev/null +++ b/tests/core/test_eegbci.py @@ -0,0 +1,1522 @@ +import os +import sys +import unittest +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import List +from unittest.mock import patch + +import numpy as np +import pandas as pd +import torch + +from pyhealth.tasks.eegbci import ( + EEGBCI_LABELS, + label_family_for_run, + numeric_label_for_task, + run_type_for_run, + task_label_for_event, +) + + +class TestEEGBCIHelpers(unittest.TestCase): + def test_run_type_for_run(self): + self.assertEqual(run_type_for_run(3), "motor_execution_left_right") + self.assertEqual(run_type_for_run(4), "motor_imagery_left_right") + self.assertEqual(run_type_for_run(5), "motor_execution_fists_feet") + self.assertEqual(run_type_for_run(6), "motor_imagery_fists_feet") + self.assertEqual(run_type_for_run(14), "motor_imagery_fists_feet") + + def test_task_label_for_event_is_run_aware(self): + self.assertEqual(task_label_for_event(3, "T0"), "rest") + self.assertEqual(task_label_for_event(3, "T1"), "execute_left_fist") + self.assertEqual(task_label_for_event(3, "T2"), "execute_right_fist") + self.assertEqual(task_label_for_event(4, "T1"), "imagine_left_fist") + self.assertEqual(task_label_for_event(4, "T2"), "imagine_right_fist") + self.assertEqual(task_label_for_event(5, "T1"), "execute_both_fists") + self.assertEqual(task_label_for_event(5, "T2"), "execute_both_feet") + self.assertEqual(task_label_for_event(6, "T1"), "imagine_both_fists") + self.assertEqual(task_label_for_event(6, "T2"), "imagine_both_feet") + + def test_label_family_and_numeric_labels(self): + self.assertEqual(label_family_for_run(3), "motor_execution") + self.assertEqual(label_family_for_run(4), "motor_imagery") + self.assertEqual(numeric_label_for_task("rest"), 0) + self.assertEqual(numeric_label_for_task("execute_left_fist"), 1) + self.assertEqual(numeric_label_for_task("imagine_both_feet"), 8) + + def test_invalid_run_and_event_raise_clear_errors(self): + with self.assertRaisesRegex(ValueError, "Unsupported EEGBCI run"): + run_type_for_run(2) + with self.assertRaisesRegex(ValueError, "Unsupported EEGBCI event"): + task_label_for_event(3, "BAD") + + def test_select_eegbci_channels_compat16(self): + from pyhealth.tasks.eegbci import EEGBCI_COMPAT_CHANNELS, select_eegbci_channels + + ch_names = list(EEGBCI_COMPAT_CHANNELS) + ["EXTRA"] + data = np.arange(len(ch_names) * 100, dtype=float).reshape(len(ch_names), 100) + selected, selected_names = select_eegbci_channels(data, ch_names, "compat16") + self.assertEqual(selected.shape, (16, 100)) + self.assertEqual(selected_names, list(EEGBCI_COMPAT_CHANNELS)) + np.testing.assert_allclose(selected[0], data[0]) + + def test_select_eegbci_channels_all(self): + from pyhealth.tasks.eegbci import select_eegbci_channels + + data = np.ones((64, 50)) + ch_names = [f"CH{i}" for i in range(64)] + selected, selected_names = select_eegbci_channels(data, ch_names, "all") + self.assertEqual(selected.shape, (64, 50)) + self.assertEqual(selected_names, ch_names) + + def test_select_eegbci_channels_missing_channel_raises(self): + from pyhealth.tasks.eegbci import select_eegbci_channels + + with self.assertRaisesRegex(ValueError, "Missing EEGBCI channels"): + select_eegbci_channels(np.ones((2, 20)), ["C3", "C4"], "compat16") + + def test_normalize_signal_95th_percentile(self): + from pyhealth.tasks.eegbci import normalize_signal + + signal = np.array([[0.0, 1.0, 2.0, 100.0], [0.0, -2.0, 2.0, 4.0]]) + normalized = normalize_signal(signal, "95th_percentile") + self.assertEqual(normalized.shape, signal.shape) + self.assertLess(np.max(np.abs(normalized[0])), 2.0) + + def test_compute_band_powers_detects_alpha_sinusoid(self): + from pyhealth.tasks.eegbci import compute_band_powers + + sfreq = 200.0 + times = np.arange(0, 2, 1 / sfreq) + alpha = np.sin(2 * np.pi * 10 * times) + data = np.stack([alpha, alpha]) + features = compute_band_powers(data, sfreq) + self.assertEqual(features["dominant_band"], "alpha") + self.assertGreater(features["alpha_relative"], 0.5) + self.assertGreater(features["alpha_beta_ratio"], 1.0) + + def test_compute_band_powers_detects_beta_sinusoid(self): + from pyhealth.tasks.eegbci import compute_band_powers + + sfreq = 200.0 + times = np.arange(0, 2, 1 / sfreq) + beta = np.sin(2 * np.pi * 20 * times) + data = np.stack([beta, beta]) + features = compute_band_powers(data, sfreq) + self.assertEqual(features["dominant_band"], "beta") + self.assertGreater(features["beta_relative"], 0.5) + + def test_interpret_band_profile_returns_cautious_metadata(self): + from pyhealth.tasks.eegbci import interpret_band_profile + + interpretation = interpret_band_profile( + { + "dominant_band": "alpha", + "alpha_relative": 0.65, + "beta_relative": 0.10, + "theta_relative": 0.10, + "gamma_relative": 0.05, + "alpha_beta_ratio": 6.5, + "theta_beta_ratio": 1.0, + } + ) + self.assertEqual(interpretation["brain_state_hypothesis"], "relaxed_or_idle") + self.assertIn(interpretation["confidence"], {"low", "medium", "high"}) + self.assertIn("consistent with", interpretation["interpretation"]) + self.assertNotIn( + "This is exploratory signal metadata", interpretation["interpretation"] + ) + self.assertNotIn("clinical diagnosis", interpretation["interpretation"]) + + +from pyhealth.datasets.eegbci import EEGBCIDataset + + +class TestEEGBCIDataset(unittest.TestCase): + def _set_metadata_identity(self, ds): + ds.selection_key = ds._build_selection_key() + ds.metadata_file_name = ds._metadata_file_name() + + def test_prepare_metadata_with_existing_files(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + edf = root / "files" / "eegmmidb" / "1.0.0" / "S001" / "S001R03.edf" + edf.parent.mkdir(parents=True) + edf.write_bytes(b"") + + ds = EEGBCIDataset.__new__(EEGBCIDataset) + ds.root = str(root) + ds.subjects = [1] + ds.runs = [3] + ds.download = False + self._set_metadata_identity(ds) + ds.prepare_metadata() + + csv_path = root / ds.metadata_file_name + self.assertTrue(csv_path.exists()) + df = pd.read_csv(csv_path) + self.assertEqual(len(df), 1) + self.assertEqual(df.loc[0, "patient_id"], "S001") + self.assertEqual(df.loc[0, "record_id"], "R03") + self.assertEqual(df.loc[0, "subject_id"], 1) + self.assertEqual(df.loc[0, "run"], 3) + self.assertEqual(df.loc[0, "run_type"], "motor_execution_left_right") + self.assertEqual(df.loc[0, "source"], "physionet_eegbci") + + def test_selection_inputs_are_normalized_for_stable_identity(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + for subject, run in [(1, 3), (1, 4), (2, 3), (2, 4)]: + edf = ( + root + / "files" + / "eegmmidb" + / "1.0.0" + / f"S{subject:03d}" + / f"S{subject:03d}R{run:02d}.edf" + ) + edf.parent.mkdir(parents=True, exist_ok=True) + edf.write_bytes(b"") + + first = EEGBCIDataset( + root=str(root), + subjects=[2, 1, 1], + runs=[4, 3, 4], + download=False, + ) + second = EEGBCIDataset( + root=str(root), + subjects=[1, 2], + runs=[3, 4], + download=False, + ) + + self.assertEqual(first.subjects, [1, 2]) + self.assertEqual(first.runs, [3, 4]) + self.assertEqual(first.selection_key, second.selection_key) + self.assertEqual(first.dataset_name, second.dataset_name) + + def test_prepare_metadata_uses_selection_specific_files(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + first = root / "files" / "eegmmidb" / "1.0.0" / "S001" / "S001R03.edf" + second = root / "files" / "eegmmidb" / "1.0.0" / "S002" / "S002R04.edf" + first.parent.mkdir(parents=True) + second.parent.mkdir(parents=True) + first.write_bytes(b"") + second.write_bytes(b"") + + ds_first = EEGBCIDataset( + root=str(root), subjects=[1], runs=[3], download=False + ) + ds_second = EEGBCIDataset( + root=str(root), subjects=[2], runs=[4], download=False + ) + + first_csv = root / ds_first.metadata_file_name + second_csv = root / ds_second.metadata_file_name + self.assertNotEqual(first_csv, second_csv) + self.assertTrue(first_csv.exists()) + self.assertTrue(second_csv.exists()) + + first_df = pd.read_csv(first_csv) + second_df = pd.read_csv(second_csv) + self.assertEqual(first_df.loc[0, "subject_id"], 1) + self.assertEqual(first_df.loc[0, "run"], 3) + self.assertEqual(second_df.loc[0, "subject_id"], 2) + self.assertEqual(second_df.loc[0, "run"], 4) + + def test_find_local_edf_checks_canonical_mne_path_first(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + canonical = root / "files" / "eegmmidb" / "1.0.0" / "S001" / "S001R03.edf" + fallback = root / "other" / "S001R03.edf" + canonical.parent.mkdir(parents=True) + fallback.parent.mkdir(parents=True) + canonical.write_bytes(b"") + fallback.write_bytes(b"") + + ds = EEGBCIDataset.__new__(EEGBCIDataset) + ds.root = str(root) + + self.assertEqual(ds._find_local_edf(1, 3), canonical) + + def test_prepare_metadata_download_uses_mne_loader(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + fake_path = root / "S001R04.edf" + fake_path.write_bytes(b"") + ds = EEGBCIDataset.__new__(EEGBCIDataset) + ds.root = str(root) + ds.subjects = [1] + ds.runs = [4] + ds.download = True + self._set_metadata_identity(ds) + + with patch( + "pyhealth.datasets.eegbci.mne.datasets.eegbci.load_data", + return_value=[str(fake_path)], + ) as load_data: + ds.prepare_metadata() + + load_data.assert_called_once_with(1, [4], path=str(root), update_path=False) + df = pd.read_csv(root / ds.metadata_file_name) + self.assertEqual(df.loc[0, "record_id"], "R04") + self.assertEqual(df.loc[0, "run_type"], "motor_imagery_left_right") + + def test_prepare_metadata_missing_local_file_raises(self): + with tempfile.TemporaryDirectory() as tmp: + ds = EEGBCIDataset.__new__(EEGBCIDataset) + ds.root = tmp + ds.subjects = [1] + ds.runs = [3] + ds.download = False + self._set_metadata_identity(ds) + with self.assertRaisesRegex(FileNotFoundError, "download=True"): + ds.prepare_metadata() + + def test_default_task_returns_pattern_discovery(self): + from pyhealth.tasks.eegbci import EEGBCIPatternDiscovery + + ds = EEGBCIDataset.__new__(EEGBCIDataset) + self.assertIsInstance(ds.default_task, EEGBCIPatternDiscovery) + + def test_dataset_set_task_offline_integration(self): + import mne + from pyhealth.tasks.eegbci import EEGBCI_COMPAT_CHANNELS, EEGMotorImageryEEGBCI + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + edf = root / "files" / "eegmmidb" / "1.0.0" / "S001" / "S001R03.edf" + edf.parent.mkdir(parents=True) + edf.write_bytes(b"") + sfreq = 200.0 + raw = mne.io.RawArray( + np.ones((16, int(sfreq * 2))), + mne.create_info( + list(EEGBCI_COMPAT_CHANNELS), sfreq=sfreq, ch_types=["eeg"] * 16 + ), + verbose="error", + ) + raw.set_annotations( + mne.Annotations(onset=[0.0], duration=[2.0], description=["T1"]) + ) + dataset = EEGBCIDataset( + root=str(root), + subjects=[1], + runs=[3], + download=False, + cache_dir=root / "cache", + ) + + with patch("pyhealth.tasks.eegbci.mne.io.read_raw_edf", return_value=raw): + sample_dataset = dataset.set_task( + EEGMotorImageryEEGBCI( + compute_stft=False, resample_rate=None, bandpass_filter=None + ), + num_workers=1, + ) + + self.assertEqual(len(sample_dataset), 1) + sample = sample_dataset[0] + self.assertEqual(sample["task_label"], "execute_left_fist") + self.assertEqual(sample["eegbci_label"], 1) + self.assertEqual(tuple(sample["signal"].shape), (16, 400)) + + +from pyhealth.tasks.eegbci import EEGBCIPatternDiscovery, EEGMotorImageryEEGBCI + + +@dataclass +class _EEGBCIEvent: + signal_file: str + record_id: str = "R03" + subject_id: int = 1 + run: int = 3 + run_type: str = "motor_execution_left_right" + source: str = "physionet_eegbci" + + +class _EEGBCIPatient: + def __init__(self, patient_id: str, events: List[_EEGBCIEvent]): + self.patient_id = patient_id + self._events = events + + def get_events(self, event_type=None) -> List[_EEGBCIEvent]: + if event_type not in (None, "records"): + return [] + return self._events + + +class TestEEGBCITasks(unittest.TestCase): + def test_task_schema_attributes(self): + task = EEGMotorImageryEEGBCI() + self.assertEqual(task.task_name, "EEGBCI_motor_imagery") + self.assertEqual(task.input_schema, {"signal": "tensor", "stft": "tensor"}) + self.assertEqual(task.output_schema, {"label": "multiclass"}) + + def test_task_schema_without_stft(self): + task = EEGMotorImageryEEGBCI(compute_stft=False) + self.assertEqual(task.input_schema, {"signal": "tensor"}) + + def test_pattern_discovery_schema_attributes(self): + task = EEGBCIPatternDiscovery(compute_stft=False) + self.assertEqual(task.task_name, "EEGBCI_pattern_discovery") + self.assertEqual(task.input_schema, {"signal": "tensor"}) + + def test_iter_annotation_windows_uses_full_2s_windows(self): + import mne + from pyhealth.tasks.eegbci import iter_annotation_windows + + sfreq = 200.0 + raw = mne.io.RawArray( + np.zeros((2, int(sfreq * 6))), + mne.create_info(["C3", "C4"], sfreq=sfreq, ch_types=["eeg", "eeg"]), + verbose="error", + ) + raw.set_annotations( + mne.Annotations(onset=[0.5, 2.0], duration=[1.0, 3.0], description=["T0", "T1"]) + ) + windows = iter_annotation_windows(raw, run=3, window_size=2.0) + self.assertEqual(len(windows), 1) + self.assertEqual(windows[0]["event_code"], "T1") + self.assertEqual(windows[0]["task_label"], "execute_left_fist") + self.assertEqual(windows[0]["start_sample"], 400) + self.assertEqual(windows[0]["end_sample"], 800) + + def test_motor_imagery_task_returns_samples_from_raw(self): + import mne + + sfreq = 200.0 + raw = mne.io.RawArray( + np.ones((16, int(sfreq * 5))), + mne.create_info( + list( + __import__( + "pyhealth.tasks.eegbci", fromlist=["EEGBCI_COMPAT_CHANNELS"] + ).EEGBCI_COMPAT_CHANNELS + ), + sfreq=sfreq, + ch_types=["eeg"] * 16, + ), + verbose="error", + ) + raw.set_annotations(mne.Annotations(onset=[0.0], duration=[2.0], description=["T1"])) + patient = _EEGBCIPatient("S001", [_EEGBCIEvent(signal_file="dummy.edf")]) + task = EEGMotorImageryEEGBCI(compute_stft=False, resample_rate=None, bandpass_filter=None) + + with patch("pyhealth.tasks.eegbci.mne.io.read_raw_edf", return_value=raw): + samples = task(patient) + + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertEqual(sample["patient_id"], "S001") + self.assertEqual(sample["record_id"], "R03") + self.assertEqual(sample["event_code"], "T1") + self.assertEqual(sample["task_label"], "execute_left_fist") + self.assertEqual(sample["label"], 1) + self.assertEqual(sample["eegbci_label"], 1) + self.assertEqual(tuple(sample["signal"].shape), (16, 400)) + + def test_stft_uses_current_sample_rate(self): + import mne + from pyhealth.tasks.eegbci import EEGBCI_COMPAT_CHANNELS + + sfreq = 100.0 + raw = mne.io.RawArray( + np.ones((16, int(sfreq * 2))), + mne.create_info( + list(EEGBCI_COMPAT_CHANNELS), sfreq=sfreq, ch_types=["eeg"] * 16 + ), + verbose="error", + ) + raw.set_annotations( + mne.Annotations(onset=[0.0], duration=[2.0], description=["T1"]) + ) + patient = _EEGBCIPatient("S001", [_EEGBCIEvent(signal_file="dummy.edf")]) + task = EEGMotorImageryEEGBCI(resample_rate=None, bandpass_filter=None) + + with patch("pyhealth.tasks.eegbci.mne.io.read_raw_edf", return_value=raw): + with patch( + "pyhealth.models.tfm_tokenizer.get_stft_torch", + return_value=torch.zeros((1, 16, 50, 1)), + ) as get_stft: + samples = task(patient) + + self.assertEqual(len(samples), 1) + self.assertEqual(get_stft.call_args.kwargs["resampling_rate"], 100) + + def test_pattern_discovery_adds_bandpower_metadata(self): + import mne + from pyhealth.tasks.eegbci import EEGBCI_COMPAT_CHANNELS + + sfreq = 200.0 + times = np.arange(0, 2, 1 / sfreq) + alpha = np.sin(2 * np.pi * 10 * times) + raw = mne.io.RawArray( + np.tile(alpha, (16, 1)), + mne.create_info(list(EEGBCI_COMPAT_CHANNELS), sfreq=sfreq, ch_types=["eeg"] * 16), + verbose="error", + ) + raw.set_annotations(mne.Annotations(onset=[0.0], duration=[2.0], description=["T0"])) + patient = _EEGBCIPatient("S001", [_EEGBCIEvent(signal_file="dummy.edf")]) + task = EEGBCIPatternDiscovery(compute_stft=False, resample_rate=None, bandpass_filter=None) + + with patch("pyhealth.tasks.eegbci.mne.io.read_raw_edf", return_value=raw): + samples = task(patient) + + self.assertEqual(len(samples), 1) + sample = samples[0] + self.assertEqual(sample["bandpower"]["dominant_band"], "alpha") + self.assertEqual(sample["brain_state_hypothesis"], "relaxed_or_idle") + self.assertIn("interpretation", sample) + + +class TestEEGBCIMomentReportHelpers(unittest.TestCase): + def _moment_row(self, **overrides): + row = { + "patient_id": "S001", + "record_id": "R03", + "subject_id": 1, + "run": 3, + "run_type": "motor_execution_left_right", + "trial_id": "S001_R03_T0_0", + "event_code": "T0", + "task_label": "rest", + "label_family": "rest", + "label": 0, + "eegbci_label": 0, + "model_label": 0, + "start_time": 0.0, + "end_time": 2.0, + "dominant_band": "alpha", + "delta_relative": 0.05, + "theta_relative": 0.10, + "alpha_relative": 0.55, + "beta_relative": 0.20, + "gamma_relative": 0.10, + "alpha_beta_ratio": 2.75, + "theta_beta_ratio": 0.50, + } + row.update(overrides) + return row + + def _sample(self, **overrides): + sample = { + "patient_id": "S001", + "record_id": "R03", + "subject_id": 1, + "run": 3, + "run_type": "motor_execution_left_right", + "trial_id": "S001_R03_T0_0", + "event_code": "T0", + "task_label": "rest", + "label_family": "rest", + "label": 0, + "eegbci_label": 0, + "start_time": 0.0, + "end_time": 2.0, + "brain_state_hypothesis": "relaxed_or_idle", + "confidence": "medium", + "quality_flags": "", + "interpretation": "Alpha-dominant profile.", + "bandpower": { + "dominant_band": "alpha", + "alpha_beta_ratio": 2.75, + "theta_beta_ratio": 0.50, + "delta_power": 0.05, + "theta_power": 0.10, + "alpha_power": 0.55, + "beta_power": 0.20, + "gamma_power": 0.10, + "delta_relative": 0.05, + "theta_relative": 0.10, + "alpha_relative": 0.55, + "beta_relative": 0.20, + "gamma_relative": 0.10, + }, + } + sample.update(overrides) + return sample + + def test_analysis_version_constant(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ANALYSIS_VERSION + + self.assertEqual(ANALYSIS_VERSION, "eegbci_pattern_moment_report_v1") + + def test_parse_int_list_strips_whitespace(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import parse_int_list + + self.assertEqual(parse_int_list("1, 2, 4-6"), [1, 2, 4, 5, 6]) + + def test_parse_int_list_rejects_descending_ranges(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import parse_int_list + + with self.assertRaisesRegex(ValueError, "Range start must be <= range end"): + parse_int_list("5-3") + + def test_build_rest_baselines_uses_rest_rows_only(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import build_rest_baselines + + rows = [ + self._moment_row(task_label="rest", subject_id=1, run=3, alpha_relative=0.50), + self._moment_row( + task_label="execute_left_fist", subject_id=1, run=3, alpha_relative=0.90 + ), + self._moment_row(task_label="rest", subject_id=1, run=4, alpha_relative=0.70), + ] + + baselines = build_rest_baselines(rows) + + self.assertAlmostEqual( + baselines["same_subject_run"][(1, 3)]["alpha_relative"], 0.50 + ) + self.assertAlmostEqual( + baselines["same_subject_all_runs"][1]["alpha_relative"], 0.60 + ) + self.assertAlmostEqual(baselines["global_rest"]["alpha_relative"], 0.60) + + def test_build_rest_baselines_handles_no_rest_rows(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import build_rest_baselines + + rows = [ + self._moment_row( + task_label="execute_left_fist", label_family="motor_execution" + ) + ] + + baselines = build_rest_baselines(rows) + + self.assertEqual(baselines["same_subject_run"], {}) + self.assertEqual(baselines["same_subject_all_runs"], {}) + self.assertIsNone(baselines["global_rest"]) + + def test_render_summary_reports_rest_baseline_source_rows(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + rows = [ + self._moment_row(task_label="rest"), + self._moment_row( + task_label="execute_left_fist", label_family="motor_execution" + ), + self._moment_row(task_label="rest", run=4), + ] + + summary = render_summary( + rows, + { + "subjects": [1], + "runs": [3, 4], + "max_windows": None, + "baseline_row_count": 2, + "output_was_capped": False, + }, + ) + + self.assertIn("- Baseline source rows: 2", summary) + + def test_annotate_rest_fallback_scopes(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + ) + + rows = [ + self._moment_row(task_label="rest", subject_id=1, run=3, alpha_relative=0.50), + self._moment_row(task_label="rest", subject_id=1, run=4, alpha_relative=0.70), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + subject_id=1, + run=3, + alpha_relative=0.80, + ), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + subject_id=1, + run=5, + alpha_relative=0.80, + ), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + subject_id=2, + run=8, + alpha_relative=0.80, + ), + ] + + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + + self.assertEqual(annotated[2]["rest_reference_scope"], "same_subject_run") + self.assertEqual(annotated[3]["rest_reference_scope"], "same_subject_all_runs") + self.assertEqual(annotated[4]["rest_reference_scope"], "global_rest") + + def test_derive_state_hypothesis_detects_profiles(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import derive_state_hypothesis + + cases = [ + ( + self._moment_row( + alpha_relative=0.60, + beta_relative=0.12, + gamma_relative=0.05, + alpha_beta_ratio=5.0, + ), + "idle_alpha_profile", + ), + ( + self._moment_row( + alpha_relative=0.12, + beta_relative=0.48, + gamma_relative=0.16, + alpha_beta_ratio=0.25, + ), + "sensorimotor_engagement_profile", + ), + ( + self._moment_row( + delta_relative=0.42, + theta_relative=0.36, + alpha_relative=0.08, + beta_relative=0.08, + ), + "slow_wave_dominant_pattern", + ), + ( + self._moment_row( + gamma_relative=0.48, alpha_relative=0.10, beta_relative=0.12 + ), + "possible_artifact_profile", + ), + ( + self._moment_row( + delta_relative=0.18, + theta_relative=0.20, + alpha_relative=0.22, + beta_relative=0.21, + gamma_relative=0.19, + alpha_beta_ratio=1.05, + ), + "mixed_ambiguous_profile", + ), + ] + + for row, expected in cases: + with self.subTest(expected=expected): + result = derive_state_hypothesis(row) + self.assertEqual(result["state_hypothesis"], expected) + self.assertIn(result["state_confidence"], {"low", "medium", "high"}) + self.assertGreaterEqual(result["evidence_score"], 0.0) + self.assertLessEqual(result["evidence_score"], 1.0) + self.assertIn("alpha=", result["evidence_summary"]) + + def test_state_confidence_requires_margin(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + STATE_CONFIDENCE_RANK, + derive_state_hypothesis, + ) + + clear = derive_state_hypothesis( + self._moment_row( + alpha_relative=0.70, + beta_relative=0.10, + gamma_relative=0.04, + alpha_beta_ratio=6.0, + ) + ) + weaker = derive_state_hypothesis( + self._moment_row( + alpha_relative=0.40, + beta_relative=0.22, + gamma_relative=0.10, + alpha_beta_ratio=2.0, + ) + ) + + self.assertEqual(clear["state_hypothesis"], weaker["state_hypothesis"]) + self.assertGreater( + STATE_CONFIDENCE_RANK[clear["state_confidence"]], + STATE_CONFIDENCE_RANK[weaker["state_confidence"]], + ) + + def test_task_state_relation_table_is_deterministic(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + derive_task_state_relation, + ) + + cases = [ + ("rest", "rest", "idle_alpha_profile", "supports_label"), + ("rest", "rest", "mixed_ambiguous_profile", "ambiguous"), + ("rest", "rest", "possible_artifact_profile", "not_applicable"), + ( + "execute_left_fist", + "motor_execution", + "sensorimotor_engagement_profile", + "supports_label", + ), + ( + "imagine_left_fist", + "motor_imagery", + "sensorimotor_engagement_profile", + "adds_detail", + ), + ("execute_left_fist", "motor_execution", "idle_alpha_profile", "disagrees"), + ( + "imagine_left_fist", + "motor_imagery", + "slow_wave_dominant_pattern", + "adds_detail", + ), + ] + + for task_label, label_family, state, expected in cases: + with self.subTest(state=state, label_family=label_family): + result = derive_task_state_relation( + self._moment_row( + task_label=task_label, + label_family=label_family, + state_hypothesis=state, + ) + ) + self.assertEqual(result["task_state_relation"], expected) + self.assertIn(result["task_state_confidence"], {"low", "medium", "high"}) + self.assertGreater(len(result["task_state_rationale"]), 20) + + def test_quality_booleans_are_parseable(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import derive_quality_columns + + flags = derive_quality_columns( + self._moment_row( + state_hypothesis="possible_artifact_profile", + state_confidence="low", + quality_flags="low_confidence; high_gamma", + ) + ) + + self.assertTrue(flags["is_low_confidence"]) + self.assertTrue(flags["is_possible_artifact"]) + self.assertFalse(flags["is_mixed_or_ambiguous"]) + + def test_quality_booleans_do_not_depend_on_string_parsing_only(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import derive_quality_columns + + flags = derive_quality_columns( + self._moment_row( + state_hypothesis="mixed_ambiguous_profile", + state_confidence="medium", + quality_flags="", + ) + ) + + self.assertTrue(flags["is_mixed_or_ambiguous"]) + + def test_quality_booleans_do_not_conflate_legacy_low_confidence(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import derive_quality_columns + + flags = derive_quality_columns( + self._moment_row( + state_hypothesis="idle_alpha_profile", + state_confidence="medium", + quality_flags="low_confidence", + ) + ) + + self.assertFalse(flags["is_low_confidence"]) + + def test_annotate_moment_rows_adds_required_fields(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + ANALYSIS_VERSION, + MOMENT_REPORT_COLUMNS, + annotate_moment_rows, + build_rest_baselines, + ) + + rows = [ + self._moment_row(task_label="rest", alpha_relative=0.50, beta_relative=0.20), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + alpha_relative=0.20, + beta_relative=0.45, + ), + ] + + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + + for annotated_row in annotated: + for column in MOMENT_REPORT_COLUMNS: + self.assertIn(column, annotated_row) + row = annotated[1] + self.assertEqual(row["analysis_version"], ANALYSIS_VERSION) + self.assertIn( + row["state_hypothesis"], + { + "idle_alpha_profile", + "sensorimotor_engagement_profile", + "slow_wave_dominant_pattern", + "possible_artifact_profile", + "mixed_ambiguous_profile", + }, + ) + self.assertIn("rest_alpha_relative_delta", row) + self.assertAlmostEqual(row["rest_alpha_relative_delta"], -0.30) + self.assertIn("task_state_relation", row) + self.assertIn("task_state_rationale", row) + self.assertIn("is_low_confidence", row) + + def test_annotate_moment_rows_marks_unavailable_rest(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + ) + + rows = [ + self._moment_row( + task_label="execute_left_fist", label_family="motor_execution" + ) + ] + + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + + self.assertEqual(annotated[0]["rest_reference_scope"], "unavailable") + for band in ("delta", "theta", "alpha", "beta", "gamma"): + self.assertEqual(annotated[0][f"rest_{band}_relative_delta"], "") + + def test_rest_delta_values_are_band_specific(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + ) + + rows = [ + self._moment_row( + task_label="rest", + delta_relative=0.10, + theta_relative=0.20, + alpha_relative=0.30, + beta_relative=0.25, + gamma_relative=0.15, + ), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + delta_relative=0.15, + theta_relative=0.18, + alpha_relative=0.25, + beta_relative=0.35, + gamma_relative=0.07, + ), + ] + + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + + self.assertAlmostEqual(annotated[1]["rest_delta_relative_delta"], 0.05) + self.assertAlmostEqual(annotated[1]["rest_theta_relative_delta"], -0.02) + self.assertAlmostEqual(annotated[1]["rest_alpha_relative_delta"], -0.05) + self.assertAlmostEqual(annotated[1]["rest_beta_relative_delta"], 0.10) + self.assertAlmostEqual(annotated[1]["rest_gamma_relative_delta"], -0.08) + + def test_annotate_moment_rows_adds_report_interpretation(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + ) + + row = self._moment_row() + + annotated = annotate_moment_rows([row], build_rest_baselines([row])) + + self.assertIn("consistent with", annotated[0]["interpretation"]) + self.assertIn("task label", annotated[0]["interpretation"]) + self.assertIn(annotated[0]["state_hypothesis"], annotated[0]["interpretation"]) + + def test_annotate_moment_rows_does_not_mutate_input_rows(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + ) + + rows = [self._moment_row()] + original = [dict(row) for row in rows] + + annotate_moment_rows(rows, build_rest_baselines(rows)) + + self.assertEqual(rows, original) + + def test_select_representative_windows_is_deterministic(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + select_representative_windows, + ) + + rows = [ + self._moment_row( + subject_id=2, + run=4, + start_time=6.0, + state_hypothesis="idle_alpha_profile", + state_confidence="medium", + evidence_score=0.80, + ), + self._moment_row( + subject_id=1, + run=3, + start_time=4.0, + state_hypothesis="idle_alpha_profile", + state_confidence="medium", + evidence_score=0.80, + ), + self._moment_row( + subject_id=1, + run=3, + start_time=8.0, + state_hypothesis="sensorimotor_engagement_profile", + state_confidence="high", + evidence_score=0.90, + ), + self._moment_row( + subject_id=1, + run=3, + start_time=10.0, + state_hypothesis="mixed_ambiguous_profile", + state_confidence="low", + evidence_score=0.12, + ), + self._moment_row( + subject_id=1, + run=3, + start_time=12.0, + state_hypothesis="idle_alpha_profile", + task_state_relation="disagrees", + state_confidence="medium", + evidence_score=0.70, + ), + ] + + selected = select_representative_windows(rows) + + self.assertEqual(selected["cards"]["strongest_idle_like"]["subject_id"], 1) + self.assertEqual( + selected["cards"]["strongest_motor_engaged"]["state_hypothesis"], + "sensorimotor_engagement_profile", + ) + self.assertEqual(selected["cards"]["most_ambiguous"]["start_time"], 10.0) + self.assertEqual( + selected["cards"]["strongest_task_state_disagreement"][ + "task_state_relation" + ], + "disagrees", + ) + self.assertIn("strongest_artifact_like", selected["absent"]) + + def test_select_representative_windows_picks_lowest_evidence_ambiguous(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + select_representative_windows, + ) + + rows = [ + self._moment_row( + subject_id=2, + run=3, + start_time=4.0, + state_hypothesis="mixed_ambiguous_profile", + state_confidence="low", + evidence_score=0.20, + ), + self._moment_row( + subject_id=1, + run=3, + start_time=8.0, + state_hypothesis="mixed_ambiguous_profile", + state_confidence="low", + evidence_score=0.10, + ), + ] + + selected = select_representative_windows(rows) + + self.assertEqual(selected["cards"]["most_ambiguous"]["subject_id"], 1) + + def test_select_representative_windows_picks_strongest_disagreement(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + select_representative_windows, + ) + + rows = [ + self._moment_row( + subject_id=1, + run=3, + start_time=4.0, + state_hypothesis="idle_alpha_profile", + task_state_relation="disagrees", + state_confidence="medium", + evidence_score=0.50, + ), + self._moment_row( + subject_id=2, + run=3, + start_time=6.0, + state_hypothesis="idle_alpha_profile", + task_state_relation="disagrees", + state_confidence="medium", + evidence_score=0.80, + ), + ] + + selected = select_representative_windows(rows) + + self.assertEqual( + selected["cards"]["strongest_task_state_disagreement"]["subject_id"], 2 + ) + + def test_render_summary_contains_required_sections_and_limitations(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + ANALYSIS_VERSION, + annotate_moment_rows, + build_rest_baselines, + render_summary, + ) + + rows = [ + self._moment_row( + task_label="execute_left_fist", label_family="motor_execution" + ) + ] + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + summary = render_summary( + annotated, + { + "subjects": [1], + "runs": [3], + "max_windows": 1, + "baseline_row_count": 1, + "output_was_capped": True, + }, + ) + + self.assertIn(ANALYSIS_VERSION, summary.splitlines()[2]) + for heading in [ + "## Executive Result", + "## Run Configuration", + "## Window Coverage", + "## Moment-State Summary", + "## Task Label x State Matrix", + "## Rest-Normalized Bandpower Summary", + "## Confidence and Quality Audit", + "## Representative Windows", + "## Limitations", + "## Next Checks", + ]: + self.assertIn(heading, summary) + self.assertIn("No rest baseline was available", summary) + self.assertIn("Output was capped by `--max-windows`", summary) + self.assertNotIn( + "Brain-state hypotheses are exploratory signal metadata", + summary.splitlines()[2], + ) + + def test_render_summary_handles_empty_rows(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + summary = render_summary( + [], + { + "subjects": [1], + "runs": [3], + "max_windows": 0, + "baseline_row_count": 0, + "output_was_capped": True, + }, + ) + + self.assertIn("No windows were produced", summary) + self.assertIn("## Limitations", summary) + + def test_render_summary_reports_all_low_confidence_and_same_state(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + rows = [ + self._moment_row( + state_hypothesis="mixed_ambiguous_profile", + state_confidence="low", + evidence_score=0.10, + task_state_relation="ambiguous", + task_state_confidence="low", + rest_reference_scope="unavailable", + is_low_confidence=True, + is_possible_artifact=False, + is_mixed_or_ambiguous=True, + ), + self._moment_row( + start_time=2.0, + state_hypothesis="mixed_ambiguous_profile", + state_confidence="low", + evidence_score=0.12, + task_state_relation="ambiguous", + task_state_confidence="low", + rest_reference_scope="unavailable", + is_low_confidence=True, + is_possible_artifact=False, + is_mixed_or_ambiguous=True, + ), + ] + + summary = render_summary( + rows, + { + "subjects": [1], + "runs": [3], + "max_windows": None, + "baseline_row_count": 2, + "output_was_capped": False, + }, + ) + + self.assertIn("Every window is low confidence", summary) + self.assertIn("Every window maps to the same state", summary) + + def test_render_summary_reports_task_state_matrix(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + rows = [ + self._moment_row( + task_label="rest", + state_hypothesis="idle_alpha_profile", + state_confidence="medium", + evidence_score=0.60, + task_state_relation="supports_label", + task_state_confidence="medium", + rest_reference_scope="same_subject_run", + ), + self._moment_row( + task_label="execute_left_fist", + label_family="motor_execution", + state_hypothesis="sensorimotor_engagement_profile", + state_confidence="medium", + evidence_score=0.70, + task_state_relation="supports_label", + task_state_confidence="medium", + rest_reference_scope="same_subject_run", + ), + ] + + summary = render_summary( + rows, + { + "subjects": [1], + "runs": [3], + "max_windows": None, + "baseline_row_count": 2, + "output_was_capped": False, + }, + ) + + self.assertIn("rest x idle_alpha_profile: 1", summary) + self.assertIn( + "execute_left_fist x sensorimotor_engagement_profile: 1", summary + ) + + def test_render_summary_includes_representative_window_details(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + row = self._moment_row( + state_hypothesis="idle_alpha_profile", + state_confidence="medium", + evidence_score=0.75, + task_state_relation="supports_label", + task_state_confidence="medium", + task_state_rationale="The idle-like alpha profile is consistent with rest.", + rest_reference_scope="same_subject_run", + rest_delta_relative_delta=0.01, + rest_theta_relative_delta=0.02, + rest_alpha_relative_delta=0.03, + rest_beta_relative_delta=-0.02, + rest_gamma_relative_delta=-0.01, + is_low_confidence=False, + is_possible_artifact=False, + is_mixed_or_ambiguous=False, + ) + + summary = render_summary( + [row], + { + "subjects": [1], + "runs": [3], + "max_windows": None, + "baseline_row_count": 1, + "output_was_capped": False, + }, + ) + + for text in [ + "Subject 1 run 3 trial S001_R03_T0_0", + "Task: rest from 0.0s to 2.0s", + "State: idle_alpha_profile", + "Dominant band: alpha", + "Rest deltas:", + "Task relation: supports_label", + "low_confidence=False", + "Rationale: The idle-like alpha profile", + ]: + self.assertIn(text, summary) + + def test_render_summary_moves_nonclinical_warning_to_limitations(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import render_summary + + summary = render_summary( + [self._moment_row(state_hypothesis="idle_alpha_profile")], + { + "subjects": [1], + "runs": [3], + "max_windows": None, + "baseline_row_count": 1, + "output_was_capped": False, + }, + ) + + opening = "\n".join(summary.splitlines()[:6]) + limitations = summary.split("## Limitations", 1)[1] + self.assertNotIn("clinical findings", opening) + self.assertIn("clinical findings", limitations) + + def test_summary_text_does_not_repeat_old_row_level_caveat(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + annotate_moment_rows, + build_rest_baselines, + render_summary, + ) + + rows = [ + self._moment_row( + interpretation="This is exploratory signal metadata, not a diagnosis." + ) + ] + annotated = annotate_moment_rows(rows, build_rest_baselines(rows)) + + summary = render_summary( + annotated, + { + "subjects": [1], + "runs": [3], + "max_windows": None, + "baseline_row_count": 1, + "output_was_capped": False, + }, + ) + + self.assertNotIn("This is exploratory signal metadata", summary) + + def test_moment_report_columns_are_declared(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import ( + MOMENT_REPORT_COLUMNS, + OUTPUT_COLUMNS, + ) + + for column in [ + "patient_id", + "task_label", + "alpha_relative", + "analysis_version", + "state_hypothesis", + "state_confidence", + "evidence_score", + "evidence_summary", + "rest_reference_scope", + "rest_alpha_relative_delta", + "task_state_relation", + "task_state_rationale", + "task_state_confidence", + "interpretation", + "is_low_confidence", + "is_possible_artifact", + "is_mixed_or_ambiguous", + ]: + self.assertIn(column, OUTPUT_COLUMNS) + self.assertIn("analysis_version", MOMENT_REPORT_COLUMNS) + + def test_output_columns_remove_legacy_task_fields(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import OUTPUT_COLUMNS + + for legacy_column in [ + "brain_state_hypothesis", + "confidence", + "quality_flags", + "legacy_brain_state_hypothesis", + "legacy_confidence", + "legacy_quality_flags", + "legacy_interpretation", + ]: + self.assertNotIn(legacy_column, OUTPUT_COLUMNS) + + def test_empty_dataframe_uses_output_columns(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import OUTPUT_COLUMNS + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "empty.csv" + pd.DataFrame([], columns=OUTPUT_COLUMNS).to_csv(path, index=False) + + df = pd.read_csv(path) + + self.assertEqual(len(df), 0) + self.assertEqual(list(df.columns), list(OUTPUT_COLUMNS)) + + def test_main_max_windows_zero_writes_empty_artifacts(self): + from examples.eeg.eegbci import eegbci_pattern_discovery as example + + class FakeDataset: + def __init__(self, *args, **kwargs): + pass + + def set_task(self, task): + return [self_sample] + + self_sample = self._sample() + with tempfile.TemporaryDirectory() as tmp: + argv = [ + "eegbci_pattern_discovery.py", + "--subjects", + "1", + "--runs", + "3", + "--max-windows", + "0", + "--output-dir", + tmp, + ] + with patch.object(sys, "argv", argv), patch.object( + example, "EEGBCIDataset", FakeDataset + ): + example.main() + + csv_path = Path(tmp) / "eegbci_pattern_windows.csv" + summary_path = Path(tmp) / "eegbci_pattern_summary.md" + df = pd.read_csv(csv_path) + summary = summary_path.read_text(encoding="utf-8") + + self.assertEqual(len(df), 0) + self.assertEqual(list(df.columns), list(example.OUTPUT_COLUMNS)) + self.assertIn("No windows were produced", summary) + self.assertIn("Output was capped by `--max-windows`", summary) + + def test_main_baseline_uses_uncapped_rows(self): + from examples.eeg.eegbci import eegbci_pattern_discovery as example + + first = self._sample( + task_label="execute_left_fist", + label_family="motor_execution", + alpha_beta_ratio=0.5, + bandpower={ + **self._sample()["bandpower"], + "dominant_band": "beta", + "alpha_relative": 0.20, + "beta_relative": 0.45, + "alpha_beta_ratio": 0.5, + }, + ) + rest = self._sample( + task_label="rest", + start_time=2.0, + bandpower={ + **self._sample()["bandpower"], + "alpha_relative": 0.50, + "beta_relative": 0.20, + }, + ) + + class FakeDataset: + def __init__(self, *args, **kwargs): + pass + + def set_task(self, task): + return [first, rest] + + with tempfile.TemporaryDirectory() as tmp: + argv = [ + "eegbci_pattern_discovery.py", + "--subjects", + "1", + "--runs", + "3", + "--max-windows", + "1", + "--output-dir", + tmp, + ] + with patch.object(sys, "argv", argv), patch.object( + example, "EEGBCIDataset", FakeDataset + ): + example.main() + + df = pd.read_csv(Path(tmp) / "eegbci_pattern_windows.csv") + + self.assertEqual(len(df), 1) + self.assertEqual(df.loc[0, "rest_reference_scope"], "same_subject_run") + self.assertAlmostEqual(df.loc[0, "rest_alpha_relative_delta"], -0.30) + + def test_main_writes_analysis_version_to_every_csv_row(self): + from examples.eeg.eegbci import eegbci_pattern_discovery as example + + samples = [self._sample(), self._sample(start_time=2.0, trial_id="second")] + + class FakeDataset: + def __init__(self, *args, **kwargs): + pass + + def set_task(self, task): + return samples + + with tempfile.TemporaryDirectory() as tmp: + argv = [ + "eegbci_pattern_discovery.py", + "--subjects", + "1", + "--runs", + "3", + "--output-dir", + tmp, + ] + with patch.object(sys, "argv", argv), patch.object( + example, "EEGBCIDataset", FakeDataset + ): + example.main() + + df = pd.read_csv(Path(tmp) / "eegbci_pattern_windows.csv") + + self.assertEqual(len(df), 2) + self.assertTrue((df["analysis_version"] == example.ANALYSIS_VERSION).all()) + + def test_parse_int_list_rejects_invalid_input_loudly(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import parse_int_list + + with self.assertRaises(ValueError): + parse_int_list("a") + with self.assertRaises(ValueError): + parse_int_list("3-a") + + def test_parse_int_list_accepts_ranges_and_singletons(self): + from examples.eeg.eegbci.eegbci_pattern_discovery import parse_int_list + + self.assertEqual(parse_int_list("1,3-5"), [1, 3, 4, 5]) + + +@unittest.skipUnless( + os.environ.get("PYHEALTH_RUN_REAL_EEGBCI") == "1", + "Set PYHEALTH_RUN_REAL_EEGBCI=1 to download and test real EEGBCI data.", +) +class TestEEGBCIRealDataSmoke(unittest.TestCase): + def test_real_eegbci_subject_1_run_3_pattern_discovery(self): + with tempfile.TemporaryDirectory() as tmp: + dataset = EEGBCIDataset(root=tmp, subjects=[1], runs=[3], download=True) + sample_dataset = dataset.set_task( + EEGBCIPatternDiscovery(compute_stft=False, window_size=2.0) + ) + self.assertGreater(len(sample_dataset), 0) + sample = sample_dataset[0] + self.assertIn("signal", sample) + self.assertEqual(sample["signal"].shape[0], 16) + self.assertIn(sample["task_label"], set(EEGBCI_LABELS)) + self.assertIn("bandpower", sample) + self.assertIn("brain_state_hypothesis", sample)