refactor: consolidate shared components and reorganize schema modules

  - Move common functions (risk categorization, calibration, gompertz model) to helpers.py
  - Consolidate schema modules into algorithm-specific files (phenoage.py, score2.py)
  - Rename coefficient classes for clarity and consistency
  - Update imports across all compute modules to use new schema structure
  - Remove redundant schema files (core.py, markers.py, units.py)

Author: fbraza

tests/test_phenoage.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from vitals.phenoage import compute
+from vitals.models import phenoage
 
 OUT_FILEPATH = Path(__file__).parent / "inputs" / "phenoage"
 
@@ -25,7 +25,7 @@
 )
 def test_phenoage(filename, expected):
     # Get the actual fixture value using request.getfixturevalue
-    age, pred_age, accl_age = compute.biological_age(OUT_FILEPATH / filename)
+    age, pred_age, accl_age = phenoage.compute(OUT_FILEPATH / filename)
     expected_age, expected_pred_age, expected_accl_age = expected
 
     assert age == expected_age

tests/test_score2.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from vitals.score2 import compute
+from vitals.models import score2
 
 OUT_FILEPATH = Path(__file__).parent / "inputs" / "score2"
 
@@ -26,9 +26,7 @@
 )
 def test_score2(filename, expected):
     # Get the actual fixture value using request.getfixturevalue
-    age, pred_risk, pred_risk_category = compute.cardiovascular_risk(
-        OUT_FILEPATH / filename
-    )
+    age, pred_risk, pred_risk_category = score2.compute(OUT_FILEPATH / filename)
     expected_age, expected_risk, expected_category = expected
 
     assert age == expected_age

tests/test_score2_diabetes.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from vitals.score2_diabetes import compute
+from vitals.models import score2_diabetes
 
 OUT_FILEPATH = Path(__file__).parent / "inputs" / "score2_diabetes"
 
@@ -26,7 +26,7 @@ def test_score2_diabetes(filename, expected):
     They need to be calculated using MDCalc and updated before running tests.
     """
     # Get the actual fixture value
-    age, pred_risk, pred_risk_category = compute.cardiovascular_risk(
+    age, pred_risk, pred_risk_category = score2_diabetes.compute(
         OUT_FILEPATH / filename
     )
     expected_age, expected_risk, expected_category = expected

vitals/biomarkers/helpers.py

@@ -1,13 +1,15 @@
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, TypedDict, TypeVar
+from typing import Any, Literal, TypeAlias, TypedDict, TypeVar
 
+import numpy as np
 from pydantic import BaseModel
 
-from vitals.schemas.units import PhenoageUnits, Score2DiabetesUnits, Score2Units
+from vitals.schemas import phenoage, score2
 
+RiskCategory: TypeAlias = Literal["Low to moderate", "High", "Very high"]
 Biomarkers = TypeVar("Biomarkers", bound=BaseModel)
-Units = PhenoageUnits | Score2Units | Score2DiabetesUnits
+Units = phenoage.Units | score2.Units | score2.UnitsDiabetes
 
 
 class ConversionInfo(TypedDict):
@@ -198,3 +200,57 @@ def extract_biomarkers_from_json(
         extracted_values[field_name] = value
 
     return biomarker_class(**extracted_values)
+
+
+def determine_risk_category(age: float, calibrated_risk: float) -> RiskCategory:
+    """
+    Determine cardiovascular risk category based on age and calibrated risk percentage.
+
+    Args:
+        age: Patient's age in years
+        calibrated_risk: Calibrated 10-year CVD risk as a percentage
+
+    Returns:
+        Risk stratification category
+    """
+    if age < 50:
+        if calibrated_risk < 2.5:
+            return "Low to moderate"
+        elif calibrated_risk < 7.5:
+            return "High"
+        else:
+            return "Very high"
+    else:  # age 50-69
+        if calibrated_risk < 5:
+            return "Low to moderate"
+        elif calibrated_risk < 10:
+            return "High"
+        else:
+            return "Very high"
+
+
+def apply_calibration(uncalibrated_risk: float, scale1: float, scale2: float) -> float:
+    """
+    Apply regional calibration to uncalibrated risk estimate.
+
+    Args:
+        uncalibrated_risk: Raw risk estimate from the Cox model
+        scale1: First calibration scale parameter
+        scale2: Second calibration scale parameter
+
+    Returns:
+        Calibrated 10-year CVD risk as a percentage
+    """
+    return float(
+        (1 - np.exp(-np.exp(scale1 + scale2 * np.log(-np.log(1 - uncalibrated_risk)))))
+        * 100
+    )
+
+
+def gompertz_mortality_model(weighted_risk_score: float) -> float:
+    params = phenoage.Gompertz()
+    return 1 - np.exp(
+        -np.exp(weighted_risk_score)
+        * (np.exp(120 * params.lambda_) - 1)
+        / params.lambda_
+    )

vitals/phenoage/__init__.py vitals/models/__init__.pyvitals/phenoage/__init__.py renamed to vitals/models/__init__.py

@@ -1,52 +1,12 @@
 from pathlib import Path
 
 import numpy as np
-from pydantic import BaseModel
 
 from vitals.biomarkers import helpers
-from vitals.schemas.markers import PhenoageMarkers
-from vitals.schemas.units import PhenoageUnits
+from vitals.schemas.phenoage import Gompertz, LinearModel, Markers, Units
 
 
-class LinearModel(BaseModel):
-    """
-    Coefficients used to calculate the PhenoAge from Levine et al 2018
-    """
-
-    intercept: float = -19.9067
-    albumin: float = -0.0336
-    creatinine: float = 0.0095
-    glucose: float = 0.1953
-    log_crp: float = 0.0954
-    lymphocyte_percent: float = -0.0120
-    mean_cell_volume: float = 0.0268
-    red_cell_distribution_width: float = 0.3306
-    alkaline_phosphatase: float = 0.00188
-    white_blood_cell_count: float = 0.0554
-    age: float = 0.0804
-
-
-class Gompertz(BaseModel):
-    """
-    Parameters of the Gompertz distribution for PhenoAge computation
-    """
-
-    lambda_: float = 0.0192
-    coef1: float = 141.50225
-    coef2: float = -0.00553
-    coef3: float = 0.090165
-
-
-def __gompertz_mortality_model(weighted_risk_score: float) -> float:
-    __params = Gompertz()
-    return 1 - np.exp(
-        -np.exp(weighted_risk_score)
-        * (np.exp(120 * __params.lambda_) - 1)
-        / __params.lambda_
-    )
-
-
-def biological_age(filepath: str | Path) -> tuple[float, float, float]:
+def compute(filepath: str | Path) -> tuple[float, float, float]:
     """
     The Phenoage score is calculated as a weighted (coefficients available in Levine et al 2018)
     linear combination of these variables, which was then transformed into units of years using 2 parametric
@@ -57,14 +17,14 @@ def biological_age(filepath: str | Path) -> tuple[float, float, float]:
     # Extract biomarkers from JSON file
     biomarkers = helpers.extract_biomarkers_from_json(
         filepath=filepath,
-        biomarker_class=PhenoageMarkers,
-        biomarker_units=PhenoageUnits(),
+        biomarker_class=Markers,
+        biomarker_units=Units(),
     )
 
     age = biomarkers.age
     coef = LinearModel()
 
-    if isinstance(biomarkers, PhenoageMarkers):
+    if isinstance(biomarkers, Markers):
         weighted_risk_score = (
             coef.intercept
             + (coef.albumin * biomarkers.albumin)
@@ -81,7 +41,9 @@ def biological_age(filepath: str | Path) -> tuple[float, float, float]:
             + (coef.white_blood_cell_count * biomarkers.white_blood_cell_count)
             + (coef.age * biomarkers.age)
         )
-        gompertz = __gompertz_mortality_model(weighted_risk_score=weighted_risk_score)
+        gompertz = helpers.gompertz_mortality_model(
+            weighted_risk_score=weighted_risk_score
+        )
         model = Gompertz()
         pred_age = (
             model.coef1 + np.log(model.coef2 * np.log(1 - gompertz)) / model.coef3
@@ -90,17 +52,3 @@ def biological_age(filepath: str | Path) -> tuple[float, float, float]:
         return (age, pred_age, accl_age)
     else:
         raise ValueError(f"Invalid biomarker class used: {biomarkers}")
-
-
-# if __name__ == "__main__":
-#     from pathlib import Path
-#     input_dir = Path("tests/outputs")
-#     output_dir = Path("tests/outputs")
-
-#     for input_file in input_dir.glob("*.json"):
-#         if "patient" not in str(input_file):
-#             continue
-
-#         # Update biomarker data
-#         age, pred_age, accl_age = biological_age(str(input_file))
-#         print(f"Chrono Age: {age} ::: Predicted Age: {pred_age} ::: Accel {accl_age}")

vitals/phenoage/compute.py vitals/models/phenoage.pyvitals/phenoage/compute.py renamed to vitals/models/phenoage.py

@@ -10,19 +10,19 @@
 import numpy as np
 
 from vitals.biomarkers import helpers
-from vitals.schemas.coefficients import Score2FemaleCoefficients, Score2MaleCoefficients
-from vitals.schemas.core import (
+from vitals.schemas.score2 import (
     BaselineSurvival,
     CalibrationScales,
-    RiskCategory,
-    apply_calibration,
-    determine_risk_category,
+    FemaleCoefficientsBaseModel,
+    MaleCoefficientsBaseModel,
+    Markers,
+    Units,
 )
-from vitals.schemas.markers import Score2Markers
-from vitals.schemas.units import Score2Units
 
 
-def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategory]:
+def compute(
+    filepath: str | Path,
+) -> tuple[float, float, helpers.RiskCategory]:
     """
     Calculate the 10-year cardiovascular disease risk using the SCORE2 algorithm.
 
@@ -46,11 +46,11 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     # Extract biomarkers from JSON file
     biomarkers = helpers.extract_biomarkers_from_json(
         filepath=filepath,
-        biomarker_class=Score2Markers,
-        biomarker_units=Score2Units(),
+        biomarker_class=Markers,
+        biomarker_units=Units(),
     )
 
-    if not isinstance(biomarkers, Score2Markers):
+    if not isinstance(biomarkers, Markers):
         raise ValueError(f"Invalid biomarker class used: {biomarkers}")
 
     age: float = biomarkers.age
@@ -73,14 +73,14 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     baseline_survival_model = BaselineSurvival()
     calibration_scales = CalibrationScales()
 
-    coef: Score2MaleCoefficients | Score2FemaleCoefficients
+    coef: MaleCoefficientsBaseModel | FemaleCoefficientsBaseModel
     if is_male:
-        coef = Score2MaleCoefficients()
+        coef = MaleCoefficientsBaseModel()
         baseline_survival = baseline_survival_model.male
         scale1 = calibration_scales.male_scale1
         scale2 = calibration_scales.male_scale2
     else:
-        coef = Score2FemaleCoefficients()
+        coef = FemaleCoefficientsBaseModel()
         baseline_survival = baseline_survival_model.female
         scale1 = calibration_scales.female_scale1
         scale2 = calibration_scales.female_scale2
@@ -101,9 +101,13 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     uncalibrated_risk: float = 1 - np.power(baseline_survival, np.exp(linear_pred))
 
     # Apply calibration for Belgium (Low Risk region)
-    calibrated_risk: float = apply_calibration(uncalibrated_risk, scale1, scale2)
+    calibrated_risk: float = helpers.apply_calibration(
+        uncalibrated_risk, scale1, scale2
+    )
 
     # Determine risk category based on age
-    risk_category: RiskCategory = determine_risk_category(age, calibrated_risk)
+    risk_category: helpers.RiskCategory = helpers.determine_risk_category(
+        age, calibrated_risk
+    )
 
     return (age, round(calibrated_risk, 2), risk_category)

vitals/score2/compute.py vitals/models/score2.pyvitals/score2/compute.py renamed to vitals/models/score2.py

@@ -11,22 +11,19 @@
 import numpy as np
 
 from vitals.biomarkers import helpers
-from vitals.schemas.coefficients import (
-    Score2DiabetesFemaleCoefficients,
-    Score2DiabetesMaleCoefficients,
-)
-from vitals.schemas.core import (
+from vitals.schemas.score2 import (
     BaselineSurvival,
     CalibrationScales,
-    RiskCategory,
-    apply_calibration,
-    determine_risk_category,
+    FemaleCoefficientsDiabeticModel,
+    MaleCoefficientsDiabeticModel,
+    MarkersDiabetes,
+    UnitsDiabetes,
 )
-from vitals.schemas.markers import Score2DiabetesMarkers
-from vitals.schemas.units import Score2DiabetesUnits
 
 
-def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategory]:
+def compute(
+    filepath: str | Path,
+) -> tuple[float, float, helpers.RiskCategory]:
     """
     Calculate the 10-year cardiovascular disease risk using the SCORE2-Diabetes algorithm.
 
@@ -51,11 +48,11 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     # Extract biomarkers from JSON file
     biomarkers = helpers.extract_biomarkers_from_json(
         filepath=filepath,
-        biomarker_class=Score2DiabetesMarkers,
-        biomarker_units=Score2DiabetesUnits(),
+        biomarker_class=MarkersDiabetes,
+        biomarker_units=UnitsDiabetes(),
     )
 
-    if not isinstance(biomarkers, Score2DiabetesMarkers):
+    if not isinstance(biomarkers, MarkersDiabetes):
         raise ValueError(f"Invalid biomarker class used: {biomarkers}")
 
     age: float = biomarkers.age
@@ -86,14 +83,14 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     baseline_survival_model = BaselineSurvival()
     calibration_scales = CalibrationScales()
 
-    coef: Score2DiabetesMaleCoefficients | Score2DiabetesFemaleCoefficients
+    coef: MaleCoefficientsDiabeticModel | FemaleCoefficientsDiabeticModel
     if is_male:
-        coef = Score2DiabetesMaleCoefficients()
+        coef = MaleCoefficientsDiabeticModel()
         baseline_survival = baseline_survival_model.male
         scale1 = calibration_scales.male_scale1
         scale2 = calibration_scales.male_scale2
     else:
-        coef = Score2DiabetesFemaleCoefficients()
+        coef = FemaleCoefficientsDiabeticModel()
         baseline_survival = baseline_survival_model.female
         scale1 = calibration_scales.female_scale1
         scale2 = calibration_scales.female_scale2
@@ -122,9 +119,13 @@ def cardiovascular_risk(filepath: str | Path) -> tuple[float, float, RiskCategor
     uncalibrated_risk: float = 1 - np.power(baseline_survival, np.exp(linear_pred))
 
     # Apply calibration for Belgium (Low Risk region)
-    calibrated_risk: float = apply_calibration(uncalibrated_risk, scale1, scale2)
+    calibrated_risk: float = helpers.apply_calibration(
+        uncalibrated_risk, scale1, scale2
+    )
 
     # Determine risk category based on age
-    risk_category: RiskCategory = determine_risk_category(age, calibrated_risk)
+    risk_category: helpers.RiskCategory = helpers.determine_risk_category(
+        age, calibrated_risk
+    )
 
     return (age, round(calibrated_risk, 2), risk_category)