diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..0c43d7e --- /dev/null +++ b/.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 diff --git a/.gitignore b/.gitignore index 2d4689a..6c8452f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,6 @@ .venv .ipynb_checkpoints +__pycache__/ +*.py[cod] +.pytest_cache/ diff --git a/README.md b/README.md index d50e4b7..2d4a75f 100644 --- a/README.md +++ b/README.md @@ -58,15 +58,24 @@ $ python evaluation.py config/deeprecon_pyfastl2lir_alpha100_vgg19_allunits.yaml ```shell # Training of decoding models -$ python train_decoder_sklearn_ridge.py config/deeprecon_pyfastl2lir_alpha100_vgg19_allunits.yaml +$ python train_decoder_sklearn_ridge.py config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml # Prediction of DNN features -$ python predict_feature.py config/deeprecon_pyfastl2lir_alpha100_vgg19_allunits.yaml +$ python predict_feature.py config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml # Evaluation -$ python evaluation.py config/deeprecon_pyfastl2lir_alpha100_vgg19_allunits.yaml +$ python evaluation.py config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml ``` +The scikit-learn Ridge decoder is trained and stored in a factorized form: the +model maps brain activity onto the training stimulus basis, and prediction +combines its coefficients with the training features. This is mathematically +identical to regressing the features directly, but the stored decoder is much +smaller and training does not scale with the feature dimension. +`predict_feature.py` therefore reads the training features +(`decoder.features.paths`, already set in the example config). See +`ridge_factorization.py` for the details. + ### Cross-validation feature decoding - Training: `cv_train_decoder_fastl2lir.py` (example for scikit-learn Ridge regression) @@ -85,6 +94,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 diff --git a/config/README.md b/config/README.md index 311ed2c..b9c2588 100644 --- a/config/README.md +++ b/config/README.md @@ -17,7 +17,7 @@ These settings are used by the decoder training script (e.g., `train_decoder_fas - `decoder.path`: Defines the path where the decoder data is located. The path can include placeholders like `${decoder.name}` and `${decoder.features.name}` to dynamically set paths based on the feature name. - `decoder.parameters`: - `alpha`: A regularization parameter set to 100. - - `chunk_axis`: Indicates that chunking is done along axis 1. + - `chunk_axis`: Indicates that chunking is done along axis 1. For the sklearn Ridge decoder it applies to the prediction step only. 2. fMRI data - This specifies fMRI data used for the decoder training. - `decoder.fmri.name`: Refers to the dataset name, here "ImageNetTraining_fmriprep_volume_native". @@ -27,7 +27,7 @@ These settings are used by the decoder training script (e.g., `train_decoder_fas 3. Features - This specifies features used for the decoder training. - `decoder.features.name`: Refers to the pre-trained DNN model used. - - `decoder.features.paths`: Provides the path to the feature data. + - `decoder.features.paths`: Provides the path to the feature data. The factorized sklearn Ridge decoder reads these at prediction time as well, as the features of the training stimuli. - `decoder.features.layers`: Specifies the layers of the DNN from which features will be extracted. ## Decoded features diff --git a/config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml b/config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml index cf3babd..20f6390 100644 --- a/config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml +++ b/config/deeprecon_sklearn_ridge_alpha100_vgg19_allunits.yaml @@ -5,6 +5,7 @@ decoder: path: ./data/feature_decoders/ImageNetTraining/${decoder.name}/${decoder.features.name} parameters: alpha: 100 + # Axis along which the feature combination is blocked at prediction time. chunk_axis: 1 fmri: @@ -34,6 +35,7 @@ decoder: features: name: caffe/VGG19 + # Features of the training stimuli; also read by `predict_feature.py`. paths: - ./data/features/ImageNetTraining/${decoder.features.name} layers: diff --git a/predict_feature.py b/predict_feature.py index 098a1b9..1a4fefd 100644 --- a/predict_feature.py +++ b/predict_feature.py @@ -1,12 +1,20 @@ -'''DNN Feature decoding - feature prediction script.''' +'''DNN Feature decoding - feature prediction script. +Decoders trained by ``train_decoder_sklearn_ridge.py`` are factorized (see +``ridge_factorization``), so prediction runs in two steps: predict one +coefficient per training stimulus, then combine those coefficients with the +training features of the layer being decoded. Decoders saved by the earlier, +direct implementation are detected automatically and predicted as before. +''' -from typing import Dict, List, Optional + +from typing import Optional, Sequence from itertools import product import os import shutil from time import time +import warnings import bdpy from bdpy.dataform import load_array, save_array @@ -14,9 +22,147 @@ from bdpy.ml import ModelTest from bdpy.pipeline.config import init_hydra_cfg from bdpy.util import makedir_ifnot -from fastl2lir import FastL2LiR import numpy as np +from ridge_factorization import ( + FEATURE_INDEX_FILE, + TrainingFeatureLoader, + brain_model_dir, + check_column_correspondence, + combine_features, + feature_statistics, + is_factorized_model_dir, + load_factorized_metadata, + layer_model_dir, + load_train_labels, + model_file_path, + resolve_feature_index_path, + save_array_atomic, + verify_train_labels, +) + + +# Prediction back ends ####################################################### + +def _training_features(model_dir: str, layer: str, + feature_loader: TrainingFeatureLoader): + """Load the training features the decoder's coefficients refer to. + + Returns ``(train_labels, features)``, one feature row per training stimulus + in the order the coefficient columns were fitted in. + """ + metadata = load_factorized_metadata(model_dir) + train_labels = load_train_labels(model_dir) + verify_train_labels(metadata, train_labels, model_dir) + return train_labels, feature_loader.get(layer, train_labels) + + +class FeatureStatistics(object): + """``y_mean``/``y_norm`` of the training features, cached per label set. + + They depend on ``(layer, ordered training labels)`` only -- never on the + ROI -- and at DeepRecon scale each computation is two passes over a 15 GB + tensor, so they are computed once and reused for every decoder directory. + """ + + def __init__(self): + self._cache = {} + self.computations = 0 # observable in tests + + def get(self, layer: str, train_labels: Sequence[str], + features: np.ndarray): + key = (layer, tuple(train_labels)) + if key not in self._cache: + self._cache[key] = feature_statistics(features) + self.computations += 1 + return self._cache[key] + + +def _materialize_feature_statistics(decoder_path: str, layer: str, + subject: str, roi: str, model_dir: str, + feature_loader: TrainingFeatureLoader, + statistics: FeatureStatistics) -> None: + """Fill in the ``y_mean``/``y_norm`` that ``evaluation.py`` reads. + + Compatibility sidecars, not parameters of the model: prediction never reads + them, and training cannot produce them because it never opens a feature + file, so the step that has the features writes them. Only when missing -- + one decoder directory belongs to one training feature configuration -- and + atomically, since parallel prediction workers can reach the same path. + """ + directory = layer_model_dir(decoder_path, layer, subject, roi) + targets = {key: os.path.join(directory, '%s.mat' % key) + for key in ('y_mean', 'y_norm')} + if all(os.path.exists(path) for path in targets.values()): + return + + train_labels, features = _training_features(model_dir, layer, + feature_loader) + values = dict(zip(('y_mean', 'y_norm'), + statistics.get(layer, train_labels, features))) + + for key, path in sorted(targets.items()): + if os.path.exists(path): + continue + try: + save_array_atomic(path, values[key], key=key, dtype=np.float32) + print('Saved %s' % path) + except Exception: + warnings.warn('Failed to save %s. The decoder directory may be ' + 'read-only; evaluation.py will not find the training ' + 'feature statistics.' % path) + + +def _predict_factorized(model_dir: str, brain: np.ndarray, + layer: str, chunk_axis: Optional[int], + feature_loader: TrainingFeatureLoader) -> np.ndarray: + '''Predict features with a factorized decoder. + + ``brain`` must already be normalized with the decoder's ``x_mean`` / + ``x_norm``. The feature normalization parameters are not needed: they + cancel algebraically (``ridge_factorization``), so the prediction is the + linear combination of the raw training features. + ''' + train_labels, features = _training_features(model_dir, layer, + feature_loader) + + test = ModelTest(None, brain) + test.model_format = 'pickle' + test.model_path = model_file_path(model_dir) + test.dtype = np.float32 + test.chunk_axis = None # the target is the stimulus basis; never chunked + + coefficients = test.run() + check_column_correspondence(np.asarray(coefficients).shape[1], + train_labels, features, model_dir) + + # np.asarray rather than .astype: the training features can be very large + # and are already float32, so this must not copy them. + return combine_features(np.asarray(coefficients, dtype=np.float32), + np.asarray(features, dtype=np.float32), + chunk_axis=chunk_axis) + + +def _predict_legacy(model_dir: str, brain: np.ndarray, + chunk_axis: Optional[int]) -> np.ndarray: + '''Predict features with a decoder saved by the direct implementation. + + One Ridge model per chunk of the feature dimensions, and the prediction is + in normalized feature space, so it has to be un-normalized here. + ''' + feat_mean = load_array(os.path.join(model_dir, 'y_mean.mat'), key='y_mean') # shape = (1, shape_features) + feat_norm = load_array(os.path.join(model_dir, 'y_norm.mat'), key='y_norm') # shape = (1, shape_features) + + test = ModelTest(None, brain) + test.model_format = 'pickle' + test.model_path = model_dir + test.dtype = np.float32 + test.chunk_axis = chunk_axis + + feat_pred = test.run() + + return feat_pred * feat_norm + feat_mean + # Main ####################################################################### @@ -31,6 +177,7 @@ def featdec_predict( excluded_labels=[], average_sample=True, chunk_axis=1, + training_features_paths=None, analysis_name="feature_prediction" ): '''Feature prediction. @@ -39,6 +186,8 @@ def featdec_predict( - fmri_data - feature_decoder_dir + - training_features_paths: directories holding the features of the training + stimuli. Required for decoders in the factorized format. Output: @@ -64,21 +213,43 @@ def featdec_predict( data_brain = {sbj: bdpy.BData(dat_file[0]) for sbj, dat_file in fmri_data.items()} + # Training features of the factorized decoders. Only the layer currently + # being decoded is held in memory (see TrainingFeatureLoader). + feature_loader = TrainingFeatureLoader( + training_features_paths or [], feature_index_file=feature_index_file) + statistics = FeatureStatistics() + # Initialize directories ------------------------------------------- makedir_ifnot(output_dir) makedir_ifnot('tmp') # Save feature index ----------------------------------------------------- if feature_index_file is not None: - feature_index_save_file = os.path.join(output_dir, 'feature_index.mat') - shutil.copy(feature_index_file, feature_index_save_file) + # The path as given wins; only when it does not resolve is the feature + # store consulted, which is where a factorized decoder's index lives + # (`TrainingFeatureLoader` hands the same resolved path to `Features`). + feature_index_source = feature_index_file + if not os.path.exists(feature_index_source) and training_features_paths: + feature_index_source = resolve_feature_index_path( + training_features_paths[0], feature_index_file) + feature_index_save_file = os.path.join(output_dir, FEATURE_INDEX_FILE) + shutil.copy(feature_index_source, feature_index_save_file) print('Saved %s' % feature_index_save_file) # Analysis loop ---------------------------------------------------- print('----------------------------------------') print('Analysis loop') + # `layer` is the outer loop so that the training features of one layer are + # read once and reused for every subject and ROI. + current_layer = None for layer, sbj, roi in product(layers, fmri_data, rois): + if layer != current_layer: + # Release the previous layer's training features before loading the + # next ones: at most one layer is ever resident. + feature_loader.release() + current_layer = layer + print('--------------------') print('Feature: %s' % layer) print('Subject: %s' % sbj) @@ -89,10 +260,36 @@ def featdec_predict( analysis_id = analysis_name + '-' + sbj + '-' + roi + '-' + layer results_dir_prediction = os.path.join(output_dir, layer, sbj, roi) + # The brain-side directory is the only discriminator: once prediction + # has run once, a factorized decoder's per-layer directory looks like a + # legacy decoder directory from the outside. + shared_dir = brain_model_dir(decoder_path, sbj, roi) + factorized = is_factorized_model_dir(shared_dir) + model_dir = (shared_dir if factorized + else layer_model_dir(decoder_path, layer, sbj, roi)) + + # Before the "already done" skip below: a run whose features were + # already decoded would otherwise skip forever and the statistics would + # never appear. Only for factorized decoders -- a legacy decoder wrote + # its own at training time -- and only when the features are available. + if factorized and training_features_paths: + _materialize_feature_statistics( + decoder_path, layer, sbj, roi, model_dir, feature_loader, + statistics) + if os.path.exists(results_dir_prediction): print('%s is already done. Skipped.' % analysis_id) continue + # Before `makedir_ifnot`: a run that cannot proceed must not leave an + # output directory behind, or the next, correct run would skip it. + if factorized and not training_features_paths: + raise ValueError( + '%s is a factorized decoder, which needs the features of the ' + 'training stimuli. Pass training_features_paths (the ' + 'decoder.features.paths of the config used for training).' + % model_dir) + makedir_ifnot(results_dir_prediction) distcomp_db = os.path.join('./tmp', analysis_name + '.db') @@ -127,16 +324,10 @@ def featdec_predict( print('Elapsed time (data preparation): %f' % (time() - start_time)) - # Model directory - # --------------- - model_dir = os.path.join(decoder_path, layer, sbj, roi, 'model') - # Preprocessing # ------------- brain_mean = load_array(os.path.join(model_dir, 'x_mean.mat'), key='x_mean') # shape = (1, n_voxels) brain_norm = load_array(os.path.join(model_dir, 'x_norm.mat'), key='x_norm') # shape = (1, n_voxels) - feat_mean = load_array(os.path.join(model_dir, 'y_mean.mat'), key='y_mean') # shape = (1, shape_features) - feat_norm = load_array(os.path.join(model_dir, 'y_norm.mat'), key='y_norm') # shape = (1, shape_features) brain = (brain - brain_mean) / brain_norm @@ -146,20 +337,15 @@ def featdec_predict( start_time = time() - test = ModelTest(None, brain) - test.model_format = 'pickle' - test.model_path = model_dir - test.dtype = np.float32 - test.chunk_axis = chunk_axis - - feat_pred = test.run() + if factorized: + feat_pred = _predict_factorized( + model_dir, brain, layer, chunk_axis, feature_loader) + else: + print('Legacy decoder format detected in %s' % model_dir) + feat_pred = _predict_legacy(model_dir, brain, chunk_axis) print('Total elapsed time (prediction): %f' % (time() - start_time)) - # Postprocessing - # -------------- - feat_pred = feat_pred * feat_norm + feat_mean - # Save results # ------------ print('Saving results') @@ -183,6 +369,8 @@ def featdec_predict( distcomp.unlock(analysis_id) + feature_loader.release() + print('%s finished.' % analysis_name) return output_dir @@ -216,6 +404,13 @@ def featdec_predict( average_sample = cfg["decoded_feature"]["parameters"]["average_sample"] excluded_labels = cfg.decoded_feature.fmri.get("exclude_labels", []) + # Features of the training stimuli. The factorized decoder references them + # instead of storing a copy, so prediction needs them too. + training_features_paths = cfg.decoded_feature.decoder.get( + "training_features_paths", None) + if training_features_paths is None: + training_features_paths = cfg["decoder"]["features"]["paths"] + featdec_predict( test_fmri_data, decoder_path, @@ -227,5 +422,6 @@ def featdec_predict( excluded_labels=excluded_labels, average_sample=average_sample, chunk_axis=cfg["decoder"]["parameters"]["chunk_axis"], + training_features_paths=training_features_paths, analysis_name=analysis_name ) diff --git a/pyproject.toml b/pyproject.toml index b879cb1..68d0a59 100644 --- a/pyproject.toml +++ b/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", +] diff --git a/ridge_factorization.py b/ridge_factorization.py new file mode 100644 index 0000000..70ea42e --- /dev/null +++ b/ridge_factorization.py @@ -0,0 +1,532 @@ +'''Factorized formulation of the sklearn Ridge feature decoder. + +Motivation +---------- + +The direct decoder regresses brain activity on DNN features and stores a +coefficient matrix of shape ``(d_out, d_in)``. For DeepRecon-scale features +that matrix is enormous (``d_out`` reaches ~10^7 per VGG-19 layer against +``d_in`` ~10^4 voxels), even though the training targets live in a space of at +most ``M`` dimensions: every stimulus is presented several times, so the target +matrix is ``M`` unique feature vectors repeated over ``N`` trials. + +Notation +-------- + +``X`` (N x p) + Raw brain data, one row per fMRI trial, with labels ``l_1 ... l_N``. +``L_1 ... L_M`` + The unique stimulus labels, ``np.unique(brain_labels)``. +``R`` (N x M) + One-hot assignment, ``R[i, j] = 1`` iff ``l_i == L_j``. Every row sums to + one, i.e. ``R @ 1 == 1``. +``F`` (M x d) + Features of the unique stimuli, rows ordered like ``L``. +``m``, ``s`` (1 x d) + Feature mean / SD over the ``M`` unique rows (``ddof=1``). + +The direct decoder fits ``Ridge(alpha)`` on ``X_n = (X - x_mean) / x_norm`` and +``Y_n = (R F - m) / s``, then un-normalizes the prediction as ``* s + m``. Note +``Y_n = R F_n`` with ``F_n = (F - m) / s``, exactly, because ``R @ 1 == 1``. + +Equivalence +----------- + +With ``fit_intercept=True`` the sklearn Ridge prediction is affine in the +target. Writing ``X_c = X_n - 1 x_bar_n`` and + + S = (X_t - 1 x_bar_n) (X_c^T X_c + alpha I)^-1 X_c^T (n_test x N) + +which depends on the brain data only, + + predict(X_t) = S Y + (1 - S 1) y_bar, y_bar = (1/N) 1^T Y + +Substituting ``Y = R F_n`` and ``y_bar = r_bar F_n`` with ``r_bar = (1/N) 1^T R``, + + predict = [ S R + (1 - S 1) r_bar ] F_n = C F_n + +and ``C = S R + (1 - S 1) r_bar`` is precisely what the *same* estimator returns +when it is fitted on ``R`` instead: + + C = Ridge(alpha).fit(X_n, R).predict(X_t) (n_test x M) + +Row sums: ``C 1 = S R 1 + (1 - S 1)(r_bar 1) = S 1 + (1 - S 1) = 1``. The final +un-normalization therefore collapses: + + (C F_n) * s + m = C (F - 1 m) + m = C F - (C 1) m + m = C F + +So the whole decoder is ``Ridge(alpha).fit(X_n, R).predict(X_t) @ F``: the +feature mean/SD normalization cancels exactly, the meaning of ``alpha`` is +unchanged (same ``X``, same penalty), and the intercept is still fitted and +applied -- just on the ``M``-dimensional target. + +What is stored shrinks from ``d_out x d_in`` to ``M x d_in`` and no longer +depends on the feature layer, so a single fit per (subject, ROI) serves every +layer. ``F`` itself is not copied into the decoder: it already exists on disk +as the training feature directory, and prediction loads it from there. + +Note also that the fit is on ``R``, which is built from the trial labels alone, +so *training never reads* ``F``: what it learns is the map from brain activity +to coefficients over the training stimuli, plus the identity and order of those +stimuli. So the artifact records those labels and nothing about the features, +and the one thing it has to guarantee is + + coefficient column i is multiplied by the feature row of training stimulus i + +Three cheap checks cover it: ``verify_train_labels`` (the stored label list is +the one the coefficients were fitted with, in order), +``TrainingFeatureLoader.get`` (exactly one feature row from exactly one store +per requested label, assembled in the requested order), and +``check_column_correspondence`` at the contraction. + +This is a statement about the *fit*, not a licence to point one decoder +directory at several feature sets: the surrounding pipeline stores one set of +feature statistics per decoder directory for ``evaluation.py``, so one training +feature configuration means one decoder directory. The layout follows from the +same two facts -- the model does not depend on the layer, the statistics do: + + ///model/ the brain-side model, once + ////model/ y_mean.mat, y_norm.mat + +The statistics are compatibility artifacts for ``evaluation.py``, not +parameters of the model -- prediction never reads them, because the feature +normalization cancels -- and they are written by ``predict_feature.py``, the +step that has the features. +''' + + +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import hashlib +import json +import os +import pickle +import uuid + +import numpy as np +import yaml +from bdpy.dataform import Features, save_array +from bdpy.util import makedir_ifnot +from sklearn.linear_model import Ridge + + +# Model artifact ############################################################## + +FORMAT_NAME = 'sklearn_ridge_factorized' +FORMAT_VERSION = 1 + +MODEL_FILE = 'model.pkl.gz' # bdpy ModelTraining pickle layout +LABELS_FILE = 'train_labels.json' +META_FILE = 'factorized.yaml' +FEATURE_INDEX_FILE = 'feature_index.mat' + +# Identifies how `labels_digest` canonicalizes its input, so the recipe can +# change later without silently comparing incomparable digests. +LABELS_DIGEST_SPEC = 'sha256/labels-v1' + +PICKLE_PROTOCOL = 4 # matches bdpy.ml.ModelTraining + +# bdpy ModelTraining/ModelTest chunk the target only when Y.ndim >= chunk_ndim + 1. +CHUNK_NDIM = 2 + + +# Brain-side model ############################################################ + +def one_hot_assignment(brain_labels: Sequence[Any], + unique_labels: Sequence[Any], + dtype: Any = np.float64) -> np.ndarray: + '''Return the assignment matrix ``R`` (n_trials x n_unique_stimuli). + + ``R[i, j] == 1`` iff trial ``i`` presented ``unique_labels[j]``. This is + the same alignment bdpy's ``Y_sort`` index expresses; expressed as a matrix + it makes ``Y == R @ F`` explicit. + ''' + unique_labels = list(unique_labels) + position = {label: j for j, label in enumerate(unique_labels)} + if len(position) != len(unique_labels): + raise ValueError('unique_labels contains duplicates') + + assignment = np.zeros((len(brain_labels), len(unique_labels)), dtype=dtype) + for i, label in enumerate(brain_labels): + try: + assignment[i, position[label]] = 1 + except KeyError: + raise ValueError('Label %r of trial %d is not in unique_labels' + % (label, i)) + return assignment + + +def normalize_brain_for_training(brain: np.ndarray, mean: np.ndarray, + std: np.ndarray) -> np.ndarray: + '''z-score the brain data the way ``bdpy.ml.ModelTraining`` does. + + Infinities produced by zero-variance voxels are replaced by zero, matching + ``ModelTraining.run()``. NaNs are deliberately left alone: the original + implementation does not clean them either. + ''' + with np.errstate(invalid='ignore', divide='ignore'): + normalized = (brain - mean) / std + normalized[np.isinf(normalized)] = 0 + return normalized + + +def fit_stimulus_ridge(brain_normalized: np.ndarray, + brain_labels: Sequence[Any], + unique_labels: Sequence[Any], + alpha: float = 100, + dtype: Any = np.float32) -> Ridge: + '''Fit ``Ridge(alpha)`` from normalized brain data onto the stimulus basis. + + The target is the one-hot assignment matrix ``R``, so the fitted model maps + a brain pattern to ``M`` coefficients, one per unique training stimulus. + It does not depend on the feature layer. + ''' + assignment = one_hot_assignment(brain_labels, unique_labels) + model = Ridge(alpha=alpha) + model.fit(brain_normalized.astype(dtype), assignment.astype(dtype)) + return model + + +# Feature-side combination #################################################### + +def combine_features(coefficients: np.ndarray, features: np.ndarray, + chunk_axis: Optional[int] = None) -> np.ndarray: + '''Compute ``coefficients @ features`` over the stimulus axis. + + ``features`` has the training features of the unique stimuli along axis 0 + and may be multidimensional (e.g. ``(M, channels, height, width)``). When + ``chunk_axis`` is given and the features are at least 3-D, the product is + evaluated one index at a time along that axis and written straight into the + output, so no intermediate copy of the whole product is created. The + contraction runs over axis 0 only, so every output element is + mathematically independent of the blocking; it is not bit-identical, since + BLAS may accumulate a width-1 slice in a different order than the full + array (a float32 rounding difference, ~1e-7 relative). + ''' + coefficients = np.asarray(coefficients) + features = np.asarray(features) + + if coefficients.ndim != 2: + raise ValueError('coefficients must be 2-D, got %d dimensions' + % coefficients.ndim) + if features.ndim < 2: + raise ValueError('features must be at least 2-D, got %d dimensions' + % features.ndim) + if coefficients.shape[1] != features.shape[0]: + raise ValueError( + 'coefficients has %d stimulus columns but features has %d rows' + % (coefficients.shape[1], features.shape[0])) + + if chunk_axis is None or features.ndim < CHUNK_NDIM + 1: + return np.tensordot(coefficients, features, axes=([1], [0])) + + out_shape = (coefficients.shape[0],) + features.shape[1:] + out = np.empty(out_shape, dtype=np.result_type(coefficients, features)) + selector: List[Any] = [slice(None)] * out.ndim + for i_chunk in range(features.shape[chunk_axis]): + chunk = np.take(features, [i_chunk], axis=chunk_axis) + selector[chunk_axis] = slice(i_chunk, i_chunk + 1) + out[tuple(selector)] = np.tensordot(coefficients, chunk, + axes=([1], [0])) + return out + + +# Feature selection ########################################################### + +def resolve_feature_index_path(features_path: str, + feature_index_file: str) -> str: + '''Where the feature index of ``features_path`` lives. + + The single resolution rule for the whole pipeline, so that the loader and + the copy of the index cannot disagree: the index belongs to the feature + store and is resolved inside ``features_path`` unless it is absolute. + ''' + if os.path.isabs(feature_index_file): + return feature_index_file + return os.path.join(features_path, feature_index_file) + + +# Feature statistics ########################################################## + +def feature_statistics(features: np.ndarray): + '''``(mean, std)`` of the training features, shaped ``(1, *shape)``. + + The same definition the direct implementation used: per-unit statistics over + the **unique** stimuli, with ``ddof=1``. + ''' + mean = np.mean(features, axis=0)[np.newaxis, :] + std = np.std(features, axis=0, ddof=1)[np.newaxis, :] + return mean, std + + +def save_array_atomic(path: str, array: np.ndarray, key: str, + dtype: Any = np.float32) -> None: + '''Write a ``.mat`` array so that concurrent writers cannot tear it. + + The feature statistics are materialized during prediction, and parallel + prediction workers against one decoder are a supported mode, so two + processes can reach the same target path. Writing to a unique temporary + name in the same directory and then ``os.replace``-ing it is atomic within a + filesystem: a reader sees either the old file or the complete new one, never + a partial write, and a failed write leaves any existing file untouched. + ''' + directory = os.path.dirname(os.path.abspath(path)) + makedir_ifnot(directory) + # The name must keep the '.mat' suffix: hdf5storage appends one otherwise, + # and the file would not be where os.replace looks for it. + temporary = os.path.join( + directory, + '.tmp-%d-%s-%s' % (os.getpid(), uuid.uuid4().hex[:8], + os.path.basename(path))) + try: + save_array(temporary, array, key=key, dtype=dtype, sparse=False) + os.replace(temporary, path) + except BaseException: + if os.path.exists(temporary): + os.unlink(temporary) + raise + + +# Stimulus correspondence ##################################################### + +def labels_digest(labels: Sequence[Any]) -> str: + '''Digest of the **ordered** training labels. + + The stored coefficients are indexed by training stimulus, so the label order + is what makes ``coeff @ F`` meaningful. The labels are already part of the + artifact, so pinning them costs nothing and does not read any feature file. + A reordered or edited ``train_labels.json`` no longer matches. + ''' + digest = hashlib.sha256() + digest.update(b'featdec-labels-v1|') + for label in labels: + digest.update(str(label).encode('utf-8')) + digest.update(b'\0') + return digest.hexdigest() + + +def verify_train_labels(metadata: Dict[str, Any], train_labels: Sequence[Any], + model_dir: str) -> None: + '''Check the stored labels are the ones the coefficients were fitted with.''' + spec = metadata.get('labels_digest_spec') + if spec != LABELS_DIGEST_SPEC: + raise ValueError( + '%s records its label digest as %r, but this code computes %r.' + % (model_dir, spec, LABELS_DIGEST_SPEC)) + + actual = labels_digest(train_labels) + if actual != metadata['labels_digest']: + raise ValueError( + 'The training labels of %s do not match the order its coefficients ' + 'were fitted with (digest %s, expected %s). Each coefficient column ' + 'belongs to one training stimulus, so a reordered or edited %s ' + 'would silently combine the wrong features.' + % (model_dir, actual[:16], metadata['labels_digest'][:16], + LABELS_FILE)) + + +def check_column_correspondence(n_columns: int, train_labels: Sequence[Any], + features: np.ndarray, + model_dir: str) -> None: + '''Assert the decoder's only requirement, at the point it is relied on. + + Coefficient column ``i`` belongs to training stimulus ``i``, so it has to + meet the feature row of that stimulus. The label order is checked by + ``verify_train_labels`` and the row assembly by + ``TrainingFeatureLoader.get``; this catches the remaining way the three can + disagree -- a model whose width is not the number of stored labels -- with a + readable message instead of a numpy shape error inside the contraction. + ''' + n_labels = len(train_labels) + n_rows = int(np.asarray(features).shape[0]) + if n_columns == n_labels == n_rows: + return + raise ValueError( + 'Coefficient column i must be combined with the feature row of ' + 'training stimulus i, but %s predicts %d coefficient(s) while %s lists ' + '%d training stimulus/stimuli and %d feature row(s) were loaded.' + % (model_dir, n_columns, LABELS_FILE, n_labels, n_rows)) + + +# Training feature access ##################################################### + +class TrainingFeatureLoader(object): + '''Load unique-stimulus training features, one layer at a time. + + Prediction needs the training features ``F`` of the layer being decoded. + Holding one layer costs the same memory the training step already needs, but + holding several would not scale, so this cache keeps **exactly one** layer + resident: requesting a different layer releases the previous array first. + ''' + + def __init__(self, features_paths: Sequence[str], + feature_index_file: Optional[str] = None): + self._paths = list(features_paths) + self._feature_index_file = feature_index_file + self._stores: Optional[List[Features]] = None + self._cache_key: Optional[Tuple[str, Tuple[Any, ...]]] = None + self._cache: Optional[np.ndarray] = None + self.load_count = 0 # number of times features were read from disk + + @property + def cached_layers(self) -> List[str]: + '''Layers currently held in memory (never more than one).''' + return [] if self._cache_key is None else [self._cache_key[0]] + + def features(self) -> List[Features]: + if self._stores is None: + if self._feature_index_file is None: + self._stores = [Features(path) for path in self._paths] + else: + self._stores = [ + Features(path, feature_index=resolve_feature_index_path( + path, self._feature_index_file)) + for path in self._paths] + return self._stores + + def get(self, layer: str, labels: Sequence[Any]) -> np.ndarray: + '''Features of ``labels`` for ``layer``, shape ``(len(labels), ...)``. + + Row ``i`` is the features of ``labels[i]`` -- which is what makes the + coefficients meaningful, so it is enforced per label rather than checked + in aggregate. ``bdpy.dataform.utils.get_multi_features`` cannot be used + for this: it appends a row for *every* store holding a label and skips + labels no store holds, so one duplicated and one missing label cancel + out in the total row count while every row after them is shifted against + its coefficient column. + ''' + key = (layer, tuple(labels)) + if self._cache_key == key: + assert self._cache is not None + return self._cache + + # Release the previous layer before allocating the next one. + self.release() + + rows = [] + for label in labels: + sources = [store for store in self.features() + if label in store.labels] + if len(sources) != 1: + raise ValueError( + 'Training stimulus %r must come from exactly one feature ' + 'store, but %d of %d store(s) provide it: %s' + % (label, len(sources), len(self._paths), + ', '.join(self._paths))) + row = sources[0].get(layer=layer, label=label) + if row.shape[0] != 1: + raise ValueError( + 'Training stimulus %r contributes %d feature rows for ' + 'layer %r; one row per stimulus is required.' + % (label, row.shape[0], layer)) + rows.append(row) + + features = np.vstack(rows) + self._cache_key = key + self._cache = features + self.load_count += 1 + return features + + def release(self) -> None: + '''Drop the cached layer.''' + self._cache = None + self._cache_key = None + + +# Serialization ############################################################### + +def brain_model_dir(decoder_path: str, subject: str, roi: str) -> str: + """Where the brain-side model of one (subject, ROI) lives.""" + return os.path.join(decoder_path, subject, roi, 'model') + + +def layer_model_dir(decoder_path: str, layer: str, subject: str, + roi: str) -> str: + """The per-layer directory: a legacy decoder, or the statistics sidecars. + + ``evaluation.py`` reads ``y_mean``/``y_norm`` from exactly this path. + """ + return os.path.join(decoder_path, layer, subject, roi, 'model') + + +def is_factorized_model_dir(model_dir: str) -> bool: + '''True if ``model_dir`` holds a factorized decoder. + + The presence of ``factorized.yaml`` is the only discriminator, and it is + checked on the **brain-side** directory. A factorized decoder's per-layer + directories contain nothing but the statistics sidecars, which from the + outside look exactly like a legacy decoder directory, so deciding by the + per-layer path would misread a decoder as legacy as soon as prediction has + run against it once. + ''' + return os.path.isfile(os.path.join(model_dir, META_FILE)) + + +def save_factorized_model(model_dir: str, model: Ridge, + train_labels: Sequence[Any], + metadata: Optional[Dict[str, Any]] = None) -> List[str]: + '''Write the factorized decoder into ``model_dir``. + + The estimator itself is stored in bdpy's ``ModelTraining`` pickle layout + (``{'model': ..., 'y_shape': ...}``) so that ``bdpy.ml.ModelTest`` reads it + without modification. ``train_labels`` records the stimulus order, i.e. the + row order of ``F`` the coefficients refer to. ``factorized.yaml`` marks the + format and is what distinguishes a factorized decoder from a legacy one. + ''' + makedir_ifnot(model_dir) + train_labels = [str(label) for label in train_labels] + + model_path = os.path.join(model_dir, MODEL_FILE) + with open(model_path, 'wb') as f: + pickle.dump({'model': model, 'y_shape': (len(train_labels),)}, f, + protocol=PICKLE_PROTOCOL) + + labels_path = os.path.join(model_dir, LABELS_FILE) + with open(labels_path, 'w') as f: + json.dump(train_labels, f, indent=1) + + meta: Dict[str, Any] = { + 'format': FORMAT_NAME, + 'version': FORMAT_VERSION, + 'model_file': MODEL_FILE, + 'train_labels_file': LABELS_FILE, + 'n_train_stimuli': len(train_labels), + } + meta.update(metadata or {}) + meta_path = os.path.join(model_dir, META_FILE) + with open(meta_path, 'w') as f: + f.write(yaml.dump(meta, default_flow_style=False, sort_keys=True)) + + return [model_path, labels_path, meta_path] + + +def load_factorized_metadata(model_dir: str) -> Dict[str, Any]: + '''Read ``factorized.yaml``.''' + with open(os.path.join(model_dir, META_FILE), 'r') as f: + meta = yaml.safe_load(f) + if not isinstance(meta, dict) or meta.get('format') != FORMAT_NAME: + raise ValueError('%s is not a %s decoder' % (model_dir, FORMAT_NAME)) + if meta.get('version') != FORMAT_VERSION: + raise ValueError( + 'Unsupported %s version %r in %s (this code writes version %d)' + % (FORMAT_NAME, meta.get('version'), model_dir, FORMAT_VERSION)) + return meta + + +def load_train_labels(model_dir: str) -> List[str]: + '''Read the stimulus order the stored coefficients refer to.''' + meta = load_factorized_metadata(model_dir) + labels_file = meta.get('train_labels_file', LABELS_FILE) + with open(os.path.join(model_dir, labels_file), 'r') as f: + labels = json.load(f) + if len(labels) != meta['n_train_stimuli']: + raise ValueError( + '%s lists %d labels but the metadata says %d' + % (labels_file, len(labels), meta['n_train_stimuli'])) + return [str(label) for label in labels] + + +def model_file_path(model_dir: str) -> str: + '''Path of the pickled brain-side Ridge model.''' + meta = load_factorized_metadata(model_dir) + return os.path.join(model_dir, meta.get('model_file', MODEL_FILE)) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ae95565 --- /dev/null +++ b/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/.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') diff --git a/tests/data/golden/sklearn_ridge_pipeline.npz b/tests/data/golden/sklearn_ridge_pipeline.npz new file mode 100644 index 0000000..fa374e9 Binary files /dev/null and b/tests/data/golden/sklearn_ridge_pipeline.npz differ diff --git a/tests/generate_golden.py b/tests/generate_golden.py new file mode 100644 index 0000000..35abb50 --- /dev/null +++ b/tests/generate_golden.py @@ -0,0 +1,93 @@ +"""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) + 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) + values = pipeline.read_norm_params(decoder_dir, subject, + roi) + values.update(pipeline.read_feature_statistics( + decoder_dir, layer, subject, roi)) + for key, value in values.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() diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/helpers/legacy_ridge.py b/tests/helpers/legacy_ridge.py new file mode 100644 index 0000000..0f14687 --- /dev/null +++ b/tests/helpers/legacy_ridge.py @@ -0,0 +1,203 @@ +"""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 + +# What the *direct* implementation stored per decoder. The factorized decoder +# keeps only the brain half (`pipeline.BRAIN_NORM_KEYS`); the feature half is a +# property of the features a prediction combines and lives next to the decoded +# features. +NORM_KEYS = ('x_mean', 'x_norm', 'y_mean', 'y_norm') + + +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) + + +def write_legacy_decoder(model_dir: str, trained: Dict[str, Any]) -> None: + """Write ``trained`` in the pre-factorization on-disk decoder format. + + One pickle per chunk (``%08d.pkl.gz``, or ``model.pkl.gz`` when the target + was not chunked) holding ``{'model': ..., 'y_shape': ...}``, the four + normalization ``.mat`` files, and an ``info.yaml`` completion marker. Used + to test that the current prediction script can still read decoders produced + by the direct implementation. + """ + import os + import pickle + + import yaml + from bdpy.dataform import save_array + + os.makedirs(model_dir, exist_ok=True) + + for key in ('x_mean', 'x_norm', 'y_mean', 'y_norm'): + save_array(os.path.join(model_dir, key + '.mat'), trained[key], + key=key, dtype=np.float32, sparse=False) + + for i, entry in enumerate(trained['models']): + name = ('%08d.pkl.gz' % i) if trained['chunking'] else 'model.pkl.gz' + with open(os.path.join(model_dir, name), 'wb') as f: + pickle.dump({'model': entry['model'], 'y_shape': entry['y_shape']}, + f, protocol=4) + + with open(os.path.join(model_dir, 'info.yaml'), 'w') as f: + f.write(yaml.dump({'_status': {'computation_id': 'legacy', + 'computation_status': 'done'}}, + default_flow_style=False)) diff --git a/tests/helpers/pipeline.py b/tests/helpers/pipeline.py new file mode 100644 index 0000000..2e1705f --- /dev/null +++ b/tests/helpers/pipeline.py @@ -0,0 +1,156 @@ +"""Run the real training / prediction scripts on a synthetic dataset. + +Shared by the tests and by ``tests/generate_golden.py`` so that the golden +fixtures and the tests exercise exactly the same code path. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence +import inspect +import os + +import numpy as np +from bdpy.dataform import load_array + +import evaluation +import predict_feature +import ridge_factorization +import train_decoder_sklearn_ridge + +BRAIN_NORM_KEYS = ('x_mean', 'x_norm') +FEATURE_NORM_KEYS = ('y_mean', 'y_norm') + + +def run_training(dataset, output_dir: str, alpha: float = 100, + subjects: Optional[Sequence[str]] = None, + rois: Optional[Dict[str, str]] = None, + analysis_name: str = 'test_training') -> str: + """Train decoders with ``train_decoder_sklearn_ridge.py``. + + No feature-side argument at all -- not even the layers: the factorized fit + is on the stimulus basis and is shared by every layer. + """ + fmri = {sbj: paths for sbj, paths in dataset.train_fmri.items() + if subjects is None or sbj in subjects} + return train_decoder_sklearn_ridge.featdec_sklearn_ridge_train( + fmri, + output_dir=output_dir, + rois=dict(dataset.rois) if rois is None else dict(rois), + label_key=dataset.label_key, + alpha=alpha, + analysis_name=analysis_name, + ) + + +def run_prediction(dataset, decoder_dir: str, output_dir: str, + chunk_axis: Optional[int] = 1, + layers: Optional[Sequence[str]] = None, + subjects: Optional[Sequence[str]] = None, + rois: Optional[Dict[str, str]] = None, + average_sample: bool = True, + excluded_labels: Sequence[str] = (), + features_paths: Optional[Sequence[str]] = None, + feature_index_file: Optional[str] = None, + analysis_name: str = 'test_prediction') -> str: + """Predict features with ``predict_feature.py``. + + ``training_features_paths`` is passed only when the script accepts it, so + this helper works both before and after the decoder is factorized. Pass + ``features_paths=[]`` to run without them. + """ + fmri = {sbj: paths for sbj, paths in dataset.test_fmri.items() + if subjects is None or sbj in subjects} + kwargs: Dict[str, Any] = dict( + output_dir=output_dir, + rois=dict(dataset.rois) if rois is None else dict(rois), + label_key=dataset.label_key, + layers=list(dataset.layers if layers is None else layers), + excluded_labels=list(excluded_labels), + average_sample=average_sample, + chunk_axis=chunk_axis, + feature_index_file=feature_index_file, + analysis_name=analysis_name, + ) + if 'training_features_paths' in inspect.signature( + predict_feature.featdec_predict).parameters: + # An explicit empty list means "run without the training features", + # which is a supported case for an already-finished prediction. + if features_paths is None: + features_paths = [dataset.train_features_dir] + kwargs['training_features_paths'] = list(features_paths) + return predict_feature.featdec_predict(fmri, decoder_dir, **kwargs) + + +def brain_model_dir(decoder_dir: str, subject: str, roi: str) -> str: + """The shared brain-side model of a factorized decoder.""" + return ridge_factorization.brain_model_dir(decoder_dir, subject, roi) + + +def layer_model_dir(decoder_dir: str, layer: str, subject: str, + roi: str) -> str: + """A legacy decoder directory, or a factorized decoder's sidecars.""" + return ridge_factorization.layer_model_dir(decoder_dir, layer, subject, roi) + + +def read_norm_params(decoder_dir: str, subject: str, roi: str, + keys: Sequence[str] = BRAIN_NORM_KEYS + ) -> Dict[str, np.ndarray]: + """Load the brain-side normalization parameters of a factorized decoder.""" + directory = brain_model_dir(decoder_dir, subject, roi) + return {key: load_array(os.path.join(directory, key + '.mat'), key=key) + for key in keys} + + +def read_feature_statistics(decoder_dir: str, layer: str, subject: str, + roi: str) -> Dict[str, np.ndarray]: + """Load the per-layer statistics sidecars ``evaluation.py`` reads.""" + directory = layer_model_dir(decoder_dir, layer, subject, roi) + return {key: load_array(os.path.join(directory, key + '.mat'), key=key) + for key in FEATURE_NORM_KEYS} + + +def read_legacy_norm_params(decoder_dir: str, layer: str, subject: str, + roi: str) -> Dict[str, np.ndarray]: + """All four parameters of a decoder in the direct (per-layer) format.""" + directory = layer_model_dir(decoder_dir, layer, subject, roi) + return {key: load_array(os.path.join(directory, key + '.mat'), key=key) + for key in BRAIN_NORM_KEYS + FEATURE_NORM_KEYS} + + +def run_evaluation(dataset, decoder_dir: str, decoded_dir: str, + true_features_dir: str, + layers: Optional[Sequence[str]] = None, + subjects: Optional[Sequence[str]] = None, + rois: Optional[Dict[str, str]] = None, + output_file: Optional[str] = None, + feature_index_file: Optional[str] = None, + average_sample: bool = True) -> str: + """Evaluate decoded features with ``evaluation.py``.""" + return evaluation.featdec_eval( + decoded_dir, + true_features_dir, + output_file=output_file or os.path.join(decoded_dir, 'evaluation.db'), + subjects=list(subjects or dataset.test_fmri), + rois=list(rois or dataset.rois), + layers=list(dataset.layers if layers is None else layers), + feature_index_file=feature_index_file, + feature_decoder_path=decoder_dir, + average_sample=average_sample, + ) + + +def read_decoded_features(output_dir: str, layer: str, subject: str, roi: str, + labels: Sequence[str]) -> np.ndarray: + """Stack the saved per-label ``.mat`` files into one array.""" + directory = os.path.join(output_dir, layer, subject, roi) + return np.vstack([ + load_array(os.path.join(directory, '%s.mat' % label), key='feat') + for label in labels]) + + +def decoded_feature_labels(output_dir: str, layer: str, subject: str, + roi: str) -> List[str]: + directory = os.path.join(output_dir, layer, subject, roi) + return sorted(os.path.splitext(f)[0] for f in os.listdir(directory) + if f.endswith('.mat')) diff --git a/tests/helpers/synthetic.py b/tests/helpers/synthetic.py new file mode 100644 index 0000000..ea7e28e --- /dev/null +++ b/tests/helpers/synthetic.py @@ -0,0 +1,274 @@ +"""Build tiny synthetic on-disk datasets for the feature-decoding tests. + +Everything is small enough to live in a temporary directory: a handful of +``bdpy.BData`` HDF5 files and a ``bdpy.dataform.Features`` tree. No part of the +real DeepRecon dataset is needed. + +The dataset intentionally reproduces the properties the decoders depend on: + +* stimuli are **repeated** across fMRI trials, with an *unbalanced* number of + repetitions per stimulus (3, 3, 2, 2, 1, 1); +* one subject's training trials are split over **two** ``.h5`` files, so + ``select_data_multi_bdatas`` / ``get_labels_multi_bdatas`` are exercised; +* two ROIs of different sizes; +* a 1-D ``fc``-like feature layer (no chunking) *and* a 3-D ``conv``-like layer, + so the ``chunk_axis`` / ``order='F'`` reshape path is exercised. The shipped + config only enables ``fc6``/``fc7``/``fc8``, which never reaches that path. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Sequence, Tuple +import os + +import numpy as np +from bdpy import BData +from bdpy.dataform import save_array + +# Feature layers: name -> feature shape of a single stimulus. +LAYER_SHAPES: Dict[str, Tuple[int, ...]] = { + 'conv_like': (3, 2, 2), + 'fc_like': (4,), +} + +N_VOXELS = 8 +ROIS: Dict[str, str] = {'VC': 'ROI_VC = 1', 'HALF': 'ROI_HALF = 1'} +ROI_SIZES: Dict[str, int] = {'VC': N_VOXELS, 'HALF': N_VOXELS // 2} +LABEL_KEY = 'stimulus_name' + +# Unbalanced repetitions: 6 unique training stimuli, 12 trials. +TRAIN_REPEATS: Tuple[int, ...] = (3, 3, 2, 2, 1, 1) +# 3 unique test stimuli, 5 trials. +TEST_REPEATS: Tuple[int, ...] = (2, 2, 1) + + +@dataclass +class SyntheticDataset: + """Paths and in-memory copies of a synthetic decoding dataset.""" + + root: str + train_fmri: Dict[str, List[str]] + test_fmri: Dict[str, List[str]] + train_features_dir: str + train_brain: Dict[str, np.ndarray] = field(default_factory=dict) + train_labels: Dict[str, List[str]] = field(default_factory=dict) + test_brain: Dict[str, np.ndarray] = field(default_factory=dict) + test_labels: Dict[str, List[str]] = field(default_factory=dict) + features: Dict[str, np.ndarray] = field(default_factory=dict) + unique_train_labels: List[str] = field(default_factory=list) + unique_test_labels: List[str] = field(default_factory=list) + rois: Dict[str, str] = field(default_factory=lambda: dict(ROIS)) + label_key: str = LABEL_KEY + layer_shapes: Dict[str, Tuple[int, ...]] = field( + default_factory=lambda: dict(LAYER_SHAPES)) + + @property + def layers(self) -> List[str]: + return sorted(self.layer_shapes) + + def brain(self, subject: str, roi: str, split: str = 'train') -> np.ndarray: + """Voxel data for one subject/ROI, as the scripts would select it.""" + data = self.train_brain if split == 'train' else self.test_brain + return data[subject][:, :ROI_SIZES[roi]] + + def labels(self, subject: str, split: str = 'train') -> List[str]: + return (self.train_labels if split == 'train' else self.test_labels)[subject] + + def feature_matrix(self, layer: str, labels: Sequence[str]) -> np.ndarray: + """Rows of ``features[layer]`` in the order of ``labels``.""" + index = [self.unique_train_labels.index(lb) for lb in labels] + return self.features[layer][index] + + +def _expand(labels: Sequence[str], repeats: Sequence[int]) -> List[str]: + """Repeat each label ``repeats[i]`` times, interleaved as in a real run.""" + trials: List[str] = [] + remaining = list(repeats) + while any(r > 0 for r in remaining): + for i, label in enumerate(labels): + if remaining[i] > 0: + trials.append(label) + remaining[i] -= 1 + return trials + + +def write_bdata(path: str, brain: np.ndarray, labels: Sequence[str], + label_to_number: Dict[str, int]) -> None: + """Write one ``BData`` HDF5 file with ROI metadata and a label vmap.""" + n_voxels = brain.shape[1] + bdata = BData() + bdata.add(np.asarray(brain, dtype=float), 'VoxelData') + bdata.add(np.array([[label_to_number[lb]] for lb in labels], dtype=float), + LABEL_KEY) + bdata.add_metadata('ROI_VC', np.ones(n_voxels), where='VoxelData') + half = np.zeros(n_voxels) + half[:ROI_SIZES['HALF']] = 1 + bdata.add_metadata('ROI_HALF', half, where='VoxelData') + bdata.add_vmap(LABEL_KEY, {v: k for k, v in label_to_number.items()}) + os.makedirs(os.path.dirname(path), exist_ok=True) + bdata.save(path) + + +def write_features(root: str, layer_shapes: Dict[str, Tuple[int, ...]], + labels: Sequence[str], + arrays: Dict[str, np.ndarray]) -> None: + """Write a ``Features`` tree ``//