Source code for modulus.sym.utils.sympy.numpy_printer

# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

Helper functions for converting sympy equations to numpy

import types
import inspect
import numpy as np
import symengine as se
import sympy as sp


[docs]def np_lambdify(f, r): """ generates a numpy function from a sympy equation Parameters ---------- f : Sympy Exp, float, int, bool or list of the previous the equation to convert to a numpy function. If float, int, or bool this gets converted to a constant function of value `f`. If f is a list then output for each element in list is is concatenated on axis -1. r : list, dict A list of the arguments for `f`. If dict then the keys of the dict are used. Returns ------- np_f : numpy function """ # possibly lambdify list of f if not isinstance(f, list): f = [f] # convert r to a list if dictionary # break up any tuples to elements in list if isinstance(r, dict): r = list(r.keys()) no_tuple_r = [] for key in r: if isinstance(key, tuple): for k in key: no_tuple_r.append(k) else: no_tuple_r.append(key) # lambidfy all functions in list lambdify_f = [] for f_i in f: # check if already a numpy function if isinstance(f_i, types.FunctionType): # add r inputs to function args = inspect.getargspec(f_i).args def lambdify_f_i(**x): return f_i(**{key: x[key] for key in args}) else: # check if already lambdified equation if (f_i, tuple(no_tuple_r)) in NP_LAMBDA_STORE.keys(): lambdify_f_i = NP_LAMBDA_STORE[(f_i, tuple(no_tuple_r))] else: # if not lambdify it try: if not isinstance(f_i, bool): f_i = float(f_i) except: pass if isinstance(f_i, (float, int)): # constant function def loop_lambda(constant): return ( lambda **x: np.zeros_like(next(iter(x.items()))[1]) + constant ) lambdify_f_i = loop_lambda(f_i) elif type(f_i) in [ type((se.Symbol("x") > 0).subs(se.Symbol("x"), 1)), type((se.Symbol("x") > 0).subs(se.Symbol("x"), -1)), bool, ]: # TODO hacky sympy boolian check def loop_lambda(constant): if constant: return lambda **x: np.ones_like( next(iter(x.items()))[1], dtype=bool ) else: return lambda **x: np.zeros_like( next(iter(x.items()))[1], dtype=bool ) lambdify_f_i = loop_lambda(f_i) else: try: # first try to compile with Symengine kk = [] for k in no_tuple_r: if isinstance(k, str): kk.append(se.Symbol(k)) else: kk.append(k) kk = [se.Symbol(name) for name in sorted([ for x in kk])] se_lambdify_f_i = se.lambdify(kk, [f_i], backend="llvm") def lambdify_f_i(**x): if len(x) == 1: v = list(x.values())[0] else: v = np.stack( [v for v in dict(sorted(x.items())).values()], axis=-1, ) out = se_lambdify_f_i(v) if isinstance(out, list): out = np.concatenate(out, axis=-1) return out except: # fall back on older SymPy compile sp_lambdify_f_i = sp.lambdify( [k for k in no_tuple_r], f_i, [NP_SYMPY_PRINTER, "numpy"] ) def lambdify_f_i(**x): v = sp_lambdify_f_i(**x) if isinstance(v, list): v = np.concatenate(v, axis=-1) return v # add new lambdified function to dictionary NP_LAMBDA_STORE[(f_i, tuple(no_tuple_r))] = lambdify_f_i # add new list of lambda functions lambdify_f.append(lambdify_f_i) # construct master lambda function for all def loop_grouped_lambda(lambdify_f): def grouped_lambda(**invar): output = [] for lambdify_f_i in lambdify_f: output.append(lambdify_f_i(**invar)) return np.concatenate(output, axis=-1) return grouped_lambda return loop_grouped_lambda(lambdify_f)

def _xor_np(x): return np.logical_xor(x) def _min_np(x): return_value = x[0] for value in x: return_value = np.minimum(return_value, value) return return_value def _max_np(x): return_value = x[0] for value in x: return_value = np.maximum(return_value, value) return return_value def _heaviside_np(x): return np.heaviside(x, 0) def _equal_np(x, y): return np.isclose(x, y) NP_SYMPY_PRINTER = { "amin": _min_np, "amax": _max_np, "Heaviside": _heaviside_np, "equal": _equal_np, "Xor": _xor_np, } SYMENGINE_BLACKLIST = [sp.Heaviside, sp.DiracDelta]

© Copyright 2023, NVIDIA Modulus Team. Last updated on Apr 27, 2023.