| """TR-level HRF regressor and z-score alignment utilities.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import nibabel as nib |
| import numpy as np |
| import pandas as pd |
|
|
|
|
| def _zscore_columns(values: np.ndarray, eps: float = 1e-8) -> np.ndarray: |
| mean = np.mean(values, axis=0, keepdims=True) |
| std = np.std(values, axis=0, keepdims=True) |
| std[std < eps] = 1.0 |
| return ((values - mean) / std).astype(np.float32) |
|
|
|
|
| def _trim_time_axis(values: np.ndarray, trim_start_tr: int, trim_end_tr: int) -> np.ndarray: |
| if trim_start_tr < 0 or trim_end_tr < 0: |
| raise ValueError("trim_start_tr and trim_end_tr must be non-negative") |
|
|
| start = int(trim_start_tr) |
| stop = values.shape[0] - int(trim_end_tr) |
|
|
| if start >= stop: |
| raise ValueError( |
| "Trimming removed all time points: " |
| f"n_tr={values.shape[0]}, trim_start={trim_start_tr}, trim_end={trim_end_tr}" |
| ) |
|
|
| return values[start:stop] |
|
|
|
|
| def _format_tr_for_filename(tr_seconds: float) -> str: |
| return f"{tr_seconds:.4f}".replace(".", "p") |
|
|
|
|
| def _assert_same_grid(img: nib.Nifti1Image, reference_img: nib.Nifti1Image, atol: float = 1e-5) -> None: |
| if img.shape[:3] != reference_img.shape: |
| raise ValueError( |
| "Image shape mismatch against analysis mask grid: " |
| f"image={img.shape[:3]}, reference={reference_img.shape}" |
| ) |
| if not np.allclose(img.affine, reference_img.affine, atol=atol): |
| raise ValueError("Image affine mismatch against analysis mask grid") |
|
|
|
|
| def _extract_masked_bold_2d( |
| bold_path: Path, |
| analysis_mask_bool: np.ndarray, |
| analysis_mask_img: nib.Nifti1Image, |
| ) -> np.ndarray: |
| img = nib.load(str(bold_path)) |
| if len(img.shape) != 4: |
| raise ValueError(f"Expected 4D BOLD image, got shape={img.shape} for {bold_path}") |
|
|
| _assert_same_grid(img=img, reference_img=analysis_mask_img) |
|
|
| data = np.asanyarray(img.dataobj) |
| |
| bold_2d = data[analysis_mask_bool, :].T.astype(np.float32) |
| return bold_2d |
|
|
|
|
| def compute_hrf_regressor_matrix( |
| word_features: np.ndarray, |
| onsets_s: np.ndarray, |
| offsets_s: np.ndarray, |
| tr_seconds: float, |
| n_volumes: int, |
| hrf_model: str, |
| ) -> np.ndarray: |
| """Build TR-level HRF-convolved regressors from word-level features.""" |
| from nilearn.glm.first_level import compute_regressor |
|
|
| if word_features.ndim != 2: |
| raise ValueError("word_features must be 2D: (n_words, hidden_dim)") |
|
|
| if onsets_s.shape[0] != word_features.shape[0] or offsets_s.shape[0] != word_features.shape[0]: |
| raise ValueError("onset/offset lengths must match n_words in word_features") |
|
|
| durations = offsets_s - onsets_s |
| frame_times = np.arange(n_volumes, dtype=np.float64) * float(tr_seconds) + 0.5 * float(tr_seconds) |
|
|
| regressors = np.zeros((int(n_volumes), int(word_features.shape[1])), dtype=np.float32) |
|
|
| for feature_index in range(word_features.shape[1]): |
| amplitudes = word_features[:, feature_index].astype(np.float64) |
| exp_condition = np.array([onsets_s, durations, amplitudes], dtype=np.float64) |
| signal, _ = compute_regressor(exp_condition, hrf_model, frame_times) |
| regressors[:, feature_index] = signal[:, 0].astype(np.float32) |
|
|
| return regressors |
|
|
|
|
| @dataclass(frozen=True) |
| class BoldAlignmentRecord: |
| """Summary for one subject-run BOLD alignment artifact.""" |
|
|
| subject: str |
| run: int |
| tr_seconds: float |
| n_volumes_raw: int |
| n_tr_aligned: int |
| n_voxels: int |
| trim_start_tr: int |
| trim_end_tr: int |
| bold_z_path: str |
|
|
|
|
| @dataclass(frozen=True) |
| class RegressorAlignmentRecord: |
| """Summary for one model-run-layer regressor artifact.""" |
|
|
| model_id: str |
| model_slug: str |
| run: int |
| tr_seconds: float |
| n_volumes_raw: int |
| n_tr_aligned: int |
| layer_idx: int |
| n_features: int |
| trim_start_tr: int |
| trim_end_tr: int |
| hrf_model: str |
| regressor_z_path: str |
|
|
|
|
| def build_and_cache_alignment_inputs( |
| manifest_df: pd.DataFrame, |
| run_events_df: pd.DataFrame, |
| feature_summary_df: pd.DataFrame, |
| analysis_mask_path: Path, |
| output_dir: Path, |
| trim_start_tr: int, |
| trim_end_tr: int, |
| hrf_model: str, |
| overwrite: bool, |
| alignment_subjects: list[str] | None = None, |
| ) -> tuple[pd.DataFrame, pd.DataFrame, dict[str, Any]]: |
| """Cache z-scored BOLD targets and HRF-convolved z-scored regressors.""" |
| if manifest_df.empty: |
| raise ValueError("Manifest is empty; cannot build alignment inputs") |
| if run_events_df.empty: |
| raise ValueError("run_events_df is empty; cannot build regressors") |
| if feature_summary_df.empty: |
| raise ValueError("feature_summary_df is empty; feature extraction must run first") |
|
|
| required_manifest_columns = {"subject", "run", "tr_seconds", "n_volumes", "derivatives_bold_path"} |
| missing_manifest = required_manifest_columns.difference(manifest_df.columns) |
| if missing_manifest: |
| raise ValueError(f"manifest_df missing required columns: {sorted(missing_manifest)}") |
|
|
| required_feature_columns = {"model_id", "model_slug", "run", "features_npz_path"} |
| missing_feature = required_feature_columns.difference(feature_summary_df.columns) |
| if missing_feature: |
| raise ValueError(f"feature_summary_df missing required columns: {sorted(missing_feature)}") |
|
|
| analysis_mask_img = nib.load(str(analysis_mask_path)) |
| analysis_mask_bool = analysis_mask_img.get_fdata() > 0.5 |
|
|
| manifest_alignment_df = manifest_df |
| requested_subjects = [value for value in (alignment_subjects or []) if str(value).strip()] |
| if requested_subjects: |
| subject_set = {str(value) for value in requested_subjects} |
| manifest_alignment_df = manifest_df[manifest_df["subject"].isin(subject_set)].copy() |
| if manifest_alignment_df.empty: |
| raise ValueError( |
| "Alignment subject filter produced no rows. " |
| f"Requested subjects={sorted(subject_set)}" |
| ) |
|
|
| output_dir = output_dir.resolve() |
| bold_output_dir = output_dir / "bold_z" |
| regressor_output_dir = output_dir / "regressors" |
| bold_output_dir.mkdir(parents=True, exist_ok=True) |
| regressor_output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| bold_rows: list[BoldAlignmentRecord] = [] |
|
|
| for row in manifest_alignment_df.itertuples(index=False): |
| subject = str(getattr(row, "subject")) |
| run = int(getattr(row, "run")) |
| tr_seconds = float(getattr(row, "tr_seconds")) |
| n_volumes = int(getattr(row, "n_volumes")) |
| bold_path = Path(str(getattr(row, "derivatives_bold_path"))) |
|
|
| bold_z_path = bold_output_dir / f"{subject}_run-{run:02d}_bold_z.npy" |
| if bold_z_path.exists() and not overwrite: |
| aligned_bold = np.load(bold_z_path) |
| else: |
| bold_2d = _extract_masked_bold_2d( |
| bold_path=bold_path, |
| analysis_mask_bool=analysis_mask_bool, |
| analysis_mask_img=analysis_mask_img, |
| ) |
|
|
| if bold_2d.shape[0] != n_volumes: |
| raise ValueError( |
| f"BOLD volume mismatch for {subject} run {run}: " |
| f"manifest={n_volumes}, loaded={bold_2d.shape[0]}" |
| ) |
|
|
| aligned_bold = _trim_time_axis( |
| values=bold_2d, |
| trim_start_tr=trim_start_tr, |
| trim_end_tr=trim_end_tr, |
| ) |
| aligned_bold = _zscore_columns(aligned_bold) |
| np.save(bold_z_path, aligned_bold) |
|
|
| bold_rows.append( |
| BoldAlignmentRecord( |
| subject=subject, |
| run=run, |
| tr_seconds=tr_seconds, |
| n_volumes_raw=n_volumes, |
| n_tr_aligned=int(aligned_bold.shape[0]), |
| n_voxels=int(aligned_bold.shape[1]), |
| trim_start_tr=int(trim_start_tr), |
| trim_end_tr=int(trim_end_tr), |
| bold_z_path=str(bold_z_path), |
| ) |
| ) |
|
|
| regressor_rows: list[RegressorAlignmentRecord] = [] |
|
|
| |
| run_config_map: dict[int, list[tuple[float, int]]] = {} |
| for run, group_df in manifest_alignment_df.groupby("run"): |
| unique_pairs = sorted( |
| {(float(tr), int(n_vol)) for tr, n_vol in zip(group_df["tr_seconds"], group_df["n_volumes"])}, |
| key=lambda pair: (pair[0], pair[1]), |
| ) |
| run_config_map[int(run)] = unique_pairs |
|
|
| for feature_row in feature_summary_df.itertuples(index=False): |
| model_id = str(getattr(feature_row, "model_id")) |
| model_slug = str(getattr(feature_row, "model_slug")) |
| run = int(getattr(feature_row, "run")) |
| features_npz_path = Path(str(getattr(feature_row, "features_npz_path"))) |
|
|
| if run not in run_config_map: |
| continue |
|
|
| run_events = run_events_df[run_events_df["run"] == run].sort_values("word_index") |
| onsets_s = run_events["onset_s"].to_numpy(dtype=np.float64) |
| offsets_s = run_events["offset_s"].to_numpy(dtype=np.float64) |
|
|
| feature_bundle = np.load(features_npz_path) |
| layer_keys = sorted( |
| [key for key in feature_bundle.files if key.startswith("layer_")], |
| key=lambda value: int(value.split("_")[1]), |
| ) |
|
|
| for tr_seconds, n_volumes in run_config_map[run]: |
| tr_tag = _format_tr_for_filename(tr_seconds) |
|
|
| model_regressor_dir = regressor_output_dir / model_slug |
| model_regressor_dir.mkdir(parents=True, exist_ok=True) |
|
|
| for layer_key in layer_keys: |
| layer_idx = int(layer_key.split("_")[1]) |
| layer_features = np.asarray(feature_bundle[layer_key], dtype=np.float32) |
|
|
| if layer_features.shape[0] != onsets_s.shape[0]: |
| raise ValueError( |
| f"Word count mismatch for model={model_id}, run={run}, layer={layer_idx}: " |
| f"features={layer_features.shape[0]}, events={onsets_s.shape[0]}" |
| ) |
|
|
| regressor_z_path = ( |
| model_regressor_dir |
| / f"run-{run:02d}_tr-{tr_tag}_nvol-{n_volumes}_layer-{layer_idx:03d}_regressor_z.npy" |
| ) |
|
|
| if regressor_z_path.exists() and not overwrite: |
| regressors_aligned = np.load(regressor_z_path) |
| else: |
| regressors = compute_hrf_regressor_matrix( |
| word_features=layer_features, |
| onsets_s=onsets_s, |
| offsets_s=offsets_s, |
| tr_seconds=tr_seconds, |
| n_volumes=n_volumes, |
| hrf_model=hrf_model, |
| ) |
| regressors_aligned = _trim_time_axis( |
| values=regressors, |
| trim_start_tr=trim_start_tr, |
| trim_end_tr=trim_end_tr, |
| ) |
| regressors_aligned = _zscore_columns(regressors_aligned) |
| np.save(regressor_z_path, regressors_aligned) |
|
|
| regressor_rows.append( |
| RegressorAlignmentRecord( |
| model_id=model_id, |
| model_slug=model_slug, |
| run=run, |
| tr_seconds=tr_seconds, |
| n_volumes_raw=n_volumes, |
| n_tr_aligned=int(regressors_aligned.shape[0]), |
| layer_idx=layer_idx, |
| n_features=int(regressors_aligned.shape[1]), |
| trim_start_tr=int(trim_start_tr), |
| trim_end_tr=int(trim_end_tr), |
| hrf_model=hrf_model, |
| regressor_z_path=str(regressor_z_path), |
| ) |
| ) |
|
|
| bold_summary_df = pd.DataFrame([asdict(row) for row in bold_rows]) |
| if not bold_summary_df.empty: |
| bold_summary_df = bold_summary_df.sort_values(["subject", "run"]).reset_index(drop=True) |
|
|
| regressor_summary_df = pd.DataFrame([asdict(row) for row in regressor_rows]) |
| if not regressor_summary_df.empty: |
| regressor_summary_df = regressor_summary_df.sort_values( |
| ["model_slug", "run", "layer_idx", "tr_seconds", "n_volumes_raw"] |
| ).reset_index(drop=True) |
|
|
| alignment_qc: dict[str, Any] = { |
| "analysis_mask_path": str(analysis_mask_path), |
| "hrf_model": hrf_model, |
| "trim_start_tr": int(trim_start_tr), |
| "trim_end_tr": int(trim_end_tr), |
| "alignment_subjects": sorted({str(value) for value in manifest_alignment_df["subject"].tolist()}), |
| "n_bold_rows": int(len(bold_summary_df)), |
| "n_regressor_rows": int(len(regressor_summary_df)), |
| } |
|
|
| if not bold_summary_df.empty: |
| alignment_qc["bold_n_tr_min"] = int(bold_summary_df["n_tr_aligned"].min()) |
| alignment_qc["bold_n_tr_max"] = int(bold_summary_df["n_tr_aligned"].max()) |
| alignment_qc["bold_n_voxels_min"] = int(bold_summary_df["n_voxels"].min()) |
| alignment_qc["bold_n_voxels_max"] = int(bold_summary_df["n_voxels"].max()) |
|
|
| if not regressor_summary_df.empty: |
| alignment_qc["regressor_n_tr_min"] = int(regressor_summary_df["n_tr_aligned"].min()) |
| alignment_qc["regressor_n_tr_max"] = int(regressor_summary_df["n_tr_aligned"].max()) |
| alignment_qc["regressor_n_features_min"] = int(regressor_summary_df["n_features"].min()) |
| alignment_qc["regressor_n_features_max"] = int(regressor_summary_df["n_features"].max()) |
|
|
| return bold_summary_df, regressor_summary_df, alignment_qc |
|
|
