feat: add required_features() and missing-data validation

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)

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,41 +1523,25 @@ 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):
         raise NotImplementedError()
 
 
 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]