对两个数组进行广义ufunc，其中一个维度不匹配

Question

我想用numpy的C接口写一个扩展，这个扩展可以处理任意维度的两个矩阵。它的功能是计算一个维度上的平均值，然后把一个结果减去另一个结果。最简单的情况是两个一维数组，格式是'(i), (j) -> ()'，这样会返回一个标量值。

如果有额外的维度，我们可以假设它们是核心维度，也就是说它们是相同的。例如，一个二维数组的格式可能是'(n, i), (n, j) -> (n)'，这里的核心维度是轴0，这样就可以在这个轴上循环执行这个函数。以下是我目前的代码：

#include <Python.h>
#include <numpy/arrayobject.h>
#include <numpy/ufuncobject.h>
#include <math.h>

static void mean_diff(char **args,
                      const npy_intp *dimensions,
                      const npy_intp* steps,
                      void* extra) {
    npy_intp i;
    npy_intp n1 = 0, n2 = 0;
    double s1 = 0.0, s2 = 0.0;
    char *in1 = args[0], *in2 = args[1], *out = args[2];
    npy_intp len1 = dimensions[0], len2 = dimensions[1];

    for (i = 0; i < len1; i++) {
        double val = *((double *)in1);
        if (!isnan(val)) {
            s1 += val;
            n1++;
        }
        in1 += steps[0];
    }

    for (i = 0; i < len2; i++) {
        double val = *((double *)in2);
        if (!isnan(val)) {
            s2 += val;
            n2++;
        }
        in2 += steps[1];
    }

    double mean1 = (n1 > 0) ? s1 / n1 : 0.0;
    double mean2 = (n2 > 0) ? s2 / n2 : 0.0;

    *((double *)out) = mean1 - mean2;
}

我遇到的问题是，即使是简单的两个一维数组，它们只考虑每个数组的第一个元素。例如：

>>> mean_diff([1.,2.,3.,4.], [2.,7.,29.,3.])
-1.0
>>> np.mean([1.,2.,3.,4.]) - np.mean([2.,7.,29.,3.])
-7.75

我猜这可能是因为抓取了错误的维度之类的原因，但我搞不清楚具体是什么问题。

Answer 1

首先要知道的是，你的函数在每次调用时会计算多次结果。计算的次数存储在 dimensions[0] 中。你指定了 gufunc 的签名 (i),(j)->()。当函数被调用时，dimensions[1] 和 dimensions[2] 将分别保存 i 和 j 的值，也就是核心操作的输入数组的长度。

举个例子，在这个代码片段中，

x = np.array([[0, 1, 2, 3, 4], [0, 0, 0, 1, 1.5]])  # shape (2, 5)
y = np.array([[-1, -3, 7], [2, 2, 2]])              # shape (2, 3)
result = mean_diff(x, y)                            # shape (2,)

你的函数会被调用 一次，此时 dimensions[0] 被设置为 2。dimensions[1] 和 dimensions[2] 分别为 5 和 3。你的外层循环会计算结果，首先是对于这对数据 ([0, 1, 2, 3, 4], [-1, -3, 7])，然后是 ([0, 0, 0, 1, 1.5], [2, 2, 2])。

args[0] 和 args[1] 将指向输入数组的第一个元素，而 args[2] 将指向输出数组的第一个元素。你的代码中有 in1 = args[0]，in2 = args[1] 和 out = args[2]，所以它必须计算 in1 和 in2 指向的数据的结果，并将结果存储在 *out 中。然后，它必须将 in1、in2 和 out 增加，以指向下一组数据，并再次进行计算。因此，你需要实现一个外层循环，针对每对输入进行计算，并增加指针以移动到下一组输入和输出。增加 in1、in2 和 out 的量存储在 steps[0]、steps[1] 和 steps[2] 中。（例如，steps[0] 是从 x 的第一行跳到第二行所需的字节数。）

然后，在这个外层循环中，你会对当前的 x 和 y 数组对进行计算（由 in1 和 in2 指向）。一般来说，你不能假设数组的元素是连续的，所以在计算均值时，你必须按适当的步长增加指向数据的指针；在这种情况下，这些步长是 steps[3] 和 steps[4]。

这里有一种实现方式。我在代码中注释了从 dimensions 和 steps 中提取的变量，以帮助澄清我上面简短描述的内容。

static void mean_diff(
    char **args,
    const npy_intp *dimensions,
    const npy_intp *steps,
    void *extra)
{
    char *in1 = args[0];
    char *in2 = args[1];
    char *out = args[2];

    npy_intp nloops = dimensions[0];  // Number of outer loops
    npy_intp len1 = dimensions[1];    // Core dimension i
    npy_intp len2 = dimensions[2];    // Core dimension j

    npy_intp step1 = steps[0];        // Outer loop step size for the first input
    npy_intp step2 = steps[1];        // Outer loop step size for the second input
    npy_intp step_out = steps[2];     // Outer loop step size for the output
    npy_intp innerstep1 = steps[3];   // Step size of elements within the first input
    npy_intp innerstep2 = steps[4];   // Step size of elements within the second input

    for (npy_intp i = 0; i < nloops;
         i++, in1 += step1, in2 += step2, out += step_out) {

        // The core calculation.  in1 and in2 point to the 1-d input arrays,
        // and out points to the output array.

        double s1 = 0.0;
        double s2 = 0.0;
        npy_intp n1 = 0;
        npy_intp n2 = 0;

        for (npy_intp j = 0; j < len1; ++j) {
            double val = *(double *)(in1 + j*innerstep1);
            if (!isnan(val)) {
                s1 += val;
                n1++;
            }
        }

        for (npy_intp j = 0; j < len2; ++j) {
            double val = *(double *)(in2 + j*innerstep2);
            if (!isnan(val)) {
                s2 += val;
                n2++;
            }
        }

        double mean1 = (n1 > 0) ? s1 / n1 : 0.0;
        double mean2 = (n2 > 0) ? s2 / n2 : 0.0;

        *((double *)out) = mean1 - mean2;
    }
}

顺便说一下，如果像这样实现的 gufunc 实际上比使用 mean 方法计算均值并减去结果要 慢，不要感到惊讶。NumPy 不断在提升性能，最近的 SIMD 实现使得简单的 ufunc 和 gufunc 实现可能比几个优化过的 NumPy 方法的组合要慢。