Add performance metric to summary file

pyproject.toml

@@ -12,7 +12,15 @@ license = { file = "LICENSE" }
 authors = [
     { name = "IHME Math Sciences", email = "[email protected]" },
 ]
-dependencies = ["numpy", "scipy", "pandas", "matplotlib", "mrtool==0.1.4", "pplkit"]
+dependencies = [
+    "numpy",
+    "scipy",
+    "pandas",
+    "matplotlib",
+    "limetr @ git+https://github.com/zhengp0/[email protected]",
+    "mrtool==0.1.4",
+    "pplkit",
+]
 
 [project.optional-dependencies]
 test = ["pytest"]

src/bopforge/continuous_pipeline/__main__.py

@@ -65,7 +65,7 @@ def fit_signal_model(dataif: DataInterface) -> None:
 
     df = functions.convert_bc_to_em(df, signal_model)
 
-    summary = functions.get_signal_model_summary(name, all_settings, df)
+    summary = functions.get_signal_model_summary(name, all_settings, df, signal_model)
 
     fig = functions.plot_signal_model(
         name,

src/bopforge/continuous_pipeline/functions.py

@@ -9,6 +9,7 @@
 from mrtool.core.utils import sample_knots
 from pandas import DataFrame
 from scipy.stats import norm
+from limetr import get_aic, get_bic, get_rmse
 
 
 def get_signal_model(settings: dict, df: DataFrame) -> MRBeRT:
@@ -138,7 +139,12 @@ def convert_bc_to_em(df: DataFrame, signal_model: MRBeRT) -> DataFrame:
     return df
 
 
-def get_signal_model_summary(name: str, all_settings: dict, df: DataFrame) -> dict:
+def get_signal_model_summary(
+    name: str,
+    all_settings: dict,
+    df: DataFrame,
+    signal_model: MRBeRT,
+) -> dict:
     """Create signal model summary.
 
     Parameters
@@ -149,6 +155,8 @@ def get_signal_model_summary(name: str, all_settings: dict, df: DataFrame) -> di
         All the settings for the pipeline.
     df
         Data frame that contains the training dataset.
+    signal_model
+        Fitted signal model for risk curve.
 
     Returns
     -------
@@ -183,6 +191,26 @@ def get_signal_model_summary(name: str, all_settings: dict, df: DataFrame) -> di
         "normalize_to_tmrel"
     ]
 
+    summary["model_performance"] = {
+        "signal_model": {
+            "aic": float(
+                np.array(
+                    [get_aic(sub_model.lt) for sub_model in signal_model.sub_models]
+                ).dot(signal_model.weights)
+            ),
+            "bic": float(
+                np.array(
+                    [get_bic(sub_model.lt) for sub_model in signal_model.sub_models]
+                ).dot(signal_model.weights)
+            ),
+            "rmse": float(
+                np.array(
+                    [get_rmse(sub_model.lt) for sub_model in signal_model.sub_models]
+                ).dot(signal_model.weights)
+            ),
+        },
+    }
+
     return summary
 
 
@@ -414,11 +442,9 @@ def get_linear_model_summary(
         summary["score"] = float("nan")
         summary["star_rating"] = 0
     else:
-        score = float(
-            ((sign * burden_of_proof)[:, index].mean(axis=1)).min()
-        )
+        score = float(((sign * burden_of_proof)[:, index].mean(axis=1)).min())
         summary["score"] = score
-        #Assign star rating based on ROS
+        # Assign star rating based on ROS
         if np.isnan(score):
             summary["star_rating"] = 0
         elif score > np.log(1 + 0.85):
@@ -445,6 +471,12 @@ def get_linear_model_summary(
     summary["pub_bias"] = int(pval < 0.05)
     summary["pub_bias_pval"] = float(pval)
 
+    summary["model_performance"]["linear_model"] = {
+        "aic": float(get_aic(linear_model.lt)),
+        "bic": float(get_bic(linear_model.lt)),
+        "rmse": float(get_rmse(linear_model.lt)),
+    }
+
     return summary
 
 

src/bopforge/dichotomous_pipeline/functions.py

@@ -6,6 +6,7 @@
 from mrtool import MRBRT, CovFinder, LinearCovModel, MRData
 from pandas import DataFrame
 from scipy.stats import norm
+from limetr import get_aic, get_bic, get_rmse
 
 
 def get_signal_model(settings: dict, df: DataFrame) -> MRBRT:
@@ -272,6 +273,14 @@ def get_linear_model_summary(
     summary["pub_bias"] = int(pval < 0.05)
     summary["pub_bias_pval"] = float(pval)
 
+    summary["model_performance"] = {
+        "linear_model": {
+            "aic": float(get_aic(linear_model.lt)),
+            "bic": float(get_bic(linear_model.lt)),
+            "rmse": float(get_rmse(linear_model.lt)),
+        }
+    }
+
     return summary