diff --git a/benchmark.py b/benchmark.py index 40ca22f..2552356 100644 --- a/benchmark.py +++ b/benchmark.py @@ -3,9 +3,15 @@ import numpy as np import torch -from utils.dataset import TVSD_Dataset +from utils.dataset import TVSD_Dataset, TVSD_TestDataset from utils.load_model import load_model -from utils.brainscore import compute_brain_score +from utils.brainscore import score_train_test + + +def _load_activations(path, device): + """Load saved activations and flatten to (n_stimuli, features).""" + a = torch.load(path, map_location=device) + return a.reshape(a.shape[0], -1).detach().cpu().float().numpy() def main(args): @@ -13,54 +19,78 @@ def main(args): print(f"Using device: {device}") model, model_name, _ = load_model(args.model_config) - # layers = [name for name, module in model.named_modules() if 'relu' not in name] - tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region) + # Train responses (single-trial train_MUA) fit the mapping; the held-out test + # responses (repetition-averaged test_MUA) are what we score against. Neuron + # selection uses the single-trial reliability (`reliab`); the noise ceiling is + # the split-half + Spearman-Brown internal consistency of the 30-rep average. + train_ds = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region) + test_ds = TVSD_TestDataset( + root_dir=args.root_dir, + monkey=args.monkey, + region=args.region, + recompute_reliability=True, + spearman_brown=True, + n_boot=args.ceiling_n_boot, + random_state=args.random_state, + ) + + single_trial_reliab = train_ds.reliability.numpy() # (C,) + Y_train_all = np.asarray(train_ds.responses) # (n_train, C) + Y_test_all = np.asarray(test_ds.responses).mean(axis=0) # (n_test, C) = test_MUA + ceiling_all = test_ds.reliability.numpy() # (C,) SB ceiling + + mask = single_trial_reliab > args.reliability_threshold + Y_train_all = Y_train_all[:, mask] + Y_test_all = Y_test_all[:, mask] + ceiling = ceiling_all[mask] + print( + f"{int(mask.sum())} neuroids retained with single-trial reliability " + f"> {args.reliability_threshold} (median ceiling {np.nanmedian(ceiling):.3f})" + ) + + train_dir = f"{args.output_dir}/activations/TVSD_train/{model_name}" + test_dir = f"{args.output_dir}/activations/TVSD_test/{model_name}" layer_scores = {} - activation_dir = f"{args.output_dir}/activations/TVSD/{model_name}" - layers = os.listdir(activation_dir) - for layer in layers: + for layer in sorted(os.listdir(train_dir)): print(f"===== EVALUATING LAYER: {layer} =========") - layer_dir = f"{activation_dir}/{layer}" - activation_path = f"{layer_dir}/activations.pt" - if not os.path.exists(activation_path): - print(f"Activations for layer {layer} not found at {activation_path}. Skipping.") - continue - activations = torch.load(f"{layer_dir}/activations.pt", map_location=device) - if activations is not None: - print(f"Layer: {layer}, Activations shape: {activations.shape}") - else: - print(f"No activations found for layer: {layer}") + train_path = f"{train_dir}/{layer}/activations.pt" + test_path = f"{test_dir}/{layer}/activations.pt" + if not (os.path.exists(train_path) and os.path.exists(test_path)): + print(f"Missing train/test activations for layer {layer}. Skipping.") continue - # float32: float16 overflows StandardScaler's in-place divide to inf. - activations = ( - activations.reshape(activations.shape[0], -1).detach().cpu().float().numpy() - ) # [B, H * W * C] - neural_responses, reliability = tvsd_dataset[: activations.shape[0]] - reliability_mask = reliability > args.reliability_threshold - neural_responses = neural_responses[:, reliability_mask] # [B, C'] + X_train = _load_activations(train_path, device) + X_test = _load_activations(test_path, device) + + # Align rows with responses in case generation was truncated (e.g. --max_batches). + n_tr = min(X_train.shape[0], Y_train_all.shape[0]) + n_te = min(X_test.shape[0], Y_test_all.shape[0]) + X_train, Y_train = X_train[:n_tr], Y_train_all[:n_tr] + X_test, Y_test = X_test[:n_te], Y_test_all[:n_te] + print(f"train X{X_train.shape} Y{Y_train.shape} | test X{X_test.shape} Y{Y_test.shape}") if args.noise_test: - activations = np.random.normal( - size=activations.shape - ) # Use random noise for null results + X_train = np.random.normal(size=X_train.shape) + X_test = np.random.normal(size=X_test.shape) if args.permutation_test: - np.random.shuffle(neural_responses) - - print( - f"{neural_responses.shape[1]} neural responses retained with reliability > {args.reliability_threshold}" - ) + np.random.shuffle(Y_test) - layer_score, layer_std = compute_brain_score( - X=activations, - Y=neural_responses, - n_splits=args.n_splits, + layer_score, layer_std = score_train_test( + X_train, + Y_train, + X_test, + Y_test, reducer=args.reducer, correlation_fn=args.correlation_fn, pca_components=args.pca_components, - preprocessed=args.preprocessed, + skip_pca=args.skip_pca, + standardize=args.standardize, + ceiling=ceiling, + ceiling_normalize=args.ceiling_normalize, + n_boot_ci=args.n_boot_ci, + random_state=args.random_state, ) layer_scores[layer] = {"score": layer_score, "std": layer_std} print(f"Score: {layer_score}, Std: {layer_std}") @@ -108,60 +138,76 @@ def main(args): "--output_dir", type=str, default=f"{os.getcwd()}/outputs", - help="Directory to save activations.", + help="Directory holding generated activations and results.", ) parser.add_argument( "--reliability_threshold", type=float, default=0.3, - help="Reliability threshold for neural responses.", + help="Keep neuroids whose single-trial reliability exceeds this (TVSD-paper selection).", ) parser.add_argument( "--reducer", type=str, default="median", choices=["mean", "median"], - help="Reduction method for brain score.", + help="Reduction across neuroids.", ) parser.add_argument( "--correlation_fn", type=str, default="pearson", choices=["pearson", "spearman"], - help="Correlation function to use for brain score computation.", + help="Correlation function for predictivity.", ) parser.add_argument( - "--n_splits", + "--pca_components", type=int, - default=5, - help="Number of splits for KFold cross-validation.", + default=100, + help="PCA components to reduce features to (fit on train, applied to test).", ) parser.add_argument( - "--pca_components", + "--skip_pca", + action="store_true", + help="Skip in-benchmark PCA (e.g. features were already reduced at generation).", + ) + parser.add_argument( + "--standardize", + action="store_true", + help="Z-score features/targets (fit on train). Off by default to match Brain-Score.", + ) + parser.add_argument( + "--ceiling_normalize", + action="store_true", + help="Normalize scores by the median noise ceiling (Brain-Score-style).", + ) + parser.add_argument( + "--ceiling_n_boot", type=int, - default=100, - help="Number of PCA components to use.", + default=30, + help="Bootstrap splits for the split-half+SB noise ceiling from test reps.", ) parser.add_argument( - "--skip_interval", + "--n_boot_ci", type=int, - default=1, - help="Skip every n-th image in the dataset.", + default=100, + help="Bootstrap resamples over test stimuli for the score's std.", ) parser.add_argument( - "--preprocessed", - action="store_true", - help="Whether the data is preprocessed (scaled and PCA applied).", + "--random_state", + type=int, + default=42, + help="Random seed for the ceiling and bootstrap.", ) parser.add_argument( "--noise_test", action="store_true", - help="Run with pure noise to test. Useful for debugging.", + help="Replace activations with pure noise (null control).", ) parser.add_argument( "--permutation_test", action="store_true", - help="Randomly permute neural responses. Useful for debugging.", + help="Randomly permute test responses (null control).", ) args = parser.parse_args() diff --git a/generate_activations.py b/generate_activations.py index 19592c7..079f0ea 100644 --- a/generate_activations.py +++ b/generate_activations.py @@ -4,7 +4,7 @@ import torch from torch.utils.data import DataLoader -from utils.dataset import TVSD_Dataset +from utils.dataset import TVSD_Dataset, TVSD_TestDataset from utils.hooks import Activations from utils.load_model import load_model, resolve_transform @@ -19,36 +19,49 @@ def main(args): layers = [name for name, module in model.named_modules() if "relu" not in name] hook_layers = [layers[i] for i in range(0, len(layers), hook_interval)] + # Activations are saved per split under TVSD_train / TVSD_test so the benchmark + # can fit on train and score on the held-out test images. + dataset_name = f"TVSD_{args.split}" activations = Activations( output_dir=args.output_dir, model_name=model_name, - dataset_name="TVSD", + dataset_name=dataset_name, pca_components=args.pca_components, ) activations.register(model, hook_layers) - tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region) + if args.split == "train": + tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region) + else: + tvsd_dataset = TVSD_TestDataset( + root_dir=args.root_dir, monkey=args.monkey, region=args.region + ) things_dataset = tvsd_dataset.get_things(things_path=args.things_path, transform=transform) things_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=False) - shuffled_things_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=True) if args.pca_components is not None: - print("Training IPCA models...") - activations.set_training_mode(True) - - for i, batch in tqdm(enumerate(shuffled_things_loader), desc="Training IPCA"): - if args.max_pca_train_batches and i >= int(args.max_pca_train_batches): - break - activations.set_batch(i) - with torch.no_grad(): - _ = model(batch.to(device)) - activations.finalize_batch_training() - - torch.cuda.empty_cache() - activations.save_ipca_models() - print("IPCA training completed.") - - print("Generating activations...") + if args.split == "train": + # Fit the IPCA basis on the train images only. + print("Training IPCA models on train split...") + activations.set_training_mode(True) + shuffled_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=True) + for i, batch in tqdm(enumerate(shuffled_loader), desc="Training IPCA"): + if args.max_pca_train_batches and i >= int(args.max_pca_train_batches): + break + activations.set_batch(i) + with torch.no_grad(): + _ = model(batch.to(device)) + activations.finalize_batch_training() + torch.cuda.empty_cache() + activations.save_ipca_models() + activations.set_training_mode(False) + print("IPCA training completed.") + else: + # Apply the train-fit IPCA to the test split (shared basis, no leakage). + train_models = f"{args.output_dir}/activations/TVSD_train/{model_name}" + activations.load_ipca_models(train_models) + + print(f"Generating {args.split} activations...") activations.set_training_mode(False) for i, batch in tqdm(enumerate(things_loader), desc="Generating activations"): @@ -92,6 +105,14 @@ def main(args): choices=["V1", "V4", "IT"], help="Which brain region to use.", ) + parser.add_argument( + "--split", + type=str, + default="train", + choices=["train", "test"], + help="Which image split to generate activations for. 'test' reuses the " + "train-fit IPCA basis (generate train first).", + ) parser.add_argument( "--things_path", type=str, diff --git a/scripts/benchmark.sh b/scripts/benchmark.sh index e2467b0..0b9f6ad 100644 --- a/scripts/benchmark.sh +++ b/scripts/benchmark.sh @@ -11,9 +11,12 @@ regions=("V1" "V4" "IT") export PYTHONPATH="$PYTHONPATH:$(pwd)" module load cuda cudnn +# Fit on train activations, score on the held-out test images. --skip_pca because +# activations were already reduced by the train-fit IPCA at generation time (the +# same basis is applied to the test images, so it is leak-free). python -u benchmark.py --model_config $1 \ --monkey monkeyF \ --region ${regions[$SLURM_ARRAY_TASK_ID]} \ - --n_splits 4 \ - --preprocessed - + --skip_pca \ + --ceiling_normalize + diff --git a/scripts/generate_activations.sh b/scripts/generate_activations.sh index f2e9929..3787757 100644 --- a/scripts/generate_activations.sh +++ b/scripts/generate_activations.sh @@ -11,8 +11,18 @@ export PYTHONPATH="$PYTHONPATH:$(pwd)" module load cuda cudnn + +# Train split first: fits the IPCA basis and saves reduced train activations. python -u generate_activations.py --model_config $1 \ + --split train \ --monkey monkeyF \ --batch_size 128 \ --pca_components 100 \ --max_pca_train_batches 4 + +# Test split: reuses the train-fit IPCA (no refitting -> leak-free shared basis). +python -u generate_activations.py --model_config $1 \ + --split test \ + --monkey monkeyF \ + --batch_size 128 \ + --pca_components 100 diff --git a/tests/test_brainscore.py b/tests/test_brainscore.py new file mode 100644 index 0000000..48c3170 --- /dev/null +++ b/tests/test_brainscore.py @@ -0,0 +1,133 @@ +"""Tests for utils.brainscore scoring helpers.""" + +import numpy as np +import pytest + +from utils.brainscore import ( + compute_brain_score, + score_train_test, + spearman_brown, + _make_splitter, +) + + +class TestSpearmanBrown: + def test_known_values(self): + # 2r/(1+r) + assert spearman_brown(0.5) == pytest.approx(2 * 0.5 / 1.5) + assert spearman_brown(0.0) == pytest.approx(0.0) + assert spearman_brown(1.0) == pytest.approx(1.0) + + def test_inflates_positive_correlation(self): + # For r in (0, 1), the SB-corrected value exceeds the raw half-split r. + for r in [0.1, 0.3, 0.6, 0.9]: + assert spearman_brown(r) > r + + def test_n_parts_generalization(self): + # n=1 is the identity. + assert spearman_brown(0.4, n=1) == pytest.approx(0.4) + + +class TestMakeSplitter: + def test_shuffle(self): + from sklearn.model_selection import ShuffleSplit + + sp = _make_splitter("shuffle", n_splits=10, train_size=0.9, random_state=0) + assert isinstance(sp, ShuffleSplit) + assert sp.get_n_splits() == 10 + + def test_kfold(self): + from sklearn.model_selection import KFold + + sp = _make_splitter("kfold", n_splits=5, train_size=0.9, random_state=0) + assert isinstance(sp, KFold) + + def test_unknown_raises(self): + with pytest.raises(ValueError, match="Unknown cv_strategy"): + _make_splitter("bogus", n_splits=5, train_size=0.9, random_state=0) + + +class TestComputeBrainScore: + def _linear_data(self): + # Y is a linear function of the first few features -> PLS predicts well. + rng = np.random.default_rng(0) + X = rng.standard_normal((120, 20)) + W = rng.standard_normal((20, 5)) + Y = X @ W + rng.standard_normal((120, 5)) * 0.05 + return X, Y + + def test_predicts_correlated_data(self): + X, Y = self._linear_data() + score, std = compute_brain_score(X, Y) + # Strong linear relationship -> high positive predictivity. + assert score > 0.8 + assert std >= 0 + + def test_ceiling_normalization_scales_score(self): + X, Y = self._linear_data() + raw, _ = compute_brain_score(X, Y) + ceiling = np.full(Y.shape[1], 0.5) + normed, _ = compute_brain_score(X, Y, ceiling=ceiling, ceiling_normalize=True) + # Dividing by a median ceiling of 0.5 doubles the score. + assert normed == pytest.approx(raw / 0.5, rel=1e-6) + + def test_ceiling_normalize_requires_ceiling(self): + X, Y = self._linear_data() + with pytest.raises(ValueError, match="requires a ceiling"): + compute_brain_score(X, Y, ceiling_normalize=True) + + def test_reproducible(self): + X, Y = self._linear_data() + s1, _ = compute_brain_score(X, Y, random_state=7) + s2, _ = compute_brain_score(X, Y, random_state=7) + assert s1 == pytest.approx(s2) + + +class TestScoreTrainTest: + def _split_data(self): + # Train and test share one linear map -> fit-on-train predicts test well. + rng = np.random.default_rng(0) + W = rng.standard_normal((20, 5)) + X_train = rng.standard_normal((500, 20)) + Y_train = X_train @ W + rng.standard_normal((500, 5)) * 0.05 + X_test = rng.standard_normal((100, 20)) + Y_test = X_test @ W + rng.standard_normal((100, 5)) * 0.05 + return X_train, Y_train, X_test, Y_test + + def test_predicts_heldout_test(self): + X_tr, Y_tr, X_te, Y_te = self._split_data() + score, std = score_train_test(X_tr, Y_tr, X_te, Y_te) + assert score > 0.8 + assert std >= 0 + + def test_ceiling_normalization_scales_score(self): + X_tr, Y_tr, X_te, Y_te = self._split_data() + raw, _ = score_train_test(X_tr, Y_tr, X_te, Y_te) + ceiling = np.full(Y_te.shape[1], 0.5) + normed, _ = score_train_test( + X_tr, Y_tr, X_te, Y_te, ceiling=ceiling, ceiling_normalize=True + ) + assert normed == pytest.approx(raw / 0.5, rel=1e-6) + + def test_ceiling_normalize_requires_ceiling(self): + X_tr, Y_tr, X_te, Y_te = self._split_data() + with pytest.raises(ValueError, match="requires a ceiling"): + score_train_test(X_tr, Y_tr, X_te, Y_te, ceiling_normalize=True) + + def test_std_reproducible_and_nonzero(self): + X_tr, Y_tr, X_te, Y_te = self._split_data() + s1, std1 = score_train_test(X_tr, Y_tr, X_te, Y_te, random_state=7) + s2, std2 = score_train_test(X_tr, Y_tr, X_te, Y_te, random_state=7) + assert s1 == pytest.approx(s2) + assert std1 == pytest.approx(std2) # deterministic bootstrap + assert std1 > 0 # bootstrap over test stimuli gives real spread + + def test_noise_features_score_near_zero(self): + # Features unrelated to responses -> predictivity near 0 on held-out test. + rng = np.random.default_rng(1) + X_tr = rng.standard_normal((500, 20)) + Y_tr = rng.standard_normal((500, 5)) + X_te = rng.standard_normal((100, 20)) + Y_te = rng.standard_normal((100, 5)) + score, _ = score_train_test(X_tr, Y_tr, X_te, Y_te) + assert abs(score) < 0.2 diff --git a/tests/test_reliability.py b/tests/test_reliability.py index 1a28f3e..5339374 100644 --- a/tests/test_reliability.py +++ b/tests/test_reliability.py @@ -213,6 +213,26 @@ def test_reliability_increasing_reps(self): # (though not strictly monotonic due to bootstrapping randomness) assert reliabilities[-1] >= reliabilities[0] - 0.1 # Allow some tolerance + def test_spearman_brown_inflates_positive_reliability(self): + """SB correction should raise a moderate positive split-half reliability.""" + np.random.seed(42) + n_reps, n_stim, n_neu = 30, 100, 5 + base = np.random.randn(n_stim, n_neu) + # Moderate noise -> split-half r safely in (0, 1), where 2r/(1+r) > r. + data = np.array([base + np.random.randn(n_stim, n_neu) * 0.8 for _ in range(n_reps)]) + + with patch.object(TVSD_TestDataset, "_get_paths", return_value=[]): + with patch.object(TVSD_TestDataset, "_get_responses", return_value=(None, None)): + dataset = TVSD_TestDataset(monkey="monkeyF", region="V1") + raw = dataset._compute_reliability(data, n_boot=20, random_state=42) + corrected = dataset._compute_reliability( + data, n_boot=20, random_state=42, spearman_brown=True + ) + + assert np.all(raw > 0) + assert np.all(corrected > raw) + assert np.all(corrected <= 1.0 + 1e-9) + class TestReliabilityIntegration: """Integration tests for reliability in dataset context.""" diff --git a/utils/brainscore.py b/utils/brainscore.py index 30d5bfe..f0dd2e0 100644 --- a/utils/brainscore.py +++ b/utils/brainscore.py @@ -1,12 +1,22 @@ import numpy as np from time import time from scipy.stats import pearsonr, spearmanr -from sklearn.model_selection import KFold +from sklearn.model_selection import KFold, ShuffleSplit from sklearn.cross_decomposition import PLSRegression from sklearn.preprocessing import StandardScaler from sklearn.decomposition import PCA +def spearman_brown(r, n: int = 2): + """Spearman-Brown correction. + + Extrapolates the reliability of a measurement composed of `n` equal parts + from the correlation `r` between single parts. With `n=2` (split-half) this + is the `2r / (1 + r)` form used by Brain-Score's InternalConsistency ceiling. + """ + return n * r / (1 + (n - 1) * r) + + def brain_score_pearsonr(Y_pred, Y_test): """ Compute the Pearson's correlation between the predicted and actual labels. @@ -61,70 +71,178 @@ def brain_score_spearman(Y_pred, Y_test): return spearman_correlations +def _make_splitter(cv_strategy, n_splits, train_size, random_state): + """Build the cross-validation splitter. + + Brain-Score's default is a ShuffleSplit with 10 splits and train_size=0.9, + which we mirror here. KFold is kept as an option for non-overlapping folds. + """ + if cv_strategy == "shuffle": + return ShuffleSplit(n_splits=n_splits, train_size=train_size, random_state=random_state) + elif cv_strategy == "kfold": + return KFold(n_splits=n_splits, shuffle=True, random_state=random_state) + else: + raise ValueError(f"Unknown cv_strategy: {cv_strategy}") + + +def _fit_predict(X_train, Y_train, X_test, pca_components=100, skip_pca=False, standardize=False): + """Fit the (optional standardize -> optional PCA -> PLS) pipeline on the + training split and predict the test split. + + The scaler and PCA are fit on ``X_train``/``Y_train`` only and applied to the + test split, so there is no train->test leakage in the mapping. Returns the + predictions on the test split, shape (n_test_stimuli, n_neuroids). + """ + if standardize: + scaler_X = StandardScaler() + X_train = scaler_X.fit_transform(X_train) + X_test = scaler_X.transform(X_test) + Y_train = StandardScaler().fit_transform(Y_train) + + if not skip_pca and pca_components is not None and X_train.shape[-1] > pca_components: + pca_X = PCA(n_components=min(X_train.shape[-1], pca_components)) + X_train = pca_X.fit_transform(X_train) + X_test = pca_X.transform(X_test) + + n_components = min(X_train.shape[-1], Y_train.shape[-1], 25) + pls_reg = PLSRegression(n_components=n_components, scale=False) + pls_reg.fit(X_train, Y_train) + return pls_reg.predict(X_test) + + +def _correlations(Y_pred, Y_test, correlation_fn="pearson"): + """Per-neuroid correlation between predicted and actual responses.""" + if correlation_fn == "pearson": + return brain_score_pearsonr(Y_pred, Y_test) + elif correlation_fn == "spearman": + return brain_score_spearman(Y_pred, Y_test) + else: + raise ValueError("Unknown correlation metric") + + +def _reduce(correlations, reducer="median"): + """Aggregate per-neuroid correlations to a single score.""" + if reducer == "median": + return np.nanmedian(correlations) + elif reducer == "mean": + return np.nanmean(correlations) + else: + raise ValueError("Unknown reducer") + + +def score_train_test( + X_train, + Y_train, + X_test, + Y_test, + reducer="median", + correlation_fn="pearson", + pca_components=100, + skip_pca=False, + standardize=False, + ceiling=None, + ceiling_normalize=False, + n_boot_ci=100, + random_state=42, +): + """Neural predictivity on a held-out test split (the principled TVSD path). + + Fit the mapping on the train split, predict the test split, correlate per + neuroid across the test stimuli, and reduce across neuroids. The train/test + split *is* the evaluation, so there is no cross-validation; the score's spread + is instead estimated by bootstrapping over the test stimuli. + + Args: + X_train, Y_train: model features / neural responses on the train split. + X_test, Y_test: model features / neural responses on the held-out test + split (for TVSD, Y_test is the repetition-averaged ``test_MUA``). + ceiling: per-neuroid noise-ceiling vector aligned to the Y columns + (for TVSD, the split-half + Spearman-Brown internal consistency of the + repetition-averaged test responses). + ceiling_normalize: if True, divide the score by the median ceiling so that + 1.0 means "predicts as well as the noise ceiling allows". + n_boot_ci: number of bootstrap resamples over test stimuli for the std. + + Returns: + (layer_score, layer_std). + """ + Y_pred = _fit_predict(X_train, Y_train, X_test, pca_components, skip_pca, standardize) + + def _scalar(pred, actual): + s = _reduce(_correlations(pred, actual, correlation_fn), reducer) + if ceiling_normalize: + if ceiling is None: + raise ValueError("ceiling_normalize=True requires a ceiling vector") + s = s / np.nanmedian(ceiling) + return s + + layer_score = _scalar(Y_pred, Y_test) + + # Spread via bootstrap over the test stimuli (rows), reusing the fixed + # predictions -- the mapping is not refit. + rng = np.random.default_rng(random_state) + n_test = Y_test.shape[0] + boot = [ + _scalar(Y_pred[idx], Y_test[idx]) + for idx in (rng.integers(0, n_test, n_test) for _ in range(n_boot_ci)) + ] + layer_std = float(np.nanstd(boot)) + + print(f"Layer score: {layer_score:.4f}, Layer std: {layer_std:.4f}") + return float(layer_score), layer_std + + def compute_brain_score( X, Y, - n_splits=4, + n_splits=10, + train_size=0.9, + cv_strategy="shuffle", reducer="median", correlation_fn="pearson", pca_components=100, - preprocessed=True, + skip_pca=False, + standardize=False, + ceiling=None, + ceiling_normalize=False, + random_state=42, ): - kf = KFold(n_splits=n_splits, shuffle=True, random_state=42) + """Cross-validated neural predictivity within a single assembly (legacy path). + + Kept for comparison; the TVSD pipeline now prefers ``score_train_test`` (fit on + the train split, score the held-out test split). For each CV split, fit the + mapping on the train fold and correlate predictions on the test fold, then + average scores across folds. + """ + splitter = _make_splitter(cv_strategy, n_splits, train_size, random_state) scores = [] times = [] - for train_index, test_index in kf.split(X): + for train_index, test_index in splitter.split(X): start_time = time() - X_train, X_test = X[train_index], X[test_index] - Y_train, Y_test = Y[train_index], Y[test_index] - - if not preprocessed: - scaler_X, scaler_Y = StandardScaler(), StandardScaler() - X_train = scaler_X.fit_transform(X_train) - Y_train = scaler_Y.fit_transform(Y_train) - X_test = scaler_X.transform(X_test) - Y_test = scaler_Y.transform(Y_test) - - if pca_components is not None and X_train.shape[-1] > pca_components: - print("Performing PCA...") - pca_X = PCA(n_components=min(X_train.shape[-1], pca_components)) - X_train = pca_X.fit_transform(X_train) - X_test = pca_X.transform(X_test) - else: - print("Skipping PCA, using original dimensions.") - - n_components = min(X_train.shape[-1], Y_train.shape[-1], 25) - print("Performing PLS regression...") - pls_reg = PLSRegression(n_components=n_components, scale=False) - pls_reg.fit(X_train, Y_train) - Y_pred = pls_reg.predict(X_test) - - if correlation_fn == "pearson": - correlations = brain_score_pearsonr( - Y_pred, Y_test - ) # we are interested in the trend and no need to interpret the results in the original scale - elif correlation_fn == "spearman": - correlations = brain_score_spearman(Y_pred, Y_test) - else: - raise ValueError("Unknown correlation metric") - - if reducer == "median": - score = np.nanmedian(correlations) - elif reducer == "mean": - score = np.nanmean(correlations) - else: - raise ValueError("Unknown reducer") - - scores.append(score) - end_time = time() - times.append(end_time - start_time) + Y_pred = _fit_predict( + X[train_index], + Y[train_index], + X[test_index], + pca_components, + skip_pca, + standardize, + ) + correlations = _correlations(Y_pred, Y[test_index], correlation_fn) + scores.append(_reduce(correlations, reducer)) + times.append(time() - start_time) layer_score = np.nanmean(scores) layer_std = np.nanstd(scores) - mean_time = np.mean(times) + if ceiling_normalize: + if ceiling is None: + raise ValueError("ceiling_normalize=True requires a ceiling vector") + ceiling_center = np.nanmedian(ceiling) + layer_score = layer_score / ceiling_center + layer_std = layer_std / ceiling_center + print( - f"Layer score: {layer_score:.4f}, Layer std: {layer_std:.4f}, Mean time per fold: {mean_time:.4f} seconds" + f"Layer score: {layer_score:.4f}, Layer std: {layer_std:.4f}, " + f"Mean time per fold: {np.mean(times):.4f} seconds" ) - return layer_score, layer_std diff --git a/utils/dataset.py b/utils/dataset.py index d0dbe53..e5d6f87 100644 --- a/utils/dataset.py +++ b/utils/dataset.py @@ -165,6 +165,7 @@ def __init__( n_boot: int = 30, n_reps_subset: Optional[int] = None, random_state: Optional[int] = None, + spearman_brown: bool = False, ): """ Args: @@ -175,11 +176,16 @@ def __init__( n_boot: Number of bootstrap splits for reliability computation (only used if recompute_reliability=True) n_reps_subset: If not None, randomly sample this many reps for reliability computation (only used if recompute_reliability=True) random_state: Random seed for reliability computation (only used if recompute_reliability=True) + spearman_brown: If True, apply the Spearman-Brown correction to the + split-half correlations, yielding a full-length reliability estimate + (i.e. Brain-Score's InternalConsistency ceiling). Only used if + recompute_reliability=True. """ self.recompute_reliability = recompute_reliability self.n_boot = n_boot self.n_reps_subset = n_reps_subset self.random_state = random_state + self.spearman_brown = spearman_brown super().__init__(root_dir=root_dir, monkey=monkey, region=region, split="test") def _get_responses(self): @@ -201,6 +207,7 @@ def _get_responses(self): n_boot=self.n_boot, n_reps_subset=self.n_reps_subset, random_state=self.random_state, + spearman_brown=self.spearman_brown, ) reliability = torch.tensor(reliability, dtype=torch.float32) else: @@ -215,6 +222,7 @@ def _compute_reliability( n_boot: int = 30, n_reps_subset: Optional[int] = None, random_state: Optional[int] = None, + spearman_brown: bool = False, ) -> np.ndarray: """ Compute split-half reliability via bootstrapped correlations. @@ -224,6 +232,9 @@ def _compute_reliability( n_boot: number of bootstrap splits n_reps_subset: if not None, randomly sample this many reps instead of using all random_state: random seed for reproducibility + spearman_brown: if True, apply the Spearman-Brown correction (2r/(1+r)) + to each split-half correlation before averaging, giving a + full-length reliability estimate (noise ceiling) Returns: reliabilities: array of shape (neuroids,) with reliability scores @@ -251,6 +262,9 @@ def _compute_reliability( # correlation across stimuli r, _ = pearsonr(group1, group2) + if spearman_brown: + # 2r/(1+r): extrapolate the half-split correlation to full length + r = 2 * r / (1 + r) corrs.append(r) reliabilities[neu] = np.mean(corrs) return reliabilities diff --git a/utils/hooks.py b/utils/hooks.py index 9ed6bd7..1dae002 100644 --- a/utils/hooks.py +++ b/utils/hooks.py @@ -158,6 +158,24 @@ def flush(self): def get_ipca_models(self): return self.ipca_models + def load_ipca_models(self, models_root: str): + """Load per-layer IncrementalPCA models previously fit on another split. + + Lets a train-fit IPCA basis be applied to a different split (e.g. test) so + the two share one feature space with no train->test leakage in the reduction. + """ + if not os.path.isdir(models_root): + raise FileNotFoundError(f"IPCA model dir not found: {models_root}") + loaded = 0 + for layer_name in os.listdir(models_root): + pkl = os.path.join(models_root, layer_name, "ipca_model.pkl") + if os.path.exists(pkl): + with open(pkl, "rb") as f: + self.ipca_models[layer_name] = pickle.load(f) + loaded += 1 + print(f"[IPCA] loaded {loaded} models from {models_root}") + return loaded + def save_ipca_models(self): if not self.ipca_models: return