如何在使用 LowLevelCallable cfunc 时使用 scipy.integrate.nquad 并向其传递参数

How to use scipy.integrate.nquad and passing arguments to it while using LowLevelCallable cfunc

我非常渴望实现一个集成,我想将其结果用作 PDE 的初始数据。积分本身是标量值的,需要在整个 space $\mathbb{R}^3$ 上进行评估。首先,我不期望非常快速、超级容易的集成,但希望它能在合理的时间内完成。因此,我决定在 LowLvelCallable 选项中使用 nquad。我正在尝试在 Python.

中执行我的代码

scipy 的 LowLevelCallable 选项以较少的开销调用 cfuncs,以节省每个单独集成步骤的时间。现在,我在 quad is used instead of nquad, Sample 中找到了一个用于我的目的的示例实现。据我所知,这个示例的想法是使用装饰器来“欺骗” quad 的签名,这需要输入 sig = (types.CPointer(types.double), types.CPointer(types.void)) 形式的 cfunc 来绕过存在的问题当代码以 Python.

编写时,不可能从 viod 类型中提取数据

效果很好,但现在我需要为 nquad 调整这种方法,原则上应该没有任何改变,唯一的区别是 nquad 需要 [=20= 形式的签名].不幸的是,事实证明仅仅编辑 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]

Unfortunately, it turns out that just to edit the signature of the cfunc doesn't do the trick, … My hope is that someone can help me out and saying me what my mistake is.

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

    @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的末尾;您必须查看此值的来源。