Use rtol and atol with solve_ivp

overreact/api.py

@@ -503,7 +503,7 @@ def get_drc(
     compounds,
     y0,
     t_span=None,
-    method="BDF",
+    method="Radau",
     qrrho=True,
     scale="l mol-1 s-1",
     temperature=298.15,

overreact/cli.py

@@ -73,6 +73,8 @@ def __init__(
         qrrho=True,  # TODO(schneiderfelipe): change to use_qrrho
         temperature=298.15,
         bias=0.0,
+        rtol=1e-5,
+        atol=1e-9,
         box_style=box.SIMPLE,
     ):
         self.model = model
@@ -82,6 +84,8 @@ def __init__(
         self.qrrho = qrrho
         self.temperature = temperature
         self.bias = bias
+        self.rtol = rtol
+        self.atol = atol
         self.box_style = box_style
 
     def __rich_console__(self, console, options):
@@ -542,6 +546,8 @@ def main():
         ),
         action="store_false",
     )
+    parser.add_argument("--rtol", type=float, default=1e-5)
+    parser.add_argument("--atol", type=float, default=1e-9)
     parser.add_argument(
         "--plot",
         help=(
@@ -581,6 +587,8 @@ def main():
 - QRRHO?         = {args.qrrho}
 - Temperature    = {args.temperature} K
 - Pressure       = {args.pressure} Pa
+- Rel. Tol.      = {args.rtol}
+- Abs. Tol.      = {args.atol}
 - Bias           = {args.bias / constants.kcal} kcal/mol
 
 Parsing and calculating…
@@ -604,6 +612,8 @@ def main():
         qrrho=args.qrrho,
         temperature=args.temperature,
         bias=args.bias,
+        rtol=args.rtol,
+        atol=args.atol,
     )
     console.print(report, justify="left")
 

overreact/simulate.py

@@ -20,15 +20,14 @@
 if _found_jax:
     import jax.numpy as jnp
     from jax import jacfwd
-    from jax import jit
     from jax.config import config
 
     config.update("jax_enable_x64", True)
 else:
     jnp = np
 
 
-def get_y(dydt, y0, t_span=None, method="BDF"):
+def get_y(dydt, y0, t_span=None, method="Radau", rtol=1e-5, atol=1e-9):
     """Simulate a reaction scheme from its rate function.
 
     This uses scipy's ``solve_ivp`` under the hood.
@@ -48,6 +47,8 @@ def get_y(dydt, y0, t_span=None, method="BDF"):
         Integration method to use. See `scipy.integrade.solve_ivp` for details.
         Kinetics problems are very often stiff and, as such, "RK45" is
         normally unsuited. "Radau", "BDF" or "LSODA" are good choices.
+    rtol, atol : array-like
+        See `scipy.integrade.solve_ivp` for details.
 
     Returns
     -------
@@ -114,7 +115,17 @@ def get_y(dydt, y0, t_span=None, method="BDF"):
     if hasattr(dydt, "jac"):
         jac = dydt.jac
 
-    res = _solve_ivp(dydt, t_span, y0, method=method, dense_output=True, jac=jac)
+    # TODO(schneiderfelipe): log solve_ivp stuff.
+    res = _solve_ivp(
+        dydt,
+        t_span,
+        y0,
+        method=method,
+        dense_output=True,
+        rtol=rtol,
+        atol=atol,
+        jac=jac,
+    )
     y = res.sol
 
     def r(t):
@@ -213,7 +224,7 @@ def _dydt(t, y, k=k_adj, M=M):
         return jnp.dot(A, r)
 
     if _found_jax:
-        _dydt = jit(_dydt)
+        _dydt = _dydt
 
         def _jac(t, y, k=k_adj, M=M):
             # _jac(t, y)[i, j] == d f_i / d y_j