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

Feb. 22, 2019, 2:31 a.m. edited Oct. 28, 2019, 12:14 p.m.

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

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

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

まず,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"]

numpyndarrayを導入します.前者は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(())
}

  1. rust-numpyは,以前 (v0.2.1) rust-cpythonに沿っていたようですが,今はPyO3の方をメインでサポートしているようです