Source code for lv_sens.ode

r"""Adjoint method for the Lotka-Volterra ODE (predator-prey) using analytic model
derivatives.

Forward ODE:

.. math::

    \frac{dx}{dt} = \alpha x - \beta xy \\
    \frac{dy}{dt} = xy - \delta y
"""

import numpy as np
from numba import njit  # type: ignore[import-untyped]

from lv_sens.types import NDArray_f64


[docs] @njit # type: ignore def LV(x: NDArray_f64, theta: NDArray_f64, dt: float) -> NDArray_f64: r"""Propagator of Lotka-Volterra system discretized with semi-implicit Euler method. Continuous ODE: .. math:: \frac{dx}{dt} = \alpha x - \beta xy \\ \frac{dy}{dt} = xy - \delta y Time-discretized 'propagator' form: .. math:: u_{n+1} = \mathcal{N}(u_n) with the state :math:`u_n = (x_n^T,\ y_n^T)` and .. math:: \mathcal{N}(u_n) = \left[I - \tau \begin{pmatrix}\alpha & -\beta x\\y & -\delta\end{pmatrix}\right]^{-1} u_n Args: x: current state at time t theta: parameter dt: time step Returns: state at the next time step """ assert len(theta) == 3 A = np.identity(2) - dt * np.array( [[theta[0], -theta[1] * x[0]], [x[1], -theta[2]]] ) x = np.linalg.solve(A, x) return x
[docs] @njit # type: ignore def LV_dx(x: NDArray_f64, theta: NDArray_f64, dt: float) -> NDArray_f64: """Partial derivative of ODE propagation operator with respect to the state, x. Args: x: current state at time t theta: parameter dt: time step Returns: derivative of function evaluated in (x, theta, t) """ alpha, beta, delta = theta[0], theta[1], theta[2] y = x[1] x = x[0] N_dx = np.array( [ [ -beta * dt**2 * x * y * (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 + (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ], [ -beta * dt**2 * y**2 * (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 + dt * y / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ], ] ) N_dy = np.array( [ [ -beta * dt**2 * x**2 * (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 - beta * dt * x / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ], [ -beta * dt**2 * x * y * (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 + (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ], ] ) return np.hstack((N_dx, N_dy))
[docs] @njit # type: ignore def LV_da(x: NDArray_f64, theta: NDArray_f64, dt: float) -> NDArray_f64: """Partial derivative of ODE propagation operator with respect to the parameter, theta. Args: x: current state at time t theta: parameter dt: time step Returns: derivative of function evaluated in (x, theta, t) """ alpha, beta, delta = theta[0], theta[1], theta[2] y = x[1] x = x[0] N_a1 = np.array( [ [ x * (delta * dt**2 + dt) * (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 ], [ -dt * y / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) + y * (delta * dt**2 + dt) * (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 ], ] ) N_a2 = np.array( [ [ -(dt**2) * x**2 * y * (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 - dt * x * y / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ], [ -(dt**2) * x * y**2 * (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 ], ] ) N_a3 = np.array( [ [ dt * x / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) + x * (alpha * dt**2 - dt) * (-beta * dt * y + delta * dt + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 ], [ y * (alpha * dt**2 - dt) * (-alpha * dt + dt * x + 1) / ( -alpha * delta * dt**2 - alpha * dt + beta * dt**2 * x * y + delta * dt + 1 ) ** 2 ], ] ) return np.hstack((N_a1, N_a2, N_a3))
[docs] @njit # type: ignore def forward(theta: float, dt: float, N: int, x0: list[float]) -> NDArray_f64: """Solve ODE. Args: theta: parameter dt: time step N: number of timesteps x0: initial condition Returns: state trajectory (all timesteps) """ x = np.zeros((len(x0), N)) x[:, 0] = x0 for n in range(1, N): x[:, n] = LV(x[:, n - 1], theta, dt) return x