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]
def _integrate_runge_kutta_tgls_jit(f, fjac, time, ic, tg_ic, time_direction, write_steps, b, c, a, adjoint, inverse, boundary): 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]
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 signature``f(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's `ic` time. 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 condition 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`. 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 determine 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` is the 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` are the saved 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 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)
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 signature``f(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 signature``J(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's `ic` time. 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 condition 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`. tg_ic: None or ~numpy.ndarray(float), optional Initial condition 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`. 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. adjoint: bool, optional 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 determine 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` is the 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` are the 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` are the 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 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)
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)