4

我正在使用该模块sympy.code.autowrap的make_routine、autowrap和函数来生成编译函数。ufuncify我希望能够用来sympy.codegen.ast生成要包装的代码。例如：

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

这会产生以下 C 代码，看起来不错：

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也理解这一点，例如，以下产生的结果正是您所期望的。

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

我还可以生成一个包装精美的函数，它的作用基本相同，如下所示：

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
)

但是有什么方法可以使用ast模块来指定要包装的函数autowrap吗？看来autowrap,make_routine和ufuncify不能接受任意代码。他们采用的表达式可以是直接sympy表达式列表，也可以是表达式列表sympy.Eq（在这种情况下，包装函数将每个等式的 RHS 分配给其 LHS 并将每个 LHS 作为输出参数返回）。

问题是我希望能够包装更多涉及的功能，包括 for 循环和分支等。该ast模块提供了编写此类函数的工具。但据我所知，autowrap根本ast不互相玩。

任何关于如何实现这一点的想法将不胜感激。谢谢。

4

0 回答 0