from abc import ABC, abstractmethod
from typing import Callable, Sequence, Dict, TypeVar, Generic, Optional, Union
import numpy as np
import sympy as sp
from pararealml.differential_equation import DifferentialEquation, Lhs
SymbolMapArg = TypeVar('SymbolMapArg')
SymbolMapValue = TypeVar('SymbolMapValue')
SymbolMapFunction = Callable[[SymbolMapArg], SymbolMapValue]
[docs]class SymbolMapper(ABC, Generic[SymbolMapArg, SymbolMapValue]):
"""
A class for mapping symbolic differential equation to numerical values.
"""
def __init__(self, diff_eq: DifferentialEquation):
"""
:param diff_eq: the differential equation to create a symbol mapper for
"""
self._diff_eq = diff_eq
self._symbol_map = self.create_symbol_map()
eq_sys = diff_eq.symbolic_equation_system
self._rhs_functions: Dict[
Optional[Lhs],
Callable[[SymbolMapArg], Sequence[SymbolMapValue]]
] = {None: self.create_rhs_map_function(range(len(eq_sys.rhs)))}
for lhs_type in Lhs:
self._rhs_functions[lhs_type] = self.create_rhs_map_function(
eq_sys.equation_indices_by_type(lhs_type))
[docs] @abstractmethod
def t_map_function(self) -> SymbolMapFunction:
"""
Returns a function for mapping t to a numerical value.
"""
[docs] @abstractmethod
def y_map_function(self, y_ind: int) -> SymbolMapFunction:
"""
Returns a function for mapping a component of y to a numerical value.
:param y_ind: the component of y to return a map for
:return: the mapper function for y
"""
[docs] @abstractmethod
def x_map_function(self, x_axis: int) -> SymbolMapFunction:
"""
Returns a function for mapping a component of x to a numerical value.
:param x_axis: the component of x to return a map for
:return: the mapper function for x
"""
[docs] @abstractmethod
def y_gradient_map_function(
self,
y_ind: int,
x_axis: int) -> SymbolMapFunction:
"""
Returns a function for mapping a component of the gradient of y to a
numerical value.
:param y_ind: the component of y whose gradient to return a map for
:param x_axis: the x-axis denoting the element of the gradient to
return a map for
:return: the mapper function for the gradient of y
"""
[docs] @abstractmethod
def y_hessian_map_function(
self,
y_ind: int,
x_axis1: int,
x_axis2: int) -> SymbolMapFunction:
"""
Returns a function for mapping a component of the Hessian of y to a
numerical value.
:param y_ind: the component of y whose Hessian to return a map for
:param x_axis1: the first x-axis denoting the element of the gradient
to return a map for
:param x_axis2: the second x-axis denoting the element of the gradient
to return a map for
:return: the mapper function for the Hessian of y
"""
[docs] @abstractmethod
def y_divergence_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_]) -> SymbolMapFunction:
"""
Returns a function for mapping the divergence of a set of components of
y to a numerical value.
:param y_indices: the components of y whose divergence to return a map
for
:param indices_contiguous: whether the indices are contiguous
:return: the mapper function for the divergence of y
"""
[docs] @abstractmethod
def y_curl_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_],
curl_ind: int) -> SymbolMapFunction:
"""
Returns a function for mapping the curl of a set of components of y to
a numerical value.
:param y_indices: the components of y whose curl to return a map for
:param indices_contiguous: whether the indices are contiguous
:param curl_ind: the index of the component of the curl to map
:return: the mapper function for the curl of y
"""
[docs] @abstractmethod
def y_laplacian_map_function(self, y_ind: int) -> SymbolMapFunction:
"""
Returns a function for mapping a component of the element-wise scalar
Laplacian of y to a numerical value.
:param y_ind: the component of y whose Laplacian to return a mp for
:return: the mapper function for the Laplacian of y
"""
[docs] @abstractmethod
def y_vector_laplacian_map_function(
self,
y_indices: Sequence[int],
indices_contiguous: Union[bool, np.bool_],
vector_laplacian_ind: int) -> SymbolMapFunction:
"""
Returns a function for mapping the vector Laplacian of a set of
components of y to a numerical value.
:param y_indices: the components of y whose vector Laplacian to return
a map for
:param indices_contiguous: whether the indices are contiguous
:param vector_laplacian_ind: the index of the component of the vector
Laplacian to map
:return: the mapper function for the vector Laplacian of y
"""
[docs] def create_symbol_map(self) -> Dict[sp.Symbol, SymbolMapFunction]:
"""
Creates a dictionary linking the symbols present in the differential
equation instance associated with the symbol mapper to a set of
functions used to map the symbols to numerical values.
"""
symbol_map = {}
x_dimension = self._diff_eq.x_dimension
eq_sys = self._diff_eq.symbolic_equation_system
all_symbols = set.union(*[rhs.free_symbols for rhs in eq_sys.rhs])
for symbol in all_symbols:
symbol_name_tokens = symbol.name.split('_')
prefix = symbol_name_tokens[0]
indices = [int(ind) for ind in symbol_name_tokens[1:]] \
if len(symbol_name_tokens) > 1 else []
if prefix == 't':
symbol_map[symbol] = self.t_map_function()
elif prefix == 'y':
symbol_map[symbol] = self.y_map_function(*indices)
elif prefix == 'x':
symbol_map[symbol] = self.x_map_function(*indices)
elif prefix == 'y-gradient':
symbol_map[symbol] = self.y_gradient_map_function(*indices)
elif prefix == 'y-hessian':
symbol_map[symbol] = self.y_hessian_map_function(*indices)
elif prefix == 'y-laplacian':
symbol_map[symbol] = self.y_laplacian_map_function(*indices)
else:
indices_contiguous = np.all([
indices[i] == indices[i + 1] - 1
for i in range(len(indices) - 1)
])
if prefix == 'y-divergence':
symbol_map[symbol] = self.y_divergence_map_function(
indices, indices_contiguous)
elif prefix == 'y-curl':
symbol_map[symbol] = \
self.y_curl_map_function(
indices, indices_contiguous, 0) \
if x_dimension == 2 else \
self.y_curl_map_function(
indices[:-1], indices_contiguous, indices[-1])
elif prefix == 'y-vector-laplacian':
self.y_vector_laplacian_map_function(
indices[:-1], indices_contiguous, indices[-1])
return symbol_map
[docs] def create_rhs_map_function(
self,
indices: Sequence[int]
) -> Callable[[SymbolMapArg], Sequence[SymbolMapValue]]:
"""
Creates a function for evaluating the right-hand sides of the equations
denoted by the provided indices.
:param indices: the indices of the equations within the differential
equation system whose evaluation function is to be created
:return: a function that returns the numerical value of the right-hand
sides given a substitution argument
"""
rhs = self._diff_eq.symbolic_equation_system.rhs
selected_rhs = []
selected_rhs_symbols = set()
for i in indices:
rhs_i = rhs[i]
selected_rhs.append(rhs_i)
selected_rhs_symbols.update(rhs_i.free_symbols)
subst_functions = \
[self._symbol_map[symbol] for symbol in selected_rhs_symbols]
rhs_lambda = sp.lambdify([selected_rhs_symbols], selected_rhs, 'numpy')
def rhs_map_function(arg: SymbolMapArg) -> Sequence[SymbolMapValue]:
return rhs_lambda(
[subst_function(arg) for subst_function in subst_functions])
return rhs_map_function
[docs] def map(self,
arg: SymbolMapArg,
lhs_type: Optional[Lhs] = None) -> Sequence[SymbolMapValue]:
"""
Evaluates the right-hand side of the differential equation system
given the map argument.
:param arg: the map argument that the numerical values of the
right-hand sides depend on
:param lhs_type: the left-hand type of the equations whose right-hand
sides are to be evaluated; if None, the whole differential equation
system's right-hand side is evaluated
:return: the numerical value of the right-hand side of the differential
equation as a sequence of map values where each element corresponds
to an equation within the system
"""
return self._rhs_functions[lhs_type](arg)