from ..core import ti_or_fallback as ti, is_taichi_used, TaichiCompatible, TaichiConvertible
from ..typing import arrf64
from .viv_model import VIVModel
from typing import Callable, ClassVar, Generic, Sequence, TypeVar, Protocol, runtime_checkable
from numpy import zeros_like, asarray
@runtime_checkable
class TaichiCompatibleInCollection(TaichiCompatible, Protocol):
"""A protocol for models that are compatible with Taichi and can be used in a collection of models."""
n_params: ClassVar[int]
"""Number of parameters required to define the model."""
@property
def params(self) -> Sequence:
"""Parameters required to define the model, in the order expected by the model constructor."""
...
T = TypeVar("T", bound=VIVModel) # All the models in the collection must be the same type of VIVModel
[docs]
@ti.data_oriented
class Collection(VIVModel, Generic[T]):
"""
A collection of uncoupled VIV models. The state of the collection is the
concatenation of the states of the individual models, and the right-hand
side of the collection is the concatenation of the RHS of the individual models.
.. Important::
As for now, a Collection must be **instantiated** inside a Taichi
context in order to benefit from Taichi acceleration.
Parameters
----------
models : Sequence[VIVModel]
A sequence of VIV models to be included in the collection. All models must be
of the same type.
"""
def __init__(self, models: Sequence[T]) -> None:
self.models = models
self.n_nodes = len(self.models)
if is_taichi_used():
self.models_taichi = []
for model in self.models:
if isinstance(model, TaichiCompatibleInCollection):
self.models_taichi.append(model)
elif isinstance(model, TaichiConvertible):
self.models_taichi.append(model.to_taichi())
else:
raise ValueError("The collection contains models that cannot be used with Taichi acceleration")
self.model_type_taichi = type(self.models_taichi[0])
self.params = ti.field(ti.f64, shape=(self.n_nodes, self.model_type_taichi.n_params))
for i, mt in enumerate(self.models_taichi):
params = mt.params
for j, param in enumerate(params):
self.params[i, j] = param
self._call = self._call_taichi
else:
self._call = self._call_python
self.state_size = self.models[0].state_size * self.n_nodes
def __getitem__(self, idx: int) -> T:
return self.models[idx]
[docs]
def rhs(self, t: float, state: arrf64, handle: Callable = (lambda *args, **kwargs: None)) -> arrf64:
"""
Compute the right-hand side of all the models in the collection.
Args:
t: Time.
state: Concatenated state array indexed of shape (n_nodes, state_size).
handle: A callable function for handling the model evaluation.
Returns:
arrf64: Concatenated right-hand side array indexed of shape (n_nodes, state_size).
"""
return self._call(t, self.validate_state(state), handle)
[docs]
def validate_state(self, state: arrf64) -> arrf64:
"""
Validate the shape of the state vector.
Args:
state: State vector.
Returns:
arrf64: Validated state vector.
Raises:
ValueError: If the state vector has an invalid shape.
"""
array = asarray(state, dtype=float)
if array.shape[0] != self.n_nodes:
raise ValueError(f"Incompatible state shape: expected first dimension {self.n_nodes}, got {array.shape[0]}")
self.models[0].validate_state(array[0, ...])
return array
def _call_python(self, t: float, state: arrf64, handle: Callable) -> arrf64:
diff = zeros_like(state)
for i, model in enumerate(self.models):
diff[i] = model(t, state[i], handle)
return diff
def _call_taichi(self, t: float, state: arrf64, handle: Callable) -> arrf64:
state_flat = state.flatten()
out = zeros_like(state_flat)
self._call_taichi_kernel(t, state_flat, handle, out)
return out.reshape(state.shape)
@ti.kernel
def _call_taichi_kernel(
self,
t : ti.f64, # type: ignore
state_flat: ti.types.ndarray(dtype=ti.f64, ndim=1), # type: ignore
handle : ti.template(), # type: ignore
out : ti.types.ndarray(dtype=ti.f64, ndim=1) # type: ignore
):
rhs = self._call_taichi_func(t, state_flat, handle)
for i in range(out.shape[0]):
out[i] = rhs[i]
@ti.func
def _call_taichi_func(
self,
t : ti.f64, # type: ignore
state_flat: ti.template(), # type: ignore
handle : ti.template() # type: ignore
) -> ti.types.ndarray(dtype=ti.f64, ndim=1): # type: ignore
n_nodes = ti.static(self.n_nodes)
state_size = ti.static(self.model_type_taichi.state_size)
n_params = ti.static(self.model_type_taichi.n_params)
out = ti.Vector.zero(ti.f64, n_nodes * state_size)
node_state = ti.Vector.zero(ti.f64, state_size)
params = ti.Vector.zero(ti.f64, n_params)
for node_idx in range(n_nodes):
for state_idx in range(state_size):
node_state[state_idx] = state_flat[node_idx * state_size + state_idx]
for state_idx in range(n_params):
params[state_idx] = self.params[node_idx, state_idx]
rhs_i = self._call_taichi_func_single(t, node_state, node_idx, handle, params)
for state_idx in range(state_size):
out[node_idx * state_size + state_idx] = rhs_i[state_idx]
return out
@ti.func
def _call_taichi_func_single(
self,
t : float, # type: ignore
node_state : ti.template(), # type: ignore
node_idx : int, # type: ignore
handle : ti.template(), # type: ignore
params : ti.template() # type: ignore
) -> ti.types.ndarray(dtype=ti.f64, ndim=1): # type: ignore
return self.model_type_taichi._call_taichi_func_param(
t,
node_state,
handle,
*params
)