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 __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