如何在使用 LowLevelCallable cfunc 时使用 scipy.integrate.nquad 并将参数传递给它

Posted 2023-03-24

【中文标题】如何在使用 LowLevelCallable cfunc 时使用 scipy.integrate.nquad 并将参数传递给它

【英文标题】：How to use scipy.integrate.nquad and passing arguments to it while using LowLevelCallable cfunc

【发布时间】：2021-12-19 08:24:29

【问题描述】：

我非常绝望地尝试实现一个集成，我想将其用作 PDE 的初始数据。积分本身是标量值，需要在整个空间 $\mathbbR^3$ 上求值。首先，我不期望一个非常快速超级简单的集成，但希望它可以在合理的时间内完成。因此我决定使用nquad 连接LowLvelCallable 选项。我正在尝试在 Python 中实现我的代码。

scipy 的 LowLevelCallable 选项以更少的开销调用 cfuncs，以节省每个单独集成步骤的时间。现在我为我的目的找到了一个示例实现，其中使用quad 而不是nquad、Sample。据我了解，此示例的想法是使用装饰器来“欺骗”quad 的签名，这需要以sig = (types.CPointer(types.double), types.CPointer(types.void)) 的形式输入 cfunc 以绕过没有当代码用 Python 编写时，可以从 viod 类型中提取数据。

这很好用，但现在我需要为nquad 调整这种方法，原则上什么都不应该改变，唯一的区别是nquad 需要sig=(types.intc, types.CPointer(types.double), types.CPointer(types.void)) 形式的签名。不幸的是，事实证明，仅仅编辑 cfunc 的签名并不能解决问题，因此在没有任何高性能计算和 C 语言背景的情况下，这变成了一个非常讨厌的尝试和错误游戏。我希望有人可以帮助我并告诉我我的错误是什么。

这是我的代码示例

import numpy as np
import numba as nb
from numba import types
from scipy import integrate, LowLevelCallable
import ctypes
from math import sqrt, exp

#define some constants
c = 137
sigma = 10
p_boost = 100
m = 1
gamma = sqrt(1 + p_boost ** 2 / c ** 2)
beta_x = p_boost / (gamma * m * c)

x = np.linspace(-4, 4, 32, endpoint = False)
y = np.linspace(-4, 4, 32, endpoint = False)
z = np.linspace(-4, 4, 32, endpoint = False)

#define arg data type
args_dtype = types.Record.make_c_struct([
    ('x_coord', types.float64),
    ('y_coord', types.float64),
    ('z_coord', types.float64),])


def create_jit_integrand_function(integrand_function,args,args_dtype):
    jitted_function = nb.njit(integrand_function)

    @nb.cfunc(types.intc,types.float64(types.float64,types.CPointer(args_dtype)))
    def wrapped(n,x1,user_data_p):
        #Array of structs
        user_data = nb.carray(user_data_p, 1)

        #Extract arg data
        x_co=user_data[0].x_co
        y_co=user_data[0].y_co
        z_co=user_data[0].z_co

        return jitted_function(n,x1, x_co, y_co, z_co)
    return wrapped

def function_using_arrays(n ,p, x_co, y_co, z_co):
    p_square = 0.0
    for i in range(n):
        p_square += p[i]**2
    rel_energy = sqrt(p_square + 1)
    px_new = beta_x * gamma * rel_energy / c + gamma * p[0]
    p_square_new =  px_new**2 + p[1]**2 + p[2]**2
    rel_energy_new = sqrt(p_square_new +1)
    return exp(-p_square_new/(2*sigma**2) - 1j * (p[0]*x_co + p[1]*y_co + p[2]*z_co)) * (rel_energy_new - beta_x * p[boost_index] + 0j)**(1/2) * rel_energy_new ** (-1)


def jit_with_dummy_data(args,args_dtype):
    func=create_jit_integrand_function(function_using_arrays,args,args_dtype)
    return func

def do_integrate_w_arrays(func,args,ranges):
    integrand_func=LowLevelCallable(func.ctypes,user_data=args.ctypes.data_as(ctypes.c_void_p))
    return integrate.nquad(integrand_func, ranges)


output = np.zeros((len(x),len(y),len(z)), dtype=complex)
for i in range(len(x)):
    for j in range(len(y)):
        for k in range(len(z)):
            #creating args
            x_co = x[i]
            y_co = y[j]
            z_co = z[k]
            #hand over args
            args=np.array((x_co, y_co, z_co),dtype=args_dtype)
            func=jit_with_dummy_data(args,args_dtype)
            #perform integration
            ranges = ((-np.inf, np.inf), (-np.inf, np.inf), (-np.inf, np.inf))
            output[i,j,k] = do_integrate_w_arrays(func,args,ranges)[0]

【参考方案1】：

不幸的是，事实证明仅仅编辑 cfunc 的签名并不能解决问题，......我希望有人可以帮助我并告诉我我的错误是什么。

错误是在函数签名之前插入了额外的整数参数，而不是插入到参数列表中；正确：

    @nb.cfunc(types.float64(types.intc, types.float64, types.CPointer(args_dtype)))

此更正揭示了wrapped 中的命名不匹配；正确：

        x_co=user_data[0].x_coord
        y_co=user_data[0].y_coord
        z_co=user_data[0].z_coord

该更正允许程序继续处理下一个错误：

NameError: name 'boost_index' is not defined

boost_index用在function_using_arrays的末尾；你得看看这个值是从哪里来的。