use crate::array::Array;
use crate::error::Result;
use num_traits::{Float, NumCast, One, Zero};
use std::ops::{Add, Div, Mul, Sub};
use crate::math::{from_nd_array, to_nd_array_f32, to_nd_array_f64};
use scirs2_core::simd_ops::SimdUnifiedOps;
pub fn argmax<T>(array: &Array<T>, axis: Option<usize>, keepdims: bool) -> Result<Array<usize>>
where
T: PartialOrd + Clone + Zero,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot find argmax of empty array".to_string(),
));
}
match axis {
None => {
let data = array.to_vec();
let mut max_idx = 0;
let mut max_val = &data[0];
for (i, val) in data.iter().enumerate().skip(1) {
if val > max_val {
max_val = val;
max_idx = i;
}
}
Ok(Array::from_vec(vec![max_idx]))
}
Some(ax) => {
if ax >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
ax,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[ax];
let mut out_shape = shape.clone();
if keepdims {
out_shape[ax] = 1;
} else {
out_shape.remove(ax);
}
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![0_usize; out_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
for out_idx in 0..out_size {
let mut indices = vec![0; array.ndim()];
let mut temp = out_idx;
for i in 0..array.ndim() {
if i < ax {
let dim_size = shape[i];
indices[i] = temp % dim_size;
temp /= dim_size;
} else if i > ax || (i == ax && keepdims) {
let dim_idx = if keepdims { i } else { i - 1 };
if dim_idx < out_shape.len() {
let dim_size = out_shape[dim_idx];
indices[i] = temp % dim_size;
temp /= dim_size;
}
}
}
let mut max_idx = 0;
let mut max_val = None;
for j in 0..axis_size {
indices[ax] = j;
let val = array.get(&indices)?;
if max_val.as_ref().is_none_or(|mv| &val > mv) {
max_val = Some(val);
max_idx = j;
}
}
result_data[out_idx] = max_idx;
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn argmin<T>(array: &Array<T>, axis: Option<usize>, keepdims: bool) -> Result<Array<usize>>
where
T: PartialOrd + Clone + Zero,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot find argmin of empty array".to_string(),
));
}
match axis {
None => {
let data = array.to_vec();
let mut min_idx = 0;
let mut min_val = &data[0];
for (i, val) in data.iter().enumerate().skip(1) {
if val < min_val {
min_val = val;
min_idx = i;
}
}
Ok(Array::from_vec(vec![min_idx]))
}
Some(ax) => {
if ax >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
ax,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[ax];
let mut out_shape = shape.clone();
if keepdims {
out_shape[ax] = 1;
} else {
out_shape.remove(ax);
}
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![0_usize; out_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
for out_idx in 0..out_size {
let mut indices = vec![0; array.ndim()];
let mut temp = out_idx;
for i in 0..array.ndim() {
if i < ax {
let dim_size = shape[i];
indices[i] = temp % dim_size;
temp /= dim_size;
} else if i > ax || (i == ax && keepdims) {
let dim_idx = if keepdims { i } else { i - 1 };
if dim_idx < out_shape.len() {
let dim_size = out_shape[dim_idx];
indices[i] = temp % dim_size;
temp /= dim_size;
}
}
}
let mut min_idx = 0;
let mut min_val = None;
for j in 0..axis_size {
indices[ax] = j;
let val = array.get(&indices)?;
if min_val.as_ref().is_none_or(|mv| &val < mv) {
min_val = Some(val);
min_idx = j;
}
}
result_data[out_idx] = min_idx;
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn argsort<T>(
array: &Array<T>,
axis: Option<isize>,
_kind: Option<&str>,
_order: Option<&[&str]>,
) -> Result<Array<usize>>
where
T: PartialOrd + Clone + Zero,
{
let axis = if let Some(ax) = axis {
if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
}
} else {
let data = array.to_vec();
let mut indices: Vec<usize> = (0..data.len()).collect();
indices.sort_by(|&a, &b| {
data[a]
.partial_cmp(&data[b])
.unwrap_or(std::cmp::Ordering::Equal)
});
return Ok(Array::from_vec(indices));
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
let result_shape = shape.clone();
let total_size: usize = shape.iter().product();
let mut result_data = vec![0_usize; total_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
let n_sorts = total_size / axis_size;
for sort_idx in 0..n_sorts {
let mut values_with_indices: Vec<(T, usize)> = Vec::with_capacity(axis_size);
let mut base_indices = vec![0; array.ndim()];
let mut temp = sort_idx;
let _dim_idx = 0;
for i in 0..array.ndim() {
if i != axis {
let size = shape[i];
base_indices[i] = temp % size;
temp /= size;
}
}
for j in 0..axis_size {
base_indices[axis] = j;
let val = array.get(&base_indices)?;
values_with_indices.push((val, j));
}
values_with_indices
.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
for (k, (_, idx)) in values_with_indices.into_iter().enumerate() {
base_indices[axis] = k;
let flat_idx = base_indices
.iter()
.enumerate()
.map(|(i, &idx)| idx * strides[i])
.sum::<usize>();
result_data[flat_idx] = idx;
}
}
Ok(Array::from_vec(result_data).reshape(&result_shape))
}
pub fn around<T>(
array: &Array<T>,
decimals: Option<i32>,
out: Option<&mut Array<T>>,
) -> Result<Array<T>>
where
T: Float + Clone,
{
let decimals = decimals.unwrap_or(0);
let multiplier = T::from(10.0_f64.powi(decimals)).expect("power of 10 should be representable");
let result = array.map(|x| {
if decimals >= 0 {
(x * multiplier).round() / multiplier
} else {
let divisor =
T::from(10.0_f64.powi(-decimals)).expect("power of 10 should be representable");
(x / divisor).round() * divisor
}
});
if let Some(out_array) = out {
if out_array.shape() != result.shape() {
return Err(crate::error::NumRs2Error::ShapeMismatch {
expected: result.shape(),
actual: out_array.shape(),
});
}
*out_array = result;
Ok(out_array.clone())
} else {
Ok(result)
}
}
pub fn cumsum<T>(
array: &Array<T>,
axis: Option<isize>,
_out: Option<&mut Array<T>>,
) -> Result<Array<T>>
where
T: Float + Clone + Add<Output = T> + Send + Sync + 'static,
{
use scirs2_core::parallel_ops::*;
const PARALLEL_THRESHOLD: usize = 1000;
if array.is_empty() {
return Ok(array.clone());
}
match axis {
None => {
if array.len() >= 64 {
if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f64>() {
let nd = to_nd_array_f64(unsafe {
std::mem::transmute::<&Array<T>, &Array<f64>>(array)
});
let result = f64::simd_cumsum(&nd.view());
return Ok(unsafe {
std::mem::transmute::<Array<f64>, Array<T>>(from_nd_array(
result,
&array.shape(),
))
});
}
if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() {
let nd = to_nd_array_f32(unsafe {
std::mem::transmute::<&Array<T>, &Array<f32>>(array)
});
let result = f32::simd_cumsum(&nd.view());
return Ok(unsafe {
std::mem::transmute::<Array<f32>, Array<T>>(from_nd_array(
result,
&array.shape(),
))
});
}
}
let data = array.to_vec();
let mut result = Vec::with_capacity(data.len());
let mut sum = T::zero();
for val in data {
sum = sum + val;
result.push(sum);
}
Ok(Array::from_vec(result))
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let data = array.to_vec();
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
let axis_size = shape[axis];
let n_iterations = array.size() / axis_size;
if is_parallel_enabled() && n_iterations >= PARALLEL_THRESHOLD {
let result_data: Vec<T> = (0..n_iterations)
.into_par_iter()
.flat_map(|iter_idx| {
let mut indices = vec![0; shape.len()];
let mut temp = iter_idx;
for i in (0..shape.len()).rev() {
if i != axis {
indices[i] = temp % shape[i];
temp /= shape[i];
}
}
let mut local_results = Vec::with_capacity(axis_size);
let mut sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
let flat_idx = indices
.iter()
.enumerate()
.map(|(i, &idx)| idx * strides[i])
.sum::<usize>();
sum = sum + data[flat_idx];
local_results.push((flat_idx, sum));
}
local_results
})
.collect::<Vec<(usize, T)>>()
.into_iter()
.fold(vec![T::zero(); array.size()], |mut acc, (idx, val)| {
acc[idx] = val;
acc
});
Ok(Array::from_vec(result_data).reshape(&shape))
} else {
let mut result_data = vec![T::zero(); array.size()];
for iter_idx in 0..n_iterations {
let mut indices = vec![0; array.ndim()];
let mut temp = iter_idx;
for i in (0..array.ndim()).rev() {
if i != axis {
indices[i] = temp % shape[i];
temp /= shape[i];
}
}
let mut sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
let flat_idx = indices
.iter()
.enumerate()
.map(|(i, &idx)| idx * strides[i])
.sum::<usize>();
sum = sum + data[flat_idx];
result_data[flat_idx] = sum;
}
}
Ok(Array::from_vec(result_data).reshape(&shape))
}
}
}
}
pub fn cumprod<T>(
array: &Array<T>,
axis: Option<isize>,
_out: Option<&mut Array<T>>,
) -> Result<Array<T>>
where
T: Float + Clone + Mul<Output = T> + Send + Sync,
{
use scirs2_core::parallel_ops::*;
const PARALLEL_THRESHOLD: usize = 1000;
if array.is_empty() {
return Ok(array.clone());
}
match axis {
None => {
let data = array.to_vec();
let mut result = Vec::with_capacity(data.len());
let mut prod = T::one();
for val in data {
prod = prod * val;
result.push(prod);
}
Ok(Array::from_vec(result))
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let data = array.to_vec();
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
let axis_size = shape[axis];
let n_iterations = array.size() / axis_size;
if is_parallel_enabled() && n_iterations >= PARALLEL_THRESHOLD {
let result_data: Vec<T> = (0..n_iterations)
.into_par_iter()
.flat_map(|iter_idx| {
let mut indices = vec![0; shape.len()];
let mut temp = iter_idx;
for i in (0..shape.len()).rev() {
if i != axis {
indices[i] = temp % shape[i];
temp /= shape[i];
}
}
let mut local_results = Vec::with_capacity(axis_size);
let mut prod = T::one();
for j in 0..axis_size {
indices[axis] = j;
let flat_idx = indices
.iter()
.enumerate()
.map(|(i, &idx)| idx * strides[i])
.sum::<usize>();
prod = prod * data[flat_idx];
local_results.push((flat_idx, prod));
}
local_results
})
.collect::<Vec<(usize, T)>>()
.into_iter()
.fold(vec![T::zero(); array.size()], |mut acc, (idx, val)| {
acc[idx] = val;
acc
});
Ok(Array::from_vec(result_data).reshape(&shape))
} else {
let mut result_data = vec![T::zero(); array.size()];
for iter_idx in 0..n_iterations {
let mut indices = vec![0; array.ndim()];
let mut temp = iter_idx;
for i in (0..array.ndim()).rev() {
if i != axis {
indices[i] = temp % shape[i];
temp /= shape[i];
}
}
let mut prod = T::one();
for j in 0..axis_size {
indices[axis] = j;
let flat_idx = indices
.iter()
.enumerate()
.map(|(i, &idx)| idx * strides[i])
.sum::<usize>();
prod = prod * data[flat_idx];
result_data[flat_idx] = prod;
}
}
Ok(Array::from_vec(result_data).reshape(&shape))
}
}
}
}
pub fn mean<T>(array: &Array<T>, axis: Option<isize>, keepdims: bool) -> Result<Array<T>>
where
T: Float + Clone + Add<Output = T> + Div<Output = T> + NumCast,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot compute mean of empty array".to_string(),
));
}
match axis {
None => {
let data = array.to_vec();
let sum = data.iter().fold(T::zero(), |acc, x| acc + *x);
let count = T::from(data.len()).expect("data length should be representable");
let mean_val = sum / count;
if keepdims {
let shape = vec![1; array.ndim()];
Ok(Array::from_vec(vec![mean_val]).reshape(&shape))
} else {
Ok(Array::from_vec(vec![mean_val]))
}
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
let axis_size_t = T::from(axis_size).expect("axis_size should be representable");
let mut out_shape = shape.clone();
if keepdims {
out_shape[axis] = 1;
} else {
out_shape.remove(axis);
}
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![T::zero(); out_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
for out_idx in 0..out_size {
let mut indices = vec![0; array.ndim()];
let mut temp = out_idx;
for i in 0..array.ndim() {
if i < axis {
let dim_size = shape[i];
indices[i] = temp % dim_size;
temp /= dim_size;
} else if i > axis || (i == axis && keepdims) {
let dim_idx = if keepdims { i } else { i - 1 };
if dim_idx < out_shape.len() {
let dim_size = out_shape[dim_idx];
indices[i] = temp % dim_size;
temp /= dim_size;
}
}
}
let mut sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
sum = sum + array.get(&indices)?;
}
result_data[out_idx] = sum / axis_size_t;
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn prod<T>(
array: &Array<T>,
axis: Option<isize>,
keepdims: bool,
initial: Option<T>,
) -> Result<Array<T>>
where
T: Float + Clone + Mul<Output = T>,
{
if array.is_empty() && initial.is_none() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot compute product of empty array without initial value".to_string(),
));
}
let init_val = initial.unwrap_or_else(|| T::one());
match axis {
None => {
let data = array.to_vec();
let product = data.iter().fold(init_val, |acc, x| acc * *x);
if keepdims {
let shape = vec![1; array.ndim()];
Ok(Array::from_vec(vec![product]).reshape(&shape))
} else {
Ok(Array::from_vec(vec![product]))
}
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
let mut out_shape = shape.clone();
if keepdims {
out_shape[axis] = 1;
} else {
out_shape.remove(axis);
}
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![T::zero(); out_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
for out_idx in 0..out_size {
let mut indices = vec![0; array.ndim()];
let mut temp = out_idx;
for i in 0..array.ndim() {
if i < axis {
let dim_size = shape[i];
indices[i] = temp % dim_size;
temp /= dim_size;
} else if i > axis || (i == axis && keepdims) {
let dim_idx = if keepdims { i } else { i - 1 };
if dim_idx < out_shape.len() {
let dim_size = out_shape[dim_idx];
indices[i] = temp % dim_size;
temp /= dim_size;
}
}
}
let mut product = init_val;
for j in 0..axis_size {
indices[axis] = j;
product = product * array.get(&indices)?;
}
result_data[out_idx] = product;
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn resize<T>(array: &Array<T>, new_shape: &[usize]) -> Result<Array<T>>
where
T: Clone + Zero,
{
let new_size: usize = new_shape.iter().product();
if new_size == 0 || array.is_empty() {
return Ok(Array::zeros(new_shape));
}
let old_data = array.to_vec();
let old_size = old_data.len();
let mut new_data = Vec::with_capacity(new_size);
for i in 0..new_size {
new_data.push(old_data[i % old_size].clone());
}
Ok(Array::from_vec(new_data).reshape(new_shape))
}
pub fn std<T>(
array: &Array<T>,
axis: Option<isize>,
ddof: usize,
keepdims: bool,
) -> Result<Array<T>>
where
T: Float
+ Clone
+ Add<Output = T>
+ Sub<Output = T>
+ Mul<Output = T>
+ Div<Output = T>
+ NumCast,
{
let variance = var(array, axis, ddof, keepdims)?;
Ok(variance.map(|x| x.sqrt()))
}
pub fn var<T>(
array: &Array<T>,
axis: Option<isize>,
ddof: usize,
keepdims: bool,
) -> Result<Array<T>>
where
T: Float
+ Clone
+ Add<Output = T>
+ Sub<Output = T>
+ Mul<Output = T>
+ Div<Output = T>
+ NumCast,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot compute variance of empty array".to_string(),
));
}
let mean_arr = mean(array, axis, keepdims)?;
match axis {
None => {
let data = array.to_vec();
let mean_val = mean_arr.to_vec()[0];
let sum_sq = data
.iter()
.map(|x| {
let diff = *x - mean_val;
diff * diff
})
.fold(T::zero(), |acc, x| acc + x);
let n = data.len();
if n <= ddof {
return Err(crate::error::NumRs2Error::InvalidOperation(format!(
"Degrees of freedom {} >= number of elements {}",
ddof, n
)));
}
let divisor = T::from(n - ddof).expect("n-ddof should be representable");
let var_val = sum_sq / divisor;
if keepdims {
let shape = vec![1; array.ndim()];
Ok(Array::from_vec(vec![var_val]).reshape(&shape))
} else {
Ok(Array::from_vec(vec![var_val]))
}
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
if axis_size <= ddof {
return Err(crate::error::NumRs2Error::InvalidOperation(format!(
"Degrees of freedom {} >= axis size {}",
ddof, axis_size
)));
}
let divisor =
T::from(axis_size - ddof).expect("axis_size-ddof should be representable");
let mut out_shape = shape.clone();
if keepdims {
out_shape[axis] = 1;
} else {
out_shape.remove(axis);
}
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![T::zero(); out_size];
let mut strides = vec![1; array.ndim()];
for i in (0..array.ndim() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
let mean_vec = mean_arr.to_vec();
for out_idx in 0..out_size {
let mut indices = vec![0; array.ndim()];
let mut temp = out_idx;
for i in 0..array.ndim() {
if i < axis {
let dim_size = shape[i];
indices[i] = temp % dim_size;
temp /= dim_size;
} else if i > axis || (i == axis && keepdims) {
let dim_idx = if keepdims { i } else { i - 1 };
if dim_idx < out_shape.len() {
let dim_size = out_shape[dim_idx];
indices[i] = temp % dim_size;
temp /= dim_size;
}
}
}
let mean_val = mean_vec[out_idx];
let mut sum_sq = T::zero();
for j in 0..axis_size {
indices[axis] = j;
let val = array.get(&indices)?;
let diff = val - mean_val;
sum_sq = sum_sq + (diff * diff);
}
result_data[out_idx] = sum_sq / divisor;
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn clip<T>(array: &Array<T>, min: Option<T>, max: Option<T>) -> Result<Array<T>>
where
T: PartialOrd + Clone,
{
Ok(array.map(|x| {
let mut val = x;
if let Some(ref min_val) = min {
if &val < min_val {
val = min_val.clone();
}
}
if let Some(ref max_val) = max {
if &val > max_val {
val = max_val.clone();
}
}
val
}))
}
pub fn skew<T>(array: &Array<T>, axis: Option<isize>) -> Result<Array<T>>
where
T: Float
+ Clone
+ Add<Output = T>
+ Sub<Output = T>
+ Mul<Output = T>
+ Div<Output = T>
+ NumCast,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot compute skewness of empty array".to_string(),
));
}
match axis {
None => {
let data = array.to_vec();
let n = data.len();
if n < 2 {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Skewness requires at least 2 elements".to_string(),
));
}
let n_t = T::from(n).ok_or_else(|| {
crate::error::NumRs2Error::ConversionError("Cannot convert n to T".to_string())
})?;
let mean_val = data.iter().fold(T::zero(), |acc, &x| acc + x) / n_t;
let variance = data.iter().fold(T::zero(), |acc, &x| {
let d = x - mean_val;
acc + d * d
}) / n_t;
let std_val = variance.sqrt();
if std_val == T::zero() {
return Ok(Array::from_vec(vec![T::zero()]));
}
let third_moment = data.iter().fold(T::zero(), |acc, &x| {
let d = x - mean_val;
acc + d * d * d
}) / n_t;
let skew_val = third_moment / (std_val * std_val * std_val);
Ok(Array::from_vec(vec![skew_val]))
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
if axis_size < 2 {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Skewness requires at least 2 elements along the axis".to_string(),
));
}
let axis_size_t = T::from(axis_size).ok_or_else(|| {
crate::error::NumRs2Error::ConversionError(
"Cannot convert axis_size to T".to_string(),
)
})?;
let mut out_shape = shape.clone();
out_shape.remove(axis);
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![T::zero(); out_size];
for out_idx in 0..out_size {
let mut indices = vec![0usize; array.ndim()];
let mut temp = out_idx;
for i in (0..array.ndim()).rev() {
if i == axis {
continue;
}
let dim_idx = if i < axis { i } else { i - 1 };
if dim_idx < out_shape.len() {
indices[i] = temp % out_shape[dim_idx];
temp /= out_shape[dim_idx];
}
}
let mut sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
sum = sum + array.get(&indices)?;
}
let mean_val = sum / axis_size_t;
let mut var_sum = T::zero();
let mut third_sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
let d = array.get(&indices)? - mean_val;
var_sum = var_sum + d * d;
third_sum = third_sum + d * d * d;
}
let variance = var_sum / axis_size_t;
let std_val = variance.sqrt();
result_data[out_idx] = if std_val == T::zero() {
T::zero()
} else {
(third_sum / axis_size_t) / (std_val * std_val * std_val)
};
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
pub fn kurtosis<T>(array: &Array<T>, axis: Option<isize>) -> Result<Array<T>>
where
T: Float
+ Clone
+ Add<Output = T>
+ Sub<Output = T>
+ Mul<Output = T>
+ Div<Output = T>
+ NumCast,
{
if array.is_empty() {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Cannot compute kurtosis of empty array".to_string(),
));
}
let three = T::from(3u8).ok_or_else(|| {
crate::error::NumRs2Error::ConversionError("Cannot convert 3 to T".to_string())
})?;
match axis {
None => {
let data = array.to_vec();
let n = data.len();
if n < 2 {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Kurtosis requires at least 2 elements".to_string(),
));
}
let n_t = T::from(n).ok_or_else(|| {
crate::error::NumRs2Error::ConversionError("Cannot convert n to T".to_string())
})?;
let mean_val = data.iter().fold(T::zero(), |acc, &x| acc + x) / n_t;
let variance = data.iter().fold(T::zero(), |acc, &x| {
let d = x - mean_val;
acc + d * d
}) / n_t;
let std_val = variance.sqrt();
if std_val == T::zero() {
return Ok(Array::from_vec(vec![T::zero()]));
}
let fourth_moment = data.iter().fold(T::zero(), |acc, &x| {
let d = x - mean_val;
let d2 = d * d;
acc + d2 * d2
}) / n_t;
let sigma4 = variance * variance;
let kurt_val = fourth_moment / sigma4 - three;
Ok(Array::from_vec(vec![kurt_val]))
}
Some(ax) => {
let axis = if ax < 0 {
(array.ndim() as isize + ax) as usize
} else {
ax as usize
};
if axis >= array.ndim() {
return Err(crate::error::NumRs2Error::DimensionMismatch(format!(
"Axis {} out of bounds for array of dimension {}",
axis,
array.ndim()
)));
}
let shape = array.shape();
let axis_size = shape[axis];
if axis_size < 2 {
return Err(crate::error::NumRs2Error::InvalidOperation(
"Kurtosis requires at least 2 elements along the axis".to_string(),
));
}
let axis_size_t = T::from(axis_size).ok_or_else(|| {
crate::error::NumRs2Error::ConversionError(
"Cannot convert axis_size to T".to_string(),
)
})?;
let mut out_shape = shape.clone();
out_shape.remove(axis);
if out_shape.is_empty() {
out_shape.push(1);
}
let out_size: usize = out_shape.iter().product();
let mut result_data = vec![T::zero(); out_size];
for out_idx in 0..out_size {
let mut indices = vec![0usize; array.ndim()];
let mut temp = out_idx;
for i in (0..array.ndim()).rev() {
if i == axis {
continue;
}
let dim_idx = if i < axis { i } else { i - 1 };
if dim_idx < out_shape.len() {
indices[i] = temp % out_shape[dim_idx];
temp /= out_shape[dim_idx];
}
}
let mut sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
sum = sum + array.get(&indices)?;
}
let mean_val = sum / axis_size_t;
let mut var_sum = T::zero();
let mut fourth_sum = T::zero();
for j in 0..axis_size {
indices[axis] = j;
let d = array.get(&indices)? - mean_val;
let d2 = d * d;
var_sum = var_sum + d2;
fourth_sum = fourth_sum + d2 * d2;
}
let variance = var_sum / axis_size_t;
let sigma4 = variance * variance;
result_data[out_idx] = if sigma4 == T::zero() {
T::zero()
} else {
fourth_sum / axis_size_t / sigma4 - three
};
}
Ok(Array::from_vec(result_data).reshape(&out_shape))
}
}
}
#[cfg(test)]
mod skew_kurtosis_tests {
use super::*;
use crate::array::Array;
#[test]
fn test_skew_symmetric_data() {
let data = Array::from_vec(vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]);
let s = skew(&data, None).expect("skew should succeed");
assert!(
s.to_vec()[0].abs() < 1e-10,
"symmetric distribution skewness should be ~0, got {}",
s.to_vec()[0]
);
}
#[test]
fn test_skew_right_skewed() {
let data = Array::from_vec(vec![1.0_f64, 2.0, 3.0, 4.0, 10.0]);
let s = skew(&data, None).expect("skew should succeed");
assert!(
s.to_vec()[0] > 0.0,
"right-skewed data should have positive skewness"
);
}
#[test]
fn test_skew_constant_array_returns_zero() {
let data = Array::from_vec(vec![5.0_f64, 5.0, 5.0, 5.0]);
let s = skew(&data, None).expect("constant array skew should succeed");
assert_eq!(s.to_vec()[0], 0.0);
}
#[test]
fn test_skew_too_few_elements_errors() {
let data = Array::from_vec(vec![42.0_f64]);
assert!(
skew(&data, None).is_err(),
"skew with 1 element should return error"
);
}
#[test]
fn test_kurtosis_uniform_like_data() {
let data = Array::from_vec(vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]);
let k = kurtosis(&data, None).expect("kurtosis should succeed");
assert!(
k.to_vec()[0] < 0.0,
"uniform-like data should have negative excess kurtosis, got {}",
k.to_vec()[0]
);
}
#[test]
fn test_kurtosis_constant_array_returns_zero() {
let data = Array::from_vec(vec![3.0_f64, 3.0, 3.0, 3.0]);
let k = kurtosis(&data, None).expect("constant array kurtosis should succeed");
assert_eq!(k.to_vec()[0], 0.0);
}
#[test]
fn test_kurtosis_too_few_elements_errors() {
let data = Array::from_vec(vec![99.0_f64]);
assert!(
kurtosis(&data, None).is_err(),
"kurtosis with 1 element should return error"
);
}
#[test]
fn test_skew_empty_array_errors() {
let data: Array<f64> = Array::from_vec(vec![]);
assert!(
skew(&data, None).is_err(),
"skew on empty array should fail"
);
}
#[test]
fn test_kurtosis_empty_array_errors() {
let data: Array<f64> = Array::from_vec(vec![]);
assert!(
kurtosis(&data, None).is_err(),
"kurtosis on empty array should fail"
);
}
}