import numpy as np
import numpy.typing
import scipy # type: ignore[import-untyped]
from numba import njit # type: ignore[import-untyped]
from lv_sens import ode
from lv_sens.types import NDArray_f64
[docs]
@njit # type: ignore
def J(
theta: NDArray_f64, xdata: NDArray_f64, dt: float, N: int, x0: NDArray_f64
) -> float:
"""Compute cost function for a specific parameter a by first solving the forward problem.
Args:
theta: parameter
xdata: "measurement" data
dt: time step
N: number of timesteps
x0: initial condition of forward solve
Returns:
cost function value
"""
x = ode.forward(theta, dt, N, x0)
cost: float = 0.5 * np.sum(np.square(x - xdata))
return cost
[docs]
@njit # type: ignore
def grad_fd(
theta: NDArray_f64,
xdata: NDArray_f64,
dt: float,
N: int,
x0: NDArray_f64,
h: float = 1e-7,
) -> NDArray_f64:
"""Approximate the cost function gradient with finite differences.
Args:
theta: parameter
xdata: "measurement" data
dt: time step size
N: number of timesteps
x0: initial condition of forward solve
h: finite differente step size (parameter)
Returns:
cost function gradient approximation
"""
dx = np.identity(len(theta)) * h
dJ = np.array(
[
(J(theta + dx[:, i], xdata, dt, N, x0) - J(theta, xdata, dt, N, x0)) / h
for i in range(len(theta))
]
)
return dJ
[docs]
@njit # type: ignore
def grad_adj(
theta: NDArray_f64, xdata: NDArray_f64, dt: float, N: int, x0: NDArray_f64
) -> NDArray_f64:
"""Compute the cost function gradient via the adjoint.
Args:
theta: parameter
xdata: 'measurement' data
dt: time step size
N: number of timesteps
x0: initial condition of forward solve
Returns:
cost function gradient
"""
x = ode.forward(theta, dt, N, x0)
# partial derivative of discrepancy function j_n(x_n) w.r.t. x_n
djndx = -(xdata - x)
# adjoint state
adj = np.zeros_like(x)
adj[:, -1] = -djndx[:, -1]
for i in range(N - 2, 0, -1):
adj[:, i] = adj[:, i + 1] @ ode.LV_dx(x[:, i], theta, dt) - djndx[:, i]
grad_J = np.zeros(len(theta))
for i in range(N - 1):
grad_J += -adj[:, i + 1] @ ode.LV_da(x[:, i], theta, dt)
# grad_J = -np.inner(adj[:, 1:], LV_da(x[:, :-1], theta, dt))
return grad_J
[docs]
def optimize(
theta0: NDArray_f64,
x_data: NDArray_f64,
dt: float,
N: int,
x0: NDArray_f64,
grad_type: str,
) -> scipy.optimize.OptimizeResult:
r"""Optimize the inverse problem
.. math::
\hat\theta = \textrm{argmin}\, J(\theta)
for cost function
.. math::
J(\theta) = \frac{1}{2}\sum_{i=1}^N |x_i(\theta) - y_i|
where :math:`x` is the solution of the ODE :py:func:`~.ode.LV`
at time step :math:`i` and :math:`y_i` is a noisy measurement of :math:`x_i`.
Args:
theta0: parameter initial guess
x_data: data array
dt: time step
N: number of time steps
x0: initial guess
grad_type: gradient type of the cost function
Returns:
scipy result object
"""
if grad_type == "adjoint":
jac = grad_adj
elif grad_type == "fd":
jac = grad_fd
else:
jac = None
bfgs_gtol = 1e-3
args = (x_data, dt, N, x0)
res = scipy.optimize.minimize(
J,
theta0,
args=args,
method="BFGS",
jac=jac,
options={"disp": True, "gtol": bfgs_gtol},
)
return res