diff --git a/README.md b/README.md index 714d7d3..943f667 100644 --- a/README.md +++ b/README.md @@ -49,27 +49,36 @@ $ python predict_feature_fastl2lir.py config/deeprecon_pyfastl2lir_alpha100_vgg1 $ python evaluation.py config/deeprecon_pyfastl2lir_alpha100_vgg19_allunits.yaml ``` -### Decoding with generic regression models +### Decoding with scikit-learn Ridge regression -- Training: `train_decoder_sklearn_ridge.py` (example for scikit-learn Ridge regression) +- Training: `train_decoder_sklearn_ridge.py` - Test (prediction): `predict_feature.py` - Evaluation: `evaluation.py` - Example config file: [deeprecon_sklearn_ridge_alpha100_vgg19_allunits](config/deeprecon_sklearn_ridge_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) +- Training: `cv_train_decoder_fastl2lir.py` - Test (prediction): `cv_predict_feature_fastl2lir.py` - Evaluation: `cv_evaluation.py` - Example config file: [deeprecon_cv_pyfastl2lir_alpha100_vgg19_allunits](config/deeprecon_cv_pyfastl2lir_alpha100_vgg19_allunits.yaml) 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-