Rust ndarray 数组的初等函数数学运算

Question

我只是想为 Rust ndarray 做基本的数学运算（例如，sin、exp、log、sqrt ...）。但是，通过阅读 ndarray 的文档，我没有找到任何有用的示例。

比如说：

extern crate ndarray;

use ndarray as nd;

fn main() {
    let matrix = nd::array![[1., 2., 3.], [9., 8., 7.]];
    let result = some_math(matrix);
    println!("{}", result)
}

fn some_math(...) {
    //Here I would like to do elementwise exp() and sqrt
    sqrt(exp(...))
    // Using f64::exp would fail.
}

如何有效地实现这样的some_math？我当然可以通过循环遍历矩阵的元素来进行元素操作，但这听起来不太好，我不想这样做。

在 python 的numpy 中，这就是np.sqrt(np.exp(matrix))。我的意思是 Rust 确实是一门很棒的语言，但是，即使是简单的代数也很不方便（缺乏适当的生态系统）。

---

更新：ndarray 正在进行pull request。如果这被接受，那么您可以简单地执行matrix.exp().sqrt() 等。

ndarray-doc 中有一个非常 hidden page 告诉如何进行此类数学运算。

---

一些相关问题：12

Answer 1

如何高效地实现这样的some_math？

你可以使用mapv_into():

use ndarray as nd;
use ndarray::Array2;

fn some_math(matrix: Array2<f64>) -> Array2<f64> {
    // np.sqrt(np.exp(matrix)) would literally translate to equivalent to
    // matrix.mapv_into(f64::exp).mapv_into(f64::sqrt)
    // but this version iterates over the matrix just once
    matrix.mapv_into(|v| v.exp().sqrt())
}

fn main() {
    let matrix = nd::array![[1., 2., 3.], [9., 8., 7.]];
    let result = some_math(matrix);
    println!("{:?}", result)
}

Playground

这应该让您的性能与numpy 相当，但您应该确定。

要使用对大型阵列有意义的多核，您需要启用 crate 的 rayon 功能并使用 par_mapv_inplace()：

fn some_math(mut matrix: Array2<f64>) -> Array2<f64> {
    matrix.par_mapv_inplace(|v| v.exp().sqrt());
    matrix
}

（不能在 Playground 上编译，因为 Playground 的 ndarray 不包含 rayon 功能。）

请注意，在上述示例中，如果感觉更自然，您可以将 v.exp().sqrt() 替换为 f64::sqrt(f64::exp(v))。

---

编辑：我对时序很好奇，所以我决定做一个简单（且不科学）的基准测试——创建一个随机的 10_000x10_000 数组并将 np.sqrt(np.sqrt(array)) 与 Rust 等效项进行比较。

用于基准测试的 Python 代码：

import numpy as np
import time

matrix = np.random.rand(10000, 10000)

t0 = time.time()
np.sqrt(np.exp(matrix))
t1 = time.time()

print(t1 - t0)

锈代码：

use std::time::Instant;
use ndarray::Array2;
use ndarray_rand::{RandomExt, rand_distr::Uniform};

fn main() {
    let matrix: Array2<f64> = Array2::random((10000, 10000), Uniform::new(0., 1.));
    let t0 = Instant::now();
    let _result = matrix.mapv_into(|v| v.exp().sqrt());
    let elapsed = t0.elapsed();
    println!("{}", elapsed.as_secs_f64());
}

在我的古老桌面系统的实验中，Python 需要 3.7 秒 来计算，而 Rust 需要 2.5 秒。将 mapv_into() 替换为 par_mapv_inplace() 使 Rust 速度大大加快，现在时钟速度 0.5 秒，比同等 Python 快 7.4 倍。

单线程的 Rust 版本更快是有道理的，因为它只迭代整个数组一次，而 Python 会迭代两次。如果我们删除 sqrt() 操作，Python 的时钟速度为 2.8 秒，而 Rust 在 2.4 秒时仍然稍快（并且仍然是 0.5 秒并行）。我不确定是否可以在不使用 numba 之类的情况下优化 Python 版本。事实上，在不因手动进行低级计算而降低性能的情况下调整代码的能力是像 Rust 这样的编译语言的好处。

多线程版本是我不知道如何在 Python 中复制的东西，但是知道 numba 的人可以做到并进行比较。