Friday 4 December 2020

Can `sympy.codegen.ast` be used with `sympy.utilities.autowrap`?

I'm working on producing compiled functions using sympy.code.autowrap, specifically the make_routine, autowrap, and ufuncify functions of that module. I would like to be able to use sympy.codegen.ast to generate code to be wrapped. For instance:

import sympy as sy
from sympy.codegen.ast import (Assignment, CodeBlock, Return, Element, 
                               Pointer, Declaration, float64,
                               FunctionPrototype, FunctionDefinition)
zlist = sy.symarray('z', (5,), real=True)
d1expz = sy.apply_finite_diff(1, x_list=range(-2,3), y_list=zlist)
d2expz = sy.apply_finite_diff(2, x_list=range(-2,3), y_list=zlist)
d1p = Pointer('d1p', float64)
d2p = Pointer('d2p', float64)
d1p0 = Element('d1p', '0')
d2p0 = Element('d2p', '0')
d12fpa = FunctionPrototype(float64, 'd12funca', [d1p, d2p] + list(zlist))
body12a = [
    CodeBlock(
        Assignment(d1p0, d1expz),
        Assignment(d2p0, d2expz),
    ).cse(),
]
d12funca = FunctionDefinition.from_FunctionPrototype(
    d12fpa, body12a
)
print(sy.ccode(d12funca))

This results in the following C code, which looks good:

double d12funca(double * d1p, double * d2p, double z_0, double z_1, double z_2, double z_3, double z_4){
   x0 = (1.0/12.0)*z_0;
   x1 = -1.0/12.0*z_4;
   d1p[0] = x0 + x1 - 2.0/3.0*z_1 + (2.0/3.0)*z_3;
   d2p[0] = -x0 + x1 + (4.0/3.0)*z_1 - 5.0/2.0*z_2 + (4.0/3.0)*z_3;
}

sympy.utilities.codegen.codegen also understands this, e.g. the following produces just the kind of results you'd expect.

[(c_name, c_code), (h_name, c_header)] = codegen(
    ('d12func',
    [
        sy.Eq(sy.Symbol('d1'), d1expz),
        sy.Eq(sy.Symbol('d2'), d2expz)
    ]),
    language='c'
)

I can also produce a nicely wrapped function that does essentially the same thing as follows:

from sympy.utilities.autowrap import autowrap
autowrap_dir = <some path>
derivs = sy.Matrix(
    [[
        d1expz,
        d2expz
    ]]
)
dvraw = autowrap(
    expr=derivs,
    args=list(zlist),
    backend='Cython',
    tempdir=autowrap_dir,
    verbose=True
)

But is there any way to use the ast module to specify functions to be wrapped with autowrap? It appears that autowrap, make_routine, and ufuncify cannot accept arbitrary code. The expressions they take can be either lists of straightforward sympy expressions, or lists of sympy.Eq expressions (in which case the wrapped function assigns the RHS of each equality to its LHS and returns each LHS as an output argument).

The problem is that I would like to be able to wrap more involved functions including such things as for loops and branching. The ast module provides the tools to write such functions. But as far as I can tell, autowrap and ast don't play with each other at all.

Any thoughts on how to accomplish this would be appreciated. Thanks.



from Can `sympy.codegen.ast` be used with `sympy.utilities.autowrap`?

No comments:

Post a Comment