rust-ndarrayを使った行列演算

Rust

Rustの行列演算ライブラリ rust-ndarray についての紹介です。ほとんどが備忘録的なコード例となります。

rust-ndarray

Rustの行列演算ライブラリはいくつかあります。rust-ndarrayもその1つで、概ねプリミティブな行列演算のみをサポートしたものです。

github.com

ndarrayクレートを追加しておきます。

[dependencies]
ndarray = "0.13.0"

ベクトル

最初にベクトル演算です。Array1型でやりとりします。

  • 定数で初期化する場合、関数arr1で生成
  • ブラケット [i] で各要素にアクセス
  • sliceメソッドにマクロ s! の値を渡すことでスライスライクに複数要素へアクセスできる
use ndarray::prelude::*;

fn main() {
    // ベクトルの定義
    let vec1 = arr1(&[1, 2, 3]);
    println!("{}", vec1);  // [1, 2, 3]
    println!("{:?}", vec1.shape());  // [3]

    // 0番目の値を参照
    println!("{}", vec1[0]);  // 1

    // 1〜2番目の値にアクセス
    println!("{}", vec1.slice(s![1..3]));  // [2, 3]

    // ベクトル同士の和
    let vec2 = arr1(&[4, 5, 6]);
    let v = &vec1 + &vec2;
    println!("{}", v);  // [5, 7, 9]

    // ベクトル同士の積
    let v = &vec1 * &vec2;
    println!("{}", v);  // [4, 10, 18]
}

行列

次に行列です。行列の場合は、Array2型となります。

  • i行j列の値にはmat[[i, j]]でアクセス
  • *はアダマール積 (要素ごとの積)で、内積はdotメソッドで計算する
  • 転置はtメソッド
use ndarray::prelude::*;

fn main() {
    // 行列の定義
    let mat1 = arr2(&[
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ]);
    println!("{}", mat1);
    // [[1, 2, 3],
    //  [4, 5, 6],
    //  [7, 8, 9]]
    println!("{:?}", mat1.shape());  // [3, 3]

    // 1行目・2列目の値にアクセス
    println!("{}", mat1[[1, 2]]);  // 6

    // 行列のアダマール積
    let mat2 = arr2(&[
        [1, 1, 1],
        [1, 2, 4],
        [1, 3, 9]
    ]);

    let m = &mat1 * &mat2;
    println!("{}", m);
    // [[1, 2, 3],
    //  [4, 10, 24],
    //  [7, 24, 81]]

    // 転置
    let m = &mat1.t();
    println!("{}", m);
    // [[1, 4, 7],
    //  [2, 5, 8],
    //  [3, 6, 9]]

    // 行列の内積
    let m = mat1.dot(&mat2);
    println!("{}", m);
    // [[6, 14, 36],
    //  [15, 32, 78],
    //  [24, 50, 120]]

    // 行列とベクトルの積 (v = Ax, このときxは列ベクトル扱い)
    let x = arr1(&[1, 2, 3]);
    let v = mat1.dot(&x);
    println!("{}", v);  // [14, 32, 50]

    // ベクトルと行列の積 (v = xA, このときxは行ベクトル扱い)
    let v = &x.dot(&mat1);
    println!("{}", v);  // [30, 36, 42]
}

高度な機能が無いので、必要に応じて自前で実装する必要があります。例えば、行列式の計算は以下のように行います。

use ndarray::prelude::*;
use ndarray::Array2;

fn determinant(mat: &mut Array2<f64>) -> Option<f64> {
    if mat.shape()[0] != mat.shape()[1] {
        return None;
    }

    let n = mat.shape()[0];

    // 上三角行列を作る
    for i in 0..(n - 1) {
        if mat[[i, i]] == 0.0 {
            for j in (i + 1)..n {
                if mat[[j, i]] != 0.0 {
                    // 置換行列を作って行交換
                    let mut perm_mat: Array2<f64> = Array2::eye(n);
                    perm_mat[[i, i]] = 0.0;
                    perm_mat[[j, i]] = 1.0;
                    perm_mat[[j, j]] = 0.0;
                    perm_mat[[i, j]] = 1.0;

                    let sign = if (j - i) % 2 == 0 { 1.0 } else { -1.0 };

                    *mat = mat.dot(&perm_mat) * sign;
                }
            }

            // 他の全ての行が0なら終了
            break;
        }

        for j in (i + 1)..n {
            let c = mat[[j, i]] / mat[[i, i]];
            for k in 0..n {
                mat[[j, k]] = mat[[j, k]] - c * mat[[i, k]];
            }
        }
    }

    // 対角の値の積
    let det_mat = mat.diag().fold(1.0, |prod, x| prod * x);

    Some(det_mat)
}

fn main() {
    let mut mat = arr2(&[
        [2.0, 1.0, 3.0, 2.0],
        [6.0, 6.0, 10.0, 7.0],
        [2.0, 7.0, 6.0, 6.0],
        [4.0, 5.0, 10.0, 9.0]
    ]);
    let det = determinant(&mut mat);
    println!("{:?}", det);  // Some(-12.0)
}

まとめ

rust-ndarrayを使った行列演算について紹介しました。逆行列や特異値分解などが必要な場合、以下のndarray-linalgなどを検討したほうが良いかもしれません。

github.com