CODEWITHSUNDEEP
pythonnumpyeigenpybind11
rxFt20
rxFt20

Reputation: 316

Returning a list or tuple of arrays from pybind11 wrapping eigen

I have a c++ function using eigen, which is wrapped using pybind11 so that I can call it from python. A simple version of the intended function returns an Eigen::MatrixXd type, which pybind successfully converts to a 2D numpy array.

I would like this function to be able to return either a list or tuple of such matrices, or a 3D numpy array.

I am somewhat of a novice with c++, and the documentation for pybind does not (as far as I am able to understand) provide any direction. A mock example is below:

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);
}

I would like test2 to return a list or tuple of arrays.

In 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)

The array y is as expected, but ys only contains the array corresponding to the first coefficient of a2.

How should I modify test2 to correctly return more than one array? A 3D array would also be acceptable.

Upvotes: 5

Views: 10431

Answers (2)

Ashot Matevosyan
Ashot Matevosyan

Reputation: 2479

I would suggest to return std::tuple but by moving local objects:

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

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

In the PYBIND11_MODULE I'm not quite sure which is right:

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

or

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

I've tested both work as needed, i.e. wthout copying and allocating more memory during returning.

Upvotes: 4

Olaf
Olaf

Reputation: 606

I have used Eigen before but I am not an expert, so others might be able to improve this solution.

#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);
}

The vector is converted into a list of numpy arrays by pybind11. Results:

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.]])]

Upvotes: 5

Related Questions