diff --git a/encoding/models/folding.py b/encoding/models/folding.py index 23db8f1..60c27f1 100644 --- a/encoding/models/folding.py +++ b/encoding/models/folding.py @@ -12,6 +12,7 @@ def create_folds( chunk_length: Optional[int] = None, trim_size: Optional[int] = None, groups: Optional[np.ndarray] = None, + seed: int = 0, ) -> List[Tuple[List[int], List[int]]]: """Create train/test indices according to the specified folding strategy @@ -24,6 +25,7 @@ def create_folds( trim_size: Number of TRs to trim from beginning and end of test folds/chunks (used for 'chunked_trimmed' and 'kfold_trimmed' fold types) groups: Group labels for group-based folding + seed: Random seed for chunk shuffling, for reproducible folds Returns: List of (train_indices, test_indices) tuples @@ -31,17 +33,21 @@ def create_folds( print(f"THIS IS THE FOLDING TYPE: {fold_type} inside folding.py") if fold_type == "chunked": # Current chunked implementation with random shuffling - return create_chunked_folds(n_samples, n_folds, chunk_length, shuffle=True) + return create_chunked_folds( + n_samples, n_folds, chunk_length, shuffle=True, seed=seed + ) elif fold_type == "chunked_trimmed": # Chunked with trimmed test chunks to reduce autocorrelation if trim_size is None: trim_size = 5 # Default trim size of 5 TRs return create_chunked_folds_trimmed( - n_samples, n_folds, chunk_length, trim_size, shuffle=True + n_samples, n_folds, chunk_length, trim_size, shuffle=True, seed=seed ) elif fold_type == "chunked_contiguous": # Chunked but maintaining original order - return create_chunked_folds(n_samples, n_folds, chunk_length, shuffle=False) + return create_chunked_folds( + n_samples, n_folds, chunk_length, shuffle=False, seed=seed + ) elif fold_type == "kfold": # Regular KFold with shuffle=False for contiguity kf = KFold(n_splits=n_folds, shuffle=False) @@ -65,7 +71,11 @@ def create_folds( def create_chunked_folds( - n_samples: int, n_folds: int, chunk_length: int, shuffle: bool = True + n_samples: int, + n_folds: int, + chunk_length: int, + shuffle: bool = True, + seed: int = 0, ) -> List[Tuple[List[int], List[int]]]: """Create KFold splits that respect chunk structure of fMRI data. @@ -74,6 +84,7 @@ def create_chunked_folds( n_folds: Number of folds to create chunk_length: Length of chunks shuffle: Whether to shuffle chunks or keep them in original order + seed: Random seed for chunk shuffling, for reproducible folds Returns: List of (train_indices, test_indices) tuples @@ -83,7 +94,7 @@ def create_chunked_folds( chunk_indices = list(range(n_complete_chunks)) if shuffle: - random.shuffle(chunk_indices) + random.Random(seed).shuffle(chunk_indices) # Divide chunks evenly into n_folds chunks_per_fold = n_complete_chunks // n_folds @@ -92,7 +103,9 @@ def create_chunked_folds( logging.warning( "Not enough chunks for the requested folds, falling back to regular KFold" ) - kf = KFold(n_splits=n_folds, shuffle=shuffle) + kf = KFold( + n_splits=n_folds, shuffle=shuffle, random_state=seed if shuffle else None + ) return list(kf.split(range(n_samples))) # Create exactly n_folds splits @@ -130,6 +143,7 @@ def create_chunked_folds_trimmed( chunk_length: int, trim_size: int = 5, shuffle: bool = True, + seed: int = 0, ) -> List[Tuple[List[int], List[int]]]: """Create KFold splits with trimmed test chunks to reduce autocorrelation effects. @@ -144,6 +158,7 @@ def create_chunked_folds_trimmed( chunk_length: Length of chunks trim_size: Number of time points to trim from beginning and end of test chunks shuffle: Whether to shuffle chunks or keep them in original order + seed: Random seed for chunk shuffling, for reproducible folds Returns: List of (train_indices, test_indices) tuples @@ -153,7 +168,7 @@ def create_chunked_folds_trimmed( chunk_indices = list(range(n_complete_chunks)) if shuffle: - random.shuffle(chunk_indices) + random.Random(seed).shuffle(chunk_indices) # Divide chunks evenly into n_folds chunks_per_fold = n_complete_chunks // n_folds diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..2f6e689 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,311 @@ +"""Shared fixtures for the synthetic end-to-end pipeline tests. + +Builds a small assembly with the same structural conventions as +``data/assembly_lebel_uts03.pkl`` and a signal planted at known FIR lags. +Conventions (verified against all 25 stories of that file): + +* ``len(split_indices) + 1 == len(tr_times)`` -- split_indices are *cut points* + into the word axis, consumed by ``np.split``, not a per-word TR label. +* ``brain_data.shape[0] == len(tr_times) - 15`` -- 10 TRs trimmed at the head and + 5 at the tail upstream of the assembly. +* ``len(words) == len(stimuli) == len(data_times) == n_words`` +* ``word_rates.shape == (n_trs, 1)`` +""" + +import os + +# Must happen before anything imports matplotlib/wandb (encoding.trainer pulls in +# encoding.plotting.plotting_utils, which imports matplotlib and nilearn at module +# scope). +os.environ.setdefault("MPLBACKEND", "Agg") +os.environ.setdefault("WANDB_MODE", "disabled") +os.environ.setdefault("WANDB_SILENT", "true") +os.environ.setdefault("WANDB_DISABLED", "true") + +import hashlib +from dataclasses import dataclass, field +from typing import Dict, List, Tuple + +import numpy as np +import pytest + +from encoding.assembly.assemblies import SimpleNeuroidAssembly +from encoding.assembly.story_data import StoryData +from encoding.features.base import BaseFeatureExtractor +from encoding.features.embeddings import StaticEmbeddingFeatureExtractor + +N_VERTICES = 20484 # fsaverage5, both hemispheres +N_FEATURES = 4 +N_SIGNAL_VERTICES = 10 # vertices 0..9 carry the planted signal +PLANTED_LAGS = (2, 3) +FIR_DELAYS = [1, 2, 3, 4] +TR = 1.5 +BRAIN_TRIM_HEAD = 10 +BRAIN_TRIM_TAIL = 5 +BRAIN_TRIM_TOTAL = BRAIN_TRIM_HEAD + BRAIN_TRIM_TAIL + +# Same trimming used in train_simple.py for the LeBel dataset. +TRIMMING_CONFIG = { + "train_features_start": 10, + "train_features_end": -5, + "train_targets_start": 0, + "train_targets_end": None, + "test_features_start": 50, + "test_features_end": -5, + "test_targets_start": 40, + "test_targets_end": None, +} + + +def token_vector(token: str, dim: int = N_FEATURES) -> np.ndarray: + """Map a token to a fixed pseudo-random vector. + + Uses blake2b rather than ``hash()`` so the mapping is stable across processes + (Python salts str hashes per interpreter run). + """ + digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest() + seed = int.from_bytes(digest, "little") + return np.random.default_rng(seed).standard_normal(dim) + + +class SyntheticTokenFeatureExtractor(StaticEmbeddingFeatureExtractor): + """Word-level extractor that plugs into the real factory dispatch. + + Subclasses ``StaticEmbeddingFeatureExtractor`` so the factory feeds it + ``assembly.get_words()[idx]``; the parent ``__init__`` is bypassed because it + requires a gensim vector file. The class name must not contain "wordrate", or + ``AbstractTrainer._should_downsample`` would skip downsampling for it. + """ + + def __init__(self, dim: int = N_FEATURES): + BaseFeatureExtractor.__init__(self, {"model_name": "synthetic-tokens"}) + self.dim = dim + + def extract_features(self, stimuli, **kwargs) -> np.ndarray: + if isinstance(stimuli, str): + stimuli = stimuli.split() + return np.vstack([token_vector(tok, self.dim) for tok in stimuli]) + + +def reference_legacy_average(word_features: np.ndarray, counts: np.ndarray) -> np.ndarray: + """Per-TR mean of word features, from counts rather than np.split; empty TRs stay zero.""" + out = np.zeros((len(counts), word_features.shape[1])) + pos = 0 + for i, c in enumerate(counts): + if c: + out[i] = word_features[pos:pos + c].mean(axis=0) + pos += c + return out + + +def reference_make_delayed(features: np.ndarray, delays) -> np.ndarray: + """Causal FIR expansion: block d of row t holds features[t - d], else zeros.""" + n_t, n_d = features.shape + blocks = [] + for d in delays: + block = np.zeros((n_t, n_d)) + block[d:] = features[:-d] + blocks.append(block) + return np.hstack(blocks) + + +@dataclass +class StoryTruth: + n_trs: int + n_words: int + counts: np.ndarray # words per TR, length n_trs + word_features: np.ndarray # (n_words, N_FEATURES) + tr_features: np.ndarray # (n_trs, N_FEATURES), the correct downsampling + empty_trs: np.ndarray # indices of TRs containing no words + + +@dataclass +class GroundTruth: + stories: Dict[str, StoryTruth] = field(default_factory=dict) + lag_weights: Dict[int, np.ndarray] = field(default_factory=dict) # lag -> (N_FEATURES, N_SIGNAL) + signal_vertices: np.ndarray = field(default_factory=lambda: np.arange(N_SIGNAL_VERTICES)) + noise_vertices: np.ndarray = field(default_factory=lambda: np.arange(N_SIGNAL_VERTICES, N_VERTICES)) + + +def synthetic_assembly( + n_stories: int = 3, + n_trs: Tuple[int, ...] = (200, 185, 210), + n_vertices: int = N_VERTICES, + n_features: int = N_FEATURES, + n_signal_vertices: int = N_SIGNAL_VERTICES, + planted_lags: Tuple[int, ...] = PLANTED_LAGS, + noise_ratio: float = 0.5, + max_words_per_tr: int = 12, + seed: int = 0, +) -> SimpleNeuroidAssembly: + """Build a small assembly with a signal planted at known FIR lags. + + Ground truth is attached as ``assembly.truth``. Brain data for the first + ``n_signal_vertices`` is + + brain[j, v] = sum_{d in planted_lags} F[j + BRAIN_TRIM_HEAD - d] @ lag_weights[d][:, v] + + where ``F`` is the legacy-averaged TR-level features; other vertices are noise. + The ``+ BRAIN_TRIM_HEAD`` offset matches ``TRIMMING_CONFIG``: the trainer pairs + ``features[10:-5]`` with ``targets[0:]``, so brain row ``j`` is aligned with + feature row ``j + 10``, whose FIR block for delay ``d`` holds ``F[j + 10 - d]``. + + ``noise_ratio`` is noise std as a fraction of each signal vertex's std; 0.5 + gives a ceiling of r ~= 0.89. + """ + if n_stories < 2: + raise ValueError("need at least 2 stories for a train/test split") + if max(planted_lags) >= BRAIN_TRIM_HEAD: + raise ValueError("planted lags must be < BRAIN_TRIM_HEAD to stay in range") + + rng = np.random.default_rng(seed) + lag_weights = { + lag: rng.standard_normal((n_features, n_signal_vertices)) + for lag in planted_lags + } + + truth = GroundTruth( + lag_weights=lag_weights, + signal_vertices=np.arange(n_signal_vertices), + noise_vertices=np.arange(n_signal_vertices, n_vertices), + ) + + story_data_list: List[StoryData] = [] + + for s in range(n_stories): + story_rng = np.random.default_rng(seed + 1000 + s) + name = f"synthstory{s}" + n_tr = n_trs[s % len(n_trs)] + + counts = story_rng.integers(0, max_words_per_tr, size=n_tr) + # Mirror the real data: silence at the very start and end of a run. + counts[:3] = 0 + counts[-2:] = 0 + n_words = int(counts.sum()) + + # n_tr - 1 cut points, so np.split yields n_tr chunks. + split_indices = np.cumsum(counts)[:-1].astype(np.int64).tolist() + + words = [f"w{s}_{i}" for i in range(n_words)] + stimuli = [" ".join(words[max(0, i - 5):i + 1]) for i in range(n_words)] + + tr_times = np.arange(n_tr, dtype=np.float64) * TR + + data_times = np.empty(n_words, dtype=np.float64) + pos = 0 + for i, c in enumerate(counts): + if c: + offsets = np.sort(story_rng.uniform(0, TR, size=c)) + data_times[pos:pos + c] = tr_times[i] + offsets + pos += c + + word_features = np.vstack([token_vector(w, n_features) for w in words]) + tr_features = reference_legacy_average(word_features, counts) + + n_brain = n_tr - BRAIN_TRIM_TOTAL + j = np.arange(n_brain) + signal = np.zeros((n_brain, n_signal_vertices)) + for lag in planted_lags: + signal += tr_features[j + BRAIN_TRIM_HEAD - lag] @ lag_weights[lag] + + brain = np.empty((n_brain, n_vertices), dtype=np.float64) + signal_std = signal.std(axis=0) + brain[:, :n_signal_vertices] = signal + story_rng.standard_normal( + (n_brain, n_signal_vertices) + ) * (signal_std * noise_ratio) + brain[:, n_signal_vertices:] = story_rng.standard_normal( + (n_brain, n_vertices - n_signal_vertices) + ) * float(np.median(signal_std)) + + truth.stories[name] = StoryTruth( + n_trs=n_tr, + n_words=n_words, + counts=counts, + word_features=word_features, + tr_features=tr_features, + empty_trs=np.flatnonzero(counts == 0), + ) + + story_data_list.append( + StoryData( + name=name, + brain_data=brain, + stimuli=stimuli, + split_indices=split_indices, + tr_times=tr_times, + data_times=data_times, + temporal_baseline=None, + word_rates=counts.reshape(-1, 1).astype(np.float64), + words=words, + mask_indices=None, + audio_path=f"/nonexistent/audio/{name}.wav", + ) + ) + + assembly = SimpleNeuroidAssembly(story_data_list, validation_method="outer") + assembly.truth = truth + return assembly + + +class _NullLogger: + """Stand-in for WandBLogger/TensorBoardLogger that records scalars in memory.""" + + def __init__(self, *args, **kwargs): + self.scalars: Dict[str, float] = {} + + def log_scalar(self, name, value, step=None): + self.scalars[name] = value + + def log_image(self, name, figure, step=None): + pass + + def log_histogram(self, name, values, step=None): + pass + + +class _NullBrainPlotter: + """Stand-in for BrainPlotter; the real one downloads fsaverage and renders 3D.""" + + def __init__(self, logger=None): + self.logger = logger + self.calls = 0 + + def log_plots(self, *args, **kwargs): + self.calls += 1 + + +@pytest.fixture(autouse=True) +def _no_external_side_effects(monkeypatch): + """Keep tests offline: no wandb run, no nilearn surface download, no GUI backend.""" + import matplotlib + matplotlib.use("Agg", force=True) + + import encoding.trainer as trainer_mod + + monkeypatch.setattr(trainer_mod, "WandBLogger", _NullLogger) + monkeypatch.setattr(trainer_mod, "TensorBoardLogger", _NullLogger) + monkeypatch.setattr(trainer_mod, "BrainPlotter", _NullBrainPlotter) + + # setup_logger calls wandb.init() directly before constructing WandBLogger. + import wandb + monkeypatch.setattr(wandb, "init", lambda *a, **k: None, raising=False) + monkeypatch.setattr(wandb, "log", lambda *a, **k: None, raising=False) + monkeypatch.setattr(wandb, "finish", lambda *a, **k: None, raising=False) + + import encoding.plotting.plotting_utils as plotting_utils + monkeypatch.setattr( + plotting_utils.datasets, + "fetch_surf_fsaverage", + lambda *a, **k: pytest.fail("test tried to download fsaverage"), + raising=False, + ) + + +@pytest.fixture(scope="module") +def assembly() -> SimpleNeuroidAssembly: + return synthetic_assembly() + + +@pytest.fixture(scope="module") +def truth(assembly) -> GroundTruth: + return assembly.truth diff --git a/tests/fixtures/golden_lebel_uts03.json b/tests/fixtures/golden_lebel_uts03.json new file mode 100644 index 0000000..b55e1c2 --- /dev/null +++ b/tests/fixtures/golden_lebel_uts03.json @@ -0,0 +1,157 @@ +{ + "fixture": { + "file": "assembly_lebel_uts03.pkl", + "md5": "9bec82d87d528ca0e68e421c3c73136e", + "source": "Google Drive link in README", + "env_var": "LITCODER_LEBEL_ASSEMBLY" + }, + "run": { + "date": "2026-10-12", + "commit": "ec31915a05daede3edfd31ae9966bec556e5c5a9", + "script": "train_simple.py", + "machine": "macOS, Apple Silicon, MPS (linalg_svd falls back to CPU)", + "env": "conda env 'litcoder'", + "fit_seconds": 14.5, + "fold_seed": 0 + }, + "environment": { + "note": "Tolerances are chosen strict when this fingerprint matches the running environment exactly, loose otherwise. See tolerances.rationale.", + "fingerprint": { + "platform_system": "Darwin", + "platform_machine": "arm64", + "torch": "2.13.0", + "numpy": "2.4.6", + "scipy": "1.18.0", + "device": "mps:0", + "torch_threads": 4 + } + }, + "config": { + "extractor": "wordrate", + "fir_delays": [ + 1, + 2, + 3, + 4 + ], + "folding_type": "chunked", + "n_inner_folds": 5, + "chunk_length": 20, + "alphas": "np.logspace(-1, 8, 10) (fit_predict default)", + "downsampler": "default (rect; wordrate is not downsampled)", + "trimming_config": { + "train_features_start": 10, + "train_features_end": -5, + "train_targets_start": 0, + "train_targets_end": null, + "test_features_start": 50, + "test_features_end": -5, + "test_targets_start": 40, + "test_targets_end": null + } + }, + "shapes": { + "X_train": [ + 8673, + 4 + ], + "Y_train": [ + 8673, + 20484 + ], + "X_test": [ + 251, + 4 + ], + "Y_test": [ + 251, + 20484 + ] + }, + "metrics": { + "median_score": 0.042643629, + "n_significant": 2867, + "n_voxels": 20484 + }, + "quantiles": { + "1": -0.1695338935, + "5": -0.1018701419, + "25": -0.0154135684, + "50": 0.042643629, + "75": 0.1100961417, + "90": 0.1946224123, + "95": 0.2598253787, + "99": 0.4205047786 + }, + "threshold_counts": { + "0.1": 5721, + "0.2": 1924 + }, + "tolerances": { + "quantile_atol_strict": 1e-05, + "quantile_atol_loose": 0.002, + "count_rtol_strict": 0.01, + "count_rtol_loose": 0.05, + "rationale": "Derived from measurement, not from the two historical runs. Across four configurations of the same input (two processes on CPU, one on MPS, one with torch 2.11/numpy 1.26/scipy 1.13) the largest spread of any asserted quantile was 2.4e-07, and n_significant was identical (2867) in all four. Strict quantile atol of 1e-5 is ~40x that observed spread and ~1x the mean spacing of the correlation distribution near the median, so a single voxel changing rank cannot trip it. Loose atol of 2e-3 covers the regime where per-voxel correlations move by 1e-3; real regressions (misaligned trimming, wrong downsampler, wrong FIR window) move quantiles by 1e-2 or more, so the loose band still catches them with an order of magnitude to spare. Counts are banded rather than exact because FDR thresholding is an order statistic with a cascade (see measured_variance.bh_cutoff)." + }, + "measured_variance": { + "method": "Same persisted train/test matrices fed to NestedCVModel.fit_predict under four configurations.", + "cpu_process_1_vs_process_2": { + "bitwise_identical": true, + "max_abs_dr": 0.0, + "delta_median": 0.0, + "alpha_flips": 0, + "significance_flips": 0 + }, + "cpu_vs_mps": { + "bitwise_identical": false, + "max_abs_dr": 6.654e-06, + "p99_abs_dr": 4.321e-07, + "delta_median": 0.0, + "alpha_flips": 0, + "significance_flips": 0 + }, + "cross_library_torch_2_11_numpy_1_26_scipy_1_13": { + "bitwise_identical": false, + "max_abs_dr": 1.067e-07, + "delta_median": 4.717e-09, + "alpha_flips": 0, + "significance_flips": 0 + }, + "max_quantile_spread_across_all_four": 2.384e-07, + "correlation_spacing_at_median": { + "adjacent_gap": 4.247e-07, + "mean_gap_within_100_ranks": 9.449e-06, + "max_gap_within_100_ranks": 6.757e-05 + }, + "bh_cutoff": { + "alpha": 0.05, + "p_at_cutoff": 0.006986395, + "bh_step_alpha_over_m": 2.441e-06, + "gap_to_next_p_value": 2.414e-05, + "voxels_within_one_bh_index_step": 1, + "n_significant_swing_for_uniform_r_shift": { + "1e-05": 0, + "0.0001": 10, + "0.001": 59 + } + }, + "threshold_count_swing_for_uniform_r_shift": { + "r_gt_0.1": { + "1e-05": 0, + "0.0001": 16, + "0.001": 128 + }, + "r_gt_0.2": { + "1e-05": 1, + "0.0001": 4, + "0.001": 44 + } + } + }, + "regeneration": { + "command": "LITCODER_LEBEL_ASSEMBLY=/path/to/assembly_lebel_uts03.pkl python tests/test_golden_lebel.py --update", + "when": "Run this only when a change to the pipeline is intended. It rewrites metrics, quantiles, threshold_counts, shapes and the environment fingerprint from a fresh run, and leaves tolerances and notes alone. Review the diff: a change you did not expect is the test doing its job." + }, + "notes": "Baseline generated locally rather than on PACE. An earlier run recorded on 2026-08 (median 0.043, 2853 significant voxels) is NOT comparable to this one and must not be used to justify a tolerance. That run predates the commit 'seed fold shuffling for reproducible splits'; before it, create_chunked_folds called bare random.shuffle() on the global unseeded RNG, so the inner-CV folds differed on every invocation and the fit was a different model each time. The gap between the two recorded numbers is therefore not evidence of environmental variance, and the August run is not reproducible by current code on any machine." +} diff --git a/tests/test_golden_lebel.py b/tests/test_golden_lebel.py new file mode 100644 index 0000000..3048f79 --- /dev/null +++ b/tests/test_golden_lebel.py @@ -0,0 +1,361 @@ +"""Characterization test for the real LeBel assembly. + +Pins the numbers ``train_simple.py`` produces against a 2.8 GB fixture that is not +in the repository. Skips when the fixture is absent. + +Running it:: + + export LITCODER_LEBEL_ASSEMBLY=/path/to/assembly_lebel_uts03.pkl + pytest tests/test_golden_lebel.py -v + +Regenerating the golden values after an *intended* change:: + + LITCODER_LEBEL_ASSEMBLY=/path/to/assembly_lebel_uts03.pkl \ + python tests/test_golden_lebel.py --update + +Tolerances and the measurements behind them are in the JSON under +``tolerances.rationale`` and ``measured_variance``. +""" + +import datetime +import hashlib +import json +import os +import platform +import subprocess +import sys +import warnings +from pathlib import Path + +import numpy as np +import pytest + +GOLDEN_PATH = Path(__file__).parent / "fixtures" / "golden_lebel_uts03.json" +ENV_VAR = "LITCODER_LEBEL_ASSEMBLY" + +# Same trimming train_simple.py uses for LeBel. +TRIMMING_CONFIG = { + "train_features_start": 10, + "train_features_end": -5, + "train_targets_start": 0, + "train_targets_end": None, + "test_features_start": 50, + "test_features_end": -5, + "test_targets_start": 40, + "test_targets_end": None, +} +FIR_DELAYS = [1, 2, 3, 4] + + +def load_golden() -> dict: + with open(GOLDEN_PATH) as f: + return json.load(f) + + +def skip_reason(golden: dict) -> str: + raw = os.environ.get(ENV_VAR) + where = golden["fixture"]["source"] + name = golden["fixture"]["file"] + md5 = golden["fixture"]["md5"] + if not raw: + return ( + f"{ENV_VAR} is not set, so the real LeBel assembly is unavailable.\n" + f" This test is opt-in and CI is expected to skip it.\n" + f" To run it: export {ENV_VAR}=/path/to/{name}\n" + f" The file is ~2.8 GB (md5 {md5}); get it from: {where}." + ) + return ( + f"{ENV_VAR} is set to {raw!r}, but no file exists there.\n" + f" Point it at {name} (~2.8 GB, md5 {md5}).\n" + f" Get it from: {where}." + ) + + +def file_md5(path: Path, chunk: int = 8 << 20) -> str: + h = hashlib.md5() + with open(path, "rb") as f: + for block in iter(lambda: f.read(chunk), b""): + h.update(block) + return h.hexdigest() + + +def environment_fingerprint() -> dict: + import scipy + import torch + + if torch.backends.mps.is_available(): + device = "mps:0" + elif torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + + return { + "platform_system": platform.system(), + "platform_machine": platform.machine(), + "torch": torch.__version__, + "numpy": np.__version__, + "scipy": scipy.__version__, + "device": device, + "torch_threads": torch.get_num_threads(), + } + + +def run_pipeline(assembly_path: str, results_dir: str) -> dict: + """Run what train_simple.py runs, with logging and plotting stubbed. + + Calls ``fit_predict`` with no model kwargs, as ``train_simple.py`` calls + ``trainer.train()`` with none -- the defaults are part of what is characterized. + """ + from conftest import _NullBrainPlotter, _NullLogger + import encoding.trainer as trainer_mod + from encoding.assembly.assembly_loader import load_assembly + from encoding.downsample.downsampling import Downsampler + from encoding.features.factory import FeatureExtractorFactory + from encoding.models.nested_cv import NestedCVModel + + assembly = load_assembly(assembly_path) + extractor = FeatureExtractorFactory.create_extractor( + modality="wordrate", + model_name="wordrate", + config={}, + cache_dir=str(Path(results_dir) / "cache"), + ) + trainer = trainer_mod.AbstractTrainer( + assembly=assembly, + feature_extractors=[extractor], + downsampler=Downsampler(), + model=NestedCVModel(model_name="ridge_regression"), + fir_delays=FIR_DELAYS, + trimming_config=dict(TRIMMING_CONFIG), + use_train_test_split=True, + logger_backend="wandb", + dataset_type="lebel", + results_dir=results_dir, + downsample_config={}, + ) + assert isinstance(trainer.experiment_logger, _NullLogger) + assert isinstance(trainer.brain_plotter, _NullBrainPlotter) + + data = trainer.structure_data( + trainer.apply_fir_delays(trainer.extract_and_downsample_features()) + ) + metrics, weights, alphas = trainer.model.fit_predict( + features=data["Rstim"], + targets=data["Rresp"], + X_test=data["Pstim"], + y_test=data["Presp"], + ) + return { + "shapes": { + "X_train": list(data["Rstim"].shape), + "Y_train": list(data["Rresp"].shape), + "X_test": list(data["Pstim"].shape), + "Y_test": list(data["Presp"].shape), + }, + "correlations": np.asarray(metrics["correlations"]), + "n_significant": int(metrics["n_significant"]), + "median_score": float(metrics["median_score"]), + } + + +@pytest.fixture(scope="module") +def golden() -> dict: + return load_golden() + + +@pytest.fixture(scope="module") +def verified_assembly_path(golden) -> str: + """Skip when the fixture is absent; fail when it is the wrong file.""" + raw = os.environ.get(ENV_VAR) + if not raw or not Path(raw).is_file(): + pytest.skip(skip_reason(golden)) + + path = Path(raw) + actual = file_md5(path) + expected = golden["fixture"]["md5"] + if actual != expected: + pytest.fail( + f"Assembly at {path} has md5 {actual}; golden values were generated " + f"against md5 {expected}.\n" + f"If the new file is intentional, regenerate with:\n" + f" {golden['regeneration']['command']}" + ) + return str(path) + + +@pytest.fixture(scope="module") +def result(verified_assembly_path, golden, tmp_path_factory) -> dict: + """Module-scoped, so conftest's function-scoped autouse patching is not yet + active; patch explicitly for the duration of the run.""" + with pytest.MonkeyPatch.context() as mp: + from conftest import _NullBrainPlotter, _NullLogger + import encoding.trainer as trainer_mod + import wandb + + mp.setattr(trainer_mod, "WandBLogger", _NullLogger) + mp.setattr(trainer_mod, "BrainPlotter", _NullBrainPlotter) + mp.setattr(wandb, "init", lambda *a, **k: None, raising=False) + + results_dir = tmp_path_factory.mktemp("golden_results") + return run_pipeline(verified_assembly_path, str(results_dir)) + + +@pytest.fixture(scope="module") +def tolerances(golden) -> dict: + """Pick the strict or loose set by fingerprint, and say which was used.""" + current = environment_fingerprint() + recorded = golden["environment"]["fingerprint"] + strict = current == recorded + + tol = golden["tolerances"] + chosen = { + "strict": strict, + "quantile_atol": tol["quantile_atol_strict"] if strict else tol["quantile_atol_loose"], + "count_rtol": tol["count_rtol_strict"] if strict else tol["count_rtol_loose"], + } + + banner = ( + f"\n[golden_lebel] tolerance set: {'STRICT' if strict else 'LOOSE'} " + f"(quantile atol={chosen['quantile_atol']:g}, count rtol={chosen['count_rtol']:g})" + ) + if not strict: + diff = { + k: (recorded.get(k), current.get(k)) + for k in sorted(set(recorded) | set(current)) + if recorded.get(k) != current.get(k) + } + banner += ( + f"\n[golden_lebel] environment differs from the recorded baseline " + f"(recorded -> current): {diff}" + ) + warnings.warn(banner.strip(), stacklevel=2) + print(banner) + return chosen + + +def test_golden_json_is_well_formed(golden): + for key in ("fixture", "shapes", "metrics", "quantiles", "tolerances", "regeneration"): + assert key in golden + assert golden["metrics"]["median_score"] == pytest.approx(golden["quantiles"]["50"]) + + +def test_shapes_match_exactly(result, golden): + assert result["shapes"] == {k: list(v) for k, v in golden["shapes"].items()} + + +def test_correlation_quantiles(result, golden, tolerances): + correlations = result["correlations"] + assert correlations.shape == (golden["metrics"]["n_voxels"],) + + atol = tolerances["quantile_atol"] + failures = [] + for q_str, expected in sorted(golden["quantiles"].items(), key=lambda kv: float(kv[0])): + actual = float(np.percentile(correlations, float(q_str))) + if abs(actual - expected) > atol: + failures.append(f" p{q_str}: expected {expected:.10f}, got {actual:.10f} " + f"(delta {actual - expected:+.3e}, atol {atol:g})") + assert not failures, ( + "Correlation quantiles moved:\n" + "\n".join(failures) + + f"\nIf intended, regenerate with:\n {golden['regeneration']['command']}" + ) + + +def test_n_significant_within_band(result, golden, tolerances): + """Banded, not exact: voxels near the BH cutoff flip on tiny correlation changes.""" + expected = golden["metrics"]["n_significant"] + actual = result["n_significant"] + rtol = tolerances["count_rtol"] + lo, hi = expected * (1 - rtol), expected * (1 + rtol) + assert lo <= actual <= hi, ( + f"n_significant={actual} outside [{lo:.0f}, {hi:.0f}] (expected {expected}, rtol {rtol:g}).\n" + f"If intended, regenerate with:\n {golden['regeneration']['command']}" + ) + + +def test_threshold_counts_within_band(result, golden, tolerances): + correlations = result["correlations"] + rtol = tolerances["count_rtol"] + for cut_str, expected in golden["threshold_counts"].items(): + actual = int((correlations > float(cut_str)).sum()) + lo, hi = expected * (1 - rtol), expected * (1 + rtol) + assert lo <= actual <= hi, ( + f"voxels with r > {cut_str}: {actual} outside [{lo:.0f}, {hi:.0f}] " + f"(expected {expected}, rtol {rtol:g}).\n" + f"If intended, regenerate with:\n {golden['regeneration']['command']}" + ) + + +def regenerate() -> int: + """Rewrite the measured fields of the golden JSON; tolerances and notes are left alone.""" + import tempfile + + golden = load_golden() + raw = os.environ.get(ENV_VAR) + if not raw or not Path(raw).is_file(): + print(skip_reason(golden), file=sys.stderr) + return 1 + + actual_md5 = file_md5(Path(raw)) + if actual_md5 != golden["fixture"]["md5"]: + print( + f"note: fixture md5 changed ({golden['fixture']['md5']} -> {actual_md5}); " + f"recording the new one.", + file=sys.stderr, + ) + + import encoding.trainer as trainer_mod + import wandb + from conftest import _NullBrainPlotter, _NullLogger + + trainer_mod.WandBLogger = _NullLogger + trainer_mod.BrainPlotter = _NullBrainPlotter + wandb.init = lambda *a, **k: None + + with tempfile.TemporaryDirectory() as tmp: + result = run_pipeline(raw, tmp) + + correlations = result["correlations"] + golden["fixture"]["md5"] = actual_md5 + golden["run"]["date"] = datetime.date.today().isoformat() + try: + golden["run"]["commit"] = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parent, + capture_output=True, + text=True, + check=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError): + pass + golden["environment"]["fingerprint"] = environment_fingerprint() + golden["shapes"] = result["shapes"] + golden["metrics"]["median_score"] = round(float(np.median(correlations)), 10) + golden["metrics"]["n_significant"] = result["n_significant"] + golden["metrics"]["n_voxels"] = int(correlations.size) + golden["quantiles"] = { + q: round(float(np.percentile(correlations, float(q))), 10) + for q in golden["quantiles"] + } + golden["threshold_counts"] = { + cut: int((correlations > float(cut)).sum()) for cut in golden["threshold_counts"] + } + + with open(GOLDEN_PATH, "w") as f: + json.dump(golden, f, indent=2) + f.write("\n") + + print(f"updated {GOLDEN_PATH}") + print(f" median_score = {golden['metrics']['median_score']}") + print(f" n_significant = {golden['metrics']['n_significant']}") + print("Review the diff before committing.") + return 0 + + +if __name__ == "__main__": + if "--update" in sys.argv: + sys.path.insert(0, str(Path(__file__).parent)) + sys.exit(regenerate()) + print(__doc__) + print("Run the tests with pytest; pass --update to regenerate the golden values.") + sys.exit(0) diff --git a/tests/test_pipeline_synthetic.py b/tests/test_pipeline_synthetic.py new file mode 100644 index 0000000..a8cd327 --- /dev/null +++ b/tests/test_pipeline_synthetic.py @@ -0,0 +1,436 @@ +"""End-to-end tests for the encoding pipeline against a synthetic assembly. + +Checks structure, recovery of a signal planted at FIR lags 2 and 3 in vertices +0..9, and that the held-out story cannot leak into training. Runs offline. +""" + +import numpy as np +import pytest + +from conftest import ( + BRAIN_TRIM_TAIL, + BRAIN_TRIM_TOTAL, + FIR_DELAYS, + N_FEATURES, + N_SIGNAL_VERTICES, + N_VERTICES, + PLANTED_LAGS, + TRIMMING_CONFIG, + SyntheticTokenFeatureExtractor, + reference_legacy_average, + reference_make_delayed, + _NullBrainPlotter, + _NullLogger, +) + +from encoding.downsample.downsampling import Downsampler +from encoding.features.FIR_expander import FIR +from encoding.models.nested_cv import NestedCVModel +from encoding.trainer import AbstractTrainer + +MODEL_KWARGS = dict(alphas=[1.0, 10.0, 100.0, 1000.0], n_inner_folds=2, use_gpu=False) + + +def make_trainer(assembly, results_dir, **overrides): + """Build an AbstractTrainer wired the way train_simple.py wires it.""" + kwargs = dict( + assembly=assembly, + feature_extractors=[SyntheticTokenFeatureExtractor()], + downsampler=Downsampler(), + model=NestedCVModel(model_name="ridge_regression"), + fir_delays=FIR_DELAYS, + trimming_config=dict(TRIMMING_CONFIG), + use_train_test_split=True, + logger_backend="wandb", + dataset_type="synthetic", + results_dir=str(results_dir), + downsample_config={"method": "legacy_average"}, + ) + kwargs.update(overrides) + return AbstractTrainer(**kwargs) + + +@pytest.fixture(scope="module") +def pipeline_data(assembly, tmp_path_factory): + """Module-scoped, so conftest's function-scoped autouse patching is not yet + active; patch explicitly for the duration of trainer construction.""" + results_dir = tmp_path_factory.mktemp("results") + with pytest.MonkeyPatch.context() as mp: + import encoding.trainer as trainer_mod + import wandb + + mp.setattr(trainer_mod, "WandBLogger", _NullLogger) + mp.setattr(trainer_mod, "BrainPlotter", _NullBrainPlotter) + mp.setattr(wandb, "init", lambda *a, **k: None, raising=False) + + trainer = make_trainer(assembly, results_dir) + features = trainer.extract_and_downsample_features() + delayed = trainer.apply_fir_delays(features) + data = trainer.structure_data(delayed) + + return { + "trainer": trainer, + "features": features, + "delayed": delayed, + "data": data, + "results_dir": results_dir, + } + + +@pytest.fixture(scope="module") +def fitted(pipeline_data): + data = pipeline_data["data"] + model = pipeline_data["trainer"].model + metrics, weights, alphas = model.fit_predict( + features=data["Rstim"], + targets=data["Rresp"], + X_test=data["Pstim"], + y_test=data["Presp"], + **MODEL_KWARGS, + ) + return {"metrics": metrics, "weights": weights, "alphas": alphas} + + +def test_fixture_matches_lebel_conventions(assembly, truth): + total_brain_rows = 0 + for name in assembly.stories: + sd = assembly.story_data[name] + st = truth.stories[name] + + assert len(sd.split_indices) + 1 == len(sd.tr_times) + assert sd.brain_data.shape[0] == st.n_trs - BRAIN_TRIM_TOTAL + assert len(sd.words) == st.n_words + assert len(sd.stimuli) == st.n_words + assert len(sd.data_times) == st.n_words + assert sd.word_rates.shape == (st.n_trs, 1) + total_brain_rows += sd.brain_data.shape[0] + + assert assembly.shape == (total_brain_rows, N_VERTICES) + + +def test_legacy_average_matches_independent_reference(assembly, truth): + ds = Downsampler() + for name in assembly.stories: + sd = assembly.story_data[name] + st = truth.stories[name] + + out = ds.downsample( + data=st.word_features, + data_times=sd.data_times, + tr_times=sd.tr_times, + method="legacy_average", + split_indices=sd.split_indices, + ) + assert out.shape == (st.n_trs, N_FEATURES) + np.testing.assert_allclose(out, reference_legacy_average(st.word_features, st.counts)) + + +def test_legacy_average_leaves_empty_trs_as_zero_rows(assembly, truth): + """Empty TRs come back as exact zero rows, not NaN.""" + ds = Downsampler() + name = assembly.stories[0] + sd = assembly.story_data[name] + st = truth.stories[name] + + out = ds.downsample( + data=st.word_features, + data_times=sd.data_times, + tr_times=sd.tr_times, + method="legacy_average", + split_indices=sd.split_indices, + ) + + assert st.empty_trs.size > 0, "fixture should contain some silent TRs" + assert np.all(out[st.empty_trs] == 0.0) + assert np.isfinite(out).all() + + occupied = np.setdiff1d(np.arange(st.n_trs), st.empty_trs) + assert np.any(out[occupied] != 0.0) + + +@pytest.mark.parametrize("method", ["legacy_average", "legacy_sum", "legacy_last"]) +def test_legacy_downsamplers_return_one_row_per_tr(assembly, truth, method): + ds = Downsampler() + name = assembly.stories[0] + sd = assembly.story_data[name] + st = truth.stories[name] + + out = ds.downsample( + data=st.word_features, + data_times=sd.data_times, + tr_times=sd.tr_times, + method=method, + split_indices=sd.split_indices, + ) + assert out.shape == (st.n_trs, N_FEATURES) + assert out.shape[0] == len(sd.split_indices) + 1 + + +@pytest.mark.xfail( + strict=True, + reason="non-legacy downsamplers read split_indices as per-word TR labels " + "and return ~n_words rows", +) +@pytest.mark.parametrize("method", ["average", "sum", "last"]) +def test_nonlegacy_downsamplers_return_one_row_per_tr(assembly, truth, method): + """Known bug: non-legacy methods read split_indices as per-word TR labels and + return max(split_indices) + 1 rows (~n_words) instead of n_trs.""" + ds = Downsampler() + name = assembly.stories[0] + sd = assembly.story_data[name] + st = truth.stories[name] + + out = ds.downsample( + data=st.word_features, + data_times=sd.data_times, + tr_times=sd.tr_times, + method=method, + split_indices=sd.split_indices, + ) + assert out.shape == (st.n_trs, N_FEATURES) + + +@pytest.mark.parametrize( + "method", ["legacy_average", "legacy_sum", "legacy_last", "average", "sum", "last"] +) +def test_downsamplers_require_split_indices(assembly, truth, method): + ds = Downsampler() + sd = assembly.story_data[assembly.stories[0]] + st = truth.stories[assembly.stories[0]] + + with pytest.raises(ValueError, match="split_indices"): + ds.downsample( + data=st.word_features, + data_times=sd.data_times, + tr_times=sd.tr_times, + method=method, + ) + + +def test_fir_expansion_matches_reference(assembly, truth): + st = truth.stories[assembly.stories[0]] + delayed = FIR.make_delayed(st.tr_features, FIR_DELAYS) + + assert delayed.shape == (st.n_trs, N_FEATURES * len(FIR_DELAYS)) + np.testing.assert_allclose(delayed, reference_make_delayed(st.tr_features, FIR_DELAYS)) + + for block, d in enumerate(FIR_DELAYS): + cols = slice(block * N_FEATURES, (block + 1) * N_FEATURES) + assert np.all(delayed[:d, cols] == 0.0) + np.testing.assert_allclose(delayed[d:, cols], st.tr_features[:-d]) + + +def test_structure_data_shapes_and_story_split(pipeline_data, assembly, truth): + features = pipeline_data["features"] + delayed = pipeline_data["delayed"] + data = pipeline_data["data"] + + for name in assembly.stories: + st = truth.stories[name] + assert features[name].shape == (st.n_trs, N_FEATURES) + assert delayed[name].shape == (st.n_trs, N_FEATURES * len(FIR_DELAYS)) + + train_stories = assembly.stories[:-1] + test_story = assembly.stories[-1] + + expected_train = sum(truth.stories[s].n_trs - BRAIN_TRIM_TOTAL for s in train_stories) + expected_test = truth.stories[test_story].n_trs - ( + TRIMMING_CONFIG["test_features_start"] + BRAIN_TRIM_TAIL + ) + + assert data["Rstim"].shape == (expected_train, N_FEATURES * len(FIR_DELAYS)) + assert data["Rresp"].shape == (expected_train, N_VERTICES) + assert data["Pstim"].shape == (expected_test, N_FEATURES * len(FIR_DELAYS)) + assert data["Presp"].shape == (expected_test, N_VERTICES) + + # _create_train_test_split does not check that features and targets line up. + assert data["Rstim"].shape[0] == data["Rresp"].shape[0] + assert data["Pstim"].shape[0] == data["Presp"].shape[0] + assert np.isfinite(data["Rstim"]).all() + assert np.isfinite(data["Pstim"]).all() + + +def test_recovers_planted_signal_vertices(fitted, truth): + correlations = np.asarray(fitted["metrics"]["correlations"]) + assert correlations.shape == (N_VERTICES,) + + signal_r = correlations[truth.signal_vertices] + # Ceiling is ~0.89 at the fixture's default noise_ratio of 0.5. + assert np.median(signal_r) > 0.75 + assert signal_r.min() > 0.6 + + +def test_noise_vertices_are_not_predicted(fitted, truth): + correlations = np.asarray(fitted["metrics"]["correlations"]) + noise_r = np.abs(correlations[truth.noise_vertices]) + + assert np.median(noise_r) < 0.12 + assert np.percentile(noise_r, 99) < 0.35 + signal_r = correlations[truth.signal_vertices] + assert signal_r.min() > np.percentile(noise_r, 99.9) + + +def test_fdr_selects_the_planted_vertices(fitted, truth): + significant = np.asarray(fitted["metrics"]["significant_mask"], dtype=bool) + assert significant[truth.signal_vertices].all() + assert significant.sum() < 5 * N_SIGNAL_VERTICES + + +def test_recovered_kernel_peaks_at_planted_lags(fitted, truth): + weights = fitted["weights"] + assert weights.shape == (N_FEATURES * len(FIR_DELAYS), N_VERTICES) + + # (n_delays, n_features, n_vertices) -- make_delayed hstacks blocks in delay order. + by_lag = weights.reshape(len(FIR_DELAYS), N_FEATURES, N_VERTICES) + signal_norms = { + delay: float(np.linalg.norm(by_lag[i][:, truth.signal_vertices])) + for i, delay in enumerate(FIR_DELAYS) + } + + planted = [signal_norms[d] for d in PLANTED_LAGS] + empty = [signal_norms[d] for d in FIR_DELAYS if d not in PLANTED_LAGS] + + assert min(planted) > 3 * max(empty), ( + f"kernel did not peak at lags {PLANTED_LAGS}: {signal_norms}" + ) + + noise_norm = float(np.linalg.norm(by_lag[:, :, truth.noise_vertices])) + per_vertex_noise = noise_norm / np.sqrt(len(truth.noise_vertices)) + per_vertex_signal = min(planted) / np.sqrt(N_SIGNAL_VERTICES) + assert per_vertex_signal > per_vertex_noise + + +def test_trainer_train_runs_offline_and_reports_metrics(assembly, tmp_path): + trainer = make_trainer(assembly, tmp_path) + assert isinstance(trainer.experiment_logger, _NullLogger) + assert isinstance(trainer.brain_plotter, _NullBrainPlotter) + + metrics = trainer.train(**MODEL_KWARGS) + + correlations = np.asarray(metrics["correlations"]) + assert np.median(correlations[:N_SIGNAL_VERTICES]) > 0.75 + assert metrics["median_score"] == pytest.approx(float(np.median(correlations))) + assert 0 < metrics["n_significant"] < 5 * N_SIGNAL_VERTICES + + assert "median_correlation" in trainer.experiment_logger.scalars + assert trainer.brain_plotter.calls == 1 + assert any(tmp_path.rglob("hyperparams.json")) + + +def test_fir_is_causal(assembly, truth): + """Perturbing TR t never changes delayed rows before t.""" + st = truth.stories[assembly.stories[0]] + base = FIR.make_delayed(st.tr_features, FIR_DELAYS) + + perturb_at = st.n_trs // 2 + bumped = st.tr_features.copy() + bumped[perturb_at] += 1000.0 + after = FIR.make_delayed(bumped, FIR_DELAYS) + + changed_rows = np.flatnonzero(np.any(base != after, axis=1)) + assert changed_rows.size > 0 + assert changed_rows.min() > perturb_at, ( + "a future timepoint leaked into an earlier delayed row" + ) + assert set(changed_rows.tolist()) == {perturb_at + d for d in FIR_DELAYS} + + +def test_held_out_story_cannot_influence_training_matrix(pipeline_data, assembly): + """NaN-poisoning the held-out story's features and brain data leaves the + training matrices bit-identical.""" + trainer = pipeline_data["trainer"] + baseline = pipeline_data["data"] + test_story = assembly.stories[-1] + + poisoned = {k: v.copy() for k, v in pipeline_data["delayed"].items()} + poisoned[test_story][:] = np.nan + + story = assembly.story_data[test_story] + clean_brain = story.brain_data + try: + story.brain_data = np.full_like(clean_brain, np.nan) + data = trainer.structure_data(poisoned) + finally: + story.brain_data = clean_brain + + assert np.isfinite(data["Rstim"]).all() + assert np.isfinite(data["Rresp"]).all() + np.testing.assert_array_equal(data["Rstim"], baseline["Rstim"]) + np.testing.assert_array_equal(data["Rresp"], baseline["Rresp"]) + + assert not np.array_equal(data["Pstim"], baseline["Pstim"]) + assert np.isnan(data["Presp"]).all() + + +def test_nan_features_are_silently_zeroed(pipeline_data, assembly): + """_create_train_test_split applies np.nan_to_num to features but not targets, so + NaN features become zeros and NaN targets stay NaN; neither raises.""" + trainer = pipeline_data["trainer"] + test_story = assembly.stories[-1] + + poisoned = {k: v.copy() for k, v in pipeline_data["delayed"].items()} + poisoned[test_story][:] = np.nan + + data = trainer.structure_data(poisoned) + + assert np.all(data["Pstim"] == 0.0) + assert np.isfinite(data["Presp"]).all() + + +def test_each_split_is_zscored_independently(pipeline_data, assembly): + """Each story is z-scored within its split; the test split uses its own + statistics, not the training set's.""" + trainer = pipeline_data["trainer"] + delayed = pipeline_data["delayed"] + baseline = pipeline_data["data"] + + scaled = {k: v.copy() for k, v in delayed.items()} + for story in assembly.stories[:-1]: + scaled[story] = scaled[story] * 50.0 + 7.0 + + data = trainer.structure_data(scaled) + + np.testing.assert_allclose(data["Pstim"], baseline["Pstim"]) + np.testing.assert_allclose(data["Rstim"], baseline["Rstim"], atol=1e-8) + + assert np.allclose(baseline["Pstim"].mean(axis=0), 0.0, atol=1e-8) + assert np.allclose(baseline["Pstim"].std(axis=0), 1.0, atol=1e-6) + + +def test_determinism(pipeline_data, monkeypatch): + """Compares correlation vectors because a different fold seed barely moves the + median (~5e-6) and leaves n_significant unchanged on this fixture.""" + data = pipeline_data["data"] + model = pipeline_data["trainer"].model + + def fit(): + metrics, _, alphas = model.fit_predict( + features=data["Rstim"], + targets=data["Rresp"], + X_test=data["Pstim"], + y_test=data["Presp"], + **MODEL_KWARGS, + ) + return np.asarray(metrics["correlations"]), np.asarray(alphas) + + first, first_alphas = fit() + second, second_alphas = fit() + + assert np.allclose(first, second, rtol=0, atol=1e-12) + np.testing.assert_array_equal(first_alphas, second_alphas) + + import encoding.models.nested_cv as nested_cv_mod + from encoding.models.folding import create_folds as real_create_folds + + # fit_predict takes no seed; _find_best_alphas calls create_folds with the default. + monkeypatch.setattr( + nested_cv_mod, + "create_folds", + lambda *a, **k: real_create_folds(*a, **{**k, "seed": 7}), + ) + other, other_alphas = fit() + + assert not np.allclose(first, other, rtol=0, atol=1e-12), ( + "fold seed 7 gave the same correlations as seed 0" + ) + assert (first_alphas != other_alphas).any(), "no voxel changed its selected alpha" diff --git a/train_simple.py b/train_simple.py index 463bba9..6b46caa 100644 --- a/train_simple.py +++ b/train_simple.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import logging +import os from encoding.assembly.assembly_loader import load_assembly from encoding.features.factory import FeatureExtractorFactory @@ -14,7 +15,10 @@ def main(): logger = logging.getLogger(__name__) # 1) Load the packaged assembly - assembly_path = "/storage/coda1/p-aivanova7/0/shared/litcoder_core/scripts/assembly_lebel_uts03.pkl" + assembly_path = os.environ.get( + "LITCODER_LEBEL_ASSEMBLY", + "/storage/coda1/p-aivanova7/0/shared/litcoder_core/scripts/assembly_lebel_uts03.pkl", + ) logger.info(f"Loading assembly from {assembly_path}") assembly = load_assembly(assembly_path)