From 34e6584ad8aad9e8603c0a22dffc737612a55aa1 Mon Sep 17 00:00:00 2001
From: Tom Szendrey <T.Szendrey@hotmail.com>
Date: Tue, 30 Jul 2024 12:59:47 -0400
Subject: [PATCH] Ts/107 add solver_options param (#109)

* rename step size as we have two.
Add only to extra

* Adding solver step size for calibrate

* desc update

* check for solver method too

* always pop step size out of extras.

* extra_options

* extras suck. unrequired but dont double send.

* good call Dan
---
 service/models/operations/calibrate.py | 19 ++++++++++++++++++-
 service/models/operations/simulate.py  | 25 +++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/service/models/operations/calibrate.py b/service/models/operations/calibrate.py
index bb9ddd6..c49fe90 100644
--- a/service/models/operations/calibrate.py
+++ b/service/models/operations/calibrate.py
@@ -38,6 +38,12 @@ class CalibrateExtra(BaseModel):
     solver_method: str = Field(
         "dopri5", description="Optional field for CIEMSS calibration", example="dopri5"
     )
+    # https://github.com/ciemss/pyciemss/blob/main/pyciemss/integration_utils/interface_checks.py
+    solver_step_size: float = Field(
+        None,
+        description="Step size required if solver method is euler.",
+        example=1.0,
+    )
 
 
 class Calibrate(OperationRequest):
@@ -74,6 +80,15 @@ def hook(progress, _loss):
                     logging.info(f"Calibration is {progress}% complete")
                 return None
 
+        extra_options = self.extra.dict()
+        solver_options = {}
+        step_size = extra_options.pop(
+            "solver_step_size"
+        )  # Need to pop this out of extra.
+        solver_method = extra_options.pop("solver_method")
+        if step_size is not None and solver_method == "euler":
+            solver_options["step_size"] = step_size
+
         return {
             "model_path_or_json": amr_path,
             "start_time": self.timespan.start,
@@ -82,8 +97,10 @@ def hook(progress, _loss):
             "data_path": dataset_path,
             "static_parameter_interventions": static_interventions,
             "progress_hook": hook,
+            "solver_method": solver_method,
+            "solver_options": solver_options,
             # "visual_options": True,
-            **self.extra.dict(),
+            **extra_options,
         }
 
     class Config:
diff --git a/service/models/operations/simulate.py b/service/models/operations/simulate.py
index 2aeeebb..5e3ca60 100644
--- a/service/models/operations/simulate.py
+++ b/service/models/operations/simulate.py
@@ -21,6 +21,17 @@ class SimulateExtra(BaseModel):
         description="id from a previous calibration",
         example=None,
     )
+    solver_method: str = Field(
+        "dopri5",
+        description="Optional field for CIEMSS calibration",
+        example="dopri5",
+    )
+    # https://github.com/ciemss/pyciemss/blob/main/pyciemss/integration_utils/interface_checks.py
+    solver_step_size: float = Field(
+        None,
+        description="Step size required if solver method is euler.",
+        example=1.0,
+    )
 
 
 class Simulate(OperationRequest):
@@ -28,7 +39,7 @@ class Simulate(OperationRequest):
     model_config_id: str = Field(..., example="ba8da8d4-047d-11ee-be56")
     timespan: Timespan = Timespan(start=0, end=90)
     policy_intervention_id: str = Field(None, example="ba8da8d4-047d-11ee-be56")
-    step_size: float = 1.0
+    logging_step_size: float = 1.0
     extra: SimulateExtra = Field(
         None,
         description="optional extra system specific arguments for advanced use cases",
@@ -51,14 +62,24 @@ def gen_pyciemss_args(self, job_id):
             extra_options.pop("inferred_parameters"), job_id
         )
 
+        solver_options = {}
+        step_size = extra_options.pop(
+            "solver_step_size"
+        )  # Need to pop this out of extra.
+        solver_method = extra_options.pop("solver_method")
+        if step_size is not None and solver_method == "euler":
+            solver_options["step_size"] = step_size
+
         return {
             "model_path_or_json": amr_path,
-            "logging_step_size": self.step_size,
+            "logging_step_size": self.logging_step_size,
             "start_time": self.timespan.start,
             "end_time": self.timespan.end,
             "static_parameter_interventions": static_interventions,
             "dynamic_parameter_interventions": dynamic_interventions,
             "inferred_parameters": inferred_parameters,
+            "solver_method": solver_method,
+            "solver_options": solver_options,
             **extra_options,
         }