Source code for vivyd.excitation.excitation_group

from ..core import is_taichi_used, ti_or_fallback as ti, TaichiCompatible
from ..typing import arrf64
from .excitation import Excitation

from numpy import asarray, zeros




@ti.data_oriented
class ExcitationGroup(TaichiCompatible):
    """
    A class for representing a group of excitations. More importantly, this
    class provides a handle comprising all the excitations in the group, which
    can then be passed to the solver in order to have better control of the
    model's excitations.
    
    Parameters
    ----------
    *excitations : Excitation
        One or more `Excitation` objects representing the group of excitations.
    """
    
    def __init__(self, *excitations: Excitation):
        
        self._excitations = tuple(excitations)
        self._n = len(self._excitations)
        


    def __call__(self, t: float) -> arrf64:
        """Shorthand for :meth:`interpolate`."""
        return self.interpolate(t)

    


    def interpolate(self, t: float) -> arrf64:
        """
        Interpolate the excitations in the group at a given time `t`.
        
        Parameters
        ----------
        t : float
            The time at which to interpolate the excitations.
        
        Returns
        -------
        arrf64
            The interpolated excitations at time `t`. Each element in the array
            is the interpolated value of the corresponding excitation in the
            group.
        """
        
        if is_taichi_used():
            _call = self._call_taichi
        else:
            _call = self._call_python
        return _call(t)

    
    def _call_python(self, t: float) -> arrf64:
        return asarray([excitation(t) for excitation in self._excitations])
    
    def _call_taichi(self, t: float) -> arrf64:
        out = zeros(self._n, dtype=float)
        self._call_taichi_kernel(t, out)
        return out

    @ti.func
    def _call_taichi_func(self, t: float):
        
        n = ti.static(self._n)

        out = ti.Vector.zero(ti.f64, n)

        for i in ti.static(range(n)):
            out[i] = self._excitations[i]._call_taichi_func(t)

        return out

    @ti.kernel
    def _call_taichi_kernel(
        self,
        t: float,
        out: ti.types.ndarray(dtype=ti.f64, ndim=1)     # type: ignore
    ):
        out_ = self._call_taichi_func(t)
        for i in range(ti.static(self._n)):
            out[i] = out_[i]
            
    @property
    def handle(self):
        """Get the handle for this excitation group, which can be passed to the solver."""
        return self._call_taichi_func if is_taichi_used() else self._call_python