Nonlinear coupled ODEs (Lotka-Volterra)

Download this notebook

In this example the dae module is used to solve the Lotka-Volterra problem, a set of nonlinear ODEs

import jax
import jax.numpy as jnp
from autopdex import dae

jax.config.update("jax_enable_x64", True)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

The dae module can solve ordinary differential equations expressed in an implicit form. Here, we have the following set of equations:

\(0 = -\frac{du}{dt} + αu − βuv\)

\(0 = -\frac{dv}{dt} −γv + δuv\)

Where:

u(t): prey population

v(t): predator population

α,β,γ,δ: model parameters

For autopdex, we can prepare this in a JAX-traceable function as follows:

def implicit_ode(q_fun, t, settings):
  # 'q_fun' is a function of time that returns the state variables accessible via their keywords.
  q_t_fun = jax.jacfwd(q_fun)
  q = q_fun(t)
  q_t = q_t_fun(t)

  u = q['u']
  v = q['v']
  u_t = q_t['u']
  v_t = q_t['v']

  # Here, we hardcode the parameters, but we could also load them from the settings dictionary and take derivatives with respect to them
  α, β, γ, δ = (0.1, 0.02, 0.4, 0.02)

  # Define the residuals of the system of ODEs
  res_u = u_t - (α * u - β * u * v)
  res_v = v_t - (-γ * v + δ * u * v)
  return jnp.array([res_u, res_v])
#

As in the PDE modules, the time stepping manager uses the dictionaries ‘settings’ and ‘static_settings’ in order to set up the problem. Here, we define the Lotka-Voltera-system as the ‘dae’ to be solved. Further, we chose the integrators for the different fields.

static_settings = {
    'dae': implicit_ode,
    'time integrators': {
        # 'u': dae.ForwardEuler(),
        # 'v': dae.ForwardEuler(),
        # 'u': dae.BackwardEuler(),
        # 'v': dae.BackwardEuler(),
        # 'u': dae.AdamsBashforth(4),
        # 'v': dae.AdamsBashforth(4),
        # 'u': dae.AdamsMoulton(1),
        # 'v': dae.AdamsMoulton(1),
        # 'u': dae.BackwardDiffFormula(3),
        # 'v': dae.BackwardDiffFormula(3),
        # 'u': dae.DiagonallyImplicitRungeKutta(3),
        # 'v': dae.DiagonallyImplicitRungeKutta(3),
        # 'u': dae.Kvaerno(5),
        # 'v': dae.Kvaerno(5),
        # 'u': dae.GaussLegendreRungeKutta(14),
        # 'v': dae.GaussLegendreRungeKutta(14),
        # 'u': dae.DormandPrince(5),
        # 'v': dae.DormandPrince(5),
        'u': dae.ExplicitRungeKutta(11),
        'v': dae.ExplicitRungeKutta(11),
    },
    'verbose': 0,
}

Next, we have to define the policies for time stepping and data saving.

manager = dae.TimeSteppingManager(
    static_settings,
    save_policy=dae.SaveEquidistantPolicy(),
    step_size_controller=dae.ConstantStepSizeController()
    # step_size_controller=dae.PIDController(rtol=1e-6, atol=1e-9)
    # step_size_controller=dae.RootIterationController(max_step_size = 2.)
)

After specifying the initial values, end time, initial time increment and maximal number of time steps, we can run the time stepping procedure.

dofs_0 = {
    'u': jnp.array([10.0]),
    'v': jnp.array([10.0]),
}
t_max = 140.0
num_time_steps = 2000
result = manager.run(dofs_0, t_max / num_time_steps, t_max, num_time_steps)

Progress: 3%, Time: 4.27e+00, dt: 7.00e-02, iterations: 1
Progress: 8%, Time: 1.12e+01, dt: 7.00e-02, iterations: 1
Progress: 13%, Time: 1.82e+01, dt: 7.00e-02, iterations: 1
Progress: 18%, Time: 2.52e+01, dt: 7.00e-02, iterations: 1
Progress: 23%, Time: 3.22e+01, dt: 7.00e-02, iterations: 1
Progress: 28%, Time: 3.92e+01, dt: 7.00e-02, iterations: 1
Progress: 33%, Time: 4.62e+01, dt: 7.00e-02, iterations: 1
Progress: 38%, Time: 5.32e+01, dt: 7.00e-02, iterations: 1
Progress: 43%, Time: 6.02e+01, dt: 7.00e-02, iterations: 1
Progress: 48%, Time: 6.73e+01, dt: 7.00e-02, iterations: 1
Progress: 53%, Time: 7.43e+01, dt: 7.00e-02, iterations: 1
Progress: 58%, Time: 8.13e+01, dt: 7.00e-02, iterations: 1
Progress: 63%, Time: 8.83e+01, dt: 7.00e-02, iterations: 1
Progress: 68%, Time: 9.53e+01, dt: 7.00e-02, iterations: 1
Progress: 73%, Time: 1.02e+02, dt: 7.00e-02, iterations: 1
Progress: 78%, Time: 1.09e+02, dt: 7.00e-02, iterations: 1
Progress: 83%, Time: 1.16e+02, dt: 7.00e-02, iterations: 1
Progress: 88%, Time: 1.23e+02, dt: 7.00e-02, iterations: 1
Progress: 93%, Time: 1.30e+02, dt: 7.00e-02, iterations: 1
Progress: 98%, Time: 1.37e+02, dt: 7.00e-02, iterations: 1

Here, we extract the data accumulated during time stepping and visualize it.

# Extract data
print(result.q)
history = result.history
u_sol = history.q['u']
v_sol = history.q['v']
ts = history.t

# Plot the results
import matplotlib.pyplot as plt
plt.plot(ts, u_sol, label="Prey", marker="o")
plt.plot(ts, v_sol, label="Predator", marker="o")
plt.xlabel("Time")
plt.ylabel("Population")
plt.title("Lotka-Volterra system")
plt.legend()
plt.grid(True)
plt.show()

{'u': Array([10.12422874], dtype=float64), 'v': Array([10.24250256], dtype=float64)}