1use std::fmt;
68use std::ops::{Index, IndexMut, Range, RangeFrom, RangeFull, RangeTo};
69use std::sync::Arc;
70
71use crate::enums::error::MinarrowError;
72use crate::enums::shape_dim::ShapeDim;
73use crate::structs::buffer::Buffer;
74#[cfg(all(feature = "views", feature = "select"))]
75use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
76use crate::traits::type_unions::Float;
77use crate::traits::{concatenate::Concatenate, shape::Shape};
78use crate::{Array, ArrowType, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
79
80#[cfg(feature = "matrix")]
81use crate::structs::matrix::{Matrix, aligned_stride};
82#[cfg(feature = "views")]
83use crate::structs::views::ndarray_view::NdArrayV;
84#[cfg(feature = "dlpack")]
85use crate::ffi::dlpack::{
86 export_to_dlpack, export_to_dlpack_versioned, DLPackTensor, DLPackTensorVersioned,
87};
88
89#[derive(Clone)]
100pub struct NdArray<T> {
101 pub(crate) data: Arc<Buffer<T>>,
102 pub(crate) dims: NdDims,
103 pub name: Option<String>,
104}
105
106impl<T: Float> PartialEq for NdArray<T> {
109 fn eq(&self, other: &Self) -> bool {
110 if self.dims.shape() != other.dims.shape() {
111 return false;
112 }
113 self.into_iter()
114 .zip(other.into_iter())
115 .all(|(a, b)| a == b)
116 }
117}
118
119impl<T: Float> NdArray<T> {
122 pub fn new(shape: &[usize]) -> Self {
126 let dims = NdDims::from_shape(shape);
127 let total = buffer_len(dims.shape(), dims.strides());
128 let mut v = Vec64::with_capacity(total);
129 v.0.resize(total, T::default());
130 NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
131 }
132
133 pub fn new_named(shape: &[usize], name: impl Into<String>) -> Self {
135 let mut arr = Self::new(shape);
136 arr.name = Some(name.into());
137 arr
138 }
139
140 pub fn from_slice(data: &[T], shape: &[usize]) -> Self {
148 let logical_len: usize = shape.iter().product();
149 assert_eq!(
150 data.len(), logical_len,
151 "NdArray::from_slice: data length {} does not match shape product {}",
152 data.len(), logical_len
153 );
154 let dims = NdDims::from_shape(shape);
155 NdArray {
156 data: Arc::new(Buffer::from_slice(data)),
157 dims,
158 name: None,
159 }
160 }
161
162 pub fn from_vec64(data: Vec64<T>, shape: &[usize]) -> Self {
167 let logical_len: usize = shape.iter().product();
168 assert_eq!(
169 data.len(), logical_len,
170 "NdArray::from_vec64: data length {} does not match shape product {}",
171 data.len(), logical_len
172 );
173 let dims = NdDims::from_shape(shape);
174 NdArray {
175 data: Arc::new(Buffer::from_vec64(data)),
176 dims,
177 name: None,
178 }
179 }
180
181 pub fn from_buffer(data: Buffer<T>, shape: &[usize], strides: &[usize]) -> Self {
186 let required = buffer_len(shape, strides);
187 assert!(
188 data.len() >= required,
189 "NdArray::from_buffer: buffer has {} elements but shape requires {}",
190 data.len(), required
191 );
192 let dims = NdDims::from_shape_and_strides(shape, strides);
193 NdArray { data: Arc::new(data), dims, name: None }
194 }
195
196 pub fn fill(shape: &[usize], value: T) -> Self {
198 let dims = NdDims::from_shape(shape);
199 let total = buffer_len(dims.shape(), dims.strides());
200 let mut v = Vec64::with_capacity(total);
201 v.0.resize(total, value);
202 NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
203 }
204
205 pub fn ones(shape: &[usize]) -> Self {
207 Self::fill(shape, T::one())
208 }
209
210 pub fn eye(n: usize) -> Self {
212 let mut arr = Self::new(&[n, n]);
213 let stride = arr.dims.strides()[1];
214 let buf = Arc::make_mut(&mut arr.data).as_mut_slice();
215 for i in 0..n {
216 buf[i * stride + i] = T::one();
217 }
218 arr
219 }
220
221 pub fn linspace(start: T, end: T, n: usize) -> Self {
223 assert!(n >= 2, "linspace requires at least 2 points");
224 let step = (end - start) / T::from(n - 1).unwrap();
225 let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
226 NdArray {
227 data: Arc::new(Buffer::from_vec64(v)),
228 dims: NdDims::from_shape(&[n]),
229 name: None,
230 }
231 }
232
233 pub fn arange(start: T, step: T, n: usize) -> Self {
236 let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
237 NdArray {
238 data: Arc::new(Buffer::from_vec64(v)),
239 dims: NdDims::from_shape(&[n]),
240 name: None,
241 }
242 }
243
244 #[inline]
248 pub fn ndim(&self) -> usize { self.dims.ndim() }
249
250 #[inline]
252 pub fn shape(&self) -> &[usize] { self.dims.shape() }
253
254 #[inline]
256 pub fn strides(&self) -> &[usize] { self.dims.strides() }
257
258 #[inline]
260 pub fn len(&self) -> usize { self.dims.len() }
261
262 #[inline]
266 pub fn n_obs(&self) -> usize {
267 assert!(self.ndim() > 0, "n_obs() requires an axis 0");
268 self.dims.shape()[0]
269 }
270
271 #[inline]
273 pub fn is_empty(&self) -> bool { self.len() == 0 }
274
275 pub fn is_contiguous(&self) -> bool {
278 let shape = self.dims.shape();
279 let strides = self.dims.strides();
280 let mut expected = 1;
281 for d in 0..shape.len() {
282 if strides[d] != expected {
283 return false;
284 }
285 expected *= shape[d];
286 }
287 true
288 }
289
290 pub fn has_nan(&self) -> bool {
292 self.into_iter().any(|v| v.is_nan())
293 }
294
295 pub fn nan_count(&self) -> usize {
297 self.into_iter().filter(|v| v.is_nan()).count()
298 }
299
300 #[inline]
304 pub fn get(&self, indices: &[usize]) -> T {
305 self.data.as_slice()[self.offset_of(indices)]
306 }
307
308 #[inline(always)]
313 pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
314 let strides = self.dims.strides();
315 let mut off = 0;
316 for d in 0..indices.len() {
317 off += indices[d] * strides[d];
318 }
319 unsafe { *self.data.as_slice().get_unchecked(off) }
321 }
322
323 #[inline]
336 pub fn set(&mut self, indices: &[usize], value: T) {
337 let off = self.offset_of(indices);
338 Arc::make_mut(&mut self.data).as_mut_slice()[off] = value;
339 }
340
341 #[inline]
343 pub(crate) fn offset_of(&self, indices: &[usize]) -> usize {
344 offset_of_impl(indices, self.dims.shape(), self.dims.strides())
345 }
346
347 #[inline]
351 pub fn set_name(&mut self, name: impl Into<String>) {
352 self.name = Some(name.into());
353 }
354
355 #[inline]
357 pub fn as_slice(&self) -> &[T] {
358 self.data.as_slice()
359 }
360
361 #[inline]
364 pub fn as_mut_slice(&mut self) -> &mut [T] {
365 Arc::make_mut(&mut self.data).as_mut_slice()
366 }
367
368 pub fn fill_with(&mut self, value: T) {
370 if self.is_contiguous() {
372 Arc::make_mut(&mut self.data).as_mut_slice().fill(value);
373 return;
374 }
375 let offsets: Vec<usize> = {
377 let shape = self.dims.shape();
378 let strides = self.dims.strides();
379 let mut result = Vec::with_capacity(self.len());
380 let mut indices = vec![0usize; shape.len()];
381 for _ in 0..self.len() {
382 result.push(indices.iter().zip(strides).map(|(&i, &s)| i * s).sum());
383 let mut carry = true;
384 for d in 0..shape.len() {
385 if carry {
386 indices[d] += 1;
387 if indices[d] < shape[d] { carry = false; } else { indices[d] = 0; }
388 }
389 }
390 }
391 result
392 };
393 let buf = Arc::make_mut(&mut self.data).as_mut_slice();
394 for off in offsets { buf[off] = value; }
395 }
396
397 #[inline]
403 pub fn col(&self, col: usize) -> &[T] {
404 let shape = self.dims.shape();
405 assert_eq!(shape.len(), 2, "col() requires a 2D array");
406 assert!(col < shape[1], "Column index out of bounds");
407 assert_eq!(
408 self.dims.strides()[0], 1,
409 "col() requires unit stride on axis 0; call to_contiguous() first"
410 );
411 let stride = self.dims.strides()[1];
412 let start = col * stride;
413 &self.data.as_slice()[start..start + shape[0]]
414 }
415
416 #[inline]
421 pub fn col_mut(&mut self, col: usize) -> &mut [T] {
422 let shape = self.dims.shape();
423 assert_eq!(shape.len(), 2, "col_mut() requires a 2D array");
424 assert!(col < shape[1], "Column index out of bounds");
425 assert_eq!(
426 self.dims.strides()[0], 1,
427 "col_mut() requires unit stride on axis 0; call to_contiguous() first"
428 );
429 let stride = self.dims.strides()[1];
430 let n_rows = shape[0];
431 let start = col * stride;
432 &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
433 }
434
435 pub fn columns(&self) -> Vec<&[T]> {
438 let shape = self.dims.shape();
439 assert_eq!(shape.len(), 2, "columns() requires a 2D array");
440 assert_eq!(
441 self.dims.strides()[0], 1,
442 "columns() requires unit stride on axis 0; call to_contiguous() first"
443 );
444 let stride = self.dims.strides()[1];
445 let n_rows = shape[0];
446 let buf = self.data.as_slice();
447 (0..shape[1])
448 .map(|c| &buf[c * stride..c * stride + n_rows])
449 .collect()
450 }
451
452 pub fn columns_mut(&mut self) -> Vec<&mut [T]> {
455 let shape = self.dims.shape();
456 assert_eq!(shape.len(), 2, "columns_mut() requires a 2D array");
457 assert_eq!(
458 self.dims.strides()[0], 1,
459 "columns_mut() requires unit stride on axis 0; call to_contiguous() first"
460 );
461 let n_rows = shape[0];
462 let n_cols = shape[1];
463 let stride = self.dims.strides()[1];
464 let ptr = Arc::make_mut(&mut self.data).as_mut_slice().as_mut_ptr();
465 let mut result = Vec::with_capacity(n_cols);
466 for c in 0..n_cols {
467 let start = c * stride;
468 unsafe {
471 let col_ptr = ptr.add(start);
472 result.push(std::slice::from_raw_parts_mut(col_ptr, n_rows));
473 }
474 }
475 result
476 }
477
478 #[cfg(feature = "views")]
484 pub fn as_view(&self) -> NdArrayV<T> {
485 NdArrayV::from_ndarray(self.clone())
486 }
487
488 #[cfg(feature = "views")]
498 pub fn obs(&self, idx: usize) -> NdArrayV<T> {
499 self.as_view().obs(idx)
500 }
501
502 #[inline]
506 pub fn m(&self) -> i32 {
507 assert_eq!(self.ndim(), 2, "m() requires a 2D array");
508 self.dims.shape()[0] as i32
509 }
510
511 #[inline]
513 pub fn n(&self) -> i32 {
514 assert_eq!(self.ndim(), 2, "n() requires a 2D array");
515 self.dims.shape()[1] as i32
516 }
517
518 #[inline]
520 pub fn lda(&self) -> i32 {
521 assert_eq!(self.ndim(), 2, "lda() requires a 2D array");
522 self.dims.strides()[1] as i32
523 }
524
525 pub fn reshape(&self, new_shape: &[usize]) -> Result<NdArray<T>, MinarrowError> {
532 let new_len: usize = new_shape.iter().product();
533 if new_len != self.len() {
534 return Err(MinarrowError::ShapeError {
535 message: format!(
536 "reshape: cannot reshape array of size {} into shape {:?}",
537 self.len(), new_shape
538 ),
539 });
540 }
541 let flat: Vec64<T> = self.into_iter().collect();
542 let mut result = NdArray::from_slice(&flat, new_shape);
543 result.name = self.name.clone();
544 Ok(result)
545 }
546
547 pub fn transpose(&self) -> NdArray<T> {
553 let shape = self.dims.shape();
554 let ndim = shape.len();
555 if ndim <= 1 {
556 let mut result = self.to_contiguous();
557 result.name = self.name.clone();
558 return result;
559 }
560 if ndim == 2 {
561 let (n_rows, n_cols) = (shape[0], shape[1]);
562 let new_shape = [n_cols, n_rows];
563 let mut result = NdArray::new(&new_shape);
564 let src_stride = self.dims.strides()[1];
565 let dst_stride = result.dims.strides()[1];
566 let src = self.data.as_slice();
567 let dst = Arc::make_mut(&mut result.data).as_mut_slice();
568 for r in 0..n_rows {
569 for c in 0..n_cols {
570 dst[r * dst_stride + c] = src[c * src_stride + r];
571 }
572 }
573 result.name = self.name.clone();
574 return result;
575 }
576
577 let new_shape: Vec<usize> = shape.iter().rev().copied().collect();
581 let rev_strides: Vec<usize> = self.dims.strides().iter().rev().copied().collect();
582 let src = self.data.as_slice();
583 let total = self.len();
584 let mut buf = Vec64::with_capacity(total);
585 let mut indices = vec![0usize; ndim];
586 for _ in 0..total {
587 let offset: usize = indices.iter()
588 .zip(rev_strides.iter())
589 .map(|(&i, &s)| i * s)
590 .sum();
591 buf.push(src[offset]);
592 let mut carry = true;
593 for d in 0..ndim {
594 if carry {
595 indices[d] += 1;
596 if indices[d] < new_shape[d] {
597 carry = false;
598 } else {
599 indices[d] = 0;
600 }
601 }
602 }
603 }
604 NdArray {
605 data: Arc::new(Buffer::from_vec64(buf)),
606 dims: NdDims::from_shape(&new_shape),
607 name: self.name.clone(),
608 }
609 }
610
611 pub fn flatten(&self) -> NdArray<T> {
613 let flat: Vec64<T> = self.into_iter().collect();
614 let n = flat.len();
615 NdArray {
616 data: Arc::new(Buffer::from_vec64(flat)),
617 dims: NdDims::from_shape(&[n]),
618 name: self.name.clone(),
619 }
620 }
621
622 pub fn to_contiguous(&self) -> NdArray<T> {
625 if self.is_contiguous() {
626 return self.clone();
627 }
628 let mut result = NdArray::from_slice(
629 &self.into_iter().collect::<Vec64<T>>(),
630 self.shape(),
631 );
632 result.name = self.name.clone();
633 result
634 }
635
636 #[cfg(all(feature = "views", feature = "select"))]
655 pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
656 assert_eq!(
657 selection.len(), self.ndim(),
658 "slice(): expected {} axes, got {}", self.ndim(), selection.len()
659 );
660
661 let shape = self.dims.shape();
662 let strides = self.dims.strides();
663
664 let mut new_offset: usize = 0;
666 let mut new_shape = Vec::with_capacity(self.ndim());
667 let mut new_strides = Vec::with_capacity(self.ndim());
668
669 for (d, sel) in selection.iter().enumerate() {
670 let (start, end, collapse) = sel.resolve_axis(shape[d]);
671 assert!(
672 end <= shape[d],
673 "slice(): end {} out of bounds for axis {} (size {})", end, d, shape[d]
674 );
675 new_offset += start * strides[d];
676 if !collapse {
677 new_shape.push(end - start);
678 new_strides.push(strides[d]);
679 }
680 }
681
682 NdArrayV::new(self.clone(), new_offset, &new_shape, &new_strides)
683 }
684
685 pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray<T> {
691 let flat: Vec64<T> = self.into_iter().map(f).collect();
692 let mut result = NdArray::from_slice(&flat, self.shape());
693 result.name = self.name.clone();
694 result
695 }
696
697 pub fn apply_mut(&mut self, f: impl Fn(T) -> T) {
701 if self.is_contiguous() {
702 for v in Arc::make_mut(&mut self.data).as_mut_slice() {
703 *v = f(*v);
704 }
705 return;
706 }
707 let shape = self.dims.shape().to_vec();
710 let strides = self.dims.strides().to_vec();
711 let total = self.len();
712 let buf = Arc::make_mut(&mut self.data).as_mut_slice();
713 let mut indices = vec![0usize; shape.len()];
714 for _ in 0..total {
715 let offset: usize = indices.iter().zip(strides.iter()).map(|(&i, &s)| i * s).sum();
716 buf[offset] = f(buf[offset]);
717 let mut carry = true;
718 for d in 0..shape.len() {
719 if carry {
720 indices[d] += 1;
721 if indices[d] < shape[d] {
722 carry = false;
723 } else {
724 indices[d] = 0;
725 }
726 }
727 }
728 }
729 }
730
731 #[cfg(feature = "views")]
736 pub fn apply_axis(&self, axis: usize, mut f: impl FnMut(NdArrayV<T>) -> T) -> NdArray<T> {
737 let shape = self.dims.shape();
738 let strides = self.dims.strides();
739 let ndim = shape.len();
740 assert!(ndim >= 2, "apply_axis requires a 2D or higher array");
741 assert!(axis < ndim, "apply_axis: axis {} out of bounds for {}D array", axis, ndim);
742
743 let out_shape: Vec<usize> = shape
744 .iter()
745 .enumerate()
746 .filter(|(d, _)| *d != axis)
747 .map(|(_, &s)| s)
748 .collect();
749 let out_dims: Vec<usize> = (0..ndim).filter(|&d| d != axis).collect();
750 let total: usize = out_shape.iter().product();
751
752 let lane_shape = [shape[axis]];
753 let lane_strides = [strides[axis]];
754
755 let mut flat = Vec64::with_capacity(total);
758 let mut indices = vec![0usize; out_shape.len()];
759 for _ in 0..total {
760 let offset: usize = indices
761 .iter()
762 .zip(out_dims.iter())
763 .map(|(&i, &d)| i * strides[d])
764 .sum();
765 let lane = NdArrayV::new(self.clone(), offset, &lane_shape, &lane_strides);
766 flat.push(f(lane));
767 let mut carry = true;
768 for d in 0..out_shape.len() {
769 if carry {
770 indices[d] += 1;
771 if indices[d] < out_shape[d] {
772 carry = false;
773 } else {
774 indices[d] = 0;
775 }
776 }
777 }
778 }
779 let mut result = NdArray::from_slice(&flat, &out_shape);
780 result.name = self.name.clone();
781 result
782 }
783
784 #[cfg(feature = "dlpack")]
794 pub fn to_dlpack(self) -> DLPackTensor {
795 export_to_dlpack(self)
796 }
797
798 #[cfg(feature = "dlpack")]
802 pub fn to_dlpack_versioned(self) -> DLPackTensorVersioned {
803 export_to_dlpack_versioned(self)
804 }
805
806 #[cfg(feature = "parallel_proc")]
811 pub fn par_iter(&self) -> rayon::slice::Iter<'_, T>
812 where
813 T: Send + Sync,
814 {
815 use rayon::prelude::*;
816 self.data.as_slice().par_iter()
817 }
818
819 #[cfg(feature = "parallel_proc")]
822 pub fn par_chunks(&self, chunk_size: usize) -> rayon::slice::Chunks<'_, T>
823 where
824 T: Send + Sync,
825 {
826 use rayon::prelude::*;
827 self.data.as_slice().par_chunks(chunk_size)
828 }
829
830 #[cfg(feature = "views")]
833 pub fn iter_obs(&self) -> impl Iterator<Item = (usize, NdArrayV<T>)> + '_ {
834 assert!(self.ndim() >= 2, "iter_obs() requires a 2D or higher array");
835 let n_obs = self.shape()[0];
836 (0..n_obs).map(move |i| (i, self.obs(i)))
837 }
838
839 #[cfg(all(feature = "parallel_proc", feature = "views"))]
842 pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator<Item = (usize, NdArrayV<T>)> + '_
843 where
844 T: Send + Sync,
845 {
846 use rayon::prelude::*;
847 assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher array");
848 let n_obs = self.dims.shape()[0];
849 (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i)))
850 }
851
852 pub(crate) fn iter_axis0_run(&self, run_idx: usize) -> impl ExactSizeIterator<Item = T> + '_ {
856 assert!(self.ndim() > 0, "axis-0 iteration requires an axis 0");
857 let n_runs: usize = self.shape()[1..].iter().product();
858 assert!(run_idx < n_runs, "axis-0 run {} out of bounds ({})", run_idx, n_runs);
859
860 let mut rem = run_idx;
861 let mut offset = 0;
862 for d in 1..self.ndim() {
863 offset += (rem % self.shape()[d]) * self.strides()[d];
864 rem /= self.shape()[d];
865 }
866 let stride = self.strides()[0];
867 (0..self.shape()[0]).map(move |i| self.data.as_slice()[offset + i * stride])
868 }
869}
870
871impl NdArray<f64> {
874 pub fn to_table(self, fields: Option<Vec<Field>>) -> Result<Table, MinarrowError> {
883 let shape = self.dims.shape();
884 if shape.len() != 2 {
885 return Err(MinarrowError::ShapeError {
886 message: format!("to_table requires a 2D array, got {}D", shape.len()),
887 });
888 }
889 let n_cols = shape[1];
890 let fields = match fields {
891 Some(fields) => {
892 if fields.len() != n_cols {
893 return Err(MinarrowError::ShapeError {
894 message: format!(
895 "to_table: expected {} fields for {} columns, got {}",
896 n_cols, n_cols, fields.len()
897 ),
898 });
899 }
900 for (index, field) in fields.iter().enumerate() {
901 if field.dtype != ArrowType::Float64 {
902 return Err(MinarrowError::TypeError {
903 from: "Field",
904 to: "Float64 NdArray column",
905 message: Some(format!(
906 "to_table: field {} ('{}') has dtype {:?}, expected Float64",
907 index, field.name, field.dtype
908 )),
909 });
910 }
911 if field.nullable {
912 return Err(MinarrowError::NullError {
913 message: Some(format!(
914 "to_table: field {} ('{}') is nullable, but NdArray columns are non-nullable",
915 index, field.name
916 )),
917 });
918 }
919 }
920 fields
921 }
922 None => (0..n_cols)
923 .map(|i| Field::new(format!("col_{}", i), ArrowType::Float64, false, None))
924 .collect(),
925 };
926 let this = if self.dims.strides()[0] == 1 { self } else { self.to_contiguous() };
929 let n_rows = this.dims.shape()[0];
930 let stride = this.dims.strides()[1];
931 let name = this.name;
932 let buf = this.data.as_slice();
933
934 let mut cols = Vec::with_capacity(n_cols);
935 for (i, field) in fields.into_iter().enumerate() {
936 let col_start = i * stride;
937 let col: Buffer<f64> = Buffer::from_slice(&buf[col_start..col_start + n_rows]);
938 let float_arr = FloatArray::new(col, None);
939 let array = Array::NumericArray(NumericArray::Float64(Arc::new(float_arr)));
940 cols.push(FieldArray::new(field, array));
941 }
942
943 Ok(Table::new(name.unwrap_or_default(), Some(cols)))
944 }
945
946 pub fn to_array(self) -> Result<Array, MinarrowError> {
950 if self.ndim() != 1 {
951 return Err(MinarrowError::ShapeError {
952 message: format!("to_array requires a 1D array, got {}D", self.ndim()),
953 });
954 }
955 let buffer = if self.dims.strides()[0] == 1 && self.data.len() == self.len() {
956 Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone())
957 } else {
958 Buffer::from_vec64((&self).into_iter().collect())
959 };
960 let float_arr = FloatArray::new(buffer, None);
961 Ok(Array::NumericArray(NumericArray::Float64(Arc::new(float_arr))))
962 }
963
964 #[cfg(feature = "matrix")]
971 pub fn to_matrix(self) -> Result<Matrix, MinarrowError> {
972 let shape = self.dims.shape();
973 if shape.len() != 2 {
974 return Err(MinarrowError::ShapeError {
975 message: format!("to_matrix requires a 2D array, got {}D", shape.len()),
976 });
977 }
978 let n_rows = shape[0];
979 let n_cols = shape[1];
980 let strides = self.dims.strides();
981
982 if strides[0] == 1 && strides[1] == aligned_stride(n_rows) {
983 return Ok(Matrix {
984 n_rows,
985 n_cols,
986 stride: strides[1],
987 data: Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone()),
988 name: self.name,
989 });
990 }
991
992 let name = self.name.clone();
993 let compact: Vec64<f64> = self.into_iter().collect();
994 Ok(Matrix::from_f64_unaligned(&compact, n_rows, n_cols, name))
995 }
996}
997
998#[derive(Clone, PartialEq)]
1007pub(crate) enum NdDims {
1008 D1 { shape: [usize; 1], strides: [usize; 1] },
1009 D2 { shape: [usize; 2], strides: [usize; 2] },
1010 D3 { shape: [usize; 3], strides: [usize; 3] },
1011 D4 { shape: [usize; 4], strides: [usize; 4] },
1012 D5 { shape: [usize; 5], strides: [usize; 5] },
1013 Dn { shape: Box<[usize]>, strides: Box<[usize]> },
1014}
1015
1016impl NdDims {
1017 pub(crate) fn from_shape(shape: &[usize]) -> Self {
1020 let strides = col_major_strides(shape);
1021 Self::from_shape_and_strides(shape, &strides)
1022 }
1023
1024 pub(crate) fn from_shape_and_strides(shape: &[usize], strides: &[usize]) -> Self {
1026 assert_eq!(
1027 shape.len(),
1028 strides.len(),
1029 "NdArray: shape rank {} does not match strides rank {}",
1030 shape.len(),
1031 strides.len()
1032 );
1033 match shape.len() {
1034 1 => NdDims::D1 {
1035 shape: [shape[0]],
1036 strides: [strides[0]],
1037 },
1038 2 => NdDims::D2 {
1039 shape: [shape[0], shape[1]],
1040 strides: [strides[0], strides[1]],
1041 },
1042 3 => NdDims::D3 {
1043 shape: [shape[0], shape[1], shape[2]],
1044 strides: [strides[0], strides[1], strides[2]],
1045 },
1046 4 => NdDims::D4 {
1047 shape: [shape[0], shape[1], shape[2], shape[3]],
1048 strides: [strides[0], strides[1], strides[2], strides[3]],
1049 },
1050 5 => NdDims::D5 {
1051 shape: [shape[0], shape[1], shape[2], shape[3], shape[4]],
1052 strides: [strides[0], strides[1], strides[2], strides[3], strides[4]],
1053 },
1054 _ => NdDims::Dn {
1055 shape: shape.into(),
1056 strides: strides.into(),
1057 },
1058 }
1059 }
1060
1061 #[inline]
1063 pub(crate) fn ndim(&self) -> usize {
1064 match self {
1065 NdDims::D1 { .. } => 1,
1066 NdDims::D2 { .. } => 2,
1067 NdDims::D3 { .. } => 3,
1068 NdDims::D4 { .. } => 4,
1069 NdDims::D5 { .. } => 5,
1070 NdDims::Dn { shape, .. } => shape.len(),
1071 }
1072 }
1073
1074 #[inline]
1076 pub(crate) fn shape(&self) -> &[usize] {
1077 match self {
1078 NdDims::D1 { shape, .. } => shape,
1079 NdDims::D2 { shape, .. } => shape,
1080 NdDims::D3 { shape, .. } => shape,
1081 NdDims::D4 { shape, .. } => shape,
1082 NdDims::D5 { shape, .. } => shape,
1083 NdDims::Dn { shape, .. } => shape,
1084 }
1085 }
1086
1087 #[inline]
1089 pub(crate) fn strides(&self) -> &[usize] {
1090 match self {
1091 NdDims::D1 { strides, .. } => strides,
1092 NdDims::D2 { strides, .. } => strides,
1093 NdDims::D3 { strides, .. } => strides,
1094 NdDims::D4 { strides, .. } => strides,
1095 NdDims::D5 { strides, .. } => strides,
1096 NdDims::Dn { strides, .. } => strides,
1097 }
1098 }
1099
1100 #[inline]
1102 pub(crate) fn len(&self) -> usize {
1103 self.shape().iter().product()
1104 }
1105}
1106
1107#[macro_export]
1121macro_rules! nd {
1122 ($($sel:expr),+ $(,)?) => {
1123 &[$(&$sel as &dyn $crate::traits::selection::DataSelector),+]
1124 };
1125}
1126
1127pub(crate) fn col_major_strides(shape: &[usize]) -> Vec<usize> {
1138 let mut strides = Vec::with_capacity(shape.len());
1139 if shape.is_empty() {
1140 return strides;
1141 }
1142 strides.push(1);
1143 for d in 1..shape.len() {
1144 strides.push(strides[d - 1] * shape[d - 1]);
1145 }
1146 strides
1147}
1148
1149pub(crate) fn buffer_len(shape: &[usize], strides: &[usize]) -> usize {
1151 if shape.iter().any(|&d| d == 0) {
1152 return 0;
1153 }
1154 let max_offset: usize = shape.iter()
1157 .zip(strides.iter())
1158 .map(|(&s, &st)| (s - 1) * st)
1159 .sum();
1160 max_offset + 1
1161}
1162
1163impl<'a, T: Float> IntoIterator for &'a NdArray<T> {
1172 type Item = T;
1173 type IntoIter = NdArrayIter<'a, T>;
1174
1175 #[inline]
1176 fn into_iter(self) -> NdArrayIter<'a, T> {
1177 let shape = self.dims.shape();
1178 let strides = self.dims.strides();
1179 if shape.is_empty() {
1180 return NdArrayIter {
1181 buf: self.data.as_slice(),
1182 n_inner: 1,
1183 inner_stride: 1,
1184 run_offsets: vec![0],
1185 run_idx: 0,
1186 inner_idx: 0,
1187 total: 1,
1188 yielded: 0,
1189 };
1190 }
1191 let n_inner = shape[0];
1192
1193 let n_runs: usize = shape[1..].iter().product();
1195
1196 let mut run_offsets = Vec::with_capacity(n_runs);
1201 if shape.len() <= 1 {
1202 run_offsets.push(0);
1203 } else {
1204 let outer_shape = &shape[1..];
1205 let outer_strides = &strides[1..];
1206 let mut outer_indices = vec![0usize; outer_shape.len()];
1207 for _ in 0..n_runs {
1208 let off: usize = outer_indices.iter()
1209 .zip(outer_strides.iter())
1210 .map(|(&i, &s)| i * s)
1211 .sum();
1212 run_offsets.push(off);
1213 let mut carry = true;
1215 for d in 0..outer_shape.len() {
1216 if carry {
1217 outer_indices[d] += 1;
1218 if outer_indices[d] < outer_shape[d] {
1219 carry = false;
1220 } else {
1221 outer_indices[d] = 0;
1222 }
1223 }
1224 }
1225 }
1226 }
1227
1228 NdArrayIter {
1229 buf: self.data.as_slice(),
1230 n_inner,
1231 inner_stride: strides[0],
1232 run_offsets,
1233 run_idx: 0,
1234 inner_idx: 0,
1235 total: self.len(),
1236 yielded: 0,
1237 }
1238 }
1239}
1240
1241impl<T: Float> IntoIterator for NdArray<T> {
1243 type Item = T;
1244 type IntoIter = std::vec::IntoIter<T>;
1245
1246 #[inline]
1247 fn into_iter(self) -> std::vec::IntoIter<T> {
1248 let v: Vec<T> = (&self).into_iter().collect();
1249 v.into_iter()
1250 }
1251}
1252
1253pub struct NdArrayIter<'a, T> {
1260 pub(crate) buf: &'a [T],
1261 pub(crate) n_inner: usize,
1262 pub(crate) inner_stride: usize,
1263 pub(crate) run_offsets: Vec<usize>,
1264 pub(crate) run_idx: usize,
1265 pub(crate) inner_idx: usize,
1266 pub(crate) total: usize,
1267 pub(crate) yielded: usize,
1268}
1269
1270impl<'a, T: Float> Iterator for NdArrayIter<'a, T> {
1271 type Item = T;
1272
1273 #[inline]
1274 fn next(&mut self) -> Option<T> {
1275 if self.yielded >= self.total { return None; }
1276
1277 let val = self.buf[self.run_offsets[self.run_idx] + self.inner_idx * self.inner_stride];
1278 self.yielded += 1;
1279 self.inner_idx += 1;
1280 if self.inner_idx >= self.n_inner {
1281 self.inner_idx = 0;
1282 self.run_idx += 1;
1283 }
1284 Some(val)
1285 }
1286
1287 #[inline]
1288 fn size_hint(&self) -> (usize, Option<usize>) {
1289 let r = self.total - self.yielded;
1290 (r, Some(r))
1291 }
1292}
1293
1294impl<'a, T: Float> ExactSizeIterator for NdArrayIter<'a, T> {}
1295
1296#[inline]
1305pub(crate) fn offset_of_impl(indices: &[usize], shape: &[usize], strides: &[usize]) -> usize {
1306 assert_eq!(
1307 indices.len(),
1308 shape.len(),
1309 "NdArray: {} indices for a {}D array",
1310 indices.len(),
1311 shape.len()
1312 );
1313 let mut offset = 0;
1314 for d in 0..shape.len() {
1315 assert!(
1316 indices[d] < shape[d],
1317 "NdArray: index {} out of bounds for dim {} (size {})",
1318 indices[d], d, shape[d]
1319 );
1320 offset += indices[d] * strides[d];
1321 }
1322 offset
1323}
1324
1325impl<T: Float> Shape for NdArray<T> {
1330 fn shape(&self) -> ShapeDim {
1331 match self.dims.ndim() {
1332 0 => ShapeDim::Rank0(1),
1333 1 => ShapeDim::Rank1(self.dims.shape()[0]),
1334 2 => ShapeDim::Rank2 {
1335 rows: self.dims.shape()[0],
1336 cols: self.dims.shape()[1],
1337 },
1338 _ => ShapeDim::RankN(self.dims.shape().to_vec()),
1339 }
1340 }
1341}
1342
1343impl<T: Float> Concatenate for NdArray<T> {
1344 fn concat(self, other: Self) -> Result<Self, MinarrowError> {
1346 let s1 = self.dims.shape();
1347 let s2 = other.dims.shape();
1348 if s1.len() != s2.len() {
1349 return Err(MinarrowError::IncompatibleTypeError {
1350 from: "NdArray",
1351 to: "NdArray",
1352 message: Some(format!(
1353 "Cannot concatenate {}D and {}D arrays", s1.len(), s2.len()
1354 )),
1355 });
1356 }
1357 if s1.is_empty() {
1358 return Err(MinarrowError::ShapeError {
1359 message: "Cannot concatenate rank-zero arrays along axis 0".to_string(),
1360 });
1361 }
1362 for d in 1..s1.len() {
1363 if s1[d] != s2[d] {
1364 return Err(MinarrowError::IncompatibleTypeError {
1365 from: "NdArray",
1366 to: "NdArray",
1367 message: Some(format!(
1368 "Dimension {} mismatch: {} vs {}", d, s1[d], s2[d]
1369 )),
1370 });
1371 }
1372 }
1373
1374 let mut new_shape: Vec<usize> = s1.to_vec();
1375 new_shape[0] += s2[0];
1376
1377 let name = match (&self.name, &other.name) {
1379 (Some(a), Some(b)) => Some(format!("{}+{}", a, b)),
1380 (Some(a), None) => Some(a.clone()),
1381 (None, Some(b)) => Some(b.clone()),
1382 (None, None) => None,
1383 };
1384
1385 if s1.len() == 1 {
1387 let mut flat = Vec64::with_capacity(self.len() + other.len());
1388 flat.extend(&self);
1389 flat.extend(&other);
1390 let mut result = NdArray::from_slice(&flat, &new_shape);
1391 result.name = name;
1392 return Ok(result);
1393 }
1394
1395 if s1.len() == 2 && self.dims.strides()[0] == 1 && other.dims.strides()[0] == 1 {
1397 let new_dims = NdDims::from_shape(&new_shape);
1398 let new_stride = new_dims.strides()[1];
1399 let total = buffer_len(&new_shape, new_dims.strides());
1400 let mut buf = Vec64::with_capacity(total);
1401 buf.0.resize(total, T::default());
1402 let dst = buf.as_mut_slice();
1403 for c in 0..s1[1] {
1404 let dst_start = c * new_stride;
1405 dst[dst_start..dst_start + s1[0]].copy_from_slice(self.col(c));
1406 dst[dst_start + s1[0]..dst_start + s1[0] + s2[0]].copy_from_slice(other.col(c));
1407 }
1408 return Ok(NdArray {
1409 data: Arc::new(Buffer::from_vec64(buf)),
1410 dims: new_dims,
1411 name,
1412 });
1413 }
1414
1415 let n1 = s1[0];
1418 let n2 = s2[0];
1419 let run_len = n1 + n2;
1420 let n_runs: usize = s1[1..].iter().product();
1421 let new_dims = NdDims::from_shape(&new_shape);
1422 let total = buffer_len(&new_shape, new_dims.strides());
1423 let mut buf = Vec64::with_capacity(total);
1424 buf.0.resize(total, T::default());
1425 let dst = buf.as_mut_slice();
1426 let mut it_a = (&self).into_iter();
1427 let mut it_b = (&other).into_iter();
1428 for r in 0..n_runs {
1429 let base = r * run_len;
1430 for i in 0..n1 {
1431 dst[base + i] = it_a.next().unwrap();
1432 }
1433 for i in 0..n2 {
1434 dst[base + n1 + i] = it_b.next().unwrap();
1435 }
1436 }
1437 Ok(NdArray {
1438 data: Arc::new(Buffer::from_vec64(buf)),
1439 dims: new_dims,
1440 name,
1441 })
1442 }
1443}
1444
1445#[cfg(all(feature = "views", feature = "select"))]
1451impl<T: Float> AxisSelection for NdArray<T> {
1452 type View = NdArrayV<T>;
1453
1454 fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
1455 self.slice(selection)
1456 }
1457
1458 fn get_axis_count(&self) -> usize {
1459 self.ndim()
1460 }
1461}
1462
1463#[cfg(all(feature = "views", feature = "select"))]
1469impl<T: Float> RowSelection for NdArray<T> {
1470 type View = NdArrayV<T>;
1471
1472 fn r<S: DataSelector>(&self, selection: S) -> NdArrayV<T> {
1473 assert!(self.ndim() > 0, "row selection requires an axis 0");
1474 let n_obs = self.shape()[0];
1475 let indices = selection.resolve_indices(n_obs);
1476 if selection.is_contiguous() {
1477 let start = indices.first().copied().unwrap_or(0);
1478 let ranges: Vec<Range<usize>> = std::iter::once(start..start + indices.len())
1479 .chain(self.shape()[1..].iter().map(|&n| 0..n))
1480 .collect();
1481 let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect();
1482 return self.slice(&refs);
1483 }
1484 NdArrayV::from_ndarray(gather_obs_impl(
1485 &indices,
1486 self.shape(),
1487 self.name.clone(),
1488 |idx| self.get(idx),
1489 ))
1490 }
1491
1492 fn get_row_count(&self) -> usize {
1493 self.n_obs()
1494 }
1495}
1496
1497#[cfg(all(feature = "views", feature = "select"))]
1502pub(crate) fn gather_obs_impl<T: Float>(
1503 indices: &[usize],
1504 shape: &[usize],
1505 name: Option<String>,
1506 get: impl Fn(&[usize]) -> T,
1507) -> NdArray<T> {
1508 let mut out_shape = shape.to_vec();
1509 out_shape[0] = indices.len();
1510 let total: usize = out_shape.iter().product();
1511
1512 let mut flat = Vec64::with_capacity(total);
1513 let ndim = shape.len();
1514 let mut idx = vec![0usize; ndim];
1515 let inner_runs: usize = shape[1..].iter().product::<usize>().max(1);
1516 for _ in 0..inner_runs {
1517 for &obs in indices {
1518 idx[0] = obs;
1519 flat.push(get(&idx));
1520 }
1521 let mut carry = true;
1523 for d in 1..ndim {
1524 if carry {
1525 idx[d] += 1;
1526 if idx[d] < shape[d] {
1527 carry = false;
1528 } else {
1529 idx[d] = 0;
1530 }
1531 }
1532 }
1533 }
1534 let mut result = NdArray::from_slice(&flat, &out_shape);
1535 result.name = name;
1536 result
1537}
1538
1539impl<T: Float> Index<usize> for NdArray<T> {
1547 type Output = [T];
1548
1549 #[inline]
1550 fn index(&self, idx: usize) -> &[T] {
1551 let shape = self.dims.shape();
1552 let strides = self.dims.strides();
1553 match shape.len() {
1554 0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1555 1 => {
1556 assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
1557 &self.data.as_slice()[idx..idx + 1]
1558 }
1559 2 => {
1560 assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
1561 let start = idx * strides[1];
1562 &self.data.as_slice()[start..start + shape[0]]
1563 }
1564 n => {
1565 assert!(
1567 self.is_contiguous(),
1568 "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
1569 );
1570 let last = n - 1;
1571 assert!(idx < shape[last], "index out of bounds for axis {}", last);
1572 let start = idx * strides[last];
1573 &self.data.as_slice()[start..start + strides[last]]
1574 }
1575 }
1576 }
1577}
1578
1579impl<T: Float> IndexMut<usize> for NdArray<T> {
1580 #[inline]
1581 fn index_mut(&mut self, idx: usize) -> &mut [T] {
1582 let shape = self.dims.shape().to_vec();
1583 let strides = self.dims.strides().to_vec();
1584 match shape.len() {
1585 0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1586 1 => {
1587 assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
1588 &mut Arc::make_mut(&mut self.data).as_mut_slice()[idx..idx + 1]
1589 }
1590 2 => {
1591 assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
1592 let start = idx * strides[1];
1593 let n_rows = shape[0];
1594 &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
1595 }
1596 n => {
1597 assert!(
1598 self.is_contiguous(),
1599 "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
1600 );
1601 let last = n - 1;
1602 assert!(idx < shape[last], "index out of bounds for axis {}", last);
1603 let start = idx * strides[last];
1604 &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + strides[last]]
1605 }
1606 }
1607 }
1608}
1609
1610impl<T: Float> Index<Range<usize>> for NdArray<T> {
1620 type Output = [T];
1621
1622 #[inline]
1623 fn index(&self, range: Range<usize>) -> &[T] {
1624 let shape = self.dims.shape();
1625 let strides = self.dims.strides();
1626 match shape.len() {
1627 0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1628 1 => &self.data.as_slice()[range],
1629 _ => {
1630 assert!(
1631 self.is_contiguous(),
1632 "range indexing requires a contiguous layout, use slice() for strided access"
1633 );
1634 let last = shape.len() - 1;
1635 assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
1636 let start = range.start * strides[last];
1637 let end = range.end * strides[last];
1638 &self.data.as_slice()[start..end]
1639 }
1640 }
1641 }
1642}
1643
1644impl<T: Float> IndexMut<Range<usize>> for NdArray<T> {
1645 #[inline]
1646 fn index_mut(&mut self, range: Range<usize>) -> &mut [T] {
1647 let shape = self.dims.shape().to_vec();
1648 let strides = self.dims.strides().to_vec();
1649 match shape.len() {
1650 0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1651 1 => &mut Arc::make_mut(&mut self.data).as_mut_slice()[range],
1652 _ => {
1653 assert!(
1654 self.is_contiguous(),
1655 "range indexing requires a contiguous layout, use slice() for strided access"
1656 );
1657 let last = shape.len() - 1;
1658 assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
1659 let start = range.start * strides[last];
1660 let end = range.end * strides[last];
1661 &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..end]
1662 }
1663 }
1664 }
1665}
1666
1667impl<T: Float> Index<RangeFrom<usize>> for NdArray<T> {
1668 type Output = [T];
1669
1670 #[inline]
1671 fn index(&self, range: RangeFrom<usize>) -> &[T] {
1672 assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
1673 let last_dim = self.dims.shape().len() - 1;
1674 let end = self.dims.shape()[last_dim];
1675 &self[range.start..end]
1676 }
1677}
1678
1679impl<T: Float> Index<RangeTo<usize>> for NdArray<T> {
1680 type Output = [T];
1681
1682 #[inline]
1683 fn index(&self, range: RangeTo<usize>) -> &[T] {
1684 &self[0..range.end]
1685 }
1686}
1687
1688impl<T: Float> Index<RangeFull> for NdArray<T> {
1689 type Output = [T];
1690
1691 #[inline]
1692 fn index(&self, _: RangeFull) -> &[T] {
1693 assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
1694 let last_dim = self.dims.shape().len() - 1;
1695 let end = self.dims.shape()[last_dim];
1696 &self[0..end]
1697 }
1698}
1699
1700impl<T: Float> Index<()> for NdArray<T> {
1703 type Output = T;
1704 #[inline]
1705 fn index(&self, (): ()) -> &T {
1706 &self.data.as_slice()[self.offset_of(&[])]
1707 }
1708}
1709
1710impl<T: Float> Index<(usize,)> for NdArray<T> {
1711 type Output = T;
1712 #[inline]
1713 fn index(&self, (i,): (usize,)) -> &T {
1714 &self.data.as_slice()[self.offset_of(&[i])]
1715 }
1716}
1717
1718impl<T: Float> Index<(usize, usize)> for NdArray<T> {
1719 type Output = T;
1720 #[inline]
1721 fn index(&self, (i, j): (usize, usize)) -> &T {
1722 &self.data.as_slice()[self.offset_of(&[i, j])]
1723 }
1724}
1725
1726impl<T: Float> Index<(usize, usize, usize)> for NdArray<T> {
1727 type Output = T;
1728 #[inline]
1729 fn index(&self, (i, j, k): (usize, usize, usize)) -> &T {
1730 &self.data.as_slice()[self.offset_of(&[i, j, k])]
1731 }
1732}
1733
1734impl<T: Float> Index<(usize, usize, usize, usize)> for NdArray<T> {
1735 type Output = T;
1736 #[inline]
1737 fn index(&self, (i, j, k, l): (usize, usize, usize, usize)) -> &T {
1738 &self.data.as_slice()[self.offset_of(&[i, j, k, l])]
1739 }
1740}
1741
1742impl<T: Float> Index<(usize, usize, usize, usize, usize)> for NdArray<T> {
1743 type Output = T;
1744 #[inline]
1745 fn index(&self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &T {
1746 &self.data.as_slice()[self.offset_of(&[i, j, k, l, m])]
1747 }
1748}
1749
1750impl<T: Float> IndexMut<(usize,)> for NdArray<T> {
1751 #[inline]
1752 fn index_mut(&mut self, (i,): (usize,)) -> &mut T {
1753 let off = self.offset_of(&[i]);
1754 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1755 }
1756}
1757
1758impl<T: Float> IndexMut<()> for NdArray<T> {
1759 #[inline]
1760 fn index_mut(&mut self, (): ()) -> &mut T {
1761 let off = self.offset_of(&[]);
1762 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1763 }
1764}
1765
1766impl<T: Float> IndexMut<(usize, usize)> for NdArray<T> {
1767 #[inline]
1768 fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut T {
1769 let off = self.offset_of(&[i, j]);
1770 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1771 }
1772}
1773
1774impl<T: Float> IndexMut<(usize, usize, usize)> for NdArray<T> {
1775 #[inline]
1776 fn index_mut(&mut self, (i, j, k): (usize, usize, usize)) -> &mut T {
1777 let off = self.offset_of(&[i, j, k]);
1778 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1779 }
1780}
1781
1782impl<T: Float> IndexMut<(usize, usize, usize, usize)> for NdArray<T> {
1783 #[inline]
1784 fn index_mut(&mut self, (i, j, k, l): (usize, usize, usize, usize)) -> &mut T {
1785 let off = self.offset_of(&[i, j, k, l]);
1786 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1787 }
1788}
1789
1790impl<T: Float> IndexMut<(usize, usize, usize, usize, usize)> for NdArray<T> {
1791 #[inline]
1792 fn index_mut(&mut self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &mut T {
1793 let off = self.offset_of(&[i, j, k, l, m]);
1794 &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1795 }
1796}
1797
1798impl<T: Float> From<&[T]> for NdArray<T> {
1802 fn from(data: &[T]) -> Self {
1803 NdArray {
1804 data: Arc::new(Buffer::from_slice(data)),
1805 dims: NdDims::from_shape(&[data.len()]),
1806 name: None,
1807 }
1808 }
1809}
1810
1811impl<T: Float> From<Vec64<T>> for NdArray<T> {
1813 fn from(v: Vec64<T>) -> Self {
1814 let n = v.len();
1815 NdArray {
1816 data: Arc::new(Buffer::from_vec64(v)),
1817 dims: NdDims::from_shape(&[n]),
1818 name: None,
1819 }
1820 }
1821}
1822
1823impl<T: Float> From<&[Vec<T>]> for NdArray<T> {
1825 fn from(columns: &[Vec<T>]) -> Self {
1826 let n_cols = columns.len();
1827 if n_cols == 0 {
1828 return NdArray::new(&[0, 0]);
1829 }
1830 let n_rows = columns[0].len();
1831 for col in columns {
1832 assert_eq!(col.len(), n_rows, "Column length mismatch");
1833 }
1834 let shape = [n_rows, n_cols];
1835 let dims = NdDims::from_shape(&shape);
1836 let stride = dims.strides()[1];
1837 let total = buffer_len(&shape, dims.strides());
1838 let mut buf = Vec64::with_capacity(total);
1839 buf.0.resize(total, T::default());
1840 for (c, col) in columns.iter().enumerate() {
1841 let start = c * stride;
1842 buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col);
1843 }
1844 NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
1845 }
1846}
1847
1848impl<T: Float> From<&[FloatArray<T>]> for NdArray<T> {
1850 fn from(columns: &[FloatArray<T>]) -> Self {
1851 let n_cols = columns.len();
1852 if n_cols == 0 {
1853 return NdArray::new(&[0, 0]);
1854 }
1855 let n_rows = columns[0].data.len();
1856 for col in columns {
1857 assert_eq!(col.data.len(), n_rows, "Column length mismatch");
1858 }
1859 let shape = [n_rows, n_cols];
1860 let dims = NdDims::from_shape(&shape);
1861 let stride = dims.strides()[1];
1862 let total = buffer_len(&shape, dims.strides());
1863 let mut buf = Vec64::with_capacity(total);
1864 buf.0.resize(total, T::default());
1865 for (c, col) in columns.iter().enumerate() {
1866 let start = c * stride;
1867 buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col.data.as_slice());
1868 }
1869 NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
1870 }
1871}
1872
1873#[cfg(feature = "matrix")]
1877impl From<Matrix> for NdArray<f64> {
1878 fn from(mat: Matrix) -> Self {
1879 let shape = [mat.n_rows, mat.n_cols];
1880 let strides = [1, mat.stride];
1881 NdArray {
1882 data: Arc::new(mat.data),
1883 dims: NdDims::from_shape_and_strides(&shape, &strides),
1884 name: mat.name,
1885 }
1886 }
1887}
1888
1889impl TryFrom<&Table> for NdArray<f64> {
1891 type Error = MinarrowError;
1892
1893 fn try_from(table: &Table) -> Result<Self, Self::Error> {
1894 let n_cols = table.n_cols();
1895 let n_rows = table.n_rows;
1896 if n_cols == 0 {
1897 return Ok(NdArray::new(&[0, 0]));
1898 }
1899
1900 let shape = [n_rows, n_cols];
1901 let dims = NdDims::from_shape(&shape);
1902 let stride = dims.strides()[1];
1903 let total = buffer_len(&shape, dims.strides());
1904 let mut buf = Vec64::with_capacity(total);
1905 buf.0.resize(total, 0.0);
1906
1907 for (col_idx, fa) in table.cols.iter().enumerate() {
1908 let numeric = fa.array.try_num().map_err(|_| MinarrowError::TypeError {
1909 from: "non-numeric",
1910 to: "Float64",
1911 message: Some(format!("column {} is not numeric", col_idx)),
1912 })?;
1913 let f64_arr = numeric.try_f64()?;
1914 if f64_arr.data.len() != n_rows {
1915 return Err(MinarrowError::ColumnLengthMismatch {
1916 col: col_idx,
1917 expected: n_rows,
1918 found: f64_arr.data.len(),
1919 });
1920 }
1921
1922 let start = col_idx * stride;
1923 let src = f64_arr.data.as_slice();
1924 let dst = &mut buf.as_mut_slice()[start..start + n_rows];
1925
1926 match f64_arr.null_mask.as_ref() {
1928 Some(mask) => {
1929 for i in 0..n_rows {
1930 dst[i] = if mask.get(i) { src[i] } else { f64::NAN };
1931 }
1932 }
1933 None => dst.copy_from_slice(src),
1934 }
1935 }
1936
1937 let name = if table.name.is_empty() { None } else { Some(table.name.clone()) };
1938 Ok(NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name })
1939 }
1940}
1941
1942impl TryFrom<Table> for NdArray<f64> {
1943 type Error = MinarrowError;
1944 fn try_from(table: Table) -> Result<Self, Self::Error> {
1945 NdArray::try_from(&table)
1946 }
1947}
1948
1949impl<T: Float> fmt::Debug for NdArray<T> {
1952 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1953 write!(
1954 f, "NdArray{}: {:?} [{}D, col-major]",
1955 self.name.as_deref().map_or(String::new(), |n| format!(" '{}'", n)),
1956 self.dims.shape(),
1957 self.ndim(),
1958 )?;
1959 if self.ndim() == 0 {
1960 write!(f, "\n{:8.4}", self.get(&[]).to_f64().unwrap_or(f64::NAN))?;
1961 } else if self.ndim() == 2 {
1962 let shape = self.dims.shape();
1963 let max_rows = shape[0].min(6);
1964 let max_cols = shape[1].min(8);
1965 for r in 0..max_rows {
1966 write!(f, "\n[")?;
1967 for c in 0..max_cols {
1968 write!(f, " {:8.4}", self.get(&[r, c]).to_f64().unwrap_or(f64::NAN))?;
1969 if c < max_cols - 1 { write!(f, ",")?; }
1970 }
1971 if shape[1] > 8 { write!(f, " ...")?; }
1972 write!(f, " ]")?;
1973 }
1974 if shape[0] > 6 { write!(f, "\n...")?; }
1975 } else if self.ndim() == 1 {
1976 let n = self.dims.shape()[0].min(10);
1977 write!(f, "\n[")?;
1978 for i in 0..n {
1979 write!(f, " {:8.4}", self.get(&[i]).to_f64().unwrap_or(f64::NAN))?;
1980 if i < n - 1 { write!(f, ",")?; }
1981 }
1982 if self.dims.shape()[0] > 10 { write!(f, " ...")?; }
1983 write!(f, " ]")?;
1984 }
1985 Ok(())
1986 }
1987}
1988
1989#[cfg(test)]
1990mod tests {
1991 use super::*;
1992 use crate::StringArray;
1993 use crate::structs::bitmask::Bitmask;
1994
1995 #[cfg(all(feature = "views", feature = "select"))]
1998 #[test]
1999 fn axis_selection_rank10() {
2000 let shape = [3, 4, 2, 5, 3, 2, 4, 3, 2, 5];
2001 let len: usize = shape.iter().product();
2002 let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
2003 let a = NdArray::from_slice(&data, &shape);
2004
2005 let v = a.s(nd![1..3, 0..2, 1, 2..5, .., 0..1, 1..3, 2, 0..2, 3..5]);
2007 assert_eq!(v.shape(), &[2, 2, 3, 3, 1, 2, 2, 2]);
2008 assert_eq!(
2009 v.get(&[0, 0, 0, 0, 0, 0, 0, 0]),
2010 a.get(&[1, 0, 1, 2, 0, 0, 1, 2, 0, 3])
2011 );
2012 assert_eq!(
2013 v.get(&[1, 1, 2, 2, 0, 1, 1, 1]),
2014 a.get(&[2, 1, 1, 4, 2, 0, 2, 2, 1, 4])
2015 );
2016 }
2017
2018 #[cfg(all(feature = "views", feature = "select"))]
2019 #[test]
2020 fn axis_selection_runtime_rank() {
2021 let mut shape = vec![1usize; 100];
2023 shape[0] = 3;
2024 shape[10] = 4;
2025 shape[50] = 5;
2026 shape[99] = 2;
2027 let len: usize = shape.iter().product();
2028 let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
2029 let a = NdArray::from_slice(&data, &shape);
2030
2031 let mut sels: Vec<Box<dyn DataSelector>> =
2033 shape.iter().map(|&n| Box::new(0..n) as Box<dyn DataSelector>).collect();
2034 sels[0] = Box::new(1..3);
2035 sels[10] = Box::new(2usize);
2036 sels[50] = Box::new(1..4);
2037 let refs: Vec<&dyn DataSelector> = sels.iter().map(|s| s.as_ref()).collect();
2038
2039 let v = a.s(&refs);
2040 assert_eq!(v.ndim(), 99);
2041 assert_eq!(v.shape()[0], 2);
2042 assert_eq!(v.shape()[49], 3);
2043 assert_eq!(v.shape()[98], 2);
2044
2045 let mut view_idx = vec![0usize; 99];
2046 view_idx[0] = 1;
2047 view_idx[49] = 2;
2048 view_idx[98] = 1;
2049 let mut source_idx = vec![0usize; 100];
2050 source_idx[0] = 2;
2051 source_idx[10] = 2;
2052 source_idx[50] = 3;
2053 source_idx[99] = 1;
2054 assert_eq!(v.get(&view_idx), a.get(&source_idx));
2055 }
2056
2057 #[cfg(all(feature = "views", feature = "select"))]
2058 #[test]
2059 fn axis_selection_trait() {
2060 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2061 let v = a.s(nd![1..3, 1]);
2063 assert_eq!(v.shape(), &[2]);
2064 assert_eq!(v.get(&[0]), 5.0);
2065 assert_eq!(v.get(&[1]), 6.0);
2066 let sub = a.s(nd![0..3, 0..2]).s(nd![2, 0..2]);
2068 assert_eq!(sub.shape(), &[2]);
2069 assert_eq!(sub.get(&[0]), 3.0);
2070 assert_eq!(sub.get(&[1]), 6.0);
2071 assert_eq!(a.get_axis_count(), 2);
2072 assert_eq!(sub.get_axis_count(), 1);
2073 }
2074
2075 #[cfg(all(feature = "views", feature = "select"))]
2076 #[test]
2077 fn row_selection_contiguous() {
2078 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2079 let v = a.r(1..3);
2080 assert_eq!(v.shape(), &[2, 2]);
2081 assert_eq!(v.get(&[0, 0]), 2.0);
2082 assert_eq!(v.get(&[1, 1]), 6.0);
2083 let single = a.row(2);
2085 assert_eq!(single.shape(), &[1, 2]);
2086 assert_eq!(single.get(&[0, 1]), 6.0);
2087 }
2088
2089 #[cfg(all(feature = "views", feature = "select"))]
2090 #[test]
2091 fn row_selection_gathers_indices() {
2092 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2093 let v = a.r(&[2, 0]);
2094 assert_eq!(v.shape(), &[2, 2]);
2095 assert_eq!(v.get(&[0, 0]), 3.0);
2097 assert_eq!(v.get(&[0, 1]), 6.0);
2098 assert_eq!(v.get(&[1, 0]), 1.0);
2099 assert_eq!(v.get(&[1, 1]), 4.0);
2100 }
2101
2102 #[test]
2103 fn apply_maps_elements() {
2104 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2105 a.set_name("m");
2106 let b = a.apply(|x| x * 10.0);
2107 assert_eq!(b.get(&[1, 1]), 40.0);
2108 assert_eq!(b.name.as_deref(), Some("m"));
2109 assert_eq!(a.get(&[1, 1]), 4.0);
2111 }
2112
2113 #[test]
2114 fn apply_mut_in_place() {
2115 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2116 a.apply_mut(|x| x + 0.5);
2117 assert_eq!((&a).into_iter().collect::<Vec<f64>>(), vec![1.5, 2.5, 3.5]);
2118 }
2119
2120 #[cfg(feature = "matrix")]
2121 #[test]
2122 fn apply_mut_non_contiguous_touches_logical_only() {
2123 let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None);
2126 let mut a = NdArray::from(mat);
2127 assert!(!a.is_contiguous());
2128 a.apply_mut(|x| x * 2.0);
2129 assert_eq!(a.get(&[0, 0]), 2.0);
2130 assert_eq!(a.get(&[2, 1]), 12.0);
2131 }
2132
2133 #[cfg(feature = "views")]
2134 #[test]
2135 fn apply_axis_collapses_axis() {
2136 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2137 let col_sums = a.apply_axis(0, |lane| (&lane).into_iter().sum());
2139 assert_eq!(col_sums.shape(), &[2]);
2140 assert_eq!(col_sums.get(&[0]), 6.0);
2141 assert_eq!(col_sums.get(&[1]), 15.0);
2142 let row_sums = a.apply_axis(1, |lane| (&lane).into_iter().sum());
2144 assert_eq!(row_sums.shape(), &[3]);
2145 assert_eq!(row_sums.get(&[0]), 5.0);
2146 assert_eq!(row_sums.get(&[2]), 9.0);
2147 }
2148
2149 #[cfg(feature = "views")]
2150 #[test]
2151 fn apply_axis_3d() {
2152 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2153 let a = NdArray::from_slice(&data, &[2, 3, 4]);
2154 let maxes = a.apply_axis(1, |lane| {
2155 (&lane).into_iter().fold(f64::MIN, f64::max)
2156 });
2157 assert_eq!(maxes.shape(), &[2, 4]);
2158 assert_eq!(maxes.get(&[0, 0]), a.get(&[0, 2, 0]));
2160 assert_eq!(maxes.get(&[1, 3]), a.get(&[1, 2, 3]));
2161 }
2162
2163 #[test]
2168 fn new_zeroed_1d() {
2169 let a = NdArray::<f64>::new(&[5]);
2170 assert_eq!(a.ndim(), 1);
2171 assert_eq!(a.shape(), &[5]);
2172 assert_eq!(a.len(), 5);
2173 assert!(!a.is_empty());
2174 for v in &a { assert_eq!(v, 0.0); }
2175 }
2176
2177 #[test]
2178 fn new_zeroed_2d() {
2179 let a = NdArray::<f64>::new(&[3, 4]);
2180 assert_eq!(a.ndim(), 2);
2181 assert_eq!(a.shape(), &[3, 4]);
2182 assert_eq!(a.len(), 12);
2183 for v in &a { assert_eq!(v, 0.0); }
2184 }
2185
2186 #[test]
2187 fn new_zeroed_3d() {
2188 let a = NdArray::<f64>::new(&[2, 3, 4]);
2189 assert_eq!(a.ndim(), 3);
2190 assert_eq!(a.shape(), &[2, 3, 4]);
2191 assert_eq!(a.len(), 24);
2192 }
2193
2194 #[test]
2195 fn new_zeroed_5d() {
2196 let a = NdArray::<f64>::new(&[2, 3, 4, 5, 6]);
2197 assert_eq!(a.ndim(), 5);
2198 assert_eq!(a.len(), 720);
2199 }
2200
2201 #[test]
2202 fn new_zeroed_6d_heap() {
2203 let a = NdArray::<f64>::new(&[2, 3, 2, 2, 2, 2]);
2204 assert_eq!(a.ndim(), 6);
2205 assert_eq!(a.len(), 96);
2206 }
2207
2208 #[test]
2209 fn new_named() {
2210 let a = NdArray::<f64>::new_named(&[3, 3], "covariance");
2211 assert_eq!(a.name.as_deref(), Some("covariance"));
2212 }
2213
2214 #[test]
2215 fn from_slice_1d() {
2216 let data = [1.0, 2.0, 3.0, 4.0, 5.0];
2217 let a = NdArray::from_slice(&data, &[5]);
2218 assert_eq!(a.len(), 5);
2219 let vals: Vec<f64> = (&a).into_iter().collect();
2220 assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
2221 }
2222
2223 #[test]
2224 fn from_slice_2d_column_major() {
2225 let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2226 let a = NdArray::from_slice(&data, &[3, 2]);
2227 assert_eq!(a.get(&[0, 0]), 1.0);
2228 assert_eq!(a.get(&[1, 0]), 2.0);
2229 assert_eq!(a.get(&[2, 0]), 3.0);
2230 assert_eq!(a.get(&[0, 1]), 4.0);
2231 assert_eq!(a.get(&[1, 1]), 5.0);
2232 assert_eq!(a.get(&[2, 1]), 6.0);
2233 }
2234
2235 #[test]
2236 fn fill_and_ones() {
2237 let a = NdArray::fill(&[2, 3], 7.0);
2238 assert_eq!(a.len(), 6);
2239 for v in &a { assert_eq!(v, 7.0); }
2240
2241 let b = NdArray::<f64>::ones(&[4]);
2242 for v in &b { assert_eq!(v, 1.0); }
2243 }
2244
2245 #[test]
2246 fn eye_identity() {
2247 let a = NdArray::<f64>::eye(3);
2248 assert_eq!(a.shape(), &[3, 3]);
2249 assert_eq!(a.get(&[0, 0]), 1.0);
2250 assert_eq!(a.get(&[1, 1]), 1.0);
2251 assert_eq!(a.get(&[2, 2]), 1.0);
2252 assert_eq!(a.get(&[0, 1]), 0.0);
2253 assert_eq!(a.get(&[1, 0]), 0.0);
2254 }
2255
2256 #[test]
2257 fn linspace_basic() {
2258 let a = NdArray::<f64>::linspace(0.0, 1.0, 5);
2259 assert_eq!(a.shape(), &[5]);
2260 assert_eq!(a.get(&[0]), 0.0);
2261 assert_eq!(a.get(&[4]), 1.0);
2262 assert!((a.get(&[2]) - 0.5).abs() < 1e-15);
2263 }
2264
2265 #[test]
2266 fn arange_basic() {
2267 let a = NdArray::arange(0.0, 0.5, 4);
2268 assert_eq!(a.shape(), &[4]);
2269 assert_eq!(a.get(&[0]), 0.0);
2270 assert_eq!(a.get(&[1]), 0.5);
2271 assert_eq!(a.get(&[2]), 1.0);
2272 assert_eq!(a.get(&[3]), 1.5);
2273 }
2274
2275 #[test]
2276 fn from_vec64_moves_data() {
2277 let data: Vec64<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into();
2278 let a = NdArray::from_vec64(data, &[3, 2]);
2279 assert_eq!(a.shape(), &[3, 2]);
2280 assert_eq!(a.get(&[2, 1]), 6.0);
2281 }
2282
2283 #[cfg(feature = "views")]
2284 #[test]
2285 fn iter_obs_walks_observations() {
2286 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2287 let rows: Vec<(usize, Vec<f64>)> = a
2288 .iter_obs()
2289 .map(|(i, v)| (i, (&v).into_iter().collect()))
2290 .collect();
2291 assert_eq!(rows.len(), 3);
2292 assert_eq!(rows[0], (0, vec![1.0, 4.0]));
2293 assert_eq!(rows[2], (2, vec![3.0, 6.0]));
2294 }
2295
2296 #[test]
2297 fn equality_ignores_name() {
2298 let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2299 let mut b = a.clone();
2300 b.name = Some("named".to_string());
2301 assert_eq!(a, b);
2302 }
2303
2304 #[test]
2305 fn name_survives_reshape_and_concat() {
2306 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
2307 a.name = Some("px".to_string());
2308 assert_eq!(a.reshape(&[2, 2]).unwrap().name.as_deref(), Some("px"));
2309
2310 let mut b = NdArray::from_slice(&[5.0, 6.0], &[2]);
2311 b.name = Some("qty".to_string());
2312 let joined = a.concat(b).unwrap();
2313 assert_eq!(joined.name.as_deref(), Some("px+qty"));
2314 }
2315
2316 #[test]
2321 fn get_set_1d() {
2322 let mut a = NdArray::new(&[3]);
2323 a.set(&[0], 10.0);
2324 a.set(&[1], 20.0);
2325 a.set(&[2], 30.0);
2326 assert_eq!(a.get(&[0]), 10.0);
2327 assert_eq!(a.get(&[1]), 20.0);
2328 assert_eq!(a.get(&[2]), 30.0);
2329 }
2330
2331 #[test]
2332 fn tuple_index_1d() {
2333 let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
2334 assert_eq!(a[(0,)], 10.0);
2335 assert_eq!(a[(1,)], 20.0);
2336 assert_eq!(a[(2,)], 30.0);
2337 }
2338
2339 #[test]
2340 fn tuple_index_2d() {
2341 let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2342 let a = NdArray::from_slice(&data, &[3, 2]);
2343 assert_eq!(a[(0, 0)], 1.0);
2344 assert_eq!(a[(2, 0)], 3.0);
2345 assert_eq!(a[(0, 1)], 4.0);
2346 assert_eq!(a[(2, 1)], 6.0);
2347 }
2348
2349 #[test]
2350 fn tuple_index_3d() {
2351 let data: Vec<f64> = (1..=12).map(|x| x as f64).collect();
2352 let a = NdArray::from_slice(&data, &[2, 3, 2]);
2353 assert_eq!(a[(0, 0, 0)], 1.0);
2354 assert_eq!(a[(1, 0, 0)], 2.0);
2355 assert_eq!(a[(0, 1, 0)], 3.0);
2356 assert_eq!(a[(1, 1, 0)], 4.0);
2357 assert_eq!(a[(0, 2, 0)], 5.0);
2358 assert_eq!(a[(1, 2, 0)], 6.0);
2359 assert_eq!(a[(0, 0, 1)], 7.0);
2360 assert_eq!(a[(1, 2, 1)], 12.0);
2361 }
2362
2363 #[test]
2364 fn index_mut_2d() {
2365 let mut a = NdArray::new(&[2, 2]);
2366 a[(0, 0)] = 1.0;
2367 a[(1, 0)] = 2.0;
2368 a[(0, 1)] = 3.0;
2369 a[(1, 1)] = 4.0;
2370 assert_eq!(a[(0, 0)], 1.0);
2371 assert_eq!(a[(1, 1)], 4.0);
2372 }
2373
2374 #[test]
2379 fn iter_1d() {
2380 let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
2381 let vals: Vec<f64> = (&a).into_iter().collect();
2382 assert_eq!(vals, vec![10.0, 20.0, 30.0]);
2383 }
2384
2385 #[test]
2386 fn iter_2d_column_major_order() {
2387 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2388 let vals: Vec<f64> = (&a).into_iter().collect();
2389 assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
2390 }
2391
2392 #[test]
2393 fn iter_2d_compact_layout() {
2394 let data: Vec<f64> = (1..=30).map(|x| x as f64).collect();
2397 let a = NdArray::from_slice(&data, &[10, 3]);
2398 assert_eq!(a.strides()[1], 10);
2399 let vals: Vec<f64> = (&a).into_iter().collect();
2400 assert_eq!(vals.len(), 30);
2401 assert_eq!(&vals[..10], &data[..10]);
2402 assert_eq!(&vals[10..20], &data[10..20]);
2403 assert_eq!(&vals[20..30], &data[20..30]);
2404 }
2405
2406 #[test]
2407 fn iter_3d_column_major_order() {
2408 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2409 let a = NdArray::from_slice(&data, &[2, 3, 4]);
2410 let vals: Vec<f64> = (&a).into_iter().collect();
2411 assert_eq!(vals.len(), 24);
2412 assert_eq!(vals, data);
2413 }
2414
2415 #[test]
2416 fn iter_exact_size() {
2417 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2418 let iter = (&a).into_iter();
2419 assert_eq!(iter.len(), 6);
2420 }
2421
2422 #[test]
2423 fn consuming_into_iter() {
2424 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2425 let vals: Vec<f64> = a.into_iter().collect();
2426 assert_eq!(vals, vec![1.0, 2.0, 3.0]);
2427 }
2428
2429 #[test]
2434 fn rank_zero_scalar_semantics() {
2435 let mut a = NdArray::from_slice(&[5.0], &[]);
2436 assert_eq!(a.ndim(), 0);
2437 assert!(a.shape().is_empty());
2438 assert!(a.strides().is_empty());
2439 assert_eq!(a.len(), 1);
2440 assert!(!a.is_empty());
2441 assert!(a.is_contiguous());
2442 assert_eq!(a[()], 5.0);
2443 assert_eq!((&a).into_iter().collect::<Vec<_>>(), vec![5.0]);
2444 assert_eq!(Shape::shape(&a), ShapeDim::Rank0(1));
2445
2446 a[()] = 7.0;
2447 assert_eq!(a[()], 7.0);
2448 assert!(a.transpose().shape().is_empty());
2449 assert_eq!(a.reshape(&[1]).unwrap().get(&[0]), 7.0);
2450 }
2451
2452 #[cfg(all(feature = "views", feature = "select"))]
2453 #[test]
2454 fn selecting_every_axis_can_produce_a_scalar_view() {
2455 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2456 let scalar = a.slice(&[&1usize, &1usize]);
2457 assert!(scalar.shape().is_empty());
2458 assert!(scalar.strides().is_empty());
2459 assert_eq!(scalar.len(), 1);
2460 assert_eq!(scalar.get(&[]), 4.0);
2461 assert_eq!((&scalar).into_iter().collect::<Vec<_>>(), vec![4.0]);
2462 assert_eq!(Shape::shape(&scalar), ShapeDim::Rank0(1));
2463 }
2464
2465 #[test]
2466 fn is_contiguous_default() {
2467 let a = NdArray::<f64>::new(&[3, 4]);
2468 assert!(a.is_contiguous());
2469 }
2470
2471 #[test]
2472 fn shape_trait_1d() {
2473 let a = NdArray::<f64>::new(&[5]);
2474 assert_eq!(Shape::shape(&a), ShapeDim::Rank1(5));
2475 }
2476
2477 #[test]
2478 fn shape_trait_2d() {
2479 let a = NdArray::<f64>::new(&[3, 4]);
2480 assert_eq!(Shape::shape(&a), ShapeDim::Rank2 { rows: 3, cols: 4 });
2481 }
2482
2483 #[test]
2484 fn shape_trait_3d() {
2485 let a = NdArray::<f64>::new(&[2, 3, 4]);
2486 assert_eq!(Shape::shape(&a), ShapeDim::RankN(vec![2, 3, 4]));
2487 }
2488
2489 #[test]
2494 fn has_nan_false() {
2495 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2496 assert!(!a.has_nan());
2497 assert_eq!(a.nan_count(), 0);
2498 }
2499
2500 #[test]
2501 fn has_nan_true() {
2502 let a = NdArray::from_slice(&[1.0, f64::NAN, 3.0], &[3]);
2503 assert!(a.has_nan());
2504 assert_eq!(a.nan_count(), 1);
2505 }
2506
2507 #[test]
2508 fn has_nan_2d() {
2509 let a = NdArray::from_slice(&[1.0, 2.0, f64::NAN, 4.0, 5.0, f64::NAN], &[3, 2]);
2510 assert!(a.has_nan());
2511 assert_eq!(a.nan_count(), 2);
2512 }
2513
2514 #[test]
2519 fn col_access() {
2520 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2521 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2522 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2523 }
2524
2525 #[test]
2526 fn col_mut_access() {
2527 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2528 a.col_mut(0)[1] = 99.0;
2529 assert_eq!(a.col(0), &[1.0, 99.0, 3.0]);
2530 }
2531
2532 #[test]
2533 fn columns_access() {
2534 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2535 let cols = a.columns();
2536 assert_eq!(cols.len(), 2);
2537 assert_eq!(cols[0], &[1.0, 2.0, 3.0]);
2538 assert_eq!(cols[1], &[4.0, 5.0, 6.0]);
2539 }
2540
2541 #[test]
2542 fn columns_mut_access() {
2543 let mut a = NdArray::new(&[3, 2]);
2544 {
2545 let mut cols = a.columns_mut();
2546 cols[0].copy_from_slice(&[1.0, 2.0, 3.0]);
2547 cols[1].copy_from_slice(&[4.0, 5.0, 6.0]);
2548 }
2549 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2550 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2551 }
2552
2553 #[cfg(feature = "views")]
2554 #[test]
2555 fn obs_access() {
2556 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2557 let v = a.as_view();
2558 let obs0: Vec<f64> = (&v.obs(0)).into_iter().collect();
2559 let obs1: Vec<f64> = (&v.obs(1)).into_iter().collect();
2560 let obs2: Vec<f64> = (&v.obs(2)).into_iter().collect();
2561 assert_eq!(obs0, vec![1.0, 4.0]);
2562 assert_eq!(obs1, vec![2.0, 5.0]);
2563 assert_eq!(obs2, vec![3.0, 6.0]);
2564
2565 assert_eq!((&a.obs(1)).into_iter().collect::<Vec<f64>>(), vec![2.0, 5.0]);
2567
2568 let b = NdArray::from_slice(&(1..=24).map(|x| x as f64).collect::<Vec<_>>(), &[2, 3, 4]);
2570 let obs0_3d = b.obs(0);
2571 assert_eq!(obs0_3d.shape(), &[3, 4]);
2572 }
2573
2574 #[test]
2575 fn col_access_2d() {
2576 let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
2577 let a = NdArray::from_slice(&data, &[10, 2]);
2578 assert_eq!(a.col(0).len(), 10);
2579 assert_eq!(a.col(1).len(), 10);
2580 assert_eq!(a.col(0), &data[..10]);
2581 assert_eq!(a.col(1), &data[10..20]);
2582 }
2583
2584 #[test]
2589 fn blas_params() {
2590 let a = NdArray::<f64>::new(&[10, 5]);
2591 assert_eq!(a.m(), 10);
2592 assert_eq!(a.n(), 5);
2593 assert_eq!(a.lda(), 10);
2594 }
2595
2596 #[test]
2597 fn blas_params_aligned_rows() {
2598 let a = NdArray::<f64>::new(&[8, 3]);
2599 assert_eq!(a.m(), 8);
2600 assert_eq!(a.n(), 3);
2601 assert_eq!(a.lda(), 8);
2602 }
2603
2604 #[test]
2609 fn compact_strides_2d() {
2610 for n_rows in 1..=20 {
2611 let a = NdArray::<f64>::new(&[n_rows, 3]);
2612 assert_eq!(a.strides()[1], n_rows);
2613 assert!(a.is_contiguous());
2614 }
2615 }
2616
2617 #[test]
2618 fn compact_strides_3d() {
2619 let a = NdArray::<f64>::new(&[10, 3, 4]);
2620 let strides = a.strides();
2621 assert_eq!(strides[0], 1);
2622 assert_eq!(strides[1], 10);
2623 assert_eq!(strides[2], 10 * 3);
2624 assert!(a.is_contiguous());
2625 }
2626
2627 #[test]
2632 fn reshape_1d_to_2d() {
2633 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[6]);
2634 let b = a.reshape(&[3, 2]).unwrap();
2635 assert_eq!(b.shape(), &[3, 2]);
2636 assert_eq!(b.get(&[0, 0]), 1.0);
2637 assert_eq!(b.get(&[1, 0]), 2.0);
2638 assert_eq!(b.get(&[2, 0]), 3.0);
2639 assert_eq!(b.get(&[0, 1]), 4.0);
2640 }
2641
2642 #[test]
2643 fn reshape_size_mismatch() {
2644 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2645 assert!(a.reshape(&[2, 2]).is_err());
2646 }
2647
2648 #[test]
2649 fn transpose_2d() {
2650 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2651 let t = a.transpose();
2652 assert_eq!(t.shape(), &[2, 3]);
2653 assert_eq!(t.get(&[0, 0]), 1.0);
2654 assert_eq!(t.get(&[1, 0]), 4.0);
2655 assert_eq!(t.get(&[0, 1]), 2.0);
2656 assert_eq!(t.get(&[1, 1]), 5.0);
2657 assert_eq!(t.get(&[0, 2]), 3.0);
2658 assert_eq!(t.get(&[1, 2]), 6.0);
2659 }
2660
2661 #[test]
2662 fn transpose_3d_reverses_axes() {
2663 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2664 let a = NdArray::from_slice(&data, &[2, 3, 4]);
2665 let t = a.transpose();
2666 assert_eq!(t.shape(), &[4, 3, 2]);
2667 assert!(t.is_contiguous());
2668 for i in 0..2 {
2670 for j in 0..3 {
2671 for k in 0..4 {
2672 assert_eq!(t.get(&[k, j, i]), a.get(&[i, j, k]));
2673 }
2674 }
2675 }
2676 }
2677
2678 #[test]
2679 fn transpose_1d_copies_through() {
2680 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2681 let t = a.transpose();
2682 assert_eq!(t.shape(), &[3]);
2683 assert_eq!(t, a);
2684 }
2685
2686 #[test]
2687 fn flatten() {
2688 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2689 let flat = a.flatten();
2690 assert_eq!(flat.shape(), &[6]);
2691 let vals: Vec<f64> = (&flat).into_iter().collect();
2692 assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
2693 }
2694
2695 #[test]
2696 fn to_contiguous_noop() {
2697 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2698 let b = a.to_contiguous();
2699 assert_eq!(a, b);
2700 }
2701
2702 #[test]
2703 fn fill_with_contiguous() {
2704 let mut a = NdArray::new(&[3, 2]);
2705 a.fill_with(42.0);
2706 for v in &a { assert_eq!(v, 42.0); }
2707 }
2708
2709 #[test]
2714 fn from_f64_slice() {
2715 let a = NdArray::from(&[1.0, 2.0, 3.0][..]);
2716 assert_eq!(a.ndim(), 1);
2717 assert_eq!(a.len(), 3);
2718 assert_eq!(a[(0,)], 1.0);
2719 }
2720
2721 #[test]
2722 fn from_vec64() {
2723 let v: Vec64<f64> = vec![10.0, 20.0].into_iter().collect();
2724 let a = NdArray::from(v);
2725 assert_eq!(a.ndim(), 1);
2726 assert_eq!(a[(0,)], 10.0);
2727 assert_eq!(a[(1,)], 20.0);
2728 }
2729
2730 #[test]
2731 fn from_column_vecs() {
2732 let cols = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
2733 let a = NdArray::from(cols.as_slice());
2734 assert_eq!(a.shape(), &[3, 2]);
2735 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2736 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2737 }
2738
2739 #[test]
2740 fn from_float_arrays() {
2741 let c0 = FloatArray::from_slice(&[1.0, 2.0]);
2742 let c1 = FloatArray::from_slice(&[3.0, 4.0]);
2743 let a = NdArray::from([c0, c1].as_slice());
2744 assert_eq!(a.shape(), &[2, 2]);
2745 assert_eq!(a.col(0), &[1.0, 2.0]);
2746 assert_eq!(a.col(1), &[3.0, 4.0]);
2747 }
2748
2749 #[test]
2750 fn from_buffer_explicit_strides() {
2751 let mut buf = Vec64::with_capacity(16);
2752 buf.0.resize(16, 0.0);
2753 buf[0] = 1.0; buf[1] = 2.0; buf[2] = 3.0;
2754 buf[8] = 4.0; buf[9] = 5.0; buf[10] = 6.0;
2755 let a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
2756 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2757 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2758 }
2759
2760 #[cfg(feature = "matrix")]
2765 #[test]
2766 fn from_matrix() {
2767 let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, Some("m".into()));
2768 let a = NdArray::from(mat);
2769 assert_eq!(a.shape(), &[3, 2]);
2770 assert_eq!(a.name.as_deref(), Some("m"));
2771 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2772 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2773 }
2774
2775 #[cfg(feature = "matrix")]
2776 #[test]
2777 fn to_matrix_roundtrip() {
2778 let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2779 let a = NdArray::from_slice(&data, &[3, 2]);
2780 let mat = a.to_matrix().unwrap();
2781 assert_eq!(mat.n_rows, 3);
2782 assert_eq!(mat.n_cols, 2);
2783 assert_eq!(mat.col(0), &[1.0, 2.0, 3.0]);
2784 assert_eq!(mat.col(1), &[4.0, 5.0, 6.0]);
2785 }
2786
2787 #[cfg(feature = "matrix")]
2788 #[test]
2789 fn to_matrix_non_2d_fails() {
2790 let a = NdArray::new(&[5]);
2791 assert!(a.to_matrix().is_err());
2792 }
2793
2794 fn make_numeric_table() -> Table {
2799 let c0 = FieldArray::from_arr("x", Array::NumericArray(
2800 NumericArray::Float64(Arc::new(FloatArray::from_slice(&[1.0, 2.0, 3.0])))
2801 ));
2802 let c1 = FieldArray::from_arr("y", Array::NumericArray(
2803 NumericArray::Float64(Arc::new(FloatArray::from_slice(&[4.0, 5.0, 6.0])))
2804 ));
2805 Table::new("data".to_string(), Some(vec![c0, c1]))
2806 }
2807
2808 #[test]
2809 fn try_from_table() {
2810 let table = make_numeric_table();
2811 let a = NdArray::try_from(&table).unwrap();
2812 assert_eq!(a.shape(), &[3, 2]);
2813 assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2814 assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2815 assert_eq!(a.name.as_deref(), Some("data"));
2816 }
2817
2818 #[test]
2819 fn try_from_table_with_nulls_converts_to_nan() {
2820 let mut mask = Bitmask::new_set_all(3, true);
2821 mask.set(1, false);
2822 let arr = FloatArray::new(Buffer::from_slice(&[10.0, 0.0, 30.0]), Some(mask));
2823 let c0 = FieldArray::from_arr("v", Array::NumericArray(
2824 NumericArray::Float64(Arc::new(arr))
2825 ));
2826 let table = Table::new("nulls".to_string(), Some(vec![c0]));
2827 let a = NdArray::try_from(&table).unwrap();
2828 assert_eq!(a.get(&[0, 0]), 10.0);
2829 assert!(a.get(&[1, 0]).is_nan());
2830 assert_eq!(a.get(&[2, 0]), 30.0);
2831 }
2832
2833 #[test]
2834 fn try_from_table_coerces_unparseable_text_to_nan() {
2835 let c0 = FieldArray::from_arr("name", Array::from_string32(
2839 StringArray::from_slice(&["a", "b"])
2840 ));
2841 let table = Table::new("text".to_string(), Some(vec![c0]));
2842 let a = NdArray::try_from(&table).unwrap();
2843 assert_eq!(a.shape(), &[2, 1]);
2844 assert!(a.get(&[0, 0]).is_nan());
2845 assert!(a.get(&[1, 0]).is_nan());
2846 }
2847
2848 #[test]
2849 fn to_table_roundtrip() {
2850 let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2851 let a = NdArray::from_slice(&data, &[3, 2]);
2852 let fields = vec![
2853 Field::new("x", ArrowType::Float64, false, None),
2854 Field::new("y", ArrowType::Float64, false, None),
2855 ];
2856 let table = a.to_table(Some(fields)).unwrap();
2857 assert_eq!(table.n_rows(), 3);
2858 assert_eq!(table.n_cols(), 2);
2859 assert_eq!(table.col_names(), vec!["x", "y"]);
2860 let col0 = table.cols[0].array.num().f64();
2861 assert_eq!(col0.data.as_slice(), &[1.0, 2.0, 3.0]);
2862 }
2863
2864 #[test]
2865 fn to_table_generated_names() {
2866 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2867 let table = a.to_table(None).unwrap();
2868 assert_eq!(table.col_names(), vec!["col_0", "col_1"]);
2869 }
2870
2871 #[test]
2872 fn to_table_rejects_incompatible_field_metadata() {
2873 let wrong_dtype = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
2874 Field::new("value", ArrowType::Float32, false, None),
2875 ]));
2876 assert!(matches!(wrong_dtype, Err(MinarrowError::TypeError { .. })));
2877
2878 let nullable = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
2879 Field::new("value", ArrowType::Float64, true, None),
2880 ]));
2881 assert!(matches!(nullable, Err(MinarrowError::NullError { .. })));
2882 }
2883
2884 #[test]
2885 fn to_table_non_2d_fails() {
2886 let a = NdArray::new(&[5]);
2887 assert!(a.to_table(None).is_err());
2888 }
2889
2890 #[test]
2891 fn to_array_1d() {
2892 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2893 let arr = a.to_array().unwrap();
2894 let f = arr.num().f64();
2895 assert_eq!(f.data.as_slice(), &[1.0, 2.0, 3.0]);
2896 }
2897
2898 #[test]
2899 fn to_array_materialises_strided_logical_values() {
2900 let a = NdArray::from_buffer(
2901 Buffer::from_slice(&[1.0, 99.0, 2.0, 99.0, 3.0]),
2902 &[3],
2903 &[2],
2904 );
2905 let arr = a.to_array().unwrap();
2906 assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
2907 }
2908
2909 #[test]
2910 fn to_array_ignores_unused_backing_elements() {
2911 let a = NdArray::from_buffer(
2912 Buffer::from_slice(&[1.0, 2.0, 3.0, 99.0]),
2913 &[3],
2914 &[1],
2915 );
2916 let arr = a.to_array().unwrap();
2917 assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
2918 }
2919
2920 #[test]
2921 fn to_array_non_1d_fails() {
2922 let a = NdArray::new(&[2, 3]);
2923 assert!(a.to_array().is_err());
2924 }
2925
2926 #[test]
2931 fn concat_1d() {
2932 let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2933 let b = NdArray::from_slice(&[3.0, 4.0, 5.0], &[3]);
2934 let c = a.concat(b).unwrap();
2935 assert_eq!(c.shape(), &[5]);
2936 let vals: Vec<f64> = (&c).into_iter().collect();
2937 assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
2938 }
2939
2940 #[test]
2941 fn concat_2d() {
2942 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2943 let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]);
2944 let c = a.concat(b).unwrap();
2945 assert_eq!(c.shape(), &[5, 2]);
2946 assert_eq!(c.col(0), &[1.0, 2.0, 5.0, 6.0, 7.0]);
2947 assert_eq!(c.col(1), &[3.0, 4.0, 8.0, 9.0, 10.0]);
2948 }
2949
2950 #[test]
2951 fn concat_dimension_mismatch_fails() {
2952 let a = NdArray::<f64>::new(&[3, 2]);
2953 let b = NdArray::new(&[3, 3]);
2954 assert!(a.concat(b).is_err());
2955 }
2956
2957 #[test]
2958 fn concat_rank_mismatch_fails() {
2959 let a = NdArray::<f64>::new(&[3]);
2960 let b = NdArray::new(&[3, 2]);
2961 assert!(a.concat(b).is_err());
2962 }
2963
2964 #[test]
2969 fn clone_and_eq() {
2970 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2971 let b = a.clone();
2972 assert_eq!(a, b);
2973 }
2974
2975 #[test]
2976 fn ne_different_data() {
2977 let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2978 let b = NdArray::from_slice(&[1.0, 3.0], &[2]);
2979 assert_ne!(a, b);
2980 }
2981
2982 #[test]
2983 fn ne_different_shape() {
2984 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
2985 let b = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2986 assert_ne!(a, b);
2987 }
2988
2989 #[test]
2994 fn debug_1d() {
2995 let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2996 let s = format!("{:?}", a);
2997 assert!(s.contains("[3]"));
2998 assert!(s.contains("1D"));
2999 }
3000
3001 #[test]
3002 fn debug_2d_named() {
3003 let a = NdArray::<f64>::new_named(&[2, 3], "test");
3004 let s = format!("{:?}", a);
3005 assert!(s.contains("'test'"));
3006 assert!(s.contains("[2, 3]"));
3007 assert!(s.contains("2D"));
3008 }
3009
3010 #[test]
3015 fn empty_array() {
3016 let a = NdArray::new(&[0, 5]);
3017 assert!(a.is_empty());
3018 assert_eq!(a.len(), 0);
3019 let vals: Vec<f64> = (&a).into_iter().collect();
3020 assert!(vals.is_empty());
3021 }
3022
3023 #[test]
3024 fn single_element() {
3025 let a = NdArray::from_slice(&[42.0], &[1]);
3026 assert_eq!(a.len(), 1);
3027 assert_eq!(a[(0,)], 42.0);
3028 let vals: Vec<f64> = (&a).into_iter().collect();
3029 assert_eq!(vals, vec![42.0]);
3030 }
3031
3032 #[test]
3033 fn single_element_2d() {
3034 let a = NdArray::from_slice(&[42.0], &[1, 1]);
3035 assert_eq!(a[(0, 0)], 42.0);
3036 }
3037
3038 #[test]
3039 fn large_array_iteration_count() {
3040 let n = 1000;
3041 let a = NdArray::<f64>::ones(&[n, 100]);
3042 let count = (&a).into_iter().count();
3043 assert_eq!(count, n * 100);
3044 }
3045
3046 #[test]
3051 fn bracket_index_1d() {
3052 let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3053 assert_eq!(a[0], [10.0]);
3054 assert_eq!(a[2], [30.0]);
3055 }
3056
3057 #[test]
3058 fn bracket_index_2d_column() {
3059 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3060 assert_eq!(a[0], [1.0, 2.0, 3.0]);
3061 assert_eq!(a[1], [4.0, 5.0, 6.0]);
3062 }
3063
3064 #[test]
3065 fn bracket_index_2d_chained() {
3066 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3067 assert_eq!(a[0][0], 1.0);
3069 assert_eq!(a[0][2], 3.0);
3070 assert_eq!(a[1][0], 4.0);
3071 assert_eq!(a[1][2], 6.0);
3072 }
3073
3074 #[test]
3075 fn bracket_index_mut_2d() {
3076 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3077 a[0][1] = 99.0;
3078 assert_eq!(a[0][1], 99.0);
3079 assert_eq!(a[0], [1.0, 99.0, 3.0]);
3080 }
3081
3082 #[test]
3087 fn range_index_1d() {
3088 let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3089 assert_eq!(a[1..3], [20.0, 30.0]);
3090 }
3091
3092 #[test]
3093 fn range_index_2d_columns() {
3094 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3097 let slab = &a[0..2];
3098 assert_eq!(slab, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3099 let col1 = &a[1..2];
3100 assert_eq!(col1, &[4.0, 5.0, 6.0]);
3101 }
3102
3103 #[cfg(feature = "matrix")]
3104 #[test]
3105 #[should_panic(expected = "contiguous")]
3106 fn range_index_non_contiguous_panics() {
3107 let mat = Matrix::from_f64_unaligned(
3110 &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
3111 );
3112 let a = NdArray::from(mat);
3113 if a.is_contiguous() {
3114 panic!("premise failed: expected non-contiguous import");
3117 }
3118 let _ = &a[0..2];
3119 }
3120
3121 #[cfg(feature = "matrix")]
3122 #[test]
3123 fn to_matrix_repacks_compact_layout() {
3124 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3126 let mat = a.to_matrix().unwrap();
3127 assert_eq!(mat.n_rows, 3);
3128 assert_eq!(mat.n_cols, 2);
3129 assert_eq!(mat.stride, 8);
3130 assert_eq!(&mat.data.as_slice()[..3], &[1.0, 2.0, 3.0]);
3131 assert_eq!(&mat.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
3132 }
3133
3134 #[cfg(feature = "matrix")]
3135 #[test]
3136 fn matrix_roundtrip_via_contiguous() {
3137 let mat = Matrix::from_f64_unaligned(
3140 &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
3141 );
3142 let a = NdArray::from(mat);
3143 assert!(!a.is_contiguous());
3144 assert_eq!(a.get(&[2, 1]), 6.0);
3145 let compact = a.to_contiguous();
3146 assert!(compact.is_contiguous());
3147 assert_eq!(&compact[0..2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3148 let back = compact.to_matrix().unwrap();
3149 assert_eq!(back.stride, 8);
3150 assert_eq!(&back.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
3151 }
3152
3153 #[test]
3154 fn range_from_index() {
3155 let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3156 assert_eq!(a[2..], [30.0, 40.0]);
3157 }
3158
3159 #[test]
3160 fn range_to_index() {
3161 let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3162 assert_eq!(a[..2], [10.0, 20.0]);
3163 }
3164
3165 #[test]
3166 fn range_full_index() {
3167 let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3168 assert_eq!(a[..], [10.0, 20.0, 30.0]);
3169 }
3170
3171 #[cfg(feature = "views")]
3176 #[test]
3177 fn slice_1d_single_index() {
3178 let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3179 let v = a.slice(&[&1]);
3180 assert!(v.shape().is_empty());
3181 assert_eq!(v[()], 20.0);
3182 }
3183
3184 #[cfg(feature = "views")]
3185 #[test]
3186 fn slice_1d_range() {
3187 let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3188 let v = a.slice(&[&(1..3)]);
3189 assert_eq!(v.shape(), &[2]);
3190 assert_eq!(v[(0,)], 20.0);
3191 assert_eq!(v[(1,)], 30.0);
3192 }
3193
3194 #[cfg(feature = "views")]
3195 #[test]
3196 fn slice_2d_row_range_single_col() {
3197 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3198 let v = a.slice(nd![0..2, 1]);
3200 assert_eq!(v.shape(), &[2]);
3201 assert_eq!(v[(0,)], 4.0);
3202 assert_eq!(v[(1,)], 5.0);
3203 }
3204
3205 #[cfg(feature = "views")]
3206 #[test]
3207 fn slice_2d_single_row_col_range() {
3208 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3209 let v = a.slice(nd![1, 0..2]);
3211 assert_eq!(v.shape(), &[2]);
3212 let vals: Vec<f64> = (&v).into_iter().collect();
3214 assert_eq!(vals, vec![2.0, 5.0]);
3215 }
3216
3217 #[cfg(feature = "views")]
3218 #[test]
3219 fn slice_2d_both_ranges() {
3220 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3221 let v = a.slice(nd![0..2, 0..2]);
3223 assert_eq!(v.shape(), &[2, 2]);
3224 assert_eq!(v[(0, 0)], 1.0);
3225 assert_eq!(v[(1, 0)], 2.0);
3226 assert_eq!(v[(0, 1)], 4.0);
3227 assert_eq!(v[(1, 1)], 5.0);
3228 }
3229
3230 #[cfg(feature = "views")]
3231 #[test]
3232 fn slice_2d_both_indices_scalar() {
3233 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3234 let v = a.slice(nd![2, 1]);
3236 assert!(v.shape().is_empty());
3237 assert_eq!(v[()], 6.0);
3238 }
3239
3240 #[cfg(feature = "views")]
3241 #[test]
3242 fn slice_3d_mixed() {
3243 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3245 let a = NdArray::from_slice(&data, &[2, 3, 4]);
3246 let v = a.slice(nd![0..2, 1, 0..2]);
3248 assert_eq!(v.shape(), &[2, 2]);
3249 assert_eq!(v[(0, 0)], 3.0);
3251 assert_eq!(v[(1, 0)], 4.0);
3252 assert_eq!(v[(0, 1)], 9.0);
3253 assert_eq!(v[(1, 1)], 10.0);
3254 }
3255
3256 #[cfg(feature = "views")]
3257 #[test]
3258 fn slice_with_nd_macro() {
3259 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3260 let a = NdArray::from_slice(&data, &[2, 3, 4]);
3261 let v = a.slice(nd![0..2, 0..3, 0..4]);
3262 assert_eq!(v.shape(), &[2, 3, 4]);
3263 assert_eq!(v.len(), 24);
3264 }
3265
3266 #[cfg(feature = "views")]
3267 #[test]
3268 fn slice_preserves_data_through_iteration() {
3269 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3270 let v = a.slice(nd![1..3, 0..2]);
3271 let vals: Vec<f64> = (&v).into_iter().collect();
3273 assert_eq!(vals, vec![2.0, 3.0, 5.0, 6.0]);
3274 }
3275
3276 #[cfg(feature = "views")]
3277 #[test]
3278 fn slice_column_window() {
3279 let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
3281 let a = NdArray::from_slice(&data, &[10, 2]);
3282 let v = a.slice(nd![2..5, 1]);
3283 assert_eq!(v.shape(), &[3]);
3284 assert_eq!(v[(0,)], 13.0);
3285 assert_eq!(v[(1,)], 14.0);
3286 assert_eq!(v[(2,)], 15.0);
3287 }
3288
3289 #[test]
3294 #[should_panic(expected = "Column index out of bounds")]
3295 fn col_out_of_bounds_panics() {
3296 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3297 let _ = a.col(2);
3298 }
3299
3300 #[test]
3301 #[should_panic(expected = "indices for a 2D array")]
3302 fn get_rank_mismatch_panics() {
3303 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3304 let _ = a.get(&[0]);
3305 }
3306
3307 #[test]
3308 #[should_panic(expected = "out of bounds for dim")]
3309 fn get_index_out_of_bounds_panics() {
3310 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3311 let _ = a.get(&[2, 0]);
3312 }
3313
3314 #[test]
3315 #[should_panic(expected = "out of bounds for dim")]
3316 fn set_out_of_bounds_panics() {
3317 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3318 a.set(&[0, 5], 9.0);
3319 }
3320
3321 #[cfg(all(feature = "views", feature = "select"))]
3322 #[test]
3323 #[should_panic(expected = "expected 2 axes, got 1")]
3324 fn slice_wrong_axis_count_panics() {
3325 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3326 let _ = a.slice(nd![0..2]);
3327 }
3328
3329 #[cfg(all(feature = "views", feature = "select"))]
3330 #[test]
3331 #[should_panic(expected = "range 0..100 out of bounds")]
3332 fn slice_range_out_of_bounds_panics() {
3333 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3334 let _ = a.slice(nd![0..100, 0..2]);
3335 }
3336
3337 #[cfg(all(feature = "views", feature = "select"))]
3338 #[test]
3339 #[should_panic(expected = "index 5 out of bounds")]
3340 fn slice_single_index_out_of_bounds_panics() {
3341 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3342 let _ = a.slice(nd![5, 0..2]);
3343 }
3344
3345 #[test]
3346 #[should_panic(expected = "at least 2 points")]
3347 fn linspace_requires_two_points() {
3348 let _ = NdArray::<f64>::linspace(0.0, 1.0, 1);
3349 }
3350
3351 #[cfg(all(feature = "views", feature = "select"))]
3352 #[test]
3353 fn slice_one_element_index_array_collapses() {
3354 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3355 let idx: &[usize] = &[1];
3356 let v = a.slice(nd![idx, 0..2]);
3357 assert_eq!(v.shape(), &[2]);
3358 assert_eq!(v[(0,)], 2.0);
3359 assert_eq!(v[(1,)], 5.0);
3360 }
3361
3362 #[test]
3363 fn get_unchecked_matches_get() {
3364 let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3365 let a = NdArray::from_slice(&data, &[2, 3, 4]);
3366 for i in 0..2 {
3367 for j in 0..3 {
3368 for k in 0..4 {
3369 let idx = [i, j, k];
3370 assert_eq!(unsafe { a.get_unchecked(&idx) }, a.get(&idx));
3372 }
3373 }
3374 }
3375 }
3376
3377 #[test]
3382 fn concat_3d_interleaves_axis0() {
3383 let da: Vec<f64> = (1..=8).map(|x| x as f64).collect();
3384 let db: Vec<f64> = (9..=16).map(|x| x as f64).collect();
3385 let a = NdArray::from_slice(&da, &[2, 2, 2]);
3386 let b = NdArray::from_slice(&db, &[2, 2, 2]);
3387 let c = a.clone().concat(b.clone()).unwrap();
3388 assert_eq!(c.shape(), &[4, 2, 2]);
3389 for i in 0..4 {
3390 for j in 0..2 {
3391 for k in 0..2 {
3392 let expected = if i < 2 {
3393 a.get(&[i, j, k])
3394 } else {
3395 b.get(&[i - 2, j, k])
3396 };
3397 assert_eq!(c.get(&[i, j, k]), expected);
3398 }
3399 }
3400 }
3401 }
3402
3403 #[test]
3404 fn concat_non_contiguous_operand() {
3405 let a = NdArray::from_buffer(
3408 Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]),
3409 &[2, 2],
3410 &[2, 1],
3411 );
3412 assert!(!a.is_contiguous());
3413 let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]);
3414 let c = a.concat(b).unwrap();
3415 assert_eq!(c.shape(), &[4, 2]);
3416 assert_eq!(c.col(0), &[1.0, 3.0, 5.0, 6.0]);
3417 assert_eq!(c.col(1), &[2.0, 4.0, 7.0, 8.0]);
3418 }
3419
3420 #[cfg(feature = "views")]
3425 #[test]
3426 fn set_after_view_copy_on_write() {
3427 let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3428 let v = a.as_view();
3429 a.set(&[0, 0], 99.0);
3430 assert_eq!(a.get(&[0, 0]), 99.0);
3433 assert_eq!(v.get(&[0, 0]), 1.0);
3434 }
3435
3436 #[test]
3437 fn fill_with_non_contiguous_logical_only() {
3438 let mut buf = Vec64::with_capacity(11);
3440 buf.0.resize(11, 0.0);
3441 let mut a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
3442 assert!(!a.is_contiguous());
3443 a.fill_with(7.0);
3444 for v in &a { assert_eq!(v, 7.0); }
3445 assert_eq!(a.as_slice()[3], 0.0);
3447 assert_eq!(a.as_slice()[7], 0.0);
3448 }
3449
3450 #[test]
3455 fn range_from_index_2d() {
3456 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3457 assert_eq!(&a[1..], &[4.0, 5.0, 6.0]);
3458 }
3459
3460 #[test]
3461 fn range_to_index_2d() {
3462 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3463 assert_eq!(&a[..1], &[1.0, 2.0, 3.0]);
3464 }
3465
3466 #[test]
3467 fn range_full_index_2d() {
3468 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3469 assert_eq!(&a[..], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3470 }
3471
3472 #[cfg(all(feature = "views", feature = "select"))]
3477 #[test]
3478 fn row_selection_single_index_window() {
3479 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3480 assert_eq!(a.n_obs(), 3);
3481 let v = a.r(1usize);
3482 assert_eq!(v.shape(), &[1, 2]);
3483 assert_eq!(v.get(&[0, 0]), 2.0);
3484 assert_eq!(v.get(&[0, 1]), 5.0);
3485 assert_eq!(v.source.len(), a.len());
3487 assert_eq!(unsafe { v.get_unchecked(&[0, 1]) }, 5.0);
3489 }
3490
3491 #[cfg(feature = "size")]
3492 #[test]
3493 fn est_bytes_covers_buffer() {
3494 use crate::traits::byte_size::ByteSize;
3495 let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3496 assert!(a.est_bytes() >= 6 * size_of::<f64>());
3497 }
3498}