対象の配列と同じサイズのゼロ配列を作成する ー Rust-Ndarray例文集（第１回）

やりたいこと

任意の次元・大きさの配列が与えられたときに、それと全く同じサイズのゼロ配列を作りたい。つまり、Shape=[3, 4, 5]の３次元配列が与えられた際に、Shape=[3, 4, 5]の３次元ゼロ配列を取得したい。

ようは、Numpyのzeros_like関数みたいなのを使いたい。

解決策

対象の配列の次元が決まっているとき

配列の次元が決まっているなら、以下のように次元とサイズを取得して、それを並べてやれば実装可能だが、Traitにする際など、N次元配列に一般化したい場合はこの方法は使えない。

let a = /* ３次元配列 */;
let sh = a.shape();
let z = Array::zeros((sh[0], sh[1], sh[2]));

一般化された実装

以下のように書くと、一般化できる。

let z = Array::zeros(a.dim().clone());

dim()関数は配列の次元情報を返すので、それをzeros()関数に与える。ただし、zeros()関数は引数の所有権をmoveするので、clone()を忘れないようにする。

活用事例

配列の要素すべてに特定の操作を行い、その結果を新しい配列として返すメソッドを実装することを考える。

今回は全要素の平方根を取るsqrt()関数を持った、SqrtArr Traitを一般のN次元配列に実装してみた。この関数が行うのは以下の作業である。

1. 答えを格納する場所として、selfと同じ次元・大きさのゼロ配列ansを作成する。
2. selfの全要素について１つずつ平方根を求め、ansに格納する。
3. ansを返す。

use ndarray::*;

trait SqrtArr {
    fn sqrt(&self) -> Self;
}

impl<D: Dimension> SqrtArr for Array<f64, D> {
    fn sqrt(&self) -> Self {
        // 1.
        let mut ans: Array<f64, D> = Array::zeros(self.dim().clone());

        // 2.
        Zip::from(&mut ans)
            .and(self)
            .for_each(|a, &b| *a = b.sqrt());

        // 3.
        ans
    }
}

fn main() {
    let a = arr2(&[
        [ 0.0,  1.0,  4.0],
        [ 9.0, 16.0, 25.0],
    ]);

    println!("a.sqrt() = ");
    println!("{:?}", a.sqrt());
}

コードについて

このシリーズで取り扱ったコードは、

GitHub - doraneko94/ndarray-tutorial: RustのNdarrayクレートの実用例について書いていく。

RustのNdarrayクレートの実用例について書いていく。. Contribute to doraneko94/ndarray-tutorial development by creating an account on GitHub.

github.com

にて公開されている。