deeplearning/modulus/modulus-sym-v130/_modules/modulus/sym/utils/sympy/numpy_printer.html
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
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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
NP_LAMBDA_STORE = {}
[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([x.name 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, axis=None):
return_value = x[0]
for value in x:
return_value = np.minimum(return_value, value)
return return_value
def _max_np(x, axis=None):
return_value = x[0]
for value in x:
return_value = np.maximum(return_value, value)
return return_value
def _heaviside_np(x, x2=0):
# force x2 to 0
x2 = 0
return np.heaviside(x, x2)
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]