sympy

Printing functions and their derivatives in SymPy as function calls in the target language

I have this. It prints Math.sin(Math.pow(JSFunctions.f(x), 2)). The user_functions dictionary instructs jscode to format expressions such as f(x) as calls to JSFunctions.f. I am targetting Javascript but I believe it should be the same with other targets.

from sympy import *
from sympy.printing.jscode import jscode

x = Symbol('x')
f = Function('f')
g = sin(f(x) ** 2)

code = jscode(g, user_functions={ 'f': 'JSFunctions.f' })

print(code)

Now, instead of generating the code of g, I want to generate the code for Dg, the first derivative of g. Conceptually, I want to do the following but it does not work because user_functions must have function names as keys.

from sympy import *
from sympy.printing.jscode import jscode

x = Symbol('x')
f = Function('f')
g = sin(f(x) ** 2)
Dg = diff(g, x)

code = jscode(Dg, user_functions={ f: 'JSFunctions.f', diff(f(x), x): 'JSFunctions.Df' })

print(code)

I have been able to work around this by introducing another symbolic function Df, and using it to replace the occurrences of diff(f(x)), x) in Dg. This prints 2*JSFunctions.Df(x)*JSFunctions.f(x)*Math.cos(Math.pow(JSFunctions.f(x), 2)).

from sympy import *
from sympy.printing.jscode import jscode

x = Symbol('x')
f = Function('f')
g = sin(f(x) ** 2)
Dg = diff(g, x)
Df = Function('Df')

code = jscode(Dg.replace(diff(f(x), x), Df(x)), user_functions={ 'f': 'JSFunctions.f', 'Df': 'JSFunctions.Df' })

print(code)

Are there other ways to achieve the same output? Ideally, built-in ways, but I am open to implementing my own custom printer, possibly as a subclass of JavascriptCodePrinter.

Edit

Well, I guess another workaround is replace f and its derivative as symbols and then generate some extra code in advance to initialize the corresponding variables.

from sympy import *
from sympy.printing.jscode import jscode

x = Symbol('x')
f = Function('f')
g = sin(f(x) ** 2)
Dg = diff(g, x)

code = ""
code += "const f_x = JSFunctions.f(x);\n"
code += "const Df_x = JSFunctions.Df(x);\n"
code += jscode(Dg.replace(diff(f(x), x), Symbol('Df_x')).replace(f(x), Symbol('f_x')))

print(code)

I would still love a cleaner way to emit code in these cases.

Solution

  • Here is how you could modify the printer:

    from sympy.printing.jscode import JavascriptCodePrinter
from sympy.core.function import AppliedUndef

class MyJavascriptCodePrinter(JavascriptCodePrinter):
    def _print_Function(self, expr):
        if isinstance(expr, AppliedUndef):
            return "JSFunctions." + expr.func.__name__ + f"({", ".join([str(a) for a in expr.args])})"
        return super()._print_Function(expr)
        
    def _print_Derivative(self, expr):
        if isinstance(expr.args[0], AppliedUndef):
            # Usually you would do this:
            # return "JSFunctions.D" + self._print(expr.args[0])
            # But in this particular case you'd have to do something like this
            return "JSFunctions.D" + expr.args[0].func.__name__ + f"({", ".join([str(a) for a in expr.args[0].args])})"
        raise NotImplementedError

MyJavascriptCodePrinter().doprint(Dg)

# '2*JSFunctions.f(x)*Math.cos(Math.pow(JSFunctions.f(x), 2))*JSFunctions.Df(x)'

    But you might end up spending a lot of time customizing it for particular cases not shown in your example. I'd probably go with your last attempt, because it seems to be less messy and easier to read...