如何在自定义ufunc中调用NumPy随机数生成器？

嘿，我来帮你搞定这个问题！首先给个明确结论：NumPy的随机数C API是存在的，只不过它属于内部开发接口，没有放在对外的主API文档里，所以容易找不到。接下来我会一步步帮你修正代码，同时解释清楚怎么用这个API来保证你的ufunc随机结果可复现。

一、先搞清楚头文件和核心依赖 #

你之前写的#include <numpy/random.h>是不对的，NumPy随机数C API的头文件分散在numpy/random/目录下，核心需要这几个：

#include <numpy/arrayobject.h>
#include <numpy/random/bitgen.h>
#include <numpy/random/distributions.h>

其中bitgen.h定义了随机数生成器的核心结构体bitgen_t，distributions.h里包含了各种分布的采样函数（比如你需要的均匀分布）。

二、调整你的参数结构体，加入随机数生成器指针 #

要保证可复现性，绝对不能用全局随机状态——很容易被其他代码干扰。正确的做法是让Python层传递一个NumPy随机生成器的实例指针到C层。所以先修改你的basic_gate_argument_t：

typedef struct {
    int act;
    bitgen_t *rng; // 存储NumPy随机生成器的底层指针
} basic_gate_argument_t;

三、修改Ufunc内核函数，正确调用随机采样API #

你之前写的random_uniform(0, 1)是不存在的，NumPy的C API里对应的函数是random_uniform，但需要传入生成器指针。调整后的内核函数如下：

static void ufunc_M(char **args, npy_intp *dimensions, npy_intp *steps, void *data) {
    basic_gate_argument_t argument = *((basic_gate_argument_t *)data);
    PYQCS_GATE_GENERIC_SETUP;
    npy_intp i;
    npy_double amplitude_1 = 0;

    // 你的振幅计算逻辑保持不变
    for(i = 0; i < ndim; i++) {
        if(i & (1 << argument.act)) {
            amplitude_1 += qm_in[i].real * qm_in[i].real;
            amplitude_1 += qm_in[i].imag * qm_in[i].imag;
        }
    }

    // 用NumPy的C API生成[0,1)范围内的均匀随机数
    // 参数依次是：下界、上界、随机生成器指针
    npy_double rand = random_uniform(0.0, 1.0, argument.rng);

    // 后面的逻辑你该怎么写就怎么写
    ...
    *measured_out = 1 << argument.act;
}

四、Python层怎么传递随机生成器？ #

在Python代码里，你需要从NumPy的新随机生成器（比如np.random.default_rng()）中获取底层的C指针，然后传递给你的扩展。举个例子：

import numpy as np
from your_extension_module import ufunc_M

def measured_gate(act, rng=None):
    # 默认创建一个新的生成器，或者用用户传入的
    if rng is None:
        rng = np.random.default_rng()
    # 获取底层bitgen的C指针
    bitgen_ptr = rng.bitgen._cpointer
    # 构造你的参数结构体（这里需要用ctypes或者扩展的包装逻辑）
    args = basic_gate_argument_t(act=act, rng=bitgen_ptr)
    # 调用你的ufunc，把参数传进去
    return ufunc_M(..., data=args)

这样用户在Python层设置种子（比如rng = np.random.default_rng(42)），你的ufunc每次运行的随机结果就完全可复现了。

几个关键提醒 #

1. 别用全局随机状态：旧的np.random.rand()这类全局函数的状态是共享的，其他代码一调用就会打乱你的随机序列，用传递bitgen_t指针的方式能完全隔离状态。
2. 版本兼容性：这个新的随机数C API是NumPy 1.17之后才稳定的，如果你要兼容旧版本NumPy，可能需要额外处理，但现在大部分用户都用新版本了。
3. 找例子的话：直接看NumPy源码里的numpy/random目录，里面的C文件全是用这些API实现的，比文档还直观。

内容的提问来源于stack exchange，提问作者LittleByBlue