PyO3とrust-numpyを使ってPythonからNumPyをRustに渡して操作する

#NumPy  #PyO3  #Rust  #シミュレーション  #Python 

前回の記事でPyO3を使ってPythonで使えるライブラリを作りました．今回はPythonを使って科学計算をする上で非常によく使われるNumPyをRustに渡して操作します．そのためにPyO3に加えて使うライブラリがrust-numpy¹です．

前回のプログラムを流用します．見ていない方はここから．

まず，Cargo.tomlに

[package]
name = "python_rust_example"
version = "0.1.0"

[lib]
name = "python_rust_example"
crate-type = ["cdylib"]

[dependencies]
numpy = "0.4.0"
ndarray = "0.12"

[dependencies.pyo3]
version = "0.5.2"
features = ["extension-module"]

とnumpyとndarrayを導入します．前者はRustでNumPyを扱えるようにするライブラリ（rust-numpyのこと），後者はRustにおけるNumPyのような線形代数ライブラリです．

次に，src/lib.rsに

#![feature(proc_macro)]

extern crate ndarray;
extern crate numpy;
extern crate pyo3;

use ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
use numpy::{IntoPyArray, PyArrayDyn};
use pyo3::prelude::*;

#[pymodinit]
fn python_rust_example(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "add_double")]
    fn add_double_py(_py: Python, a: f64, b: f64) -> PyResult<f64> {
        Ok(a + b)
    }

    #[pyfn(m, "add_numpy")]
    fn add_numpy_py(py: Python, x: &PyArrayDyn<f64>, y: &PyArrayDyn<f64>) -> PyResult<Py<PyArrayDyn<f64>>> {
        let x = x.as_array();
        let y = y.as_array();
        Ok((&x + &y).into_pyarray(py).to_owned())
    }

    #[pyfn(m, "mult_numpy")]
    fn mult_numpy_py(_py: Python, a: f64, x: &PyArrayDyn<f64>) -> PyResult<()> {
        let mut x = x.as_array_mut();
        x *= a;
        Ok(())
    }

    Ok(())
}

とadd_numpy()とmult_numpy()の2つをPythonから呼べるように新たに作ります．ここで，add_numpy()は2つのNumPy配列を受け取ってその和を新しいNumPy配列で返す関数です．一方，mult_numpy()は小数値とNumPy配列を受け取って，NumPy配列にその小数値を掛ける，つまり，配列自体の中身を書き換える関数です．

そして前回同様にビルドして，ライブラリをPython側で使えるようにtarget/release/から配置します．それからPythonで，

>>> from python_rust_example import add_numpy, mult_numpy
>>> import numpy as np
>>> a = np.arange(10, dtype=np.float)
>>> b = np.arange(10, dtype=np.float) * 2
>>> a
array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.])
>>> b
array([  0.,   2.,   4.,   6.,   8.,  10.,  12.,  14.,  16.,  18.])
>>>
>>> add_numpy(a, b)
array([  0.,   3.,   6.,   9.,  12.,  15.,  18.,  21.,  24.,  27.])
>>>
>>> mult_numpy(5, b)
>>> b
array([  0.,  10.,  20.,  30.,  40.,  50.,  60.,  70.,  80.,  90.])

と動作を確認します．なお，ちゃんとdtype=np.floatでないと型が違うというエラーが出るので気をつけましょう．

2019-10-28 追記

現在（rustc 1.40.0-nightly (95f437b3c 2019-10-27)）はこちらが動きました．（上述したものでも動くとは思いますが，一応最新のレポジトリのサンプルと合わせた）

[dependencies]
numpy = "0.7.0"
ndarray = "0.12"

[dependencies.pyo3]
version = "0.8"
features = ["extension-module"]

extern crate ndarray;
extern crate numpy;
extern crate pyo3;

use ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
use numpy::{IntoPyArray, PyArrayDyn};
use pyo3::prelude::*;

#[pymodule]
fn edge_adjacency(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "add_double")]
    fn add_double_py(_py: Python, a: f64, b: f64) -> PyResult<f64> {
        Ok(a + b)
    }

    #[pyfn(m, "add_numpy")]
    fn add_numpy_py(py: Python, x: &PyArrayDyn<f64>, y: &PyArrayDyn<f64>) -> PyResult<Py<PyArrayDyn<f64>>> {
        let x = x.as_array();
        let y = y.as_array();
        Ok((&x + &y).into_pyarray(py).to_owned())
    }

    #[pyfn(m, "mult_numpy")]
    fn mult_numpy_py(_py: Python, a: f64, x: &PyArrayDyn<f64>) -> PyResult<()> {
        let mut x = x.as_array_mut();
        x *= a;
        Ok(())
    }

    Ok(())
}