From 04ec61713044c7063b6468dfbff1946d4e42b608 Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 18:00:09 -0700 Subject: [PATCH 1/6] feat: add RequiredFeatures type and MissingFeaturesError --- biolearn/features.py | 55 ++++++++++++++++++++++++++++++++++ biolearn/test/test_features.py | 42 ++++++++++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 biolearn/features.py create mode 100644 biolearn/test/test_features.py diff --git a/biolearn/features.py b/biolearn/features.py new file mode 100644 index 0000000..1037c02 --- /dev/null +++ b/biolearn/features.py @@ -0,0 +1,55 @@ +"""Stable description of the features a model requires, plus the validator +that enforces them. + +Introduced in the 1.0 line. The return type of ``Model.required_features()`` +and ``MissingFeaturesError`` are stable public API: breaking either requires a +major version bump. +""" + +from collections.abc import Mapping +from dataclasses import dataclass + + +VALID_LAYERS = ( + "dnam", + "rna", + "protein_olink", + "protein_alamar", + "clinical", +) + + +@dataclass(frozen=True) +class RequiredFeatures(Mapping): + """What a clock consumes, as a frozen value. + + Behaves as the documented mapping ``{"layer", "features", "metadata"}`` + (so ``rf["features"]`` and ``dict(rf)`` work) and as an object with + ``.layer`` / ``.features`` / ``.metadata`` attributes. All three keys are + always present; ``features`` and ``metadata`` default to empty tuples. + """ + + layer: str + features: tuple = () + metadata: tuple = () + + def __getitem__(self, key): + return { + "layer": self.layer, + "features": list(self.features), + "metadata": list(self.metadata), + }[key] + + def __iter__(self): + return iter(("layer", "features", "metadata")) + + def __len__(self): + return 3 + + +class MissingFeaturesError(ValueError): + """Raised when a GeoData lacks features or metadata a clock requires. + + Subclasses ``ValueError`` so existing ``except ValueError`` handlers keep + working while callers can catch this specific case. + """ diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py new file mode 100644 index 0000000..b37e581 --- /dev/null +++ b/biolearn/test/test_features.py @@ -0,0 +1,42 @@ +import pytest + +from biolearn.features import ( + RequiredFeatures, + MissingFeaturesError, + VALID_LAYERS, +) + + +def test_required_features_behaves_as_object(): + rf = RequiredFeatures("dnam", ("cg1", "cg2")) + assert rf.layer == "dnam" + assert rf.features == ("cg1", "cg2") + assert rf.metadata == () + + +def test_required_features_behaves_as_mapping(): + rf = RequiredFeatures("dnam", ("cg1", "cg2")) + assert rf["layer"] == "dnam" + assert rf["features"] == ["cg1", "cg2"] + assert rf["metadata"] == [] + assert set(rf.keys()) == {"layer", "features", "metadata"} + assert dict(rf) == { + "layer": "dnam", + "features": ["cg1", "cg2"], + "metadata": [], + } + + +def test_metadata_key_always_present_when_empty(): + rf = RequiredFeatures("clinical", ("albumin",)) + assert "metadata" in rf + assert rf["metadata"] == [] + + +def test_missing_features_error_is_value_error(): + assert issubclass(MissingFeaturesError, ValueError) + + +def test_valid_layers_contains_expected(): + assert "dnam" in VALID_LAYERS + assert "clinical" in VALID_LAYERS From eab15c5c4d64616bd0fd8ea06dd68328488bf3be Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 18:14:49 -0700 Subject: [PATCH 2/6] feat: add layer-aware required-feature validator --- biolearn/features.py | 81 ++++++++++++++++++++++++++++++++++ biolearn/test/test_features.py | 67 ++++++++++++++++++++++++++++ 2 files changed, 148 insertions(+) diff --git a/biolearn/features.py b/biolearn/features.py index 1037c02..a66296f 100644 --- a/biolearn/features.py +++ b/biolearn/features.py @@ -53,3 +53,84 @@ class MissingFeaturesError(ValueError): Subclasses ``ValueError`` so existing ``except ValueError`` handlers keep working while callers can catch this specific case. """ + + +# layer -> (GeoData attribute, axis carrying feature names). The ONLY place +# that knows the index-vs-columns asymmetry: dnam/rna/protein are +# features-as-rows (features in .index); clinical is samples-as-rows +# (biomarkers in .columns), matching metadata. +_LAYER_FRAME = { + "dnam": ("dnam", "index"), + "rna": ("rna", "index"), + "protein_olink": ("protein_olink", "index"), + "protein_alamar": ("protein_alamar", "index"), + "clinical": ("clinical", "columns"), +} + +_LAYER_NOUN = { + "dnam": "CpG sites", + "rna": "genes", + "protein_olink": "proteins", + "protein_alamar": "proteins", + "clinical": "clinical markers", +} + +_MISSING_PREVIEW_LIMIT = 5 + + +def _available_features(geo_data, layer): + attr, axis = _LAYER_FRAME[layer] + frame = getattr(geo_data, attr, None) + if frame is None: + return set() + return set(getattr(frame, axis)) + + +def validate_required_features(model, geo_data): + """Raise MissingFeaturesError if geo_data lacks what the model requires.""" + req = model.required_features() + + available = _available_features(geo_data, req.layer) + missing_features = [f for f in req.features if f not in available] + + metadata = getattr(geo_data, "metadata", None) + metadata_cols = set(metadata.columns) if metadata is not None else set() + missing_metadata = [m for m in req.metadata if m not in metadata_cols] + + if missing_features or missing_metadata: + raise MissingFeaturesError( + _format_missing(model, req, missing_features, missing_metadata) + ) + + +def _format_missing(model, req, missing_features, missing_metadata): + details = getattr(model, "details", None) or {} + name = details.get("name") if isinstance(details, dict) else None + label = f" for model '{name}'" if name else "" + messages = [] + + if missing_features: + missing = sorted(missing_features) + noun = _LAYER_NOUN.get(req.layer, "features") + preview = ", ".join(missing[:_MISSING_PREVIEW_LIMIT]) + remaining = len(missing) - _MISSING_PREVIEW_LIMIT + if remaining > 0: + preview = ( + f"showing first {_MISSING_PREVIEW_LIMIT}: {preview} " + f"(+{remaining} more)" + ) + messages.append( + f"Missing required {noun}{label} " + f"({len(missing)}/{len(req.features)}): {preview}. " + f"Provide {req.layer} data with these {noun} or use an " + f"imputation method that includes them." + ) + + if missing_metadata: + cols = ", ".join(sorted(missing_metadata)) + messages.append( + f"Missing required metadata{label}: {cols}. " + f"Provide these columns in geo_data.metadata." + ) + + return " ".join(messages) diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py index b37e581..381741b 100644 --- a/biolearn/test/test_features.py +++ b/biolearn/test/test_features.py @@ -40,3 +40,70 @@ def test_missing_features_error_is_value_error(): def test_valid_layers_contains_expected(): assert "dnam" in VALID_LAYERS assert "clinical" in VALID_LAYERS + + +import types + +import pandas as pd + +from biolearn.features import validate_required_features + + +class _StubModel: + def __init__(self, req, name=None): + self._req = req + self.details = {"name": name} if name else {} + + def required_features(self): + return self._req + + +def _geo(dnam=None, rna=None, clinical=None, metadata=None): + return types.SimpleNamespace( + dnam=dnam, + rna=rna, + protein_olink=None, + protein_alamar=None, + clinical=clinical, + metadata=metadata, + ) + + +def test_validate_passes_when_dnam_present(): + dnam = pd.DataFrame({"S1": [0.1, 0.2]}, index=["cg1", "cg2"]) + model = _StubModel(RequiredFeatures("dnam", ("cg1", "cg2"))) + assert validate_required_features(model, _geo(dnam=dnam)) is None + + +def test_validate_dnam_missing_raises_with_cpg_message(): + dnam = pd.DataFrame({"S1": [0.1]}, index=["cg1"]) + model = _StubModel(RequiredFeatures("dnam", ("cg1", "cg2")), name="Clock") + with pytest.raises( + MissingFeaturesError, match=r"Missing required CpG sites.*cg2" + ): + validate_required_features(model, _geo(dnam=dnam)) + + +def test_validate_clinical_uses_columns_axis(): + # clinical is samples-as-rows, biomarkers-as-columns + clinical = pd.DataFrame({"albumin": [4.0], "glucose": [5.0]}, index=["S1"]) + model = _StubModel(RequiredFeatures("clinical", ("albumin", "crp"))) + with pytest.raises(MissingFeaturesError, match=r"crp"): + validate_required_features(model, _geo(clinical=clinical)) + + +def test_validate_metadata_missing_raises(): + dnam = pd.DataFrame({"S1": [0.1]}, index=["cg1"]) + meta = pd.DataFrame({"sex": ["m"]}, index=["S1"]) + model = _StubModel(RequiredFeatures("dnam", ("cg1",), ("age",))) + with pytest.raises( + MissingFeaturesError, match=r"Missing required metadata.*age" + ): + validate_required_features(model, _geo(dnam=dnam, metadata=meta)) + + +def test_validate_missing_layer_frame_is_empty_set(): + # layer present but geo_data has None for it -> all features missing + model = _StubModel(RequiredFeatures("dnam", ("cg1",))) + with pytest.raises(MissingFeaturesError): + validate_required_features(model, _geo(dnam=None)) From d6c55f2c336066ce8e0516baa6a727b3a6b1fff4 Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 19:25:07 -0700 Subject: [PATCH 3/6] feat: route LinearMethylationModel through required-feature validator --- biolearn/model.py | 47 ++++++++++++++----------------- biolearn/test/test_features.py | 51 ++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 26 deletions(-) diff --git a/biolearn/model.py b/biolearn/model.py index e4edf3a..891034b 100644 --- a/biolearn/model.py +++ b/biolearn/model.py @@ -16,6 +16,7 @@ from biolearn.data_library import GeoData from biolearn.dunedin_pace import dunedin_pace_normalization from biolearn.util import get_data_file +from biolearn.features import RequiredFeatures, validate_required_features def anti_trafo(x, adult_age=20): @@ -1447,6 +1448,8 @@ def methylation_sites(self): class LinearModel: + _LAYER = "dnam" + def __init__( self, coefficient_file_or_df, @@ -1490,9 +1493,9 @@ def no_transform(_): ) def predict(self, geo_data): + self._validate_inputs(geo_data) matrix_data = self._get_data_matrix(geo_data) matrix_data = self.preprocess(matrix_data) - self._validate_required_features(matrix_data) matrix_data.loc["intercept"] = 1 # Join the coefficients and dnam_data on the index @@ -1514,7 +1517,13 @@ def predict(self, geo_data): # Return as a DataFrame return result.apply(self.transform).to_frame(name="Predicted") - def _validate_required_features(self, matrix_data): + def required_features(self): + features = tuple( + index for index in self.coefficients.index if index != "intercept" + ) + return RequiredFeatures(self._LAYER, features) + + def _validate_inputs(self, geo_data): return def _get_data_matrix(self, geo_data): @@ -1522,33 +1531,11 @@ def _get_data_matrix(self, geo_data): class LinearMethylationModel(LinearModel): - _MISSING_CPG_PREVIEW_LIMIT = 5 - def _get_data_matrix(self, geo_data): return geo_data.dnam - def _validate_required_features(self, matrix_data): - required_cpgs = self.methylation_sites() - missing_cpgs = sorted(set(required_cpgs) - set(matrix_data.index)) - if not missing_cpgs: - return - - model_name = self.details.get("name") - model_label = f" for model '{model_name}'" if model_name else "" - preview_limit = self._MISSING_CPG_PREVIEW_LIMIT - preview = ", ".join(missing_cpgs[:preview_limit]) - remaining = len(missing_cpgs) - preview_limit - if remaining > 0: - preview = ( - f"showing first {preview_limit}: {preview} (+{remaining} more)" - ) - - raise ValueError( - "Missing required CpG sites" - f"{model_label} ({len(missing_cpgs)}/{len(required_cpgs)}): " - f"{preview}. " - "Provide methylation data with these CpGs or use an imputation method that includes them." - ) + def _validate_inputs(self, geo_data): + validate_required_features(self, geo_data) def methylation_sites(self): unique_vars = set(self.coefficients.index) - {"intercept"} @@ -1657,8 +1644,16 @@ def methylation_sites(self): ).iloc[:, 0] return list(self.center_.index) + def required_features(self): + # Coefficients are principal components, not CpGs; describe the CpG + # inputs instead. PC clocks tolerate missing CpGs (intersect and + # center-fill), so this is descriptive only, never hard-validated. + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class LinearTranscriptomicModel(LinearModel): + _LAYER = "rna" + def _get_data_matrix(self, geo_data): return geo_data.rna diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py index 381741b..3a82e11 100644 --- a/biolearn/test/test_features.py +++ b/biolearn/test/test_features.py @@ -107,3 +107,54 @@ def test_validate_missing_layer_frame_is_empty_set(): model = _StubModel(RequiredFeatures("dnam", ("cg1",))) with pytest.raises(MissingFeaturesError): validate_required_features(model, _geo(dnam=None)) + + +from biolearn.model import ( + LinearMethylationModel, + LinearTranscriptomicModel, +) + + +def _linear_methylation_model(): + coeffs = pd.DataFrame( + {"CoefficientTraining": [0.5, 0.5, 1.0]}, + index=["cg1", "cg2", "intercept"], + ) + return LinearMethylationModel(coeffs, transform=lambda x: x, name="TestLM") + + +def test_linear_methylation_required_features(): + model = _linear_methylation_model() + rf = model.required_features() + assert rf.layer == "dnam" + assert set(rf.features) == {"cg1", "cg2"} + assert rf.metadata == () + assert "intercept" not in rf.features + + +def test_linear_transcriptomic_layer_is_rna(): + coeffs = pd.DataFrame( + {"CoefficientTraining": [0.5, 1.0]}, index=["GENE1", "intercept"] + ) + model = LinearTranscriptomicModel(coeffs, transform=lambda x: x) + rf = model.required_features() + assert rf.layer == "rna" + assert set(rf.features) == {"GENE1"} + + +def test_linear_methylation_predict_raises_on_missing_cpg(): + model = _linear_methylation_model() + dnam = pd.DataFrame({"S1": [0.1]}, index=["cg1"]) + geo = _geo(dnam=dnam) + with pytest.raises( + MissingFeaturesError, match=r"Missing required CpG sites.*cg2" + ): + model.predict(geo) + + +def test_linear_methylation_missing_cpg_is_value_error(): + # Backward compatibility: existing `except ValueError` still catches it. + model = _linear_methylation_model() + dnam = pd.DataFrame({"S1": [0.1]}, index=["cg1"]) + with pytest.raises(ValueError): + model.predict(_geo(dnam=dnam)) From 94d78f4167af7a8fe4f2be805b1e45e4861c5ad7 Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 19:37:26 -0700 Subject: [PATCH 4/6] feat: add descriptive required_features to tolerant clocks --- biolearn/model.py | 21 +++++++++++++++++++++ biolearn/test/test_features.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/biolearn/model.py b/biolearn/model.py index 891034b..0d605cb 100644 --- a/biolearn/model.py +++ b/biolearn/model.py @@ -1293,6 +1293,9 @@ def from_definition(cls, clock_definition): def methylation_sites(self): return list(self.reference) + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + def quantile_normalize(df): rank_mean = ( @@ -1446,6 +1449,9 @@ def solve_qp(meth_vector, deconv_reference): def methylation_sites(self): return list(self.reference.index) + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class LinearModel: _LAYER = "dnam" @@ -1929,6 +1935,9 @@ def predict(self, geo_data): def methylation_sites(self): return list(self.coefficients.index) + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class EpiTOC2Model: def __init__(self, reference_file): @@ -1968,6 +1977,9 @@ def predict(self, geo_data): def methylation_sites(self): return list(self.CpG_names) + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class HurdleAPIModel: """ @@ -2231,6 +2243,9 @@ def methylation_sites(self): """Return list of required CpG sites for imputation compatibility.""" return self.required_cpgs if self.required_cpgs else [] + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class GPAgeModel: """Gaussian Process regression model for age prediction (GP-age clock).""" @@ -2281,6 +2296,9 @@ def predict(self, geo_data): def methylation_sites(self): return self._sites + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + class ImputationDecorator: def __init__(self, clock, imputation_method): @@ -2433,6 +2451,9 @@ def methylation_sites(self): """Return list of required CpG sites""" return self.cpg_sites + def required_features(self): + return RequiredFeatures("dnam", tuple(self.methylation_sites())) + def single_sample_clock(clock_function, data): return clock_function(data).iloc[0, 0] diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py index 3a82e11..fbd443d 100644 --- a/biolearn/test/test_features.py +++ b/biolearn/test/test_features.py @@ -158,3 +158,33 @@ def test_linear_methylation_missing_cpg_is_value_error(): dnam = pd.DataFrame({"S1": [0.1]}, index=["cg1"]) with pytest.raises(ValueError): model.predict(_geo(dnam=dnam)) + + +from biolearn.model_gallery import ModelGallery +from biolearn.data_library import GeoData +from biolearn.util import get_test_data_file + + +def test_tolerant_clock_runs_when_required_cpg_missing(): + # EpiTOC2 is tolerant by design (uses the CpG intersection). Dropping a + # required CpG must NOT raise; required_features() is descriptive only. + data = GeoData.load_csv( + get_test_data_file("testset/"), "testset", validate=False + ) + model = ModelGallery().get("EpiTOC2", imputation_method="none") + present_required = [ + cpg for cpg in model.methylation_sites() if cpg in data.dnam.index + ] + assert present_required, "test fixture lacks EpiTOC2 CpGs" + + partial = data.copy() + partial.dnam = data.dnam.drop(index=present_required[:1]) + result = model.predict(partial) # must not raise + assert result is not None + + +def test_tolerant_clock_reports_dnam_layer(): + model = ModelGallery().get("EpiTOC2", imputation_method="none") + rf = model.required_features() + assert rf.layer == "dnam" + assert len(rf.features) > 0 From 76e9bdbee7d4d6338d9ead1f4c3eaea5fc76e1fe Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 19:49:25 -0700 Subject: [PATCH 5/6] feat: add required_features to grimage and proteomic clocks --- biolearn/model.py | 15 +++++++++++++++ biolearn/test/test_features.py | 16 ++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/biolearn/model.py b/biolearn/model.py index 0d605cb..6ea14c4 100644 --- a/biolearn/model.py +++ b/biolearn/model.py @@ -1804,6 +1804,13 @@ def methylation_sites(self): unique_vars = set(filtered_df["var"]) - {"Intercept", "Age", "Female"} return list(unique_vars) + def required_features(self): + # Grimage keeps its own runtime age/sex checks; these are declared + # for introspection only. + return RequiredFeatures( + "dnam", tuple(self.methylation_sites()), ("age", "sex") + ) + class LinearMultipartProteomicModel: def __init__( @@ -1872,6 +1879,14 @@ def predict(self, geo_data): def methylation_sites(self): return [] + def required_features(self): + proteins = tuple( + protein + for protein in self.coefficients["Protein"].unique() + if str(protein).lower() != "intercept" + ) + return RequiredFeatures("protein_olink", proteins) + class SexEstimationModel: def __init__(self, coeffecient_file, **details): diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py index fbd443d..f3195a0 100644 --- a/biolearn/test/test_features.py +++ b/biolearn/test/test_features.py @@ -188,3 +188,19 @@ def test_tolerant_clock_reports_dnam_layer(): rf = model.required_features() assert rf.layer == "dnam" assert len(rf.features) > 0 + + +def test_grimage_declares_age_and_sex_metadata(): + model = ModelGallery().get("GrimAgeV1", imputation_method="none") + rf = model.required_features() + assert rf.layer == "dnam" + assert set(rf.metadata) == {"age", "sex"} + + +def test_proteomic_clock_reports_protein_layer(): + model = ModelGallery().get( + "OrganAgeChronological", imputation_method="none" + ) + rf = model.required_features() + assert rf.layer == "protein_olink" + assert "intercept" not in [str(f).lower() for f in rf.features] From 4bb450ccf8f4844fa8e9927952839c10159f1e41 Mon Sep 17 00:00:00 2001 From: Marc Balestreri Date: Sat, 11 Jul 2026 20:21:19 -0700 Subject: [PATCH 6/6] test: smoke-test required_features across the gallery --- biolearn/features.py | 1 - biolearn/test/test_features.py | 22 ++++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/biolearn/features.py b/biolearn/features.py index a66296f..798fd3f 100644 --- a/biolearn/features.py +++ b/biolearn/features.py @@ -9,7 +9,6 @@ from collections.abc import Mapping from dataclasses import dataclass - VALID_LAYERS = ( "dnam", "rna", diff --git a/biolearn/test/test_features.py b/biolearn/test/test_features.py index f3195a0..0c3c4aa 100644 --- a/biolearn/test/test_features.py +++ b/biolearn/test/test_features.py @@ -204,3 +204,25 @@ def test_proteomic_clock_reports_protein_layer(): rf = model.required_features() assert rf.layer == "protein_olink" assert "intercept" not in [str(f).lower() for f in rf.features] + + +def test_every_gallery_model_exposes_required_features(): + gallery = ModelGallery() + for name, model_def in gallery.model_definitions.items(): + # HurdleAPIModel needs API credentials just to instantiate; skip it. + if model_def["model"]["type"] == "HurdleAPIModel": + continue + model = gallery.get(name) # default imputation may wrap in decorator + rf = model.required_features() + assert rf.layer in VALID_LAYERS, name + assert set(rf.keys()) == {"layer", "features", "metadata"}, name + assert isinstance(rf["features"], list), name + assert isinstance(rf["metadata"], list), name + + +def test_required_features_forwards_through_imputation_decorator(): + model = ModelGallery().get("Horvathv1", imputation_method="averaging") + # ImputationDecorator wraps the clock; __getattr__ forwards the call. + rf = model.required_features() + assert rf.layer == "dnam" + assert len(rf.features) > 0