From fe63d5e4a2aa3083bb5a46b0fa1775a66a6e8f01 Mon Sep 17 00:00:00 2001 From: Kenya Otsuka Date: Thu, 10 Sep 2026 01:57:15 +0000 Subject: [PATCH 1/2] refactor: Factorize the sklearn Ridge decoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The decoder regressed brain activity directly onto DNN features and pickled the fitted Ridge, so the stored coefficient matrix was d_out x n_voxels: ~120 GB for one VGG19 conv1_1 decoder of one (subject, ROI) at DeepRecon scale, and one Ridge solve per layer and per feature chunk. None of that is intrinsic. Each stimulus is presented in several trials, so the training target is not one independent vector per trial but the features of the M unique stimuli, repeated: Y = R F with R the one-hot trial-to-stimulus assignment. Ridge prediction is linear in the target and R has unit row sums, so Ridge(a).fit(X, R F).predict(Xt) == Ridge(a).fit(X, R).predict(Xt) @ F exactly, including the intercept, and the feature mean/SD round trip cancels rather than being dropped. The full derivation is in the module docstring of ridge_factorization.py; tests/test_ridge_factorization.py pins it in float64 at rtol=1e-9 across alphas, repetition patterns, solvers and chunk axes. So the decoder is trained and stored as brain -> training stimulus basis, and prediction combines the coefficients with F. The stored model is M x n_voxels, it no longer depends on the feature layer, and the repeated trials are *not* averaged: all N trials enter the fit exactly as before, only the target is re-expressed. Artifact layout. The model is the same for every layer, so it is stored once per (subject, ROI); the training feature statistics do depend on the layer, and evaluation.py reads them from where it always has: ///model/ model.pkl.gz, x_mean.mat, x_norm.mat, info.yaml ////model/ y_mean.mat, y_norm.mat model.pkl.gz is the decoder, whole: the fitted estimator, the ordered training labels its coefficient columns are indexed by, the format name and version, and the provenance of the fit, in one pickle. bdpy's ModelTest reads only 'model' and 'y_shape' from such a file and ignores the rest, so the extra keys ride along at no cost -- and there is no second file that could disagree with the coefficients. y_mean/y_norm are compatibility artifacts for the evaluation pipeline, not parameters of the model -- prediction never reads them, because the normalization cancels. Training cannot produce them (it opens no feature file), so predict_feature.py writes them on its first run and never overwrites what is there: a decoder directory carries the statistics of one training feature configuration, and another feature set or feature index is another decoder directory. The fit is on R, built from the trial labels, so training takes no feature-side argument at all: no feature paths, no feature index, no chunk_axis (which only blocks the coefficient-feature product in memory at prediction), not even the layer names. Over a DeepRecon run that removes ~3 TB of reads that only ever recomputed identical statistics. What the artifact does pin is the correspondence the coefficients need: coefficient column i is multiplied by the feature row of stimulus i The ordered labels sit in the same pickle as the coefficients they index, so the two cannot be reordered, renamed or replaced against each other. What is left is the assembly of F: TrainingFeatureLoader.get requires exactly one feature row from exactly one store per requested label (bdpy's get_multi_features appends a row for every store holding a label and skips labels no store holds, so a duplicate and a missing label cancel out in the total row count while every row after them is shifted), and check_column_correspondence compares the widths at the contraction. Prediction resolves the decoder by looking at the shared directory first: a factorized decoder's per-layer directory holds nothing but the statistics sidecars, so from the outside it looks exactly like a legacy decoder directory. A model.pkl.gz on the shared path is the discriminator -- the direct implementation stores its model per layer, never there -- and the format tag inside it is checked when it is loaded; decoders written by the direct implementation keep their layout and their previous code path, unchanged. Deliberate behavior change: a feature unit that is constant across the training stimuli made the direct implementation fail (0/0 = NaN targets that ModelTraining's inf clean-up does not catch, which scikit-learn rejects) and now decodes exactly, since C·1 = 1. That is a bug fix, not a consequence of the refactor, and tests/test_zero_variance_features.py pins both behaviors side by side. Also fixes feature_index_file, which could not run at all: the feature stores resolved the index inside the feature directory while the copy into the decoder resolved it against the current directory. It is now prediction-side unit selection, resolved once through resolve_feature_index_path and copied next to the decoded features. The same unresolved-path bug remains in train_decoder_fastl2lir.py, predict_feature_fastl2lir.py, evaluation.py and the two cv_* scripts, left alone here. No bdpy changes and no change to evaluation.py or cv_evaluation.py. The baseline the factorized output is checked against is the reference implementation in tests/helpers/legacy_ridge.py, which tests/test_legacy_reference.py pins against the pre-factorization scripts. The benchmarks that produced the numbers quoted in the pull request are kept outside the repository. --- README.md | 15 +- config/README.md | 4 +- ...sklearn_ridge_alpha100_vgg19_allunits.yaml | 2 + predict_feature.py | 248 ++++++++- ridge_factorization.py | 470 ++++++++++++++++++ tests/helpers/legacy_ridge.py | 38 ++ tests/helpers/pipeline.py | 87 +++- tests/helpers/synthetic.py | 53 ++ tests/test_chunking.py | 172 ++++--- tests/test_evaluation.py | 54 ++ tests/test_factorized_artifacts.py | 122 +++++ tests/test_lazy_feature_statistics.py | 312 ++++++++++++ tests/test_legacy_model_compat.py | 203 ++++++++ tests/test_legacy_reference.py | 17 +- tests/test_normalization.py | 122 +++++ tests/test_predict_feature.py | 18 +- tests/test_ridge_factorization.py | 235 +++++++++ tests/test_serialization.py | 59 +++ tests/test_stimulus_correspondence.py | 247 +++++++++ tests/test_train_decoder_sklearn_ridge.py | 149 +++--- tests/test_training_feature_loading.py | 141 ++++++ tests/test_zero_variance_features.py | 49 +- train_decoder_sklearn_ridge.py | 251 +++++----- 23 files changed, 2742 insertions(+), 326 deletions(-) create mode 100644 ridge_factorization.py create mode 100644 tests/test_evaluation.py create mode 100644 tests/test_factorized_artifacts.py create mode 100644 tests/test_lazy_feature_statistics.py create mode 100644 tests/test_legacy_model_compat.py create mode 100644 tests/test_normalization.py create mode 100644 tests/test_ridge_factorization.py create mode 100644 tests/test_serialization.py create mode 100644 tests/test_stimulus_correspondence.py create mode 100644 tests/test_training_feature_loading.py diff --git a/README.md b/README.md index 714d7d3..2be2389 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) 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..d1347ff 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 Any, Dict, 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,146 @@ 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, + layer_model_dir, + load_factorized_model, + resolve_feature_index_path, + save_array_atomic, +) + + +# Prediction back ends ####################################################### + +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 = {} + + 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) + return self._cache[key] + + +def _feature_statistics_paths(decoder_path: str, layer: str, subject: str, + roi: str) -> Dict[str, str]: + directory = layer_model_dir(decoder_path, layer, subject, roi) + return {key: os.path.join(directory, '%s.mat' % key) + for key in ('y_mean', 'y_norm')} + + +def missing_feature_statistics(decoder_path: str, layer: str, subject: str, + roi: str) -> bool: + """Whether ``evaluation.py``'s sidecars still have to be produced. + + Checked before the decoder is read, so an already-complete run touches + neither the model nor the features. + """ + return not all(os.path.exists(path) for path in _feature_statistics_paths( + decoder_path, layer, subject, roi).values()) + + +def _materialize_feature_statistics(decoder_path: str, layer: str, + subject: str, roi: str, + train_labels: Sequence[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. + + Takes the ordered training labels rather than the decoder: the statistics + are a property of the features, not of the estimator. + """ + targets = _feature_statistics_paths(decoder_path, layer, subject, roi) + features = feature_loader.get(layer, train_labels) + 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(artifact: Dict[str, Any], 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. + + The estimator is called directly rather than through ``bdpy.ml.ModelTest``: + the target is the stimulus basis, so there is nothing to chunk or reshape, + and the decoder has already been read once for this iteration. The float32 + cast is what ``ModelTest`` applies to ``X`` before predicting. + ''' + train_labels = artifact['train_labels'] + features = feature_loader.get(layer, train_labels) + + coefficients = artifact['model'].predict( + np.asarray(brain).astype(np.float32)) + 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 +176,7 @@ def featdec_predict( excluded_labels=[], average_sample=True, chunk_axis=1, + training_features_paths=None, analysis_name="feature_prediction" ): '''Feature prediction. @@ -39,6 +185,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 +212,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 +259,43 @@ 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)) + + # The decoder is read at most once per iteration, and only if this + # iteration still has something to produce: fully done -> no read at + # all, sidecars or prediction (or both) -> exactly one. + artifact = None + + # 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 + and missing_feature_statistics(decoder_path, layer, sbj, roi)): + artifact = load_factorized_model(model_dir) + _materialize_feature_statistics( + decoder_path, layer, sbj, roi, artifact['train_labels'], + 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 +330,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 +343,17 @@ 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: + if artifact is None: + artifact = load_factorized_model(model_dir) + feat_pred = _predict_factorized( + artifact, 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 +377,8 @@ def featdec_predict( distcomp.unlock(analysis_id) + feature_loader.release() + print('%s finished.' % analysis_name) return output_dir @@ -216,6 +412,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 +430,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/ridge_factorization.py b/ridge_factorization.py new file mode 100644 index 0000000..e6f99ad --- /dev/null +++ b/ridge_factorization.py @@ -0,0 +1,470 @@ +'''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 + +The ordered labels travel *inside* ``model.pkl.gz``, next to the coefficients +they index, so they cannot drift apart. Two cheap checks cover the rest: +``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/ model.pkl.gz, x_mean.mat, + x_norm.mat, info.yaml + ////model/ y_mean.mat, y_norm.mat + +``model.pkl.gz`` is the whole factorized model: the fitted estimator, the +ordered training labels, the format tag and the provenance metadata, in bdpy's +own ``{'model', 'y_shape'}`` pickle layout extended with the extra keys (bdpy's +``ModelTest`` reads only those two and ignores the rest). + +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 os +import pickle +import uuid + +import numpy as np +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 +FEATURE_INDEX_FILE = 'feature_index.mat' + +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 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 labels are stored with the + coefficients and the row assembly is ``TrainingFeatureLoader.get``'s job; + this catches what is left -- 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 it ' + 'records %d training stimulus/stimuli and %d feature row(s) were ' + 'loaded.' % (model_dir, n_columns, 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 + + 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 + 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 model_file_path(model_dir: str) -> str: + """Path of the factorized model artifact.""" + return os.path.join(model_dir, MODEL_FILE) + + +def is_factorized_model_dir(model_dir: str) -> bool: + '''True if ``model_dir`` holds a factorized decoder. + + The discriminator is a ``model.pkl.gz`` on the **brain-side** path + ``///model/``, which a decoder written by the direct + implementation never has: its models live under + ``///model/``. Deciding by the per-layer path would + misread a factorized decoder as legacy as soon as prediction has written + the statistics sidecars there. What the file *contains* is validated when + it is loaded, not here -- this has to stay cheap enough to call per + iteration. + ''' + return os.path.isfile(model_file_path(model_dir)) + + +def save_factorized_model(model_dir: str, model: Ridge, + train_labels: Sequence[Any], + metadata: Optional[Dict[str, Any]] = None) -> str: + '''Write the factorized decoder as a single artifact. + + Everything the decoder is lives in one pickle: the fitted estimator, the + ordered training labels (the row order of ``F`` its columns refer to), the + format tag and whatever provenance ``metadata`` carries. The layout is + bdpy's own ``{'model': ..., 'y_shape': ...}`` with extra keys, which + ``bdpy.ml.ModelTest`` ignores, so an old reader still works. + + Keeping the labels *with* the coefficients is the point: they cannot be + reordered, renamed or replaced independently of the model they index. + ''' + makedir_ifnot(model_dir) + train_labels = [str(label) for label in train_labels] + + artifact: Dict[str, Any] = { + 'model': model, + 'y_shape': (len(train_labels),), + 'format': FORMAT_NAME, + 'version': FORMAT_VERSION, + 'train_labels': train_labels, + } + artifact.update(metadata or {}) + + model_path = model_file_path(model_dir) + with open(model_path, 'wb') as f: + pickle.dump(artifact, f, protocol=PICKLE_PROTOCOL) + return model_path + + +def load_factorized_model(model_dir: str) -> Dict[str, Any]: + '''Read the factorized decoder, checking it is one this code understands.''' + with open(model_file_path(model_dir), 'rb') as f: + artifact = pickle.load(f) + + if not isinstance(artifact, dict) or artifact.get('format') != FORMAT_NAME: + raise ValueError('%s is not a %s decoder' % (model_dir, FORMAT_NAME)) + if artifact.get('version') != FORMAT_VERSION: + raise ValueError( + 'Unsupported %s version %r in %s (this code writes version %d)' + % (FORMAT_NAME, artifact.get('version'), model_dir, + FORMAT_VERSION)) + return artifact diff --git a/tests/helpers/legacy_ridge.py b/tests/helpers/legacy_ridge.py index 48c30b1..0ca677f 100644 --- a/tests/helpers/legacy_ridge.py +++ b/tests/helpers/legacy_ridge.py @@ -33,6 +33,12 @@ 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)``.""" @@ -158,3 +164,35 @@ 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 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 index 3ffc0a6..340913b 100644 --- a/tests/helpers/pipeline.py +++ b/tests/helpers/pipeline.py @@ -13,30 +13,32 @@ import numpy as np from bdpy.dataform import load_array +import evaluation import predict_feature +import ridge_factorization import train_decoder_sklearn_ridge -NORM_KEYS = ('x_mean', 'x_norm', 'y_mean', 'y_norm') +BRAIN_NORM_KEYS = ('x_mean', 'x_norm') +FEATURE_NORM_KEYS = ('y_mean', 'y_norm') def run_training(dataset, output_dir: str, alpha: float = 100, - chunk_axis: Optional[int] = 1, - layers: Optional[Sequence[str]] = None, 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``.""" + """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, - [dataset.train_features_dir], 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), alpha=alpha, - chunk_axis=chunk_axis, analysis_name=analysis_name, ) @@ -48,11 +50,14 @@ def run_prediction(dataset, decoder_dir: str, output_dir: str, 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. + 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} @@ -64,23 +69,75 @@ def run_prediction(dataset, decoder_dir: str, output_dir: str, 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: - kwargs['training_features_paths'] = [dataset.train_features_dir] + # 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 model_dir(decoder_dir: str, layer: str, subject: str, roi: str) -> str: - return os.path.join(decoder_dir, layer, subject, roi, 'model') +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_norm_params(decoder_dir: str, layer: str, subject: str, - roi: str) -> Dict[str, np.ndarray]: - directory = model_dir(decoder_dir, layer, subject, roi) +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 NORM_KEYS} + 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, diff --git a/tests/helpers/synthetic.py b/tests/helpers/synthetic.py index 60fb3a9..ea7e28e 100644 --- a/tests/helpers/synthetic.py +++ b/tests/helpers/synthetic.py @@ -219,3 +219,56 @@ def make_dataset(root: str, seed: int = 0, unique_train_labels=unique_train_labels, unique_test_labels=unique_test_labels, ) + + +def add_layer(dataset: SyntheticDataset, name: str, shape: Tuple[int, ...], + seed: int = 1) -> np.ndarray: + """Add one more feature layer to an existing dataset's feature tree.""" + rng = np.random.RandomState(seed) + n_stimuli = len(dataset.unique_train_labels) + array = rng.randn(n_stimuli, *shape).astype(np.float32) + write_features(dataset.train_features_dir, {name: shape}, + dataset.unique_train_labels, {name: array}) + dataset.layer_shapes[name] = shape + dataset.features[name] = array + return array + + +def write_test_features(dataset: SyntheticDataset, seed: int = 21) -> str: + """Write a ``Features`` tree for the *test* stimuli and return its path. + + ``evaluation.py`` compares the decoded features against the true features + of the test stimuli, which the training dataset does not contain. + """ + root = os.path.join(dataset.root, 'features', 'test') + write_feature_tree(root, dataset.layer_shapes, dataset.unique_test_labels, + seed=seed) + return root + + +def write_feature_index(path: str, + index_by_layer: Dict[str, Sequence[int]]) -> None: + """Write a feature-index ``.mat`` as ``Features(feature_index=...)`` expects. + + bdpy reads it as ``hdf5storage.loadmat(path)['index']`` and then indexes the + result by layer name, so ``index`` is a mapping from layer to a 1-D array of + zero-based column indices into the C-order-flattened features. + """ + import hdf5storage + + payload = {layer: np.asarray(index).ravel() + for layer, index in index_by_layer.items()} + os.makedirs(os.path.dirname(path) or '.', exist_ok=True) + hdf5storage.savemat(path, {'index': payload}, format='7.3', + oned_as='column', store_python_metadata=True) + + +def write_feature_tree(root: str, layer_shapes: Dict[str, Tuple[int, ...]], + labels: Sequence[str], seed: int = 0 + ) -> Dict[str, np.ndarray]: + """Write an independent ``Features`` tree and return its arrays.""" + rng = np.random.RandomState(seed) + arrays = {layer: rng.randn(len(labels), *shape).astype(np.float32) + for layer, shape in layer_shapes.items()} + write_features(root, layer_shapes, labels, arrays) + return arrays diff --git a/tests/test_chunking.py b/tests/test_chunking.py index d7472f1..40de131 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -1,10 +1,15 @@ """Chunking of the feature dimension. -bdpy chunks the target along ``chunk_axis`` when the feature array is at least -3-D (``Y.ndim >= chunk_ndim + 1``, ``chunk_ndim=2``), fitting one model per -index along that axis and flattening/unflattening each chunk with -``order='F'``. The shipped config only enables ``fc6``/``fc7``/``fc8``, which -are 2-D and never reach this path, so it is pinned here explicitly. +Chunking used to be a *training* concern: bdpy fitted one Ridge per index along +``chunk_axis`` and flattened each chunk with ``order='F'``. In the factorized +decoder the fitted model no longer touches the feature dimensions at all, so +chunking is purely a prediction-time memory device: the linear combination with +the training features is evaluated one index at a time along ``chunk_axis``. + +Either way ``chunk_axis`` must not change the result, and each feature position +must depend only on its own training features. The shipped config only enables +``fc6``/``fc7``/``fc8``, which are 2-D and never reach this path, so it is +pinned explicitly here. """ from __future__ import annotations @@ -13,91 +18,129 @@ import pickle import numpy as np +import pytest +from ridge_factorization import combine_features from tests.helpers import pipeline ALPHA = 100 def _model_files(decoder_dir, layer, subject='sub-01', roi='VC'): - directory = pipeline.model_dir(decoder_dir, layer, subject, roi) + directory = pipeline.brain_model_dir(decoder_dir, subject, roi) return sorted(f for f in os.listdir(directory) if f.endswith('.pkl.gz')) -def test_multidimensional_features_are_chunked(dataset, decoder_dir): - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, chunk_axis=1, - layers=['conv_like']) +def test_one_model_per_layer_regardless_of_feature_dimensionality(dataset, + decoder_dir): + """No per-chunk model files: the stored model is the stimulus-basis Ridge.""" + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) - n_chunks = dataset.layer_shapes['conv_like'][0] - assert _model_files(decoder_dir, 'conv_like') == [ - '%08d.pkl.gz' % i for i in range(n_chunks)] + for layer in dataset.layers: + assert _model_files(decoder_dir, layer) == ['model.pkl.gz'] -def test_two_dimensional_features_are_not_chunked(dataset, decoder_dir): - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, chunk_axis=1, - layers=['fc_like']) +def test_stored_model_targets_the_stimulus_basis(dataset, decoder_dir): + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) - assert _model_files(decoder_dir, 'fc_like') == ['model.pkl.gz'] + n_stimuli = len(np.unique(dataset.labels('sub-01', 'train'))) + for layer in dataset.layers: + directory = pipeline.brain_model_dir(decoder_dir, 'sub-01', 'VC') + with open(os.path.join(directory, 'model.pkl.gz'), 'rb') as f: + payload = pickle.load(f) + assert payload['y_shape'] == (n_stimuli,) + n_voxels = dataset.brain('sub-01', 'VC').shape[1] + assert payload['model'].coef_.shape == (n_stimuli, n_voxels) -def test_chunk_axis_none_trains_a_single_model(dataset, decoder_dir): - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, chunk_axis=None, - layers=['conv_like']) +@pytest.mark.parametrize('chunk_axis', [1, 2, None]) +def test_chunk_axis_does_not_change_the_prediction(dataset, tmp_path, + chunk_axis): + """``chunk_axis`` blocks the computation; it must not move the numbers.""" + layer, subject, roi = 'conv_like', 'sub-01', 'VC' - assert _model_files(decoder_dir, 'conv_like') == ['model.pkl.gz'] + reference_decoder = str(tmp_path / 'decoders_ref') + reference_decoded = str(tmp_path / 'decoded_ref') + pipeline.run_training(dataset, reference_decoder, alpha=ALPHA, + analysis_name='train_ref') + pipeline.run_prediction(dataset, reference_decoder, reference_decoded, + chunk_axis=None, layers=[layer], + analysis_name='predict_ref') + reference = pipeline.read_decoded_features( + reference_decoded, layer, subject, roi, dataset.unique_test_labels) + + decoder_dir = str(tmp_path / ('decoders_%s' % chunk_axis)) + decoded_dir = str(tmp_path / ('decoded_%s' % chunk_axis)) + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, + analysis_name='train_%s' % chunk_axis) + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=chunk_axis, layers=[layer], + analysis_name='predict_%s' % chunk_axis) + actual = pipeline.read_decoded_features( + decoded_dir, layer, subject, roi, dataset.unique_test_labels) + + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize('chunk_axis', [1, 2, 3]) +def test_combine_features_matches_the_unblocked_contraction(chunk_axis): + """Blocking contracts over the same axis, to within rounding. + + The contraction runs over the stimulus axis only, so every output element is + mathematically independent of the blocking. It is not guaranteed to be + bit-identical: BLAS may accumulate a width-1 slice in a different order than + the full array, which shows up in float32 (see the float32 case below). + """ + rng = np.random.RandomState(0) + coefficients = rng.randn(5, 7) + features = rng.randn(7, 3, 4, 2) + whole = combine_features(coefficients, features, chunk_axis=None) + blocked = combine_features(coefficients, features, chunk_axis=chunk_axis) -def test_saved_chunk_shape_and_fortran_order(dataset, decoder_dir): - """Each chunk records ``y_shape`` with the chunk axis kept at size 1.""" - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, chunk_axis=1, - layers=['conv_like']) + assert blocked.shape == whole.shape + np.testing.assert_allclose(blocked, whole, rtol=1e-12, atol=1e-12) - shape = dataset.layer_shapes['conv_like'] - directory = pipeline.model_dir(decoder_dir, 'conv_like', 'sub-01', 'VC') - for name in _model_files(decoder_dir, 'conv_like'): - with open(os.path.join(directory, name), 'rb') as f: - payload = pickle.load(f) - assert payload['y_shape'] == (1,) + shape[1:] - # The fitted model is flat over the chunk's units. - assert payload['model'].coef_.shape[0] == int(np.prod(shape[1:])) +def test_combine_features_blocking_error_stays_at_rounding_level(): + """At realistic float32 sizes the two orders differ only by rounding.""" + rng = np.random.RandomState(0) + coefficients = rng.randn(50, 600).astype(np.float32) + features = rng.randn(600, 512, 14, 14).astype(np.float32) + + whole = combine_features(coefficients, features, chunk_axis=None) + blocked = combine_features(coefficients, features, chunk_axis=1) + + np.testing.assert_allclose(blocked, whole, rtol=1e-5, + atol=1e-5 * np.max(np.abs(whole))) + + +def test_combine_features_ignores_chunk_axis_for_2d_features(): + """2-D features are below bdpy's chunking threshold, as before.""" + rng = np.random.RandomState(0) + coefficients = rng.randn(4, 6) + features = rng.randn(6, 9) + + np.testing.assert_array_equal( + combine_features(coefficients, features, chunk_axis=1), + combine_features(coefficients, features, chunk_axis=None)) -def test_chunked_prediction_matches_unchunked(dataset, tmp_path): - """``chunk_axis`` is a memory device: it must not change the result.""" - layer, subject, roi = 'conv_like', 'sub-01', 'VC' - results = {} - for tag, chunk_axis in (('chunked1', 1), ('chunked2', 2), ('whole', None)): - decoder_dir = str(tmp_path / ('decoders_%s' % tag)) - decoded_dir = str(tmp_path / ('decoded_%s' % tag)) - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, - chunk_axis=chunk_axis, layers=[layer], - analysis_name='train_%s' % tag) - pipeline.run_prediction(dataset, decoder_dir, decoded_dir, - chunk_axis=chunk_axis, layers=[layer], - analysis_name='predict_%s' % tag) - results[tag] = pipeline.read_decoded_features( - decoded_dir, layer, subject, roi, dataset.unique_test_labels) - - np.testing.assert_allclose(results['chunked1'], results['whole'], - rtol=1e-4, atol=1e-5) - np.testing.assert_allclose(results['chunked2'], results['whole'], - rtol=1e-4, atol=1e-5) - - -def test_each_chunk_depends_only_on_its_own_features(dataset, tmp_path): +def test_each_feature_channel_depends_only_on_its_own_features(dataset, + tmp_path): """A per-position shift of one channel reappears at exactly those positions. - This is what makes the ``order='F'`` reshape round trip observable. The - perturbation has to differ *within* the chunk: a uniform shift of the whole - channel would still pass if the flatten and unflatten disagreed about C + The perturbation has to differ *within* the channel. A uniform shift of the + whole channel would still pass if a flatten/unflatten pair disagreed about C versus Fortran order, because a transposition of identical values is - invisible. With distinct values per position, a mismatched pair would move - the response to the wrong position inside the channel. + invisible; with distinct values per position, a mismatch would move the + response to the wrong position inside the channel. A constant per-position shift passes through the decoder exactly: it moves ``y_mean`` by ``delta`` and leaves ``y_norm`` untouched, so the normalized - target is unchanged and the un-normalization adds ``delta`` back. + target is unchanged and the un-normalization adds ``delta`` back. In the + factorized decoder the same follows from the stimulus coefficients summing + to one. """ import shutil @@ -107,8 +150,7 @@ def test_each_chunk_depends_only_on_its_own_features(dataset, tmp_path): baseline_decoder = str(tmp_path / 'decoders_base') baseline_decoded = str(tmp_path / 'decoded_base') - pipeline.run_training(dataset, baseline_decoder, alpha=ALPHA, chunk_axis=1, - layers=[layer], analysis_name='train_base') + pipeline.run_training(dataset, baseline_decoder, alpha=ALPHA, analysis_name='train_base') pipeline.run_prediction(dataset, baseline_decoder, baseline_decoded, chunk_axis=1, layers=[layer], analysis_name='predict_base') @@ -134,7 +176,6 @@ def test_each_chunk_depends_only_on_its_own_features(dataset, tmp_path): perturbed_decoder = str(tmp_path / 'decoders_perturbed') perturbed_decoded = str(tmp_path / 'decoded_perturbed') pipeline.run_training(dataset, perturbed_decoder, alpha=ALPHA, - chunk_axis=1, layers=[layer], analysis_name='train_perturbed') pipeline.run_prediction(dataset, perturbed_decoder, perturbed_decoded, chunk_axis=1, layers=[layer], @@ -151,5 +192,4 @@ def test_each_chunk_depends_only_on_its_own_features(dataset, tmp_path): # Nothing else moves. other = np.delete(shift, 1, axis=1) - np.testing.assert_allclose(other, np.zeros_like(other), - rtol=0, atol=1e-4) + np.testing.assert_allclose(other, np.zeros_like(other), rtol=0, atol=1e-4) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 0000000..67b5401 --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,54 @@ +"""The factorized pipeline still satisfies the evaluation file contract. + +``evaluation.py`` is not changed by the factorization, and it reads the +training-feature statistics from +``////model/y_{mean,norm}.mat``. The factorized +decoder cannot write them at training time -- it never opens a feature file -- +so ``predict_feature.py`` fills them in. This test exists to catch that +contract breaking: nothing else in the suite runs ``evaluation.py``. +""" + +from __future__ import annotations + +import os + +import numpy as np + +from tests.helpers import pipeline, synthetic + +ALPHA = 100 +CHUNK_AXIS = 1 +LAYER = 'fc_like' +ONE_SUBJECT = ['sub-01'] +ONE_ROI = {'VC': 'ROI_VC = 1'} + + +def test_train_predict_evaluate_end_to_end(dataset, tmp_path): + """Train, predict, then run the unchanged ``evaluation.py`` over it.""" + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, + subjects=ONE_SUBJECT, rois=ONE_ROI) + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=CHUNK_AXIS, layers=[LAYER], + subjects=ONE_SUBJECT, rois=ONE_ROI) + + # The statistics are where evaluation.py looks them up, unchanged: + # ////model/. + for key in ('y_mean', 'y_norm'): + assert os.path.isfile(os.path.join( + pipeline.layer_model_dir(decoder_dir, LAYER, 'sub-01', 'VC'), + '%s.mat' % key)) + + output = pipeline.run_evaluation( + dataset, decoder_dir, decoded_dir, synthetic.write_test_features(dataset), + layers=[LAYER], subjects=ONE_SUBJECT, rois=ONE_ROI) + + from evaluation import ResultsStore + results = ResultsStore(output) + for metric in ('profile_correlation', 'pattern_correlation', + 'identification_accuracy'): + value = results.get(layer=LAYER, subject='sub-01', roi='VC', + metric=metric) + assert value is not None and np.asarray(value).size > 0, metric diff --git a/tests/test_factorized_artifacts.py b/tests/test_factorized_artifacts.py new file mode 100644 index 0000000..93423a2 --- /dev/null +++ b/tests/test_factorized_artifacts.py @@ -0,0 +1,122 @@ +"""The factorized decoder artifact: contents, layer independence, robustness.""" + +from __future__ import annotations + +import os +import pickle + +import numpy as np +import pytest + +import ridge_factorization +from tests.helpers import pipeline, synthetic + +ALPHA = 100 +ONE_SUBJECT = ['sub-01'] +ONE_ROI = {'VC': 'ROI_VC = 1'} + + +def test_the_artifact_carries_the_model_and_its_stimulus_order(dataset, + decoder_dir): + """One file holds everything the decoder is. + + The coefficients are indexed by training stimulus, so the ordered labels + live in the same pickle: they cannot be reordered, renamed or replaced + independently of the model they index. Nothing about the features is + recorded -- training never reads one -- and not even a layer, since the + model is shared by all of them. + """ + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + + directory = pipeline.brain_model_dir(decoder_dir, 'sub-01', 'VC') + artifact = ridge_factorization.load_factorized_model(directory) + expected = [str(lb) for lb in np.unique(dataset.labels('sub-01', 'train'))] + + assert artifact['format'] == ridge_factorization.FORMAT_NAME + assert artifact['version'] == ridge_factorization.FORMAT_VERSION + assert artifact['train_labels'] == expected + assert artifact['y_shape'] == (len(expected),) + assert artifact['model'].coef_.shape[0] == len(expected) + + assert artifact['alpha'] == float(ALPHA) + assert artifact['subject'] == 'sub-01' + assert artifact['roi'] == 'VC' + assert artifact['dtype'] == 'float32' + assert artifact['n_trials'] == len(dataset.labels('sub-01', 'train')) + assert artifact['n_voxels'] == dataset.brain('sub-01', 'VC').shape[1] + + for absent in ('layer', 'chunk_axis', 'feature_source', 'feature_shape', + 'feature_digest', 'feature_index_file'): + assert absent not in artifact, absent + + +def test_the_brain_side_model_is_stored_once(dataset, decoder_dir, + decoded_dir): + """The stimulus-basis Ridge does not depend on the layer, so it is not + copied under every layer -- the per-layer directories that prediction + creates hold only the statistics ``evaluation.py`` reads. + """ + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + pipeline.run_prediction(dataset, decoder_dir, decoded_dir) + + for subject in dataset.train_fmri: + for roi in dataset.rois: + assert sorted(os.listdir( + pipeline.brain_model_dir(decoder_dir, subject, roi))) == [ + 'info.yaml', 'model.pkl.gz', 'x_mean.mat', 'x_norm.mat'] + for layer in dataset.layers: + sidecars = pipeline.layer_model_dir(decoder_dir, layer, + subject, roi) + assert sorted(os.listdir(sidecars)) == ['y_mean.mat', + 'y_norm.mat'] + + +def test_model_size_does_not_grow_with_the_feature_dimension(dataset, + decoder_dir, + decoded_dir, + tmp_path): + """The stored model scales with n_train_stimuli x n_voxels only. + + This is the whole point of the factorization: the direct implementation + stored n_features x n_voxels coefficients, so a 50x larger layer meant a + 50x larger model file. Here the layer does not enter the model at all -- + adding one changes the sidecars and nothing else. + """ + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + before = os.path.getsize(os.path.join( + pipeline.brain_model_dir(decoder_dir, 'sub-01', 'VC'), 'model.pkl.gz')) + + small_units = int(np.prod(dataset.layer_shapes['fc_like'])) + synthetic.add_layer(dataset, 'big_like', (small_units * 50,)) + + wider_dir = str(tmp_path / 'decoders_with_a_big_layer') + pipeline.run_training(dataset, wider_dir, alpha=ALPHA, + analysis_name='training_wide') + after = os.path.getsize(os.path.join( + pipeline.brain_model_dir(wider_dir, 'sub-01', 'VC'), 'model.pkl.gz')) + assert after == before + + pipeline.run_prediction(dataset, wider_dir, decoded_dir, + layers=['fc_like', 'big_like']) + statistics = { + layer: pipeline.read_feature_statistics(wider_dir, layer, 'sub-01', + 'VC')['y_mean'].size + for layer in ('fc_like', 'big_like')} + assert statistics['big_like'] == 50 * statistics['fc_like'] + + +def test_the_format_version_is_checked(dataset, decoder_dir): + """An artifact from a future version is refused, not misread.""" + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, + subjects=ONE_SUBJECT, rois=ONE_ROI) + + directory = pipeline.brain_model_dir(decoder_dir, 'sub-01', 'VC') + path = ridge_factorization.model_file_path(directory) + with open(path, 'rb') as f: + artifact = pickle.load(f) + artifact['version'] = ridge_factorization.FORMAT_VERSION + 1 + with open(path, 'wb') as f: + pickle.dump(artifact, f) + + with pytest.raises(ValueError, match='Unsupported'): + ridge_factorization.load_factorized_model(directory) diff --git a/tests/test_lazy_feature_statistics.py b/tests/test_lazy_feature_statistics.py new file mode 100644 index 0000000..ac1a4a1 --- /dev/null +++ b/tests/test_lazy_feature_statistics.py @@ -0,0 +1,312 @@ +"""The training-feature statistics: compatibility sidecars, written lazily. + +``evaluation.py`` normalizes the pattern correlation with the per-unit mean and +SD of the training features, read from +``////model/``. The factorized decoder cannot +produce them: training never opens a feature file, and prediction does not need +them either (the feature normalization cancels). So ``predict_feature.py`` -- +the step that has the features -- writes them there, purely for the unchanged +evaluation pipeline. + +Three things make that work: the write happens *before* the "results directory +already exists" skip, so a run whose features were already decoded still gets +them; it is atomic, because parallel prediction workers are a supported mode and +two runs with different analysis names are not excluded by the DistComp lock; +and the statistics are computed once per ``(layer, ordered training labels)`` +and then written into every decoder directory that needs them. Existing files +are never overwritten: one decoder directory belongs to one training feature +configuration. +""" + +from __future__ import annotations + +import os +import threading + +import numpy as np +import pytest +from bdpy.dataform import load_array, save_array + +import predict_feature +import ridge_factorization +from ridge_factorization import save_array_atomic +from tests.helpers import pipeline + +ALPHA = 100 +CHUNK_AXIS = 1 +LAYER = 'fc_like' +ONE_SUBJECT = ['sub-01'] +ONE_ROI = {'VC': 'ROI_VC = 1'} +KEYS = ('y_mean', 'y_norm') + + +def _train(dataset, decoder_dir, **kwargs): + return pipeline.run_training( + dataset, decoder_dir, alpha=ALPHA, subjects=ONE_SUBJECT, rois=ONE_ROI, **kwargs) + + +def _predict(dataset, decoder_dir, decoded_dir, **kwargs): + return pipeline.run_prediction( + dataset, decoder_dir, decoded_dir, chunk_axis=CHUNK_AXIS, + layers=[LAYER], subjects=ONE_SUBJECT, rois=ONE_ROI, **kwargs) + + +def _statistics_paths(decoder_dir, roi='VC'): + directory = pipeline.layer_model_dir(decoder_dir, LAYER, 'sub-01', roi) + return {key: os.path.join(directory, '%s.mat' % key) for key in KEYS} + + +# The materialization ########################################################## + +def test_a_second_prediction_run_does_not_rewrite_the_statistics(dataset, + tmp_path): + """The cheap existence check comes first, so nothing is recomputed.""" + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + _train(dataset, decoder_dir) + _predict(dataset, decoder_dir, decoded_dir) + + paths = _statistics_paths(decoder_dir) + before = {key: os.stat(path).st_mtime_ns for key, path in paths.items()} + + _predict(dataset, decoder_dir, decoded_dir, analysis_name='again') + + after = {key: os.stat(path).st_mtime_ns for key, path in paths.items()} + assert after == before + + +def test_the_statistics_are_materialized_past_the_already_done_skip( + dataset, tmp_path, capsys): + """The reason the write runs before the skip. + + A prediction whose features were already decoded would otherwise skip + forever, leaving ``evaluation.py`` permanently unable to run against it. + """ + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + _train(dataset, decoder_dir) + _predict(dataset, decoder_dir, decoded_dir) + + paths = _statistics_paths(decoder_dir) + for path in paths.values(): + os.unlink(path) + + decoded_before = sorted( + (name, os.stat(os.path.join(decoded_dir, LAYER, 'sub-01', 'VC', + name)).st_mtime_ns) + for name in os.listdir(os.path.join(decoded_dir, LAYER, 'sub-01', + 'VC'))) + + capsys.readouterr() + _predict(dataset, decoder_dir, decoded_dir, analysis_name='rerun') + + # Every (layer, subject, ROI) hit the skip ... + assert 'is already done. Skipped.' in capsys.readouterr().out + decoded_after = sorted( + (name, os.stat(os.path.join(decoded_dir, LAYER, 'sub-01', 'VC', + name)).st_mtime_ns) + for name in os.listdir(os.path.join(decoded_dir, LAYER, 'sub-01', + 'VC'))) + assert decoded_after == decoded_before + + # ... and the statistics came back anyway, with the right values. + feat = dataset.feature_matrix(LAYER, dataset.unique_train_labels) + np.testing.assert_allclose( + load_array(paths['y_mean'], key='y_mean'), + np.mean(feat, axis=0)[np.newaxis].astype(np.float32), + rtol=1e-5, atol=1e-6) + np.testing.assert_allclose( + load_array(paths['y_norm'], key='y_norm'), + np.std(feat, axis=0, ddof=1)[np.newaxis].astype(np.float32), + rtol=1e-5, atol=1e-6) + + +def test_a_finished_prediction_reruns_without_the_training_features(dataset, + tmp_path): + """Nothing left to decode, so nothing to verify or materialize either.""" + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + _train(dataset, decoder_dir) + _predict(dataset, decoder_dir, decoded_dir) + + _predict(dataset, decoder_dir, decoded_dir, features_paths=[], + analysis_name='no_features') + + +def test_an_unwritable_output_directory_warns_instead_of_failing(dataset, + tmp_path, + monkeypatch): + """Decoding must still work when the statistics cannot be written. + + The write is made to fail here rather than by permission bits, since the + test suite may run as a user for whom no permission bit denies anything. + """ + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + _train(dataset, decoder_dir) + + def read_only(*args, **kwargs): + raise IOError('Read-only file system') + + monkeypatch.setattr(predict_feature, 'save_array_atomic', read_only) + with pytest.warns(UserWarning, match='read-only'): + _predict(dataset, decoder_dir, decoded_dir) + + assert not os.path.exists(_statistics_paths(decoder_dir)['y_mean']) + assert pipeline.read_decoded_features( + decoded_dir, LAYER, 'sub-01', 'VC', + dataset.unique_test_labels).shape[0] == len(dataset.unique_test_labels) + + +def test_the_statistics_are_computed_once_per_layer_and_label_set( + dataset, decoder_dir, decoded_dir, monkeypatch): + """Not once per decoder: they do not depend on the ROI. + + At DeepRecon scale each computation is two passes over a 15 GB tensor, and + there are 9 ROIs per subject. + """ + computed = [] + original = ridge_factorization.feature_statistics + + def spy(features): + computed.append(features.shape) + return original(features) + + monkeypatch.setattr(predict_feature, 'feature_statistics', spy) + + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=CHUNK_AXIS) + + # Both synthetic subjects see the same stimuli, so one computation per + # layer covers every (subject, ROI) of the run ... + assert len(dataset.rois) > 1 and len(dataset.test_fmri) > 1 + assert len(computed) == len(dataset.layers) + + # ... while every decoder directory still gets its own copy, because that + # is where evaluation.py looks. + for layer in dataset.layers: + for subject in dataset.test_fmri: + for roi in dataset.rois: + assert pipeline.read_feature_statistics(decoder_dir, layer, + subject, roi) + + +def test_existing_statistics_are_never_overwritten(dataset, tmp_path): + """One decoder directory belongs to one training feature configuration. + + Re-pointing it at another ``Y_tr`` or feature index is not a supported + workflow, so prediction fills the sidecars in when they are missing and + leaves them alone otherwise -- it never decides that what is there is + wrong. + """ + decoder_dir = str(tmp_path / 'decoders') + _train(dataset, decoder_dir) + _predict(dataset, decoder_dir, str(tmp_path / 'decoded')) + + paths = _statistics_paths(decoder_dir) + planted = np.full((1, 3), 7.0, dtype=np.float32) + for path, key in ((paths[key], key) for key in KEYS): + save_array(path, planted, key=key, dtype=np.float32, sparse=False) + + _predict(dataset, decoder_dir, str(tmp_path / 'decoded_again'), + analysis_name='again') + + for key in KEYS: + np.testing.assert_array_equal(load_array(paths[key], key=key), planted) + + +# save_array_atomic ############################################################ + +def _temp_files(directory): + return [name for name in os.listdir(directory) if name.startswith('.tmp')] + + +def test_atomic_save_writes_a_loadable_file(tmp_path): + """The temp name has to keep the '.mat' suffix that hdf5storage appends.""" + path = str(tmp_path / 'y_mean.mat') + array = np.arange(6, dtype=np.float32)[np.newaxis] + + save_array_atomic(path, array, key='y_mean', dtype=np.float32) + + np.testing.assert_array_equal(load_array(path, key='y_mean'), array) + assert _temp_files(str(tmp_path)) == [] + + +def test_atomic_save_replaces_an_existing_file(tmp_path): + path = str(tmp_path / 'y_mean.mat') + save_array_atomic(path, np.zeros((1, 6), dtype=np.float32), key='y_mean') + + replacement = np.arange(6, dtype=np.float32)[np.newaxis] + save_array_atomic(path, replacement, key='y_mean') + + np.testing.assert_array_equal(load_array(path, key='y_mean'), replacement) + assert _temp_files(str(tmp_path)) == [] + + +def test_a_failed_atomic_save_leaves_the_original_intact(tmp_path, + monkeypatch): + path = str(tmp_path / 'y_mean.mat') + original = np.arange(6, dtype=np.float32)[np.newaxis] + save_array_atomic(path, original, key='y_mean') + + def broken(*args, **kwargs): + raise IOError('disk full') + + monkeypatch.setattr(ridge_factorization, 'save_array', broken) + with pytest.raises(IOError): + save_array_atomic(path, np.zeros((1, 6), dtype=np.float32), + key='y_mean') + + np.testing.assert_array_equal(load_array(path, key='y_mean'), original) + assert _temp_files(str(tmp_path)) == [] + + +def test_a_partially_written_temp_file_is_cleaned_up(tmp_path, monkeypatch): + """The failure that would otherwise leave rubbish next to the target.""" + path = str(tmp_path / 'y_mean.mat') + real_save = ridge_factorization.save_array + + def write_then_fail(save_file, array, **kwargs): + real_save(save_file, array, **kwargs) + raise IOError('interrupted after the write') + + monkeypatch.setattr(ridge_factorization, 'save_array', write_then_fail) + with pytest.raises(IOError): + save_array_atomic(path, np.zeros((1, 6), dtype=np.float32), + key='y_mean') + + assert not os.path.exists(path) + assert _temp_files(str(tmp_path)) == [] + + +def test_concurrent_atomic_saves_leave_one_valid_file(tmp_path): + """Two writers past the existence check, as parallel workers can be. + + Each writer builds its own temp name, so the loser's rename is the only + thing that races, and a reader sees one complete file either way. + """ + path = str(tmp_path / 'y_norm.mat') + candidates = [np.full((1, 6), value, dtype=np.float32) + for value in (1.0, 2.0)] + start = threading.Barrier(len(candidates)) + errors = [] + + def write(array): + try: + start.wait() + save_array_atomic(path, array, key='y_norm') + except BaseException as error: # pragma: no cover - reported below + errors.append(error) + + threads = [threading.Thread(target=write, args=(array,)) + for array in candidates] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert errors == [] + loaded = load_array(path, key='y_norm') + assert any(np.array_equal(loaded, array) for array in candidates) + assert _temp_files(str(tmp_path)) == [] diff --git a/tests/test_legacy_model_compat.py b/tests/test_legacy_model_compat.py new file mode 100644 index 0000000..3a9d330 --- /dev/null +++ b/tests/test_legacy_model_compat.py @@ -0,0 +1,203 @@ +"""Reading decoders saved by the direct (pre-factorization) implementation. + +New training runs write the factorized format, but large decoders produced by +the previous implementation may already exist, so ``predict_feature.py`` keeps a +compatibility path. The format is detected from the presence of a +``model.pkl.gz`` on the brain-side path, which only a factorized decoder has. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import ridge_factorization +from ridge_factorization import is_factorized_model_dir +from tests.helpers import legacy_ridge, pipeline + +ALPHA = 100 +CHUNK_AXIS = 1 + + +def build_legacy_decoder(dataset, decoder_dir, layers=None, + subjects=('sub-01',), rois=('VC',), + chunk_axis=CHUNK_AXIS): + """Train with the reference implementation and save the legacy layout.""" + trained = {} + for layer in (layers or dataset.layers): + for subject in subjects: + for roi in rois: + brain_labels = dataset.labels(subject, 'train') + feat_labels = list(np.unique(brain_labels)) + entry = legacy_ridge.legacy_train( + dataset.brain(subject, roi, 'train'), brain_labels, + dataset.feature_matrix(layer, feat_labels), feat_labels, + alpha=ALPHA, chunk_axis=chunk_axis) + legacy_ridge.write_legacy_decoder( + pipeline.layer_model_dir(decoder_dir, layer, subject, roi), + entry) + trained[(layer, subject, roi)] = entry + return trained + + +def test_format_detection(dataset, tmp_path): + factorized_dir = str(tmp_path / 'factorized') + pipeline.run_training(dataset, factorized_dir, alpha=ALPHA) + assert is_factorized_model_dir( + pipeline.brain_model_dir(factorized_dir, 'sub-01', 'VC')) + + # A legacy decoder has no model on the brain-side path -- that is the + # discriminator, and it is the path prediction asks about. + legacy_dir = str(tmp_path / 'legacy') + build_legacy_decoder(dataset, legacy_dir, layers=['fc_like']) + assert not is_factorized_model_dir( + pipeline.brain_model_dir(legacy_dir, 'sub-01', 'VC')) + + # And a legacy model pickle is refused rather than misread if something + # does hand it to the factorized loader. + with pytest.raises(ValueError, match='not a sklearn_ridge_factorized'): + ridge_factorization.load_factorized_model( + pipeline.layer_model_dir(legacy_dir, 'fc_like', 'sub-01', 'VC')) + + +@pytest.mark.parametrize('layer', ['fc_like', 'conv_like']) +def test_legacy_decoder_predictions_are_unchanged(dataset, tmp_path, layer): + """A legacy decoder still decodes exactly as it used to, chunked or not.""" + decoder_dir = str(tmp_path / 'legacy_decoders') + decoded_dir = str(tmp_path / 'legacy_decoded') + + trained = build_legacy_decoder(dataset, decoder_dir, layers=[layer]) + + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=CHUNK_AXIS, layers=[layer], + subjects=['sub-01'], rois={'VC': 'ROI_VC = 1'}) + + test_brain, test_labels = legacy_ridge.average_test_brain( + dataset.brain('sub-01', 'VC', 'test'), dataset.labels('sub-01', 'test')) + expected = legacy_ridge.legacy_predict(trained[(layer, 'sub-01', 'VC')], + test_brain) + actual = pipeline.read_decoded_features(decoded_dir, layer, 'sub-01', 'VC', + test_labels) + + np.testing.assert_allclose(actual, expected.astype(np.float32), + rtol=1e-5, atol=1e-6) + + +def test_legacy_decoder_needs_no_training_features(dataset, tmp_path): + """The legacy path is self-contained: it must not require ``F`` on disk.""" + import predict_feature + + decoder_dir = str(tmp_path / 'legacy_decoders') + decoded_dir = str(tmp_path / 'legacy_decoded') + build_legacy_decoder(dataset, decoder_dir, layers=['fc_like']) + + predict_feature.featdec_predict( + {'sub-01': dataset.test_fmri['sub-01']}, decoder_dir, + output_dir=decoded_dir, + rois={'VC': 'ROI_VC = 1'}, + label_key=dataset.label_key, + layers=['fc_like'], + average_sample=True, + chunk_axis=CHUNK_AXIS, + training_features_paths=None, + analysis_name='legacy_without_features') + + names = pipeline.decoded_feature_labels(decoded_dir, 'fc_like', 'sub-01', + 'VC') + assert names == dataset.unique_test_labels + + +def test_factorized_decoder_requires_training_features(dataset, tmp_path): + """A clear error beats a confusing one -- and it must leave no output. + + The check runs before the output directory is created, so the failed run + cannot make the next, correct one print "already done" and skip. + """ + import os + + import predict_feature + + decoder_dir = str(tmp_path / 'decoders') + decoded_dir = str(tmp_path / 'decoded') + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + + def predict(training_features_paths, analysis_name): + return predict_feature.featdec_predict( + {'sub-01': dataset.test_fmri['sub-01']}, decoder_dir, + output_dir=decoded_dir, + rois={'VC': 'ROI_VC = 1'}, + label_key=dataset.label_key, + layers=['fc_like'], + average_sample=True, + chunk_axis=CHUNK_AXIS, + training_features_paths=training_features_paths, + analysis_name=analysis_name) + + with pytest.raises(ValueError, match='training_features_paths'): + predict(None, 'factorized_without_features') + + assert not os.path.exists(os.path.join(decoded_dir, 'fc_like', 'sub-01', + 'VC')) + + predict([dataset.train_features_dir], 'factorized_with_features') + + assert pipeline.decoded_feature_labels( + decoded_dir, 'fc_like', 'sub-01', 'VC') == dataset.unique_test_labels + + +def test_legacy_prediction_still_copies_the_feature_index(dataset, tmp_path): + """The pre-factorization contract: given an index, it lands in the output. + + The direct implementation copied ``feature_index_file`` as given, and a + legacy decoder is predicted without training features, so the copy cannot + depend on them. + """ + import os + + from tests.helpers import synthetic + + decoder_dir = str(tmp_path / 'legacy_decoders') + decoded_dir = str(tmp_path / 'legacy_decoded') + build_legacy_decoder(dataset, decoder_dir, layers=['fc_like']) + + # Written into the working directory, which is what the old code resolved + # against (conftest chdirs into the test's own tmp_path). + synthetic.write_feature_index('index_local.mat', {'fc_like': [0, 2]}) + + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=CHUNK_AXIS, layers=['fc_like'], + subjects=['sub-01'], rois={'VC': 'ROI_VC = 1'}, + features_paths=[], + feature_index_file='index_local.mat', + analysis_name='legacy_with_index') + + assert os.path.isfile(os.path.join(decoded_dir, 'feature_index.mat')) + + +def test_legacy_and_factorized_decoders_agree(dataset, tmp_path): + """The compatibility path and the new path give the same numbers.""" + layer = 'conv_like' + + legacy_dir = str(tmp_path / 'legacy_decoders') + legacy_decoded = str(tmp_path / 'legacy_decoded') + build_legacy_decoder(dataset, legacy_dir, layers=[layer]) + pipeline.run_prediction(dataset, legacy_dir, legacy_decoded, + chunk_axis=CHUNK_AXIS, layers=[layer], + subjects=['sub-01'], rois={'VC': 'ROI_VC = 1'}, + analysis_name='predict_legacy') + legacy_pred = pipeline.read_decoded_features( + legacy_decoded, layer, 'sub-01', 'VC', dataset.unique_test_labels) + + new_dir = str(tmp_path / 'new_decoders') + new_decoded = str(tmp_path / 'new_decoded') + pipeline.run_training(dataset, new_dir, alpha=ALPHA, + subjects=['sub-01'], rois={'VC': 'ROI_VC = 1'}, + analysis_name='train_new') + pipeline.run_prediction(dataset, new_dir, new_decoded, + chunk_axis=CHUNK_AXIS, layers=[layer], + subjects=['sub-01'], rois={'VC': 'ROI_VC = 1'}, + analysis_name='predict_new') + new_pred = pipeline.read_decoded_features( + new_decoded, layer, 'sub-01', 'VC', dataset.unique_test_labels) + + np.testing.assert_allclose(new_pred, legacy_pred, rtol=1e-4, atol=1e-5) diff --git a/tests/test_legacy_reference.py b/tests/test_legacy_reference.py index dbb680e..7bcef23 100644 --- a/tests/test_legacy_reference.py +++ b/tests/test_legacy_reference.py @@ -57,8 +57,7 @@ def reference_prediction(dataset, layer, subject, roi, alpha=ALPHA, def test_reference_matches_scripts(dataset, decoder_dir, decoded_dir, alpha, chunk_axis, average_sample): """The reference reproduces the shipped scripts, in every configuration.""" - pipeline.run_training(dataset, decoder_dir, alpha=alpha, - chunk_axis=chunk_axis) + pipeline.run_training(dataset, decoder_dir, alpha=alpha) pipeline.run_prediction(dataset, decoder_dir, decoded_dir, chunk_axis=chunk_axis, average_sample=average_sample) @@ -70,17 +69,21 @@ def test_reference_matches_scripts(dataset, decoder_dir, decoded_dir, alpha, dataset, layer, subject, roi, alpha=alpha, chunk_axis=chunk_axis, average_sample=average_sample) + # The factorized path reaches the same numbers through a + # different product, so float32 rounding differs; a small + # alpha leaves larger coefficients and makes it visible. + rtol, atol = (1e-5, 1e-6) if alpha >= 100 else (1e-4, 1e-5) actual = pipeline.read_decoded_features( decoded_dir, layer, subject, roi, test_labels) np.testing.assert_allclose( - actual, expected.astype(np.float32), - rtol=1e-5, atol=1e-6, + actual, expected.astype(np.float32), rtol=rtol, atol=atol, err_msg='prediction mismatch for %s/%s/%s' % (layer, subject, roi)) - saved = pipeline.read_norm_params(decoder_dir, layer, subject, - roi) - for key in pipeline.NORM_KEYS: + saved = pipeline.read_norm_params(decoder_dir, subject, roi) + saved.update(pipeline.read_feature_statistics( + decoder_dir, layer, subject, roi)) + for key in legacy_ridge.NORM_KEYS: np.testing.assert_allclose( saved[key], trained[key].astype(np.float32), rtol=1e-6, atol=1e-7, diff --git a/tests/test_normalization.py b/tests/test_normalization.py new file mode 100644 index 0000000..c73229e --- /dev/null +++ b/tests/test_normalization.py @@ -0,0 +1,122 @@ +"""Normalization semantics of the factorized decoder. + +Brain normalization is unchanged. Feature normalization is where the +factorization removes work: the round trip ``(Y - y_mean) / y_norm`` at training +followed by ``* y_norm + y_mean`` at prediction cancels exactly (see +``ridge_factorization``), so the factorized prediction combines the raw training +features. ``y_mean.mat`` / ``y_norm.mat`` still exist in the decoder's +per-layer directory because ``evaluation.py`` normalizes the pattern +correlation with them, but nothing in the decoding needs them, so +``predict_feature.py`` writes them as compatibility sidecars (see +``tests/test_lazy_feature_statistics.py``). +""" + +from __future__ import annotations + +import os + +import numpy as np + +from bdpy.dataform import load_array, save_array +from ridge_factorization import normalize_brain_for_training +from tests.helpers import pipeline + +ALPHA = 100 +CHUNK_AXIS = 1 + + +def test_brain_normalization_zeroes_infinities(): + """Matches ``ModelTraining.run()``: infinities become zero, NaNs do not.""" + brain = np.array([[1.0, 2.0, 3.0], + [1.0, 4.0, 3.0]]) + mean = np.array([[0.0, 3.0, 3.0]]) + std = np.array([[0.0, 1.0, 0.0]]) + + normalized = normalize_brain_for_training(brain, mean, std) + + # Column 0: nonzero deviation over zero SD -> inf -> replaced by 0. + np.testing.assert_array_equal(normalized[:, 0], [0.0, 0.0]) + # Column 1: ordinary z-score. + np.testing.assert_allclose(normalized[:, 1], [-1.0, 1.0]) + # Column 2: 0/0 -> NaN, deliberately left as is. + assert np.all(np.isnan(normalized[:, 2])) + + +def test_feature_statistics_are_materialized_by_prediction( + dataset, decoder_dir, decoded_dir): + """Training cannot compute them -- it never reads the features. + + ``evaluation.py`` reads them from ``///model/``, so + the definition has to stay exactly what the direct implementation wrote: + per-unit statistics over the unique stimuli, with ``ddof=1``. + """ + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + + for layer in dataset.layers: + for subject in dataset.train_fmri: + for roi in dataset.rois: + assert not os.path.exists(os.path.join( + pipeline.layer_model_dir(decoder_dir, layer, subject, roi), + 'y_mean.mat')) + + pipeline.run_prediction(dataset, decoder_dir, decoded_dir, + chunk_axis=CHUNK_AXIS) + + for layer in dataset.layers: + for subject in dataset.test_fmri: + feat = dataset.feature_matrix( + layer, list(np.unique(dataset.labels(subject, 'train')))) + for roi in dataset.rois: + params = pipeline.read_feature_statistics(decoder_dir, layer, + subject, roi) + np.testing.assert_allclose( + params['y_mean'], + np.mean(feat, axis=0)[np.newaxis].astype(np.float32), + rtol=1e-5, atol=1e-6) + np.testing.assert_allclose( + params['y_norm'], + np.std(feat, axis=0, ddof=1)[np.newaxis].astype(np.float32), + rtol=1e-5, atol=1e-6) + + +def test_prediction_does_not_depend_on_the_feature_statistics(dataset, + tmp_path): + """Corrupting ``y_mean``/``y_norm`` must not change what is decoded. + + A direct demonstration that the feature normalization cancels: the + factorized prediction never reads these files. They are corrupted in the + decoder and the features are decoded again into a fresh directory, so the + second run really does see them. + """ + decoder_dir = str(tmp_path / 'decoders') + + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) + pipeline.run_prediction(dataset, decoder_dir, str(tmp_path / 'intact'), + chunk_axis=CHUNK_AXIS, + analysis_name='predict_intact') + + intact = {layer: pipeline.read_decoded_features( + str(tmp_path / 'intact'), layer, 'sub-01', 'VC', + dataset.unique_test_labels) for layer in dataset.layers} + + for layer in dataset.layers: + for subject in dataset.test_fmri: + for roi in dataset.rois: + directory = pipeline.layer_model_dir(decoder_dir, layer, + subject, roi) + for key in ('y_mean', 'y_norm'): + path = os.path.join(directory, '%s.mat' % key) + array = load_array(path, key=key) + save_array(path, array * 3.0 + 17.0, key=key, + dtype=np.float32, sparse=False) + + pipeline.run_prediction(dataset, decoder_dir, str(tmp_path / 'corrupted'), + chunk_axis=CHUNK_AXIS, + analysis_name='predict_corrupted') + + for layer in dataset.layers: + np.testing.assert_array_equal( + pipeline.read_decoded_features(str(tmp_path / 'corrupted'), layer, + 'sub-01', 'VC', + dataset.unique_test_labels), + intact[layer]) diff --git a/tests/test_predict_feature.py b/tests/test_predict_feature.py index 006b10a..40094f9 100644 --- a/tests/test_predict_feature.py +++ b/tests/test_predict_feature.py @@ -19,8 +19,7 @@ def test_output_layout_and_dtype(dataset, decoder_dir, decoded_dir): - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, - chunk_axis=CHUNK_AXIS) + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) pipeline.run_prediction(dataset, decoder_dir, decoded_dir, chunk_axis=CHUNK_AXIS) @@ -44,8 +43,7 @@ def test_output_layout_and_dtype(dataset, decoder_dir, decoded_dir): def test_average_sample_averages_repetitions(dataset, decoder_dir, decoded_dir): """With ``average_sample`` on, one file per unique stimulus is written.""" - pipeline.run_training(dataset, decoder_dir, alpha=ALPHA, - chunk_axis=CHUNK_AXIS) + pipeline.run_training(dataset, decoder_dir, alpha=ALPHA) pipeline.run_prediction(dataset, decoder_dir, decoded_dir, chunk_axis=CHUNK_AXIS, average_sample=True) @@ -57,8 +55,7 @@ def test_average_sample_averages_repetitions(dataset, decoder_dir, def test_single_trial_output_names(dataset, decoder_dir, decoded_dir): """With ``average_sample`` off, files are ``sample%06d-