from ..typing import arrf64
from ..core import ti_or_fallback as ti, is_taichi_used, TaichiCompatible, TaichiConvertible
from .compatibility import SolverCompatible
from abc import ABC, abstractmethod
from typing import Callable
from numpy import zeros
[docs]
@ti.data_oriented
class Solver(ABC):
"""
Abstract base class for numerical solvers of ordinary differential equations
(ODEs). The class defines the common interface and structure for all
solvers in the `vivyd.solvers` module, including both pure Python and
Taichi-accelerated implementations.
"""
def __init__(self, verbose: bool=False):
self.verbose = verbose
[docs]
def run(
self,
model : SolverCompatible,
state0: arrf64,
t_tab : arrf64,
handle: Callable | None = None
) -> arrf64:
"""
Run the integration process.
Parameters
----------
model: SolverCompatible
The right-hand side of the system of ODEs, which takes the current time and state as arguments.
state0: arrf64
The initial state of the system.
t_tab: arrf64
The array of time points used for integration.
handle: Callable | None
A function that is called inside the model. Default is `None`.
Returns
-------
arrf64:
An array containing the state of the system at each time point in
``t_tab``. Its shape is ``(len(t_tab), *state0.shape)``.
Raises
------
TypeError:
If Taichi acceleration is used but the model does not support Taichi
solvers.
"""
if is_taichi_used():
self._run = self._integrate_taichi
else:
self._run = self._integrate_python
if handle is None:
handle = lambda t, state: None
return self._run(model, t_tab, state0, handle)
@abstractmethod
def _integrate_python(
self,
model : SolverCompatible,
t_tab : arrf64,
state0 : arrf64,
handle : Callable
) -> arrf64: ...
def _integrate_taichi(
self,
model : SolverCompatible,
t_tab : arrf64,
state0 : arrf64,
handle : Callable
) -> arrf64:
if not isinstance(model, TaichiCompatible):
if not isinstance(model, TaichiConvertible):
raise TypeError("Model does not support Taichi solvers")
else:
taichi_model = model.to_taichi()
if isinstance(taichi_model, SolverCompatible):
model = taichi_model
else:
raise TypeError("Model's Taichi conversion does not return a SolverCompatible")
state0_flat = state0.flatten()
out = zeros((len(t_tab), len(state0_flat)), dtype=float)
self._integrate_taichi_kernel(model, t_tab, state0_flat, handle, out)
return out.reshape((len(t_tab), *state0.shape))
@ti.kernel
def _integrate_taichi_kernel(
self,
model : ti.template(), # type: ignore
t_tab : ti.types.ndarray(dtype=ti.f64, ndim=1), # type: ignore
state0: ti.types.ndarray(dtype=ti.f64, ndim=1), # type: ignore
handle: ti.template(), # type: ignore
out : ti.types.ndarray(dtype=ti.f64, ndim=2), # type: ignore
):
self._integrate_taichi_func(model, t_tab, state0, handle, out)
@abstractmethod
@ti.func
def _integrate_taichi_func(
self,
model : ti.template(), # type: ignore
t_tab : ti.types.ndarray(dtype=ti.f64, ndim=1), # type: ignore
state0: ti.types.ndarray(dtype=ti.f64, ndim=1), # type: ignore
handle: ti.template(), # type: ignore
out : ti.types.ndarray(dtype=ti.f64, ndim=2), # type: ignore
):
...