# Source code for qgs.integrators.integrate

"""
Integrate module
================

Module with the function to integrate the ordinary differential equations

.. math:: \\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x})

of the model and its linearized version.

Description of the module functions
-----------------------------------

Two main functions:

* :obj:integrate_runge_kutta
* :obj:integrate_runge_kutta_tgls

"""

# TODO : - test the provided functions with try before proceeding

from numba import njit
import numpy as np
from qgs.functions.util import reverse

[docs]def integrate_runge_kutta(f, t0, t, dt, ic=None, forward=True, write_steps=1, b=None, c=None, a=None):
"""
Integrate the ordinary differential equations (ODEs)

.. math:: \\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x})

with a specified Runge-Kutta method_. The function :math:\\boldsymbol{f} should
be a Numba_ jitted function. This function must have a signature f(t, x) where x is
the state value and t is the time.

.. _Runge-Kutta method: https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
.. _Numba: https://numba.pydata.org/

Parameters
----------
f: callable
The Numba_-jitted function :math:\\boldsymbol{f}.
Should have the signaturef(t, x) where x is the state value and t is the time.
t0: float
Initial time of the time integration. Corresponds to the initial condition.
Important if the ODEs are non-autonomous.
t: float
Final time of the time integration. Corresponds to the final condition.
Important if the ODEs are non-autonomous.
dt: float
Timestep of the integration.
ic: None or ~numpy.ndarray(float), optional
Initial (or final) conditions of the system. Can be a 1D or a 2D array:

* 1D: Provide a single initial condition.
Should be of shape (n_dim,) where n_dim = :math:\\mathrm{dim}(\\boldsymbol{x}).
* 2D: Provide an ensemble of initial condition.
Should be of shape (n_traj, n_dim) where n_dim = :math:\\mathrm{dim}(\\boldsymbol{x}),
and where n_traj is the number of initial conditions.

If None, use a zero initial condition. Default to None.
If the forward argument is False, it specifies final conditions.
forward: bool, optional
Whether to integrate the ODEs forward or backward in time. In case of backward integration, the
initial condition ic becomes a final condition. Default to forward integration.
write_steps: int, optional
Save the state of the integration in memory every write_steps steps. The other intermediary
steps are lost. It determines the size of the returned objects. Default is 1.
Set to 0 to return only the final state.
b: None or ~numpy.ndarray, optional
Vector of coefficients :math:b_i of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.
c: None or ~numpy.ndarray, optional
Matrix of coefficients :math:c_{i,j} of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.
a: None or ~numpy.ndarray, optional
Vector of coefficients :math:a_i of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.

Returns
-------
time, traj: ~numpy.ndarray
The result of the integration:

* **time:** Time at which the state of the system was saved. Array of shape (n_step,) where
n_step is the number of saved states of the integration.
* **traj:** Saved dynamical system states. 3D array of shape (n_traj, n_dim, n_steps). If n_traj = 1,
a 2D array of shape (n_dim, n_steps) is returned instead.

Examples
--------

>>> from numba import njit
>>> import numpy as np
>>> from qgs.integrators.integrate import integrate_runge_kutta
>>> a = 0.25
>>> F = 16.
>>> G = 3.
>>> b = 6.
>>> # Lorenz 84 example
>>> @njit
... def fL84(t, x):
...     xx = -x[1] ** 2 - x[2] ** 2 - a * x[0] + a * F
...     yy = x[0] * x[1] - b * x[0] * x[2] - x[1] + G
...     zz = b * x[0] * x[1] + x[0] * x[2] - x[2]
...     return np.array([xx, yy, zz])
>>> # no ic
>>> # write_steps is 1 by default
>>> tt, traj = integrate_runge_kutta(fL84, t0=0., t=10., dt=0.1)  # 101 steps
>>> print(traj.shape)
(3, 101)
>>> # 1 ic
>>> ic = 0.1 * np.random.randn(3)
>>> tt, traj = integrate_runge_kutta(fL84, t0=0., t=10., dt=0.1, ic=ic)  # 101 steps
>>> print(ic.shape)
(3,)
>>> print(traj.shape)
(3, 101)
>>> # 4 ic
>>> ic = 0.1 * np.random.randn(4, 3)
>>> tt, traj = integrate_runge_kutta(fL84, t0=0., t=10., dt=0.1, ic=ic)  # 101 steps
>>> print(ic.shape)
(4, 3)
>>> print(traj.shape)
(4, 3, 101)
"""

if ic is None:
i = 1
while True:
ic = np.zeros(i)
try:
x = f(0., ic)
except:
i += 1
else:
break

i = len(f(0., ic))
ic = np.zeros(i)

if len(ic.shape) == 1:
ic = ic.reshape((1, -1))

# Default is RK4
if a is None and b is None and c is None:
c = np.array([0., 0.5, 0.5, 1.])
b = np.array([1./6, 1./3, 1./3, 1./6])
a = np.zeros((len(c), len(b)))
a[1, 0] = 0.5
a[2, 1] = 0.5
a[3, 2] = 1.

if forward:
time_direction = 1
else:
time_direction = -1

time = np.concatenate((np.arange(t0, t, dt), np.full((1,), t)))

recorded_traj = _integrate_runge_kutta_jit(f, time, ic, time_direction, write_steps, b, c, a)

if write_steps > 0:
if forward:
if time[::write_steps][-1] == time[-1]:
return time[::write_steps], np.squeeze(recorded_traj)
else:
return np.concatenate((time[::write_steps], np.full((1,), t))), np.squeeze(recorded_traj)
else:
rtime = reverse(time[::-write_steps])
if rtime[0] == time[0]:
return rtime, np.squeeze(recorded_traj)
else:
return np.concatenate((np.full((1,), t0), rtime)), np.squeeze(recorded_traj)
else:
return time[-1], np.squeeze(recorded_traj)

@njit
def _integrate_runge_kutta_jit(f, time, ic, time_direction, write_steps, b, c, a):

n_traj = ic.shape[0]
n_dim = ic.shape[1]

s = len(b)

if write_steps == 0:
n_records = 1
else:
tot = time[::write_steps]
n_records = len(tot)
if tot[-1] != time[-1]:
n_records += 1

recorded_traj = np.zeros((n_traj, n_dim, n_records))
if time_direction == -1:
directed_time = reverse(time)
else:
directed_time = time

for i_traj in range(n_traj):
y = ic[i_traj].copy()
k = np.zeros((s, n_dim))
iw = 0
for ti, (tt, dt) in enumerate(zip(directed_time[:-1], np.diff(directed_time))):

if write_steps > 0 and np.mod(ti, write_steps) == 0:
recorded_traj[i_traj, :, iw] = y
iw += 1

k.fill(0.)
for i in range(s):
y_s = y + dt * a[i] @ k
k[i] = f(tt + c[i] * dt, y_s)
y_new = y + dt * b @ k
y = y_new

recorded_traj[i_traj, :, -1] = y

return recorded_traj[:, :, ::time_direction]

@njit
def _tangent_linear_system(fjac, t, xs, x, adjoint):
return fjac(t, xs).transpose() @ x
else:
return fjac(t, xs) @ x

# a function that return always zero
@njit
def _zeros_func(t, x):
return np.zeros_like(x)

[docs]def integrate_runge_kutta_tgls(f, fjac, t0, t, dt, ic=None, tg_ic=None,
forward=True, adjoint=False, inverse=False, boundary=None,
write_steps=1, b=None, c=None, a=None):
"""Integrate simultaneously the ordinary differential equations (ODEs)

.. math:: \dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x})

and its tangent linear model, i.e. the linearized ODEs

.. math :: \dot{\\boldsymbol{\delta x}} = \\boldsymbol{\mathrm{J}}(t, \\boldsymbol{x}) \cdot \\boldsymbol{\delta x}

where :math:\\boldsymbol{\mathrm{J}} = \\frac{\partial \\boldsymbol{f}}{\partial \\boldsymbol{x}} is the
Jacobian matrix of :math:\\boldsymbol{f}, with a specified Runge-Kutta method_.
To solve this equation, one has to integrate simultaneously both ODEs.

The function :math:\\boldsymbol{f} and :math:\\boldsymbol{J} should
be Numba_ jitted functions. These functions must have a signature f(t, x) and J(t, x) where x is
the state value and t is the time.

.. _Runge-Kutta method: https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
.. _Numba: https://numba.pydata.org/
.. _fundamental matrix of solutions: https://en.wikipedia.org/wiki/Fundamental_matrix_(linear_differential_equation)

Parameters
----------
f: callable
The Numba_-jitted function :math:\\boldsymbol{f}.
Should have the signaturef(t, x) where x is the state value and t is the time.
fjac: callable
The Numba_-jitted Jacobian :math:\\boldsymbol{J}.
Should have the signatureJ(t, x) where x is the state value and t is the time.
t0: float
Initial time of the time integration. Corresponds to the initial conditions.
t: float
Final time of the time integration. Corresponds to the final conditions.
dt: float
Timestep of the integration.
ic: None or ~numpy.ndarray(float), optional
Initial (or final) conditions of the ODEs :math:\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x}).
Can be a 1D or a 2D array:

* 1D: Provide a single initial condition.
Should be of shape (n_dim,) where n_dim = :math:\mathrm{dim}(\\boldsymbol{x}).
* 2D: Provide an ensemble of initial condition.
Should be of shape (n_traj, n_dim) where n_dim = :math:\mathrm{dim}(\\boldsymbol{x}),
and where n_traj is the number of initial conditions.

If None, use a zero initial condition. Default to None.
If the forward argument is False, it specifies final conditions.
tg_ic: None or ~numpy.ndarray(float), optional
Initial (or final) conditions of the linear ODEs
:math:\dot{\\boldsymbol{\delta x}} = \\boldsymbol{\mathrm{J}}(t, \\boldsymbol{x}) \cdot \\boldsymbol{\delta x}. \n
Can be a 1D, a 2D or a 3D array:

* 1D: Provide a single initial condition. This initial condition of the linear ODEs will be the same used for each
initial condition ic of the ODEs :math:\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x})
Should be of shape (n_dim,) where n_dim = :math:\mathrm{dim}(\\boldsymbol{x}).
* 2D: Two sub-cases:

+ If tg_ic.shape[0]=ic.shape[0], assumes that each initial condition ic[i] of :math:\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x}),
correspond to a different initial condition tg_ic[i].
+ Else, assumes and integrate an ensemble of n_tg_traj initial condition of the linear ODEs for each
initial condition of :math:\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x}).
* 3D: An array of shape (n_traj, n_tg_traj, n_dim) which provide an ensemble of n_tg_ic initial conditions
specific to each of the n_traj initial conditions of :math:\dot{\\boldsymbol{x}} = \\boldsymbol{f}(t, \\boldsymbol{x}).

If None, use the identity matrix as initial condition, returning the fundamental matrix of solutions_ of the
linear ODEs.
Default to None.
If the forward argument is False, it specifies final conditions.
forward: bool, optional
Whether to integrate the ODEs forward or backward in time. In case of backward integration, the
initial condition ic becomes a final condition. Default to forward integration.
Wheter to integrate the tangent :math:\dot{\\boldsymbol{\delta x}} = \\boldsymbol{\mathrm{J}}(t, \\boldsymbol{x}) \cdot \\boldsymbol{\delta x}
or the adjoint linear model :math:\dot{\\boldsymbol{\delta x}} = \\boldsymbol{\mathrm{J}}^T(t, \\boldsymbol{x}) \cdot \\boldsymbol{\delta x}.
Integrate the tangent model by default.
inverse: bool, optional
Wheter or not to invert the Jacobian matrix
:math:\\boldsymbol{\mathrm{J}}(t, \\boldsymbol{x}) \\rightarrow \\boldsymbol{\mathrm{J}}^{-1}(t, \\boldsymbol{x}).
False by default.
boundary: None or callable, optional
Allow to add a inhomogeneous term to linear ODEs:
:math:\dot{\\boldsymbol{\delta x}} = \\boldsymbol{\mathrm{J}}(t, \\boldsymbol{x}) \cdot \\boldsymbol{\delta x} + \Psi(t, \\boldsymbol{x}).
The boundary :math:\Psi should have the same signature as :math:\\boldsymbol{\mathrm{J}}, i.e.  func(t, x).
If None, don't add anything (homogeneous case). None by default.
write_steps: int, optional
Save the state of the integration in memory every write_steps steps. The other intermediary
steps are lost. It determines the size of the returned objects. Default is 1.
Set to 0 to return only the final state.
b: None or ~numpy.ndarray, optional
Vector of coefficients :math:b_i of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.
c: None or ~numpy.ndarray, optional
Matrix of coefficients :math:c_{i,j} of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.
a: None or ~numpy.ndarray, optional
Vector of coefficients :math:a_i of the Runge-Kutta method_ .
If None, use the classic RK4 method coefficients. Default to None.

Returns
-------
time, traj, tg_traj: ~numpy.ndarray
The result of the integration:

* **time:** Time at which the state of the system was saved. Array of shape (n_step,) where
n_step is the number of saved states of the integration.
* **traj:** Saved states of the ODEs. 3D array of shape (n_traj, n_dim, n_steps). If n_traj = 1,
a 2D array of shape (n_dim, n_steps) is returned instead.
* **tg_traj:** Saved states of the linear ODEs.
Depending on the input initial conditions of both ODEs,
it is at maximum a 4D array of shape (n_traj, n_tg_traj n_dim, n_steps).
If one of the dimension is 1, it is squeezed.

Examples
--------

>>> from numba import njit
>>> import numpy as np
>>> from qgs.integrators.integrate import integrate_runge_kutta_tgls
>>> a = 0.25
>>> F = 16.
>>> G = 3.
>>> b = 6.
>>> # Lorenz 84 example
>>> @njit
... def fL84(t, x):
...     xx = -x[1] ** 2 - x[2] ** 2 - a * x[0] + a * F
...     yy = x[0] * x[1] - b * x[0] * x[2] - x[1] + G
...     zz = b * x[0] * x[1] + x[0] * x[2] - x[2]
...     return np.array([xx, yy, zz])
>>> @njit
... def DfL84(t, x):
...    return np.array([[     -a        , -2. * x[1], -2. * x[2]],
...                      [x[1] - b * x[2], -1. + x[0], -b * x[0]],
...                      [b * x[1] + x[2],  b * x[0], -1. + x[0]]])
>>> # 4 ic, no tg_ic (fundamental matrix computation of an ensemble of ic)
>>> ic = 0.1 * np.random.randn(4, 3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, write_steps=1) # 101 steps
>>> print(ic.shape)
(4, 3)
>>> print(traj.shape)
(4, 3, 101)
>>> print(dtraj.shape)
(4, 3, 3, 101)
>>> # 1 ic, 1 tg_ic (one ic and its tg_ic)
>>> ic = 0.1 * np.random.randn(3)
>>> tg_ic = 0.001 * np.random.randn(3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, tg_ic=tg_ic) # 101 steps
>>> print(ic.shape)
(3,)
>>> print(tg_ic.shape)
(3,)
>>> print(traj.shape)
(3, 101)
>>> print(dtraj.shape)
(3, 101)
>>> # 4 ic, 1 same tg_ic (an ensemble of ic with the same tg_ic)
>>> ic = 0.1 * np.random.randn(4, 3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, tg_ic=tg_ic) # 101 steps
>>> print(ic.shape)
(4, 3)
>>> print(tg_ic.shape)
(3,)
>>> print(traj.shape)
(4, 3, 101)
>>> print(dtraj.shape)
(4, 3, 101)
>>> # 1 ic, 4 tg_ic (an ic with an ensemble of tg_ic in its tangent space)
>>> ic = 0.1 * np.random.randn(3)
>>> tg_ic = 0.001 * np.random.randn(4, 3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, tg_ic=tg_ic) # 101 steps
>>> print(ic.shape)
(3,)
>>> print(tg_ic.shape)
(4, 3)
>>> print(traj.shape)
(3, 101)
>>> print(dtraj.shape)
(4, 3, 101)
>>> # 2 ic, same 4 tg_ic (an ensemble of 2 ic, both with the same ensemble
>>> # of tg_ic in their tangent space)
>>> ic = 0.1 * np.random.randn(2, 3)
>>> tg_ic = 0.001 * np.random.randn(4, 3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, tg_ic=tg_ic) # 101 steps
>>> print(ic.shape)
(2, 3)
>>> print(tg_ic.shape)
(4, 3)
>>> print(traj.shape)
(2, 3, 101)
>>> print(dtraj.shape)
(2, 4, 3, 101)
>>> # 2 ic, 4 different tg_ic (an ensemble of 2 ic, with different ensemble
>>> # of tg_ic in their tangent space)
>>> ic = 0.1 * np.random.randn(2, 3)
>>> tg_ic = 0.001 * np.random.randn(2, 4, 3)
>>> tt, traj, dtraj = integrate_runge_kutta_tgls(fL84, DfL84, t0=0., t=10., dt=0.1,
...                                              ic=ic, tg_ic=tg_ic) # 101 steps
>>> print(ic.shape)
(2, 3)
>>> print(tg_ic.shape)
(2, 4, 3)
>>> print(traj.shape)
(2, 3, 101)
>>> print(dtraj.shape)
(2, 4, 3, 101)

"""

if ic is None:
i = 1
while True:
ic = np.zeros(i)
try:
x = f(0., ic)
except:
i += 1
else:
break

i = len(f(0., ic))
ic = np.zeros(i)

if len(ic.shape) == 1:
ic = ic.reshape((1, -1))

n_traj = ic.shape[0]

if tg_ic is None:
tg_ic = np.eye(ic.shape[1])

tg_ic_sav = tg_ic.copy()

if len(tg_ic.shape) == 1:
tg_ic = tg_ic.reshape((1, -1, 1))
ict = tg_ic.copy()
for i in range(n_traj-1):
ict = np.concatenate((ict, tg_ic))
tg_ic = ict
elif len(tg_ic.shape) == 2:
if tg_ic.shape[0] == n_traj:
tg_ic = tg_ic[..., np.newaxis]
else:
tg_ic = tg_ic[np.newaxis, ...]
tg_ic = np.swapaxes(tg_ic, 1, 2)
ict = tg_ic.copy()
for i in range(n_traj-1):
ict = np.concatenate((ict, tg_ic))
tg_ic = ict
elif len(tg_ic.shape) == 3:
if tg_ic.shape[1] != ic.shape[1]:
tg_ic = np.swapaxes(tg_ic, 1, 2)

# Default is RK4
if a is None and b is None and c is None:
c = np.array([0., 0.5, 0.5, 1.])
b = np.array([1./6, 1./3, 1./3, 1./6])
a = np.zeros((len(c), len(b)))
a[1, 0] = 0.5
a[2, 1] = 0.5
a[3, 2] = 1.

if forward:
time_direction = 1
else:
time_direction = -1

time = np.concatenate((np.arange(t0, t, dt), np.full((1,), t)))

if boundary is None:
boundary = _zeros_func

inv = 1.
if inverse:
inv *= -1.

recorded_traj, recorded_fmatrix = _integrate_runge_kutta_tgls_jit(f, fjac, time, ic, tg_ic,
time_direction, write_steps,
b, c, a, adjoint, inv, boundary)

if len(tg_ic_sav.shape) == 2:
if recorded_fmatrix.shape[1:3] != tg_ic_sav.shape:
recorded_fmatrix = np.swapaxes(recorded_fmatrix, 1, 2)

elif len(tg_ic_sav.shape) == 3:
if tg_ic_sav.shape[1] != ic.shape[1]:
if recorded_fmatrix.shape[:3] != tg_ic_sav.shape:
recorded_fmatrix = np.swapaxes(recorded_fmatrix, 1, 2)

if write_steps > 0:
if forward:
if time[::write_steps][-1] == time[-1]:
return time[::write_steps], np.squeeze(recorded_traj), np.squeeze(recorded_fmatrix)
else:
return np.concatenate((time[::write_steps], np.full((1,), t))), np.squeeze(recorded_traj),\
np.squeeze(recorded_fmatrix)
else:
rtime = reverse(time[::-write_steps])
if rtime[0] == time[0]:
return rtime, np.squeeze(recorded_traj), np.squeeze(recorded_fmatrix)
else:
return np.concatenate((np.full((1,), t0), rtime)), np.squeeze(recorded_traj),\
np.squeeze(recorded_fmatrix)

else:
return time[-1], np.squeeze(recorded_traj), np.squeeze(recorded_fmatrix)

@njit
def _integrate_runge_kutta_tgls_jit(f, fjac, time, ic, tg_ic, time_direction, write_steps, b, c, a,

n_traj = ic.shape[0]
n_dim = ic.shape[1]

s = len(b)

if write_steps == 0:
n_records = 1
else:
tot = time[::write_steps]
n_records = len(tot)
if tot[-1] != time[-1]:
n_records += 1
recorded_traj = np.zeros((n_traj, n_dim, n_records))
recorded_fmatrix = np.zeros((n_traj, tg_ic.shape[1], tg_ic.shape[2], n_records))
if time_direction == -1:
directed_time = reverse(time)
else:
directed_time = time

for i_traj in range(n_traj):
y = ic[i_traj].copy()
fm = tg_ic[i_traj].copy()
recorded_traj[i_traj, :, 0] = ic[i_traj]
recorded_fmatrix[i_traj, :, :, 0] = tg_ic[i_traj]
k = np.zeros((s, n_dim))
km = np.zeros((s, tg_ic.shape[1], tg_ic.shape[2]))
iw = 0
for ti, (tt, dt) in enumerate(zip(directed_time[:-1], np.diff(directed_time))):

if write_steps > 0 and np.mod(ti, write_steps) == 0:
recorded_traj[i_traj, :, iw] = y
recorded_fmatrix[i_traj, :, :, iw] = fm
iw += 1

k.fill(0.)
km.fill(0.)
for i in range(s):
y_s = y + dt * a[i] @ k
k[i] = f(tt + c[i] * dt, y_s)
km_s = fm.copy()
for j in range(len(a[i])):
km_s += dt * a[i, j] * km[j]
hom = inverse * _tangent_linear_system(fjac, tt + c[i] * dt, y_s, km_s, adjoint)
inhom = boundary(tt + c[i] * dt, y_s)
km[i] = (hom.T + inhom.T).T
y_new = y + dt * b @ k
fm_new = fm.copy()
for j in range(len(b)):
fm_new += dt * b[j] * km[j]
y = y_new
fm = fm_new

recorded_traj[i_traj, :, -1] = y
recorded_fmatrix[i_traj, :, :, -1] = fm

return recorded_traj[:, :, ::time_direction], recorded_fmatrix[:, :, :, ::time_direction]

if __name__ == "__main__":
import matplotlib.pyplot as plt
from scipy.integrate import odeint

@njit
def f(t, x):
return - np.array([1., 2., 3.]) * x

def fr(x, t):
return f(t, x)

ic = np.random.randn(6).reshape(2, 3)

time, r = integrate_runge_kutta(f, 0., 10., 0.01, ic=ic, write_steps=3)

t = np.arange(0., 10., 0.01)
t = np.concatenate((t[::3], np.full((1,), 10.)))
rl = list()
for i in range(ic.shape[0]):
rl.append(odeint(fr, ic[i], t).T)

plt.figure()
for i in range(ic.shape[0]):
p, = plt.plot(time, r[i, 0])
c = p.get_color()
plt.plot(t, rl[i][0], color=c, ls='--')
for j in range(1, ic.shape[1]):
p, = plt.plot(time, r[i, j], color=c)
plt.plot(t, rl[i][j], color=c, ls='--')

timet, rt = integrate_runge_kutta(f, 0., 10., 0.01, ic=ic, forward=False, write_steps=3)
rlt = list()
for i in range(ic.shape[0]):
rlt.append(odeint(fr, ic[i], reverse(t)).T)

plt.figure()
for i in range(ic.shape[0]):
p, = plt.plot(timet, rt[i, 0])
c = p.get_color()
plt.plot(t, reverse(rlt[i][0]), color=c, ls='--')
for j in range(1, ic.shape[1]):
p, = plt.plot(timet, rt[i, j], color=c)
plt.plot(t, reverse(rlt[i][j]), color=c, ls='--')

tt, re = integrate_runge_kutta(f, 0., 10., 0.01, ic=ic, write_steps=0)
print(tt)
print(r[0, :, -1], re[0])
plt.show(block=False)

a = 0.25
F = 16.
G = 3.
b = 6.

@njit
def DfL84(t, x):
return np.array([[     -a        , -2. * x[1], -2. * x[2]],
[x[1] - b * x[2], -1. + x[0], -b * x[0]],
[b * x[1] + x[2],  b * x[0], -1. + x[0]]])

@njit
def fL84(t, x):
xx = -x[1] ** 2 - x[2] ** 2 - a * x[0] + a * F
yy = x[0] * x[1] - b * x[0] * x[2] - x[1] + G
zz = b * x[0] * x[1] + x[0] * x[2] - x[2]
return np.array([xx, yy, zz])

def fL84r(x, t):
return fL84(t, x)

tt, ic_L84 = integrate_runge_kutta(fL84, 0., 10000., 0.01, write_steps=0)

print(ic_L84)

ttt1, irkt1, fm_irkt1 = integrate_runge_kutta_tgls(fL84, DfL84, 0., 0.01, 0.001, ic=ic_L84, write_steps=1)
bttt1, birkt1, bfm_irkt1 = integrate_runge_kutta_tgls(fL84, DfL84, 0., 0.01, 0.001, ic=irkt1[:, -1],
forward=False, write_steps=1, adjoint=True, inverse=True)

plt.figure()
for i in range(len(ic_L84)):
for j in range(len(ic_L84)):
p, = plt.plot(ttt1, fm_irkt1[i, j])
c = p.get_color()
plt.plot(bttt1, bfm_irkt1[i, j], ls='--', color=c)

plt.figure()
for i in range(len(ic_L84)):
p, = plt.plot(ttt1, irkt1[i])
c = p.get_color()
plt.plot(bttt1, birkt1[i], ls='--', color=c)

vec = np.random.randn(3)
vec = vec/np.linalg.norm(vec)
for i in range(1,13):
lam = 2.**(-i)

dy = lam * vec

tx, dx = integrate_runge_kutta(fL84, 0., 0.001, 0.001, ic=ic_L84, write_steps=0)
txp, dxp = integrate_runge_kutta(fL84, 0., 0.001, 0.001, ic=ic_L84+dy, write_steps=0)

dy1 = dxp - dx

txtl, dxtl, dytl = integrate_runge_kutta_tgls(fL84, DfL84, 0., 0.001, 0.001, ic=ic_L84, tg_ic=dy, write_steps=0)

print('Perturbation size: ', dy @ dy)
print('Resulting difference in trajectory: (eps ~ 2^-'+str(i)+')')
print('diff: ', dy1 @ dy1)
print('tl: ', dytl @ dytl)
print('ratio: ', (dy1 @ dy1) / (dytl @ dytl))
`