diff --git a/biolearn/features.py b/biolearn/features.py new file mode 100644 index 0000000..798fd3f --- /dev/null +++ b/biolearn/features.py @@ -0,0 +1,135 @@ +"""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. + """ + + +# 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/model.py b/biolearn/model.py index e4edf3a..6ea14c4 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): @@ -1292,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 = ( @@ -1445,8 +1449,13 @@ 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" + def __init__( self, coefficient_file_or_df, @@ -1490,9 +1499,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 +1523,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 +1537,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 +1650,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 @@ -1803,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__( @@ -1871,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): @@ -1934,6 +1950,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): @@ -1973,6 +1992,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: """ @@ -2236,6 +2258,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).""" @@ -2286,6 +2311,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): @@ -2438,6 +2466,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 new file mode 100644 index 0000000..0c3c4aa --- /dev/null +++ b/biolearn/test/test_features.py @@ -0,0 +1,228 @@ +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 + + +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 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 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 + + +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] + + +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