python - 将 NumPy 中定义的函数转换为 SymPy

Question

我在 numpy 中定义了一个函数，我想将其转换为 sympy，因此我可以将其应用于符号 sympy 变量。尝试将 numpy 函数直接应用于 sympy 变量失败:

import numpy as np
import sympy as sp

def np_fun(a):
    return np.array([np.sin(a), np.cos(a)])

x = sp.symbols('x')
sp_fun = np_fun(x)

我得到了错误

AttributeError: 'Symbol' object has no attribute 'sin'

我的下一个想法是将 numpy 函数转换为 sympy，但我找不到这样做的方法。我知道我可以通过将函数定义为 sympy 表达式来使这段代码工作:

sp_fun = sp.Array([sp.sin(x), sp.cos(x)])

但我使用正弦/余弦函数作为一个简单的例子。我使用的实际函数已经在 numpy 中定义了，而且要复杂得多，因此重写它会非常繁琐。

Answer 1

原则上，您可以直接修改函数的 ast(“抽象语法树”)，尽管在实践中它可能会变得很复杂。无论如何，这是针对您的简单示例的操作方法:

这从源创建一个 ast 并派生自 NodeTransformer 类以就地修改 ast。节点转换器有一个通用的访问方法，它遍历一个节点及其子树，委托(delegate)给派生类中的节点特定访问者。在这里，我们将所有名称 np 更改为 sp，然后将这些属性更改为以前的 np 现在 sp 拼写不同。您必须将所有这些差异添加到 translate 字典中。

最后，我们从 ast 编译回代码对象并执行它以使修改后的函数可用。

import ast, inspect
import numpy as np
import sympy as sp

def f(a):
    return np.array([np.sin(a), np.cos(a)])

z = ast.parse(inspect.getsource(f))

translate = {'array': 'Array'}

class np_to_sp(ast.NodeTransformer):
    def visit_Name(self, node):
        if node.id=='np':
            node = ast.copy_location(ast.Name(id='sp', ctx=node.ctx), node)
        return node
    def visit_Attribute(self, node):
        self.generic_visit(node)
        if node.value.id=='sp' and node.attr in translate:
            fields = {k: getattr(node, k) for k in node._fields}
            fields['attr'] = translate[node.attr]
            node = ast.copy_location(ast.Attribute(**fields), node)
        return node

np_to_sp().visit(z)

exec(compile(z, '', 'exec'))

x = sp.Symbol('x')
print(f(x))

输出:

[sin(x), cos(x)]

UPDATE简单增强:修改函数调用的函数:

import ast, inspect
import numpy as np
import sympy as sp

def f(a):
    return np.array([np.sin(a), np.cos(a)])

def f2(a):
    return np.array([1, np.sin(a)])

def f3(a):
    return f(a) + f2(a)

translate = {'array': 'Array'}

class np_to_sp(ast.NodeTransformer):
    def visit_Name(self, node):
        if node.id=='np':
            node = ast.copy_location(ast.Name(id='sp', ctx=node.ctx), node)
        return node
    def visit_Attribute(self, node):
        self.generic_visit(node)
        if node.value.id=='sp' and node.attr in translate:
            fields = {k: getattr(node, k) for k in node._fields}
            fields['attr'] = translate[node.attr]
            node = ast.copy_location(ast.Attribute(**fields), node)
        return node

from types import FunctionType

for fn in f3.__code__.co_names:
    fo = globals()[fn]
    if not isinstance(fo, FunctionType):
        continue
    z = ast.parse(inspect.getsource(fo))
    np_to_sp().visit(z)
    exec(compile(z, '', 'exec'))

x = sp.Symbol('x')
print(f3(x))

打印:

[sin(x) + 1, sin(x) + cos(x)]