Можно ли использовать sympy.codegen.ast с sympy.utilities.autowrap?

Question

Я работаю над созданием скомпилированных функций, используя 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 в качестве выходного аргумента).

Проблема в том, что я хотел бы иметь возможность обернуть более сложные функции, включая такие, как циклы и ветвления. Модуль ast предоставляет инструменты для написания таких функций. Но, насколько я могу судить, autowrap и ast вообще не играют друг с другом.

Буду признателен за любые мысли о том, как это сделать sh. Спасибо.