Convert a function defined in NumPy to SymPy

Question

I have a function defined in numpy which I would like to convert to sympy, so I can apply it to symbolic sympy variables. Trying to directly apply the numpy function to a sympy variable fails:

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)

I get the error

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

My next thought was to convert the numpy function to sympy, but I couldn't find a way to do that. I know I could make this code work by just defining the function as a sympy expression:

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

But I'm using the sine/cosine function as a simple example. The actual function I'm using has already been defined in numpy, and is much more complicated, so it would be very tedious to rewrite it.

Answer 1

In principle, you could directly modify the ast ("abstract syntax tree") of the function, though in practice it might get quite hairy. Anyway, here is how to do it for your simple example:

This creates from the source an ast and derives from the NodeTransformer class to modify the ast in-place. The node transformer has a generic visit method that traverses a node and its subtree delegating to node specific visitors in derived classes. Here we change all names np to sp and afterwards change those attributes to former np now sp that spell differently. You'd have to add all such differences to the translate dict.

Finally, we compile back from the ast to a code object and execute it to make the modified function available.

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

Output:

[sin(x), cos(x)]

UPDATE simple enhancement: modify functions called by function:

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

Prints:

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