"""
Handling mathematical expressions with sympy
This module provides classes representing expressions that can be provided as
human-readable strings and are converted to :mod:`numpy` and :mod:`numba`
representations using :mod:`sympy`.
.. autosummary::
:nosignatures:
parse_number
ScalarExpression
TensorExpression
evaluate
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""
from __future__ import annotations
import builtins
import copy
import json
import logging
import numbers
import re
from abc import ABCMeta, abstractmethod
from pathlib import Path
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Mapping,
Optional,
Sequence,
Set,
Tuple,
Type,
Union,
)
import numba as nb # lgtm [py/import-and-import-from]
import numpy as np
import sympy
from sympy.core import basic
from sympy.printing.pycode import PythonCodePrinter
from sympy.utilities.lambdify import _get_namespace
from ..fields.base import DataFieldBase, FieldBase
from ..grids.boundaries.axes import BoundariesData
from ..grids.boundaries.local import BCDataError
from .cache import cached_method, cached_property
from .docstrings import fill_in_docstring
from .misc import Number, classproperty, number, number_array
from .numba import convert_scalar, jit
from .typing import NumberOrArray
try:
from numba.core.extending import overload
except ImportError:
# assume older numba module structure
from numba.extending import overload
[docs]@fill_in_docstring
def parse_number(
expression: Union[str, Number], variables: Mapping[str, Number] = None
) -> Number:
r"""return a number compiled from an expression
Warning:
{WARNING_EXEC}
Args:
expression (str or Number):
An expression that can be interpreted as a number
variables (dict):
A dictionary of values that replace variables in the expression
Returns:
Number: the calculated value
"""
from sympy.parsing import sympy_parser
if variables is None:
variables = {}
expr = sympy_parser.parse_expr(str(expression))
try:
value = number(expr.subs(variables))
except TypeError as err:
if not err.args:
err.args = ("",)
err.args = err.args + (f"Expression: `{expr}`",)
raise
return value
def _heaviside_implemention(x1, x2=0.5):
"""implementation of the Heaviside function used for numba and sympy
Args:
x1 (float): Argument of the function
x2 (float): Value returned when the argument is zero
Returns:
float: 0 if x1 is negative, 1 if x1 is positive, and x2 if x1 == 0
"""
if np.isnan(x1):
return np.nan
elif x1 == 0:
return x2
elif x1 < 0:
return 0.0
else:
return 1.0
@overload(np.heaviside)
def np_heaviside(x1, x2):
"""numba implementation of the Heaviside function"""
return _heaviside_implemention
# special functions that we want to support in expressions but that are not defined by
# sympy version 1.6 or have a different signature than expected by numba/numpy
SPECIAL_FUNCTIONS = {"Heaviside": _heaviside_implemention}
class ListArrayPrinter(PythonCodePrinter):
"""special sympy printer returning arrays as lists"""
def _print_ImmutableDenseNDimArray(self, arr):
arrays = ", ".join(f"{self._print(expr)}" for expr in arr)
return f"[{arrays}]"
class NumpyArrayPrinter(PythonCodePrinter):
"""special sympy printer returning numpy arrays"""
def _print_ImmutableDenseNDimArray(self, arr):
arrays = ", ".join(f"asarray({self._print(expr)})" for expr in arr)
return f"array(broadcast_arrays({arrays}))"
ExpressionType = Union[float, str, np.ndarray, basic.Basic, "ExpressionBase"]
[docs]class ExpressionBase(metaclass=ABCMeta):
"""abstract base class for handling expressions"""
@fill_in_docstring
def __init__(
self,
expression: basic.Basic,
signature: Sequence[Union[str, List[str]]] = None,
*,
user_funcs: Dict[str, Callable] = None,
consts: Dict[str, NumberOrArray] = None,
):
"""
Warning:
{WARNING_EXEC}
Args:
expression (:class:`sympy.core.basic.Basic`):
A sympy expression or array. This could for instance be an
instance of :class:`~sympy.core.expr.Expr` or
:class:`~sympy.tensor.array.ndim_array.NDimArray`.
signature (list of str, optional):
The signature defines which variables are expected in the
expression. This is typically a list of strings identifying
the variable names. Individual names can be specified as list,
in which case any of these names can be used. The first item in
such a list is the definite name and if another name of the list
is used, the associated variable is renamed to the definite
name. If signature is `None`, all variables in `expressions`
are allowed.
user_funcs (dict, optional):
A dictionary with user defined functions that can be used in the
expression.
consts (dict, optional):
A dictionary with user defined constants that can be used in the
expression. The values of these constants should either be numbers or
:class:`~numpy.ndarray`.
"""
try:
self._sympy_expr = sympy.simplify(expression)
except TypeError:
# work-around for sympy bug (github.com/sympy/sympy/issues/19829)
self._sympy_expr = expression
self._logger = logging.getLogger(self.__class__.__name__)
self.user_funcs = {} if user_funcs is None else user_funcs
self.consts = {} if consts is None else consts
# check consistency of the arguments
self._check_signature(signature)
for name, value in self.consts.items():
if isinstance(value, FieldBase):
self._logger.warning(
f"Constant `{name}` is a field, but expressions usually require "
f"numerical arrays. Did you mean to use `{name}.data`?"
)
def __repr__(self):
return (
f'{self.__class__.__name__}("{self.expression}", ' f"signature={self.vars})"
)
def __eq__(self, other):
"""compare this expression to another one"""
if not isinstance(other, self.__class__):
return NotImplemented
# compare what the expressions depend on
if set(self.vars) != set(other.vars):
return False
# compare the auxiliary data
if self.user_funcs != other.user_funcs or self.consts != other.consts:
return False
# compare the expressions themselves by checking their difference
diff = sympy.simplify(self._sympy_expr - other._sympy_expr)
if isinstance(self._sympy_expr, sympy.NDimArray):
return diff == sympy.Array(np.zeros(self._sympy_expr.shape))
else:
return diff == 0
@classproperty
def _reserved_symbols(self) -> Set[str]:
"""set: reserved sympy symbols that should not be used in expressions"""
try:
# try returning a cached version of the list
return ExpressionBase._reserved_symbols_cache # type: ignore
except AttributeError:
# the cache was not present => load list from resources
module_path = Path(__file__).resolve().parent
resource_path = module_path / "resources" / "reserved_sympy_symbols.json"
try:
with open(resource_path) as f:
ExpressionBase._reserved_symbols_cache = set(json.load(f)) # type: ignore
except (IOError, OSError):
# cannot read the file, so return a minimal list
ExpressionBase._reserved_symbols_cache = {"E", "I", "pi"} # type: ignore
return ExpressionBase._reserved_symbols_cache # type: ignore
[docs] @classmethod
def check_reserved_symbols(
cls, symbols: Iterable[str], strict: bool = True
) -> None:
"""throws an error if reserved symbols are found
Args:
symbols (iterable):
A sequence or set of strings with symbols to check.
strict (bool):
Flag determining whether an exception is raised
"""
symbol_set = {s.lower() for s in symbols}
reserved_symbols = symbol_set & ScalarExpression._reserved_symbols
if any(reserved_symbols):
if len(reserved_symbols) == 1:
name = reserved_symbols.pop()
msg = f"Cannot use reserved symbol `{name}` as field name"
else:
msg = f"Cannot use reserved symbols {reserved_symbols} as field names"
if strict:
raise ValueError(msg)
else:
logging.getLogger(cls.__name__).warning(msg)
@property
def _free_symbols(self) -> Set:
"""return symbols that appear in the expression and are not in self.consts"""
return {
sym for sym in self._sympy_expr.free_symbols if sym.name not in self.consts
}
@property
def constant(self) -> bool:
"""bool: whether the expression is a constant"""
return len(self._free_symbols) == 0
@property
def complex(self) -> bool:
"""bool: whether the expression contains the imaginary unit I"""
return sympy.I in self._sympy_expr.atoms()
@property
@abstractmethod
def shape(self) -> Tuple[int, ...]:
pass
def _check_signature(self, signature: Sequence[Union[str, List[str]]] = None):
"""validate the variables of the expression against the signature"""
# get arguments of the expressions
if self.constant:
# constant expression do not depend on any variables
args: Set[str] = set()
if signature is None:
signature = []
else:
# general expressions might have a variable
args = set(str(s).split("[")[0] for s in self._free_symbols)
if signature is None:
# create signature from arguments
signature = list(sorted(args))
self._logger.debug(f"Expression arguments: {args}")
# check whether signature contains reserved symbols
sig_elements = []
for sig in signature:
if isinstance(sig, str):
sig_elements.append(sig)
else:
sig_elements.extend(sig)
self.check_reserved_symbols(sig_elements, strict=False)
# check whether variables are in signature
self.vars: Any = []
found = set()
for sig in signature:
sig_list = [sig] if isinstance(sig, str) else sig
# use the first item as the variable name
arg_name = sig_list[0]
self.vars.append(arg_name)
# check whether this part of the signature is present
for arg in args:
if arg in sig_list:
if arg != arg_name: # synonym has been used
old = sympy.symbols(arg)
new = sympy.symbols(arg_name)
self._sympy_expr = self._sympy_expr.subs(old, new)
self._logger.info(f'Renamed variable "{old}"->"{new}"')
found.add(arg)
break
args = set(args) - found
if len(args) > 0:
raise RuntimeError(
f"Arguments {args} were not defined in expression signature {signature}"
)
@property
def expression(self) -> str:
"""str: the expression in string form"""
# turn numerical values into easily readable text
if isinstance(self._sympy_expr, sympy.NDimArray):
expr = self._sympy_expr.applyfunc(lambda x: x.evalf(chop=True))
else:
expr = self._sympy_expr.evalf(chop=True)
return str(expr.xreplace({n: float(n) for n in expr.atoms(sympy.Float)}))
@property
def rank(self) -> int:
"""int: the rank of the expression"""
return len(self.shape)
[docs] def depends_on(self, variable: str) -> bool:
"""determine whether the expression depends on `variable`
Args:
variable (str): the name of the variable to check for
Returns:
bool: whether the variable appears in the expression
"""
if self.constant:
return False
else:
return any(variable == str(symbol) for symbol in self._free_symbols)
def _get_function(
self,
single_arg: bool = False,
user_funcs: Dict[str, Callable] = None,
prepare_compilation: bool = False,
) -> Callable[..., NumberOrArray]:
"""return function evaluating expression
Args:
single_arg (bool):
Determines whether the returned function accepts all variables
in a single argument as an array or whether all variables need
to be supplied separately
user_funcs (dict):
Additional functions that can be used in the expression
prepare_compilation (bool):
Determines whether all user functions are marked with
:func:`numba.extending.register_jitable` to prepare for compilation.
Returns:
function: the function
"""
# collect all the user functions
user_functions = self.user_funcs.copy()
if user_funcs is not None:
user_functions.update(user_funcs)
user_functions.update(SPECIAL_FUNCTIONS)
if prepare_compilation:
# transform the user functions, so they can be compiled using numba
def compile_func(func):
if isinstance(func, np.ufunc):
# this is a work-around that allows to compile numpy ufuncs
return jit(lambda *args: func(*args))
else:
return jit(func)
user_functions = {k: compile_func(v) for k, v in user_functions.items()}
# initialize the printer that deals with numpy arrays correctly
if prepare_compilation:
printer_class: Type[PythonCodePrinter] = ListArrayPrinter
else:
printer_class = NumpyArrayPrinter
printer = printer_class(
{
"fully_qualified_modules": False,
"inline": True,
"allow_unknown_functions": True,
"user_functions": {k: k for k in user_functions},
}
)
# determine the list of variables that the function depends on
variables = (self.vars,) if single_arg else tuple(self.vars)
constants = tuple(self.consts)
# turn the expression into a callable function
self._logger.info("Parse sympy expression `%s`", self._sympy_expr)
func = sympy.lambdify(
variables + constants,
self._sympy_expr,
modules=[user_functions, "numpy"],
printer=printer,
)
# Apply the constants if there are any. Note that we use this pattern of a
# partial function instead of replacing the constants in the sympy expression
# directly since sympy does not work well with numpy arrays.
if constants:
const_values = tuple(self.consts[c] for c in constants) # @UnusedVariable
if prepare_compilation:
func = jit(func)
# TOOD: support keyword arguments
def result(*args):
return func(*args, *const_values)
else:
result = func
return result
@cached_method()
def _get_function_cached(
self, single_arg: bool = False, prepare_compilation: bool = False
) -> Callable[..., NumberOrArray]:
"""return function evaluating expression
Args:
single_arg (bool):
Determines whether the returned function accepts all variables
in a single argument as an array or whether all variables need
to be supplied separately
prepare_compilation (bool):
Determines whether all user functions are marked with
:func:`numba.extending.register_jitable` to prepare for compilation.
Returns:
function: the function
"""
return self._get_function(single_arg, prepare_compilation=prepare_compilation)
def __call__(self, *args, **kwargs) -> NumberOrArray:
"""return the value of the expression for the given values"""
return self._get_function_cached(single_arg=False)(*args, **kwargs)
[docs] @cached_method()
def get_compiled(self, single_arg: bool = False) -> Callable[..., NumberOrArray]:
"""return numba function evaluating expression
Args:
single_arg (bool): Determines whether the returned function accepts
all variables in a single argument as an array or whether all
variables need to be supplied separately
Returns:
function: the compiled function
"""
# compile the actual expression
func = self._get_function_cached(
single_arg=single_arg, prepare_compilation=True
)
return jit(func) # type: ignore
[docs]class ScalarExpression(ExpressionBase):
"""describes a mathematical expression of a scalar quantity"""
shape: Tuple[int, ...] = tuple()
@fill_in_docstring
def __init__(
self,
expression: ExpressionType = 0,
signature: Optional[Sequence[Union[str, List[str]]]] = None,
*,
user_funcs: Optional[Dict[str, Callable]] = None,
consts: Optional[Dict[str, NumberOrArray]] = None,
allow_indexed: bool = False,
):
"""
Warning:
{WARNING_EXEC}
Args:
expression (str or float):
The expression, which is either a number or a string that sympy
can parse
signature (list of str):
The signature defines which variables are expected in the expression.
This is typically a list of strings identifying the variable names.
Individual names can be specified as lists, in which case any of these
names can be used. The first item in such a list is the definite name
and if another name of the list is used, the associated variable is
renamed to the definite name. If signature is `None`, all variables in
`expressions` are allowed.
user_funcs (dict, optional):
A dictionary with user defined functions that can be used in the
expression
consts (dict, optional):
A dictionary with user defined constants that can be used in the
expression. The values of these constants should either be numbers or
:class:`~numpy.ndarray`.
allow_indexed (bool):
Whether to allow indexing of variables. If enabled, array variables are
allowed to be indexed using square bracket notation.
"""
self.allow_indexed = allow_indexed
# parse the expression
if isinstance(expression, ScalarExpression):
# copy constructor
sympy_expr = copy.copy(expression._sympy_expr)
if signature is None:
signature = expression.vars
self.allow_indexed = expression.allow_indexed
if user_funcs is None:
user_funcs = expression.user_funcs
else:
user_funcs.update(expression.user_funcs)
if consts is None:
consts = expression.consts
else:
consts.update(expression.consts)
elif callable(expression):
# expression is some other callable -> not allowed anymore
raise TypeError("Expression must be a string and not a function")
elif isinstance(expression, numbers.Number):
# expression is a simple number
sympy_expr = sympy.Float(expression)
elif bool(expression):
# parse expression as a string
expression = self._prepare_expression(str(expression))
# parse the expression using sympy
from sympy.parsing import sympy_parser
sympy_expr = sympy_parser.parse_expr(expression)
else:
# expression is empty, False or None => set it to zero
sympy_expr = sympy.Float(0)
super().__init__(
expression=sympy_expr,
signature=signature,
user_funcs=user_funcs,
consts=consts,
)
[docs] def copy(self) -> ScalarExpression:
"""return a copy of the current expression"""
# __init__ copies all relevant attributes
return self.__class__(self)
@property
def value(self) -> Number:
"""float: the value for a constant expression"""
if self.constant:
try:
# try simply evaluating the expression as a number
value = number(self._sympy_expr.evalf())
except TypeError:
# This can fail if user_funcs are supplied, which would not be replaced
# in the numeric implementation above. We thus also try to call the
# expression without any arguments
value = number(self()) # type: ignore
# Note that this may fail when the expression is actually constant, but
# has a signature that forces it to depend on some arguments. However,
# we feel this situation should not be very common, so we do not (yet)
# deal with it.
return value
else:
raise TypeError("Only constant expressions have a defined value")
@property
def is_zero(self) -> bool:
"""bool: returns whether the expression is zero"""
return self.constant and self.value == 0
def __bool__(self) -> bool:
"""tests whether the expression is nonzero"""
return not self.constant or self.value != 0
def __eq__(self, other):
if not isinstance(other, self.__class__):
return NotImplemented
return super().__eq__(other) and self.allow_indexed == other.allow_indexed
def _prepare_expression(self, expression: str) -> str:
"""replace indexed variables, if allowed
Args:
expression (str):
An expression string that might contain variables that are indexed using
square brackets. If this is the case, they are rewritten using the
sympy object `IndexedBase`.
"""
if self.allow_indexed:
return re.sub(r"(\w+)(\[\w+\])", r"IndexedBase(\1)\2", expression)
else:
return expression
def _var_indexed(self, var: str) -> bool:
"""checks whether the variable `var` is used in an indexed form"""
from sympy.tensor.indexed import Indexed
return any(
isinstance(s, Indexed) and s.base.name == var for s in self._free_symbols
)
[docs] def differentiate(self, var: str) -> ScalarExpression:
"""return the expression differentiated with respect to var"""
if self.constant:
# return empty expression
return ScalarExpression(
expression=0, signature=self.vars, allow_indexed=self.allow_indexed
)
if self.allow_indexed:
if self._var_indexed(var):
raise NotImplementedError("Cannot differentiate with respect to vector")
# turn variable into sympy object and treat an indexed variable separately
var_expr = self._prepare_expression(var)
if "[" in var:
from sympy.parsing import sympy_parser
var_symbol = sympy_parser.parse_expr(var_expr)
else:
var_symbol = sympy.Symbol(var_expr)
return ScalarExpression(
self._sympy_expr.diff(var_symbol),
signature=self.vars,
allow_indexed=self.allow_indexed,
user_funcs=self.user_funcs,
)
@cached_property()
def derivatives(self) -> TensorExpression:
"""differentiate the expression with respect to all variables"""
if self.constant:
# return empty expression
dim = len(self.vars)
expression = sympy.Array(np.zeros(dim), shape=(dim,))
return TensorExpression(expression=expression, signature=self.vars)
if self.allow_indexed:
if any(self._var_indexed(var) for var in self.vars):
raise RuntimeError(
"Cannot calculate gradient for expressions with indexed variables"
)
grad = sympy.Array([self._sympy_expr.diff(sympy.Symbol(v)) for v in self.vars])
return TensorExpression(
sympy.simplify(grad), signature=self.vars, user_funcs=self.user_funcs
)
[docs]class TensorExpression(ExpressionBase):
"""describes a mathematical expression of a tensorial quantity"""
@fill_in_docstring
def __init__(
self,
expression: ExpressionType,
signature: Optional[Sequence[Union[str, List[str]]]] = None,
*,
user_funcs: Optional[Dict[str, Callable]] = None,
consts: Optional[Dict[str, NumberOrArray]] = None,
):
"""
Warning:
{WARNING_EXEC}
Args:
expression (str or float):
The expression, which is either a number or a string that sympy
can parse
signature (list of str):
The signature defines which variables are expected in the
expression. This is typically a list of strings identifying
the variable names. Individual names can be specified as list,
in which case any of these names can be used. The first item in
such a list is the definite name and if another name of the list
is used, the associated variable is renamed to the definite
name. If signature is `None`, all variables in `expressions`
are allowed.
user_funcs (dict, optional):
A dictionary with user defined functions that can be used in the
expression.
consts (dict, optional):
A dictionary with user defined constants that can be used in the
expression. The values of these constants should either be numbers or
:class:`~numpy.ndarray`.
"""
from sympy.tensor.array.ndim_array import ImmutableNDimArray
# parse the expression
if isinstance(expression, TensorExpression):
# copy constructor
sympy_expr = copy.copy(expression._sympy_expr)
if user_funcs is None:
user_funcs = expression.user_funcs
else:
user_funcs.update(expression.user_funcs)
elif isinstance(expression, (np.ndarray, list, tuple)):
# expression is a constant array
sympy_expr = sympy.Array(sympy.sympify(expression))
elif isinstance(expression, ImmutableNDimArray):
# expression is an array of sympy expressions
sympy_expr = expression
else:
# parse expression as a string
expression = str(expression)
# parse the expression using sympy
from sympy.parsing import sympy_parser
sympy_expr = sympy.Array(sympy_parser.parse_expr(expression))
super().__init__(
expression=sympy_expr,
signature=signature,
user_funcs=user_funcs,
consts=consts,
)
def __repr__(self):
if self.shape == (0,):
# work-around for sympy bug (github.com/sympy/sympy/issues/19829)
return f'{self.__class__.__name__}("[]", signature={self.vars})'
else:
return super().__repr__()
@property
def shape(self) -> Tuple[int, ...]:
"""tuple: the shape of the tensor"""
return self._sympy_expr.shape # type: ignore
def __getitem__(self, index):
expr = self._sympy_expr[index]
if isinstance(expr, sympy.Array):
return TensorExpression(
expr, signature=self.vars, user_funcs=self.user_funcs
)
else:
return ScalarExpression(
expr, signature=self.vars, user_funcs=self.user_funcs
)
@property
def value(self):
"""the value for a constant expression"""
if self.constant:
try:
# try simply evaluating the expression as a number
value = number_array(self._sympy_expr.tolist())
except TypeError:
# This can fail if user_funcs are supplied, which would not be replaced
# in the numeric implementation above. We thus also try to call the
# expression without any arguments
value = number_array(self())
# Note that this may fail when the expression is actually constant, but
# has a signature that forces it to depend on some arguments. However,
# we feel this situation should not be very common, so we do not (yet)
# deal with it.
return value
else:
raise TypeError("Only constant expressions have a defined value")
[docs] def differentiate(self, var: str) -> TensorExpression:
"""return the expression differentiated with respect to var"""
if self.constant:
derivative = np.zeros(self.shape)
else:
derivative = self._sympy_expr.diff(sympy.Symbol(var))
return TensorExpression(derivative, self.vars, user_funcs=self.user_funcs)
@cached_property()
def derivatives(self) -> TensorExpression:
"""differentiate the expression with respect to all variables"""
shape = (len(self.vars),) + self.shape
if self.constant:
# return empty expression
derivatives = sympy.Array(np.zeros(shape), shape)
else:
# perform the derivatives with respect to all variables
dx = sympy.Array([sympy.Symbol(s) for s in self.vars])
derivatives = sympy.derive_by_array(self._sympy_expr, dx)
return TensorExpression(derivatives, self.vars, user_funcs=self.user_funcs)
[docs] def get_compiled_array(
self, single_arg: bool = True
) -> Callable[[np.ndarray, Optional[np.ndarray]], np.ndarray]:
"""compile the tensor expression such that a numpy array is returned
Args:
single_arg (bool):
Whether the compiled function expects all arguments as a single array
or whether they are supplied individually.
"""
assert isinstance(self._sympy_expr, sympy.Array), "Expression must be an array"
variables = ", ".join(v for v in self.vars)
shape = self._sympy_expr.shape
if nb.config.DISABLE_JIT:
# special path used by coverage test without jitting. This can be
# removed once the `convert_scalar` wrapper is obsolete
lines = [
f" out[{str(idx + (...,))[1:-1]}] = {self._sympy_expr[idx]}"
for idx in np.ndindex(*self._sympy_expr.shape)
]
else:
lines = [
f" out[{str(idx + (...,))[1:-1]}] = "
f"convert_scalar({self._sympy_expr[idx]})"
for idx in np.ndindex(*self._sympy_expr.shape)
]
# TODO: replace the np.ndindex with np.ndenumerate eventually. This does not
# work with numpy 1.18, so we have the work around using np.ndindex
# TODO: We should also support constants similar to ScalarExpressions. They
# could be written in separate lines and prepended to the actual code. However,
# we would need to make sure to print numpy arrays correctly.
if variables:
# the expression takes variables as input
if single_arg:
# the function takes a single input array
first_dim = 0 if len(self.vars) == 1 else 1
code = "def _generated_function(arr, out=None):\n"
code += f" arr = asarray(arr)\n"
code += f" {variables} = arr\n"
code += f" if out is None:\n"
code += f" out = empty({shape} + arr.shape[{first_dim}:])\n"
else:
# the function takes each variables as an argument
code = f"def _generated_function({variables}, out=None):\n"
code += f" if out is None:\n"
code += f" out = empty({shape} + shape({self.vars[0]}))\n"
else:
# the expression is constant
if single_arg:
code = "def _generated_function(arr=None, out=None):\n"
else:
code = "def _generated_function(out=None):\n"
code += f" if out is None:\n"
code += f" out = empty({shape})\n"
code += "\n".join(lines) + "\n"
code += " return out"
self._logger.debug("Code for `get_compiled_array`: %s", code)
namespace = _get_namespace("numpy")
namespace["convert_scalar"] = convert_scalar
namespace["builtins"] = builtins
namespace.update(self.user_funcs)
local_vars: Dict[str, Any] = {}
exec(code, namespace, local_vars)
function = local_vars["_generated_function"]
return jit(function) # type: ignore
[docs]@fill_in_docstring
def evaluate(
expression: str,
fields: Dict[str, DataFieldBase],
*,
bc: BoundariesData = "auto_periodic_neumann",
bc_ops: Dict[str, BoundariesData] = None,
user_funcs: Dict[str, Callable] = None,
consts: Dict[str, NumberOrArray] = None,
label: str = None,
) -> DataFieldBase:
"""evaluate an expression involving fields
Warning:
{WARNING_EXEC}
Args:
expression (str):
The expression, which is parsed by :mod:`sympy`. The expression may contain
variables (i.e., fields and spatial coordinates of the grid), standard local
mathematical operators defined by sympy, and the operators defined in the
:mod:`pde` package. Note that operators need to be specified with their full
name, i.e., `laplace` for a scalar Laplacian and `vector_laplace` for a
Laplacian operating on a vector field. Moreover, the dot product between
two vector fields can be denoted by using `dot(field1, field2)` in the
expression, and `outer(field1, field2)` calculates an outer product.
More information can be found in the
:ref:`expression documentation <documentation-expressions>`.
fields (dict):
Dictionary of the fields involved in the expression.
bc:
Boundary conditions for the operators used in the expression. The conditions
here are applied to all operators that do not have a specialized condition
given in `bc_ops`.
{ARG_BOUNDARIES}
bc_ops (dict):
Special boundary conditions for some operators. The keys in this dictionary
specify the operator to which the boundary condition will be applied.
user_funcs (dict, optional):
A dictionary with user defined functions that can be used in the expressions
in `rhs`.
consts (dict, optional):
A dictionary with user defined constants that can be used in the expression.
These can be either scalar numbers or fields defined on the same grid as the
actual simulation.
label (str):
Name of the field that is returned.
Returns:
:class:`pde.fields.base.DataFieldBase`: The resulting field. The rank of the
returned field (and thus the precise class) is determined automatically.
"""
from sympy.core.function import AppliedUndef
from ..fields import VectorField
logger = logging.getLogger("evaluate")
# validate input
if consts is None:
consts = {}
# turn the expression strings into sympy expressions
expr = ScalarExpression(expression, user_funcs=user_funcs, consts=consts)
# determine undefined functions in the expression
operators = {
func.__class__.__name__
for func in expr._sympy_expr.atoms(AppliedUndef)
if func.__class__.__name__ not in expr.user_funcs
}
# setup boundary conditions
if bc_ops is None:
bcs: Dict[str, Any] = {"*": bc}
else:
bcs = dict(bc_ops)
if "*" in bcs and bc != "auto_periodic_neumann":
logger.warning("Found default BCs in `bcs` and `bc_ops`")
bcs["*"] = bc # append default boundary conditions
# check whether all fields have the same grid
grid = None
for field in fields.values():
if grid is None:
grid = field.grid
else:
field.grid.assert_grid_compatible(grid)
if grid is None:
raise ValueError("No fields given")
# prepare the differential operators
# check whether PDE has variables with same names as grid axes
name_overlap = set(fields) & set(grid.axes)
if name_overlap:
raise ValueError(f"Coordinate {name_overlap} cannot be used as field name")
# obtain the (differential) operators
ops: Dict[str, Callable] = {}
for func in operators:
if func == "dot" or func == "inner":
# add dot product between two vector fields. This can for instance
# appear when two gradients of scalar fields need to be multiplied
ops[func] = VectorField(grid).make_dot_operator(backend="numpy")
elif func == "outer":
# generate an operator that calculates an outer product
ops[func] = VectorField(grid).make_outer_prod_operator(backend="numpy")
else:
# determine boundary conditions for this operator and variable
for bc_key, bc in bcs.items():
if bc_key == func or bc_key == "*":
break # found a matching boundary condition
else:
raise RuntimeError(
f"Could not find suitable boundary condition for function `{func}`"
)
# Tell the user what BC we chose for a given operator
logger.info("Using BC `%s` for operator `%s` in expression", bc, func)
# create the function evaluating the operator
try:
ops[func] = grid.make_operator(func, bc=bc)
except BCDataError:
# wrong data was supplied for the boundary condition
raise
except ValueError:
# any other exception should signal that the operator is not defined, so
# we (almost) silently assume that sympy defines the operator
logger.info("Assuming that sympy knows undefined operator `%s`", func)
# obtain the function to calculate the right hand side
signature = tuple(fields.keys()) + ("none", "bc_args")
# check whether this function depends on additional input
if any(expr.depends_on(c) for c in grid.axes):
# expression has a spatial dependence, too
# extend the signature
signature += tuple(grid.axes)
# inject the spatial coordinates into the expression for the rhs
extra_args = tuple( # @UnusedVariable
grid.cell_coords[..., i] for i in range(grid.num_axes)
)
else:
# expression only depends on the actual variables
extra_args = tuple() # @UnusedVariable
# check whether all variables are accounted for
extra_vars = set(expr.vars) - set(signature)
if extra_vars:
extra_vars_str = ", ".join(sorted(extra_vars))
raise RuntimeError(f"Undefined variable in expression: {extra_vars_str}")
expr.vars = signature
logger.info("Expression has signature %s", signature)
# extract input field data and calculate result
field_data = [field.data for field in fields.values()]
# calculate the result of the expression
func = expr._get_function(single_arg=False, user_funcs=ops)
result_data = func(*field_data, None, {}, *extra_args)
# turn result into a proper field
result_rank = result_data.ndim - grid.num_axes
result_cls = DataFieldBase.get_class_by_rank(result_rank)
return result_cls(grid, result_data, label=label)
__all__ = [
"ExpressionBase",
"ScalarExpression",
"TensorExpression",
"parse_number",
"evaluate",
]