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

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/compute.py

@@ -1,49 +1,9 @@
 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
-
-
-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_
-    )
+from vitals.schemas.phenoage import Gompertz, LinearModel, Markers, Units
 
 
 def biological_age(filepath: str | Path) -> tuple[float, float, float]:
@@ -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/schemas/core.py

@@ -0,0 +1,62 @@
+from pydantic import BaseModel
+
+
+class Markers(BaseModel):
+    """Processed PhenoAge biomarkers with standardized units."""
+
+    albumin: float
+    creatinine: float
+    glucose: float
+    crp: float
+    lymphocyte_percent: float
+    mean_cell_volume: float
+    red_cell_distribution_width: float
+    alkaline_phosphatase: float
+    white_blood_cell_count: float
+    age: float
+
+
+class Units(BaseModel):
+    """
+    The expected unit to be used for phenoage computation
+    """
+
+    albumin: str = "g/L"
+    creatinine: str = "umol/L"
+    glucose: str = "mmol/L"
+    crp: str = "mg/dL"
+    lymphocyte_percent: str = "%"
+    mean_cell_volume: str = "fL"
+    red_cell_distribution_width: str = "%"
+    alkaline_phosphatase: str = "U/L"
+    white_blood_cell_count: str = "1000 cells/uL"
+    age: str = "years"
+
+
+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