def lambdify(args, expr):
""" Returns function for fast calculation of numerical values
A modified version of sympy's lambdify that will find 'aliased'
Functions and substitute them appropriately at evaluation time.
See ``sympy.lambdify`` for more detail.
Parameters
----------
args : object or sequence of objects
May well be sympy Symbols
expr : expression
The expression is anything that can be passed to the sympy
lambdify function, meaning anything that gives valid code from
``str(expr)``.
Examples
--------
>>> x = sympy.Symbol('x')
>>> f = lambdify(x, x**2)
>>> f(3)
9
"""
n = _imp_namespace(expr)
# There was a bug in sympy such that dictionaries passed in as first
# namespaces to lambdify, before modules, would get overwritten by
# later calls to lambdify. The next two lines are to get round this
# bug.
from sympy.utilities.lambdify import _get_namespace
np_ns = _get_namespace('numpy').copy()
return sympy.lambdify(args, expr, modules=(n, np_ns))
def _set_namespace(self, namespaces):
"""Set the name space for use when calling eval. This needs to contain all the relvant functions for mapping from symbolic python to the numerical python. It also contains variables, cached portions etc."""
self.namespace = {}
for m in namespaces[::-1]:
buf = _get_namespace(m)
self.namespace.update(buf)
self.namespace.update(self.__dict__)
def __init__(self, args, expr):
''' Initialize lambdifier
This lambdifier only allows there to be one input argument, in
fact, because of the __call__ signature.
Parameters
----------
args : object or sequence of objects
May well be sympy Symbols
expr : expression
Examples
--------
>>> x = sympy.Symbol('x')
>>> f = lambdify(x, x**2)
>>> f(3)
9
'''
if isinstance(expr, sympy.FunctionClass):
# NNB t is undefined at this point
expr = expr(t)
n = {}
_add_aliases_to_namespace(n, expr)
self.n = n.copy()
from sympy.utilities.lambdify import _get_namespace
for k, v in _get_namespace('numpy').items():
self.n[k] = v
self._f = sympy.lambdify(args, expr, self.n)
self.expr = expr
def _set_namespace(self, namespaces):
"""Set the name space for use when calling eval. This needs to contain all the relvant functions for mapping from symbolic python to the numerical python. It also contains variables, cached portions etc."""
self.namespace = {}
for m in namespaces[::-1]:
buf = _get_namespace(m)
self.namespace.update(buf)
self.namespace.update(self.__dict__)
def __init__(self, args, expr):
if isinstance(expr, sympy.FunctionClass):
expr = expr(t)
n = {}
_add_aliases_to_namespace(n, expr)
self.n = n.copy()
from sympy.utilities.lambdify import _get_namespace
for k, v in _get_namespace('numpy').items():
self.n[k] = v
self._f = sympy.lambdify(args, expr, self.n)
self.expr = expr
def get_compiled_array(self) -> Callable:
"""compile the tensor expression such that a numpy array is returned
Note that the input to the returned function must be a single 1d array
with exactly as many entries as there are variables in the expression.
"""
assert isinstance(self._sympy_expr, sympy.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]}] = {val}"
for idx, val in np.ndenumerate(self._sympy_expr)
]
else:
lines = [
f" out[{str(idx + (...,))[1:-1]}] = convert_scalar({val})"
for idx, val in np.ndenumerate(self._sympy_expr)
]
if variables:
# the expression takes variables as input
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 expression is constant
code = "def _generated_function(arr=None, 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
def get_compiled_array(self, single_arg: bool = True) -> Callable:
"""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)
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]}] = {val}"
for idx, val in np.ndenumerate(self._sympy_expr)
]
else:
lines = [
f" out[{str(idx + (...,))[1:-1]}] = convert_scalar({val})"
for idx, val in np.ndenumerate(self._sympy_expr)
]
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
def cse_lambdify(args, expr, **kwargs):
'''
Wrapper for sympy.lambdify which makes use of common subexpressions.
'''
# Note:
# This was expected to speed up the evaluation of the created functions.
# However performance gain is only at ca. 5%
# check input expression
if type(expr) == str:
raise TypeError('Not implemented for string input expression!')
# check given expression
try:
check_expression(expr)
except TypeError as err:
raise NotImplementedError("Only sympy expressions are allowed, yet")
# get sequence of symbols from input arguments
if type(args) == str:
args = sp.symbols(args, seq=True)
elif hasattr(args, '__iter__'):
# this may kill assumptions
args = [sp.Symbol(str(a)) for a in args]
if not hasattr(args, '__iter__'):
args = (args, )
# get the common subexpressions
cse_pairs, red_exprs = sp.cse(expr, symbols=sp.numbered_symbols('r'))
if len(red_exprs) == 1:
red_exprs = red_exprs[0]
# check if sympy found any common subexpressions
if not cse_pairs:
# if not, use standard lambdify
return sp.lambdify(args, expr, **kwargs)
# now we are looking for those arguments that are part of the reduced expression(s)
shortcuts = zip(*cse_pairs)[0]
atoms = sp.Set(red_exprs).atoms()
cse_args = [arg for arg in tuple(args) + tuple(shortcuts) if arg in atoms]
# next, we create a function that evaluates the reduced expression
cse_expr = red_exprs
# if dummify is set to False then sympy.lambdify still returns a numpy.matrix
# regardless of the possibly passed module dictionary {'ImmutableMatrix' : numpy.array}
if kwargs.get('dummify') == False:
kwargs['dummify'] = True
reduced_exprs_fnc = sp.lambdify(args=cse_args, expr=cse_expr, **kwargs)
# get the function that evaluates the replacement pairs
modules = kwargs.get('modules')
if modules is None:
modules = ['math', 'numpy', 'sympy']
namespaces = []
if isinstance(modules, (dict, str)) or not hasattr(modules, '__iter__'):
namespaces.append(modules)
else:
namespaces += list(modules)
nspace = {}
for m in namespaces[::-1]:
nspace.update(_get_namespace(m))
eval_pairs_fnc = make_cse_eval_function(input_args=args,
replacement_pairs=cse_pairs,
ret_filter=cse_args,
namespace=nspace)
# now we can wrap things together
def cse_fnc(*args):
cse_args_evaluated = eval_pairs_fnc(args)
return reduced_exprs_fnc(*cse_args_evaluated)
return cse_fnc
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
def cse_lambdify(args, expr, **kwargs):
'''
Wrapper for sympy.lambdify which makes use of common subexpressions.
'''
# Note:
# This was expected to speed up the evaluation of the created functions.
# However performance gain is only at ca. 5%
# check input expression
if type(expr) == str:
raise TypeError('Not implemented for string input expression!')
# check given expression
try:
check_expression(expr)
except TypeError as err:
raise NotImplementedError("Only sympy expressions are allowed, yet")
# get sequence of symbols from input arguments
if type(args) == str:
args = sp.symbols(args, seq=True)
elif hasattr(args, '__iter__'):
# this may kill assumptions
args = [sp.Symbol(str(a)) for a in args]
if not hasattr(args, '__iter__'):
args = (args,)
# get the common subexpressions
cse_pairs, red_exprs = sp.cse(expr, symbols=sp.numbered_symbols('r'))
if len(red_exprs) == 1:
red_exprs = red_exprs[0]
# check if sympy found any common subexpressions
if not cse_pairs:
# if not, use standard lambdify
return sp.lambdify(args, expr, **kwargs)
# now we are looking for those arguments that are part of the reduced expression(s)
shortcuts = zip(*cse_pairs)[0]
atoms = sp.Set(red_exprs).atoms()
cse_args = [arg for arg in tuple(args) + tuple(shortcuts) if arg in atoms]
# next, we create a function that evaluates the reduced expression
cse_expr = red_exprs
# if dummify is set to False then sympy.lambdify still returns a numpy.matrix
# regardless of the possibly passed module dictionary {'ImmutableMatrix' : numpy.array}
if kwargs.get('dummify') == False:
kwargs['dummify'] = True
reduced_exprs_fnc = sp.lambdify(args=cse_args, expr=cse_expr, **kwargs)
# get the function that evaluates the replacement pairs
modules = kwargs.get('modules')
if modules is None:
modules = ['math', 'numpy', 'sympy']
namespaces = []
if isinstance(modules, (dict, str)) or not hasattr(modules, '__iter__'):
namespaces.append(modules)
else:
namespaces += list(modules)
nspace = {}
for m in namespaces[::-1]:
nspace.update(_get_namespace(m))
eval_pairs_fnc = make_cse_eval_function(input_args=args,
replacement_pairs=cse_pairs,
ret_filter=cse_args,
namespace=nspace)
# now we can wrap things together
def cse_fnc(*args):
cse_args_evaluated = eval_pairs_fnc(args)
return reduced_exprs_fnc(*cse_args_evaluated)
return cse_fnc
def cse_lambdify(args, expr, **kwargs):
"""
Wrapper for sympy.lambdify which makes use of common subexpressions.
Parameters
----------
args : iterable
expr : sympy expression or iterable of sympy expression
return callable
"""
# Notes:
# This was expected to speed up the evaluation of the created functions.
# However performance gain is only at ca. 5%
# constant expressions are handled as well
# check given expression
try:
expr = preprocess_expression(expr)
except TypeError as err:
raise NotImplementedError("Only (sequences of) sympy expressions are allowed, yet")
# get sequence of symbols from input arguments
if type(args) == str:
args = sp.symbols(args, seq=True)
elif hasattr(args, '__iter__'):
# this may kill assumptions
# TODO: find out why this is done an possbly remove
args = [sp.Symbol(str(a)) for a in args]
if not hasattr(args, '__iter__'):
args = (args,)
# get the common subexpressions
symbol_generator = sp.numbered_symbols('r')
cse_pairs, red_exprs = sp.cse(expr, symbols=symbol_generator)
# Note: cse always returns a list because expr might be a sequence of expressions
# However we want only one expression back if we put one in
# (a matrix-object is covered by this)
if len(red_exprs) == 1:
red_exprs = red_exprs[0]
# check if sympy found any common subexpressions
# typically cse_pairs looks like [(r0, cos(x1)), (r1, sin(x1))], ...
if not cse_pairs:
# add a meaningless mapping r0 |-→ 0 to avoid empty list
cse_pairs = [(symbol_generator.next(), 0)]
# now we are looking for those arguments that are part of the reduced expression(s)
# find out the shortcut-symbols
shortcuts = zip(*cse_pairs)[0]
atoms = sp.Set(red_exprs).atoms(sp.Symbol)
cse_args = [arg for arg in tuple(args) + tuple(shortcuts) if arg in atoms]
assert isinstance(cse_pairs[0][0], sp.Symbol)
if len(cse_args) == 0:
# this happens if expr is constant
cse_args = [cse_pairs[0][0]]
# next, we create a function that evaluates the reduced expression
cse_expr = red_exprs
# if dummify is set to False then sympy.lambdify still returns a numpy.matrix
# regardless of the possibly passed module dictionary {'ImmutableMatrix' : numpy.array}
if kwargs.get('dummify') == False:
kwargs['dummify'] = True
reduced_exprs_fnc = sp.lambdify(args=cse_args, expr=cse_expr, **kwargs)
# get the function that evaluates the replacement pairs
modules = kwargs.get('modules')
if modules is None:
modules = ['math', 'numpy', 'sympy']
namespaces = []
if isinstance(modules, (dict, str)) or not hasattr(modules, '__iter__'):
namespaces.append(modules)
else:
namespaces += list(modules)
nspace = {}
for m in namespaces[::-1]:
nspace.update(_get_namespace(m))
eval_pairs_fnc = make_cse_eval_function(input_args=args,
replacement_pairs=cse_pairs,
ret_filter=cse_args,
namespace=nspace)
# now we can wrap things together
def cse_fnc(*args):
# this function is intended only for scalar args
# vectorization is handled by `broadcasting_wrapper`
for a in args:
assert isinstance(a, Number)
cse_args_evaluated = eval_pairs_fnc(args)
return reduced_exprs_fnc(*cse_args_evaluated)
# later we might need the information how many scalar args this function expects
cse_fnc.args_info = args
return cse_fnc
