610 atol and rtol cannot be specified in solver options (#612)

* added rtol, atol, to interface

* added tests for optimize and calibration

* passes linter

* atol and rtol moved to inside solver_options dict, updated tests

* formatting and linting passing

pyciemss/interfaces.py

@@ -93,7 +93,7 @@ def ensemble_sample(
         - If performance is incredibly slow, we suggest using `euler` to debug.
           If using `euler` results in faster simulation, the issue is likely that the model is stiff.
     solver_options: Dict[str, Any]
-        - Options to pass to the solver. See torchdiffeq' `odeint` method for more details.
+        - Options to pass to the solver (including atol and rtol). See torchdiffeq' `odeint` method for more details.
     start_time: float
         - The start time of the model. This is used to align the `start_state` from the
           AMR model with the simulation timepoints.
@@ -121,6 +121,10 @@ def ensemble_sample(
     """
     check_solver(solver_method, solver_options)
 
+    # Get tolerances for solver
+    rtol = solver_options.pop("rtol", 1e-7)  # default = 1e-7
+    atol = solver_options.pop("atol", 1e-9)  # default = 1e-9
+
     with torch.no_grad():
         if dirichlet_alpha is None:
             dirichlet_alpha = torch.ones(len(model_paths_or_jsons))
@@ -138,7 +142,9 @@ def ensemble_sample(
             raise ValueError("num_samples must be a positive integer")
 
         def wrapped_model():
-            with TorchDiffEq(method=solver_method, options=solver_options):
+            with TorchDiffEq(
+                rtol=rtol, atol=atol, method=solver_method, options=solver_options
+            ):
                 solution = model(
                     torch.as_tensor(start_time),
                     torch.as_tensor(end_time),
@@ -233,7 +239,7 @@ def ensemble_calibrate(
         - If performance is incredibly slow, we suggest using `euler` to debug.
             If using `euler` results in faster simulation, the issue is likely that the model is stiff.
     solver_options: Dict[str, Any]
-        - Options to pass to the solver. See torchdiffeq' `odeint` method for more details.
+        - Options to pass to the solver (including atol and rtol). See torchdiffeq' `odeint` method for more details.
     start_time: float
         - The start time of the model. This is used to align the `start_state` from the
             AMR model with the simulation timepoints.
@@ -280,6 +286,10 @@ def ensemble_calibrate(
     if not (isinstance(num_iterations, int) and num_iterations > 0):
         raise ValueError("num_iterations must be a positive integer")
 
+    # Get tolerances for solver
+    rtol = solver_options.pop("rtol", 1e-7)  # default = 1e-7
+    atol = solver_options.pop("atol", 1e-9)  # default = 1e-9
+
     def autoguide(model):
         guide = pyro.infer.autoguide.AutoGuideList(model)
         guide.append(
@@ -314,7 +324,9 @@ def autoguide(model):
     def wrapped_model():
         obs = condition(data=_data)(_noise_model)
 
-        with TorchDiffEq(method=solver_method, options=solver_options):
+        with TorchDiffEq(
+            rtol=rtol, atol=atol, method=solver_method, options=solver_options
+        ):
             solution = model(
                 torch.as_tensor(start_time),
                 torch.as_tensor(data_timepoints[-1]),
@@ -384,7 +396,8 @@ def sample(
             - If performance is incredibly slow, we suggest using `euler` to debug.
               If using `euler` results in faster simulation, the issue is likely that the model is stiff.
         solver_options: Dict[str, Any]
-            - Options to pass to the solver. See torchdiffeq' `odeint` method for more details.
+            - Options to pass to the solver (including atol and rtol).
+              See torchdiffeq' `odeint` method for more details.
         start_time: float
             - The start time of the model. This is used to align the `start_state` from the
               AMR model with the simulation timepoints.
@@ -449,6 +462,10 @@ def sample(
 
     check_solver(solver_method, solver_options)
 
+    # Get tolerances for solver
+    rtol = solver_options.pop("rtol", 1e-7)  # default = 1e-7
+    atol = solver_options.pop("atol", 1e-9)  # default = 1e-9
+
     with torch.no_grad():
         model = CompiledDynamics.load(model_path_or_json)
 
@@ -492,7 +509,9 @@ def sample(
 
         def wrapped_model():
             with ParameterInterventionTracer():
-                with TorchDiffEq(method=solver_method, options=solver_options):
+                with TorchDiffEq(
+                    rtol=rtol, atol=atol, method=solver_method, options=solver_options
+                ):
                     with contextlib.ExitStack() as stack:
                         for handler in intervention_handlers:
                             stack.enter_context(handler)
@@ -602,7 +621,8 @@ def calibrate(
             - If performance is incredibly slow, we suggest using `euler` to debug.
               If using `euler` results in faster simulation, the issue is likely that the model is stiff.
         - solver_options: Dict[str, Any]
-            - Options to pass to the solver. See torchdiffeq' `odeint` method for more details.
+            - Options to pass to the solver (including atol and rtol).
+              See torchdiffeq' `odeint` method for more details.
         - start_time: float
             - The start time of the model. This is used to align the `start_state` from the
               AMR model with the simulation timepoints.
@@ -668,6 +688,10 @@ def calibrate(
 
     check_solver(solver_method, solver_options)
 
+    # Get tolerances for solver
+    rtol = solver_options.pop("rtol", 1e-7)  # default = 1e-7
+    atol = solver_options.pop("atol", 1e-9)  # default = 1e-9
+
     pyro.clear_param_store()
 
     model = CompiledDynamics.load(model_path_or_json)
@@ -740,7 +764,9 @@ def wrapped_model():
         obs = condition(data=_data)(_noise_model)
 
         with StaticBatchObservation(data_timepoints, observation=obs):
-            with TorchDiffEq(method=solver_method, options=solver_options):
+            with TorchDiffEq(
+                rtol=rtol, atol=atol, method=solver_method, options=solver_options
+            ):
                 with contextlib.ExitStack() as stack:
                     for handler in intervention_handlers:
                         stack.enter_context(handler)
@@ -834,7 +860,8 @@ def optimize(
             - If performance is incredibly slow, we suggest using `euler` to debug.
               If using `euler` results in faster simulation, the issue is likely that the model is stiff.
         solver_options: Dict[str, Any]
-            - Options to pass to the solver. See torchdiffeq' `odeint` method for more details.
+            - Options to pass to the solver (including atol and rtol).
+              See torchdiffeq' `odeint` method for more details.
         start_time: float
             - The start time of the model. This is used to align the `start_state` from the
               AMR model with the simulation timepoints.

pyciemss/ouu/ouu.py

@@ -97,6 +97,8 @@ def __init__(
         self.u_bounds = u_bounds
         self.risk_bound = risk_bound  # used for defining penalty
         warnings.simplefilter("always", UserWarning)
+        self.rtol = self.solver_options.pop("rtol", 1e-7)  # default = 1e-7
+        self.atol = self.solver_options.pop("atol", 1e-9)  # default = 1e-9
 
     def __call__(self, x):
         if np.any(x - self.u_bounds[0, :] < 0.0) or np.any(
@@ -144,7 +146,10 @@ def propagate_uncertainty(self, x):
                 def wrapped_model():
                     with ParameterInterventionTracer():
                         with TorchDiffEq(
-                            method=self.solver_method, options=self.solver_options
+                            rtol=self.rtol,
+                            atol=self.atol,
+                            method=self.solver_method,
+                            options=self.solver_options,
                         ):
                             with contextlib.ExitStack() as stack:
                                 for handler in static_parameter_intervention_handlers:

tests/test_interfaces.py

@@ -87,29 +87,56 @@ def setup_calibrate(model_fixture, start_time, end_time, logging_step_size):
     "num_iterations": 2,
 }
 
+RTOL = [1e-6, 1e-4]
+ATOL = [1e-8, 1e-6]
+
 
 @pytest.mark.parametrize("sample_method", SAMPLE_METHODS)
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("start_time", START_TIMES)
 @pytest.mark.parametrize("end_time", END_TIMES)
 @pytest.mark.parametrize("logging_step_size", LOGGING_STEP_SIZES)
 @pytest.mark.parametrize("num_samples", NUM_SAMPLES)
+@pytest.mark.parametrize("rtol", RTOL)
+@pytest.mark.parametrize("atol", ATOL)
 def test_sample_no_interventions(
-    sample_method, model, start_time, end_time, logging_step_size, num_samples
+    sample_method,
+    model,
+    start_time,
+    end_time,
+    logging_step_size,
+    num_samples,
+    rtol,
+    atol,
 ):
     model_url = model.url
 
     with pyro.poutine.seed(rng_seed=0):
         result1 = sample_method(
-            model_url, end_time, logging_step_size, num_samples, start_time=start_time
+            model_url,
+            end_time,
+            logging_step_size,
+            num_samples,
+            start_time=start_time,
+            solver_options={"rtol": rtol, "atol": atol},
         )["unprocessed_result"]
     with pyro.poutine.seed(rng_seed=0):
         result2 = sample_method(
-            model_url, end_time, logging_step_size, num_samples, start_time=start_time
+            model_url,
+            end_time,
+            logging_step_size,
+            num_samples,
+            start_time=start_time,
+            solver_options={"rtol": rtol, "atol": atol},
         )["unprocessed_result"]
 
     result3 = sample_method(
-        model_url, end_time, logging_step_size, num_samples, start_time=start_time
+        model_url,
+        end_time,
+        logging_step_size,
+        num_samples,
+        start_time=start_time,
+        solver_options={"rtol": rtol, "atol": atol},
     )["unprocessed_result"]
 
     for result in [result1, result2, result3]:
@@ -364,8 +391,10 @@ def test_calibrate_no_kwargs(
 @pytest.mark.parametrize("start_time", START_TIMES)
 @pytest.mark.parametrize("end_time", END_TIMES)
 @pytest.mark.parametrize("logging_step_size", LOGGING_STEP_SIZES)
+@pytest.mark.parametrize("rtol", RTOL)
+@pytest.mark.parametrize("atol", ATOL)
 def test_calibrate_deterministic(
-    model_fixture, start_time, end_time, logging_step_size
+    model_fixture, start_time, end_time, logging_step_size, rtol, atol
 ):
     model_url = model_fixture.url
     (
@@ -381,6 +410,7 @@ def test_calibrate_deterministic(
         "data_mapping": model_fixture.data_mapping,
         "start_time": start_time,
         "deterministic_learnable_parameters": deterministic_learnable_parameters,
+        "solver_options": {"rtol": rtol, "atol": atol},
         **CALIBRATE_KWARGS,
     }
 
@@ -400,7 +430,10 @@ def test_calibrate_deterministic(
         assert torch.allclose(param_value, param_sample_2[param_name])
 
     result = sample(
-        *sample_args, **sample_kwargs, inferred_parameters=inferred_parameters
+        *sample_args,
+        **sample_kwargs,
+        inferred_parameters=inferred_parameters,
+        solver_options={"rtol": rtol, "atol": atol},
     )["unprocessed_result"]
 
     check_result_sizes(result, start_time, end_time, logging_step_size, 1)
@@ -563,7 +596,9 @@ def test_output_format(
 @pytest.mark.parametrize("start_time", START_TIMES)
 @pytest.mark.parametrize("end_time", END_TIMES)
 @pytest.mark.parametrize("num_samples", NUM_SAMPLES)
-def test_optimize(model_fixture, start_time, end_time, num_samples):
+@pytest.mark.parametrize("rtol", RTOL)
+@pytest.mark.parametrize("atol", ATOL)
+def test_optimize(model_fixture, start_time, end_time, num_samples, rtol, atol):
     logging_step_size = 1.0
     model_url = model_fixture.url
 
@@ -581,7 +616,7 @@ def __call__(self, x):
     optimize_kwargs = {
         **model_fixture.optimize_kwargs,
         "solver_method": "euler",
-        "solver_options": {"step_size": 0.1},
+        "solver_options": {"step_size": 0.1, "rtol": rtol, "atol": atol},
         "start_time": start_time,
         "n_samples_ouu": int(2),
         "maxiter": 1,