adds pexprb4 impl in python

Author: wgurecky

ormatex_py/ode_exp.py

@@ -35,8 +35,8 @@
 
 class ExpRBIntegrator(IntegrateSys):
 
-    _valid_methods = {"exprb2": 2, "exprb3": 3, "epi2": 2, "epi3": 3,
-                      "exprb2_dense": 2, "exprb2_pfd": 2,
+    _valid_methods = {"exprb2": 2, "exprb3": 3, "pexprb4": 4, "epi2": 2, "epi3": 3,
+                      "exprb2_dense": 2, "exprb2_pfd": 2, "exp_pfd": 1,
                       "exprb2_pfd_rs": 2, "exp_pfd_rs": 1}
 
     def __init__(self, sys: OdeSys, t0: float, y0: jax.Array, method="epi2", **kwargs):
@@ -78,7 +78,7 @@ def _remf(self, tr: float, yr: jax.Array,
         jac_yd = sys_jac_lop_yt(yr - yt)
         return frhs_yr - frhs_yt - jac_yd - v*dt
 
-    def _phi2v_nonauto(self, sys_jac_lop, dt):
+    def _phi2v_nonauto(self, sys_jac_lop, dt, c=1.0):
         r"""
         For rosenbrock exp integrators, this method computes
         the correction term for nonautonomous systems.
@@ -96,10 +96,10 @@ def _phi2v_nonauto(self, sys_jac_lop, dt):
             fytt = sys_jac_lop._fdt()
             # check for nonautonomous system
             if jnp.linalg.norm(fytt, ord=jax.numpy.inf) > self.tol_fdt:
-                return (dt**2.)*phi_linop(sys_jac_lop, dt, fytt, 2, self.max_krylov_dim, self.iom), fytt
+                return (c**2.)*(dt**2.)*phi_linop(sys_jac_lop, c*dt, fytt, 2, self.max_krylov_dim, self.iom), fytt
         return 0., 0.
 
-    def _step_exprb2(self, dt: float) -> StepResult:
+    def _step_exprb2(self, dt: float, frhs_kwargs: dict) -> StepResult:
         """
         Computes the solution update by:
         y_{t+1} = y_t + dt*\varphi_1(dt*J_t)F(t, y_t)+
@@ -109,7 +109,7 @@ def _step_exprb2(self, dt: float) -> StepResult:
         """
         t = self.t
         yt = self.y_hist[0]
-        sys_jac_lop = self.sys.fjac(t, yt)
+        sys_jac_lop = self.sys.fjac(t, yt, frhs_kwargs=frhs_kwargs)
         fyt = sys_jac_lop._frhs_cached()
         phi2_v, _0 = self._phi2v_nonauto(sys_jac_lop, dt)
         y_new = yt \
@@ -181,19 +181,22 @@ def _step_epi3(self, dt: float, frhs_kwargs: dict) -> StepResult:
 
         return StepResult(t+dt, dt, y_new, y_err)
 
-    def _step_exprb3(self, dt: float) -> StepResult:
-        """
+    def _step_exprb3(self, dt: float, frhs_kwargs: dict) -> StepResult:
+        r"""
         Computes the solution update by:
-        y_{t+1} = y_t + dt*\varphi_1(dt*J_t)F(t, y_t) +
-            2*dt*\varphi_3(dt*J_t)R_2 +
-            dt**2*\varphi_2(dt*J_t)F'(t, y_t)
 
-        doi: https://doi.org/10.1137/080717717
+        .. math::
+
+            y_{t+1} = y_t + dt*\varphi_1(dt*J_t)F(t, y_t) +
+                2*dt*\varphi_3(dt*J_t)R_2 +
+                dt**2*\varphi_2(dt*J_t)F'(t, y_t)
+
+        Ref: doi: https://doi.org/10.1137/080717717
         """
         t = self.t
         yt = self.y_hist[0] # y_t
 
-        sys_jac_lop = self.sys.fjac(t, yt)
+        sys_jac_lop = self.sys.fjac(t, yt, frhs_kwargs=frhs_kwargs)
         fyt = sys_jac_lop._frhs_cached()
 
         # phi2_v is nonzero for nonautonomous systems
@@ -215,6 +218,62 @@ def _step_exprb3(self, dt: float) -> StepResult:
 
         return StepResult(t+dt, dt, y_new, y_err)
 
+    def _step_pexprb4(self, dt: float, frhs_kwargs: dict) -> StepResult:
+        r"""
+        Computes the solution update by:
+
+        .. math::
+
+            U_{n2} = u_n + (0.5)dt\varphi_1((0.5)dt J_t)F(t, y_t) +
+                0.5^2 dt^2 \varphi_2(0.5 dt*J_t)F'(t, y_t)
+            U_{n3} = u_n + (1.0)dt\varphi_1((1.0)dt J_t)F(t, y_t) +
+                dt^2 \varphi_2(dt*J_t)F'(t, y_t)
+
+            y_{t+1} = y_t + dt*\varphi_1(dt*J_t)F(t, y_t) +
+                dt^2 \varphi_2(dt*J_t)F'(t, y_t) +
+                dt [16 \varphi_3(dt J_t) - 48 \varphi_4(dt J_t)] R_2(U_{n2}) +
+                dt [-2 \varphi_3(dt J_t) + 12 \varphi_4(dt J_t)] R_3(U_{n3})
+
+        Where $`U_{n2}`$ and $`U_{n3}`$ can be computed in parallel.
+
+        Ref: V. T. Luan. and A. Ostermann. Parallel Exponential rosenbrock methods.
+        Computers and Mathmatics with Applications, v71. 2016.
+        """
+        t = self.t
+        yt = self.y_hist[0] # y_t
+
+        sys_jac_lop = self.sys.fjac(t, yt, frhs_kwargs=frhs_kwargs)
+        fyt = sys_jac_lop._frhs_cached()
+
+        # butcher tableau coeffs
+        c_2 = 0.5
+        c_3 = 1.0
+
+        # compute U_{n2}
+        t_2 = t + c_2*dt
+        phi2_v_2, v_2 = self._phi2v_nonauto(sys_jac_lop, dt, c=c_2)
+        y_2 = yt \
+            + c_2*dt*phi_linop(sys_jac_lop, c_2*dt, fyt, 1, self.max_krylov_dim, self.iom) \
+            + phi2_v_2
+        r_2 = self._remf(t_2, y_2, fyt, sys_jac_lop, v=v_2)
+
+        # compute U_{n3}
+        t_3 = t + c_3*dt
+        phi2_v, v = self._phi2v_nonauto(sys_jac_lop, dt, c=c_3)
+        y_3 = yt \
+            + c_3*dt*phi_linop(sys_jac_lop, c_3*dt, fyt, 1, self.max_krylov_dim, self.iom) \
+            + phi2_v
+        r_3 = self._remf(t_3, y_3, fyt, sys_jac_lop, v=v)
+
+        # compute final update
+        vb0 = jnp.zeros(yt.shape)
+        b2_b3 = kiops_fixedsteps(
+            sys_jac_lop, dt, [vb0, vb0, vb0, dt*(16.0*r_2-2.0*r_3), dt*(-48.0*r_2+12.0*r_3)],
+            max_krylov_dim=self.max_krylov_dim, iom=self.iom)
+        y_new = y_3 + b2_b3
+
+        y_err = -1.0
+        return StepResult(t+dt, dt, y_new, y_err)
 
     @jax.jit
     def _step_exprb2_jit(t, yt, dt, sys):
@@ -323,22 +382,43 @@ def _step_exp_pfd_rs(self, dt: float) -> StepResult:
         y_err = -1.
         return StepResult(t+dt, dt, y_new, y_err)
 
+    def _step_exp_pfd(self, dt: float, frhs_kwargs: dict) -> StepResult:
+        r"""
+        Computes the solution update by:
+        y_{t+1} = \varphi_0(dt*L)*y0
+        where L is a dense matrix and computing varphi
+        using cauchy contour integral approach with quadrature rule
+        NOTE: Only useful for pure linear systems
+        """
+        t = self.t
+        yt = self.y_hist[0]
+        J = np.asarray(self.sys.fjac(t, yt, frhs_kwargs=frhs_kwargs).dense())
+
+        phi0J_yt = f_phi_k_pfd(J*dt, yt, 0, self.pfd_method)
+        y_new = jnp.asarray(phi0J_yt.flatten())
+        y_err = -1.
+        return StepResult(t+dt, dt, y_new, y_err)
+
     def step(self, dt: float, frhs_kwargs: dict={}) -> StepResult:
         if self.method == "exprb2":
-            return self._step_exprb2(dt)
+            return self._step_exprb2(dt, frhs_kwargs)
         elif self.method == "exprb3":
-            return self._step_exprb3(dt)
+            return self._step_exprb3(dt, frhs_kwargs)
+        elif self.method == "pexprb4":
+            return self._step_pexprb4(dt, frhs_kwargs)
         elif self.method == "epi2":
             return self._step_epi2(dt, frhs_kwargs)
         elif self.method == "epi3":
             if len(self.y_hist) >= 2:
                 return self._step_epi3(dt, frhs_kwargs)
             else:
-                return self._step_epi2(dt, frhs_kwargs)
+                return self._step_exprb3(dt, frhs_kwargs)
         elif self.method == "exprb2_dense":
             return self._step_exprb2_dense(dt)
         elif self.method == "exprb2_pfd":
             return self._step_exprb2_pfd(dt, frhs_kwargs)
+        elif self.method == "exp_pfd":
+            return self._step_exp_pfd(dt, frhs_kwargs)
         elif self.method == "exprb2_pfd_rs" and HAS_ORMATEX_RUST:
             return self._step_exprb2_pfd_rs(dt)
         elif self.method == "exp_pfd_rs" and HAS_ORMATEX_RUST:

ormatex_py/ode_sys.py

@@ -258,7 +258,8 @@ class AdJacLinOp(SysJacLinOp):
 
     def __init__(self, t, u, frhs: eqx.Module, frhs_kwargs: dict={}, **kwargs):
         super().__init__(t, u, frhs, frhs_kwargs, **kwargs)
-        self._frhs_u, self._fjac_u = jax.linearize(partial(self._frhs, **self._frhs_kwargs), self._t, self._u)
+        self._frhs_u, self._fjac_u = jax.linearize(
+                partial(self._frhs, **self._frhs_kwargs), self._t, self._u)
 
     @jax.jit
     def _matvec(self, v: jax.Array) -> jax.Array:

ormatex_py/progression/bateman_sys.py

@@ -4,6 +4,7 @@
 import jax
 import numpy as np
 from jax import numpy as jnp
+from itertools import cycle
 
 from ormatex_py import integrate_wrapper
 from ormatex_py.ode_sys import OdeSys, OdeSplitSys, MatrixLinOp, CustomJacLinOp
@@ -173,7 +174,7 @@ def analytic_bateman_single_parent(t, batmat, n0):
     return np.asarray(N, dtype=np.float64)
 
 
-def analytic_bateman_s3(method="epi2", do_plot=True, dt=10.0, tf=1000., pfd_method="cram_16"):
+def analytic_bateman_s3(method="epi2", do_plot=True, dt=10.0, tf=1000., pfd_method="cram_16", true_analytic=False):
     jax.config.update("jax_enable_x64", True)
     keymap = ["c_0", "c_1", "c_2"]
     decay_lib_sp = {
@@ -185,15 +186,31 @@ def analytic_bateman_s3(method="epi2", do_plot=True, dt=10.0, tf=1000., pfd_meth
     n0 = 1.0
     t0 = 0.0
     t = np.arange(t0, tf+dt, dt)
+    # test_ode_sys = TestBatemanSysJac(keymap, decay_lib_sp)
+    test_ode_sys = TestBatemanSysNonlinFeed(keymap, decay_lib_sp)
+    y0 = jnp.array([n0, 0.0, 0.0])
+
     # analytic result
-    y_true = analytic_bateman_single_parent(t, bmat, n0)
+    if true_analytic:
+        y_true = analytic_bateman_single_parent(t, bmat, n0)
+    else:
+        # use a fine timestep size with exprb3 integrator
+        dt_fine = 0.001
+        nsteps_fine = int((tf - t0) / dt_fine)
+        fine_res = integrate_wrapper.integrate(test_ode_sys, y0, t0, dt_fine, nsteps_fine, 'dopri5', tol_fdt=1e-25)
+        t_fine = np.asarray(fine_res.t_res)
+        y_fine = np.asarray(fine_res.y_res)
+        y_true = np.array([
+                np.interp(t, t_fine, y_fine[:, 0]),
+                np.interp(t, t_fine, y_fine[:, 1]),
+                np.interp(t, t_fine, y_fine[:, 2])
+                ]).T
 
     # compute numerical result
-    test_ode_sys = TestBatemanSysJac(keymap, decay_lib_sp)
-    y0 = jnp.array([n0, 0.0, 0.0])
     nsteps = int((tf - t0) / dt)
     res = integrate_wrapper.integrate(
-            test_ode_sys, y0, t0, dt, nsteps, method, max_krylov_dim=12, iom=12, pfd_method=pfd_method)
+            test_ode_sys, y0, t0, dt, nsteps, method, pfd_method=pfd_method, tol_fdt=1e-25,
+            max_krylov_dim=100, iom=100)
     t_res, y_res = res.t_res, res.y_res
     t_res = np.asarray(t_res)
     y_res = np.asarray(y_res)
@@ -203,7 +220,7 @@ def analytic_bateman_s3(method="epi2", do_plot=True, dt=10.0, tf=1000., pfd_meth
         plt.figure()
         plt.xscale('log')
         plt.yscale('log')
-        plt.ylim((1e-8, 10.0))
+        plt.ylim((1e-8, 80.0))
         plt.xlim((1.0, tf))
         # numerical
         plt.plot(t_res+1.0, y_res[:, 0]+1e-16, label="c_0")
@@ -225,9 +242,14 @@ def analytic_bateman_s3(method="epi2", do_plot=True, dt=10.0, tf=1000., pfd_meth
 
 
 def run_sweep():
-    methods = ["epi2", "epi3", "exprb3", "exp2_dense", "exp3_dense",
-               "exprb2_dense", "exprb2_pfd_rs", "exp_pfd_rs",
-               "implicit_euler", "implicit_esdirk3", "implicit_esdirk4"]
+    methods = ["epi2_leja_re", "epi3", "exprb3", "exprb2_rs", "exprb3_rs", "epi3_rs",
+               "exprb2_pfd", "exp_pfd",
+               "implicit_euler", "implicit_esdirk3",
+               ]
+    methods = ["exprb2", "exprb3", "exprb3_rs", "epi3", "epi3_rs",
+               "implicit_euler", "implicit_esdirk3",
+              ]
+    methods = ["pexprb4", "exprb3", "exprb2"]
     dts = [1., 2., 5., 10., 25., 50.]
     tf = 100.
     nspecies = 3
@@ -250,10 +272,13 @@ def run_sweep():
 
     # error vs time step size for each method
     plt.figure()
+    markers = ['x', 'p']
+    marker_cycler = cycle(markers)
     for method in methods:
         dt = mae_dict[method][:, 0]
-        s3_err = mae_dict[method][:, 2]
-        plt.plot(dt, s3_err, '-o', label=method)
+        s3_err = mae_dict[method][:, 3]
+        mk = next(marker_cycler)
+        plt.plot(dt, s3_err, marker=mk, ls='--', label=method, alpha=0.75)
     plt.yscale("log")
     plt.xscale("log")
     plt.grid(ls='--')
@@ -265,8 +290,49 @@ def run_sweep():
     plt.savefig("bateman_ex_1_converg.png")
     for method in methods:
         for dt in [10., 25.]:
-            analytic_bateman_s3(method, dt=dt, tf=500., do_plot=True)
+            analytic_bateman_s3(method, dt=dt, tf=200., do_plot=True)
+
+class TestBatemanSysNonlinFeed(OdeSys):
+    bat_mat: jax.Array
+    feed_period: float
+    feed_scale: float
+
+    def __init__(self, keymap, decay_lib, *args, **kwargs):
+        bmat = gen_bateman_matrix(keymap, decay_lib)
+        print("Bateman test system to solve:")
+        print(bmat)
+        self.bat_mat = bmat
+        self.feed_period = 0.2e-1
+        self.feed_scale = 1.0e-2
+        super().__init__()
 
+    def _nonlin_feed(self, t, u, **kwargs):
+        # linear-in-time feed example
+        # feed_rate = jnp.clip(1e-6*t, 0.0, 100.0)
+        # quadratic-in-time feed example
+        # feed_rate = -1e-7*(t-500.)**2+1e-7*500**2
+        # feed_rate = jnp.clip(-1e-7*(t-200.)**2+1e-7*500**2, 0.0, 100.0)
+        # cubic-in-time feed example
+        feed_rate = jnp.clip(2e-10*(t-200)**3-1e-7*(t-200.)**2+1e-7*500**2, 0.0, 100.0)
+        fdr = jnp.zeros_like(u)
+        fdr = fdr.at[0].set(feed_rate)
+        return fdr
+
+    def _nonlin_sink(self, t, u, **kwargs):
+        # quadratic-in-time sink example
+        sink_rate = jnp.clip(-1e-6*((t-100)**2)+0.01, 0.0, 1.0)
+        # linear-in-time sink example
+        # sink_rate = jnp.clip(1e-5*t, 0.0, 1.0)
+        s = jnp.zeros_like(u)
+        s = s.at[2].set(-sink_rate)
+        return s
+
+    @jax.jit
+    def _frhs(self, t: float, u: jax.Array, **kwargs) -> jax.Array:
+        # res = self.bat_mat @ u + self._nonlin_sink(t, u, **kwargs)
+        # res = self.bat_mat @ u + self._nonlin_feed(t, u, **kwargs)
+        res = self.bat_mat @ u + self._nonlin_sink(t, u, **kwargs) + self._nonlin_feed(t, u, **kwargs)
+        return res
 
 class TestBatemanSysJac(OdeSplitSys):
     """
@@ -331,7 +397,7 @@ def _fl(self, t: float, u: jax.Array, **kwargs):
     plt.figure()
     plt.xscale('log')
     plt.yscale('log')
-    plt.ylim((1e-14, 10.0))
+    plt.ylim((1e-14, 80.0))
     plt.plot(t_res, y_res[:, 0], label="c_0")
     plt.plot(t_res, y_res[:, 1], label="c_1")
     plt.plot(t_res, y_res[:, 2], label="c_2")

ormatex_py/progression/lotka_volterra.py

@@ -158,7 +158,7 @@ def gen_sys(autonomous=True):
         nsteps = int(t_end/dt)
 
         res = integrate_wrapper.integrate(
-                lv_sys, y0, t0, dt, nsteps, method=method, max_krylov_dim=10, iom=5, tol_fdt=0)
+                lv_sys, y0, t0, dt, nsteps, method=method, max_krylov_dim=20, iom=20, tol_fdt=0)
         t_res, y_res = res.t_res, res.y_res
         t_res = jnp.asarray(t_res)
         y_res = jnp.asarray(y_res)
@@ -222,7 +222,7 @@ def gen_sys(autonomous=True):
 
 
 def sweep_methods(autonomous=False):
-    methods = ["epi3", "exprb2", "exprb3", "exp3_dense", "exp2_dense", "implicit_euler", "implicit_esdirk3"]
+    methods = ["epi3", "exprb2", "exprb3", "pexprb4", "implicit_euler", "implicit_esdirk3"]
     plt.figure()
     for method in methods:
         err_dt, err, lb = main(method, do_plot=False, autonomous=autonomous)