use crate::array::Array;
use crate::error::{NumRs2Error, Result};
use num_traits::{Float, NumCast};
use scirs2_core::ndarray::{Array1, ArrayView1};
use scirs2_core::simd_ops::SimdUnifiedOps;
use std::fmt::{self, Debug};
const SIMD_THRESHOLD: usize = 64;
fn to_array_view<T: Clone>(arr: &Array<T>) -> Array1<T> {
Array1::from_vec(arr.to_vec())
}
fn from_array1<T: Clone + Debug + NumCast>(arr: Array1<T>, shape: &[usize]) -> Array<T> {
let data: Vec<T> = arr.into_iter().collect();
Array::from_vec(data).reshape(shape)
}
fn should_use_simd(len: usize) -> bool {
len >= SIMD_THRESHOLD
}
pub struct BinaryUfunc<F>
where
F: Fn(f64, f64) -> f64,
{
func: F,
name: &'static str,
}
pub struct UnaryUfunc<F>
where
F: Fn(f64) -> f64,
{
func: F,
name: &'static str,
}
impl<F> fmt::Debug for BinaryUfunc<F>
where
F: Fn(f64, f64) -> f64,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "BinaryUfunc({})", self.name)
}
}
impl<F> fmt::Debug for UnaryUfunc<F>
where
F: Fn(f64) -> f64,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "UnaryUfunc({})", self.name)
}
}
impl<F> BinaryUfunc<F>
where
F: Fn(f64, f64) -> f64,
{
pub fn new(func: F, name: &'static str) -> Self {
Self { func, name }
}
pub fn call(&self, a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
a.zip_with(b, |x, y| (self.func)(x, y))
}
pub fn call_scalar_right(&self, a: &Array<f64>, b: f64) -> Array<f64> {
a.map(|x| (self.func)(x, b))
}
pub fn call_scalar_left(&self, a: f64, b: &Array<f64>) -> Array<f64> {
b.map(|x| (self.func)(a, x))
}
}
impl<F> UnaryUfunc<F>
where
F: Fn(f64) -> f64,
{
pub fn new(func: F, name: &'static str) -> Self {
Self { func, name }
}
pub fn call(&self, a: &Array<f64>) -> Array<f64> {
a.map(|x| (self.func)(x))
}
}
fn get_add_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| a + b, "add")
}
fn get_subtract_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| a - b, "subtract")
}
fn get_multiply_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| a * b, "multiply")
}
fn get_divide_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| a / b, "divide")
}
fn get_power_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| a.powf(b), "power")
}
fn get_maximum_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| if a > b { a } else { b }, "maximum")
}
fn get_minimum_ufunc() -> BinaryUfunc<fn(f64, f64) -> f64> {
BinaryUfunc::new(|a, b| if a < b { a } else { b }, "minimum")
}
pub fn add(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_add(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_add_ufunc().call(a, b)
}
pub fn subtract(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_sub(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_subtract_ufunc().call(a, b)
}
pub fn multiply(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_mul(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_multiply_ufunc().call(a, b)
}
pub fn divide(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_div(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_divide_ufunc().call(a, b)
}
pub fn power(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_pow(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_power_ufunc().call(a, b)
}
pub fn maximum(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_max(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_maximum_ufunc().call(a, b)
}
pub fn minimum(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.len() == b.len() && should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_min(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
get_minimum_ufunc().call(a, b)
}
pub fn add_scalar(a: &Array<f64>, b: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = Array1::from_elem(a.len(), b);
let result = f64::simd_add(&a_nd.view(), &b_nd.view());
return from_array1(result, &a.shape());
}
get_add_ufunc().call_scalar_right(a, b)
}
pub fn subtract_scalar(a: &Array<f64>, b: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = Array1::from_elem(a.len(), b);
let result = f64::simd_sub(&a_nd.view(), &b_nd.view());
return from_array1(result, &a.shape());
}
get_subtract_ufunc().call_scalar_right(a, b)
}
pub fn multiply_scalar(a: &Array<f64>, b: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_scalar_mul(&a_nd.view(), b);
return from_array1(result, &a.shape());
}
get_multiply_ufunc().call_scalar_right(a, b)
}
pub fn divide_scalar(a: &Array<f64>, b: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_scalar_mul(&a_nd.view(), 1.0 / b);
return from_array1(result, &a.shape());
}
get_divide_ufunc().call_scalar_right(a, b)
}
pub fn power_scalar(a: &Array<f64>, b: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_powf(&a_nd.view(), b);
return from_array1(result, &a.shape());
}
get_power_ufunc().call_scalar_right(a, b)
}
pub fn negative(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_scalar_mul(&a_nd.view(), -1.0);
return from_array1(result, &a.shape());
}
a.map(|x| -x)
}
pub fn absolute(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_abs(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.abs())
}
pub fn abs(a: &Array<f64>) -> Array<f64> {
absolute(a)
}
pub fn square(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_square(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x * x)
}
pub fn sqrt(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sqrt(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.sqrt())
}
pub fn cbrt(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_cbrt(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.cbrt())
}
pub fn reciprocal(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_recip(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| 1.0 / x)
}
pub fn rsqrt(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_rsqrt(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| 1.0 / x.sqrt())
}
pub fn exp(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_exp(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.exp())
}
pub fn exp2(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_exp2(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.exp2())
}
pub fn expm1(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_exp_m1(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.exp_m1())
}
pub fn log(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_ln(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.ln())
}
pub fn ln(a: &Array<f64>) -> Array<f64> {
log(a)
}
pub fn log2(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_log2(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.log2())
}
pub fn log10(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_log10(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.log10())
}
pub fn log1p(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_ln_1p(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.ln_1p())
}
pub fn sin(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sin(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.sin())
}
pub fn cos(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_cos(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.cos())
}
pub fn tan(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_tan(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.tan())
}
pub fn sincos(a: &Array<f64>) -> (Array<f64>, Array<f64>) {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let (sin_result, cos_result) = f64::simd_sincos(&a_nd.view());
return (
from_array1(sin_result, &a.shape()),
from_array1(cos_result, &a.shape()),
);
}
(a.map(|x| x.sin()), a.map(|x| x.cos()))
}
pub fn sinc(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sinc(&a_nd.view());
return from_array1(result, &a.shape());
}
use std::f64::consts::PI;
a.map(|x| {
if x == 0.0 {
1.0
} else {
let px = PI * x;
px.sin() / px
}
})
}
pub fn arcsin(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_asin(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.asin())
}
pub fn asin(a: &Array<f64>) -> Array<f64> {
arcsin(a)
}
pub fn arccos(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_acos(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.acos())
}
pub fn acos(a: &Array<f64>) -> Array<f64> {
arccos(a)
}
pub fn arctan(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_atan(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.atan())
}
pub fn atan(a: &Array<f64>) -> Array<f64> {
arctan(a)
}
pub fn arctan2(y: &Array<f64>, x: &Array<f64>) -> Result<Array<f64>> {
if y.shape() != x.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: x.shape(),
actual: y.shape(),
});
}
if should_use_simd(y.len()) {
let y_nd = to_array_view(y);
let x_nd = to_array_view(x);
let result = f64::simd_atan2(&y_nd.view(), &x_nd.view());
return Ok(from_array1(result, &y.shape()));
}
let y_data = y.to_vec();
let x_data = x.to_vec();
let result: Vec<f64> = y_data
.iter()
.zip(x_data.iter())
.map(|(yi, xi)| yi.atan2(*xi))
.collect();
Ok(Array::from_vec(result).reshape(&y.shape()))
}
pub fn atan2(y: &Array<f64>, x: &Array<f64>) -> Result<Array<f64>> {
arctan2(y, x)
}
pub fn sinh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sinh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.sinh())
}
pub fn cosh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_cosh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.cosh())
}
pub fn tanh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_tanh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.tanh())
}
pub fn arcsinh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_asinh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.asinh())
}
pub fn asinh(a: &Array<f64>) -> Array<f64> {
arcsinh(a)
}
pub fn arccosh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_acosh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.acosh())
}
pub fn acosh(a: &Array<f64>) -> Array<f64> {
arccosh(a)
}
pub fn arctanh(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_atanh(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.atanh())
}
pub fn atanh(a: &Array<f64>) -> Array<f64> {
arctanh(a)
}
pub fn floor(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_floor(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.floor())
}
pub fn ceil(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_ceil(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.ceil())
}
pub fn round(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_round(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.round())
}
pub fn trunc(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_trunc(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.trunc())
}
pub fn fract(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_fract(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.fract())
}
pub fn sign(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sign(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x == 0.0 {
0.0
} else if x > 0.0 {
1.0
} else {
-1.0
}
})
}
pub fn clip(a: &Array<f64>, min: f64, max: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_clip(&a_nd.view(), min, max);
return from_array1(result, &a.shape());
}
a.map(|x| x.clamp(min, max))
}
pub fn copysign(mag: &Array<f64>, sign: &Array<f64>) -> Result<Array<f64>> {
if mag.shape() != sign.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: mag.shape(),
actual: sign.shape(),
});
}
if should_use_simd(mag.len()) {
let mag_nd = to_array_view(mag);
let sign_nd = to_array_view(sign);
let result = f64::simd_copysign(&mag_nd.view(), &sign_nd.view());
return Ok(from_array1(result, &mag.shape()));
}
let mag_data = mag.to_vec();
let sign_data = sign.to_vec();
let result: Vec<f64> = mag_data
.iter()
.zip(sign_data.iter())
.map(|(m, s)| m.copysign(*s))
.collect();
Ok(Array::from_vec(result).reshape(&mag.shape()))
}
pub fn deg2rad(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_to_radians(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.to_radians())
}
pub fn radians(a: &Array<f64>) -> Array<f64> {
deg2rad(a)
}
pub fn rad2deg(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_to_degrees(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.to_degrees())
}
pub fn degrees(a: &Array<f64>) -> Array<f64> {
rad2deg(a)
}
pub fn hypot(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.shape() != b.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: a.shape(),
actual: b.shape(),
});
}
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_hypot(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
let a_data = a.to_vec();
let b_data = b.to_vec();
let result: Vec<f64> = a_data
.iter()
.zip(b_data.iter())
.map(|(ai, bi)| ai.hypot(*bi))
.collect();
Ok(Array::from_vec(result).reshape(&a.shape()))
}
pub fn fma(a: &Array<f64>, b: &Array<f64>, c: &Array<f64>) -> Result<Array<f64>> {
if a.shape() != b.shape() || a.shape() != c.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: a.shape(),
actual: if a.shape() != b.shape() {
b.shape()
} else {
c.shape()
},
});
}
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let c_nd = to_array_view(c);
let result = f64::simd_fma(&a_nd.view(), &b_nd.view(), &c_nd.view());
return Ok(from_array1(result, &a.shape()));
}
let a_data = a.to_vec();
let b_data = b.to_vec();
let c_data = c.to_vec();
let result: Vec<f64> = a_data
.iter()
.zip(b_data.iter())
.zip(c_data.iter())
.map(|((ai, bi), ci)| ai.mul_add(*bi, *ci))
.collect();
Ok(Array::from_vec(result).reshape(&a.shape()))
}
pub fn lerp(a: &Array<f64>, b: &Array<f64>, t: f64) -> Result<Array<f64>> {
if a.shape() != b.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: a.shape(),
actual: b.shape(),
});
}
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_lerp(&a_nd.view(), &b_nd.view(), t);
return Ok(from_array1(result, &a.shape()));
}
let a_data = a.to_vec();
let b_data = b.to_vec();
let result: Vec<f64> = a_data
.iter()
.zip(b_data.iter())
.map(|(ai, bi)| ai + t * (bi - ai))
.collect();
Ok(Array::from_vec(result).reshape(&a.shape()))
}
pub fn logaddexp(a: &Array<f64>, b: &Array<f64>) -> Result<Array<f64>> {
if a.shape() != b.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: a.shape(),
actual: b.shape(),
});
}
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
let result = f64::simd_logaddexp(&a_nd.view(), &b_nd.view());
return Ok(from_array1(result, &a.shape()));
}
let a_data = a.to_vec();
let b_data = b.to_vec();
let result: Vec<f64> = a_data
.iter()
.zip(b_data.iter())
.map(|(ai, bi)| {
let max = ai.max(*bi);
max + ((-(*ai - max).abs()).exp() + (-(*bi - max).abs()).exp()).ln()
})
.collect();
Ok(Array::from_vec(result).reshape(&a.shape()))
}
pub fn dot(a: &Array<f64>, b: &Array<f64>) -> Result<f64> {
if a.len() != b.len() {
return Err(NumRs2Error::ShapeMismatch {
expected: vec![a.len()],
actual: vec![b.len()],
});
}
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let b_nd = to_array_view(b);
return Ok(f64::simd_dot(&a_nd.view(), &b_nd.view()));
}
let a_data = a.to_vec();
let b_data = b.to_vec();
Ok(a_data
.iter()
.zip(b_data.iter())
.map(|(ai, bi)| ai * bi)
.sum())
}
pub fn norm_l2(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_norm(&a_nd.view());
}
let data = a.to_vec();
data.iter().map(|x| x * x).sum::<f64>().sqrt()
}
pub fn norm_l1(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_norm_l1(&a_nd.view());
}
let data = a.to_vec();
data.iter().map(|x| x.abs()).sum::<f64>()
}
pub fn sum(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_sum(&a_nd.view());
}
a.to_vec().iter().sum()
}
pub fn mean(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_mean(&a_nd.view());
}
let data = a.to_vec();
data.iter().sum::<f64>() / data.len() as f64
}
pub fn var(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_variance(&a_nd.view());
}
let data = a.to_vec();
let n = data.len() as f64;
let mean = data.iter().sum::<f64>() / n;
data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n
}
pub fn std(a: &Array<f64>) -> f64 {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
return f64::simd_std(&a_nd.view());
}
var(a).sqrt()
}
pub fn erf(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_erf(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
let t = 1.0 / (1.0 + 0.5 * x.abs());
let tau = t
* (-x * x - 1.26551223
+ t * (1.00002368
+ t * (0.37409196
+ t * (0.09678418
+ t * (-0.18628806
+ t * (0.27886807
+ t * (-1.13520398
+ t * (1.48851587
+ t * (-0.82215223 + t * 0.17087277)))))))))
.exp();
if x >= 0.0 {
1.0 - tau
} else {
tau - 1.0
}
})
}
pub fn erfc(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_erfc(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| 1.0 - erf(&Array::from_vec(vec![x])).to_vec()[0])
}
pub fn gamma(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_gamma(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x <= 0.0 && x.floor() == x {
f64::INFINITY
} else {
let g = 7;
let c = [
0.99999999999980993,
676.5203681218851,
-1259.1392167224028,
771.32342877765313,
-176.61502916214059,
12.507343278686905,
-0.13857109526572012,
9.9843695780195716e-6,
1.5056327351493116e-7,
];
if x < 0.5 {
std::f64::consts::PI / ((std::f64::consts::PI * x).sin() * gamma_single(1.0 - x))
} else {
let x = x - 1.0;
let mut y = c[0];
for i in 1..=g + 1 {
y += c[i] / (x + i as f64);
}
let t = x + g as f64 + 0.5;
(2.0 * std::f64::consts::PI).sqrt() * t.powf(x + 0.5) * (-t).exp() * y
}
}
})
}
fn gamma_single(x: f64) -> f64 {
let g = 7;
let c = [
0.99999999999980993,
676.5203681218851,
-1259.1392167224028,
771.32342877765313,
-176.61502916214059,
12.507343278686905,
-0.13857109526572012,
9.9843695780195716e-6,
1.5056327351493116e-7,
];
let x = x - 1.0;
let mut y = c[0];
for i in 1..=g + 1 {
y += c[i] / (x + i as f64);
}
let t = x + g as f64 + 0.5;
(2.0 * std::f64::consts::PI).sqrt() * t.powf(x + 0.5) * (-t).exp() * y
}
pub fn lgamma(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_ln_gamma(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| gamma_single(x).abs().ln())
}
pub fn gammaln(a: &Array<f64>) -> Array<f64> {
lgamma(a)
}
pub fn digamma(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_digamma(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x <= 0.0 {
f64::NAN
} else if x < 6.0 {
digamma_scalar(x + 1.0) - 1.0 / x
} else {
let inv_x = 1.0 / x;
let inv_x2 = inv_x * inv_x;
x.ln() - 0.5 * inv_x - inv_x2 * (1.0 / 12.0 - inv_x2 * (1.0 / 120.0 - inv_x2 / 252.0))
}
})
}
fn digamma_scalar(x: f64) -> f64 {
if x < 6.0 {
digamma_scalar(x + 1.0) - 1.0 / x
} else {
let inv_x = 1.0 / x;
let inv_x2 = inv_x * inv_x;
x.ln() - 0.5 * inv_x - inv_x2 * (1.0 / 12.0 - inv_x2 * (1.0 / 120.0 - inv_x2 / 252.0))
}
}
pub fn sigmoid(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_sigmoid(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| 1.0 / (1.0 + (-x).exp()))
}
pub fn logit(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_logit(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| (x / (1.0 - x)).ln())
}
pub fn gelu(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_gelu(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
let sqrt_2_over_pi = (2.0 / std::f64::consts::PI).sqrt();
0.5 * x * (1.0 + (sqrt_2_over_pi * (x + 0.044715 * x.powi(3))).tanh())
})
}
pub fn swish(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_swish(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x / (1.0 + (-x).exp()))
}
pub fn silu(a: &Array<f64>) -> Array<f64> {
swish(a)
}
pub fn softplus(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_softplus(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x > 20.0 {
x
} else if x < -20.0 {
x.exp()
} else {
(1.0 + x.exp()).ln()
}
})
}
pub fn mish(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_mish(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
let sp = if x > 20.0 { x } else { (1.0 + x.exp()).ln() };
x * sp.tanh()
})
}
pub fn elu(a: &Array<f64>, alpha: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_elu(&a_nd.view(), alpha);
return from_array1(result, &a.shape());
}
a.map(|x| if x > 0.0 { x } else { alpha * (x.exp() - 1.0) })
}
pub fn selu(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_selu(&a_nd.view());
return from_array1(result, &a.shape());
}
let lambda = 1.0507009873554804934193349852946;
let alpha = 1.6732632423543772848170429916717;
a.map(|x| lambda * if x > 0.0 { x } else { alpha * (x.exp() - 1.0) })
}
pub fn hardsigmoid(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_hardsigmoid(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x <= -3.0 {
0.0
} else if x >= 3.0 {
1.0
} else {
x / 6.0 + 0.5
}
})
}
pub fn hardswish(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_hardswish(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| {
if x <= -3.0 {
0.0
} else if x >= 3.0 {
x
} else {
x * (x + 3.0) / 6.0
}
})
}
pub fn relu(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_relu(&a_nd.view());
return from_array1(result, &a.shape());
}
a.map(|x| x.max(0.0))
}
pub fn leaky_relu(a: &Array<f64>, alpha: f64) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_leaky_relu(&a_nd.view(), alpha);
return from_array1(result, &a.shape());
}
a.map(|x| if x > 0.0 { x } else { alpha * x })
}
pub fn prelu(a: &Array<f64>, alpha: &Array<f64>) -> Result<Array<f64>> {
if a.shape() != alpha.shape() {
return Err(NumRs2Error::ShapeMismatch {
expected: a.shape(),
actual: alpha.shape(),
});
}
let a_data = a.to_vec();
let alpha_data = alpha.to_vec();
let result: Vec<f64> = a_data
.iter()
.zip(alpha_data.iter())
.map(|(x, alpha)| if *x > 0.0 { *x } else { alpha * x })
.collect();
Ok(Array::from_vec(result).reshape(&a.shape()))
}
pub fn log_softmax(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_log_softmax(&a_nd.view());
return from_array1(result, &a.shape());
}
let data = a.to_vec();
let max_val = data.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let shifted: Vec<f64> = data.iter().map(|x| x - max_val).collect();
let log_sum_exp = shifted.iter().map(|x| x.exp()).sum::<f64>().ln();
let result: Vec<f64> = shifted.iter().map(|x| x - log_sum_exp).collect();
Array::from_vec(result).reshape(&a.shape())
}
pub fn softmax(a: &Array<f64>) -> Array<f64> {
if should_use_simd(a.len()) {
let a_nd = to_array_view(a);
let result = f64::simd_softmax(&a_nd.view());
return from_array1(result, &a.shape());
}
let data = a.to_vec();
let max_val = data.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let exps: Vec<f64> = data.iter().map(|x| (x - max_val).exp()).collect();
let sum: f64 = exps.iter().sum();
let result: Vec<f64> = exps.iter().map(|x| x / sum).collect();
Array::from_vec(result).reshape(&a.shape())
}
pub fn smoothstep(edge0: f64, edge1: f64, x: &Array<f64>) -> Array<f64> {
if should_use_simd(x.len()) {
let x_nd = to_array_view(x);
let result = f64::simd_smoothstep(edge0, edge1, &x_nd.view());
return from_array1(result, &x.shape());
}
x.map(|v| {
let t = ((v - edge0) / (edge1 - edge0)).clamp(0.0, 1.0);
t * t * (3.0 - 2.0 * t)
})
}
pub fn smootherstep(edge0: f64, edge1: f64, x: &Array<f64>) -> Array<f64> {
if should_use_simd(x.len()) {
let x_nd = to_array_view(x);
let result = f64::simd_smootherstep(edge0, edge1, &x_nd.view());
return from_array1(result, &x.shape());
}
x.map(|v| {
let t = ((v - edge0) / (edge1 - edge0)).clamp(0.0, 1.0);
t * t * t * (t * (t * 6.0 - 15.0) + 10.0)
})
}