test: Add pytest infrastructure and regression tests for the sklearn Ridge decoder

.github/workflows/test.yml

@@ -0,0 +1,27 @@
+name: tests
+
+on:
+  push:
+    branches: [main, dev]
+  pull_request:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      # bdpy declares Requires-Python <3.12.
+      - name: Set up Python
+        run: uv python install 3.11
+
+      - name: Install dependencies
+        run: uv sync --group dev --python 3.11
+
+      - name: Run tests
+        run: uv run --python 3.11 pytest

.gitignore

@@ -4,3 +4,6 @@
 .venv
 .ipynb_checkpoints
 
+__pycache__/
+*.py[cod]
+.pytest_cache/

README.md

@@ -85,6 +85,22 @@ $ python cv_predict_feature_fastl2lir.py config/deeprecon_cv_pyfastl2lir_alpha10
 $ python cv_evaluation.py config/deeprecon_cv_pyfastl2lir_alpha100_vgg19_allunits.yaml
 ```
 
+### Tests
+
+The test suite runs on small synthetic data generated on the fly; no downloaded
+dataset is required.
+
+```shell
+# Install the test dependencies and run the suite
+$ uv sync --group dev
+$ uv run pytest
+```
+
+`tests/data/golden/` holds regression fixtures recording the numerical output of
+the decoding pipeline. Regenerate them with
+`uv run python -m tests.generate_golden` only when the expected output is meant
+to change, and say so explicitly in the commit message.
+
 ## References
 
 - Horikawa and Kamitani (2017) Generic decoding of seen and imagined objects using hierarchical visual features. *Nature Communications* 8:15037. https://www.nature.com/articles/ncomms15037

pyproject.toml

@@ -12,8 +12,29 @@ dependencies = [
     "matplotlib",
     "numpy<2.0.0",
     "pyyaml",
+    "scikit-learn",
     "tqdm",
 ]
 
+[dependency-groups]
+dev = [
+    "pytest>=7.0",
+]
+
 [tool.uv]
 package = false
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"
+# The analysis scripts are top-level modules rather than an installed package,
+# so the repository root has to be importable. `pythonpath` states that
+# explicitly instead of relying on pytest's implicit rootdir insertion.
+pythonpath = ["."]
+markers = [
+    "slow: tests that take more than a few seconds",
+]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::PendingDeprecationWarning",
+]

tests/conftest.py

@@ -0,0 +1,45 @@
+"""Shared fixtures for the feature-decoding tests."""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from tests.helpers import synthetic
+
+GOLDEN_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                          'data', 'golden')
+
+
+@pytest.fixture(autouse=True)
+def _run_in_tmp_path(tmp_path, monkeypatch):
+    """Run every test in its own directory.
+
+    Both scripts create ``./tmp/<analysis_name>.db`` (the DistComp sqlite lock
+    database) relative to the current working directory.
+    """
+    monkeypatch.chdir(tmp_path)
+
+
+@pytest.fixture
+def dataset(tmp_path):
+    """A small synthetic dataset with repeated stimuli."""
+    return synthetic.make_dataset(str(tmp_path / 'data'), seed=0)
+
+
+@pytest.fixture
+def dataset_with_constant_unit(tmp_path):
+    """Synthetic dataset where one feature unit has zero training variance."""
+    return synthetic.make_dataset(str(tmp_path / 'data'), seed=0,
+                                  constant_feature_unit=True)
+
+
+@pytest.fixture
+def decoder_dir(tmp_path):
+    return str(tmp_path / 'feature_decoders')
+
+
+@pytest.fixture
+def decoded_dir(tmp_path):
+    return str(tmp_path / 'decoded_features')

tests/generate_golden.py

@@ -0,0 +1,91 @@
+"""Regenerate the golden regression fixtures for the sklearn Ridge decoder.
+
+The fixtures record the normalization parameters and decoded features produced
+by ``train_decoder_sklearn_ridge.py`` + ``predict_feature.py`` on a fixed-seed
+synthetic dataset.  They were generated from the *direct* (pre-factorization)
+implementation and are what the factorized implementation must reproduce.
+
+Only run this when the expected numerical output is intended to change, and say
+so explicitly in the commit message.
+
+Note that regeneration is not bit-for-bit stable across environments: the
+pipeline runs in ``float32`` and the Ridge solve's summation order depends on
+how BLAS partitions the work, which moves the last bits (observed ~5e-8
+relative). That is why the tests compare against these fixtures with a
+tolerance rather than exactly, and why regenerating them without an intended
+behavior change only adds noise.
+
+    uv run python -m tests.generate_golden
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import tempfile
+
+import numpy as np
+
+from tests.helpers import pipeline, synthetic
+
+ALPHA = 100
+CHUNK_AXIS = 1
+SEED = 0
+GOLDEN_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                           'data', 'golden', 'sklearn_ridge_pipeline.npz')
+
+
+def build(work_dir: str) -> dict:
+    """Run the full pipeline in ``work_dir`` and collect everything of interest."""
+    cwd = os.getcwd()
+    os.chdir(work_dir)
+    try:
+        dataset = synthetic.make_dataset(os.path.join(work_dir, 'data'),
+                                         seed=SEED)
+        decoder_dir = os.path.join(work_dir, 'feature_decoders')
+        decoded_dir = os.path.join(work_dir, 'decoded_features')
+
+        pipeline.run_training(dataset, decoder_dir, alpha=ALPHA,
+                              chunk_axis=CHUNK_AXIS)
+        pipeline.run_prediction(dataset, decoder_dir, decoded_dir,
+                                chunk_axis=CHUNK_AXIS)
+
+        golden = {
+            '_alpha': np.array(ALPHA),
+            '_chunk_axis': np.array(CHUNK_AXIS),
+            '_seed': np.array(SEED),
+            '_layers': np.array(dataset.layers),
+            '_subjects': np.array(sorted(dataset.train_fmri)),
+            '_rois': np.array(sorted(dataset.rois)),
+            '_test_labels': np.array(dataset.unique_test_labels),
+        }
+        for layer in dataset.layers:
+            for subject in sorted(dataset.train_fmri):
+                for roi in sorted(dataset.rois):
+                    prefix = '%s|%s|%s|' % (layer, subject, roi)
+                    for key, value in pipeline.read_norm_params(
+                            decoder_dir, layer, subject, roi).items():
+                        golden[prefix + key] = value
+                    golden[prefix + 'pred'] = pipeline.read_decoded_features(
+                        decoded_dir, layer, subject, roi,
+                        dataset.unique_test_labels)
+        return golden
+    finally:
+        os.chdir(cwd)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('-o', '--output', default=GOLDEN_FILE)
+    args = parser.parse_args()
+
+    with tempfile.TemporaryDirectory() as work_dir:
+        golden = build(work_dir)
+
+    os.makedirs(os.path.dirname(args.output), exist_ok=True)
+    np.savez_compressed(args.output, **golden)
+    print('Saved %s (%d arrays)' % (args.output, len(golden)))
+
+
+if __name__ == '__main__':
+    main()

tests/helpers/legacy_ridge.py

@@ -0,0 +1,164 @@
+"""Reference implementation of the *direct* sklearn Ridge feature decoder.
+
+This mirrors, in plain numpy + scikit-learn, exactly what
+``train_decoder_sklearn_ridge.py`` + ``predict_feature.py`` did before the
+decoder was factorized, including every detail bdpy's
+``ModelTraining``/``ModelTest`` contribute:
+
+* brain z-scoring with ``ddof=1`` followed by ``X[np.isinf(X)] = 0``;
+* feature z-scoring computed over the **unique** stimuli;
+* one Ridge per index along ``chunk_axis`` when the feature array is at least
+  3-D (``Y.ndim >= chunk_ndim + 1`` with bdpy's default ``chunk_ndim=2``);
+* the chunk taken with ``np.take(..., [i], axis=chunk_axis)`` so the axis is
+  kept with size 1, and the normalization parameters sliced the same way;
+* ``Y = Y[feat_index]`` applied *after* normalization -- this is the repetition
+  expansion that turns M unique feature rows into N trial-aligned rows;
+* ``reshape(n, -1, order='F')`` before fitting and the inverse reshape after
+  predicting;
+* ``float32`` casts for both X and Y;
+* ``np.concatenate(..., axis=chunk_axis)`` to reassemble the chunks;
+* un-normalization as ``y_pred * y_norm + y_mean``.
+
+It exists so that the factorized implementation can be compared against the old
+math directly, and so that the old math itself is pinned by tests against the
+real scripts (see ``tests/test_legacy_reference.py``).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Sequence
+import copy
+
+import numpy as np
+from sklearn.linear_model import Ridge
+
+CHUNK_NDIM = 2  # bdpy ModelTraining default
+
+
+def brain_norm_params(brain: np.ndarray):
+    """``(mean, std)`` of the brain data, shaped ``(1, n_voxels)``."""
+    mean = np.mean(brain, axis=0)[np.newaxis, :]
+    std = np.std(brain, axis=0, ddof=1)[np.newaxis, :]
+    return mean, std
+
+
+def feature_norm_params(feat: np.ndarray):
+    """``(mean, std)`` of the features, shaped ``(1, *feature_shape)``."""
+    mean = np.mean(feat, axis=0)[np.newaxis, :]
+    std = np.std(feat, axis=0, ddof=1)[np.newaxis, :]
+    return mean, std
+
+
+def assignment_index(brain_labels: Sequence[str],
+                     feat_labels: Sequence[str]) -> np.ndarray:
+    """Row of ``feat`` matching each brain trial (the ``Y_sort`` index)."""
+    return np.array([np.where(np.array(feat_labels) == bl)
+                     for bl in brain_labels]).flatten()
+
+
+def normalize_brain_for_training(brain: np.ndarray, mean: np.ndarray,
+                                 std: np.ndarray) -> np.ndarray:
+    """z-score as ``ModelTraining`` does it, including the ``inf`` clean-up."""
+    with np.errstate(invalid='ignore', divide='ignore'):
+        x = (brain - mean) / std
+    x[np.isinf(x)] = 0
+    return x
+
+
+def use_chunking(feat: np.ndarray, chunk_axis: Optional[int]) -> bool:
+    if chunk_axis is None:
+        return False
+    return feat.ndim >= CHUNK_NDIM + 1
+
+
+def legacy_train(brain: np.ndarray, brain_labels: Sequence[str],
+                 feat: np.ndarray, feat_labels: Sequence[str],
+                 alpha: float = 100, chunk_axis: Optional[int] = 1,
+                 dtype: Any = np.float32) -> Dict[str, Any]:
+    """Fit the direct brain -> feature Ridge decoder, chunk by chunk."""
+    x_mean, x_norm = brain_norm_params(brain)
+    y_mean, y_norm = feature_norm_params(feat)
+    feat_index = assignment_index(brain_labels, feat_labels)
+
+    x = normalize_brain_for_training(brain, x_mean, x_norm)
+
+    chunking = use_chunking(feat, chunk_axis)
+    chunk_index: Sequence[Optional[int]]
+    chunk_index = range(feat.shape[chunk_axis]) if chunking else [None]
+
+    models: List[Dict[str, Any]] = []
+    for i_chunk in chunk_index:
+        if chunking:
+            y = np.take(feat, [i_chunk], axis=chunk_axis)
+            chunk_mean = np.take(y_mean, [i_chunk], axis=chunk_axis)
+            chunk_norm = np.take(y_norm, [i_chunk], axis=chunk_axis)
+        else:
+            y = feat
+            chunk_mean, chunk_norm = y_mean, y_norm
+
+        with np.errstate(invalid='ignore', divide='ignore'):
+            y = (y - chunk_mean) / chunk_norm
+        y[np.isinf(y)] = 0
+        y = y[feat_index]
+
+        y_shape = y.shape[1:]
+        if y.ndim > 2:
+            y = y.reshape(y.shape[0], -1, order='F')
+
+        model = Ridge(alpha=alpha)
+        model.fit(x.astype(dtype), y.astype(dtype))
+        models.append({'model': model, 'y_shape': y_shape})
+
+    return {
+        'models': models,
+        'x_mean': x_mean, 'x_norm': x_norm,
+        'y_mean': y_mean, 'y_norm': y_norm,
+        'feat_index': feat_index,
+        'chunk_axis': chunk_axis,
+        'chunking': chunking,
+    }
+
+
+def legacy_predict(trained: Dict[str, Any], brain_test: np.ndarray,
+                   dtype: Any = np.float32) -> np.ndarray:
+    """Predict features from *raw* test brain data, as ``predict_feature`` did.
+
+    Note the deliberate asymmetry with training: the prediction script does not
+    apply the ``inf`` clean-up after z-scoring.
+    """
+    x = (brain_test - trained['x_mean']) / trained['x_norm']
+    x = x.astype(dtype)
+
+    preds = []
+    for entry in trained['models']:
+        y_pred = entry['model'].predict(x)
+        if y_pred.shape[1:] != entry['y_shape']:
+            y_pred = y_pred.reshape((y_pred.shape[0],) + entry['y_shape'],
+                                    order='F')
+        preds.append(y_pred)
+
+    if trained['chunk_axis'] is None:
+        pred = preds[0]
+    else:
+        pred = np.concatenate(preds, axis=trained['chunk_axis'])
+
+    return pred * trained['y_norm'] + trained['y_mean']
+
+
+def average_test_brain(brain: np.ndarray, labels: Sequence[str],
+                       excluded_labels: Sequence[str] = ()):
+    """Trial-average the test brain data, as ``predict_feature`` does."""
+    unique = [lb for lb in np.unique(labels) if lb not in excluded_labels]
+    averaged = np.vstack([
+        np.mean(brain[(np.array(labels) == lb).flatten(), :], axis=0)
+        for lb in unique])
+    return averaged, list(unique)
+
+
+def single_trial_labels(labels: Sequence[str]) -> List[str]:
+    """Output labels used by ``predict_feature`` when ``average_sample`` is off."""
+    return ['sample{:06}-{}'.format(i + 1, lb) for i, lb in enumerate(labels)]
+
+
+def clone_trained(trained: Dict[str, Any]) -> Dict[str, Any]:
+    return copy.deepcopy(trained)