Improve numerical stability when integrating the ODE system (#69)

* Use the default integrator in general * Optionally use JAX for speed * Use an exact Jacobian using JAX * Use rtol and atol with solve_ivp * Improve reaction rate constants for equilibria * Allow the user to choose among ODE solvers * Allow the user to choose a simulation time Co-authored-by: Felipe S. S. Schneider <[email protected]>
geem-lab · Nov 29, 2020 · 893f8cb · 893f8cb
1 parent 34be51d
commit 893f8cb
Show file tree

Hide file tree

Showing 33 changed files with 375 additions and 209 deletions.
diff --git a/INSTALL.rst b/INSTALL.rst
@@ -17,10 +17,12 @@ The package, together with the above dependencies, can be installed from
 Optionally, extra functionality is provided such as a command-line interface
 and solvent properties::
 
-    pip install 'overreact[cli,solvents]'
+    pip install 'overreact[cli,fast,solvents]'
 
-This last line installs `Rich <https://github.com/willmcgugan/rich>`_
-and `thermo <https://github.com/CalebBell/thermo>`_ as well.
-Rich is used in the command-line interface, and thermo is used
-to calculate the dynamic viscosity of solvents in the context of the
-:doc:`tutorials/collins-kimball` for diffusion-limited reactions.
+This last line installs `Rich <https://github.com/willmcgugan/rich>`_,
+`JAX <https://jax.readthedocs.io/en/latest/index.html>`_ and
+`thermo <https://github.com/CalebBell/thermo>`_ as well.
+Rich is used in the command-line interface, JAX helps speedup calculations,
+and thermo is used to calculate the dynamic viscosity of solvents in the
+context of the :doc:`tutorials/collins-kimball` for diffusion-limited
+reactions.
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/docsrc/examples/hydrogen-abstraction-drc.py b/docsrc/examples/hydrogen-abstraction-drc.py
@@ -34,7 +34,6 @@
     model.compounds,
     y0,
     scale="M-1 s-1",
-    method="BDF",
     num=500,
     dx=2000.0,
     order=5,

diff --git a/docsrc/examples/hydrogen-abstraction-microkinetics.py b/docsrc/examples/hydrogen-abstraction-microkinetics.py
@@ -33,7 +33,7 @@
 print(y0)
 
 dydt = api.get_dydt(model.scheme, k_eck)
-y, r = api.get_y(dydt, y0=y0, method="Radau")
+y, r = api.get_y(dydt, y0=y0)
 
 print(model.scheme.compounds)
 print(y(y.t_max))

diff --git a/docsrc/notebooks/#1 Solving complex kinetics schemes.ipynb b/docsrc/notebooks/#1 Solving complex kinetics schemes.ipynb
diff --git a/docsrc/notebooks/#2 A mini language for chemical reactions.ipynb b/docsrc/notebooks/#2 A mini language for chemical reactions.ipynb
@@ -62,29 +62,29 @@
      "output_type": "execute_result",
      "data": {
       "text/plain": [
-       "(['A', 'B', 'TS*', 'S', 'E', 'ES', 'ES*', 'P'],\n",
-       " ['A + 2 B -> TS*',\n",
+       "(('A', 'B', 'TS*', 'S', 'E', 'ES', 'ES*', 'P'),\n",
+       " ('A + 2 B -> TS*',\n",
        "  'TS* -> S',\n",
        "  'E + S -> ES',\n",
        "  'ES -> E + S',\n",
        "  'ES -> ES*',\n",
-       "  'ES* -> E + P'],\n",
-       " [[-1.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [-2.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [1.0, -1.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [0.0, 1.0, -1.0, 1.0, 0.0, 0.0],\n",
-       "  [0.0, 0.0, -1.0, 1.0, 0.0, 1.0],\n",
-       "  [0.0, 0.0, 1.0, -1.0, -1.0, 0.0],\n",
-       "  [0.0, 0.0, 0.0, 0.0, 1.0, -1.0],\n",
-       "  [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]],\n",
-       " [[-1.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [-2.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [1.0, -1.0, 0.0, 0.0, 0.0, 0.0],\n",
-       "  [0.0, 1.0, -1.0, 0.0, 0.0, 0.0],\n",
-       "  [0.0, 0.0, -1.0, 0.0, 0.0, 1.0],\n",
-       "  [0.0, 0.0, 1.0, 0.0, -1.0, 0.0],\n",
-       "  [0.0, 0.0, 0.0, 0.0, 1.0, -1.0],\n",
-       "  [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]])"
+       "  'ES* -> E + P'),\n",
+       " ((-1.0, 0.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (-2.0, 0.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (1.0, -1.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (0.0, 1.0, -1.0, 1.0, 0.0, 0.0),\n",
+       "  (0.0, 0.0, -1.0, 1.0, 0.0, 1.0),\n",
+       "  (0.0, 0.0, 1.0, -1.0, -1.0, 0.0),\n",
+       "  (0.0, 0.0, 0.0, 0.0, 1.0, -1.0),\n",
+       "  (0.0, 0.0, 0.0, 0.0, 0.0, 1.0)),\n",
+       " ((-1.0, 0.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (-2.0, 0.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (1.0, -1.0, 0.0, 0.0, 0.0, 0.0),\n",
+       "  (0.0, 1.0, -1.0, 0.0, 0.0, 0.0),\n",
+       "  (0.0, 0.0, -1.0, 0.0, 0.0, 1.0),\n",
+       "  (0.0, 0.0, 1.0, 0.0, -1.0, 0.0),\n",
+       "  (0.0, 0.0, 0.0, 0.0, 1.0, -1.0),\n",
+       "  (0.0, 0.0, 0.0, 0.0, 0.0, 1.0)))"
       ]
      },
      "metadata": {},
@@ -151,7 +151,7 @@
    "source": [
     "# TODO: scipy 1.4+ will support passing args from solve_ivp, but for now let's do something rather mundane here\n",
     "def dydt(t, y):\n",
-    "    A = np.asanyarray(scheme.A)\n",
+    "    A = np.asarray(scheme.A)\n",
     "    r = k * np.prod(np.power(y, np.where(A > 0, 0, -A).T), axis=1)\n",
     "    return np.dot(scheme.A, r)"
    ]
@@ -189,7 +189,7 @@
    ],
    "source": [
     "y0 = np.array([1.5, 1., 0., 0., 0.1, 0., 0.0, 0.])\n",
-    "y, r = simulate.get_y(dydt, y0, t_span=[0.0, 1e-6], method=\"Radau\")\n",
+    "y, r = simulate.get_y(dydt, y0, t_span=[0.0, 1e-6])\n",
     "\n",
     "t = np.linspace(y.t_min, 5e-7, num=100)\n",
     "for i, compound in enumerate(scheme.compounds):\n",

diff --git a/docsrc/notebooks/#4 overreact and cclib.ipynb b/docsrc/notebooks/#4 overreact and cclib.ipynb
@@ -22,14 +22,14 @@
    "metadata": {},
    "outputs": [
     {
+     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "2625499.6394798253"
       ]
      },
-     "execution_count": 2,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 2
     }
    ],
    "source": [
@@ -52,14 +52,14 @@
    "metadata": {},
    "outputs": [
     {
+     "output_type": "execute_result",
      "data": {
       "text/plain": [
-       "Scheme(compounds=['CH4', 'Cl.', 'RC', 'RC*', 'PC', 'CH3.', 'HCl'], reactions=['CH4 + Cl. -> RC', 'RC -> CH4 + Cl.', 'RC -> RC*', 'RC* -> PC', 'PC -> CH3. + HCl', 'CH3. + HCl -> PC'], is_half_equilibrium=[True, True, False, False, True, True], A=[[-1.0, 1.0, 0.0, 0.0, 0.0, 0.0], [-1.0, 1.0, 0.0, 0.0, 0.0, 0.0], [1.0, -1.0, -1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, -1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, -1.0, 1.0], [0.0, 0.0, 0.0, 0.0, 1.0, -1.0], [0.0, 0.0, 0.0, 0.0, 1.0, -1.0]], B=[[-1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0, 0.0, -1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, -1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, -1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]])"
+       "Scheme(compounds=('CH4', 'Cl.', 'RC', 'RC*', 'PC', 'CH3.', 'HCl'), reactions=('CH4 + Cl. -> RC', 'RC -> CH4 + Cl.', 'RC -> RC*', 'RC* -> PC', 'PC -> CH3. + HCl', 'CH3. + HCl -> PC'), is_half_equilibrium=(True, True, False, False, True, True), A=((-1.0, 1.0, 0.0, 0.0, 0.0, 0.0), (-1.0, 1.0, 0.0, 0.0, 0.0, 0.0), (1.0, -1.0, -1.0, 0.0, 0.0, 0.0), (0.0, 0.0, 1.0, -1.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0, -1.0, 1.0), (0.0, 0.0, 0.0, 0.0, 1.0, -1.0), (0.0, 0.0, 0.0, 0.0, 1.0, -1.0)), B=((-1.0, 0.0, 0.0, 0.0, 0.0, 0.0), (-1.0, 0.0, 0.0, 0.0, 0.0, 0.0), (1.0, 0.0, -1.0, 0.0, 0.0, 0.0), (0.0, 0.0, 1.0, -1.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0, -1.0, 0.0), (0.0, 0.0, 0.0, 0.0, 1.0, 0.0), (0.0, 0.0, 0.0, 0.0, 1.0, 0.0)))"
       ]
      },
-     "execution_count": 3,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 3
     }
    ],
    "source": [
@@ -82,9 +82,9 @@
    "metadata": {},
    "outputs": [
     {
+     "output_type": "error",
      "ename": "AttributeError",
      "evalue": "'NoneType' object has no attribute 'parse'",
-     "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
@@ -202,7 +202,7 @@
    "outputs": [],
    "source": [
     "dydt = simulate.get_dydt(scheme, k)\n",
-    "t, y = simulate.get_y(dydt, y0=[1.0, 1.0, 0.0, 0.0, 0., 0., 0.], t_span=[0.0, 1.0e-4], method=\"Radau\")\n",
+    "t, y = simulate.get_y(dydt, y0=[1.0, 1.0, 0.0, 0.0, 0., 0., 0.], t_span=[0.0, 1.0e-4])\n",
     "\n",
     "for i, compound in enumerate(scheme.compounds):\n",
     "    if not compound.endswith(\"*\"):\n",
@@ -231,9 +231,13 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
+   "name": "python3",
+   "display_name": "Python 3.8.6 64-bit",
+   "metadata": {
+    "interpreter": {
+     "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+    }
+   }
   },
   "language_info": {
    "codemirror_mode": {
@@ -245,9 +249,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.6-final"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
diff --git a/docsrc/tutorials/1-rate_constant.rst b/docsrc/tutorials/1-rate_constant.rst
@@ -29,7 +29,7 @@ Simple concurrent first order reactions using the rates above would be:
 ...     A -> AC‡ -> C
 ... """)
 >>> dydt = simulate.get_dydt(scheme, k)
->>> y, r = simulate.get_y(dydt, y0=[1.0, 0.0, 0.0, 0.0, 0.0], method="Radau")
+>>> y, r = simulate.get_y(dydt, y0=[1.0, 0.0, 0.0, 0.0, 0.0])
 >>> t = np.linspace(y.t_min, 5.0)
 >>> plt.clf()
 >>> for i, compound in enumerate(scheme.compounds):

diff --git a/docsrc/tutorials/2-equilibrium-constants.rst b/docsrc/tutorials/2-equilibrium-constants.rst
@@ -20,7 +20,7 @@ So let's check it:
 >>> scheme.compounds, scheme.reactions
 (('Cd2p', 'MeNH2', '[Cd(MeNH2)4]2p'),
  ('Cd2p + 4 MeNH2 -> [Cd(MeNH2)4]2p', '[Cd(MeNH2)4]2p -> Cd2p + 4 MeNH2'))
->>> dydt = simulate.get_dydt(scheme, [K, 1.])
+>>> dydt = simulate.get_dydt(scheme, np.array([K, 1.]))
 >>> y, r = simulate.get_y(dydt, y0=[0., 0., 1.])
 >>> y(y.t_max)
 array([0.01608807, 0.06435228, 0.98391193])

diff --git a/docsrc/tutorials/4-curtin-hammett.rst b/docsrc/tutorials/4-curtin-hammett.rst
@@ -24,7 +24,7 @@ The selectivity is defined as the ratio between both products:
 >>> k
 array([1.445e+13, 6.212e+12, 2.905e+05, 2.310e+04])
 
->>> y, r = simulate.get_y(simulate.get_dydt(scheme, k), y0=[1.0, 0.0, 0.0, 0.0, 0.0, 0.0], method="BDF")
+>>> y, r = simulate.get_y(simulate.get_dydt(scheme, k), y0=[1.0, 0.0, 0.0, 0.0, 0.0, 0.0])
 >>> S = y(y.t_max)[scheme.compounds.index("P1")] / y(y.t_max)[scheme.compounds.index("P2")]
 >>> S
 5.408
diff --git a/docsrc/tutorials/degree-of-rate-control/degree-of-rate-control.rst b/docsrc/tutorials/degree-of-rate-control/degree-of-rate-control.rst
@@ -31,7 +31,7 @@ Now the graph:
 
 >>> import matplotlib.pyplot as plt
 >>> y0 = [1., 1000., 0, 1., 0., 0., 0., 0.]
->>> y, r = simulate.get_y(dydt, y0, [0.0, 30 * 60.0], method="BDF")
+>>> y, r = simulate.get_y(dydt, y0, [0.0, 30 * 60.0])
 >>> t = np.linspace(y.t_min, 500.0)
 >>> plt.clf()
 >>> for i, compound in enumerate(scheme.compounds):

diff --git a/docsrc/tutorials/simulation/0-basic-reaction-simulation.rst b/docsrc/tutorials/simulation/0-basic-reaction-simulation.rst
@@ -42,8 +42,9 @@ The above model translates to the following in overreact:
 overreact helps us to define the equations and solve the initial value problem.
 First, let's define the system of ordinary differential equations:
 
+>>> import numpy as np
 >>> from overreact import simulate
->>> dydt = simulate.get_dydt(scheme, [kf])
+>>> dydt = simulate.get_dydt(scheme, np.array([kf]))
 
 The returned object above is a function of concentrations and time that defines
 a set of ordinary differential equations in time:
@@ -56,9 +57,7 @@ We are going to simulate 10 seconds, starting with an initial concentration of
 1 molar of A (the concentration units evidently depend on the units of the
 reaction rate constant).
 
->>> y, r = simulate.get_y(dydt, y0=[1., 0.], method="Radau")
-
->>> import numpy as np
+>>> y, r = simulate.get_y(dydt, y0=[1., 0.])
 >>> t = np.linspace(y.t_min, 5.0)
 
 The simulation data is stored in `t` (points in time) and `y` (concentrations).
@@ -84,4 +83,4 @@ We can see that the reaction went to full completion by checking the final
 concentrations:
 
 >>> y(y.t_max)
-array([0.0000, 1.0000])
+array([0.00, 1.00])
diff --git a/docsrc/tutorials/simulation/1-basic-reaction-simulation.rst b/docsrc/tutorials/simulation/1-basic-reaction-simulation.rst
@@ -72,8 +72,9 @@ values from the table:
 We need one reaction rate constant per reaction, in a
 list with the same ordering as the ``scheme.reactions`` variable obtained above.
 
+>>> import numpy as np
 >>> k_f = (k_r + k_cat) / K_M  # using the equation above
->>> k = [k_f, k_r, k_cat]
+>>> k = np.array([k_f, k_r, k_cat])
 
 Solving the initial value problem
 ---------------------------------
@@ -83,7 +84,7 @@ units, as long as they match with the reaction rate constants) and solve the
 initial value problem. Below we set the substrate to one molar and the enzyme
 to 5% of it:
 
->>> y0 = [0.05, 1.00, 0.00, 0.00, 0.00]
+>>> y0 = np.array([0.05, 1.00, 0.00, 0.00, 0.00])
 
 One return value that ``core.parse_reactions`` has given us was the :math:`A`
 matrix, whose entry :math:`A_{ij}` stores the coefficient of the i-th compound
@@ -135,19 +136,16 @@ Fortunately, overreact is able to give you a function that, giving time and
 concentrations, calculates the derivative with respect to time:
 
 >>> dydt = simulate.get_dydt(scheme, k)
->>> dydt(0.0, y0)  # t = 0.0
+>>> dydt(0.0, y0)  # t = 0.0  # doctest: +SKIP
 array([-83333.3, -83333.3, 83333.3, 0. , 0. ])
 
 From the above we see that the equilibrium will likely be rapidly satisfied,
 while no product is being created at time zero, since there's no
 enzyme-substrate complex yet.
 
-Let's now do a one minute simulation with ``get_y`` (methods Radau or BDF are
-recommended for likely stiff equations such as those):
-
->>> y, r = simulate.get_y(dydt, y0, method="Radau")
+Let's now do a one minute simulation with ``get_y``:
 
->>> import numpy as np
+>>> y, r = simulate.get_y(dydt, y0)
 >>> t = np.linspace(y.t_min, 60.0)  # seconds
 
 We can graph concentrations over time with ``t`` and ``y``:
@@ -168,15 +166,15 @@ Text(...)
 
 .. figure:: ../../_static/michaelis-menten.png
 
-   A one minute simulation of the Michaelis-Menten model for the enzyme Pepsin,
+   A one-minute simulation of the Michaelis-Menten model for the enzyme Pepsin,
    an endopeptidase that breaks down proteins into smaller peptides. Observe
-   that the rapid equilibrium justifies the commonly applied steady state
+   that the rapid equilibrium justifies the commonly applied steady-state
    approximation.
 
 The simulation time was enough to convert all substrate into products and
 regenerate the initial enzyme molecules:
 
->>> y(y.t_max)
+>>> y(y.t_max)  # doctest: +SKIP
 array([0.05, 0.00, 0.00, 0.00, 1.00])
 
 Getting rates back
@@ -199,7 +197,7 @@ Text(...)
 
 .. figure:: ../../_static/michaelis-menten-dydt.png
 
-   Time derivative of concentrations for the one minute simulation of the Michaelis-Menten model for the enzyme Pepsin above.
+   The time derivative of concentrations for the one-minute simulation of the Michaelis-Menten model for the enzyme Pepsin above.
 
 Furthermore, we can get the turnover frequency (TOF) as:
 
@@ -218,4 +216,4 @@ Text(...)
 
 .. figure:: ../../_static/michaelis-menten-tof.png
 
-   Turnover frequency for the enzyme Pepsin above, in the one minute simulation of the Michaelis-Menten model.
+   The turnover frequency for the enzyme Pepsin above, in the one-minute simulation of the Michaelis-Menten model.
diff --git a/overreact/_thermo/__init__.py b/overreact/_thermo/__init__.py
@@ -510,7 +510,7 @@ def get_molecularity(transform):
     >>> get_molecularity([[0.], [0.]])
     array([1])
     """
-    res = np.sum(np.asanyarray(transform) < 0, axis=0)
+    res = np.sum(np.asarray(transform) < 0, axis=0)
     return np.where(res > 0, res, 1)
 
 
@@ -554,7 +554,7 @@ def get_delta(transform, property):
     ...            [ 1,  3]], [-5, 12])
     array([17, 46])
     """
-    return np.asanyarray(transform).T @ np.asanyarray(property)
+    return np.asarray(transform).T @ np.asarray(property)
 
 
 # TODO(schneiderfelipe): further test the usage of delta_moles with values
@@ -614,13 +614,13 @@ def equilibrium_constant(
     >>> equilibrium_constant(64187.263215698644, delta_moles=-2, temperature=745.0)
     0.118e-6
     """
-    temperature = np.asanyarray(temperature)
+    temperature = np.asarray(temperature)
 
     if volume is None:
         volume = molar_volume(temperature=temperature, pressure=pressure)
 
     equilibrium_constant = (
-        np.exp(-np.asanyarray(delta_freeenergy) / (constants.R * temperature))
+        np.exp(-np.asarray(delta_freeenergy) / (constants.R * temperature))
         * (volume) ** -delta_moles
     )
     logger.info(f"equilibrium constant = {equilibrium_constant}")
@@ -704,7 +704,7 @@ def change_reference_state(
     if old_reference is None:
         if volume is None:
             volume = molar_volume(
-                temperature=np.asanyarray(temperature), pressure=pressure
+                temperature=np.asarray(temperature), pressure=pressure
             )
         old_reference = 1.0 / volume
     return sign * constants.R * np.log(new_reference / old_reference)
@@ -745,5 +745,5 @@ def get_reaction_entropies(transform):
     >>> get_reaction_entropies(scheme.A)
     array([-5.763, 5.763])
     """
-    sym = _factorial(np.abs(np.asanyarray(transform)))
+    sym = _factorial(np.abs(np.asarray(transform)))
     return np.sum(np.sign(transform) * change_reference_state(sym, 1), axis=0)