from typing import Callable, Sequence, NamedTuple, Union
import numpy as np
from pararealml import Lhs
from pararealml.constrained_problem import ConstrainedProblem
from pararealml.operators.fdm.numerical_differentiator import \
NumericalDifferentiator
from pararealml.operators.symbol_mapper import SymbolMapper
[docs]class FDMSymbolMapArg(NamedTuple):
"""
The arguments to the FDM map functions.
"""
t: float
y: np.ndarray
d_y_constraint_function: Callable[[float], np.ndarray]
FDMSymbolMapFunction = Callable[[FDMSymbolMapArg], np.ndarray]
[docs]class FDMSymbolMapper(SymbolMapper[FDMSymbolMapArg, np.ndarray]):
"""
A symbol mapper implementation for the FDM operator.
"""
def __init__(
self,
cp: ConstrainedProblem,
differentiator: NumericalDifferentiator):
"""
:param cp: the constrained problem to create a symbol mapper for
:param differentiator: the numerical differentiator instance to use
"""
diff_eq = cp.differential_equation
super(FDMSymbolMapper, self).__init__(diff_eq)
self._differentiator = differentiator
self._mesh = cp.mesh
[docs] def t_map_function(self) -> FDMSymbolMapFunction:
return lambda arg: np.array([arg.t])
[docs] def y_map_function(self, y_ind: int) -> FDMSymbolMapFunction:
return lambda arg: arg.y[..., y_ind:y_ind + 1]
[docs] def x_map_function(self, x_axis: int) -> FDMSymbolMapFunction:
return lambda arg: \
self._mesh.vertex_coordinate_grids[x_axis][..., np.newaxis]
[docs] def y_gradient_map_function(
self,
y_ind: int,
x_axis: int) -> FDMSymbolMapFunction:
return lambda arg: self._differentiator.gradient(
arg.y[..., y_ind:y_ind + 1],
self._mesh,
x_axis,
arg.d_y_constraint_function(arg.t)[:, y_ind:y_ind + 1])
[docs] def y_hessian_map_function(
self,
y_ind: int,
x_axis1: int,
x_axis2: int) -> FDMSymbolMapFunction:
return lambda arg: self._differentiator.hessian(
arg.y[..., y_ind:y_ind + 1],
self._mesh,
x_axis1,
x_axis2,
arg.d_y_constraint_function(arg.t)[:, y_ind:y_ind + 1])
[docs] def y_divergence_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_]) -> FDMSymbolMapFunction:
if indices_contiguous:
return lambda arg: self._differentiator.divergence(
arg.y[..., y_indices[0]:y_indices[-1] + 1],
self._mesh,
arg.d_y_constraint_function(
arg.t)[:, y_indices[0]:y_indices[-1] + 1])
else:
return lambda arg: self._differentiator.divergence(
arg.y[..., y_indices],
self._mesh,
arg.d_y_constraint_function(arg.t)[:, y_indices])
[docs] def y_curl_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_],
curl_ind: int) -> FDMSymbolMapFunction:
if indices_contiguous:
return lambda arg: self._differentiator.curl(
arg.y[..., y_indices[0]:y_indices[-1] + 1],
self._mesh,
curl_ind,
arg.d_y_constraint_function(
arg.t)[:, y_indices[0]:y_indices[-1] + 1])
else:
return lambda arg: self._differentiator.curl(
arg.y[..., y_indices],
self._mesh,
curl_ind,
arg.d_y_constraint_function(arg.t)[:, y_indices])
[docs] def y_laplacian_map_function(
self,
y_ind: int) -> FDMSymbolMapFunction:
return lambda arg: self._differentiator.laplacian(
arg.y[..., y_ind:y_ind + 1],
self._mesh,
arg.d_y_constraint_function(arg.t)[:, y_ind:y_ind + 1])
[docs] def y_vector_laplacian_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_],
vector_laplacian_ind: int) -> FDMSymbolMapFunction:
if indices_contiguous:
return lambda arg: self._differentiator.vector_laplacian(
arg.y[..., y_indices[0]:y_indices[-1] + 1],
self._mesh,
vector_laplacian_ind,
arg.d_y_constraint_function(
arg.t)[:, y_indices[0]:y_indices[-1] + 1])
else:
return lambda arg: self._differentiator.vector_laplacian(
arg.y[..., y_indices],
self._mesh,
vector_laplacian_ind,
arg.d_y_constraint_function(arg.t)[:, y_indices])
[docs] def map_concatenated(
self,
arg: FDMSymbolMapArg,
lhs_type: Lhs) -> np.ndarray:
"""
Evaluates the right-hand side of the differential equation system
given the map argument and concatenates the resulting sequence of map
value arrays along the last axis.
"""
return np.concatenate(self.map(arg, lhs_type), axis=-1)