Use rtol and atol with solve_ivp

geem-lab · Nov 27, 2020 · 42e3b46 · 42e3b46
1 parent cdae256
commit 42e3b46
Show file tree

Hide file tree

Showing 8 changed files with 15 additions and 5 deletions.
diff --git a/docs/_static/drc.png b/docs/_static/drc.png
diff --git a/docs/_static/first-order.png b/docs/_static/first-order.png
diff --git a/docs/_static/michaelis-menten-dydt.png b/docs/_static/michaelis-menten-dydt.png
diff --git a/docs/_static/michaelis-menten-tof.png b/docs/_static/michaelis-menten-tof.png
diff --git a/docs/_static/michaelis-menten.png b/docs/_static/michaelis-menten.png
diff --git a/docs/_static/simple-first-order.png b/docs/_static/simple-first-order.png
diff --git a/overreact/api.py b/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,

diff --git a/overreact/simulate.py b/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-3, atol=1e-6):
     """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,16 @@ 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)
+    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 +223,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