Функция, принимающая ndarray и float одновременно
Я хотел бы написать функцию, вычисляющую показатель преломления некоторого материала в зависимости от длины волны. Я использую ящик ndarray для массивов, аналогичный numpy в Python.
На данный момент функция реализована следующим образом:
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: Как и просили, я включил тело функции.
@pretzelhammer: мой вопрос теперь включает тело функции. Нет жесткого требования, чтобы одна функция обрабатывала оба случая, это просто любопытство, поскольку я пытаюсь лучше познакомиться с ржавчиной и посмотреть, что возможно.
Ответы 1
Это определенно возможно, хотя это может потребовать больше работы, чем простое добавление еще одной функции. Сначала вам нужно будет определить два вспомогательных трейта, которые определяют операции mapv и powi (во многом похожие на Squared) как для f64, так и для Array1<f64>:
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> в дополнение ко многим другим ограничениям. Ошибки компилятора помогут вам точно определить, какие из них вам нужны.
Хотя это, вероятно, решит мою проблему, я заметил, что использование реализованного трейта вместо прямого использования .mapv(|x| x.powi(e)) приводит к увеличению времени выполнения почти на 30%, поэтому на данный момент я отказался от этой идеи.
Можете ли вы обновить свой ответ, чтобы показать реализации обеих версий функции? Кроме того, является ли жестким требованием, чтобы одна функция обрабатывала оба этих случая, или вам просто любопытно посмотреть, возможно ли это?