use std::fmt;
use std::ops::{Index, IndexMut, Range, RangeFrom, RangeFull, RangeTo};
use std::sync::Arc;
use crate::enums::error::MinarrowError;
use crate::enums::shape_dim::ShapeDim;
use crate::structs::buffer::Buffer;
#[cfg(all(feature = "views", feature = "select"))]
use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
use crate::traits::type_unions::Float;
use crate::traits::{concatenate::Concatenate, shape::Shape};
use crate::{Array, ArrowType, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
#[cfg(feature = "matrix")]
use crate::structs::matrix::{Matrix, aligned_stride};
#[cfg(feature = "views")]
use crate::structs::views::ndarray_view::NdArrayV;
#[cfg(feature = "dlpack")]
use crate::ffi::dlpack::{
export_to_dlpack, export_to_dlpack_versioned, DLPackTensor, DLPackTensorVersioned,
};
#[derive(Clone)]
pub struct NdArray<T> {
pub(crate) data: Arc<Buffer<T>>,
pub(crate) dims: NdDims,
pub name: Option<String>,
}
impl<T: Float> PartialEq for NdArray<T> {
fn eq(&self, other: &Self) -> bool {
if self.dims.shape() != other.dims.shape() {
return false;
}
self.into_iter()
.zip(other.into_iter())
.all(|(a, b)| a == b)
}
}
impl<T: Float> NdArray<T> {
pub fn new(shape: &[usize]) -> Self {
let dims = NdDims::from_shape(shape);
let total = buffer_len(dims.shape(), dims.strides());
let mut v = Vec64::with_capacity(total);
v.0.resize(total, T::default());
NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
}
pub fn new_named(shape: &[usize], name: impl Into<String>) -> Self {
let mut arr = Self::new(shape);
arr.name = Some(name.into());
arr
}
pub fn from_slice(data: &[T], shape: &[usize]) -> Self {
let logical_len: usize = shape.iter().product();
assert_eq!(
data.len(), logical_len,
"NdArray::from_slice: data length {} does not match shape product {}",
data.len(), logical_len
);
let dims = NdDims::from_shape(shape);
NdArray {
data: Arc::new(Buffer::from_slice(data)),
dims,
name: None,
}
}
pub fn from_vec64(data: Vec64<T>, shape: &[usize]) -> Self {
let logical_len: usize = shape.iter().product();
assert_eq!(
data.len(), logical_len,
"NdArray::from_vec64: data length {} does not match shape product {}",
data.len(), logical_len
);
let dims = NdDims::from_shape(shape);
NdArray {
data: Arc::new(Buffer::from_vec64(data)),
dims,
name: None,
}
}
pub fn from_buffer(data: Buffer<T>, shape: &[usize], strides: &[usize]) -> Self {
let required = buffer_len(shape, strides);
assert!(
data.len() >= required,
"NdArray::from_buffer: buffer has {} elements but shape requires {}",
data.len(), required
);
let dims = NdDims::from_shape_and_strides(shape, strides);
NdArray { data: Arc::new(data), dims, name: None }
}
pub fn fill(shape: &[usize], value: T) -> Self {
let dims = NdDims::from_shape(shape);
let total = buffer_len(dims.shape(), dims.strides());
let mut v = Vec64::with_capacity(total);
v.0.resize(total, value);
NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
}
pub fn ones(shape: &[usize]) -> Self {
Self::fill(shape, T::one())
}
pub fn eye(n: usize) -> Self {
let mut arr = Self::new(&[n, n]);
let stride = arr.dims.strides()[1];
let buf = Arc::make_mut(&mut arr.data).as_mut_slice();
for i in 0..n {
buf[i * stride + i] = T::one();
}
arr
}
pub fn linspace(start: T, end: T, n: usize) -> Self {
assert!(n >= 2, "linspace requires at least 2 points");
let step = (end - start) / T::from(n - 1).unwrap();
let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
NdArray {
data: Arc::new(Buffer::from_vec64(v)),
dims: NdDims::from_shape(&[n]),
name: None,
}
}
pub fn arange(start: T, step: T, n: usize) -> Self {
let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
NdArray {
data: Arc::new(Buffer::from_vec64(v)),
dims: NdDims::from_shape(&[n]),
name: None,
}
}
#[inline]
pub fn ndim(&self) -> usize { self.dims.ndim() }
#[inline]
pub fn shape(&self) -> &[usize] { self.dims.shape() }
#[inline]
pub fn strides(&self) -> &[usize] { self.dims.strides() }
#[inline]
pub fn len(&self) -> usize { self.dims.len() }
#[inline]
pub fn n_obs(&self) -> usize {
assert!(self.ndim() > 0, "n_obs() requires an axis 0");
self.dims.shape()[0]
}
#[inline]
pub fn is_empty(&self) -> bool { self.len() == 0 }
pub fn is_contiguous(&self) -> bool {
let shape = self.dims.shape();
let strides = self.dims.strides();
let mut expected = 1;
for d in 0..shape.len() {
if strides[d] != expected {
return false;
}
expected *= shape[d];
}
true
}
pub fn has_nan(&self) -> bool {
self.into_iter().any(|v| v.is_nan())
}
pub fn nan_count(&self) -> usize {
self.into_iter().filter(|v| v.is_nan()).count()
}
#[inline]
pub fn get(&self, indices: &[usize]) -> T {
self.data.as_slice()[self.offset_of(indices)]
}
#[inline(always)]
pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
let strides = self.dims.strides();
let mut off = 0;
for d in 0..indices.len() {
off += indices[d] * strides[d];
}
unsafe { *self.data.as_slice().get_unchecked(off) }
}
#[inline]
pub fn set(&mut self, indices: &[usize], value: T) {
let off = self.offset_of(indices);
Arc::make_mut(&mut self.data).as_mut_slice()[off] = value;
}
#[inline]
pub(crate) fn offset_of(&self, indices: &[usize]) -> usize {
offset_of_impl(indices, self.dims.shape(), self.dims.strides())
}
#[inline]
pub fn set_name(&mut self, name: impl Into<String>) {
self.name = Some(name.into());
}
#[inline]
pub fn as_slice(&self) -> &[T] {
self.data.as_slice()
}
#[inline]
pub fn as_mut_slice(&mut self) -> &mut [T] {
Arc::make_mut(&mut self.data).as_mut_slice()
}
pub fn fill_with(&mut self, value: T) {
if self.is_contiguous() {
Arc::make_mut(&mut self.data).as_mut_slice().fill(value);
return;
}
let offsets: Vec<usize> = {
let shape = self.dims.shape();
let strides = self.dims.strides();
let mut result = Vec::with_capacity(self.len());
let mut indices = vec![0usize; shape.len()];
for _ in 0..self.len() {
result.push(indices.iter().zip(strides).map(|(&i, &s)| i * s).sum());
let mut carry = true;
for d in 0..shape.len() {
if carry {
indices[d] += 1;
if indices[d] < shape[d] { carry = false; } else { indices[d] = 0; }
}
}
}
result
};
let buf = Arc::make_mut(&mut self.data).as_mut_slice();
for off in offsets { buf[off] = value; }
}
#[inline]
pub fn col(&self, col: usize) -> &[T] {
let shape = self.dims.shape();
assert_eq!(shape.len(), 2, "col() requires a 2D array");
assert!(col < shape[1], "Column index out of bounds");
assert_eq!(
self.dims.strides()[0], 1,
"col() requires unit stride on axis 0; call to_contiguous() first"
);
let stride = self.dims.strides()[1];
let start = col * stride;
&self.data.as_slice()[start..start + shape[0]]
}
#[inline]
pub fn col_mut(&mut self, col: usize) -> &mut [T] {
let shape = self.dims.shape();
assert_eq!(shape.len(), 2, "col_mut() requires a 2D array");
assert!(col < shape[1], "Column index out of bounds");
assert_eq!(
self.dims.strides()[0], 1,
"col_mut() requires unit stride on axis 0; call to_contiguous() first"
);
let stride = self.dims.strides()[1];
let n_rows = shape[0];
let start = col * stride;
&mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
}
pub fn columns(&self) -> Vec<&[T]> {
let shape = self.dims.shape();
assert_eq!(shape.len(), 2, "columns() requires a 2D array");
assert_eq!(
self.dims.strides()[0], 1,
"columns() requires unit stride on axis 0; call to_contiguous() first"
);
let stride = self.dims.strides()[1];
let n_rows = shape[0];
let buf = self.data.as_slice();
(0..shape[1])
.map(|c| &buf[c * stride..c * stride + n_rows])
.collect()
}
pub fn columns_mut(&mut self) -> Vec<&mut [T]> {
let shape = self.dims.shape();
assert_eq!(shape.len(), 2, "columns_mut() requires a 2D array");
assert_eq!(
self.dims.strides()[0], 1,
"columns_mut() requires unit stride on axis 0; call to_contiguous() first"
);
let n_rows = shape[0];
let n_cols = shape[1];
let stride = self.dims.strides()[1];
let ptr = Arc::make_mut(&mut self.data).as_mut_slice().as_mut_ptr();
let mut result = Vec::with_capacity(n_cols);
for c in 0..n_cols {
let start = c * stride;
unsafe {
let col_ptr = ptr.add(start);
result.push(std::slice::from_raw_parts_mut(col_ptr, n_rows));
}
}
result
}
#[cfg(feature = "views")]
pub fn as_view(&self) -> NdArrayV<T> {
NdArrayV::from_ndarray(self.clone())
}
#[cfg(feature = "views")]
pub fn obs(&self, idx: usize) -> NdArrayV<T> {
self.as_view().obs(idx)
}
#[inline]
pub fn m(&self) -> i32 {
assert_eq!(self.ndim(), 2, "m() requires a 2D array");
self.dims.shape()[0] as i32
}
#[inline]
pub fn n(&self) -> i32 {
assert_eq!(self.ndim(), 2, "n() requires a 2D array");
self.dims.shape()[1] as i32
}
#[inline]
pub fn lda(&self) -> i32 {
assert_eq!(self.ndim(), 2, "lda() requires a 2D array");
self.dims.strides()[1] as i32
}
pub fn reshape(&self, new_shape: &[usize]) -> Result<NdArray<T>, MinarrowError> {
let new_len: usize = new_shape.iter().product();
if new_len != self.len() {
return Err(MinarrowError::ShapeError {
message: format!(
"reshape: cannot reshape array of size {} into shape {:?}",
self.len(), new_shape
),
});
}
let flat: Vec64<T> = self.into_iter().collect();
let mut result = NdArray::from_slice(&flat, new_shape);
result.name = self.name.clone();
Ok(result)
}
pub fn transpose(&self) -> NdArray<T> {
let shape = self.dims.shape();
let ndim = shape.len();
if ndim <= 1 {
let mut result = self.to_contiguous();
result.name = self.name.clone();
return result;
}
if ndim == 2 {
let (n_rows, n_cols) = (shape[0], shape[1]);
let new_shape = [n_cols, n_rows];
let mut result = NdArray::new(&new_shape);
let src_stride = self.dims.strides()[1];
let dst_stride = result.dims.strides()[1];
let src = self.data.as_slice();
let dst = Arc::make_mut(&mut result.data).as_mut_slice();
for r in 0..n_rows {
for c in 0..n_cols {
dst[r * dst_stride + c] = src[c * src_stride + r];
}
}
result.name = self.name.clone();
return result;
}
let new_shape: Vec<usize> = shape.iter().rev().copied().collect();
let rev_strides: Vec<usize> = self.dims.strides().iter().rev().copied().collect();
let src = self.data.as_slice();
let total = self.len();
let mut buf = Vec64::with_capacity(total);
let mut indices = vec![0usize; ndim];
for _ in 0..total {
let offset: usize = indices.iter()
.zip(rev_strides.iter())
.map(|(&i, &s)| i * s)
.sum();
buf.push(src[offset]);
let mut carry = true;
for d in 0..ndim {
if carry {
indices[d] += 1;
if indices[d] < new_shape[d] {
carry = false;
} else {
indices[d] = 0;
}
}
}
}
NdArray {
data: Arc::new(Buffer::from_vec64(buf)),
dims: NdDims::from_shape(&new_shape),
name: self.name.clone(),
}
}
pub fn flatten(&self) -> NdArray<T> {
let flat: Vec64<T> = self.into_iter().collect();
let n = flat.len();
NdArray {
data: Arc::new(Buffer::from_vec64(flat)),
dims: NdDims::from_shape(&[n]),
name: self.name.clone(),
}
}
pub fn to_contiguous(&self) -> NdArray<T> {
if self.is_contiguous() {
return self.clone();
}
let mut result = NdArray::from_slice(
&self.into_iter().collect::<Vec64<T>>(),
self.shape(),
);
result.name = self.name.clone();
result
}
#[cfg(all(feature = "views", feature = "select"))]
pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
assert_eq!(
selection.len(), self.ndim(),
"slice(): expected {} axes, got {}", self.ndim(), selection.len()
);
let shape = self.dims.shape();
let strides = self.dims.strides();
let mut new_offset: usize = 0;
let mut new_shape = Vec::with_capacity(self.ndim());
let mut new_strides = Vec::with_capacity(self.ndim());
for (d, sel) in selection.iter().enumerate() {
let (start, end, collapse) = sel.resolve_axis(shape[d]);
assert!(
end <= shape[d],
"slice(): end {} out of bounds for axis {} (size {})", end, d, shape[d]
);
new_offset += start * strides[d];
if !collapse {
new_shape.push(end - start);
new_strides.push(strides[d]);
}
}
NdArrayV::new(self.clone(), new_offset, &new_shape, &new_strides)
}
pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray<T> {
let flat: Vec64<T> = self.into_iter().map(f).collect();
let mut result = NdArray::from_slice(&flat, self.shape());
result.name = self.name.clone();
result
}
pub fn apply_mut(&mut self, f: impl Fn(T) -> T) {
if self.is_contiguous() {
for v in Arc::make_mut(&mut self.data).as_mut_slice() {
*v = f(*v);
}
return;
}
let shape = self.dims.shape().to_vec();
let strides = self.dims.strides().to_vec();
let total = self.len();
let buf = Arc::make_mut(&mut self.data).as_mut_slice();
let mut indices = vec![0usize; shape.len()];
for _ in 0..total {
let offset: usize = indices.iter().zip(strides.iter()).map(|(&i, &s)| i * s).sum();
buf[offset] = f(buf[offset]);
let mut carry = true;
for d in 0..shape.len() {
if carry {
indices[d] += 1;
if indices[d] < shape[d] {
carry = false;
} else {
indices[d] = 0;
}
}
}
}
}
#[cfg(feature = "views")]
pub fn apply_axis(&self, axis: usize, mut f: impl FnMut(NdArrayV<T>) -> T) -> NdArray<T> {
let shape = self.dims.shape();
let strides = self.dims.strides();
let ndim = shape.len();
assert!(ndim >= 2, "apply_axis requires a 2D or higher array");
assert!(axis < ndim, "apply_axis: axis {} out of bounds for {}D array", axis, ndim);
let out_shape: Vec<usize> = shape
.iter()
.enumerate()
.filter(|(d, _)| *d != axis)
.map(|(_, &s)| s)
.collect();
let out_dims: Vec<usize> = (0..ndim).filter(|&d| d != axis).collect();
let total: usize = out_shape.iter().product();
let lane_shape = [shape[axis]];
let lane_strides = [strides[axis]];
let mut flat = Vec64::with_capacity(total);
let mut indices = vec![0usize; out_shape.len()];
for _ in 0..total {
let offset: usize = indices
.iter()
.zip(out_dims.iter())
.map(|(&i, &d)| i * strides[d])
.sum();
let lane = NdArrayV::new(self.clone(), offset, &lane_shape, &lane_strides);
flat.push(f(lane));
let mut carry = true;
for d in 0..out_shape.len() {
if carry {
indices[d] += 1;
if indices[d] < out_shape[d] {
carry = false;
} else {
indices[d] = 0;
}
}
}
}
let mut result = NdArray::from_slice(&flat, &out_shape);
result.name = self.name.clone();
result
}
#[cfg(feature = "dlpack")]
pub fn to_dlpack(self) -> DLPackTensor {
export_to_dlpack(self)
}
#[cfg(feature = "dlpack")]
pub fn to_dlpack_versioned(self) -> DLPackTensorVersioned {
export_to_dlpack_versioned(self)
}
#[cfg(feature = "parallel_proc")]
pub fn par_iter(&self) -> rayon::slice::Iter<'_, T>
where
T: Send + Sync,
{
use rayon::prelude::*;
self.data.as_slice().par_iter()
}
#[cfg(feature = "parallel_proc")]
pub fn par_chunks(&self, chunk_size: usize) -> rayon::slice::Chunks<'_, T>
where
T: Send + Sync,
{
use rayon::prelude::*;
self.data.as_slice().par_chunks(chunk_size)
}
#[cfg(feature = "views")]
pub fn iter_obs(&self) -> impl Iterator<Item = (usize, NdArrayV<T>)> + '_ {
assert!(self.ndim() >= 2, "iter_obs() requires a 2D or higher array");
let n_obs = self.shape()[0];
(0..n_obs).map(move |i| (i, self.obs(i)))
}
#[cfg(all(feature = "parallel_proc", feature = "views"))]
pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator<Item = (usize, NdArrayV<T>)> + '_
where
T: Send + Sync,
{
use rayon::prelude::*;
assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher array");
let n_obs = self.dims.shape()[0];
(0..n_obs).into_par_iter().map(move |i| (i, self.obs(i)))
}
pub(crate) fn iter_axis0_run(&self, run_idx: usize) -> impl ExactSizeIterator<Item = T> + '_ {
assert!(self.ndim() > 0, "axis-0 iteration requires an axis 0");
let n_runs: usize = self.shape()[1..].iter().product();
assert!(run_idx < n_runs, "axis-0 run {} out of bounds ({})", run_idx, n_runs);
let mut rem = run_idx;
let mut offset = 0;
for d in 1..self.ndim() {
offset += (rem % self.shape()[d]) * self.strides()[d];
rem /= self.shape()[d];
}
let stride = self.strides()[0];
(0..self.shape()[0]).map(move |i| self.data.as_slice()[offset + i * stride])
}
}
impl NdArray<f64> {
pub fn to_table(self, fields: Option<Vec<Field>>) -> Result<Table, MinarrowError> {
let shape = self.dims.shape();
if shape.len() != 2 {
return Err(MinarrowError::ShapeError {
message: format!("to_table requires a 2D array, got {}D", shape.len()),
});
}
let n_cols = shape[1];
let fields = match fields {
Some(fields) => {
if fields.len() != n_cols {
return Err(MinarrowError::ShapeError {
message: format!(
"to_table: expected {} fields for {} columns, got {}",
n_cols, n_cols, fields.len()
),
});
}
for (index, field) in fields.iter().enumerate() {
if field.dtype != ArrowType::Float64 {
return Err(MinarrowError::TypeError {
from: "Field",
to: "Float64 NdArray column",
message: Some(format!(
"to_table: field {} ('{}') has dtype {:?}, expected Float64",
index, field.name, field.dtype
)),
});
}
if field.nullable {
return Err(MinarrowError::NullError {
message: Some(format!(
"to_table: field {} ('{}') is nullable, but NdArray columns are non-nullable",
index, field.name
)),
});
}
}
fields
}
None => (0..n_cols)
.map(|i| Field::new(format!("col_{}", i), ArrowType::Float64, false, None))
.collect(),
};
let this = if self.dims.strides()[0] == 1 { self } else { self.to_contiguous() };
let n_rows = this.dims.shape()[0];
let stride = this.dims.strides()[1];
let name = this.name;
let buf = this.data.as_slice();
let mut cols = Vec::with_capacity(n_cols);
for (i, field) in fields.into_iter().enumerate() {
let col_start = i * stride;
let col: Buffer<f64> = Buffer::from_slice(&buf[col_start..col_start + n_rows]);
let float_arr = FloatArray::new(col, None);
let array = Array::NumericArray(NumericArray::Float64(Arc::new(float_arr)));
cols.push(FieldArray::new(field, array));
}
Ok(Table::new(name.unwrap_or_default(), Some(cols)))
}
pub fn to_array(self) -> Result<Array, MinarrowError> {
if self.ndim() != 1 {
return Err(MinarrowError::ShapeError {
message: format!("to_array requires a 1D array, got {}D", self.ndim()),
});
}
let buffer = if self.dims.strides()[0] == 1 && self.data.len() == self.len() {
Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone())
} else {
Buffer::from_vec64((&self).into_iter().collect())
};
let float_arr = FloatArray::new(buffer, None);
Ok(Array::NumericArray(NumericArray::Float64(Arc::new(float_arr))))
}
#[cfg(feature = "matrix")]
pub fn to_matrix(self) -> Result<Matrix, MinarrowError> {
let shape = self.dims.shape();
if shape.len() != 2 {
return Err(MinarrowError::ShapeError {
message: format!("to_matrix requires a 2D array, got {}D", shape.len()),
});
}
let n_rows = shape[0];
let n_cols = shape[1];
let strides = self.dims.strides();
if strides[0] == 1 && strides[1] == aligned_stride(n_rows) {
return Ok(Matrix {
n_rows,
n_cols,
stride: strides[1],
data: Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone()),
name: self.name,
});
}
let name = self.name.clone();
let compact: Vec64<f64> = self.into_iter().collect();
Ok(Matrix::from_f64_unaligned(&compact, n_rows, n_cols, name))
}
}
#[derive(Clone, PartialEq)]
pub(crate) enum NdDims {
D1 { shape: [usize; 1], strides: [usize; 1] },
D2 { shape: [usize; 2], strides: [usize; 2] },
D3 { shape: [usize; 3], strides: [usize; 3] },
D4 { shape: [usize; 4], strides: [usize; 4] },
D5 { shape: [usize; 5], strides: [usize; 5] },
Dn { shape: Box<[usize]>, strides: Box<[usize]> },
}
impl NdDims {
pub(crate) fn from_shape(shape: &[usize]) -> Self {
let strides = col_major_strides(shape);
Self::from_shape_and_strides(shape, &strides)
}
pub(crate) fn from_shape_and_strides(shape: &[usize], strides: &[usize]) -> Self {
assert_eq!(
shape.len(),
strides.len(),
"NdArray: shape rank {} does not match strides rank {}",
shape.len(),
strides.len()
);
match shape.len() {
1 => NdDims::D1 {
shape: [shape[0]],
strides: [strides[0]],
},
2 => NdDims::D2 {
shape: [shape[0], shape[1]],
strides: [strides[0], strides[1]],
},
3 => NdDims::D3 {
shape: [shape[0], shape[1], shape[2]],
strides: [strides[0], strides[1], strides[2]],
},
4 => NdDims::D4 {
shape: [shape[0], shape[1], shape[2], shape[3]],
strides: [strides[0], strides[1], strides[2], strides[3]],
},
5 => NdDims::D5 {
shape: [shape[0], shape[1], shape[2], shape[3], shape[4]],
strides: [strides[0], strides[1], strides[2], strides[3], strides[4]],
},
_ => NdDims::Dn {
shape: shape.into(),
strides: strides.into(),
},
}
}
#[inline]
pub(crate) fn ndim(&self) -> usize {
match self {
NdDims::D1 { .. } => 1,
NdDims::D2 { .. } => 2,
NdDims::D3 { .. } => 3,
NdDims::D4 { .. } => 4,
NdDims::D5 { .. } => 5,
NdDims::Dn { shape, .. } => shape.len(),
}
}
#[inline]
pub(crate) fn shape(&self) -> &[usize] {
match self {
NdDims::D1 { shape, .. } => shape,
NdDims::D2 { shape, .. } => shape,
NdDims::D3 { shape, .. } => shape,
NdDims::D4 { shape, .. } => shape,
NdDims::D5 { shape, .. } => shape,
NdDims::Dn { shape, .. } => shape,
}
}
#[inline]
pub(crate) fn strides(&self) -> &[usize] {
match self {
NdDims::D1 { strides, .. } => strides,
NdDims::D2 { strides, .. } => strides,
NdDims::D3 { strides, .. } => strides,
NdDims::D4 { strides, .. } => strides,
NdDims::D5 { strides, .. } => strides,
NdDims::Dn { strides, .. } => strides,
}
}
#[inline]
pub(crate) fn len(&self) -> usize {
self.shape().iter().product()
}
}
#[macro_export]
macro_rules! nd {
($($sel:expr),+ $(,)?) => {
&[$(&$sel as &dyn $crate::traits::selection::DataSelector),+]
};
}
pub(crate) fn col_major_strides(shape: &[usize]) -> Vec<usize> {
let mut strides = Vec::with_capacity(shape.len());
if shape.is_empty() {
return strides;
}
strides.push(1);
for d in 1..shape.len() {
strides.push(strides[d - 1] * shape[d - 1]);
}
strides
}
pub(crate) fn buffer_len(shape: &[usize], strides: &[usize]) -> usize {
if shape.iter().any(|&d| d == 0) {
return 0;
}
let max_offset: usize = shape.iter()
.zip(strides.iter())
.map(|(&s, &st)| (s - 1) * st)
.sum();
max_offset + 1
}
impl<'a, T: Float> IntoIterator for &'a NdArray<T> {
type Item = T;
type IntoIter = NdArrayIter<'a, T>;
#[inline]
fn into_iter(self) -> NdArrayIter<'a, T> {
let shape = self.dims.shape();
let strides = self.dims.strides();
if shape.is_empty() {
return NdArrayIter {
buf: self.data.as_slice(),
n_inner: 1,
inner_stride: 1,
run_offsets: vec![0],
run_idx: 0,
inner_idx: 0,
total: 1,
yielded: 0,
};
}
let n_inner = shape[0];
let n_runs: usize = shape[1..].iter().product();
let mut run_offsets = Vec::with_capacity(n_runs);
if shape.len() <= 1 {
run_offsets.push(0);
} else {
let outer_shape = &shape[1..];
let outer_strides = &strides[1..];
let mut outer_indices = vec![0usize; outer_shape.len()];
for _ in 0..n_runs {
let off: usize = outer_indices.iter()
.zip(outer_strides.iter())
.map(|(&i, &s)| i * s)
.sum();
run_offsets.push(off);
let mut carry = true;
for d in 0..outer_shape.len() {
if carry {
outer_indices[d] += 1;
if outer_indices[d] < outer_shape[d] {
carry = false;
} else {
outer_indices[d] = 0;
}
}
}
}
}
NdArrayIter {
buf: self.data.as_slice(),
n_inner,
inner_stride: strides[0],
run_offsets,
run_idx: 0,
inner_idx: 0,
total: self.len(),
yielded: 0,
}
}
}
impl<T: Float> IntoIterator for NdArray<T> {
type Item = T;
type IntoIter = std::vec::IntoIter<T>;
#[inline]
fn into_iter(self) -> std::vec::IntoIter<T> {
let v: Vec<T> = (&self).into_iter().collect();
v.into_iter()
}
}
pub struct NdArrayIter<'a, T> {
pub(crate) buf: &'a [T],
pub(crate) n_inner: usize,
pub(crate) inner_stride: usize,
pub(crate) run_offsets: Vec<usize>,
pub(crate) run_idx: usize,
pub(crate) inner_idx: usize,
pub(crate) total: usize,
pub(crate) yielded: usize,
}
impl<'a, T: Float> Iterator for NdArrayIter<'a, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
if self.yielded >= self.total { return None; }
let val = self.buf[self.run_offsets[self.run_idx] + self.inner_idx * self.inner_stride];
self.yielded += 1;
self.inner_idx += 1;
if self.inner_idx >= self.n_inner {
self.inner_idx = 0;
self.run_idx += 1;
}
Some(val)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let r = self.total - self.yielded;
(r, Some(r))
}
}
impl<'a, T: Float> ExactSizeIterator for NdArrayIter<'a, T> {}
#[inline]
pub(crate) fn offset_of_impl(indices: &[usize], shape: &[usize], strides: &[usize]) -> usize {
assert_eq!(
indices.len(),
shape.len(),
"NdArray: {} indices for a {}D array",
indices.len(),
shape.len()
);
let mut offset = 0;
for d in 0..shape.len() {
assert!(
indices[d] < shape[d],
"NdArray: index {} out of bounds for dim {} (size {})",
indices[d], d, shape[d]
);
offset += indices[d] * strides[d];
}
offset
}
impl<T: Float> Shape for NdArray<T> {
fn shape(&self) -> ShapeDim {
match self.dims.ndim() {
0 => ShapeDim::Rank0(1),
1 => ShapeDim::Rank1(self.dims.shape()[0]),
2 => ShapeDim::Rank2 {
rows: self.dims.shape()[0],
cols: self.dims.shape()[1],
},
_ => ShapeDim::RankN(self.dims.shape().to_vec()),
}
}
}
impl<T: Float> Concatenate for NdArray<T> {
fn concat(self, other: Self) -> Result<Self, MinarrowError> {
let s1 = self.dims.shape();
let s2 = other.dims.shape();
if s1.len() != s2.len() {
return Err(MinarrowError::IncompatibleTypeError {
from: "NdArray",
to: "NdArray",
message: Some(format!(
"Cannot concatenate {}D and {}D arrays", s1.len(), s2.len()
)),
});
}
if s1.is_empty() {
return Err(MinarrowError::ShapeError {
message: "Cannot concatenate rank-zero arrays along axis 0".to_string(),
});
}
for d in 1..s1.len() {
if s1[d] != s2[d] {
return Err(MinarrowError::IncompatibleTypeError {
from: "NdArray",
to: "NdArray",
message: Some(format!(
"Dimension {} mismatch: {} vs {}", d, s1[d], s2[d]
)),
});
}
}
let mut new_shape: Vec<usize> = s1.to_vec();
new_shape[0] += s2[0];
let name = match (&self.name, &other.name) {
(Some(a), Some(b)) => Some(format!("{}+{}", a, b)),
(Some(a), None) => Some(a.clone()),
(None, Some(b)) => Some(b.clone()),
(None, None) => None,
};
if s1.len() == 1 {
let mut flat = Vec64::with_capacity(self.len() + other.len());
flat.extend(&self);
flat.extend(&other);
let mut result = NdArray::from_slice(&flat, &new_shape);
result.name = name;
return Ok(result);
}
if s1.len() == 2 && self.dims.strides()[0] == 1 && other.dims.strides()[0] == 1 {
let new_dims = NdDims::from_shape(&new_shape);
let new_stride = new_dims.strides()[1];
let total = buffer_len(&new_shape, new_dims.strides());
let mut buf = Vec64::with_capacity(total);
buf.0.resize(total, T::default());
let dst = buf.as_mut_slice();
for c in 0..s1[1] {
let dst_start = c * new_stride;
dst[dst_start..dst_start + s1[0]].copy_from_slice(self.col(c));
dst[dst_start + s1[0]..dst_start + s1[0] + s2[0]].copy_from_slice(other.col(c));
}
return Ok(NdArray {
data: Arc::new(Buffer::from_vec64(buf)),
dims: new_dims,
name,
});
}
let n1 = s1[0];
let n2 = s2[0];
let run_len = n1 + n2;
let n_runs: usize = s1[1..].iter().product();
let new_dims = NdDims::from_shape(&new_shape);
let total = buffer_len(&new_shape, new_dims.strides());
let mut buf = Vec64::with_capacity(total);
buf.0.resize(total, T::default());
let dst = buf.as_mut_slice();
let mut it_a = (&self).into_iter();
let mut it_b = (&other).into_iter();
for r in 0..n_runs {
let base = r * run_len;
for i in 0..n1 {
dst[base + i] = it_a.next().unwrap();
}
for i in 0..n2 {
dst[base + n1 + i] = it_b.next().unwrap();
}
}
Ok(NdArray {
data: Arc::new(Buffer::from_vec64(buf)),
dims: new_dims,
name,
})
}
}
#[cfg(all(feature = "views", feature = "select"))]
impl<T: Float> AxisSelection for NdArray<T> {
type View = NdArrayV<T>;
fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
self.slice(selection)
}
fn get_axis_count(&self) -> usize {
self.ndim()
}
}
#[cfg(all(feature = "views", feature = "select"))]
impl<T: Float> RowSelection for NdArray<T> {
type View = NdArrayV<T>;
fn r<S: DataSelector>(&self, selection: S) -> NdArrayV<T> {
assert!(self.ndim() > 0, "row selection requires an axis 0");
let n_obs = self.shape()[0];
let indices = selection.resolve_indices(n_obs);
if selection.is_contiguous() {
let start = indices.first().copied().unwrap_or(0);
let ranges: Vec<Range<usize>> = std::iter::once(start..start + indices.len())
.chain(self.shape()[1..].iter().map(|&n| 0..n))
.collect();
let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect();
return self.slice(&refs);
}
NdArrayV::from_ndarray(gather_obs_impl(
&indices,
self.shape(),
self.name.clone(),
|idx| self.get(idx),
))
}
fn get_row_count(&self) -> usize {
self.n_obs()
}
}
#[cfg(all(feature = "views", feature = "select"))]
pub(crate) fn gather_obs_impl<T: Float>(
indices: &[usize],
shape: &[usize],
name: Option<String>,
get: impl Fn(&[usize]) -> T,
) -> NdArray<T> {
let mut out_shape = shape.to_vec();
out_shape[0] = indices.len();
let total: usize = out_shape.iter().product();
let mut flat = Vec64::with_capacity(total);
let ndim = shape.len();
let mut idx = vec![0usize; ndim];
let inner_runs: usize = shape[1..].iter().product::<usize>().max(1);
for _ in 0..inner_runs {
for &obs in indices {
idx[0] = obs;
flat.push(get(&idx));
}
let mut carry = true;
for d in 1..ndim {
if carry {
idx[d] += 1;
if idx[d] < shape[d] {
carry = false;
} else {
idx[d] = 0;
}
}
}
}
let mut result = NdArray::from_slice(&flat, &out_shape);
result.name = name;
result
}
impl<T: Float> Index<usize> for NdArray<T> {
type Output = [T];
#[inline]
fn index(&self, idx: usize) -> &[T] {
let shape = self.dims.shape();
let strides = self.dims.strides();
match shape.len() {
0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1 => {
assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
&self.data.as_slice()[idx..idx + 1]
}
2 => {
assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
let start = idx * strides[1];
&self.data.as_slice()[start..start + shape[0]]
}
n => {
assert!(
self.is_contiguous(),
"outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
);
let last = n - 1;
assert!(idx < shape[last], "index out of bounds for axis {}", last);
let start = idx * strides[last];
&self.data.as_slice()[start..start + strides[last]]
}
}
}
}
impl<T: Float> IndexMut<usize> for NdArray<T> {
#[inline]
fn index_mut(&mut self, idx: usize) -> &mut [T] {
let shape = self.dims.shape().to_vec();
let strides = self.dims.strides().to_vec();
match shape.len() {
0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1 => {
assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[idx..idx + 1]
}
2 => {
assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
let start = idx * strides[1];
let n_rows = shape[0];
&mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
}
n => {
assert!(
self.is_contiguous(),
"outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
);
let last = n - 1;
assert!(idx < shape[last], "index out of bounds for axis {}", last);
let start = idx * strides[last];
&mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + strides[last]]
}
}
}
}
impl<T: Float> Index<Range<usize>> for NdArray<T> {
type Output = [T];
#[inline]
fn index(&self, range: Range<usize>) -> &[T] {
let shape = self.dims.shape();
let strides = self.dims.strides();
match shape.len() {
0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1 => &self.data.as_slice()[range],
_ => {
assert!(
self.is_contiguous(),
"range indexing requires a contiguous layout, use slice() for strided access"
);
let last = shape.len() - 1;
assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
let start = range.start * strides[last];
let end = range.end * strides[last];
&self.data.as_slice()[start..end]
}
}
}
}
impl<T: Float> IndexMut<Range<usize>> for NdArray<T> {
#[inline]
fn index_mut(&mut self, range: Range<usize>) -> &mut [T] {
let shape = self.dims.shape().to_vec();
let strides = self.dims.strides().to_vec();
match shape.len() {
0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1 => &mut Arc::make_mut(&mut self.data).as_mut_slice()[range],
_ => {
assert!(
self.is_contiguous(),
"range indexing requires a contiguous layout, use slice() for strided access"
);
let last = shape.len() - 1;
assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
let start = range.start * strides[last];
let end = range.end * strides[last];
&mut Arc::make_mut(&mut self.data).as_mut_slice()[start..end]
}
}
}
}
impl<T: Float> Index<RangeFrom<usize>> for NdArray<T> {
type Output = [T];
#[inline]
fn index(&self, range: RangeFrom<usize>) -> &[T] {
assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
let last_dim = self.dims.shape().len() - 1;
let end = self.dims.shape()[last_dim];
&self[range.start..end]
}
}
impl<T: Float> Index<RangeTo<usize>> for NdArray<T> {
type Output = [T];
#[inline]
fn index(&self, range: RangeTo<usize>) -> &[T] {
&self[0..range.end]
}
}
impl<T: Float> Index<RangeFull> for NdArray<T> {
type Output = [T];
#[inline]
fn index(&self, _: RangeFull) -> &[T] {
assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
let last_dim = self.dims.shape().len() - 1;
let end = self.dims.shape()[last_dim];
&self[0..end]
}
}
impl<T: Float> Index<()> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (): ()) -> &T {
&self.data.as_slice()[self.offset_of(&[])]
}
}
impl<T: Float> Index<(usize,)> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (i,): (usize,)) -> &T {
&self.data.as_slice()[self.offset_of(&[i])]
}
}
impl<T: Float> Index<(usize, usize)> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (i, j): (usize, usize)) -> &T {
&self.data.as_slice()[self.offset_of(&[i, j])]
}
}
impl<T: Float> Index<(usize, usize, usize)> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (i, j, k): (usize, usize, usize)) -> &T {
&self.data.as_slice()[self.offset_of(&[i, j, k])]
}
}
impl<T: Float> Index<(usize, usize, usize, usize)> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (i, j, k, l): (usize, usize, usize, usize)) -> &T {
&self.data.as_slice()[self.offset_of(&[i, j, k, l])]
}
}
impl<T: Float> Index<(usize, usize, usize, usize, usize)> for NdArray<T> {
type Output = T;
#[inline]
fn index(&self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &T {
&self.data.as_slice()[self.offset_of(&[i, j, k, l, m])]
}
}
impl<T: Float> IndexMut<(usize,)> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (i,): (usize,)) -> &mut T {
let off = self.offset_of(&[i]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> IndexMut<()> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (): ()) -> &mut T {
let off = self.offset_of(&[]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> IndexMut<(usize, usize)> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut T {
let off = self.offset_of(&[i, j]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> IndexMut<(usize, usize, usize)> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (i, j, k): (usize, usize, usize)) -> &mut T {
let off = self.offset_of(&[i, j, k]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> IndexMut<(usize, usize, usize, usize)> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (i, j, k, l): (usize, usize, usize, usize)) -> &mut T {
let off = self.offset_of(&[i, j, k, l]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> IndexMut<(usize, usize, usize, usize, usize)> for NdArray<T> {
#[inline]
fn index_mut(&mut self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &mut T {
let off = self.offset_of(&[i, j, k, l, m]);
&mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
}
}
impl<T: Float> From<&[T]> for NdArray<T> {
fn from(data: &[T]) -> Self {
NdArray {
data: Arc::new(Buffer::from_slice(data)),
dims: NdDims::from_shape(&[data.len()]),
name: None,
}
}
}
impl<T: Float> From<Vec64<T>> for NdArray<T> {
fn from(v: Vec64<T>) -> Self {
let n = v.len();
NdArray {
data: Arc::new(Buffer::from_vec64(v)),
dims: NdDims::from_shape(&[n]),
name: None,
}
}
}
impl<T: Float> From<&[Vec<T>]> for NdArray<T> {
fn from(columns: &[Vec<T>]) -> Self {
let n_cols = columns.len();
if n_cols == 0 {
return NdArray::new(&[0, 0]);
}
let n_rows = columns[0].len();
for col in columns {
assert_eq!(col.len(), n_rows, "Column length mismatch");
}
let shape = [n_rows, n_cols];
let dims = NdDims::from_shape(&shape);
let stride = dims.strides()[1];
let total = buffer_len(&shape, dims.strides());
let mut buf = Vec64::with_capacity(total);
buf.0.resize(total, T::default());
for (c, col) in columns.iter().enumerate() {
let start = c * stride;
buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col);
}
NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
}
}
impl<T: Float> From<&[FloatArray<T>]> for NdArray<T> {
fn from(columns: &[FloatArray<T>]) -> Self {
let n_cols = columns.len();
if n_cols == 0 {
return NdArray::new(&[0, 0]);
}
let n_rows = columns[0].data.len();
for col in columns {
assert_eq!(col.data.len(), n_rows, "Column length mismatch");
}
let shape = [n_rows, n_cols];
let dims = NdDims::from_shape(&shape);
let stride = dims.strides()[1];
let total = buffer_len(&shape, dims.strides());
let mut buf = Vec64::with_capacity(total);
buf.0.resize(total, T::default());
for (c, col) in columns.iter().enumerate() {
let start = c * stride;
buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col.data.as_slice());
}
NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
}
}
#[cfg(feature = "matrix")]
impl From<Matrix> for NdArray<f64> {
fn from(mat: Matrix) -> Self {
let shape = [mat.n_rows, mat.n_cols];
let strides = [1, mat.stride];
NdArray {
data: Arc::new(mat.data),
dims: NdDims::from_shape_and_strides(&shape, &strides),
name: mat.name,
}
}
}
impl TryFrom<&Table> for NdArray<f64> {
type Error = MinarrowError;
fn try_from(table: &Table) -> Result<Self, Self::Error> {
let n_cols = table.n_cols();
let n_rows = table.n_rows;
if n_cols == 0 {
return Ok(NdArray::new(&[0, 0]));
}
let shape = [n_rows, n_cols];
let dims = NdDims::from_shape(&shape);
let stride = dims.strides()[1];
let total = buffer_len(&shape, dims.strides());
let mut buf = Vec64::with_capacity(total);
buf.0.resize(total, 0.0);
for (col_idx, fa) in table.cols.iter().enumerate() {
let numeric = fa.array.try_num().map_err(|_| MinarrowError::TypeError {
from: "non-numeric",
to: "Float64",
message: Some(format!("column {} is not numeric", col_idx)),
})?;
let f64_arr = numeric.try_f64()?;
if f64_arr.data.len() != n_rows {
return Err(MinarrowError::ColumnLengthMismatch {
col: col_idx,
expected: n_rows,
found: f64_arr.data.len(),
});
}
let start = col_idx * stride;
let src = f64_arr.data.as_slice();
let dst = &mut buf.as_mut_slice()[start..start + n_rows];
match f64_arr.null_mask.as_ref() {
Some(mask) => {
for i in 0..n_rows {
dst[i] = if mask.get(i) { src[i] } else { f64::NAN };
}
}
None => dst.copy_from_slice(src),
}
}
let name = if table.name.is_empty() { None } else { Some(table.name.clone()) };
Ok(NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name })
}
}
impl TryFrom<Table> for NdArray<f64> {
type Error = MinarrowError;
fn try_from(table: Table) -> Result<Self, Self::Error> {
NdArray::try_from(&table)
}
}
impl<T: Float> fmt::Debug for NdArray<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f, "NdArray{}: {:?} [{}D, col-major]",
self.name.as_deref().map_or(String::new(), |n| format!(" '{}'", n)),
self.dims.shape(),
self.ndim(),
)?;
if self.ndim() == 0 {
write!(f, "\n{:8.4}", self.get(&[]).to_f64().unwrap_or(f64::NAN))?;
} else if self.ndim() == 2 {
let shape = self.dims.shape();
let max_rows = shape[0].min(6);
let max_cols = shape[1].min(8);
for r in 0..max_rows {
write!(f, "\n[")?;
for c in 0..max_cols {
write!(f, " {:8.4}", self.get(&[r, c]).to_f64().unwrap_or(f64::NAN))?;
if c < max_cols - 1 { write!(f, ",")?; }
}
if shape[1] > 8 { write!(f, " ...")?; }
write!(f, " ]")?;
}
if shape[0] > 6 { write!(f, "\n...")?; }
} else if self.ndim() == 1 {
let n = self.dims.shape()[0].min(10);
write!(f, "\n[")?;
for i in 0..n {
write!(f, " {:8.4}", self.get(&[i]).to_f64().unwrap_or(f64::NAN))?;
if i < n - 1 { write!(f, ",")?; }
}
if self.dims.shape()[0] > 10 { write!(f, " ...")?; }
write!(f, " ]")?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::StringArray;
use crate::structs::bitmask::Bitmask;
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn axis_selection_rank10() {
let shape = [3, 4, 2, 5, 3, 2, 4, 3, 2, 5];
let len: usize = shape.iter().product();
let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
let a = NdArray::from_slice(&data, &shape);
let v = a.s(nd![1..3, 0..2, 1, 2..5, .., 0..1, 1..3, 2, 0..2, 3..5]);
assert_eq!(v.shape(), &[2, 2, 3, 3, 1, 2, 2, 2]);
assert_eq!(
v.get(&[0, 0, 0, 0, 0, 0, 0, 0]),
a.get(&[1, 0, 1, 2, 0, 0, 1, 2, 0, 3])
);
assert_eq!(
v.get(&[1, 1, 2, 2, 0, 1, 1, 1]),
a.get(&[2, 1, 1, 4, 2, 0, 2, 2, 1, 4])
);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn axis_selection_runtime_rank() {
let mut shape = vec![1usize; 100];
shape[0] = 3;
shape[10] = 4;
shape[50] = 5;
shape[99] = 2;
let len: usize = shape.iter().product();
let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
let a = NdArray::from_slice(&data, &shape);
let mut sels: Vec<Box<dyn DataSelector>> =
shape.iter().map(|&n| Box::new(0..n) as Box<dyn DataSelector>).collect();
sels[0] = Box::new(1..3);
sels[10] = Box::new(2usize);
sels[50] = Box::new(1..4);
let refs: Vec<&dyn DataSelector> = sels.iter().map(|s| s.as_ref()).collect();
let v = a.s(&refs);
assert_eq!(v.ndim(), 99);
assert_eq!(v.shape()[0], 2);
assert_eq!(v.shape()[49], 3);
assert_eq!(v.shape()[98], 2);
let mut view_idx = vec![0usize; 99];
view_idx[0] = 1;
view_idx[49] = 2;
view_idx[98] = 1;
let mut source_idx = vec![0usize; 100];
source_idx[0] = 2;
source_idx[10] = 2;
source_idx[50] = 3;
source_idx[99] = 1;
assert_eq!(v.get(&view_idx), a.get(&source_idx));
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn axis_selection_trait() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.s(nd![1..3, 1]);
assert_eq!(v.shape(), &[2]);
assert_eq!(v.get(&[0]), 5.0);
assert_eq!(v.get(&[1]), 6.0);
let sub = a.s(nd![0..3, 0..2]).s(nd![2, 0..2]);
assert_eq!(sub.shape(), &[2]);
assert_eq!(sub.get(&[0]), 3.0);
assert_eq!(sub.get(&[1]), 6.0);
assert_eq!(a.get_axis_count(), 2);
assert_eq!(sub.get_axis_count(), 1);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn row_selection_contiguous() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.r(1..3);
assert_eq!(v.shape(), &[2, 2]);
assert_eq!(v.get(&[0, 0]), 2.0);
assert_eq!(v.get(&[1, 1]), 6.0);
let single = a.row(2);
assert_eq!(single.shape(), &[1, 2]);
assert_eq!(single.get(&[0, 1]), 6.0);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn row_selection_gathers_indices() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.r(&[2, 0]);
assert_eq!(v.shape(), &[2, 2]);
assert_eq!(v.get(&[0, 0]), 3.0);
assert_eq!(v.get(&[0, 1]), 6.0);
assert_eq!(v.get(&[1, 0]), 1.0);
assert_eq!(v.get(&[1, 1]), 4.0);
}
#[test]
fn apply_maps_elements() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
a.set_name("m");
let b = a.apply(|x| x * 10.0);
assert_eq!(b.get(&[1, 1]), 40.0);
assert_eq!(b.name.as_deref(), Some("m"));
assert_eq!(a.get(&[1, 1]), 4.0);
}
#[test]
fn apply_mut_in_place() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
a.apply_mut(|x| x + 0.5);
assert_eq!((&a).into_iter().collect::<Vec<f64>>(), vec![1.5, 2.5, 3.5]);
}
#[cfg(feature = "matrix")]
#[test]
fn apply_mut_non_contiguous_touches_logical_only() {
let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None);
let mut a = NdArray::from(mat);
assert!(!a.is_contiguous());
a.apply_mut(|x| x * 2.0);
assert_eq!(a.get(&[0, 0]), 2.0);
assert_eq!(a.get(&[2, 1]), 12.0);
}
#[cfg(feature = "views")]
#[test]
fn apply_axis_collapses_axis() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let col_sums = a.apply_axis(0, |lane| (&lane).into_iter().sum());
assert_eq!(col_sums.shape(), &[2]);
assert_eq!(col_sums.get(&[0]), 6.0);
assert_eq!(col_sums.get(&[1]), 15.0);
let row_sums = a.apply_axis(1, |lane| (&lane).into_iter().sum());
assert_eq!(row_sums.shape(), &[3]);
assert_eq!(row_sums.get(&[0]), 5.0);
assert_eq!(row_sums.get(&[2]), 9.0);
}
#[cfg(feature = "views")]
#[test]
fn apply_axis_3d() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
let maxes = a.apply_axis(1, |lane| {
(&lane).into_iter().fold(f64::MIN, f64::max)
});
assert_eq!(maxes.shape(), &[2, 4]);
assert_eq!(maxes.get(&[0, 0]), a.get(&[0, 2, 0]));
assert_eq!(maxes.get(&[1, 3]), a.get(&[1, 2, 3]));
}
#[test]
fn new_zeroed_1d() {
let a = NdArray::<f64>::new(&[5]);
assert_eq!(a.ndim(), 1);
assert_eq!(a.shape(), &[5]);
assert_eq!(a.len(), 5);
assert!(!a.is_empty());
for v in &a { assert_eq!(v, 0.0); }
}
#[test]
fn new_zeroed_2d() {
let a = NdArray::<f64>::new(&[3, 4]);
assert_eq!(a.ndim(), 2);
assert_eq!(a.shape(), &[3, 4]);
assert_eq!(a.len(), 12);
for v in &a { assert_eq!(v, 0.0); }
}
#[test]
fn new_zeroed_3d() {
let a = NdArray::<f64>::new(&[2, 3, 4]);
assert_eq!(a.ndim(), 3);
assert_eq!(a.shape(), &[2, 3, 4]);
assert_eq!(a.len(), 24);
}
#[test]
fn new_zeroed_5d() {
let a = NdArray::<f64>::new(&[2, 3, 4, 5, 6]);
assert_eq!(a.ndim(), 5);
assert_eq!(a.len(), 720);
}
#[test]
fn new_zeroed_6d_heap() {
let a = NdArray::<f64>::new(&[2, 3, 2, 2, 2, 2]);
assert_eq!(a.ndim(), 6);
assert_eq!(a.len(), 96);
}
#[test]
fn new_named() {
let a = NdArray::<f64>::new_named(&[3, 3], "covariance");
assert_eq!(a.name.as_deref(), Some("covariance"));
}
#[test]
fn from_slice_1d() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0];
let a = NdArray::from_slice(&data, &[5]);
assert_eq!(a.len(), 5);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
}
#[test]
fn from_slice_2d_column_major() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let a = NdArray::from_slice(&data, &[3, 2]);
assert_eq!(a.get(&[0, 0]), 1.0);
assert_eq!(a.get(&[1, 0]), 2.0);
assert_eq!(a.get(&[2, 0]), 3.0);
assert_eq!(a.get(&[0, 1]), 4.0);
assert_eq!(a.get(&[1, 1]), 5.0);
assert_eq!(a.get(&[2, 1]), 6.0);
}
#[test]
fn fill_and_ones() {
let a = NdArray::fill(&[2, 3], 7.0);
assert_eq!(a.len(), 6);
for v in &a { assert_eq!(v, 7.0); }
let b = NdArray::<f64>::ones(&[4]);
for v in &b { assert_eq!(v, 1.0); }
}
#[test]
fn eye_identity() {
let a = NdArray::<f64>::eye(3);
assert_eq!(a.shape(), &[3, 3]);
assert_eq!(a.get(&[0, 0]), 1.0);
assert_eq!(a.get(&[1, 1]), 1.0);
assert_eq!(a.get(&[2, 2]), 1.0);
assert_eq!(a.get(&[0, 1]), 0.0);
assert_eq!(a.get(&[1, 0]), 0.0);
}
#[test]
fn linspace_basic() {
let a = NdArray::<f64>::linspace(0.0, 1.0, 5);
assert_eq!(a.shape(), &[5]);
assert_eq!(a.get(&[0]), 0.0);
assert_eq!(a.get(&[4]), 1.0);
assert!((a.get(&[2]) - 0.5).abs() < 1e-15);
}
#[test]
fn arange_basic() {
let a = NdArray::arange(0.0, 0.5, 4);
assert_eq!(a.shape(), &[4]);
assert_eq!(a.get(&[0]), 0.0);
assert_eq!(a.get(&[1]), 0.5);
assert_eq!(a.get(&[2]), 1.0);
assert_eq!(a.get(&[3]), 1.5);
}
#[test]
fn from_vec64_moves_data() {
let data: Vec64<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into();
let a = NdArray::from_vec64(data, &[3, 2]);
assert_eq!(a.shape(), &[3, 2]);
assert_eq!(a.get(&[2, 1]), 6.0);
}
#[cfg(feature = "views")]
#[test]
fn iter_obs_walks_observations() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let rows: Vec<(usize, Vec<f64>)> = a
.iter_obs()
.map(|(i, v)| (i, (&v).into_iter().collect()))
.collect();
assert_eq!(rows.len(), 3);
assert_eq!(rows[0], (0, vec![1.0, 4.0]));
assert_eq!(rows[2], (2, vec![3.0, 6.0]));
}
#[test]
fn equality_ignores_name() {
let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
let mut b = a.clone();
b.name = Some("named".to_string());
assert_eq!(a, b);
}
#[test]
fn name_survives_reshape_and_concat() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
a.name = Some("px".to_string());
assert_eq!(a.reshape(&[2, 2]).unwrap().name.as_deref(), Some("px"));
let mut b = NdArray::from_slice(&[5.0, 6.0], &[2]);
b.name = Some("qty".to_string());
let joined = a.concat(b).unwrap();
assert_eq!(joined.name.as_deref(), Some("px+qty"));
}
#[test]
fn get_set_1d() {
let mut a = NdArray::new(&[3]);
a.set(&[0], 10.0);
a.set(&[1], 20.0);
a.set(&[2], 30.0);
assert_eq!(a.get(&[0]), 10.0);
assert_eq!(a.get(&[1]), 20.0);
assert_eq!(a.get(&[2]), 30.0);
}
#[test]
fn tuple_index_1d() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
assert_eq!(a[(0,)], 10.0);
assert_eq!(a[(1,)], 20.0);
assert_eq!(a[(2,)], 30.0);
}
#[test]
fn tuple_index_2d() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let a = NdArray::from_slice(&data, &[3, 2]);
assert_eq!(a[(0, 0)], 1.0);
assert_eq!(a[(2, 0)], 3.0);
assert_eq!(a[(0, 1)], 4.0);
assert_eq!(a[(2, 1)], 6.0);
}
#[test]
fn tuple_index_3d() {
let data: Vec<f64> = (1..=12).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 2]);
assert_eq!(a[(0, 0, 0)], 1.0);
assert_eq!(a[(1, 0, 0)], 2.0);
assert_eq!(a[(0, 1, 0)], 3.0);
assert_eq!(a[(1, 1, 0)], 4.0);
assert_eq!(a[(0, 2, 0)], 5.0);
assert_eq!(a[(1, 2, 0)], 6.0);
assert_eq!(a[(0, 0, 1)], 7.0);
assert_eq!(a[(1, 2, 1)], 12.0);
}
#[test]
fn index_mut_2d() {
let mut a = NdArray::new(&[2, 2]);
a[(0, 0)] = 1.0;
a[(1, 0)] = 2.0;
a[(0, 1)] = 3.0;
a[(1, 1)] = 4.0;
assert_eq!(a[(0, 0)], 1.0);
assert_eq!(a[(1, 1)], 4.0);
}
#[test]
fn iter_1d() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals, vec![10.0, 20.0, 30.0]);
}
#[test]
fn iter_2d_column_major_order() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn iter_2d_compact_layout() {
let data: Vec<f64> = (1..=30).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[10, 3]);
assert_eq!(a.strides()[1], 10);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals.len(), 30);
assert_eq!(&vals[..10], &data[..10]);
assert_eq!(&vals[10..20], &data[10..20]);
assert_eq!(&vals[20..30], &data[20..30]);
}
#[test]
fn iter_3d_column_major_order() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals.len(), 24);
assert_eq!(vals, data);
}
#[test]
fn iter_exact_size() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let iter = (&a).into_iter();
assert_eq!(iter.len(), 6);
}
#[test]
fn consuming_into_iter() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
let vals: Vec<f64> = a.into_iter().collect();
assert_eq!(vals, vec![1.0, 2.0, 3.0]);
}
#[test]
fn rank_zero_scalar_semantics() {
let mut a = NdArray::from_slice(&[5.0], &[]);
assert_eq!(a.ndim(), 0);
assert!(a.shape().is_empty());
assert!(a.strides().is_empty());
assert_eq!(a.len(), 1);
assert!(!a.is_empty());
assert!(a.is_contiguous());
assert_eq!(a[()], 5.0);
assert_eq!((&a).into_iter().collect::<Vec<_>>(), vec![5.0]);
assert_eq!(Shape::shape(&a), ShapeDim::Rank0(1));
a[()] = 7.0;
assert_eq!(a[()], 7.0);
assert!(a.transpose().shape().is_empty());
assert_eq!(a.reshape(&[1]).unwrap().get(&[0]), 7.0);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn selecting_every_axis_can_produce_a_scalar_view() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let scalar = a.slice(&[&1usize, &1usize]);
assert!(scalar.shape().is_empty());
assert!(scalar.strides().is_empty());
assert_eq!(scalar.len(), 1);
assert_eq!(scalar.get(&[]), 4.0);
assert_eq!((&scalar).into_iter().collect::<Vec<_>>(), vec![4.0]);
assert_eq!(Shape::shape(&scalar), ShapeDim::Rank0(1));
}
#[test]
fn is_contiguous_default() {
let a = NdArray::<f64>::new(&[3, 4]);
assert!(a.is_contiguous());
}
#[test]
fn shape_trait_1d() {
let a = NdArray::<f64>::new(&[5]);
assert_eq!(Shape::shape(&a), ShapeDim::Rank1(5));
}
#[test]
fn shape_trait_2d() {
let a = NdArray::<f64>::new(&[3, 4]);
assert_eq!(Shape::shape(&a), ShapeDim::Rank2 { rows: 3, cols: 4 });
}
#[test]
fn shape_trait_3d() {
let a = NdArray::<f64>::new(&[2, 3, 4]);
assert_eq!(Shape::shape(&a), ShapeDim::RankN(vec![2, 3, 4]));
}
#[test]
fn has_nan_false() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
assert!(!a.has_nan());
assert_eq!(a.nan_count(), 0);
}
#[test]
fn has_nan_true() {
let a = NdArray::from_slice(&[1.0, f64::NAN, 3.0], &[3]);
assert!(a.has_nan());
assert_eq!(a.nan_count(), 1);
}
#[test]
fn has_nan_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, f64::NAN, 4.0, 5.0, f64::NAN], &[3, 2]);
assert!(a.has_nan());
assert_eq!(a.nan_count(), 2);
}
#[test]
fn col_access() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
}
#[test]
fn col_mut_access() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
a.col_mut(0)[1] = 99.0;
assert_eq!(a.col(0), &[1.0, 99.0, 3.0]);
}
#[test]
fn columns_access() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let cols = a.columns();
assert_eq!(cols.len(), 2);
assert_eq!(cols[0], &[1.0, 2.0, 3.0]);
assert_eq!(cols[1], &[4.0, 5.0, 6.0]);
}
#[test]
fn columns_mut_access() {
let mut a = NdArray::new(&[3, 2]);
{
let mut cols = a.columns_mut();
cols[0].copy_from_slice(&[1.0, 2.0, 3.0]);
cols[1].copy_from_slice(&[4.0, 5.0, 6.0]);
}
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "views")]
#[test]
fn obs_access() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.as_view();
let obs0: Vec<f64> = (&v.obs(0)).into_iter().collect();
let obs1: Vec<f64> = (&v.obs(1)).into_iter().collect();
let obs2: Vec<f64> = (&v.obs(2)).into_iter().collect();
assert_eq!(obs0, vec![1.0, 4.0]);
assert_eq!(obs1, vec![2.0, 5.0]);
assert_eq!(obs2, vec![3.0, 6.0]);
assert_eq!((&a.obs(1)).into_iter().collect::<Vec<f64>>(), vec![2.0, 5.0]);
let b = NdArray::from_slice(&(1..=24).map(|x| x as f64).collect::<Vec<_>>(), &[2, 3, 4]);
let obs0_3d = b.obs(0);
assert_eq!(obs0_3d.shape(), &[3, 4]);
}
#[test]
fn col_access_2d() {
let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[10, 2]);
assert_eq!(a.col(0).len(), 10);
assert_eq!(a.col(1).len(), 10);
assert_eq!(a.col(0), &data[..10]);
assert_eq!(a.col(1), &data[10..20]);
}
#[test]
fn blas_params() {
let a = NdArray::<f64>::new(&[10, 5]);
assert_eq!(a.m(), 10);
assert_eq!(a.n(), 5);
assert_eq!(a.lda(), 10);
}
#[test]
fn blas_params_aligned_rows() {
let a = NdArray::<f64>::new(&[8, 3]);
assert_eq!(a.m(), 8);
assert_eq!(a.n(), 3);
assert_eq!(a.lda(), 8);
}
#[test]
fn compact_strides_2d() {
for n_rows in 1..=20 {
let a = NdArray::<f64>::new(&[n_rows, 3]);
assert_eq!(a.strides()[1], n_rows);
assert!(a.is_contiguous());
}
}
#[test]
fn compact_strides_3d() {
let a = NdArray::<f64>::new(&[10, 3, 4]);
let strides = a.strides();
assert_eq!(strides[0], 1);
assert_eq!(strides[1], 10);
assert_eq!(strides[2], 10 * 3);
assert!(a.is_contiguous());
}
#[test]
fn reshape_1d_to_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[6]);
let b = a.reshape(&[3, 2]).unwrap();
assert_eq!(b.shape(), &[3, 2]);
assert_eq!(b.get(&[0, 0]), 1.0);
assert_eq!(b.get(&[1, 0]), 2.0);
assert_eq!(b.get(&[2, 0]), 3.0);
assert_eq!(b.get(&[0, 1]), 4.0);
}
#[test]
fn reshape_size_mismatch() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
assert!(a.reshape(&[2, 2]).is_err());
}
#[test]
fn transpose_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let t = a.transpose();
assert_eq!(t.shape(), &[2, 3]);
assert_eq!(t.get(&[0, 0]), 1.0);
assert_eq!(t.get(&[1, 0]), 4.0);
assert_eq!(t.get(&[0, 1]), 2.0);
assert_eq!(t.get(&[1, 1]), 5.0);
assert_eq!(t.get(&[0, 2]), 3.0);
assert_eq!(t.get(&[1, 2]), 6.0);
}
#[test]
fn transpose_3d_reverses_axes() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
let t = a.transpose();
assert_eq!(t.shape(), &[4, 3, 2]);
assert!(t.is_contiguous());
for i in 0..2 {
for j in 0..3 {
for k in 0..4 {
assert_eq!(t.get(&[k, j, i]), a.get(&[i, j, k]));
}
}
}
}
#[test]
fn transpose_1d_copies_through() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
let t = a.transpose();
assert_eq!(t.shape(), &[3]);
assert_eq!(t, a);
}
#[test]
fn flatten() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let flat = a.flatten();
assert_eq!(flat.shape(), &[6]);
let vals: Vec<f64> = (&flat).into_iter().collect();
assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn to_contiguous_noop() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
let b = a.to_contiguous();
assert_eq!(a, b);
}
#[test]
fn fill_with_contiguous() {
let mut a = NdArray::new(&[3, 2]);
a.fill_with(42.0);
for v in &a { assert_eq!(v, 42.0); }
}
#[test]
fn from_f64_slice() {
let a = NdArray::from(&[1.0, 2.0, 3.0][..]);
assert_eq!(a.ndim(), 1);
assert_eq!(a.len(), 3);
assert_eq!(a[(0,)], 1.0);
}
#[test]
fn from_vec64() {
let v: Vec64<f64> = vec![10.0, 20.0].into_iter().collect();
let a = NdArray::from(v);
assert_eq!(a.ndim(), 1);
assert_eq!(a[(0,)], 10.0);
assert_eq!(a[(1,)], 20.0);
}
#[test]
fn from_column_vecs() {
let cols = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
let a = NdArray::from(cols.as_slice());
assert_eq!(a.shape(), &[3, 2]);
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
}
#[test]
fn from_float_arrays() {
let c0 = FloatArray::from_slice(&[1.0, 2.0]);
let c1 = FloatArray::from_slice(&[3.0, 4.0]);
let a = NdArray::from([c0, c1].as_slice());
assert_eq!(a.shape(), &[2, 2]);
assert_eq!(a.col(0), &[1.0, 2.0]);
assert_eq!(a.col(1), &[3.0, 4.0]);
}
#[test]
fn from_buffer_explicit_strides() {
let mut buf = Vec64::with_capacity(16);
buf.0.resize(16, 0.0);
buf[0] = 1.0; buf[1] = 2.0; buf[2] = 3.0;
buf[8] = 4.0; buf[9] = 5.0; buf[10] = 6.0;
let a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "matrix")]
#[test]
fn from_matrix() {
let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, Some("m".into()));
let a = NdArray::from(mat);
assert_eq!(a.shape(), &[3, 2]);
assert_eq!(a.name.as_deref(), Some("m"));
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "matrix")]
#[test]
fn to_matrix_roundtrip() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let a = NdArray::from_slice(&data, &[3, 2]);
let mat = a.to_matrix().unwrap();
assert_eq!(mat.n_rows, 3);
assert_eq!(mat.n_cols, 2);
assert_eq!(mat.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(mat.col(1), &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "matrix")]
#[test]
fn to_matrix_non_2d_fails() {
let a = NdArray::new(&[5]);
assert!(a.to_matrix().is_err());
}
fn make_numeric_table() -> Table {
let c0 = FieldArray::from_arr("x", Array::NumericArray(
NumericArray::Float64(Arc::new(FloatArray::from_slice(&[1.0, 2.0, 3.0])))
));
let c1 = FieldArray::from_arr("y", Array::NumericArray(
NumericArray::Float64(Arc::new(FloatArray::from_slice(&[4.0, 5.0, 6.0])))
));
Table::new("data".to_string(), Some(vec![c0, c1]))
}
#[test]
fn try_from_table() {
let table = make_numeric_table();
let a = NdArray::try_from(&table).unwrap();
assert_eq!(a.shape(), &[3, 2]);
assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
assert_eq!(a.name.as_deref(), Some("data"));
}
#[test]
fn try_from_table_with_nulls_converts_to_nan() {
let mut mask = Bitmask::new_set_all(3, true);
mask.set(1, false);
let arr = FloatArray::new(Buffer::from_slice(&[10.0, 0.0, 30.0]), Some(mask));
let c0 = FieldArray::from_arr("v", Array::NumericArray(
NumericArray::Float64(Arc::new(arr))
));
let table = Table::new("nulls".to_string(), Some(vec![c0]));
let a = NdArray::try_from(&table).unwrap();
assert_eq!(a.get(&[0, 0]), 10.0);
assert!(a.get(&[1, 0]).is_nan());
assert_eq!(a.get(&[2, 0]), 30.0);
}
#[test]
fn try_from_table_coerces_unparseable_text_to_nan() {
let c0 = FieldArray::from_arr("name", Array::from_string32(
StringArray::from_slice(&["a", "b"])
));
let table = Table::new("text".to_string(), Some(vec![c0]));
let a = NdArray::try_from(&table).unwrap();
assert_eq!(a.shape(), &[2, 1]);
assert!(a.get(&[0, 0]).is_nan());
assert!(a.get(&[1, 0]).is_nan());
}
#[test]
fn to_table_roundtrip() {
let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let a = NdArray::from_slice(&data, &[3, 2]);
let fields = vec![
Field::new("x", ArrowType::Float64, false, None),
Field::new("y", ArrowType::Float64, false, None),
];
let table = a.to_table(Some(fields)).unwrap();
assert_eq!(table.n_rows(), 3);
assert_eq!(table.n_cols(), 2);
assert_eq!(table.col_names(), vec!["x", "y"]);
let col0 = table.cols[0].array.num().f64();
assert_eq!(col0.data.as_slice(), &[1.0, 2.0, 3.0]);
}
#[test]
fn to_table_generated_names() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let table = a.to_table(None).unwrap();
assert_eq!(table.col_names(), vec!["col_0", "col_1"]);
}
#[test]
fn to_table_rejects_incompatible_field_metadata() {
let wrong_dtype = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
Field::new("value", ArrowType::Float32, false, None),
]));
assert!(matches!(wrong_dtype, Err(MinarrowError::TypeError { .. })));
let nullable = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
Field::new("value", ArrowType::Float64, true, None),
]));
assert!(matches!(nullable, Err(MinarrowError::NullError { .. })));
}
#[test]
fn to_table_non_2d_fails() {
let a = NdArray::new(&[5]);
assert!(a.to_table(None).is_err());
}
#[test]
fn to_array_1d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
let arr = a.to_array().unwrap();
let f = arr.num().f64();
assert_eq!(f.data.as_slice(), &[1.0, 2.0, 3.0]);
}
#[test]
fn to_array_materialises_strided_logical_values() {
let a = NdArray::from_buffer(
Buffer::from_slice(&[1.0, 99.0, 2.0, 99.0, 3.0]),
&[3],
&[2],
);
let arr = a.to_array().unwrap();
assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
}
#[test]
fn to_array_ignores_unused_backing_elements() {
let a = NdArray::from_buffer(
Buffer::from_slice(&[1.0, 2.0, 3.0, 99.0]),
&[3],
&[1],
);
let arr = a.to_array().unwrap();
assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
}
#[test]
fn to_array_non_1d_fails() {
let a = NdArray::new(&[2, 3]);
assert!(a.to_array().is_err());
}
#[test]
fn concat_1d() {
let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
let b = NdArray::from_slice(&[3.0, 4.0, 5.0], &[3]);
let c = a.concat(b).unwrap();
assert_eq!(c.shape(), &[5]);
let vals: Vec<f64> = (&c).into_iter().collect();
assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
}
#[test]
fn concat_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]);
let c = a.concat(b).unwrap();
assert_eq!(c.shape(), &[5, 2]);
assert_eq!(c.col(0), &[1.0, 2.0, 5.0, 6.0, 7.0]);
assert_eq!(c.col(1), &[3.0, 4.0, 8.0, 9.0, 10.0]);
}
#[test]
fn concat_dimension_mismatch_fails() {
let a = NdArray::<f64>::new(&[3, 2]);
let b = NdArray::new(&[3, 3]);
assert!(a.concat(b).is_err());
}
#[test]
fn concat_rank_mismatch_fails() {
let a = NdArray::<f64>::new(&[3]);
let b = NdArray::new(&[3, 2]);
assert!(a.concat(b).is_err());
}
#[test]
fn clone_and_eq() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = a.clone();
assert_eq!(a, b);
}
#[test]
fn ne_different_data() {
let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
let b = NdArray::from_slice(&[1.0, 3.0], &[2]);
assert_ne!(a, b);
}
#[test]
fn ne_different_shape() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
let b = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
assert_ne!(a, b);
}
#[test]
fn debug_1d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
let s = format!("{:?}", a);
assert!(s.contains("[3]"));
assert!(s.contains("1D"));
}
#[test]
fn debug_2d_named() {
let a = NdArray::<f64>::new_named(&[2, 3], "test");
let s = format!("{:?}", a);
assert!(s.contains("'test'"));
assert!(s.contains("[2, 3]"));
assert!(s.contains("2D"));
}
#[test]
fn empty_array() {
let a = NdArray::new(&[0, 5]);
assert!(a.is_empty());
assert_eq!(a.len(), 0);
let vals: Vec<f64> = (&a).into_iter().collect();
assert!(vals.is_empty());
}
#[test]
fn single_element() {
let a = NdArray::from_slice(&[42.0], &[1]);
assert_eq!(a.len(), 1);
assert_eq!(a[(0,)], 42.0);
let vals: Vec<f64> = (&a).into_iter().collect();
assert_eq!(vals, vec![42.0]);
}
#[test]
fn single_element_2d() {
let a = NdArray::from_slice(&[42.0], &[1, 1]);
assert_eq!(a[(0, 0)], 42.0);
}
#[test]
fn large_array_iteration_count() {
let n = 1000;
let a = NdArray::<f64>::ones(&[n, 100]);
let count = (&a).into_iter().count();
assert_eq!(count, n * 100);
}
#[test]
fn bracket_index_1d() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
assert_eq!(a[0], [10.0]);
assert_eq!(a[2], [30.0]);
}
#[test]
fn bracket_index_2d_column() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(a[0], [1.0, 2.0, 3.0]);
assert_eq!(a[1], [4.0, 5.0, 6.0]);
}
#[test]
fn bracket_index_2d_chained() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(a[0][0], 1.0);
assert_eq!(a[0][2], 3.0);
assert_eq!(a[1][0], 4.0);
assert_eq!(a[1][2], 6.0);
}
#[test]
fn bracket_index_mut_2d() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
a[0][1] = 99.0;
assert_eq!(a[0][1], 99.0);
assert_eq!(a[0], [1.0, 99.0, 3.0]);
}
#[test]
fn range_index_1d() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
assert_eq!(a[1..3], [20.0, 30.0]);
}
#[test]
fn range_index_2d_columns() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let slab = &a[0..2];
assert_eq!(slab, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
let col1 = &a[1..2];
assert_eq!(col1, &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "matrix")]
#[test]
#[should_panic(expected = "contiguous")]
fn range_index_non_contiguous_panics() {
let mat = Matrix::from_f64_unaligned(
&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
);
let a = NdArray::from(mat);
if a.is_contiguous() {
panic!("premise failed: expected non-contiguous import");
}
let _ = &a[0..2];
}
#[cfg(feature = "matrix")]
#[test]
fn to_matrix_repacks_compact_layout() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let mat = a.to_matrix().unwrap();
assert_eq!(mat.n_rows, 3);
assert_eq!(mat.n_cols, 2);
assert_eq!(mat.stride, 8);
assert_eq!(&mat.data.as_slice()[..3], &[1.0, 2.0, 3.0]);
assert_eq!(&mat.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
}
#[cfg(feature = "matrix")]
#[test]
fn matrix_roundtrip_via_contiguous() {
let mat = Matrix::from_f64_unaligned(
&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
);
let a = NdArray::from(mat);
assert!(!a.is_contiguous());
assert_eq!(a.get(&[2, 1]), 6.0);
let compact = a.to_contiguous();
assert!(compact.is_contiguous());
assert_eq!(&compact[0..2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
let back = compact.to_matrix().unwrap();
assert_eq!(back.stride, 8);
assert_eq!(&back.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
}
#[test]
fn range_from_index() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
assert_eq!(a[2..], [30.0, 40.0]);
}
#[test]
fn range_to_index() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
assert_eq!(a[..2], [10.0, 20.0]);
}
#[test]
fn range_full_index() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
assert_eq!(a[..], [10.0, 20.0, 30.0]);
}
#[cfg(feature = "views")]
#[test]
fn slice_1d_single_index() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
let v = a.slice(&[&1]);
assert!(v.shape().is_empty());
assert_eq!(v[()], 20.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_1d_range() {
let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
let v = a.slice(&[&(1..3)]);
assert_eq!(v.shape(), &[2]);
assert_eq!(v[(0,)], 20.0);
assert_eq!(v[(1,)], 30.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_2d_row_range_single_col() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.slice(nd![0..2, 1]);
assert_eq!(v.shape(), &[2]);
assert_eq!(v[(0,)], 4.0);
assert_eq!(v[(1,)], 5.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_2d_single_row_col_range() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.slice(nd![1, 0..2]);
assert_eq!(v.shape(), &[2]);
let vals: Vec<f64> = (&v).into_iter().collect();
assert_eq!(vals, vec![2.0, 5.0]);
}
#[cfg(feature = "views")]
#[test]
fn slice_2d_both_ranges() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.slice(nd![0..2, 0..2]);
assert_eq!(v.shape(), &[2, 2]);
assert_eq!(v[(0, 0)], 1.0);
assert_eq!(v[(1, 0)], 2.0);
assert_eq!(v[(0, 1)], 4.0);
assert_eq!(v[(1, 1)], 5.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_2d_both_indices_scalar() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.slice(nd![2, 1]);
assert!(v.shape().is_empty());
assert_eq!(v[()], 6.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_3d_mixed() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
let v = a.slice(nd![0..2, 1, 0..2]);
assert_eq!(v.shape(), &[2, 2]);
assert_eq!(v[(0, 0)], 3.0);
assert_eq!(v[(1, 0)], 4.0);
assert_eq!(v[(0, 1)], 9.0);
assert_eq!(v[(1, 1)], 10.0);
}
#[cfg(feature = "views")]
#[test]
fn slice_with_nd_macro() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
let v = a.slice(nd![0..2, 0..3, 0..4]);
assert_eq!(v.shape(), &[2, 3, 4]);
assert_eq!(v.len(), 24);
}
#[cfg(feature = "views")]
#[test]
fn slice_preserves_data_through_iteration() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let v = a.slice(nd![1..3, 0..2]);
let vals: Vec<f64> = (&v).into_iter().collect();
assert_eq!(vals, vec![2.0, 3.0, 5.0, 6.0]);
}
#[cfg(feature = "views")]
#[test]
fn slice_column_window() {
let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[10, 2]);
let v = a.slice(nd![2..5, 1]);
assert_eq!(v.shape(), &[3]);
assert_eq!(v[(0,)], 13.0);
assert_eq!(v[(1,)], 14.0);
assert_eq!(v[(2,)], 15.0);
}
#[test]
#[should_panic(expected = "Column index out of bounds")]
fn col_out_of_bounds_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let _ = a.col(2);
}
#[test]
#[should_panic(expected = "indices for a 2D array")]
fn get_rank_mismatch_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let _ = a.get(&[0]);
}
#[test]
#[should_panic(expected = "out of bounds for dim")]
fn get_index_out_of_bounds_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let _ = a.get(&[2, 0]);
}
#[test]
#[should_panic(expected = "out of bounds for dim")]
fn set_out_of_bounds_panics() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
a.set(&[0, 5], 9.0);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
#[should_panic(expected = "expected 2 axes, got 1")]
fn slice_wrong_axis_count_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let _ = a.slice(nd![0..2]);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
#[should_panic(expected = "range 0..100 out of bounds")]
fn slice_range_out_of_bounds_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let _ = a.slice(nd![0..100, 0..2]);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
#[should_panic(expected = "index 5 out of bounds")]
fn slice_single_index_out_of_bounds_panics() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let _ = a.slice(nd![5, 0..2]);
}
#[test]
#[should_panic(expected = "at least 2 points")]
fn linspace_requires_two_points() {
let _ = NdArray::<f64>::linspace(0.0, 1.0, 1);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn slice_one_element_index_array_collapses() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
let idx: &[usize] = &[1];
let v = a.slice(nd![idx, 0..2]);
assert_eq!(v.shape(), &[2]);
assert_eq!(v[(0,)], 2.0);
assert_eq!(v[(1,)], 5.0);
}
#[test]
fn get_unchecked_matches_get() {
let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
let a = NdArray::from_slice(&data, &[2, 3, 4]);
for i in 0..2 {
for j in 0..3 {
for k in 0..4 {
let idx = [i, j, k];
assert_eq!(unsafe { a.get_unchecked(&idx) }, a.get(&idx));
}
}
}
}
#[test]
fn concat_3d_interleaves_axis0() {
let da: Vec<f64> = (1..=8).map(|x| x as f64).collect();
let db: Vec<f64> = (9..=16).map(|x| x as f64).collect();
let a = NdArray::from_slice(&da, &[2, 2, 2]);
let b = NdArray::from_slice(&db, &[2, 2, 2]);
let c = a.clone().concat(b.clone()).unwrap();
assert_eq!(c.shape(), &[4, 2, 2]);
for i in 0..4 {
for j in 0..2 {
for k in 0..2 {
let expected = if i < 2 {
a.get(&[i, j, k])
} else {
b.get(&[i - 2, j, k])
};
assert_eq!(c.get(&[i, j, k]), expected);
}
}
}
}
#[test]
fn concat_non_contiguous_operand() {
let a = NdArray::from_buffer(
Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]),
&[2, 2],
&[2, 1],
);
assert!(!a.is_contiguous());
let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]);
let c = a.concat(b).unwrap();
assert_eq!(c.shape(), &[4, 2]);
assert_eq!(c.col(0), &[1.0, 3.0, 5.0, 6.0]);
assert_eq!(c.col(1), &[2.0, 4.0, 7.0, 8.0]);
}
#[cfg(feature = "views")]
#[test]
fn set_after_view_copy_on_write() {
let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let v = a.as_view();
a.set(&[0, 0], 99.0);
assert_eq!(a.get(&[0, 0]), 99.0);
assert_eq!(v.get(&[0, 0]), 1.0);
}
#[test]
fn fill_with_non_contiguous_logical_only() {
let mut buf = Vec64::with_capacity(11);
buf.0.resize(11, 0.0);
let mut a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
assert!(!a.is_contiguous());
a.fill_with(7.0);
for v in &a { assert_eq!(v, 7.0); }
assert_eq!(a.as_slice()[3], 0.0);
assert_eq!(a.as_slice()[7], 0.0);
}
#[test]
fn range_from_index_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(&a[1..], &[4.0, 5.0, 6.0]);
}
#[test]
fn range_to_index_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(&a[..1], &[1.0, 2.0, 3.0]);
}
#[test]
fn range_full_index_2d() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(&a[..], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[cfg(all(feature = "views", feature = "select"))]
#[test]
fn row_selection_single_index_window() {
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert_eq!(a.n_obs(), 3);
let v = a.r(1usize);
assert_eq!(v.shape(), &[1, 2]);
assert_eq!(v.get(&[0, 0]), 2.0);
assert_eq!(v.get(&[0, 1]), 5.0);
assert_eq!(v.source.len(), a.len());
assert_eq!(unsafe { v.get_unchecked(&[0, 1]) }, 5.0);
}
#[cfg(feature = "size")]
#[test]
fn est_bytes_covers_buffer() {
use crate::traits::byte_size::ByteSize;
let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
assert!(a.est_bytes() >= 6 * size_of::<f64>());
}
}