自定义 `ufunc`：用 Python 或 C 编写自己的通用函数

好嘞！准备好开启一场关于自定义 ufunc 的奇妙冒险了吗？系好安全带，我们要出发啦！🚀

自定义 `ufunc`：化腐朽为神奇，打造你的专属 NumPy 超能力！

各位观众，欢迎来到今天的“NumPy 超能力觉醒”讲座！今天，我们要聊聊一个能让你在 NumPy 的世界里呼风唤雨的秘密武器：自定义 ufunc！

你是否曾遇到过这样的窘境：NumPy 自带的函数不够用，想要实现一些奇特的、定制化的运算，却发现无从下手？别担心！自定义 ufunc 就是你的救星，它能让你像魔法师一样，创造出属于自己的 NumPy 函数，让数据处理变得更加高效、灵活、有趣！

什么是 `ufunc`？

首先，我们来认识一下 ufunc。ufunc，全称 Universal Function，即通用函数。它是 NumPy 的核心组成部分，负责执行数组间的元素级运算。简单来说，ufunc 就像一个超级计算器，可以对数组中的每一个元素进行相同的操作，然后将结果返回到一个新的数组中。

NumPy 已经内置了大量的 ufunc，比如加法 (np.add)、减法 (np.subtract)、乘法 (np.multiply)、除法 (np.divide)、三角函数 (np.sin, np.cos, np.tan) 等等。这些 ufunc 都是经过高度优化的，能够以极快的速度处理大规模数据。

但是，人有悲欢离合，月有阴晴圆缺，NumPy 内置的 ufunc 也不是万能的。总有一些特殊的运算，需要我们自己动手来实现。这时候，自定义 ufunc 就派上用场了！

为什么要自定义 `ufunc`？

你可能会问：NumPy 已经够强大了，为什么还要自定义 ufunc 呢？原因很简单：

满足特定需求： 有些运算是 NumPy 没有提供的，比如一些特殊的数学函数、统计指标、图像处理算法等等。
提高代码效率： 对于一些复杂的运算，使用 Python 循环可能会非常慢。而自定义 ufunc 可以利用 NumPy 的底层优化，大幅提高运算速度。
代码可读性： 将复杂的运算封装成 ufunc，可以使代码更加简洁、易懂。
装X： 咳咳，好吧，这可能不是最重要的原因。但是，当你能够轻松地自定义 ufunc，并用它们解决实际问题时，那种成就感是无法言喻的！😎

如何自定义 `ufunc`？

自定义 ufunc 有两种主要方法：

使用 Python (Numba): 这是最简单的方法，只需要编写一个 Python 函数，然后使用 Numba 的 @vectorize 装饰器将其转换为 ufunc。
使用 C API: 这种方法比较复杂，需要编写 C 代码，并使用 NumPy 的 C API 将其编译成共享库。但是，它可以实现更高的性能。

接下来，我们将分别介绍这两种方法。

1. 使用 Python (Numba)

Numba 是一个 Python 的即时 (JIT) 编译器，它可以将 Python 代码编译成机器码，从而提高运行速度。Numba 的 @vectorize 装饰器可以将一个普通的 Python 函数转换为 ufunc。

步骤：

安装 Numba：
```
pip install numba
```
编写 Python 函数： 编写一个函数，用于实现你想要的运算。这个函数应该接受标量输入，并返回标量输出。
使用 @vectorize 装饰器： 在函数定义之前，加上 @vectorize 装饰器。可以指定输入和输出的数据类型。
调用 ufunc： 像调用普通的 NumPy 函数一样，调用你的自定义 ufunc。

示例：

我们来创建一个 ufunc，用于计算 sigmoid 函数：

import numpy as np
from numba import vectorize

@vectorize([ 'float64(float64)'])  # 指定输入和输出的数据类型
def sigmoid(x):
  """
  计算 sigmoid 函数：1 / (1 + exp(-x))
  """
  return 1.0 / (1.0 + np.exp(-x))

# 测试
x = np.array([-2, -1, 0, 1, 2])
y = sigmoid(x)
print(y)

代码解释：

@vectorize(['float64(float64)'])：这个装饰器告诉 Numba，我们要将 sigmoid 函数转换为 ufunc，并且输入和输出的数据类型都是 float64。
sigmoid(x)：这个函数实现了 sigmoid 函数的计算。

优点：

简单易用
代码可读性高
无需编写 C 代码

缺点：

性能不如 C API
依赖 Numba 库

2. 使用 C API

使用 C API 自定义 ufunc 可以实现更高的性能，但也更加复杂。

步骤：

编写 C 代码： 编写 C 代码，实现你想要的运算。需要定义一个函数，用于处理单个元素之间的运算。
编写初始化函数： 编写一个初始化函数，用于注册你的 ufunc。
编译 C 代码： 将 C 代码编译成共享库。
在 Python 中导入共享库： 使用 ctypes 模块或者 numpy.ctypeslib 模块导入共享库。
创建 ufunc 对象： 使用 NumPy 的 frompyfunc 函数创建一个 ufunc 对象。

示例：

我们来创建一个 ufunc，用于计算两个数的平方和：

1. C 代码 (square_sum.c):

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

// 函数：计算两个数的平方和
static void square_sum(double *in1, double *in2, double *out, int n, void *data) {
    int i;
    for (i = 0; i < n; i++) {
        out[i] = in1[i] * in1[i] + in2[i] * in2[i];
    }
}

// NumPy ufunc 包装器
static PyObject *square_sum_ufunc(PyObject *self, PyObject *args) {
    PyObject *in1_obj, *in2_obj, *out_obj;
    PyArrayObject *in1_array, *in2_array, *out_array;
    int n;

    // 解析参数
    if (!PyArg_ParseTuple(args, "OOO", &in1_obj, &in2_obj, &out_obj)) {
        return NULL;
    }

    // 将输入和输出对象转换为 NumPy 数组
    in1_array = (PyArrayObject *)PyArray_FROM_O(in1_obj);
    in2_array = (PyArrayObject *)PyArray_FROM_O(in2_obj);
    out_array = (PyArrayObject *)PyArray_FROM_O(out_obj);

    if (in1_array == NULL || in2_array == NULL || out_array == NULL) {
        Py_XDECREF(in1_array);
        Py_XDECREF(in2_array);
        Py_XDECREF(out_array);
        return NULL;
    }

    // 检查数组维度和数据类型
    if (PyArray_NDIM(in1_array) != 1 || PyArray_NDIM(in2_array) != 1 || PyArray_NDIM(out_array) != 1) {
        PyErr_SetString(PyExc_ValueError, "Input and output arrays must be 1-dimensional.");
        Py_DECREF(in1_array);
        Py_DECREF(in2_array);
        Py_DECREF(out_array);
        return NULL;
    }

    if (PyArray_TYPE(in1_array) != NPY_DOUBLE || PyArray_TYPE(in2_array) != NPY_DOUBLE || PyArray_TYPE(out_array) != NPY_DOUBLE) {
        PyErr_SetString(PyExc_ValueError, "Input and output arrays must be of type float64.");
        Py_DECREF(in1_array);
        Py_DECREF(in2_array);
        Py_DECREF(out_array);
        return NULL;
    }

    // 获取数组大小
    n = PyArray_DIM(in1_array, 0);
    if (n != PyArray_DIM(in2_array, 0) || n != PyArray_DIM(out_array, 0)) {
        PyErr_SetString(PyExc_ValueError, "Input and output arrays must have the same size.");
        Py_DECREF(in1_array);
        Py_DECREF(in2_array);
        Py_DECREF(out_array);
        return NULL;
    }

    // 调用 C 函数
    square_sum((double *)PyArray_DATA(in1_array), (double *)PyArray_DATA(in2_array), (double *)PyArray_DATA(out_array), n, NULL);

    // 释放资源
    Py_DECREF(in1_array);
    Py_DECREF(in2_array);
    Py_DECREF(out_array);

    // 返回 None
    Py_RETURN_NONE;
}

// 方法表
static PyMethodDef SquareSumMethods[] = {
    {"square_sum_ufunc",  square_sum_ufunc, METH_VARARGS, "Calculate the sum of squares of two arrays."},
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

// 模块定义
static struct PyModuleDef square_sum_module = {
    PyModuleDef_HEAD_INIT,
    "square_sum",   /* name of module */
    NULL, /* module documentation, may be NULL */
    -1,       /* size of per-interpreter state of the module,
                 or -1 if the module keeps state in global variables. */
    SquareSumMethods
};

// 模块初始化函数
PyMODINIT_FUNC
PyInit_square_sum(void)
{
    PyObject *m;

    m = PyModule_Create(&square_sum_module);
    if (m == NULL)
        return NULL;

    /* IMPORTANT: This must be called */
    import_array();

    return m;
}

2. 编译 C 代码：

gcc -shared -fPIC -o square_sum.so square_sum.c -I/usr/include/python3.x -I/usr/lib/python3.x/site-packages/numpy/core/include

替换 python3.x 为你的Python版本。

3. Python 代码：

import numpy as np
import square_sum  # 导入编译好的C模块

# 创建输入数组
x = np.array([1, 2, 3, 4], dtype=np.float64)
y = np.array([5, 6, 7, 8], dtype=np.float64)
out = np.zeros_like(x, dtype=np.float64)

# 调用 C 函数
square_sum.square_sum_ufunc(x, y, out)

# 打印结果
print(out)  # 输出: [26. 40. 58. 80.]

代码解释：

square_sum 函数：这个函数实现了两个数的平方和的计算。
PyMODINIT_FUNC PyInit_square_sum(void)：这个函数是模块的初始化函数，用于注册 square_sum 函数。
gcc -shared -fPIC -o square_sum.so square_sum.c -I/usr/include/python3.x -I/usr/lib/python3.x/site-packages/numpy/core/include：这个命令用于将 C 代码编译成共享库。
import square_sum：这个语句用于导入共享库。
square_sum.square_sum_ufunc(x, y, out)：这个语句调用了 C 函数，计算两个数组的平方和，并将结果保存到 out 数组中。

优点：

性能高
可以访问底层硬件资源

缺点：

复杂
需要编写 C 代码
可移植性差

总结：

特性	Python (Numba)	C API
易用性	高	低
性能	中	高
代码复杂度	低	高
可移植性	高	低

选择哪种方法取决于你的具体需求。如果你追求简单易用，并且对性能要求不高，那么 Python (Numba) 是一个不错的选择。如果你追求极致性能，并且愿意付出更多的努力，那么 C API 则是更好的选择。

高级技巧

广播 (Broadcasting): ufunc 支持广播机制，可以对不同形状的数组进行运算。例如，可以将一个标量与一个数组相加，或者将一个 (m, 1) 形状的数组与一个 (1, n) 形状的数组相加。
输出参数： 可以将结果保存到指定的输出数组中，而不是创建一个新的数组。这可以节省内存，并提高性能。
错误处理： 可以自定义错误处理方式，例如忽略错误、警告错误、或者抛出异常。

常见问题

ufunc 的输入和输出数据类型必须相同吗？

不一定。可以使用 @vectorize 装饰器指定不同的输入和输出数据类型。
如何调试自定义 ufunc？

可以使用 Python 的调试工具 (如 pdb) 或者 C 的调试工具 (如 gdb) 进行调试。
自定义 ufunc 会影响 NumPy 的性能吗？

如果你的 ufunc 编写不当，可能会影响 NumPy 的性能。因此，需要仔细优化你的代码。

总结

自定义 ufunc 是 NumPy 的一项强大功能，可以让你扩展 NumPy 的功能，并提高数据处理的效率。无论是使用 Python (Numba) 还是 C API，都需要仔细学习和实践，才能掌握这项技能。

希望今天的讲座能够帮助你更好地理解和使用自定义 ufunc。记住，掌握了这项技能，你就可以在 NumPy 的世界里自由驰骋，创造出属于自己的数据处理奇迹！💪

祝大家编程愉快！🎉

自定义 ufunc：化腐朽为神奇，打造你的专属 NumPy 超能力！

什么是 ufunc？

为什么要自定义 ufunc？

如何自定义 ufunc？