Source code for vivyd.excitation.excitation

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

from numpy import ascontiguousarray, asarray




@ti.data_oriented
class Excitation(TaichiCompatible):
    """
    A class for representing any signal that can be interpolated and used
    to affect models.
    
    Parameters
    ----------
    x : arrf64
        A 1D array containing the signal values at discrete time points.
    dt : float
        The time step between the discrete time points corresponding to the
        signal values in `x`.
        
    .. Important::
        To benefit from Taichi acceleration, ensure that the `Excitation` class is
        instantiated within a Taichi context.
    """
    
    def __init__(self, x: arrf64, dt: float):
        
        if x.ndim != 1:
            raise ValueError("x must be a 1D array")
                
        if is_taichi_used():
            self._x = ti.field(dtype=float, shape=len(x))
            self._x.from_numpy(ascontiguousarray(x))
        else:
            self._x = x
                
        self._dt   = dt
        self.n     = len(x)
    
    @property
    def x(self):
        return asarray(self._x)     # Numpy tries to generate a view, avoiding unnecessary copying when possible.



    def interpolate(self, t: float) -> float:
        """
        Interpolate the signal at a given time point.
        
        Parameters
        ----------
        t : float
            A time point at which to interpolate the signal.
        
        Returns
        -------
        float
            The interpolated signal value at the specified time point.
            specified time points.
        
        Raises
        ------
        ValueError
            If any time point in `t` is outside the range [0, (n-1)*dt], where
            `n` is the number of discrete time points in `x`.
        """
        
        if t < 0 or t > (self.n - 1) * self._dt:
            raise ValueError("t must be within the range [0, (n-1)*dt]")
        
        if is_taichi_used():
            _call = self._call_taichi
        else:
            _call = self._call_python
        
        return _call(t)

                


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

    
    def _call_python(self, t: float) -> float:
        
        i = t / self._dt
        
        if i == int(i):
            return self.x[int(i)]
        
        else:
            i_inf = int(i)
            i_sup = i_inf + 1
            
            x_inf = self.x[i_inf]
            x_sup = self.x[i_sup]
            
            t_mid = t - i_inf * self._dt
            return x_inf + (x_sup - x_inf) * t_mid / self._dt
    
    @ti.func
    def _call_taichi_func(
        self,
        t: float
    ) -> float:
        
        i_true = t / self._dt
        i_inf  = int(i_true)
        i_sup  = i_inf + 1
        
        return ti.select(
            ti.abs(i_true - ti.cast(i_inf, ti.f64)) < 1e-12,
            self._x[i_inf],
            self._x[i_inf] + (self._x[i_sup] - self._x[i_inf]) * (t - i_inf * self._dt) / self._dt
        )
                    
    @ti.kernel
    def _call_taichi_kernel(
        self,
        t: float
    ) -> float:
        return self._call_taichi_func(t)
        
    def _call_taichi(self, t: float) -> float:
        return self._call_taichi_kernel(t)