从包装本征的 pybind11 返回数组列表或元组

Question

我有一个使用 eigen 的 c++ 函数，它使用 pybind11 包装，以便我可以从 python 调用它。预期函数的一个简单版本 return 是一个 Eigen::MatrixXd 类型，pybind 成功地将其转换为 2D numpy 数组。

我希望此函数能够 return 此类矩阵的列表或元组，或 3D numpy 数组。

我是 C++ 的新手，pybind 的文档（据我所知）没有提供任何指导。下面是一个模拟示例：

test.cpp

#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <Eigen/Dense>

Eigen::MatrixXd test(Eigen::Ref<const Eigen::MatrixXd> x, double a)
{
  Eigen::MatrixXd y;
  y = x * a;
  return y;
}

Eigen::MatrixXd *test2(Eigen::Ref<const Eigen::MatrixXd> x, Eigen::Ref<const Eigen::VectorXd> as)
{
  Eigen::MatrixXd *ys = new Eigen::MatrixXd[as.size()];
  for(unsigned int k = 0; k < as.size(); k++){
    Eigen::MatrixXd& y = ys[k];
    y = x * as[k];
  }
  return ys;
}

namespace py = pybind11;

PYBIND11_MODULE(test, m)
{
  m.doc() = "minimal working example";
  m.def("test", &test);
  m.def("test2", &test2);
}

我想要 test2 到 return 数组的列表或元组。

在python中：

import test
import numpy as np
x = np.random.random((50, 50))
x = np.asfortranarray(x)
a = 0.1
a2 = np.array([1.0, 2.0, 3.0])
y = test.test(x, a)
ys = test.test2(x, a2)

数组y符合预期，但是ys只包含a2的第一个系数对应的数组。

我应该如何修改 test2 以正确地 return 多个数组？ 3D 数组也是可以接受的。

Answer 1

我以前使用过 Eigen，但我不是专家，所以其他人也许可以改进这个解决方案。

#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl.h>
#include <Eigen/Dense>

std::vector<Eigen::MatrixXd> 
test2(Eigen::Ref<const Eigen::MatrixXd> x, Eigen::Ref<const Eigen::VectorXd> as){
    std::vector<Eigen::MatrixXd> matrices;
    for(unsigned int k = 0; k < as.size(); k++){
        Eigen::MatrixXd ys = x * as[k];
        matrices.push_back(ys);
    }
    return matrices;
}

namespace py = pybind11;

PYBIND11_MODULE(test, m){
    m.doc() = "minimal working example";
    m.def("test2", &test2);
}

向量被pybind11转换成一个numpy数组列表。结果：

In [1]: import numpy as np; x = np.ones((2,2)); a = np.array((2., 3.)); import test

In [2]: test.test2(x, a)
Out[2]: 
[array([[2., 2.],
        [2., 2.]]), array([[3., 3.],
        [3., 3.]])]

Answer 2

我建议 return std::tuple 但是通过移动本地对象：

std::tuple<Eigen::MatrixXd,int> function(){
    ...
    int m = 4;
    Eigen::MatrixXd M = ...;
    ...

    return make_tuple(std::move(M),m);
}

在PYBIND11_MODULE中我不太确定哪个是正确的：

m.def("function", &function, py::return_value_policy::reference_internal);

或

m.def("function", &function);

我已经根据需要测试了这两项工作，即在 returning 期间没有复制和分配更多内存。

从包装本征的 pybind11 返回数组列表或元组

Returning a list or tuple of arrays from pybind11 wrapping eigen

python

numpy

eigen

pybind11