perf: cache loaded feature matrices between splits and models

docs/runyourmodel.rst

@@ -31,6 +31,7 @@ Additionally, you must define a unique model name to identify your model during
 
         is_single_drug_model = True / False # TODO: set to true if your model is a single drug model (i.e., it needs to be trained for each drug separately)
         early_stopping = True / False # TODO: set to true if you want to use a part of the validation set for early stopping
+        supports_feature_caching = True # TODO: set to false if your feature loading also initializes model state, see below
         cell_line_views = ["gene_expression", "methylation"]
         drug_views = ["fingerprints"]
 
@@ -59,6 +60,15 @@ For our provided datasets, we have other loading methods implemented in the `dre
 * ``def load_drug_ids_from_csv``
 * ``load_tissues_from_csv``
 
+Loading a feature matrix is often the most expensive part of a run, and the same matrix is needed again for
+every CV split and, for single drug models, for every drug. The pipeline therefore caches what your loaders
+return and hands the same object to later model instances, keyed by model class, data path, dataset name and
+hyperparameters. Every consumer works on a ``copy()`` of it, so this is invisible to your model as long as your
+loaders only return features. If they also initialize model state, e.g., an ontology structure that ``train``
+needs later, set ``supports_feature_caching = False`` on your class: the pipeline then calls your loaders for
+every instance, as it did before the cache existed. Setting ``DREVAL_FEATURE_CACHE=0`` in the environment
+disables the cache for all models.
+
 .. code-block:: Python
 
     def load_drug_features(self, data_path: str, dataset_name: str) -> FeatureDataset:

drevalpy/experiment.py

@@ -7,7 +7,7 @@
 import tempfile
 import warnings
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import pandas as pd
@@ -625,6 +625,82 @@ def consolidate_single_drug_model_predictions(
                     )
 
 
+#: Cache for loaded feature matrices, keyed by everything they depend on.
+#:
+#: Feature matrices depend only on (model class, hyperparameters, data path, dataset), not on the CV
+#: split and not on the drug. For single drug models the pipeline nevertheless walks the whole split
+#: loop once PER DRUG, so the same matrix was re-read from disk once per split and drug. Measured on
+#: a 545 drug CTRPv2 run with 893 genes: 17.8 s per load, 3270 loads, roughly 16 h of an 18.4 h run.
+#:
+#: Reusing one object is safe because every consumer only reads it and hands it on via .copy()
+#: (train_and_predict, train_final_model, cross_study_prediction, randomization_test).
+#: Set DREVAL_FEATURE_CACHE=0 to disable, e.g., to A/B check that results are unchanged.
+_FEATURE_CACHE: dict[tuple, FeatureDataset | None] = {}
+
+#: Kept small on purpose: an entry is a full feature matrix, and the access pattern is "same key many
+#: times in a row", so a couple of slots already give the full speedup.
+_FEATURE_CACHE_MAXSIZE = 4
+
+
+def clear_feature_cache() -> None:
+    """Drop all cached feature matrices, e.g., to free memory between datasets."""
+    _FEATURE_CACHE.clear()
+
+
+def _load_features_cached(model: DRPModel, path_data: str, dataset_name: str, kind: str) -> FeatureDataset | None:
+    """
+    Load cell line or drug features, reusing an already loaded matrix when nothing they depend on changed.
+
+    :param model: built model, i.e., build_model() was already called so the hyperparameters are set
+    :param path_data: path to the data directory, e.g., data/
+    :param dataset_name: name of the dataset, e.g., GDSC2
+    :param kind: either "cell_line" or "drug"
+    :returns: the feature dataset, or None if the model does not use this kind of feature
+    """
+    loader = model.load_cell_line_features if kind == "cell_line" else model.load_drug_features
+    # Models whose loaders also initialize model state have to run them for every instance, and the cache
+    # can be switched off entirely to check that results are unchanged.
+    if not model.supports_feature_caching or os.environ.get("DREVAL_FEATURE_CACHE", "1") == "0":
+        return loader(data_path=path_data, dataset_name=dataset_name)
+    # The hyperparameters decide which views and which gene list are loaded, so they belong in the key.
+    key = (
+        type(model).__name__,
+        kind,
+        path_data,
+        dataset_name,
+        json.dumps(getattr(model, "hyperparameters", {}), sort_keys=True, default=str),
+    )
+    if key not in _FEATURE_CACHE:
+        if len(_FEATURE_CACHE) >= _FEATURE_CACHE_MAXSIZE:
+            _FEATURE_CACHE.pop(next(iter(_FEATURE_CACHE)))
+        _FEATURE_CACHE[key] = loader(data_path=path_data, dataset_name=dataset_name)
+    return _FEATURE_CACHE[key]
+
+
+def _load_cell_line_features_cached(model: DRPModel, path_data: str, dataset_name: str) -> FeatureDataset:
+    """
+    Cached variant of ``model.load_cell_line_features``, which every model has to implement.
+
+    :param model: built model
+    :param path_data: path to the data directory, e.g., data/
+    :param dataset_name: name of the dataset, e.g., GDSC2
+    :returns: the cell line feature dataset
+    """
+    return cast(FeatureDataset, _load_features_cached(model, path_data, dataset_name, "cell_line"))
+
+
+def _load_drug_features_cached(model: DRPModel, path_data: str, dataset_name: str) -> FeatureDataset | None:
+    """
+    Cached variant of ``model.load_drug_features``, which single drug models do not provide.
+
+    :param model: built model
+    :param path_data: path to the data directory, e.g., data/
+    :param dataset_name: name of the dataset, e.g., GDSC2
+    :returns: the drug feature dataset, or None
+    """
+    return _load_features_cached(model, path_data, dataset_name, "drug")
+
+
 def load_features(
     model: DRPModel, path_data: str, dataset: DrugResponseDataset
 ) -> tuple[FeatureDataset, FeatureDataset | None]:
@@ -636,8 +712,8 @@ def load_features(
     :param dataset: dataset to load features for, e.g., GDSC2
     :returns: tuple of cell line and, potentially, drug features
     """
-    cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=dataset.dataset_name)
-    drug_features = model.load_drug_features(data_path=path_data, dataset_name=dataset.dataset_name)
+    cl_features = _load_cell_line_features_cached(model, path_data, dataset.dataset_name)
+    drug_features = _load_drug_features_cached(model, path_data, dataset.dataset_name)
     return cl_features, drug_features
 
 
@@ -1105,10 +1181,10 @@ def train_and_predict(
         raise ValueError("train_dataset must have a dataset_name")
     if cl_features is None:
         print("Loading cell line features ...")
-        cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=train_dataset.dataset_name)
+        cl_features = _load_cell_line_features_cached(model, path_data, train_dataset.dataset_name)
     if drug_features is None:
         print("Loading drug features ...")
-        drug_features = model.load_drug_features(data_path=path_data, dataset_name=train_dataset.dataset_name)
+        drug_features = _load_drug_features_cached(model, path_data, train_dataset.dataset_name)
 
     cell_lines_to_keep = cl_features.identifiers if cl_features is not None else None
     drugs_to_keep = drug_features.identifiers if drug_features is not None else None
@@ -1617,8 +1693,8 @@ def train_final_model(
     print(f"Best hyperparameters for final model: {best_hpams}")
     model.build_model(hyperparameters=best_hpams)
 
-    cl_features = model.load_cell_line_features(data_path=path_data, dataset_name=full_dataset.dataset_name)
-    drug_features = model.load_drug_features(data_path=path_data, dataset_name=full_dataset.dataset_name)
+    cl_features = _load_cell_line_features_cached(model, path_data, full_dataset.dataset_name)
+    drug_features = _load_drug_features_cached(model, path_data, full_dataset.dataset_name)
     cell_lines_to_keep = cl_features.identifiers
     drugs_to_keep = drug_features.identifiers if drug_features is not None else None
 

drevalpy/models/SparseGO/sparsego.py

@@ -415,6 +415,10 @@ class SparseGOModel(DRPModel):
     cell_line_views = ["gene_expression", "mutations"]
     drug_views = ["fingerprints"]
     early_stopping = False
+    # load_cell_line_features also builds the ontology structure (layer_connections, gene2id_mapping_ont,
+    # ontology_gene_order, gene_dim_input) that train() and predict() need, so it has to run for every
+    # instance and the features must not be served from the pipeline's cache.
+    supports_feature_caching = False
 
     def __init__(self) -> None:
         """Initialize SparseGOModel."""

drevalpy/models/drp_model.py

@@ -28,13 +28,21 @@ class DRPModel(ABC):
 
     The DRPModel class is an abstract wrapper class for drug response prediction models.
     It has a boolean attribute is_single_drug_model indicating whether it is a single drug model and a boolean
-    attribute early_stopping indicating whether early stopping is used.
+    attribute early_stopping indicating whether early stopping is used. The boolean attribute
+    supports_feature_caching tells the pipeline whether the loaded feature matrices may be reused across
+    model instances.
     """
 
     # Used in the pipeline!
     early_stopping = False
     # Then, the model is trained per drug
     is_single_drug_model = False
+    # The pipeline caches the feature matrices a model loads and hands the same object to every later
+    # instance of that model, which is what makes repeated splits and single drug runs affordable.
+    # Set this to False if load_cell_line_features or load_drug_features do more than return features,
+    # i.e., if they also initialize model state. Such a loader has to run for every instance, so the
+    # pipeline then loads the features every time, as it did before the cache existed. See SparseGOModel.
+    supports_feature_caching = True
 
     def __init__(self):
         """Initialize the DRPModel instance."""