使用 pybind11 传递包含 NumPy 数组的 Python 字典到 C++

要将包含 NumPy 数组的 Python 字典传递到 C++ 函数中,你需要使用 pybind11 的特殊处理方式。以下是完整的解决方案:

基本方法

1. 包含必要头文件

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <string>
#include <vector>
#include <unordered_map>

namespace py = pybind11;

2. 处理字典中的 NumPy 数组

void process_dict(py::dict d) {
    // 检查字典是否包含特定键
    if (d.contains("array_key")) {
        // 获取 NumPy 数组
        py::array_t<double> array = d["array_key"].cast<py::array_t<double>>();
        
        // 获取数组的缓冲区信息
        py::buffer_info buf = array.request();
        
        // 获取指向数据的指针
        double *ptr = static_cast<double *>(buf.ptr);
        
        // 获取数组形状
        size_t size = buf.size;
        std::vector<size_t> shape;
        for (auto s : buf.shape) {
            shape.push_back(s);
        }
        
        // 现在可以使用 ptr 访问数组数据
        // ...
    }
    
    // 处理其他非数组元素
    if (d.contains("other_key")) {
        std::string value = d["other_key"].cast<std::string>();
        // ...
    }
}

更通用的模板方法

如果你需要处理不同类型的 NumPy 数组,可以使用模板:

template <typename T>
void process_numpy_array(py::array_t<T> array) {
    py::buffer_info buf = array.request();
    T *ptr = static_cast<T *>(buf.ptr);
    
    // 处理数组数据...
}

void process_dict_generic(py::dict d) {
    for (auto item : d) {
        std::string key = item.first.cast<std::string>();
        py::handle value = item.second;
        
        try {
            // 尝试转换为 double 数组
            py::array_t<double> double_array = value.cast<py::array_t<double>>();
            process_numpy_array(double_array);
            continue;
        } catch (...) {}
        
        try {
            // 尝试转换为 int 数组
            py::array_t<int> int_array = value.cast<py::array_t<int>>();
            process_numpy_array(int_array);
            continue;
        } catch (...) {}
        
        // 处理其他类型
        try {
            std::string str_value = value.cast<std::string>();
            // 处理字符串...
        } catch (...) {}
    }
}

绑定到 Python

PYBIND11_MODULE(example, m) {
    m.def("process_dict", &process_dict, "Process a dictionary with numpy arrays");
    m.def("process_dict_generic", &process_dict_generic, "Generic dictionary processor");
}

Python 端使用示例

import numpy as np
import example

data = {
    "array1": np.array([1.0, 2.0, 3.0], dtype=np.float64),
    "array2": np.array([[1, 2], [3, 4]], dtype=np.int32),
    "description": "Sample data"
}

example.process_dict(data)
example.process_dict_generic(data)

注意事项

  1. 内存管理:NumPy 数组是引用计数的,只要 Python 对象存在,C++ 端访问就是安全的。

  2. 类型检查:在转换前最好检查类型,可以使用 py::isinstance<py::array_t<double>>(value)

  3. 性能考虑:对于大型数组,考虑使用 py::array_t::uncheckedpy::array_t::mutable_unchecked 来避免边界检查。

  4. 异常处理:添加适当的异常处理以确保 Python 异常能正确传播。

  5. 只读访问:如果不需要修改数组数据,使用 const 指针以避免意外修改。

高级用法:结构化数据

如果你知道字典的确切结构,可以创建更类型安全的接口:

struct MyData {
    py::array_t<double> array_data;
    std::string metadata;
    int some_value;
};

MyData parse_dict(py::dict d) {
    MyData data;
    data.array_data = d["array"].cast<py::array_t<double>>();
    data.metadata = d["meta"].cast<std::string>();
    data.some_value = d["value"].cast<int>();
    return data;
}

这样使用起来会更加安全和清晰。


使用pybind11在C++中处理NumPy数组

pybind11提供了与NumPy数组交互的功能,主要通过pybind11/numpy.h头文件实现。下面我将介绍如何在C++端处理NumPy数组,包括获取数组类型和形状信息。

基本设置

首先确保你的项目已经包含pybind11,并在代码中包含必要的头文件:

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

namespace py = pybind11;

接收NumPy数组作为参数

你可以定义一个接受NumPy数组的函数:

void process_array(py::array_t<double> array) {
    // 处理数组的代码
}

获取数组信息

1. 获取形状信息

void print_array_info(py::array_t<double> array) {
    // 获取数组缓冲区信息
    py::buffer_info buf = array.request();
    
    // 输出维度数量
    std::cout << "Number of dimensions: " << buf.ndim << std::endl;
    
    // 输出形状
    std::cout << "Shape: ";
    for (size_t i = 0; i < buf.ndim; ++i) {
        std::cout << buf.shape[i] << " ";
    }
    std::cout << std::endl;
    
    // 输出元素总数
    std::cout << "Total elements: " << buf.size << std::endl;
}

2. 获取数据类型信息

void check_array_type(py::array array) {
    py::dtype dtype = array.dtype();
    
    // 比较数据类型
    if (dtype.is(py::dtype::of<int>())) {
        std::cout << "Array is of type int" << std::endl;
    } else if (dtype.is(py::dtype::of<float>())) {
        std::cout << "Array is of type float" << std::endl;
    } else if (dtype.is(py::dtype::of<double>())) {
        std::cout << "Array is of type double" << std::endl;
    } else {
        std::cout << "Unknown array type" << std::endl;
    }
    
    // 获取类型字符串表示
    std::cout << "Type string: " << dtype.kind() << dtype.itemsize() << std::endl;
}

访问数组数据

1. 直接访问

void process_array(py::array_t<double> array) {
    py::buffer_info buf = array.request();
    double *ptr = static_cast<double *>(buf.ptr);
    
    // 假设是2D数组
    if (buf.ndim == 2) {
        size_t rows = buf.shape[0];
        size_t cols = buf.shape[1];
        
        for (size_t i = 0; i < rows; i++) {
            for (size_t j = 0; j < cols; j++) {
                // 访问元素
                double value = ptr[i * cols + j];
                // 处理元素...
            }
        }
    }
}

2. 使用更安全的访问器

void process_array_safe(py::array_t<double> array) {
    auto buf = array.mutable_unchecked<2>(); // 假设是2D数组
    
    for (py::ssize_t i = 0; i < buf.shape(0); i++) {
        for (py::ssize_t j = 0; j < buf.shape(1); j++) {
            // 直接访问元素,带边界检查
            double value = buf(i, j);
            // 修改元素
            buf(i, j) = value * 2;
        }
    }
}

创建NumPy数组并返回给Python

py::array_t<double> create_array(size_t rows, size_t cols) {
    // 创建形状向量
    std::vector<size_t> shape = {rows, cols};
    
    // 创建strides向量
    std::vector<size_t> strides = {cols * sizeof(double), sizeof(double)};
    
    // 分配内存
    double *data = new double[rows * cols];
    
    // 初始化数据
    for (size_t i = 0; i < rows; i++) {
        for (size_t j = 0; j < cols; j++) {
            data[i * cols + j] = i + j;
        }
    }
    
    // 创建数组并指定删除器
    return py::array_t<double>(
        shape,
        strides,
        data,
        py::capsule(data, [](void *data) { delete[] static_cast<double *>(data); })
    );
}

绑定函数到Python模块

PYBIND11_MODULE(example, m) {
    m.def("process_array", &process_array, "Process a NumPy array");
    m.def("print_array_info", &print_array_info, "Print array information");
    m.def("check_array_type", &check_array_type, "Check array type");
    m.def("create_array", &create_array, "Create a new array");
}

注意事项

  1. 内存管理:当从C++创建NumPy数组返回给Python时,需要确保内存管理正确。上面的例子使用了capsule来在Python不再需要数组时释放内存。

  2. 类型安全:使用array_t<T>模板可以确保传入的NumPy数组是特定类型。如果需要接受任意类型的数组,可以使用py::array

  3. 维度检查:在访问数组元素前总是检查维度数量和形状。

  4. 性能:对于高性能计算,考虑使用mutable_uncheckedunchecked访问器,它们提供了更快的访问速度但牺牲了一些安全性。

  5. 跨步访问:NumPy数组可能是非连续的(有跨步),处理时要考虑strides信息。

通过这些方法,你可以在C++中有效地处理来自Python的NumPy数组,并获取其类型和形状信息。


pybind11::array 类型详解与使用指南

pybind11::array 是 pybind11 库中用于处理 NumPy 数组的核心类型,它提供了在 C++ 和 Python 之间高效传递多维数组数据的能力。以下是关于 pybind11::array 的详细说明和使用方法。

基本概念

pybind11::array 是 pybind11 对 NumPy ndarray 的封装,允许在 C++ 中直接操作 Python 的 NumPy 数组数据,而无需复制数据。

主要特性

  1. 零拷贝操作:可以在 C++ 和 Python 之间传递数组数据而不复制
  2. 类型安全:支持多种数据类型 (dtype)
  3. 维度支持:处理任意维度的数组
  4. 缓冲区协议:实现 Python 缓冲区协议

基本用法

1. 从 Python 接收 NumPy 数组

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

namespace py = pybind11;

void process_array(py::array_t<double> arr) {
    // 获取数组信息
    py::buffer_info buf = arr.request();
    
    // 访问数据
    double *ptr = static_cast<double *>(buf.ptr);
    size_t size = buf.size;
    
    // 处理数据...
}

2. 创建并返回 NumPy 数组给 Python

py::array_t<double> create_array(size_t size) {
    // 创建未初始化的数组
    py::array_t<double> result({size});
    
    // 获取可写视图
    py::buffer_info buf = result.request();
    double *ptr = static_cast<double *>(buf.ptr);
    
    // 填充数据
    for (size_t i = 0; i < size; i++)
        ptr[i] = i * 1.5;
    
    return result;
}

多维数组处理

处理多维数组

void process_2d_array(py::array_t<double> arr) {
    py::buffer_info buf = arr.request();
    
    if (buf.ndim != 2)
        throw std::runtime_error("Number of dimensions must be 2");
    
    double *ptr = static_cast<double *>(buf.ptr);
    size_t rows = buf.shape[0];
    size_t cols = buf.shape[1];
    
    // 处理2D数组...
    for (size_t i = 0; i < rows; i++) {
        for (size_t j = 0; j < cols; j++) {
            ptr[i * cols + j] = ...; // 访问元素
        }
    }
}

创建多维数组

py::array_t<double> create_2d_array(size_t rows, size_t cols) {
    // 创建2D数组
    py::array_t<double> result({rows, cols});
    
    py::buffer_info buf = result.request();
    double *ptr = static_cast<double *>(buf.ptr);
    
    // 填充数据
    for (size_t i = 0; i < rows; i++) {
        for (size_t j = 0; j < cols; j++) {
            ptr[i * cols + j] = i + j * 0.1;
        }
    }
    
    return result;
}

数组视图和修改

获取可写视图

void modify_array(py::array_t<double> arr) {
    // 获取可写视图
    auto r = arr.mutable_unchecked<2>(); // 假设是2D数组
    
    for (py::ssize_t i = 0; i < r.shape(0); i++) {
        for (py::ssize_t j = 0; j < r.shape(1); j++) {
            r(i, j) = ...; // 修改元素
        }
    }
}

只读视图

void read_array(py::array_t<double> arr) {
    // 获取只读视图
    auto r = arr.unchecked<2>(); // 假设是2D数组
    
    for (py::ssize_t i = 0; i < r.shape(0); i++) {
        for (py::ssize_t j = 0; j < r.shape(1); j++) {
            double value = r(i, j); // 读取元素
        }
    }
}

数据类型处理

pybind11 支持多种 NumPy 数据类型:

// 创建不同类型数组的示例
py::array_t<int> create_int_array(size_t size);
py::array_t<float> create_float_array(size_t size);
py::array_t<uint8_t> create_byte_array(size_t size);

步幅和非连续数组

处理非连续数组时需要注意步幅:

void process_strided_array(py::array_t<double> arr) {
    py::buffer_info buf = arr.request();
    
    double *ptr = static_cast<double *>(buf.ptr);
    ssize_t stride = buf.strides[0] / sizeof(double); // 计算元素步幅
    
    for (ssize_t i = 0; i < buf.shape[0]; i++) {
        double value = ptr[i * stride]; // 考虑步幅访问元素
    }
}

高级用法

使用 Eigen 互操作

#include <pybind11/eigen.h>

// 将 NumPy 数组转换为 Eigen 矩阵
void process_with_eigen(py::array_t<double> arr) {
    Eigen::MatrixXd mat = py::cast<Eigen::MatrixXd>(arr);
    // 使用 Eigen 处理...
}

自定义缓冲区

// 从现有内存创建 NumPy 数组而不复制
void expose_existing_memory(double *data, size_t size) {
    // 创建数组但不拥有数据
    py::array_t<double> arr(
        {size},  // 形状
        {sizeof(double)},  // 步幅
        data  // 数据指针
    );
    
    // 返回给Python,调用者需要确保data的生命周期足够长
    return arr;
}

性能注意事项

  1. 避免不必要的拷贝:尽量使用 array_t 直接操作数据
  2. 视图检查unchecked() 比带边界检查的访问更快
  3. 维度固定:如果知道数组维度,使用模板参数如 unchecked<2> 可以提高性能
  4. 类型匹配:确保 C++ 和 Python 端数据类型一致

完整示例

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

namespace py = pybind11;

// 计算数组元素平方和
double array_sum_of_squares(py::array_t<double> input) {
    auto buf = input.request();
    if (buf.ndim != 1)
        throw std::runtime_error("Only 1D arrays are supported");
    
    double *ptr = static_cast<double *>(buf.ptr);
    double sum = 0.0;
    
    for (ssize_t i = 0; i < buf.size; i++) {
        sum += ptr[i] * ptr[i];
    }
    
    return sum;
}

// 创建并返回一个2D数组
py::array_t<double> create_2d_random_array(int rows, int cols) {
    // 分配内存并创建数组
    py::array_t<double> result({rows, cols});
    
    // 获取可写视图
    auto r = result.mutable_unchecked<2>();
    
    // 填充随机数据
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            r(i, j) = static_cast<double>(rand()) / RAND_MAX;
        }
    }
    
    return result;
}

PYBIND11_MODULE(example, m) {
    m.def("array_sum_of_squares", &array_sum_of_squares, "Calculate sum of squares");
    m.def("create_2d_random_array", &create_2d_random_array, 
          "Create a 2D array with random values",
          py::arg("rows"), py::arg("cols"));
}

常见问题

  1. 数据所有权:确保在 C++ 端访问数据时,Python 端的数组仍然存在
  2. 类型不匹配:确保 NumPy 数组的 dtype 与 C++ 类型匹配
  3. 维度不匹配:检查数组维度是否符合预期
  4. 非连续数组:处理非连续数组时要注意步幅

通过合理使用 pybind11::array,可以实现 C++ 和 Python 之间高效的数据交换,特别适合数值计算和科学计算应用。


使用 pybind11 将 C/C++ 结构体转换为 Python 字典(共享内存)

要将包含标量和 std::vector 的 C++ 结构体通过 pybind11 转换为 Python 字典,并且保持数据共享(不复制内存),可以按照以下方法实现:

基本方法

1. 定义 C++ 结构体

#include <vector>
#include <string>

struct MyStruct {
    int scalar_int;
    double scalar_double;
    std::vector<float> vec_float;
    std::vector<int> vec_int;
    std::string str;  // 字符串通常需要复制
};

2. 创建 pybind11 绑定代码

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

namespace py = pybind11;

// 将 MyStruct 转换为 Python 字典(共享内存)
py::dict struct_to_dict(MyStruct& s) {
    py::dict d;
    
    // 标量类型直接赋值(自动转换)
    d["scalar_int"] = s.scalar_int;
    d["scalar_double"] = s.scalar_double;
    
    // vector 转换为 numpy 数组(共享内存)
    d["vec_float"] = py::array_t<float>(
        s.vec_float.size(),       // 形状
        s.vec_float.data()        // 数据指针
    );
    
    d["vec_int"] = py::array_t<int>(
        s.vec_int.size(),
        s.vec_int.data()
    );
    
    // 字符串需要复制
    d["str"] = s.str;
    
    return d;
}

// 反向转换:从 Python 字典更新结构体
void dict_to_struct(const py::dict& d, MyStruct& s) {
    if (d.contains("scalar_int")) {
        s.scalar_int = d["scalar_int"].cast<int>();
    }
    if (d.contains("scalar_double")) {
        s.scalar_double = d["scalar_double"].cast<double>();
    }
    
    // 处理 numpy 数组
    if (d.contains("vec_float")) {
        py::array_t<float> arr = d["vec_float"].cast<py::array_t<float>>();
        auto buf = arr.request();
        s.vec_float.resize(buf.size);
        memcpy(s.vec_float.data(), buf.ptr, buf.size * sizeof(float));
    }
    
    if (d.contains("vec_int")) {
        py::array_t<int> arr = d["vec_int"].cast<py::array_t<int>>();
        auto buf = arr.request();
        s.vec_int.resize(buf.size);
        memcpy(s.vec_int.data(), buf.ptr, buf.size * sizeof(int));
    }
    
    if (d.contains("str")) {
        s.str = d["str"].cast<std::string>();
    }
}

PYBIND11_MODULE(my_module, m) {
    py::class_<MyStruct>(m, "MyStruct")
        .def(py::init<>())
        .def_readwrite("scalar_int", &MyStruct::scalar_int)
        .def_readwrite("scalar_double", &MyStruct::scalar_double)
        .def_readwrite("vec_float", &MyStruct::vec_float)
        .def_readwrite("vec_int", &MyStruct::vec_int)
        .def_readwrite("str", &MyStruct::str)
        .def("to_dict", &struct_to_dict)
        .def("from_dict", &dict_to_struct);
}

更高级的实现(自动属性访问)

如果你想实现更自然的 Python 字典接口,可以创建一个自定义的 dict 子类:

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

namespace py = pybind11;

class StructDict : public py::dict {
public:
    StructDict(MyStruct& s) : py::dict(), struct_ref(s) {
        update_from_struct();
    }
    
    void update_from_struct() {
        (*this)["scalar_int"] = struct_ref.scalar_int;
        (*this)["scalar_double"] = struct_ref.scalar_double;
        
        (*this)["vec_float"] = py::array_t<float>(
            struct_ref.vec_float.size(),
            struct_ref.vec_float.data()
        );
        
        (*this)["vec_int"] = py::array_t<int>(
            struct_ref.vec_int.size(),
            struct_ref.vec_int.data()
        );
        
        (*this)["str"] = struct_ref.str;
    }
    
    void update_to_struct() {
        struct_ref.scalar_int = (*this)["scalar_int"].cast<int>();
        struct_ref.scalar_double = (*this)["scalar_double"].cast<double>();
        
        // 处理 numpy 数组
        auto update_vector = [](auto& vec, const py::array& arr) {
            auto buf = arr.request();
            vec.resize(buf.size);
            memcpy(vec.data(), buf.ptr, buf.size * sizeof(typename decltype(vec)::value_type));
        };
        
        update_vector(struct_ref.vec_float, (*this)["vec_float"]);
        update_vector(struct_ref.vec_int, (*this)["vec_int"]);
        
        struct_ref.str = (*this)["str"].cast<std::string>();
    }
    
private:
    MyStruct& struct_ref;
};

PYBIND11_MODULE(my_module, m) {
    py::class_<MyStruct>(m, "MyStruct")
        .def(py::init<>())
        .def("as_dict", [](MyStruct& s) {
            return StructDict(s);
        });
}

使用示例(Python 端)

import my_module
import numpy as np

# 创建结构体实例
s = my_module.MyStruct()
s.scalar_int = 42
s.scalar_double = 3.14
s.vec_float = [1.0, 2.0, 3.0]
s.vec_int = [4, 5, 6]
s.str = "hello"

# 转换为字典(共享内存)
d = s.as_dict()

# 修改字典会影响原始结构体
d["scalar_int"] = 100
d["vec_float"][0] = 99.0  # 修改 numpy 数组

# 从字典更新结构体
s.from_dict(d)

注意事项

  1. 内存共享:对于 std::vector 转换的 numpy 数组,只有在 vector 不重新分配内存时才能保持共享。如果 vector 被修改导致重新分配内存,共享将失效。

  2. 线程安全:这种共享内存方式在多线程环境下需要额外的同步机制。

  3. 生命周期管理:确保 C++ 结构体的生命周期长于 Python 字典的使用时间。

  4. 字符串处理:字符串通常需要复制,无法直接共享内存。

  5. 复杂类型:对于更复杂的嵌套结构,可能需要更复杂的转换逻辑。

这种方法提供了高效的数据共享机制,特别适合处理大型数组数据,避免了不必要的内存复制。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐