from .solver import Solver
from ..typing import arrf64
from ..core import ti_or_fallback as ti
from numpy import zeros
from typing import Callable
[docs]
@ti.data_oriented
class RK4_38(Solver):
"""
An explicit RK4 integrator following the 3/8 rule for solving systems of
ordinary differential equations of the form
.. math::
\\dfrac{ds}{dt} = f(t, s).
The method is a simple first-order integrator that updates the state of the
system at each time step as
.. math::
k_1 &= f(t_n, s_n) \\\\
k_2 &= f\\left(t_n + \\dfrac{\\Delta t}{3}, s_n + \\dfrac{k_1}{3}\\Delta t\\right) \\\\
k_3 &= f\\left(t_n + \\dfrac{2\\Delta t}{3}, s_n + (- \\dfrac{k_1}{3} + k_2) \\Delta t\\right) \\\\
k_4 &= f\\left(t_n + \\Delta t, s_n + (k_1 - k_2 + k_3) \\Delta t\\right) \\\\
s_{n+1} &= s_n + \\dfrac{\\Delta t}{8} (k_1 + 3 k_2 + 3 k_3 + k_4)
The integrator supports both pure Python and Taichi-accelerated
implementations (as defined by the :class:`vivyd.models.TaichiCompatible`
and :class:`vivyd.models.TaichiConvertible` protocols).
"""
def _integrate_python(
self,
model : Callable,
t_tab : arrf64,
state0 : arrf64,
handle : Callable
) -> arrf64:
n = len(t_tab)
out = zeros((n, *state0.shape), dtype=float)
out[0] = state0
progress = 0.0
for i in range(1, n):
if self.verbose:
p = (i+1) / n
if p >= progress + 0.01:
progress = p
print(f"\rProgress: {int(progress*100)} %", end="")
s = out[i-1]
t = t_tab[i-1]
dt = t_tab[i] - t
k1 = model(t , s , handle=handle)
k2 = model(t + 1.0/3.0 * dt, s + (k1/3.0) * dt, handle=handle)
k3 = model(t + 2.0/3.0 * dt, s + (- k1/3.0 + k2) * dt, handle=handle)
k4 = model(t + dt, s + (k1 - k2 + k3) * dt, handle=handle)
ds = (k1 + 3.0*k2 + 3.0*k3 + k4) / 8.0
out[i] = s + ds * dt
if self.verbose:
print("\rProgress: 100 %")
return out
def _make_buffer(
self,
state0: arrf64,
t_tab: arrf64
) -> arrf64:
return zeros(state0.shape[0], dtype=float)
@ti.func
def _integrate_taichi_func(
self,
model : ti.template(), # type: ignore
t_tab : ti.types.ndarray(dtype=float, ndim=1), # type: ignore
state0: ti.types.ndarray(dtype=float, ndim=1), # type: ignore
handle: ti.template(), # type: ignore
out : ti.types.ndarray(dtype=float, ndim=2), # type: ignore
):
n_state = ti.static(model.state_size)
n_time = t_tab.shape[0]
s = ti.Vector.zero(float, n_state)
progress = 0.0
for j in range(n_state):
out[0, j] = state0[j]
ti.loop_config(serialize=True)
for i in range(1, n_time):
if ti.static(self.verbose):
p = (i+1) / n_time
if p >= progress + 0.01:
progress = p
print(f"\rProgress: {int(progress*100)} %", end="")
t = t_tab[i-1]
dt = t_tab[i] - t
for j in range(n_state):
s[j] = out[i - 1, j]
k1 = model._call_taichi_func(t , s , handle)
k2 = model._call_taichi_func(t + 1.0/3.0 * dt, s + (k1/3.0) * dt, handle)
k3 = model._call_taichi_func(t + 2.0/3.0 * dt, s + (- k1/3.0 + k2) * dt, handle)
k4 = model._call_taichi_func(t + dt, s + (k1 - k2 + k3) * dt, handle)
for j in range(n_state):
out[i, j] = out[i - 1, j] + (k1[j] + 3.0*k2[j] + 3.0*k3[j] + k4[j]) / 8.0 * dt
if ti.static(self.verbose):
print("\rProgress: 100 %")