RustでPythonモジュールをつくってみた

はじめに

最近スパースモデリングに興味があり、スパース推定法による統計モデリングを読んでいました。（まだ途中までしか読んでないけど）ちょうどRustで何か書いてみたいと思っていたのでこの本の中に登場するエラスティックネットをRustで実装しようと思いました。

EvcxrやPolarsなどのライブラリを使えば、Jupyter NotebookでRustを使ったり、 EDAや前処理もできるのですが、個人的にRustでEDAや前処理するのが辛いなと感じました。そこでRustで書いたエラスティックネットをPythonから呼び出せるようにしました。

今回つくったものはこちらです。

ライブラリ

Pythonから呼び出せるモジュールをつくるために以下のライブラリを使いました。ほかにもパラメータ推定のためにrand, ndarray-randを使っています。

Rustライブラリ
- PyO3: PythonからRustでつくったプログラムを呼び出す用
- rust-numpy: NumPy配列とRustのndarrayを変換する用
- ndarray: Rustで使える多次元配列のライブラリ
Pythonライブラリ
- setuptools-rust: RustのプログラムをPythonモジュールとしてビルドする用

Cargo.tomlの設定

プログラムのビルドと外部ライブラリを使うためCargo.tomlに以下を追加しました。 cdylibは共有ライブラリをつくるために必要らしいです。

[lib]
crate-type = ["cdylib"]
name = "gmelasticnet" # ライブラリの名前

[dependencies]
ndarray = "0.15.3"
ndarray-rand = "0.14.0"
numpy = "0.15.0"
pyo3 = { version = "0.15.1", features = ["extension-module"] }
rand = "0.8.4"

Rustの実装

RustのプログラムをPythonから使うために Rustのコードを書き換える必要がありました。 PyO3のユーザーガイドを参考にRustコードを書きました。

エラスティックネットのパラメータ推定のコードはsrc/elasticnet.rsに書きました。パラメータ推定に使うクラスElasticNetは ElasticNet構造体やそのメソッド定義にアトリビュートをつけることで実装します。自作モジュールgmelasticnetにelasticnetという名前でサブモジュールを追加するために elasticnet関数を定義しています。この関数内でElasticNetクラスをモジュールに追加しています。

// Pythonに公開する構造体
#[pyclass]
struct ElasticNet {
    scaler: StandardScaler,
    y_mean: f64,
    xy_cov: Array1<f64>,
    x_cov: Array2<f64>,
}

// Pythonに公開しないメソッド
impl ElasticNet {
    fn n_features(&self) -> usize {
        self.xy_cov.len()
    }
}

// Pythonに公開するメソッド
#[pymethods]
impl ElasticNet {
    #[new] //コンストラクタを定義
    fn new(y: PyReadonlyArray1<f64>, x: PyReadonlyArray2<f64>) -> Self {
        // NumPy配列をRustのndarrayに変換
        let y = y.as_array();
        let x = x.as_array();

        // なんかいろいろ計算する

        Self {
            scaler,
            y_mean,
            xy_cov,
            x_cov,
        }
    }

    #[args(l1 = "0.1", l2 = "0.1", max_iter = "1000", tolerance = "1e-4")] // デフォルト引数を設定できる
    fn fit(
        &self,
        l1: f64,
        l2: f64,
        max_iter: i32,
        tolerance: f64,
        random_state: Option<u64>, // Optionの引数はNoneがデフォルト値
        model: Option<&ElasticNetParams>, // たしか Option<ElasticNetParams> だとエラーになった
    ) -> ElasticNetParams {
        // なんかいろいろ計算する

        ElasticNetParams::new(self.scaler.clone(), self.y_mean, coef)
    }
}

#[pymodule]
fn elasticnet(_py: Python, m: &PyModule) -> PyResult<()> {
    // Rustの構造体をPythonのクラスとして追加
    m.add_class::<ElasticNet>()?;
    m.add_class::<ElasticNetParams>()?;
    Ok(())
}

src/lib.rsは以下のとおりです。上のsrc/elasticnet.rsに同名のelasticnet関数を定義したせいで分かりにくいですが、 elasticnet関数をm.add_wrapped(wrap_pymodule!(elasticnet))とすることで gmelasticnetにelasticnetをサブモジュールとして追加しています。

mod elasticnet;

use elasticnet::*;
use pyo3::prelude::*;
use pyo3::wrap_pymodule;

// ここの関数名はCargo.tomlで設定したライブラリの名前と同じにしないといけないみたい
#[pymodule]
fn gmelasticnet(_py: Python, m: &PyModule) -> PyResult<()> {
    // elasticnet関数を渡すことで
    // elasticnetをgmelasticnetのサブモジュールとして追加
    m.add_wrapped(wrap_pymodule!(elasticnet))?;
    Ok(())
}

ビルド

setuptools-rust のREADME.mdを参考にいろいろ設定しました。

setup.pyは下のようになりました。

from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    name="gmelasticnet",
    version="0.1.0",
    rust_extensions=[
        RustExtension("gmelasticnet.gmelasticnet", binding=Binding.PyO3)
    ],
    packages=["gmelasticnet"],
    zip_safe=False,
)

下のコマンドを実行して自作モジュールをインストールします。

python ./setup.py install

pip listで確認すると

gmelasticnet                 0.1.0

自分で作成したモジュールがインストールされていました！！

おわりに

今回作成したモジュールだとRustのコードで変更した箇所は

アトリビュートの追加
モジュール、クラスを追加する関数の定義
NumPy配列の変換

だけなので、意外と簡単にRustでPythonモジュールをつくることができました。