pybind11を使ってpipでインストール可能なPythonのパッケージをC++で実装する

pybind11を使ってpipでインストール可能なPythonのパッケージをC++で実装します。今回はFFTを実装してnumpyのfftと同じ結果になるか確認してみます。

フォルダ構成

pybind11_test
│  pyproject.toml
│  sample.cpp
│  setup.py
│  
└─minfft
        minfft.c
        minfft.h

pyproject.toml

pyproject.tomlを用意した場合、仮想環境を作成してrequiresに記載のパッケージをインストールしてからビルドを行います。必須ではありませんが、素の環境にpybind11がインストールされていないとビルドでコケるので用意しておいた方が良いでしょう。

pyproject.toml

[build-system]
requires = ["setuptools>=42", "wheel", "pybind11~=2.6.1"]
build-backend = "setuptools.build_meta"

setup.py

setup.pyを用意します（必須）。Pybind11Extensionクラスのコンストラクタ第一引数にパッケージ名、第二引数にはC/C++ソースのリストを設定します。バージョンや説明文等も情報も設定出来ますが今回は省略します。

setup.py

from setuptools import setup
from pybind11.setup_helpers import Pybind11Extension

ext_modules = [
    Pybind11Extension("pybind11_test", ["sample.cpp", "minfft/minfft.c"])
]

setup(
    name="pybind11_test",
    ext_modules=ext_modules,
    install_requires=["numpy"])

サンプルコード

FFTを実装します。今回はminfftを使用します。

引数はNumpy配列で受け取ります。詳細はこのあたりに載ってます。

サンプルコードは以下の通り。

sample.cpp

#include <pybind11/numpy.h>
#include "minfft/minfft.h"
namespace py = pybind11;

static py::array_t<std::complex<double>> fft(py::array_t<std::complex<double>> x) {
    // check ndim
    if (x.request().ndim != 1){
        throw std::runtime_error("ndim should be 1");
    }
    
    // get fft size
    int N = static_cast<int>(x.request().shape[0]);
    
    // make aux data
    minfft_aux *a = ::minfft_mkaux_dft_1d(N);
    if(a == nullptr){
        throw std::runtime_error("shape[0] should be a power of two");
    }

    // allocate return value
    py::array_t<std::complex<double>> y(x.request().shape);

    // execute fft
    ::minfft_dft(
        static_cast<minfft_cmpl*>(x.request().ptr),
        static_cast<minfft_cmpl*>(y.request().ptr),
        a);

    // free aux data
    ::minfft_free_aux(a);

    return y;
}

PYBIND11_MODULE(pybind11_test, m){
    m.def("fft", &fft);
}

ビルド

setup.pyが存在するパスをカレントディレクトリにした状態で以下のように実行します。

端末

>pip install .

動作確認

自前のFFTに適当な乱数を入力してnumpyのfft一致しているかPythonインタプリタで確認します。

端末

>>> import numpy as np
>>> import pybind11_test
>>> x = np.random.rand(8) + np.random.rand(8)*1j  
>>> np.fft.fft(x)                                  
array([ 4.65887939+3.56753418j, -0.326442  +0.90029068j,
       -0.07578485+0.56485263j,  0.37996314-0.25579869j,
       -0.98194287+0.44748422j,  0.1412921 +0.83332836j,
        0.87452551-0.31797332j,  1.02586444+0.50457297j])
>>> pybind11_test.fft(x)
array([ 4.65887939+3.56753418j, -0.326442  +0.90029068j,
       -0.07578485+0.56485263j,  0.37996314-0.25579869j,
       -0.98194287+0.44748422j,  0.1412921 +0.83332836j,
        0.87452551-0.31797332j,  1.02586444+0.50457297j])
>>>

見た感じ一致しているようです。

(おまけ)GitHubのURLを指定して直接インストールする

以下のように実行するとGitHubからダウンロードすると同時にビルドとインストールが出来ます。

端末

>pip install git+https://github.com/fukuroder/pybind11_test.git

PyPIに登録しなくてよいのは便利。