rlst/dense/
array.rs

1//! Basic Array type
2//!
3//! [Array] is the basic type for dense calculations in Rlst. The full definition
4//! `Array<ArrayImpl, NDIM>` represents a tensor with `NDIM` axes implemented through
5//! the implementeation type `ArrayImpl`.
6
7use std::ops::{Add, Sub};
8
9use aligned_vec::{AVec, CACHELINE_ALIGN, ConstAlign};
10use empty_axis::AxisPosition;
11use iterators::{
12    ArrayDefaultIteratorByRef, ArrayDefaultIteratorByValue, ArrayDefaultIteratorMut,
13    ArrayDiagIteratorByRef, ArrayDiagIteratorByValue, ArrayDiagIteratorMut, ColIterator,
14    ColIteratorMut, RowIterator, RowIteratorMut,
15};
16use itertools::izip;
17use num::{One, Zero, traits::MulAdd};
18use operators::{
19    addition::ArrayAddition, cast::ArrayCast, cmp_wise_division::CmpWiseDivision,
20    cmp_wise_product::CmpWiseProduct, coerce::CoerceArray, mul_add::MulAddImpl, negation::ArrayNeg,
21    reverse_axis::ReverseAxis, scalar_mult::ArrayScalarMult, subtraction::ArraySubtraction,
22    transpose::ArrayTranspose, unary_op::ArrayUnaryOperator,
23};
24use typenum::{Const, NonZero, PInt, ToUInt, U, Unsigned};
25
26use crate::base_types::{c32, c64};
27
28use crate::dense::array::iterators::{MultisliceIterator, MultisliceIteratorMut};
29use crate::{
30    ArrayFromContainerType, Linspace, MutableArrayImpl, RowMajorArrayFromContainerType,
31    ScalarAsArray,
32};
33
34use paste::paste;
35use rand::{Rng, SeedableRng};
36use rand_chacha::ChaCha8Rng;
37use rand_distr::{Distribution, StandardNormal, StandardUniform};
38use reference::{ArrayRef, ArrayRefMut};
39use slice::ArraySlice;
40use subview::ArraySubView;
41
42use crate::{
43    AsMatrixApply, AsMultiIndex, EvaluateArray, EvaluateRowMajorArray, Gemm, IsGreaterByOne,
44    IsGreaterZero, IsSmallerThan, Max, MemoryLayout, NumberDifference, NumberType, RandScalar,
45    RandomAccessByValue, RlstError, RlstResult, RlstScalar, SubtractEqual, TransMode,
46    dense::{
47        array::multislice::Multislice,
48        base_array::BaseArray,
49        data_container::{
50            AlignedVectorContainer, SliceContainer, SliceContainerMut, VectorContainer,
51        },
52        layout::check_multi_index_in_bounds,
53        strided_base_array::StridedBaseArray,
54    },
55    traits::{
56        accessors::{
57            RandomAccessByRef, RandomAccessMut, RawAccess, RawAccessMut, UnsafeRandom1DAccessByRef,
58            UnsafeRandom1DAccessByValue, UnsafeRandom1DAccessMut, UnsafeRandomAccessByRef,
59            UnsafeRandomAccessByValue, UnsafeRandomAccessMut,
60        },
61        base_operations::{BaseItem, ResizeInPlace, Shape, Stride},
62        data_container::ContainerType,
63    },
64};
65
66use super::{data_container::ArrayContainer, layout::row_major_stride_from_shape};
67
68pub mod empty_axis;
69pub mod iterators;
70pub mod mult_into;
71pub mod operations;
72pub mod operators;
73// pub mod rank1_array;
74pub mod flattened;
75pub mod multislice;
76pub mod reference;
77pub mod shared_array;
78pub mod slice;
79pub mod subview;
80
81/// A basic dynamically allocated array.
82pub type DynArray<Item, const NDIM: usize> = Array<BaseArray<VectorContainer<Item>, NDIM>, NDIM>;
83
84/// A basic dynamically allocated array with given alignment.
85pub type AlignedDynArray<Item, const NDIM: usize, const ALIGNMENT: usize = CACHELINE_ALIGN> =
86    Array<BaseArray<AlignedVectorContainer<Item, ALIGNMENT>, NDIM>, NDIM>;
87
88/// A basic dynamically allocated array with a given stride.
89pub type StridedDynArray<Item, const NDIM: usize> =
90    Array<StridedBaseArray<VectorContainer<Item>, NDIM>, NDIM>;
91
92/// A statically allocated array.
93///
94/// `N` is the number of elements to be reserved in the static
95/// container.
96pub type StaticArray<Item, const NDIM: usize, const N: usize> =
97    Array<BaseArray<ArrayContainer<Item, N>, NDIM>, NDIM>;
98
99/// A dynamically allocated array from a data slice.
100pub type SliceArray<'a, Item, const NDIM: usize> =
101    Array<BaseArray<SliceContainer<'a, Item>, NDIM>, NDIM>;
102
103/// A dynamically allocated array from a data slice with a given stride.
104pub type StridedSliceArray<'a, Item, const NDIM: usize> =
105    Array<StridedBaseArray<SliceContainer<'a, Item>, NDIM>, NDIM>;
106
107/// A mutable dynamically allocated array from a data slice.
108pub type SliceArrayMut<'a, Item, const NDIM: usize> =
109    Array<BaseArray<SliceContainerMut<'a, Item>, NDIM>, NDIM>;
110
111/// A mutable dynamically allocated array from a data slice with a given stride.
112pub type StridedSliceArrayMut<'a, Item, const NDIM: usize> =
113    Array<StridedBaseArray<SliceContainerMut<'a, Item>, NDIM>, NDIM>;
114
115/// The basic tuple type defining an array.
116pub struct Array<ArrayImpl, const NDIM: usize>(ArrayImpl);
117
118impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
119    /// Instantiate a new array from an `ArrayImpl` structure.
120    pub fn new(arr: ArrayImpl) -> Self {
121        assert!(
122            NDIM > 0,
123            "Array dimension must be greater than 0, got {NDIM}"
124        );
125        Self(arr)
126    }
127
128    /// Return a reference to the implementation.
129    #[inline(always)]
130    pub fn imp(&self) -> &ArrayImpl {
131        &self.0
132    }
133
134    /// Return a mutable reference to the implementation.
135    #[inline(always)]
136    pub fn imp_mut(&mut self) -> &mut ArrayImpl {
137        &mut self.0
138    }
139
140    /// Extract the inner implementation.
141    #[inline(always)]
142    pub fn into_imp(self) -> ArrayImpl {
143        self.0
144    }
145
146    /// Create an owned reference to the array.
147    #[inline(always)]
148    pub fn r(&self) -> Array<ArrayRef<'_, ArrayImpl, NDIM>, NDIM> {
149        Array::new(ArrayRef::new(self))
150    }
151
152    /// Create an owned mutable reference to the array.
153    #[inline(always)]
154    pub fn r_mut(&mut self) -> Array<ArrayRefMut<'_, ArrayImpl, NDIM>, NDIM> {
155        Array::new(ArrayRefMut::new(self))
156    }
157}
158
159impl<ArrayImpl, const NDIM: usize> BaseItem for Array<ArrayImpl, NDIM>
160where
161    ArrayImpl: BaseItem,
162{
163    type Item = ArrayImpl::Item;
164}
165
166impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> Array<ArrayImpl, NDIM> {
167    /// Return the number of elements in the array.
168    #[inline(always)]
169    pub fn number_of_elements(&self) -> usize {
170        self.0.shape().iter().product()
171    }
172}
173
174impl<Item: Copy + Default, const NDIM: usize> DynArray<Item, NDIM> {
175    /// Create a new heap allocated array from a given `shape`.
176    #[inline(always)]
177    pub fn from_shape(shape: [usize; NDIM]) -> Self {
178        let size = shape.iter().product();
179        Self::new(BaseArray::new(VectorContainer::new(size), shape))
180    }
181
182    /// Create a new heap allocated array by providing a `shape` and a vector of `data`.
183    ///
184    /// The number of elements in the vector must be compatible with the given shape.
185    /// Otherwise, an assertion error is triggered.
186    #[inline(always)]
187    pub fn from_shape_and_vec(shape: [usize; NDIM], data: Vec<Item>) -> Self {
188        assert_eq!(
189            data.len(),
190            shape.iter().product::<usize>(),
191            "Data length does not match the shape: {} != {}",
192            data.len(),
193            shape.iter().product::<usize>()
194        );
195        Self::new(BaseArray::new(VectorContainer::from_vec(data), shape))
196    }
197}
198
199impl<Item: Copy + Default, const NDIM: usize, const ALIGNMENT: usize>
200    AlignedDynArray<Item, NDIM, ALIGNMENT>
201{
202    /// Create a new heap allocated array from a given `shape`.
203    #[inline(always)]
204    pub fn from_shape(shape: [usize; NDIM]) -> Self {
205        let size = shape.iter().product();
206        Self::new(BaseArray::new(AlignedVectorContainer::new(size), shape))
207    }
208
209    /// Create a new heap allocated array by providing a `shape` and a vector of `data`.
210    ///
211    /// The number of elements in the vector must be compatible with the given shape.
212    /// Otherwise, an assertion error is triggered.
213    #[inline(always)]
214    pub fn from_shape_and_avec(
215        shape: [usize; NDIM],
216        data: AVec<Item, ConstAlign<ALIGNMENT>>,
217    ) -> Self {
218        assert_eq!(
219            data.len(),
220            shape.iter().product::<usize>(),
221            "Data length does not match the shape: {} != {}",
222            data.len(),
223            shape.iter().product::<usize>()
224        );
225        Self::new(BaseArray::new(
226            AlignedVectorContainer::from_avec(data),
227            shape,
228        ))
229    }
230}
231
232impl<Item: Copy + Default> DynArray<Item, 2> {
233    /// Create a new dense matrix with shape `shape` filled with values from the iterator `iter`.
234    pub fn from_iter_aij<Iter: Iterator<Item = ([usize; 2], Item)>>(
235        shape: [usize; 2],
236        iter: Iter,
237    ) -> DynArray<Item, 2> {
238        let mut out = DynArray::<Item, 2>::from_shape(shape);
239
240        for (index, item) in iter {
241            out[index] = item;
242        }
243
244        out
245    }
246}
247
248impl<Item: Copy + Default, const ALIGNMENT: usize> AlignedDynArray<Item, 2, ALIGNMENT> {
249    /// Create a new dense matrix with shape `shape` filled with values from the iterator `iter`.
250    pub fn from_iter_aij<Iter: Iterator<Item = ([usize; 2], Item)>>(
251        shape: [usize; 2],
252        iter: Iter,
253    ) -> AlignedDynArray<Item, 2, ALIGNMENT> {
254        let mut out = AlignedDynArray::<Item, 2, ALIGNMENT>::from_shape(shape);
255
256        for (index, item) in iter {
257            out[index] = item;
258        }
259
260        out
261    }
262}
263
264impl<Item: Copy + Default> From<&[Item]> for DynArray<Item, 1> {
265    fn from(value: &[Item]) -> Self {
266        DynArray::<Item, 1>::from_shape_and_vec([value.len()], value.to_vec())
267    }
268}
269
270impl<Item: Copy + Default, const N: usize> From<[Item; N]> for StaticArray<Item, 1, N> {
271    fn from(value: [Item; N]) -> Self {
272        StaticArray::new(BaseArray::new(value.into(), [N]))
273    }
274}
275
276impl<Item: Copy + Default> From<Vec<Item>> for DynArray<Item, 1> {
277    fn from(value: Vec<Item>) -> Self {
278        DynArray::<Item, 1>::from_shape_and_vec([value.len()], value)
279    }
280}
281
282impl<Item: Copy + Default> From<Vec<Vec<Item>>> for DynArray<Item, 2> {
283    fn from(value: Vec<Vec<Item>>) -> Self {
284        let nrows = value.len();
285        if nrows == 0 {
286            empty_array()
287        } else {
288            let ncols = value.first().unwrap().len();
289            // Check that all columns have identical length
290
291            for row in value.iter() {
292                assert_eq!(ncols, row.len(), "All rows must have equal length.");
293            }
294
295            // Now create the matrix and fill with values.
296
297            let mut out = DynArray::<Item, 2>::from_shape([nrows, ncols]);
298
299            for (row_index, row) in value.iter().enumerate() {
300                for (col_index, &elem) in row.iter().enumerate() {
301                    out[[row_index, col_index]] = elem;
302                }
303            }
304
305            out
306        }
307    }
308}
309
310impl<Item: Copy + Default, const NDIM: usize> StridedDynArray<Item, NDIM> {
311    /// Create a new heap allocated array from a given `shape` and `stride`.
312    #[inline(always)]
313    pub fn from_shape_and_stride(shape: [usize; NDIM], stride: [usize; NDIM]) -> Self {
314        let size = shape.iter().product();
315        Self::new(StridedBaseArray::new(
316            VectorContainer::new(size),
317            shape,
318            stride,
319        ))
320    }
321}
322
323impl<Item: Copy + Default, const NDIM: usize>
324    Array<StridedBaseArray<VectorContainer<Item>, NDIM>, NDIM>
325{
326    /// Create a new heap allocated row-major array.
327    #[inline(always)]
328    pub fn row_major(shape: [usize; NDIM]) -> Self {
329        let size = shape.iter().product();
330        Self::new(StridedBaseArray::new(
331            VectorContainer::new(size),
332            shape,
333            row_major_stride_from_shape(shape),
334        ))
335    }
336}
337
338impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
339where
340    ArrayImpl: Shape<NDIM>,
341{
342    /// Return the shape of an array.
343    ///
344    /// # Traits:
345    /// - [Shape]
346    #[inline(always)]
347    pub fn shape(&self) -> [usize; NDIM] {
348        self.0.shape()
349    }
350}
351
352impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
353where
354    ArrayImpl: Shape<NDIM>,
355{
356    /// Return the length of an array.
357    ///
358    /// For more than one dimension the length is the number of elements.
359    /// # Traits:
360    /// - [Shape]
361    #[inline(always)]
362    pub fn len(&self) -> usize {
363        self.0.len()
364    }
365
366    /// Return true if the array is empty.
367    ///
368    /// This is equivalent to at least one dimension being zero.
369    /// # Traits:
370    /// - [Shape]
371    #[inline(always)]
372    pub fn is_empty(&self) -> bool {
373        self.len() == 0
374    }
375}
376
377impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
378where
379    ArrayImpl: UnsafeRandomAccessByValue<NDIM>,
380{
381    /// Get an element by value.
382    ///
383    /// # Safety
384    /// `multi_index` must be in range.
385    ///
386    /// # Traits
387    /// - [UnsafeRandomAccessByValue]
388    #[inline(always)]
389    pub unsafe fn get_value_unchecked(&self, multi_index: [usize; NDIM]) -> ArrayImpl::Item {
390        unsafe { self.0.get_value_unchecked(multi_index) }
391    }
392}
393
394impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
395where
396    ArrayImpl: UnsafeRandomAccessByRef<NDIM>,
397{
398    /// Get an element by reference.
399    ///
400    /// # Safety
401    /// `multi_index` must be in range.
402    ///
403    /// # Traits
404    /// - [UnsafeRandomAccessByRef]
405    #[inline(always)]
406    pub unsafe fn get_unchecked(&self, multi_index: [usize; NDIM]) -> &ArrayImpl::Item {
407        unsafe { self.0.get_unchecked(multi_index) }
408    }
409}
410
411impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
412where
413    ArrayImpl: UnsafeRandomAccessMut<NDIM>,
414{
415    /// Get an element by mutable reference.
416    ///
417    /// # Safety
418    /// `multi_index` must be in range.
419    ///
420    /// # Traits
421    /// - [UnsafeRandomAccessMut]
422    #[inline(always)]
423    pub unsafe fn get_unchecked_mut(&mut self, multi_index: [usize; NDIM]) -> &mut ArrayImpl::Item {
424        unsafe { self.0.get_unchecked_mut(multi_index) }
425    }
426}
427
428impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
429where
430    ArrayImpl: RandomAccessByValue<NDIM>,
431{
432    /// Get an element by value.
433    ///
434    /// # Traits
435    /// - [RandomAccessByValue]
436    #[inline(always)]
437    pub fn get_value(&self, multi_index: [usize; NDIM]) -> Option<ArrayImpl::Item> {
438        self.0.get_value(multi_index)
439    }
440}
441
442impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
443where
444    ArrayImpl: RandomAccessByRef<NDIM>,
445{
446    /// Get an element by reference.
447    ///
448    /// # Traits
449    /// - [RandomAccessByRef]
450    #[inline(always)]
451    pub fn get(&self, multi_index: [usize; NDIM]) -> Option<&ArrayImpl::Item> {
452        self.0.get(multi_index)
453    }
454}
455
456impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
457where
458    ArrayImpl: RandomAccessMut<NDIM>,
459{
460    /// Get an element by mutable reference.
461    ///
462    /// # Traits
463    /// - [RandomAccessMut]
464    #[inline(always)]
465    pub fn get_mut(&mut self, multi_index: [usize; NDIM]) -> Option<&mut ArrayImpl::Item> {
466        self.0.get_mut(multi_index)
467    }
468}
469
470impl<ArrayImpl, const NDIM: usize> std::ops::Index<[usize; NDIM]> for Array<ArrayImpl, NDIM>
471where
472    ArrayImpl: RandomAccessByRef<NDIM>,
473{
474    type Output = ArrayImpl::Item;
475    #[inline(always)]
476    fn index(&self, index: [usize; NDIM]) -> &Self::Output {
477        self.0.get(index).unwrap()
478    }
479}
480
481impl<ArrayImpl, const NDIM: usize> std::ops::IndexMut<[usize; NDIM]> for Array<ArrayImpl, NDIM>
482where
483    Array<ArrayImpl, NDIM>: std::ops::Index<[usize; NDIM], Output = ArrayImpl::Item>,
484    ArrayImpl: RandomAccessMut<NDIM>,
485{
486    #[inline(always)]
487    fn index_mut(&mut self, index: [usize; NDIM]) -> &mut ArrayImpl::Item {
488        self.0.get_mut(index).unwrap()
489    }
490}
491
492impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
493where
494    ArrayImpl: RawAccess,
495{
496    /// Return the raw data as slice.
497    ///
498    /// # Traits
499    /// - [RawAccess]
500    #[inline(always)]
501    pub fn data(&self) -> Option<&[ArrayImpl::Item]> {
502        self.0.data()
503    }
504}
505
506impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
507where
508    ArrayImpl: RawAccessMut,
509{
510    /// Return the raw data as slice.
511    ///
512    /// # Traits
513    /// - [RawAccessMut]
514    #[inline(always)]
515    pub fn data_mut(&mut self) -> Option<&mut [ArrayImpl::Item]> {
516        self.0.data_mut()
517    }
518}
519
520impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
521where
522    ArrayImpl: Stride<NDIM>,
523{
524    /// Return the stride.
525    ///
526    /// # Traits
527    /// - [Stride]
528    #[inline(always)]
529    pub fn stride(&self) -> [usize; NDIM] {
530        self.0.stride()
531    }
532}
533
534impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
535where
536    ArrayImpl: Stride<NDIM> + Shape<NDIM>,
537{
538    /// Return the memory layout
539    ///
540    /// The possible memory layouts are defined in [MemoryLayout].
541    #[inline(always)]
542    pub fn memory_layout(&self) -> MemoryLayout {
543        self.0.memory_layout()
544    }
545
546    /// Return true of the memory layout is contiguous, otherwise false.
547    ///
548    /// A memory layout is contiguous if it is either row or column major.
549    pub fn is_contiguous(&self) -> bool {
550        self.0.is_contiguous()
551    }
552}
553
554impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
555where
556    ArrayImpl: ResizeInPlace<NDIM>,
557{
558    /// Resize the array to the new `shape`.
559    ///
560    /// The content of the array will be lost upon resizing.
561    ///
562    /// # Traits
563    /// - [ResizeInPlace]
564    pub fn resize_in_place(&mut self, shape: [usize; NDIM]) {
565        self.0.resize_in_place(shape)
566    }
567}
568
569impl<Item, ArrayImpl, const NDIM: usize> EvaluateArray<Item, NDIM> for Array<ArrayImpl, NDIM>
570where
571    Item: Copy + Default,
572    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM> + ContainerType,
573{
574    type OutputImpl =
575        <<ArrayImpl as ContainerType>::Type as ArrayFromContainerType>::OutputImpl<Item, NDIM>;
576
577    fn eval(&self) -> Array<Self::OutputImpl, NDIM> {
578        let mut out = <<ArrayImpl as ContainerType>::Type as ArrayFromContainerType>::new_arr::<
579            Item,
580            NDIM,
581        >(self.shape());
582
583        out.fill_from(self);
584
585        out
586    }
587}
588
589impl<Item, ArrayImpl, const NDIM: usize> EvaluateRowMajorArray<Item, NDIM>
590    for Array<ArrayImpl, NDIM>
591where
592    Item: Copy + Default,
593    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM> + ContainerType,
594{
595    type OutputImpl =
596        <<ArrayImpl as ContainerType>::Type as RowMajorArrayFromContainerType>::OutputImpl<
597            Item,
598            NDIM,
599        >;
600
601    fn eval_row_major(&self) -> Array<Self::OutputImpl, NDIM> {
602        let mut out =
603            <<ArrayImpl as ContainerType>::Type as RowMajorArrayFromContainerType>::new_row_major_arr::<
604                Item,
605                NDIM,
606            >(self.shape());
607
608        out.fill_from(self);
609
610        out
611    }
612}
613
614/// Create an empty array of given type and dimension.
615///
616/// Empty arrays serve as convenient containers for input into functions that
617/// resize an array before filling it with data.
618pub fn empty_array<Item, const NDIM: usize>() -> Array<BaseArray<VectorContainer<Item>, NDIM>, NDIM>
619where
620    Item: Copy + Default,
621{
622    let shape = [0; NDIM];
623    let container = VectorContainer::new(0);
624    Array::new(BaseArray::new(container, shape))
625}
626
627impl<'a, Item, const NDIM: usize> SliceArray<'a, Item, NDIM>
628where
629    Item: Copy + Default,
630{
631    /// Create a new array from `slice` with a given `shape`.
632    pub fn from_shape(slice: &'a [Item], shape: [usize; NDIM]) -> Self {
633        assert_eq!(
634            slice.len(),
635            shape.iter().product::<usize>(),
636            "Slice length does not match the shape: {} != {}",
637            slice.len(),
638            shape.iter().product::<usize>()
639        );
640        Array::new(BaseArray::new(SliceContainer::new(slice), shape))
641    }
642}
643
644impl<'a, Item, const NDIM: usize> SliceArrayMut<'a, Item, NDIM>
645where
646    Item: Copy + Default,
647{
648    /// Create a new array from mutable `slice` with a given `shape`.
649    pub fn from_shape(slice: &'a mut [Item], shape: [usize; NDIM]) -> Self {
650        assert_eq!(
651            slice.len(),
652            shape.iter().product::<usize>(),
653            "Slice length does not match the shape: {} != {}",
654            slice.len(),
655            shape.iter().product::<usize>()
656        );
657        Array::new(BaseArray::new(SliceContainerMut::new(slice), shape))
658    }
659}
660
661impl<'a, Item, const NDIM: usize> StridedSliceArray<'a, Item, NDIM>
662where
663    Item: Copy + Default,
664{
665    /// Create a new array from `slice` with a given `shape` and `stride`.
666    pub fn from_shape_and_stride(
667        slice: &'a [Item],
668        shape: [usize; NDIM],
669        stride: [usize; NDIM],
670    ) -> Self {
671        Array::new(StridedBaseArray::new(
672            SliceContainer::new(slice),
673            shape,
674            stride,
675        ))
676    }
677}
678
679impl<'a, Item, const NDIM: usize> StridedSliceArrayMut<'a, Item, NDIM>
680where
681    Item: Copy + Default,
682{
683    /// Create a new array from `slice` with a given `shape` and `stride`.
684    pub fn from_shape_and_stride(
685        slice: &'a mut [Item],
686        shape: [usize; NDIM],
687        stride: [usize; NDIM],
688    ) -> Self {
689        Array::new(StridedBaseArray::new(
690            SliceContainerMut::new(slice),
691            shape,
692            stride,
693        ))
694    }
695}
696
697impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> std::fmt::Debug for Array<ArrayImpl, NDIM> {
698    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
699        let shape = self.shape();
700        write!(f, "Array[").unwrap();
701        for item in shape {
702            write!(f, "{item},").unwrap();
703        }
704        write!(f, "]")
705    }
706}
707
708impl<Item, const NDIM: usize> DynArray<Item, NDIM>
709where
710    Item: Copy + Default,
711{
712    /// Create a new array and fill with values from `arr`.
713    pub fn new_from<ArrayImpl>(arr: &Array<ArrayImpl, NDIM>) -> Self
714    where
715        ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
716    {
717        let mut output = empty_array::<Item, NDIM>();
718        output.fill_from_resize(arr);
719        output
720    }
721}
722
723impl<Item, const NDIM: usize> StridedDynArray<Item, NDIM>
724where
725    Item: Copy + Default,
726{
727    /// Create a new strided array with `stride` and fill with values from `arr`.
728    pub fn new_from<ArrayImpl>(stride: [usize; NDIM], arr: &Array<ArrayImpl, NDIM>) -> Self
729    where
730        ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
731    {
732        let mut output = StridedDynArray::from_shape_and_stride(arr.shape(), stride);
733        output.fill_from(arr);
734        output
735    }
736
737    /// Create a new row-major array from existing array `arr`.
738    pub fn row_major_from<ArrayImpl>(arr: &Array<ArrayImpl, NDIM>) -> Self
739    where
740        ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
741    {
742        StridedDynArray::new_from(row_major_stride_from_shape(arr.shape()), arr)
743    }
744}
745
746impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
747where
748    ArrayImpl: UnsafeRandom1DAccessByValue + Shape<NDIM>,
749{
750    /// Iterate through the array by value.
751    ///
752    /// The iterator always proceeds in column-major order.
753    ///
754    /// # Traits
755    /// - [UnsafeRandom1DAccessByValue]
756    /// - [Shape]
757    #[inline(always)]
758    pub fn iter_value(&self) -> ArrayDefaultIteratorByValue<'_, ArrayImpl, NDIM> {
759        ArrayDefaultIteratorByValue::new(self)
760    }
761}
762
763impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
764where
765    ArrayImpl: UnsafeRandom1DAccessByRef + Shape<NDIM>,
766{
767    /// Iterate through the array by reference.
768    ///
769    /// The iterator always proceeds in column-major order.
770    ///
771    /// # Traits
772    /// - [UnsafeRandom1DAccessByRef]
773    /// - [Shape]
774    #[inline(always)]
775    pub fn iter_ref(&self) -> ArrayDefaultIteratorByRef<'_, ArrayImpl, NDIM> {
776        ArrayDefaultIteratorByRef::new(self)
777    }
778}
779
780impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
781where
782    ArrayImpl: UnsafeRandom1DAccessMut + Shape<NDIM>,
783{
784    /// Iterate through the array by mutable reference.
785    ///
786    /// The iterator always proceeds in column-major order.
787    ///
788    /// # Traits
789    /// - [UnsafeRandom1DAccessMut]
790    /// - [Shape]
791    #[inline(always)]
792    pub fn iter_mut(&mut self) -> ArrayDefaultIteratorMut<'_, ArrayImpl, NDIM> {
793        ArrayDefaultIteratorMut::new(self)
794    }
795}
796
797impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
798where
799    ArrayImpl: UnsafeRandomAccessByValue<NDIM, Item = Item> + Shape<NDIM>,
800{
801    /// Iterate through the diagonal of the array by value.
802    ///
803    /// # Traits
804    /// - [UnsafeRandomAccessByValue]
805    /// - [Shape]
806    #[inline(always)]
807    pub fn diag_iter_value(&self) -> ArrayDiagIteratorByValue<'_, ArrayImpl, NDIM> {
808        ArrayDiagIteratorByValue::new(self)
809    }
810}
811
812impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
813where
814    ArrayImpl: UnsafeRandomAccessByRef<NDIM, Item = Item> + Shape<NDIM>,
815{
816    /// Iterate through the diagonal of the array by reference.
817    ///
818    /// # Traits
819    /// - [UnsafeRandomAccessByRef]
820    /// - [Shape]
821    #[inline(always)]
822    pub fn diag_iter_ref(&self) -> ArrayDiagIteratorByRef<'_, ArrayImpl, NDIM> {
823        ArrayDiagIteratorByRef::new(self)
824    }
825}
826
827impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
828where
829    ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item> + Shape<NDIM>,
830{
831    /// Iterate through the diagonal of the array by mutable reference.
832    ///
833    /// # Traits
834    /// - [UnsafeRandomAccessByRef]
835    /// - [Shape]
836    #[inline(always)]
837    pub fn diag_iter_mut(&mut self) -> ArrayDiagIteratorMut<'_, ArrayImpl, NDIM> {
838        ArrayDiagIteratorMut::new(self)
839    }
840}
841
842impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
843where
844    ArrayImpl: UnsafeRandom1DAccessMut<Item = Item> + Shape<NDIM>,
845{
846    /// Fill from another array.
847    ///
848    /// #Traits
849    /// - [UnsafeRandomAccessMut]
850    /// - [Shape]
851    #[inline(always)]
852    pub fn fill_from<ArrayImplOther>(&mut self, other: &Array<ArrayImplOther, NDIM>)
853    where
854        ArrayImplOther: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
855    {
856        assert_eq!(self.shape(), other.shape());
857
858        for (item, other_item) in self.iter_mut().zip(other.iter_value()) {
859            *item = other_item;
860        }
861    }
862}
863
864impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
865where
866    Const<NDIM>: ToUInt,
867    <Const<NDIM> as ToUInt>::Output: Unsigned + NonZero,
868    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item>
869        + UnsafeRandomAccessByValue<NDIM, Item = Item>
870        + Shape<NDIM>,
871{
872    /// Create a new multislice iterator.
873    ///
874    /// Given an array `arr` with dimensions `[2, 3, 4, 5]`. Calling
875    /// `self.multislice_iter` with `axes = [0, 2]` returns an iterator
876    /// over values of the type `(indices, slice)`, where indices is a
877    /// [2; usize] array taking all possible axis indices in axes `0` and `2`,
878    /// and `slice` is the output of `arr.multislice([0, 2], [indices])`.
879    pub fn multislice_iter<const SDIM: usize, const OUTDIM: usize>(
880        &self,
881        axes: [usize; SDIM],
882    ) -> MultisliceIterator<'_, Item, ArrayImpl, NDIM, SDIM, OUTDIM>
883    where
884        Const<SDIM>: ToUInt,
885        Const<OUTDIM>: ToUInt,
886        <Const<SDIM> as ToUInt>::Output: Unsigned + NonZero,
887        <Const<OUTDIM> as ToUInt>::Output: Unsigned + NonZero,
888        NumberType<SDIM>: IsSmallerThan<NDIM>,
889        PInt<U<OUTDIM>>: Add<PInt<U<SDIM>>, Output = PInt<U<NDIM>>>,
890        PInt<U<NDIM>>: Sub<PInt<U<SDIM>>>,
891    {
892        MultisliceIterator::new(self, axes)
893    }
894}
895
896impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
897where
898    Const<NDIM>: ToUInt,
899    <Const<NDIM> as ToUInt>::Output: Unsigned + NonZero,
900    ArrayImpl: UnsafeRandom1DAccessMut<Item = Item>
901        + UnsafeRandomAccessMut<NDIM, Item = Item>
902        + Shape<NDIM>,
903{
904    /// Create a new mutable multislice iterator.
905    ///
906    /// Given an array `arr` with dimensions `[2, 3, 4, 5]`. Calling
907    /// `self.multislice_iter_mut` with `axes = [0, 2]` returns an iterator
908    /// over values of the type `(indices, slice)`, where indices is a
909    /// [2; usize] array taking all possible axis indices in axes `0` and `2`,
910    /// and `slice` is the output of `arr.multislice([0, 2], [indices])`.
911    pub fn multislice_iter_mut<const SDIM: usize, const OUTDIM: usize>(
912        &mut self,
913        axes: [usize; SDIM],
914    ) -> MultisliceIteratorMut<'_, Item, ArrayImpl, NDIM, SDIM, OUTDIM>
915    where
916        Const<SDIM>: ToUInt,
917        Const<OUTDIM>: ToUInt,
918        <Const<SDIM> as ToUInt>::Output: Unsigned + NonZero,
919        <Const<OUTDIM> as ToUInt>::Output: Unsigned + NonZero,
920        NumberType<SDIM>: IsSmallerThan<NDIM>,
921        PInt<U<OUTDIM>>: Add<PInt<U<SDIM>>, Output = PInt<U<NDIM>>>,
922        PInt<U<NDIM>>: Sub<PInt<U<SDIM>>>,
923    {
924        MultisliceIteratorMut::new(self, axes)
925    }
926}
927
928impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
929where
930    ArrayImpl: UnsafeRandom1DAccessMut + Shape<NDIM>,
931{
932    /// Fill from an iterator.
933    ///
934    /// Note that the array only fills as many values as the other iterator provides.
935    ///
936    /// # Traits
937    /// - [UnsafeRandom1DAccessMut]
938    /// - [Shape]
939    #[inline(always)]
940    pub fn fill_from_iter<Iter>(&mut self, iter: Iter)
941    where
942        Iter: Iterator<Item = ArrayImpl::Item>,
943    {
944        for (item, other_item) in izip!(self.iter_mut(), iter) {
945            *item = other_item;
946        }
947    }
948}
949
950impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
951where
952    ArrayImpl: ResizeInPlace<NDIM> + Shape<NDIM> + UnsafeRandom1DAccessMut,
953{
954    /// Fill from another array and resize if necessary.
955    ///
956    /// This method is especially useful together with [empty_array].
957    ///
958    /// #Traits
959    /// - [UnsafeRandom1DAccessMut]
960    /// - [Shape]
961    #[inline(always)]
962    pub fn fill_from_resize<ArrayImplOther>(&mut self, other: &Array<ArrayImplOther, NDIM>)
963    where
964        ArrayImplOther: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = ArrayImpl::Item>,
965    {
966        self.resize_in_place(other.shape());
967        self.fill_from(other);
968    }
969}
970
971impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
972where
973    ArrayImpl: UnsafeRandom1DAccessMut + Shape<NDIM>,
974{
975    /// Fill array with a given value.
976    ///
977    /// # Traits
978    /// - [UnsafeRandom1DAccessMut]
979    /// - [Shape]
980    #[inline(always)]
981    pub fn fill_with_value(&mut self, value: ArrayImpl::Item) {
982        for item in self.iter_mut() {
983            *item = value;
984        }
985    }
986}
987
988impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
989where
990    ArrayImpl: UnsafeRandom1DAccessMut + Shape<NDIM>,
991{
992    /// Set the elements of the array to zero.
993    pub fn set_zero(&mut self) {
994        for item in self.iter_mut() {
995            *item = Default::default();
996        }
997    }
998}
999
1000impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1001where
1002    ArrayImpl: UnsafeRandom1DAccessMut + UnsafeRandomAccessMut<NDIM> + Shape<NDIM>,
1003    ArrayImpl::Item: One,
1004{
1005    /// Set all off-diagonal elements to zero and the diagonal to one.
1006    pub fn set_identity(&mut self) {
1007        self.set_zero();
1008        self.diag_iter_mut().for_each(|elem| *elem = One::one());
1009    }
1010}
1011
1012impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1013where
1014    ArrayImpl: UnsafeRandomAccessByValue<NDIM> + Shape<NDIM>,
1015    ArrayImpl::Item: std::iter::Sum,
1016{
1017    /// Compute the sum of the diagonal values of the array.
1018    ///
1019    /// Note: The Item type must support the trait [std::iter::Sum].
1020    ///
1021    /// # Traits
1022    /// - [UnsafeRandomAccessByValue]
1023    /// - [Shape]
1024    #[inline(always)]
1025    pub fn trace(&self) -> ArrayImpl::Item {
1026        self.diag_iter_value().sum()
1027    }
1028}
1029
1030impl<Item, ArrayImpl> Array<ArrayImpl, 1>
1031where
1032    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
1033    ArrayImpl::Item: Conj<Output = Item>
1034        + std::ops::Mul<Output = ArrayImpl::Item>
1035        + std::ops::Add<Output = Item>,
1036{
1037    /// Compute the inner product of two 1d arrays.
1038    ///
1039    /// The elements of `other` are taken as conjugate.
1040    ///
1041    /// Returns `None` if the arrays are empty.
1042    ///
1043    /// Note: The Item type must support the traits [Conj], [std::iter::Sum] and [std::ops::Mul].
1044    ///
1045    /// # Traits
1046    /// - [UnsafeRandom1DAccessByValue]
1047    /// - [Shape]
1048    #[inline(always)]
1049    pub fn inner<ArrayImplOther>(&self, other: &Array<ArrayImplOther, 1>) -> Option<Item>
1050    where
1051        ArrayImplOther: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
1052    {
1053        assert_eq!(self.shape(), other.shape());
1054        izip!(self.iter_value(), other.iter_value())
1055            .map(|(x, y)| x * y.conj())
1056            .reduce(std::ops::Add::add)
1057    }
1058}
1059
1060impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1061where
1062    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
1063    Item: Abs,
1064    <Item as Abs>::Output: Max<Output = <Item as Abs>::Output>,
1065{
1066    /// Compute the maximum absolute value over all elements.
1067    ///
1068    /// Note: The item type must support [Abs] and the output of [Abs]
1069    /// must support [Max].
1070    ///
1071    /// The function returns `None` if the array is empty.
1072    ///
1073    /// # Traits
1074    /// - [UnsafeRandom1DAccessByValue]
1075    /// - [Shape]
1076    #[inline(always)]
1077    pub fn max_abs(&self) -> Option<<Item as Abs>::Output> {
1078        self.iter_value().map(|elem| elem.abs()).reduce(Max::max)
1079    }
1080}
1081
1082impl<Item, ArrayImpl> Array<ArrayImpl, 1>
1083where
1084    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
1085    Item: Abs,
1086    <Item as Abs>::Output: std::ops::Add<Output = <Item as Abs>::Output>,
1087{
1088    /// Compute the 1-norm of a 1d array.
1089    ///
1090    /// The function returns `None` if the array is empty.
1091    ///
1092    /// Note: The item type must support [Abs] and the output of [Abs] must
1093    /// support [std::ops::Add].
1094    ///
1095    ///
1096    /// # Traits
1097    /// - [UnsafeRandom1DAccessByValue]
1098    /// - [Shape]
1099    #[inline(always)]
1100    pub fn norm_1(&self) -> Option<<Item as Abs>::Output> {
1101        self.iter_value()
1102            .map(|elem| elem.abs())
1103            .reduce(std::ops::Add::add)
1104    }
1105}
1106
1107impl<Item, ArrayImpl> Array<ArrayImpl, 1>
1108where
1109    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
1110    Item: AbsSquare,
1111    <Item as AbsSquare>::Output: Sqrt<Output = <Item as AbsSquare>::Output>
1112        + std::ops::Add<Output = <Item as AbsSquare>::Output>,
1113{
1114    /// Compute the 2-norm of a 1d array.
1115    ///
1116    /// Note: The item type must support [AbsSquare]. The output of [AbsSquare]
1117    /// must support [std::ops::Add] and [Sqrt].
1118    ///
1119    /// # Traits
1120    /// - [UnsafeRandom1DAccessByValue]
1121    /// - [Shape]
1122    #[inline(always)]
1123    pub fn norm_2(&self) -> Option<<Item as AbsSquare>::Output> {
1124        self.iter_value()
1125            .map(|elem| elem.abs_square())
1126            .reduce(std::ops::Add::add)
1127            .map(|elem| elem.sqrt())
1128    }
1129}
1130
1131impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1132where
1133    ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
1134    Item: AbsSquare,
1135    <Item as AbsSquare>::Output: Sqrt<Output = <Item as AbsSquare>::Output>
1136        + std::ops::Add<Output = <Item as AbsSquare>::Output>
1137        + Copy
1138        + Default,
1139{
1140    /// Compute the Frobenius-Norm of a nd array.
1141    ///
1142    /// Note: The item type must support [AbsSquare]. The output of [AbsSquare]
1143    /// must support [std::ops::Add], [Sqrt], [Copy], and [Default].
1144    ///
1145    /// # Traits
1146    /// - [UnsafeRandom1DAccessByValue]
1147    /// - [Shape]
1148    #[inline(always)]
1149    pub fn norm_fro(&self) -> Option<<Item as AbsSquare>::Output> {
1150        self.r()
1151            .abs_square()
1152            .iter_value()
1153            .reduce(std::ops::Add::add)
1154            .map(|elem| elem.sqrt())
1155    }
1156}
1157
1158impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1159where
1160    ArrayImpl: UnsafeRandom1DAccessByValue + Shape<NDIM>,
1161{
1162    /// Convert an array into the new item type T.
1163    ///
1164    /// Note: It is required that `ArrayImpl::Item: Into<T>`.
1165    ///
1166    /// # Traits
1167    /// - [UnsafeRandom1DAccessByValue]
1168    /// - [Shape]
1169    #[allow(clippy::type_complexity)]
1170    pub fn into_type<T>(
1171        self,
1172    ) -> Array<
1173        ArrayUnaryOperator<ArrayImpl::Item, T, ArrayImpl, fn(ArrayImpl::Item) -> T, NDIM>,
1174        NDIM,
1175    >
1176    where
1177        ArrayImpl::Item: Into<T>,
1178    {
1179        Array::new(ArrayUnaryOperator::new(self, |item| item.into()))
1180    }
1181}
1182
1183impl<ArrayImpl> Array<ArrayImpl, 2>
1184where
1185    ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
1186{
1187    /// Return a column iterator for a 2d array.
1188    ///
1189    /// # Traits
1190    /// - [UnsafeRandomAccessByValue]
1191    /// - [Shape]
1192    pub fn col_iter(&self) -> ColIterator<'_, ArrayImpl, 2> {
1193        ColIterator::new(self)
1194    }
1195}
1196
1197impl<ArrayImpl> Array<ArrayImpl, 2>
1198where
1199    ArrayImpl: Shape<2> + UnsafeRandomAccessMut<2>,
1200{
1201    /// Return a mutable column iterator for a 2d array.
1202    ///
1203    /// # Traits
1204    /// - [UnsafeRandomAccessMut]
1205    /// - [Shape]
1206    pub fn col_iter_mut(&mut self) -> ColIteratorMut<'_, ArrayImpl, 2> {
1207        ColIteratorMut::new(self)
1208    }
1209}
1210
1211impl<ArrayImpl> Array<ArrayImpl, 2>
1212where
1213    ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
1214{
1215    /// Return a row iterator for a 2d array.
1216    ///
1217    /// # Traits
1218    /// - [UnsafeRandomAccessByValue]
1219    /// - [Shape]
1220    pub fn row_iter(&self) -> RowIterator<'_, ArrayImpl, 2> {
1221        RowIterator::new(self)
1222    }
1223}
1224
1225impl<ArrayImpl> Array<ArrayImpl, 2>
1226where
1227    ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
1228{
1229    /// Return a mutable row iterator for a 2d array.
1230    ///
1231    /// # Traits
1232    /// - [UnsafeRandomAccessMut]
1233    /// - [Shape]
1234    pub fn row_iter_mut(&mut self) -> RowIteratorMut<'_, ArrayImpl, 2> {
1235        RowIteratorMut::new(self)
1236    }
1237}
1238
1239impl<ArrayImpl> Array<ArrayImpl, 2>
1240where
1241    ArrayImpl: UnsafeRandom1DAccessByValue + Shape<2>,
1242{
1243    /// Return an iterator of the form `(i, j, data)`.
1244    ///
1245    /// Here, `i` is the row, `j` is the column, and `data` is the associated element
1246    /// returned by value.
1247    ///
1248    /// # Traits
1249    /// - [UnsafeRandom1DAccessByValue]
1250    /// - [Shape]
1251    pub fn iter_aij_value(&self) -> impl Iterator<Item = ([usize; 2], ArrayImpl::Item)> + '_ {
1252        let iter = ArrayDefaultIteratorByValue::new(self);
1253        AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), self.shape())
1254    }
1255}
1256
1257impl<ArrayImpl> Array<ArrayImpl, 2>
1258where
1259    ArrayImpl: UnsafeRandom1DAccessByRef + Shape<2>,
1260{
1261    /// Return an iterator of the form `(i, j, &data)`.
1262    ///
1263    /// Here, `i` is the row, `j` is the column, and `&data` is the reference to the associated
1264    /// element.
1265    ///
1266    /// # Traits
1267    /// - [UnsafeRandom1DAccessByRef]
1268    /// - [Shape]
1269    pub fn iter_aij_ref(&self) -> impl Iterator<Item = ([usize; 2], &ArrayImpl::Item)> + '_ {
1270        let iter = ArrayDefaultIteratorByRef::new(self);
1271        AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), self.shape())
1272    }
1273}
1274
1275impl<ArrayImpl> Array<ArrayImpl, 2>
1276where
1277    ArrayImpl: UnsafeRandom1DAccessMut + Shape<2>,
1278{
1279    /// Return an iterator of the form `(i, j, &mut data)`.
1280    ///
1281    /// Here, `i` is the row, `j` is the column, and `&data` is the mutable reference to the associated
1282    /// element.
1283    ///
1284    /// # Traits
1285    /// - [UnsafeRandom1DAccessByRef]
1286    /// - [Shape]
1287    pub fn iter_aij_mut(
1288        &mut self,
1289    ) -> impl Iterator<Item = ([usize; 2], &mut ArrayImpl::Item)> + '_ {
1290        let shape = self.shape();
1291        let iter = ArrayDefaultIteratorMut::new(self);
1292        AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), shape)
1293    }
1294}
1295
1296impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
1297    /// Cast array to type `T`.
1298    ///
1299    /// The cast is done through [num::cast::cast] and source and target types need to
1300    /// support casting through that function.
1301    pub fn cast<Target>(self) -> Array<ArrayCast<Target, ArrayImpl, NDIM>, NDIM> {
1302        Array::new(ArrayCast::new(self))
1303    }
1304}
1305
1306impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
1307    /// Coerce the array to a specific dimension.
1308    ///
1309    /// This is useful to coerce from a generic dimension parameter to
1310    /// a specific number of dimensions.
1311    pub fn coerce_dim<const CDIM: usize>(
1312        self,
1313    ) -> RlstResult<Array<CoerceArray<ArrayImpl, NDIM, CDIM>, CDIM>> {
1314        if CDIM == NDIM {
1315            // If the dimensions and item types match return a CoerceArray
1316            Ok(Array::new(CoerceArray::new(self)))
1317        } else {
1318            // Otherwise, we need to coerce the array.
1319            Err(RlstError::GeneralError(
1320                format!("Cannot coerce array: dimensions do not match {CDIM} != {NDIM}.")
1321                    .to_string(),
1322            ))
1323        }
1324    }
1325}
1326
1327impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1328where
1329    ArrayImpl: Shape<NDIM>,
1330{
1331    /// Permute axes of an array.
1332    ///
1333    /// The `permutation` gives the new ordering of the axes.
1334    pub fn permute_axes(
1335        self,
1336        permutation: [usize; NDIM],
1337    ) -> Array<ArrayTranspose<ArrayImpl, NDIM>, NDIM> {
1338        Array::new(ArrayTranspose::new(self, permutation))
1339    }
1340
1341    /// Transpose an array.
1342    ///
1343    /// The transpose of an n-dimensional array reverses the order of the axes.
1344    pub fn transpose(self) -> Array<ArrayTranspose<ArrayImpl, NDIM>, NDIM> {
1345        let mut permutation = [0; NDIM];
1346
1347        for (ind, p) in (0..NDIM).rev().enumerate() {
1348            permutation[ind] = p;
1349        }
1350
1351        self.permute_axes(permutation)
1352    }
1353}
1354
1355impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1356where
1357    Item: RandScalar + RlstScalar,
1358    ArrayImpl: UnsafeRandom1DAccessMut<Item = Item> + Shape<NDIM>,
1359    StandardNormal: Distribution<Item::Real>,
1360    StandardUniform: Distribution<Item::Real>,
1361{
1362    /// Fill an array with normally distributed random numbers.
1363    ///
1364    /// # Traits
1365    /// - [UnsafeRandom1DAccessMut]
1366    /// - [Shape]
1367    /// - Item: [RandScalar] + [RlstScalar]
1368    pub fn fill_from_standard_normal<R: Rng>(&mut self, rng: &mut R) {
1369        let dist = StandardNormal;
1370        self.iter_mut()
1371            .for_each(|val| *val = <Item>::random_scalar(rng, &dist));
1372    }
1373
1374    /// Fill an array with equally distributed random numbers.
1375    ///
1376    /// # Traits
1377    /// - [UnsafeRandom1DAccessMut]
1378    /// - [Shape]
1379    /// - Item: [RandScalar] + [RlstScalar]
1380    pub fn fill_from_equally_distributed<R: Rng>(&mut self, rng: &mut R) {
1381        let dist = StandardUniform;
1382        self.iter_mut()
1383            .for_each(|val| *val = <Item>::random_scalar(rng, &dist));
1384    }
1385
1386    /// Fill an array with equally distributed random numbers using a given `seed`.
1387    ///
1388    /// # Traits
1389    /// - [UnsafeRandom1DAccessMut]
1390    /// - [Shape]
1391    /// - Item: [RandScalar] + [RlstScalar]
1392    pub fn fill_from_seed_equally_distributed(&mut self, seed: usize) {
1393        let mut rng = ChaCha8Rng::seed_from_u64(seed as u64);
1394        let dist = StandardUniform;
1395        self.iter_mut()
1396            .for_each(|val| *val = <Item>::random_scalar(&mut rng, &dist));
1397    }
1398
1399    /// Fill an array with normally distributed random numbers using a given `seed`.
1400    ///
1401    /// # Traits
1402    /// - [UnsafeRandom1DAccessMut]
1403    /// - [Shape]
1404    /// - Item: [RandScalar] + [RlstScalar]
1405    pub fn fill_from_seed_normally_distributed(&mut self, seed: usize) {
1406        let mut rng = ChaCha8Rng::seed_from_u64(seed as u64);
1407        let dist = StandardNormal;
1408        self.iter_mut()
1409            .for_each(|val| *val = <Item>::random_scalar(&mut rng, &dist));
1410    }
1411}
1412
1413impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
1414    /// Reverse a single `axis` of the array.
1415    pub fn reverse_axis(self, axis: usize) -> Array<ReverseAxis<ArrayImpl, NDIM>, NDIM> {
1416        assert!(axis < NDIM, "Axis out of bounds");
1417        Array::new(ReverseAxis::new(self, axis))
1418    }
1419}
1420
1421impl<ArrayImpl, const ADIM: usize> Array<ArrayImpl, ADIM>
1422where
1423    ArrayImpl: Shape<ADIM>,
1424{
1425    /// Create a slice from a given array.
1426    ///
1427    /// Consider an array `arr` with shape `[a0, a1, a2, a3, ...]`. The function call
1428    /// `arr.slice(2, 3)` returns a one dimension smaller array indexed by `[a0, a1, 3, a3, ...]`.
1429    /// Hence, the dimension `2` has been fixed to always have the value `3.`
1430    ///
1431    /// # Examples
1432    ///
1433    /// If `arr` is a matrix then the first column of the matrix is obtained from
1434    /// `arr.slice(1, 0)`, while the third row of the matrix is obtained from
1435    /// `arr.slice(0, 2)`.
1436    pub fn slice<const OUTDIM: usize>(
1437        self,
1438        axis: usize,
1439        index: usize,
1440    ) -> Array<ArraySlice<ArrayImpl, ADIM, OUTDIM>, OUTDIM>
1441    where
1442        NumberType<ADIM>: IsGreaterByOne<OUTDIM>,
1443        NumberType<OUTDIM>: IsGreaterZero,
1444    {
1445        Array::new(ArraySlice::new(self, [axis, index]))
1446    }
1447
1448    /// Create a multislice from a given array.
1449    ///
1450    /// A multislice works like [Array::slice] but allows the slicing along multiple dimensions at the same time.
1451    /// The `axes` array determines which axes are fixed and the `indices` array determines the value
1452    /// in each fixed axis. The dimension `SDIM` is the number of axes along which we slice. The parameter `OUTDIM`
1453    /// is the number of dimensions after slicing.
1454    pub fn multislice<const SDIM: usize, const OUTDIM: usize>(
1455        self,
1456        axes: [usize; SDIM],
1457        indices: [usize; SDIM],
1458    ) -> Array<Multislice<ArrayImpl, ADIM, SDIM, OUTDIM>, OUTDIM>
1459    where
1460        NumberType<SDIM>: IsSmallerThan<ADIM>,
1461        NumberDifference<ADIM, SDIM>: SubtractEqual<OUTDIM>,
1462    {
1463        Array::new(Multislice::new(self, axes, indices))
1464    }
1465}
1466
1467impl<ArrayImpl> Array<ArrayImpl, 2>
1468where
1469    ArrayImpl: Shape<2>,
1470{
1471    /// Return the row with index `row_index` from a two-dimensional array.
1472    pub fn row(self, row_index: usize) -> Array<ArraySlice<ArrayImpl, 2, 1>, 1> {
1473        self.slice::<1>(0, row_index)
1474    }
1475
1476    /// Return the column with index `col_index` from a two-dimensional array.
1477    pub fn col(self, col_index: usize) -> Array<ArraySlice<ArrayImpl, 2, 1>, 1> {
1478        self.slice::<1>(1, col_index)
1479    }
1480}
1481
1482impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> Array<ArrayImpl, NDIM> {
1483    /// Move the array into a subview specified by an `offset` and `shape` of the subview.
1484    ///
1485    /// The `offset` is the starting index of the subview and the `shape` is the number of indices
1486    /// in each dimension of the subview.
1487    pub fn into_subview(
1488        self,
1489        offset: [usize; NDIM],
1490        shape: [usize; NDIM],
1491    ) -> Array<ArraySubView<ArrayImpl, NDIM>, NDIM> {
1492        Array::new(ArraySubView::new(self, offset, shape))
1493    }
1494}
1495
1496impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> MulAdd<Item, Array<ArrayImpl2, NDIM>>
1497    for Array<ArrayImpl1, NDIM>
1498where
1499    ArrayImpl1: BaseItem<Item = Item> + Shape<NDIM>,
1500    ArrayImpl2: BaseItem<Item = Item> + Shape<NDIM>,
1501    Item: MulAdd<Output = Item> + Copy,
1502{
1503    type Output = Array<MulAddImpl<ArrayImpl1, ArrayImpl2, Item, NDIM>, NDIM>;
1504    /// Compentwie form `self * a + b`, where `a` is a scalar and `b` is another array.
1505    /// The implementation depdends on the `MulAdd` trait from the `num` crate for the componets of
1506    /// the arrays.
1507    fn mul_add(
1508        self,
1509        a: Item,
1510        b: Array<ArrayImpl2, NDIM>,
1511    ) -> Array<MulAddImpl<ArrayImpl1, ArrayImpl2, Item, NDIM>, NDIM>
1512    where
1513        ArrayImpl1: BaseItem<Item = Item> + Shape<NDIM>,
1514        ArrayImpl2: BaseItem<Item = Item> + Shape<NDIM>,
1515        Item: MulAdd<Output = Item> + Copy,
1516    {
1517        Array::new(MulAddImpl::new(self, b, a))
1518    }
1519}
1520
1521impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Add<Array<ArrayImpl2, NDIM>>
1522    for Array<ArrayImpl1, NDIM>
1523where
1524    ArrayImpl1: Shape<NDIM>,
1525    ArrayImpl2: Shape<NDIM>,
1526{
1527    type Output = Array<ArrayAddition<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
1528
1529    #[inline(always)]
1530    fn add(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
1531        Array::new(ArrayAddition::new(self, rhs))
1532    }
1533}
1534
1535impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::AddAssign<Array<ArrayImpl2, NDIM>>
1536    for Array<ArrayImpl1, NDIM>
1537where
1538    ArrayImpl1: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1539    ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
1540    Item: std::ops::AddAssign,
1541{
1542    #[inline(always)]
1543    fn add_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
1544        assert_eq!(self.shape(), rhs.shape());
1545        for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
1546            *item1 += item2;
1547        }
1548    }
1549}
1550
1551impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Sub<Array<ArrayImpl2, NDIM>>
1552    for Array<ArrayImpl1, NDIM>
1553where
1554    ArrayImpl1: Shape<NDIM>,
1555    ArrayImpl2: Shape<NDIM>,
1556{
1557    type Output = Array<ArraySubtraction<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
1558
1559    fn sub(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
1560        Array::new(ArraySubtraction::new(self, rhs))
1561    }
1562}
1563
1564impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::SubAssign<Array<ArrayImpl2, NDIM>>
1565    for Array<ArrayImpl1, NDIM>
1566where
1567    ArrayImpl1: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1568    ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
1569    Item: std::ops::SubAssign,
1570{
1571    #[inline(always)]
1572    fn sub_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
1573        assert_eq!(self.shape(), rhs.shape());
1574        for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
1575            *item1 -= item2;
1576        }
1577    }
1578}
1579
1580impl<Item, ArrayImpl, const NDIM: usize> std::ops::Neg for Array<ArrayImpl, NDIM>
1581where
1582    ArrayImpl: BaseItem<Item = Item>,
1583    Item: std::ops::Neg<Output = Item>,
1584{
1585    type Output = Array<ArrayNeg<ArrayImpl, NDIM>, NDIM>;
1586
1587    fn neg(self) -> Self::Output {
1588        Array::new(ArrayNeg::new(self))
1589    }
1590}
1591
1592impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Div<Array<ArrayImpl2, NDIM>>
1593    for Array<ArrayImpl1, NDIM>
1594where
1595    ArrayImpl1: Shape<NDIM>,
1596    ArrayImpl2: Shape<NDIM>,
1597{
1598    type Output = Array<CmpWiseDivision<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
1599
1600    fn div(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
1601        Array::new(CmpWiseDivision::new(self, rhs))
1602    }
1603}
1604
1605impl<Item, ArrayImpl, const NDIM: usize> std::ops::Div<Item> for Array<ArrayImpl, NDIM>
1606where
1607    ArrayImpl: Shape<NDIM> + BaseItem<Item = Item>,
1608    Item: Recip<Output = Item> + std::ops::Div<Output = Item>,
1609{
1610    type Output = Array<ArrayScalarMult<Item, ArrayImpl, NDIM>, NDIM>;
1611
1612    fn div(self, rhs: Item) -> Self::Output {
1613        Array::new(ArrayScalarMult::new(rhs.recip(), self))
1614    }
1615}
1616
1617impl<Item, ArrayImpl, const NDIM: usize> std::ops::DivAssign<Item> for Array<ArrayImpl, NDIM>
1618where
1619    ArrayImpl: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1620    Item: std::ops::DivAssign + Copy,
1621{
1622    fn div_assign(&mut self, rhs: Item) {
1623        for item in self.iter_mut() {
1624            *item /= rhs;
1625        }
1626    }
1627}
1628
1629impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Mul<Array<ArrayImpl2, NDIM>>
1630    for Array<ArrayImpl1, NDIM>
1631where
1632    ArrayImpl1: Shape<NDIM>,
1633    ArrayImpl2: Shape<NDIM>,
1634{
1635    type Output = Array<CmpWiseProduct<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
1636
1637    fn mul(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
1638        Array::new(CmpWiseProduct::new(self, rhs))
1639    }
1640}
1641
1642impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::MulAssign<Array<ArrayImpl2, NDIM>>
1643    for Array<ArrayImpl1, NDIM>
1644where
1645    ArrayImpl1: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1646    ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
1647    Item: std::ops::MulAssign,
1648{
1649    fn mul_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
1650        for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
1651            *item1 *= item2;
1652        }
1653    }
1654}
1655
1656impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::DivAssign<Array<ArrayImpl2, NDIM>>
1657    for Array<ArrayImpl1, NDIM>
1658where
1659    ArrayImpl1: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1660    ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
1661    Item: std::ops::DivAssign,
1662{
1663    fn div_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
1664        for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
1665            *item1 /= item2;
1666        }
1667    }
1668}
1669
1670macro_rules! impl_scalar_mult {
1671    ($ScalarType:ty) => {
1672        impl<ArrayImpl, const NDIM: usize> std::ops::Mul<Array<ArrayImpl, NDIM>> for $ScalarType {
1673            type Output = Array<ArrayScalarMult<$ScalarType, ArrayImpl, NDIM>, NDIM>;
1674
1675            fn mul(self, rhs: Array<ArrayImpl, NDIM>) -> Self::Output {
1676                Array::new(ArrayScalarMult::new(self, rhs))
1677            }
1678        }
1679
1680        impl<ArrayImpl, const NDIM: usize> std::ops::Mul<$ScalarType> for Array<ArrayImpl, NDIM>
1681        where
1682            ArrayImpl: Shape<NDIM> + BaseItem<Item = $ScalarType>,
1683        {
1684            type Output = Array<ArrayScalarMult<$ScalarType, ArrayImpl, NDIM>, NDIM>;
1685
1686            fn mul(self, rhs: $ScalarType) -> Self::Output {
1687                Array::new(ArrayScalarMult::new(rhs, self))
1688            }
1689        }
1690    };
1691}
1692
1693impl_scalar_mult!(f64);
1694impl_scalar_mult!(f32);
1695impl_scalar_mult!(c64);
1696impl_scalar_mult!(c32);
1697impl_scalar_mult!(usize);
1698impl_scalar_mult!(i8);
1699impl_scalar_mult!(i16);
1700impl_scalar_mult!(i32);
1701impl_scalar_mult!(i64);
1702impl_scalar_mult!(u8);
1703impl_scalar_mult!(u16);
1704impl_scalar_mult!(u32);
1705impl_scalar_mult!(u64);
1706
1707impl<Item, ArrayImpl, const NDIM: usize> std::ops::MulAssign<Item> for Array<ArrayImpl, NDIM>
1708where
1709    ArrayImpl: Shape<NDIM> + UnsafeRandom1DAccessMut<Item = Item>,
1710    Item: std::ops::MulAssign + Copy,
1711{
1712    fn mul_assign(&mut self, rhs: Item) {
1713        for item in self.iter_mut() {
1714            *item *= rhs;
1715        }
1716    }
1717}
1718
1719impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
1720    /// Create a new array by applying the unitary operator `op` to each element of `self`.
1721    pub fn unary_op<OpItem, OpTarget, Op: Fn(OpItem) -> OpTarget>(
1722        self,
1723        op: Op,
1724    ) -> Array<ArrayUnaryOperator<OpItem, OpTarget, ArrayImpl, Op, NDIM>, NDIM> {
1725        Array::new(ArrayUnaryOperator::new(self, op))
1726    }
1727}
1728
1729impl<Item: Default + Copy + std::ops::Mul<Output = Item>, ArrayImpl, const NDIM: usize>
1730    Array<ArrayImpl, NDIM>
1731where
1732    ArrayImpl: BaseItem<Item = Item>,
1733{
1734    /// Multiple the array with a given `scalar`.
1735    ///
1736    /// Note: The `Item` type must support [std::ops::Mul].
1737    ///
1738    pub fn scalar_mul(self, scalar: Item) -> Array<ArrayScalarMult<Item, ArrayImpl, NDIM>, NDIM> {
1739        Array::new(ArrayScalarMult::new(scalar, self))
1740    }
1741}
1742
1743macro_rules! impl_unary_op_trait {
1744    ($name:ident, $method_name:ident) => {
1745        paste! {
1746
1747        use crate::traits::number_traits::$name;
1748        impl<Item: $name, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1749            where
1750                ArrayImpl: BaseItem<Item = Item>,
1751            {
1752                #[inline(always)]
1753                #[doc = "Componentwise apply `" $method_name "` to the array."]
1754                pub fn $method_name(self) -> Array<ArrayUnaryOperator<Item, <Item as $name>::Output, ArrayImpl, fn(Item) -> <Item as $name>::Output, NDIM>, NDIM> {
1755                    self.unary_op(|x| x.$method_name())
1756                }
1757            }
1758        }
1759
1760    };
1761}
1762
1763impl_unary_op_trait!(Conj, conj);
1764impl_unary_op_trait!(Abs, abs);
1765impl_unary_op_trait!(Square, square);
1766impl_unary_op_trait!(AbsSquare, abs_square);
1767impl_unary_op_trait!(Sqrt, sqrt);
1768impl_unary_op_trait!(Exp, exp);
1769impl_unary_op_trait!(Ln, ln);
1770impl_unary_op_trait!(Recip, recip);
1771impl_unary_op_trait!(Sin, sin);
1772impl_unary_op_trait!(Cos, cos);
1773impl_unary_op_trait!(Tan, tan);
1774impl_unary_op_trait!(Asin, asin);
1775impl_unary_op_trait!(Acos, acos);
1776impl_unary_op_trait!(Atan, atan);
1777impl_unary_op_trait!(Sinh, sinh);
1778impl_unary_op_trait!(Cosh, cosh);
1779impl_unary_op_trait!(Tanh, tanh);
1780impl_unary_op_trait!(Asinh, asinh);
1781impl_unary_op_trait!(Acosh, acosh);
1782impl_unary_op_trait!(Atanh, atanh);
1783
1784impl<Item, ArrayImplX, ArrayImplY, ArrayImpl>
1785    AsMatrixApply<Array<ArrayImplX, 2>, Array<ArrayImplY, 2>> for Array<ArrayImpl, 2>
1786where
1787    Item: Gemm + Copy,
1788    ArrayImplX: RawAccess<Item = Item> + Stride<2> + Shape<2>,
1789    ArrayImplY: RawAccessMut<Item = Item> + Stride<2> + Shape<2>,
1790    ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
1791{
1792    fn apply(
1793        &self,
1794        alpha: Self::Item,
1795        x: &Array<ArrayImplX, 2>,
1796        beta: Self::Item,
1797        y: &mut Array<ArrayImplY, 2>,
1798    ) {
1799        let stride_a = self.stride();
1800        let stride_x = x.stride();
1801        let stride_y = y.stride();
1802
1803        let shape_a = self.shape();
1804        let shape_x = x.shape();
1805        let shape_y = y.shape();
1806
1807        assert_eq!(
1808            shape_a[0], shape_y[0],
1809            "`y` has incompatible shape {shape_y:#?} with `self` {shape_a:#?}."
1810        );
1811
1812        assert_eq!(
1813            shape_a[1], shape_x[0],
1814            "`x` has incompatible shape {shape_x:#?} with `self` {shape_a:#?}."
1815        );
1816
1817        assert_eq!(
1818            shape_y[1], shape_x[1],
1819            "`x` has incompatible shape {shape_x:#?} with `y` {shape_y:#?}."
1820        );
1821
1822        let m = shape_a[0];
1823        let n = shape_x[1];
1824        let k = shape_a[1];
1825
1826        <Item as Gemm>::gemm(
1827            TransMode::NoTrans,
1828            TransMode::NoTrans,
1829            m,
1830            n,
1831            k,
1832            alpha,
1833            self.data().unwrap(),
1834            stride_a[0],
1835            stride_a[1],
1836            x.data().unwrap(),
1837            stride_x[0],
1838            stride_x[1],
1839            beta,
1840            y.data_mut().unwrap(),
1841            stride_y[0],
1842            stride_y[1],
1843        );
1844    }
1845}
1846
1847impl<Item, ArrayImplX, ArrayImplY, ArrayImpl>
1848    AsMatrixApply<Array<ArrayImplX, 1>, Array<ArrayImplY, 1>> for Array<ArrayImpl, 2>
1849where
1850    Item: Gemm + Copy,
1851    ArrayImplX: RawAccess<Item = Item> + Stride<1> + Shape<1>,
1852    ArrayImplY: RawAccessMut<Item = Item> + Stride<1> + Shape<1>,
1853    ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
1854{
1855    fn apply(
1856        &self,
1857        alpha: Self::Item,
1858        x: &Array<ArrayImplX, 1>,
1859        beta: Self::Item,
1860        y: &mut Array<ArrayImplY, 1>,
1861    ) {
1862        self.apply(
1863            alpha,
1864            &x.r().insert_empty_axis::<2>(AxisPosition::Back),
1865            beta,
1866            &mut y.r_mut().insert_empty_axis::<2>(AxisPosition::Back),
1867        );
1868    }
1869}
1870
1871impl<Item, ArrayImpl> Array<ArrayImpl, 2>
1872where
1873    ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
1874    Item: Gemm + Copy + Default + One + Zero,
1875{
1876    /// Compute the matrix-matrix product of `self` with `other`.
1877    pub fn dot<ArrayImplOther, const NDIM: usize>(
1878        &self,
1879        other: &Array<ArrayImplOther, NDIM>,
1880    ) -> DynArray<Item, NDIM>
1881    where
1882        ArrayImplOther: RawAccess<Item = Item> + Shape<NDIM> + Stride<NDIM>,
1883    {
1884        let a_shape = self.shape();
1885        let x_shape = other.shape();
1886
1887        let mut out = empty_array::<Item, NDIM>();
1888
1889        if NDIM == 1 {
1890            let mut out = out.r_mut().coerce_dim::<1>().unwrap();
1891            out.resize_in_place([a_shape[0]]);
1892
1893            self.apply(
1894                Item::one(),
1895                &other.r().coerce_dim::<1>().unwrap(),
1896                Item::zero(),
1897                &mut out,
1898            );
1899        } else if NDIM == 2 {
1900            let mut out = out.r_mut().coerce_dim::<2>().unwrap();
1901            out.resize_in_place([a_shape[0], x_shape[1]]);
1902
1903            self.apply(
1904                Item::one(),
1905                &other.r().coerce_dim::<2>().unwrap(),
1906                Item::zero(),
1907                &mut out,
1908            )
1909        } else {
1910            panic!("NDIM = {NDIM} not supported.");
1911        }
1912
1913        out
1914    }
1915}
1916
1917impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1918where
1919    ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item>,
1920    Item: Copy,
1921{
1922    /// Swap the elements at the positions `pos1` and `pos2`
1923    ///
1924    /// # Safety
1925    ///
1926    /// `pos1` and `pos2` must not be out of bounds.
1927    #[inline(always)]
1928    pub unsafe fn swap_unsafe(&mut self, pos1: [usize; NDIM], pos2: [usize; NDIM]) {
1929        let elem1 = *unsafe { self.get_unchecked_mut(pos1) };
1930        let elem2 = *unsafe { self.get_unchecked_mut(pos2) };
1931
1932        *unsafe { self.get_unchecked_mut(pos1) } = elem2;
1933        *unsafe { self.get_unchecked_mut(pos2) } = elem1;
1934    }
1935}
1936
1937impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
1938where
1939    ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item> + Shape<NDIM>,
1940    Item: Copy,
1941{
1942    /// Swap the elements at the positions `pos1` and `pos2`.
1943    ///
1944    /// Returns an error if either `pos1` or `pos2` is out of bounds.
1945    #[inline(always)]
1946    pub fn swap(&mut self, pos1: [usize; NDIM], pos2: [usize; NDIM]) -> RlstResult<()> {
1947        let shape = self.shape();
1948        if !check_multi_index_in_bounds(pos1, shape) {
1949            Err(RlstError::GeneralError(format!(
1950                "Index `pos1` = {:#?} out of bounds.",
1951                pos1
1952            )))
1953        } else if !check_multi_index_in_bounds(pos2, shape) {
1954            Err(RlstError::GeneralError(format!(
1955                "Index `pos2` = {:#?} out of bounds.",
1956                pos2
1957            )))
1958        } else {
1959            unsafe {
1960                self.swap_unsafe(pos1, pos2);
1961            }
1962            Ok(())
1963        }
1964    }
1965}
1966
1967impl<Item> ScalarAsArray for Item
1968where
1969    Item: Default + Copy,
1970{
1971    fn as_array<const NDIM: usize>(&self) -> StaticArray<Self, NDIM, 1> {
1972        let mut arr = Array::new(BaseArray::new(ArrayContainer::<Item, 1>::new(), [1; NDIM]));
1973        *unsafe { arr.get_unchecked_mut([0; NDIM]) } = *self;
1974        arr
1975    }
1976}
1977
1978impl<Item, const NDIM: usize> StaticArray<Item, NDIM, 1>
1979where
1980    Item: Copy + Default,
1981{
1982    /// Create an empty 1-element array of the given dimension.
1983    pub fn one_element_array() -> Self {
1984        Array::new(BaseArray::new(ArrayContainer::<Item, 1>::new(), [1; NDIM]))
1985    }
1986}
1987
1988macro_rules! impl_linspace_float {
1989    ($dtype:ty) => {
1990        impl<ArrayImpl> Linspace<$dtype> for Array<ArrayImpl, 1>
1991        where
1992            ArrayImpl: MutableArrayImpl<$dtype, 1>,
1993        {
1994            fn linspace(&mut self, a: $dtype, b: $dtype) {
1995                let n = self.len() - 1;
1996
1997                assert!(b > a, "`linspace` requires `{b} = b > a = {a} ");
1998                assert!(n > 1, "`linspace` requires `{n} = self.len() > 1");
1999
2000                let delta = (b - a) / (n as $dtype);
2001
2002                for (ind, elem) in self.iter_mut().enumerate().take(n) {
2003                    *elem = a + delta * ind as $dtype;
2004                }
2005
2006                self[[n]] = b;
2007            }
2008        }
2009    };
2010}
2011
2012impl_linspace_float!(f32);
2013impl_linspace_float!(f64);
2014
2015#[cfg(test)]
2016mod test {
2017
2018    use approx::assert_relative_eq;
2019
2020    use crate::assert_array_relative_eq;
2021
2022    use super::*;
2023
2024    #[test]
2025    pub fn test_add() {
2026        let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2027        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2028        let expected: DynArray<_, 1> = vec![2.5, 4.3, 7.0].into();
2029
2030        crate::assert_array_relative_eq!((a.r() + b.r()).eval(), expected, 1E-10);
2031    }
2032
2033    #[test]
2034    pub fn test_add_assign() {
2035        let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2036        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2037        let expected: DynArray<_, 1> = vec![2.5, 4.3, 7.0].into();
2038
2039        a += b;
2040
2041        crate::assert_array_relative_eq!(a, expected, 1E-10);
2042    }
2043
2044    #[test]
2045    pub fn test_sub() {
2046        let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2047        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2048        let expected: DynArray<_, 1> = vec![-0.5, -0.3, -1.0].into();
2049
2050        crate::assert_array_relative_eq!((a.r() - b.r()).eval(), expected, 1E-10);
2051    }
2052
2053    #[test]
2054    pub fn test_sub_assign() {
2055        let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2056        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2057        let expected: DynArray<_, 1> = vec![-0.5, -0.3, -1.0].into();
2058
2059        a -= b;
2060
2061        crate::assert_array_relative_eq!(a, expected, 1E-10);
2062    }
2063
2064    #[test]
2065    pub fn test_cmp_mul() {
2066        let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2067        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2068        let expected: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
2069
2070        crate::assert_array_relative_eq!(a.r() * b.r(), expected, 1E-10);
2071    }
2072
2073    #[test]
2074    pub fn test_cmp_mul_assign() {
2075        let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2076        let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2077        let expected: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
2078
2079        a *= b;
2080
2081        crate::assert_array_relative_eq!(a, expected, 1E-10);
2082    }
2083
2084    #[test]
2085    pub fn test_scalar_mul() {
2086        let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2087        let expected: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
2088
2089        crate::assert_array_relative_eq!(2.0_f64 * a.r(), expected, 1E-10);
2090    }
2091
2092    #[test]
2093    pub fn test_scalar_mul_assign() {
2094        let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2095        let expected: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
2096
2097        a *= 2.0;
2098
2099        crate::assert_array_relative_eq!(a.r(), expected, 1E-10);
2100    }
2101
2102    #[test]
2103    pub fn test_cmp_div() {
2104        let a: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
2105        let b: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2106        let expected: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2107
2108        crate::assert_array_relative_eq!(a.r() / b.r(), expected, 1E-10);
2109    }
2110
2111    #[test]
2112    pub fn test_cmp_div_assign() {
2113        let mut a: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
2114        let b: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2115        let expected: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
2116
2117        a /= b;
2118
2119        crate::assert_array_relative_eq!(a, expected, 1E-10);
2120    }
2121
2122    #[test]
2123    pub fn test_scalar_div() {
2124        let a: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
2125        let expected: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2126
2127        crate::assert_array_relative_eq!(a.r() / 2.0, expected, 1E-10);
2128    }
2129
2130    #[test]
2131    pub fn test_scalar_div_assign() {
2132        let mut a: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
2133        let expected: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
2134
2135        a /= 2.0;
2136
2137        crate::assert_array_relative_eq!(a.r(), expected, 1E-10);
2138    }
2139
2140    #[test]
2141    pub fn test_into_subview() {
2142        let mut rng = ChaCha8Rng::seed_from_u64(0);
2143        let mut a = DynArray::<f64, 3>::from_shape([7, 5, 9]);
2144        a.fill_from_equally_distributed(&mut rng);
2145
2146        let view = a.r().into_subview([1, 0, 3], [2, 3, 4]);
2147
2148        assert_eq!(view[[1, 2, 3]], a[[2, 2, 6]])
2149    }
2150
2151    #[test]
2152    pub fn test_apply() {
2153        let mut rng = ChaCha8Rng::seed_from_u64(0);
2154
2155        let alpha = 1.5;
2156        let beta = 3.6;
2157
2158        let mut a = DynArray::<f64, 2>::from_shape([3, 2]);
2159        a.fill_from_equally_distributed(&mut rng);
2160
2161        let mut x = DynArray::<f64, 1>::from_shape([2]);
2162        x.fill_from_equally_distributed(&mut rng);
2163
2164        let mut y = DynArray::<f64, 1>::from_shape([3]);
2165        let y2 = y.eval();
2166
2167        a.apply(alpha, &x, beta, &mut y);
2168
2169        for (row_index, row) in a.row_iter().enumerate() {
2170            assert_relative_eq!(
2171                alpha * row.inner(&x).unwrap() + beta * y2[[row_index]],
2172                y[[row_index]],
2173                epsilon = 1E-10
2174            );
2175        }
2176
2177        let mut x = DynArray::<f64, 2>::from_shape([2, 5]);
2178        x.fill_from_equally_distributed(&mut rng);
2179
2180        let mut y = DynArray::<f64, 2>::from_shape([3, 5]);
2181        let y2 = y.eval();
2182
2183        a.apply(alpha, &x, beta, &mut y);
2184
2185        for (row_index, row) in a.row_iter().enumerate() {
2186            for (col_index, col) in x.col_iter().enumerate() {
2187                assert_relative_eq!(
2188                    alpha * row.inner(&col).unwrap() + beta * y2[[row_index, col_index]],
2189                    y[[row_index, col_index]],
2190                    epsilon = 1E-10
2191                );
2192            }
2193        }
2194    }
2195
2196    #[test]
2197    pub fn test_dot() {
2198        let mut rng = ChaCha8Rng::seed_from_u64(0);
2199
2200        let mut a = DynArray::<f64, 2>::from_shape([3, 2]);
2201        a.fill_from_equally_distributed(&mut rng);
2202
2203        let mut x = DynArray::<f64, 1>::from_shape([2]);
2204        x.fill_from_equally_distributed(&mut rng);
2205
2206        let mut y = DynArray::<f64, 1>::from_shape([3]);
2207
2208        a.apply(1.0, &x, 0.0, &mut y);
2209
2210        crate::assert_array_relative_eq!(y, a.dot(&x), 1E-10);
2211
2212        let mut x = DynArray::<f64, 2>::from_shape([2, 5]);
2213        x.fill_from_equally_distributed(&mut rng);
2214
2215        let mut y = DynArray::<f64, 2>::from_shape([3, 5]);
2216
2217        a.apply(1.0, &x, 0.0, &mut y);
2218
2219        crate::assert_array_relative_eq!(y, a.dot(&x), 1E-10);
2220    }
2221
2222    #[test]
2223    pub fn test_swap() {
2224        let mut a = DynArray::<usize, 1>::from_shape([2]);
2225        a[[0]] = 1;
2226        a[[1]] = 2;
2227
2228        a.swap([0], [1]).unwrap();
2229        assert_eq!(a[[0]], 2);
2230        assert_eq!(a[[1]], 1);
2231    }
2232
2233    #[test]
2234    pub fn test_multislice_iterator() {
2235        let mut arr = DynArray::<f64, _>::from_shape([2, 3, 4, 5]);
2236        arr.fill_from_seed_equally_distributed(0);
2237
2238        let mut count = 0;
2239        for (indices, slice) in arr.r().multislice_iter::<_, 2>([2, 3]) {
2240            assert_array_relative_eq!(slice, arr.r().multislice::<_, 2>([2, 3], indices), 1E-12);
2241            count += 1;
2242        }
2243        assert_eq!(count, 20);
2244    }
2245
2246    #[test]
2247    pub fn test_multislice_iterator_mut() {
2248        let mut arr = DynArray::<f64, _>::from_shape([2, 3, 4, 5]);
2249        arr.fill_from_seed_equally_distributed(0);
2250        let expected = arr.eval();
2251
2252        let mut count = 0;
2253        for (indices, slice) in arr.r_mut().multislice_iter_mut::<_, 2>([2, 3]) {
2254            assert_array_relative_eq!(
2255                slice,
2256                expected.r().multislice::<_, 2>([2, 3], indices),
2257                1E-12
2258            );
2259            count += 1;
2260        }
2261        assert_eq!(count, 20);
2262    }
2263}