Source code for vivyd.solvers.midpoint

from .solver import Solver
from ..core import ti_or_fallback as ti
from ..typing import arrf64

from numpy import zeros
from typing import Callable




@ti.data_oriented
class Midpoint(Solver):
    """
    A midpoint integrator for solving systems of ordinary differential
    equations of the form
    
    .. math::
        \\dfrac{ds}{dt} = f(t, s).
        
    The method is a second-order integrator that updates the state of the
    system at each time step as
    
    .. math::
        t_{n+1/2} &= \\dfrac{t_{n+1} + t_n}{2}, \\\\
        s_{n+1/2} &= s_n + f(t_n, s_n) \\cdot (t_{n+1/2} - t_n), \\\\
        s_{n+1}   &= s_n + f(t_{n+1/2}, s_{n+1/2}) \\cdot (t_{n+1} - t_n).
    
    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
            
            ds = model(t, s, handle=handle)
                        
            s_mid = s + ds * (dt / 2)
            t_mid = t + dt / 2
            
            ds_mid = model(t_mid, s_mid, handle=handle)
                        
            out[i] = s + ds_mid * 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)
        s_mid = 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="")
            
            dt = t_tab[i] - t_tab[i - 1]
            
            for j in range(n_state):
                s[j] = out[i - 1, j]
            
            rhs = model._call_taichi_func(t_tab[i - 1], s, handle)

            for j in range(n_state):
                s_mid[j] = out[i - 1, j] + rhs[j] * dt / 2.0

            t_mid = t_tab[i - 1] + dt / 2.0

            rhs_mid = model._call_taichi_func(t_mid, s_mid, handle)

            for j in range(n_state):
                out[i, j] = out[i - 1, j] + rhs_mid[j] * dt
                
        if ti.static(self.verbose):
            print("\rProgress: 100 %")