Refactor SCORE2 modules and extract shared components

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.biomarkers import schemas
+from vitals.schemas import phenoage, score2
 
+RiskCategory: TypeAlias = Literal["Low to moderate", "High", "Very high"]
 Biomarkers = TypeVar("Biomarkers", bound=BaseModel)
-Units = schemas.PhenoageUnits | schemas.Score2Units | schemas.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/compute.py → vitals/models/phenoage.py

@@ -1,50 +1,12 @@
 from pathlib import Path
 
 import numpy as np
-from pydantic import BaseModel
 
-from vitals.biomarkers import helpers, schemas
+from vitals.biomarkers import helpers
+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
@@ -55,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=schemas.PhenoageMarkers,
-        biomarker_units=schemas.PhenoageUnits(),
+        biomarker_class=Markers,
+        biomarker_units=Units(),
     )
 
     age = biomarkers.age
     coef = LinearModel()
 
-    if isinstance(biomarkers, schemas.PhenoageMarkers):
+    if isinstance(biomarkers, Markers):
         weighted_risk_score = (
             coef.intercept
             + (coef.albumin * biomarkers.albumin)
@@ -79,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
@@ -88,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}")