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