adding stats list to snpe deployment

JBris · Sep 18, 2024 · 91694e0 · 91694e0
1 parent 3ab80c9
commit 91694e0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 13 deletions.
diff --git a/app/flows/run_snpe.py b/app/flows/run_snpe.py
@@ -199,6 +199,7 @@ def log_task(
     posterior_samples: torch.Tensor,
     input_parameters: RootCalibrationModel,
     observed_values: list,
+    statistics_list: list[SummaryStatisticsModel],
     names: list[str],
     limits: list[tuple],
     simulation_uuid: str,
@@ -220,6 +221,8 @@ def log_task(
             The root calibration data model.
         observed_values (list):
             The list of observed_values.
+        statistics_list (list[SummaryStatisticsModel]):
+            The list of summary statistics.
         names (list[str]):
             The parameter names.
         limits (list[tuple]):
@@ -327,12 +330,14 @@ def log_task(
         parameter_specs, input_parameters
     )
 
+    statistics_list = [statistic.dict() for statistic in statistics_list]
     parameter_intervals["inference_type"] = "summary_statistics"
     artifacts = {}
     for obj, name in [
         (inference, "inference"),
         (posterior, "posterior"),
         (parameter_intervals, "parameter_intervals"),
+        (statistics_list, "statistics_list"),
     ]:
         outfile = osp.join(outdir, f"{time_now}-{TASK}_{name}.pkl")
         artifacts[name] = outfile
@@ -380,6 +385,7 @@ def run_snpe(input_parameters: RootCalibrationModel, simulation_uuid: str) -> No
         posterior_samples,
         input_parameters,
         observed_values,
+        statistics_list,
         names,
         limits,
         simulation_uuid,

diff --git a/deeprootgen/calibration/model_versioning.py b/deeprootgen/calibration/model_versioning.py
@@ -3,6 +3,8 @@
 This module defines MLflow compatible models for versioning and deployment as microservices.
 """
 
+from typing import Any
+
 import bentoml
 import mlflow
 import numpy as np
@@ -253,14 +255,14 @@ def load_context(self, context: Context) -> None:
         """
         import joblib
 
-        loaded_data = context.artifacts["inference"]
-        self.inference = joblib.load(loaded_data)
-
-        loaded_data = context.artifacts["posterior"]
-        self.posterior = joblib.load(loaded_data)
+        def load_data(k: str) -> Any:
+            artifact = context.artifacts[k]
+            return joblib.load(artifact)
 
-        loaded_data = context.artifacts["parameter_intervals"]
-        self.parameter_intervals = joblib.load(loaded_data)
+        self.inference = load_data("inference")
+        self.posterior = load_data("posterior")
+        self.parameter_intervals = load_data("parameter_intervals")
+        self.statistics_list = load_data("statistics_list")
 
     def predict(
         self, context: Context, model_input: pd.DataFrame, params: dict | None = None
@@ -283,12 +285,14 @@ def predict(
             pd.DataFrame:
                 The model prediction.
         """
-        if (
-            self.inference is None
-            or self.posterior is None
-            or self.parameter_intervals is None
-        ):
-            raise ValueError(f"The {self.task} calibrator has not been loaded.")
+        for prop in [
+            self.inference,
+            self.posterior,
+            self.parameter_intervals,
+            self.statistics_list,
+        ]:
+            if prop is None:
+                raise ValueError(f"The {self.task} calibrator has not been loaded.")
 
         observed_values = model_input["statistic_value"].values
         posterior_samples = self.posterior.sample((50,), x=observed_values)