同时接受ndarray和float的函数

Question

我想写一个函数来计算某种材料的折射率作为波长的函数。我对数组使用ndarray crate，类似于 Python 中的 numpy。 目前该功能的实现方式如下：

fn calc_refractive_index(l: &Array1<f64>) -> Array1<f64> {
    fn refractive_index(&self, l: &Array1<f64>) -> Array1<f64> {
        let V =
            self.A[0] + self.A[1] / (1. + self.A[2] * (self.A[3] * l / self.pitch).mapv(f64::exp));
        let W =
            self.B[0] + self.B[1] / (1. + self.B[2] * (self.B[3] * l / self.pitch).mapv(f64::exp));
        (self.material.refractive_index(&l).powi(2)
            - 3. * l.powi(2) / 4. / (consts::PI * self.pitch).powi(2) * (V.powi(2) - W.powi(2)))
            .mapv(f64::sqrt)
    }
}

实现了以下特征，以缩短增强时间：

trait Squared<T> {
    fn powi(&self, e: i32) -> T;
}

impl Squared<Array1<f64>> for Array1<f64> {
    fn powi(&self, e: i32) -> Array1<f64> {
        self.mapv(|a| a.powi(e))
    }
}

但是，我可能还想只计算一个特定波长的折射率，所以我也想接受 float64 值。在不实现两个单独的功能的情况下，实现这一点的最佳方法是什么？

编辑： 该函数使用 ndarray 的语法进行平方：l.mapv(|x| x.powi(2))，不幸的是，这与 float64 不同

编辑 2： 按照要求，我包含了函数体。

Answer 1

这绝对是可能的，尽管它可能比简单地添加另一个功能需要更多的工作。您需要首先定义两个辅助特征，为f64 和Array1<f64> 指定mapv 和powi（很像Squared）操作：

trait PowI {
    fn powi(&self, e: i32) -> Self;
}

impl PowI for f64 {
    fn powi(&self, e: i32) -> Self {
        f64::powi(*self, e)
    }
}

impl PowI for Array1<f64> {
    fn powi(&self, e: i32) -> Self {
        self.mapv(|a| a.powi(e))
    }
}

trait MapV {
    fn mapv(&self, f: impl FnMut(f64) -> f64) -> Self;
}

impl MapV for f64 {
    fn mapv(&self, mut f: impl FnMut(f64) -> f64) -> Self {
        f(*self)
    }
}

impl MapV for Array1<f64> {
    fn mapv(&self, f: impl FnMut(f64) -> f64) -> Self {
        Array1::mapv(self, f)
    }
}

然后，您可以使您的折射率计算函数对实现这两个特征的类型T 通用（即T: PowI + MapV）。

请注意，您还需要在 T 上指定附加边界，以指定它可以与 Ts 和 f64s 相加、除法和相乘（我假设这是 self.pitch 的类型和元素在self.A 和self.B 数组中）。您可以通过要求 T、&'a T 甚至 f64 实现适当的 Add、Mul 和 Div 特征来做到这一点。例如，您的实现可能需要f64: Mul<&'a T, Output=T> 和T: Div<f64, Output=T> 以及许多其他边界。编译器错误将帮助您准确确定需要哪些错误。