pybind11入門(4) NumPy連携その2

その1 https://qiita.com/lucidfrontier45/items/872c38b5eec08e3800d7
その2 https://qiita.com/lucidfrontier45/items/c7cc409a3af962d100a3
その3 https://qiita.com/lucidfrontier45/items/183526df954c1d6580ba

pybind11を使ってNumPyと連携する方法を試した。今回はEigenを利用してみる。

環境準備

pybind11とEigenをそれぞれgithubなりから取得し、プロジェクトのトップディレクトリに置く。その後、以下のようにCMakeLists.txtを用意する。

cmake_minimum_required(VERSION 3.2)
project(pybind_test VERSION 0.1.0)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
add_subdirectory(pybind11)
include_directories(eigen)

pybind11_add_module(mymodule module.cpp)

コード

おまじない部分

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

template <typename T>
using RMatrix = Eigen::Matrix<T, -1, -1, Eigen::RowMajor>;

pybind11にはEigenを自動でNumpyのndarrayに変換する機能がある。これを使用するためにはpybind11/eigen.hをincludeする必要がある。また、ndarrayはメモリ配置がデフォルトではC言語で主に用いられるRowMajorなのに対し、EigenではデフォルトはfortranのColumnMajorである。ただ、両者とも指定することができるのでここではndarrayのRowMajorに揃える。毎回RowMajorを宣言するのも面倒なのでRMatrixというテンプレートクラスを作っておく。

ndarrayをEigenの行列に変換

最も簡単な方法は値渡し利用する方法である。この方法を用いるとndarrayの内容がEigenのmatrixにコピーされる。

template <typename T>
void print_array(RMatrix<T> m)
{
    std::cout << m << std::endl;
}

ndarrayを生成する

Eigenのmatrixなどを戻り値に持つ関数を作れば自動的にndarrayに変換される。しかも新しくコピーを作るのではなく、Eigenのmatrixのメモリをそのまま利用する。賢いぞpybind11!

template <typename T>
RMatrix<T> modify_array(RMatrix<T> m, T a)
{
    return m * a;
}

ndarrayを書き換える。

Eigen::Refを使用することで参照渡しができる。この方法を用いるとコピーが作成されず、Eigen::Mapを利用してndarrayのメモリ領域を直接Eigenからいじることができる。なお、この方法を使用する場合は型やメモリ配置(RowMajor, ColumnMajor)が完全に一致していることが必須条件になる。

template <typename T>
void modify_array_inplace(Eigen::Ref<RMatrix<T>> m, T a)
{
    m = m * a;
}

なお、Eigen::Ref<const RMatrix<T>>のように渡すことで読み取り専用の参照渡しになる。大きな行列を渡す場合などはコピーも大きなオーバーヘッドになるので特別な理由がない場合は常にこの方法で渡すのがいいだろう。

登録

いつも通り

PYBIND11_MODULE(mymodule, m)
{
    m.doc() = "my test module";
    m.def("print_array", &print_array<double>, "");
    m.def("modify_array", &modify_array<double>, "");
    m.def("modify_array_inplace", &modify_array_inplace<double>, "");
}

実際に試す

import numpy as np
import mymodule


x = np.arange(12).reshape(3, 4).astype(np.float64)
print("created with numpy")
print(x)
print()

print("using pybind11 to print")
mymodule.print_array(x)
print()

x2 = mymodule.modify_array(x, 2.0)
print("x2 newly created with pybind11")
print(x2)
print()

mymodule.modify_array_inplace(x, 3.0)
print("x modified inplace")
print(x)

>>>
created with numpy
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]]

using pybind11 to print
 0  1  2  3
 4  5  6  7
 8  9 10 11

x2 newly created with pybind11
[[ 0.  2.  4.  6.]
 [ 8. 10. 12. 14.]
 [16. 18. 20. 22.]]

x modified inplace
[[ 0.  3.  6.  9.]
 [12. 15. 18. 21.]
 [24. 27. 30. 33.]]