第11章 NumPy数组支持

NumPy 是 Python 科学计算的基础库。将 C++ 数据结构与 NumPy 数组无缝对接是扩展模块的核心能力。本章讲解 pybind11 的 NumPy 集成。

11.1 NumPy array protocol 概述

NumPy 数组协议是 Python 中实现高效数组数据传输的标准机制。它允许不同对象（如 list、bytes、C++ 数组、PIL Image 等）以标准方式暴露内存 buffer，供其他代码直接使用。

关键优势：零拷贝（zero-copy）。如果数据已经在内存中，协议允许使用者直接访问，无需复制。

pybind11 通过 py::array 和 py::array_t<T> 提供一流的支持。

11.2 dtype 映射

NumPy 的 dtype 指定数组元素类型。pybind11 自动映射：

C++ 类型	NumPy dtype	说明
int8_t	int8	8位有符号整数
int16_t	int16	16位有符号整数
int32_t	int32	32位有符号整数
int64_t	int64	64位有符号整数
uint8_t	uint8	8位无符号整数
uint16_t	uint16	16位无符号整数
uint32_t	uint32	32位无符号整数
uint64_t	uint64	64位无符号整数
float	float32	32位浮点
double	float64	64位浮点

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

py::array_t<double> create_array(size_t size) {
    std::vector<double> data(size);
    // 初始化数据...
    return py::array_t<double>(data);
}

PYBIND11_MODULE(example, m) {
    m.def("create_array", &create_array);
}

11.3 C++ 到 NumPy：零拷贝转换

当 C++ 数据已连续存储时，可以创建视图而非复制：

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vector>

namespace py = pybind11;

// 方案1：从现有 vector 创建（会复制）
py::array_t<int> from_vector(const std::vector<int>& vec) {
    return py::array_t<int>(vec);  // 复制数据
}

// 方案2：零拷贝——直接使用 vector 的内存
py::array_t<int> from_vector_zero_copy(const std::vector<int>& vec) {
    // shape, strides
    return py::array_t<int>(
        {static_cast<ssize_t>(vec.size())},  // shape
        {sizeof(int)},                         // strides (字节)
        vec.data()                             // 借用内存
    );
}

>>> import numpy as np
>>> from example import from_vector, from_vector_zero_copy

>>> # 方案1：复制
>>> arr = from_vector([1, 2, 3])
>>> arr
array([1, 2, 3])

>>> # 方案2：零拷贝视图
>>> vec = [10, 20, 30]
>>> arr = from_vector_zero_copy(vec)  # 依赖 vec 存活
>>> arr[0] = 999
>>> vec  # Python 端修改影响了 C++ 端！
[999, 20, 30]

关键洞察：零拷贝的关键是内存所有权。当你传入 vec.data() 时，pybind11 并不复制数据，而是创建一个「借用」视图。只要 Python 端的 array 存在，C++ 端的 vector 就不能被销毁。这需要谨慎的生命周期管理。

11.4 NumPy 到 C++：复制转换

从 Python 传入 NumPy 数组时，默认会复制数据：

py::array_t<double> process_array(py::array_t<double> input) {
    // input 已经是 py::array_t<double>
    // 但底层数据可能不是连续存储
    auto buf = input.request();
    double* ptr = static_cast<double*>(buf.ptr);

    // 处理数据...
    for (ssize_t i = 0; i < buf.size; ++i) {
        ptr[i] *= 2.0;
    }
    return input;  // 返回同一数组（已修改）
}

>>> arr = np.array([1.0, 2.0, 3.0])
>>> result = example.process_array(arr)
>>> result
array([2., 4., 6.])

11.5 连续内存假设

零拷贝要求数据是连续的（contiguous）：

py::array_t<double> process_contiguous(py::array_t<double> input) {
    // 检查是否连续
    if (!input.flags()["C_CONTIGUOUS"] && !input.flags()["F_CONTIGUOUS"]) {
        // 不是连续存储，需要复制
        py::array_t<double> contiguous = input.unchecked<double>().follow();
        // 处理 contiguous...
    }
    // ...
}

>>> arr = np.array([[1, 2], [3, 4]])  # 默认 C-contiguous
>>> arr.flags['C_CONTIGUOUS']
True

关键洞察：C++ 的 std::vector 是连续存储的，所以可以直接创建零拷贝视图。但对于多维数组（如 std::vector<std::vector<double>>），内层vector 不是连续存储，不能直接创建视图。

11.6 行主序 vs 列主序

NumPy 默认使用行主序（Row-major, C-style），最后维度变化最快：

// C++ 二维数组通常也是行主序
std::vector<std::vector<double>> matrix(3, std::vector<double>(4));

// 转换时需要处理内存布局
py::array_t<double> from_matrix(const std::vector<std::vector<double>>& matrix) {
    size_t rows = matrix.size();
    size_t cols = matrix[0].size();
    py::array_t<double> result({rows, cols});

    auto r = result.mutable_unchecked<double, 2>();
    for (size_t i = 0; i < rows; ++i) {
        for (size_t j = 0; j < cols; ++j) {
            r(i, j) = matrix[i][j];
        }
    }
    return result;
}

>>> arr = example.from_matrix([[1, 2, 3], [4, 5, 6]])
>>> arr.shape
(2, 3)
>>> arr[0, 1]
2.0

列主序（Fortran-style）需要显式指定 strides：

py::array_t<double> from_matrix_fortran(const double* data, size_t rows, size_t cols) {
    return py::array_t<double>(
        {static_cast<ssize_t>(rows), static_cast<ssize_t>(cols)},  // shape
        {sizeof(double), static_cast<ssize_t>(cols * sizeof(double))},  // strides: 列优先
        data
    );
}

11.7 数组标志

NumPy 数组有多个标志位，pybind11 可以访问：

void check_array_flags(const py::array_t<double>& arr) {
    auto flags = arr.flags();

    bool c_contiguous = flags[py::array::c_contiguous];
    bool f_contiguous = flags[py::array::f_contiguous];
    bool writeable = flags[py::array::writeable];

    // ...
}

常用标志：

• C_CONTIGUOUS (C order)
• F_CONTIGUOUS (Fortran order)
• OWNDATA (数组拥有数据所有权)
• WRITEABLE (可写)

11.8 实战示例：图像数据传递

图像处理是 NumPy 整合的典型场景：

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vector>
#include <cstdint>

namespace py = pybind11;

// 模拟读取图像：返回 RGBA 格式的图像数据
py::array_t<uint8_t> load_image(const std::string& path) {
    // 假设从文件读取图像，得到原始像素数据
    size_t width = 640;
    size_t height = 480;

    std::vector<uint8_t> pixels(height * width * 4);

    // ... 填充像素数据 ...

    // 零拷贝返回：图像数据已在 pixels 中
    return py::array_t<uint8_t>(
        {static_cast<ssize_t>(height), static_cast<ssize_t>(width), 4},
        {static_cast<ssize_t>(width * 4), 4, 1},
        pixels.data()
    );
}

// 处理图像：接收图像并应用灰度滤镜
py::array_t<uint8_t> apply_grayscale(py::array_t<uint8_t> input) {
    auto arr = input.unchecked<uint8_t, 3>();  // shape: (H, W, C)

    size_t height = arr.shape(0);
    size_t width = arr.shape(1);

    // 创建灰度输出
    py::array_t<uint8_t> output({height, width});
    auto out = output.mutable_unchecked<uint8_t, 2>();

    for (size_t i = 0; i < height; ++i) {
        for (size_t j = 0; j < width; ++j) {
            uint8_t r = arr(i, j, 0);
            uint8_t g = arr(i, j, 1);
            uint8_t b = arr(i, j, 2);
            out(i, j) = static_cast<uint8_t>(0.299 * r + 0.587 * g + 0.114 * b);
        }
    }
    return output;
}

PYBIND11_MODULE(example, m) {
    m.def("load_image", &load_image, "Load an image from file");
    m.def("apply_grayscale", &apply_grayscale, "Convert image to grayscale");
}

>>> import numpy as np
>>> from example import load_image, apply_grayscale

>>> # 加载图像（零拷贝）
>>> img = load_image("test.png")
>>> img.shape
(480, 640, 4)  # H, W, C (RGBA)

>>> # 转换为灰度图（复制）
>>> gray = apply_grayscale(img)
>>> gray.shape
(480, 640)

>>> # 显示图像
>>> import matplotlib.pyplot as plt
>>> plt.imshow(gray, cmap='gray')
>>> plt.show()

关键洞察：图像数据传输是零拷贝的经典场景。C++ 端读取图像文件后，可以直接返回 py::array_t<uint8_t> 视图，Python 端无需复制即可访问。如果 Python 端需要修改（滤镜），则会产生一份拷贝。理解何时复制、何时引用是性能优化的关键。

py::array_t 模板详解

py::array_t<T> 是处理 NumPy 数组的主要工具：

// 创建指定形状的数组
py::array_t<float> arr = py::array_t<float>::_empty({3, 4});

// 创建并指定初始化值
py::array_t<double> arr2 = py::array_t<double>::ensure({10, 20});

// 访问元素（编译时类型检查）
auto unchecked = arr.unchecked<float>();  // 只读
auto mut_unchecked = arr.mutable_unchecked<double>();  // 可写

// 请求 buffer info
auto req = arr.request();

常见陷阱

1. 内存不连续导致崩溃

// 错误：std::vector<std::vector<T>> 不是连续内存
std::vector<std::vector<double>> jagged = {{1,2}, {3,4}};
return py::array_t<double>({2, 2}, {sizeof(double)*2, sizeof(double)}, &jagged[0][0]);
// 崩溃：内存布局不符合预期

2. 生命周期管理

// 危险：返回局部变量的视图
py::array_t<double> bad_example() {
    std::vector<double> v = {1, 2, 3};
    return py::array_t<double>({3}, {sizeof(double)}, v.data());
    // v 在函数结束时销毁，视图悬空！
}

// 正确：确保数据存活
py::array_t<double> good_example() {
    static std::vector<double> v = {1, 2, 3};
    return py::array_t<double>({3}, {sizeof(double)}, v.data());
    // 或者转移到 heap
}

3. dtype 不匹配

>>> arr = np.array([1, 2, 3], dtype=np.int32)
>>> result = process_ints(arr)  # C++ 期望 float

关键洞察：pybind11 与 NumPy 的整合是科学计算扩展模块的核心。记住三个要点：

1. 零拷贝可行：当 C++ 数据连续存储且生命周期正确管理时，NumPy 可以直接访问 C++ 内存
2. 复制安全：不确定内存布局时，复制是安全的默认选择
3. dtype 必须匹配：确保 C++ 类型与 NumPy dtype 对应，否则读取错误数据