use std::ops::{Add, Sub};
use aligned_vec::{AVec, CACHELINE_ALIGN, ConstAlign};
use empty_axis::AxisPosition;
use iterators::{
ArrayDefaultIteratorByRef, ArrayDefaultIteratorByValue, ArrayDefaultIteratorMut,
ArrayDiagIteratorByRef, ArrayDiagIteratorByValue, ArrayDiagIteratorMut, ColIterator,
ColIteratorMut, RowIterator, RowIteratorMut,
};
use itertools::izip;
use num::{One, Zero, traits::MulAdd};
use operators::{
addition::ArrayAddition, cast::ArrayCast, cmp_wise_division::CmpWiseDivision,
cmp_wise_product::CmpWiseProduct, coerce::CoerceArray, mul_add::MulAddImpl, negation::ArrayNeg,
reverse_axis::ReverseAxis, scalar_mult::ArrayScalarMult, subtraction::ArraySubtraction,
transpose::ArrayTranspose, unary_op::ArrayUnaryOperator,
};
use typenum::{Const, NonZero, PInt, ToUInt, U, Unsigned};
use crate::base_types::{c32, c64};
use crate::dense::array::iterators::{MultisliceIterator, MultisliceIteratorMut};
use crate::{
ArrayFromContainerType, Linspace, MutableArrayImpl, RowMajorArrayFromContainerType,
ScalarAsArray,
};
use paste::paste;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rand_distr::{Distribution, StandardNormal, StandardUniform};
use reference::{ArrayRef, ArrayRefMut};
use slice::ArraySlice;
use subview::ArraySubView;
use crate::{
AsMatrixApply, AsMultiIndex, EvaluateArray, EvaluateRowMajorArray, Gemm, IsGreaterByOne,
IsGreaterZero, IsSmallerThan, Max, MemoryLayout, NumberDifference, NumberType, RandScalar,
RandomAccessByValue, RlstError, RlstResult, RlstScalar, SubtractEqual, TransMode,
dense::{
array::multislice::Multislice,
base_array::BaseArray,
data_container::{
AlignedVectorContainer, SliceContainer, SliceContainerMut, VectorContainer,
},
layout::check_multi_index_in_bounds,
strided_base_array::{StridedBaseArray, required_storage_len},
},
traits::{
accessors::{
RandomAccessByRef, RandomAccessMut, RawAccess, RawAccessMut,
UnsafeDisjointRandom1DAccessMut, UnsafeRandom1DAccessByRef,
UnsafeRandom1DAccessByValue, UnsafeRandomAccessByRef, UnsafeRandomAccessByValue,
UnsafeRandomAccessMut,
},
base_operations::{BaseItem, ResizeInPlace, Shape, Stride},
data_container::ContainerType,
},
};
use super::{data_container::ArrayContainer, layout::row_major_stride_from_shape};
pub mod empty_axis;
pub mod iterators;
pub mod mult_into;
pub mod operations;
pub mod operators;
pub mod flattened;
pub mod multislice;
pub mod reference;
pub mod shared_array;
pub mod slice;
pub mod subview;
pub mod zero_dim;
pub use zero_dim::ZeroDimArray;
pub type DynArray<Item, const NDIM: usize> = Array<BaseArray<VectorContainer<Item>, NDIM>, NDIM>;
pub type AlignedDynArray<Item, const NDIM: usize, const ALIGNMENT: usize = CACHELINE_ALIGN> =
Array<BaseArray<AlignedVectorContainer<Item, ALIGNMENT>, NDIM>, NDIM>;
pub type StridedDynArray<Item, const NDIM: usize> =
Array<StridedBaseArray<VectorContainer<Item>, NDIM>, NDIM>;
pub type StaticArray<Item, const NDIM: usize, const N: usize> =
Array<BaseArray<ArrayContainer<Item, N>, NDIM>, NDIM>;
pub type SliceArray<'a, Item, const NDIM: usize> =
Array<BaseArray<SliceContainer<'a, Item>, NDIM>, NDIM>;
pub type StridedSliceArray<'a, Item, const NDIM: usize> =
Array<StridedBaseArray<SliceContainer<'a, Item>, NDIM>, NDIM>;
pub type SliceArrayMut<'a, Item, const NDIM: usize> =
Array<BaseArray<SliceContainerMut<'a, Item>, NDIM>, NDIM>;
pub type StridedSliceArrayMut<'a, Item, const NDIM: usize> =
Array<StridedBaseArray<SliceContainerMut<'a, Item>, NDIM>, NDIM>;
pub struct Array<ArrayImpl, const NDIM: usize>(ArrayImpl);
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn new(arr: ArrayImpl) -> Self {
Self(arr)
}
#[inline(always)]
pub fn imp(&self) -> &ArrayImpl {
&self.0
}
#[inline(always)]
pub fn imp_mut(&mut self) -> &mut ArrayImpl {
&mut self.0
}
#[inline(always)]
pub fn into_imp(self) -> ArrayImpl {
self.0
}
#[inline(always)]
pub fn r(&self) -> Array<ArrayRef<'_, ArrayImpl, NDIM>, NDIM> {
Array::new(ArrayRef::new(self))
}
#[inline(always)]
pub fn r_mut(&mut self) -> Array<ArrayRefMut<'_, ArrayImpl, NDIM>, NDIM> {
Array::new(ArrayRefMut::new(self))
}
}
impl<ArrayImpl, const NDIM: usize> BaseItem for Array<ArrayImpl, NDIM>
where
ArrayImpl: BaseItem,
{
type Item = ArrayImpl::Item;
}
impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> Array<ArrayImpl, NDIM> {
#[inline(always)]
pub fn number_of_elements(&self) -> usize {
self.0.shape().iter().product()
}
}
impl<Item: Copy + Default, const NDIM: usize> DynArray<Item, NDIM> {
#[inline(always)]
pub fn from_shape(shape: [usize; NDIM]) -> Self {
let size = shape.iter().product();
Self::new(BaseArray::new(VectorContainer::new(size), shape))
}
#[inline(always)]
pub fn from_shape_and_vec(shape: [usize; NDIM], data: Vec<Item>) -> Self {
assert_eq!(
data.len(),
shape.iter().product::<usize>(),
"Data length does not match the shape: {} != {}",
data.len(),
shape.iter().product::<usize>()
);
Self::new(BaseArray::new(VectorContainer::from_vec(data), shape))
}
}
impl<Item: Copy + Default, const NDIM: usize, const ALIGNMENT: usize>
AlignedDynArray<Item, NDIM, ALIGNMENT>
{
#[inline(always)]
pub fn from_shape(shape: [usize; NDIM]) -> Self {
let size = shape.iter().product();
Self::new(BaseArray::new(AlignedVectorContainer::new(size), shape))
}
#[inline(always)]
pub fn from_shape_and_avec(
shape: [usize; NDIM],
data: AVec<Item, ConstAlign<ALIGNMENT>>,
) -> Self {
assert_eq!(
data.len(),
shape.iter().product::<usize>(),
"Data length does not match the shape: {} != {}",
data.len(),
shape.iter().product::<usize>()
);
Self::new(BaseArray::new(
AlignedVectorContainer::from_avec(data),
shape,
))
}
}
impl<Item: Copy + Default> DynArray<Item, 2> {
pub fn from_iter_aij<Iter: Iterator<Item = ([usize; 2], Item)>>(
shape: [usize; 2],
iter: Iter,
) -> DynArray<Item, 2> {
let mut out = DynArray::<Item, 2>::from_shape(shape);
for (index, item) in iter {
out[index] = item;
}
out
}
}
impl<Item: Copy + Default, const ALIGNMENT: usize> AlignedDynArray<Item, 2, ALIGNMENT> {
pub fn from_iter_aij<Iter: Iterator<Item = ([usize; 2], Item)>>(
shape: [usize; 2],
iter: Iter,
) -> AlignedDynArray<Item, 2, ALIGNMENT> {
let mut out = AlignedDynArray::<Item, 2, ALIGNMENT>::from_shape(shape);
for (index, item) in iter {
out[index] = item;
}
out
}
}
impl<Item: Copy + Default> From<&[Item]> for DynArray<Item, 1> {
fn from(value: &[Item]) -> Self {
DynArray::<Item, 1>::from_shape_and_vec([value.len()], value.to_vec())
}
}
impl<Item: Copy + Default, const N: usize> From<[Item; N]> for StaticArray<Item, 1, N> {
fn from(value: [Item; N]) -> Self {
StaticArray::new(BaseArray::new(value.into(), [N]))
}
}
impl<Item: Copy + Default> From<Vec<Item>> for DynArray<Item, 1> {
fn from(value: Vec<Item>) -> Self {
DynArray::<Item, 1>::from_shape_and_vec([value.len()], value)
}
}
impl<Item: Copy + Default> From<Vec<Vec<Item>>> for DynArray<Item, 2> {
fn from(value: Vec<Vec<Item>>) -> Self {
let nrows = value.len();
if nrows == 0 {
empty_array()
} else {
let ncols = value.first().unwrap().len();
for row in value.iter() {
assert_eq!(ncols, row.len(), "All rows must have equal length.");
}
let mut out = DynArray::<Item, 2>::from_shape([nrows, ncols]);
for (row_index, row) in value.iter().enumerate() {
for (col_index, &elem) in row.iter().enumerate() {
out[[row_index, col_index]] = elem;
}
}
out
}
}
}
impl<Item: Copy + Default, const NDIM: usize> StridedDynArray<Item, NDIM> {
#[inline(always)]
pub fn from_shape_and_stride(shape: [usize; NDIM], stride: [usize; NDIM]) -> Self {
let size = required_storage_len(shape, stride);
Self::new(StridedBaseArray::new(
VectorContainer::new(size),
shape,
stride,
))
}
}
impl<Item: Copy + Default, const NDIM: usize>
Array<StridedBaseArray<VectorContainer<Item>, NDIM>, NDIM>
{
#[inline(always)]
pub fn row_major(shape: [usize; NDIM]) -> Self {
let size = shape.iter().product();
Self::new(StridedBaseArray::new(
VectorContainer::new(size),
shape,
row_major_stride_from_shape(shape),
))
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM>,
{
#[inline(always)]
pub fn shape(&self) -> [usize; NDIM] {
self.0.shape()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM>,
{
#[inline(always)]
pub fn len(&self) -> usize {
self.0.len()
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessByValue<NDIM>,
{
#[inline(always)]
pub unsafe fn get_value_unchecked(&self, multi_index: [usize; NDIM]) -> ArrayImpl::Item {
unsafe { self.0.get_value_unchecked(multi_index) }
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessByRef<NDIM>,
{
#[inline(always)]
pub unsafe fn get_unchecked(&self, multi_index: [usize; NDIM]) -> &ArrayImpl::Item {
unsafe { self.0.get_unchecked(multi_index) }
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessMut<NDIM>,
{
#[inline(always)]
pub unsafe fn get_unchecked_mut(&mut self, multi_index: [usize; NDIM]) -> &mut ArrayImpl::Item {
unsafe { self.0.get_unchecked_mut(multi_index) }
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: RandomAccessByValue<NDIM>,
{
#[inline(always)]
pub fn get_value(&self, multi_index: [usize; NDIM]) -> Option<ArrayImpl::Item> {
self.0.get_value(multi_index)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: RandomAccessByRef<NDIM>,
{
#[inline(always)]
pub fn get(&self, multi_index: [usize; NDIM]) -> Option<&ArrayImpl::Item> {
self.0.get(multi_index)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: RandomAccessMut<NDIM>,
{
#[inline(always)]
pub fn get_mut(&mut self, multi_index: [usize; NDIM]) -> Option<&mut ArrayImpl::Item> {
self.0.get_mut(multi_index)
}
}
impl<ArrayImpl, const NDIM: usize> std::ops::Index<[usize; NDIM]> for Array<ArrayImpl, NDIM>
where
ArrayImpl: RandomAccessByRef<NDIM>,
{
type Output = ArrayImpl::Item;
#[inline(always)]
fn index(&self, index: [usize; NDIM]) -> &Self::Output {
self.0.get(index).unwrap()
}
}
impl<ArrayImpl, const NDIM: usize> std::ops::IndexMut<[usize; NDIM]> for Array<ArrayImpl, NDIM>
where
Array<ArrayImpl, NDIM>: std::ops::Index<[usize; NDIM], Output = ArrayImpl::Item>,
ArrayImpl: RandomAccessMut<NDIM>,
{
#[inline(always)]
fn index_mut(&mut self, index: [usize; NDIM]) -> &mut ArrayImpl::Item {
self.0.get_mut(index).unwrap()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: RawAccess,
{
#[inline(always)]
pub fn data(&self) -> Option<&[ArrayImpl::Item]> {
self.0.data()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: RawAccessMut,
{
#[inline(always)]
pub fn data_mut(&mut self) -> Option<&mut [ArrayImpl::Item]> {
self.0.data_mut()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: Stride<NDIM>,
{
#[inline(always)]
pub fn stride(&self) -> [usize; NDIM] {
self.0.stride()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: Stride<NDIM> + Shape<NDIM>,
{
#[inline(always)]
pub fn memory_layout(&self) -> MemoryLayout {
self.0.memory_layout()
}
pub fn is_contiguous(&self) -> bool {
self.0.is_contiguous()
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: ResizeInPlace<NDIM>,
{
pub fn resize_in_place(&mut self, shape: [usize; NDIM]) {
self.0.resize_in_place(shape)
}
}
impl<Item, ArrayImpl, const NDIM: usize> EvaluateArray<Item, NDIM> for Array<ArrayImpl, NDIM>
where
Item: Copy + Default,
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM> + ContainerType,
{
type OutputImpl =
<<ArrayImpl as ContainerType>::Type as ArrayFromContainerType>::OutputImpl<Item, NDIM>;
fn eval(&self) -> Array<Self::OutputImpl, NDIM> {
let mut out = <<ArrayImpl as ContainerType>::Type as ArrayFromContainerType>::new_arr::<
Item,
NDIM,
>(self.shape());
out.fill_from(self);
out
}
}
impl<Item, ArrayImpl, const NDIM: usize> EvaluateRowMajorArray<Item, NDIM>
for Array<ArrayImpl, NDIM>
where
Item: Copy + Default,
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM> + ContainerType,
{
type OutputImpl =
<<ArrayImpl as ContainerType>::Type as RowMajorArrayFromContainerType>::OutputImpl<
Item,
NDIM,
>;
fn eval_row_major(&self) -> Array<Self::OutputImpl, NDIM> {
let mut out =
<<ArrayImpl as ContainerType>::Type as RowMajorArrayFromContainerType>::new_row_major_arr::<
Item,
NDIM,
>(self.shape());
out.fill_from(self);
out
}
}
pub fn empty_array<Item, const NDIM: usize>() -> Array<BaseArray<VectorContainer<Item>, NDIM>, NDIM>
where
Item: Copy + Default,
{
let shape = [0; NDIM];
let container = VectorContainer::new(0);
Array::new(BaseArray::new(container, shape))
}
impl<'a, Item, const NDIM: usize> SliceArray<'a, Item, NDIM>
where
Item: Copy + Default,
{
pub fn from_shape(slice: &'a [Item], shape: [usize; NDIM]) -> Self {
assert_eq!(
slice.len(),
shape.iter().product::<usize>(),
"Slice length does not match the shape: {} != {}",
slice.len(),
shape.iter().product::<usize>()
);
Array::new(BaseArray::new(SliceContainer::new(slice), shape))
}
}
impl<'a, Item, const NDIM: usize> SliceArrayMut<'a, Item, NDIM>
where
Item: Copy + Default,
{
pub fn from_shape(slice: &'a mut [Item], shape: [usize; NDIM]) -> Self {
assert_eq!(
slice.len(),
shape.iter().product::<usize>(),
"Slice length does not match the shape: {} != {}",
slice.len(),
shape.iter().product::<usize>()
);
Array::new(BaseArray::new(SliceContainerMut::new(slice), shape))
}
}
impl<'a, Item, const NDIM: usize> StridedSliceArray<'a, Item, NDIM>
where
Item: Copy + Default,
{
pub fn from_shape_and_stride(
slice: &'a [Item],
shape: [usize; NDIM],
stride: [usize; NDIM],
) -> Self {
Array::new(StridedBaseArray::new(
SliceContainer::new(slice),
shape,
stride,
))
}
}
impl<'a, Item, const NDIM: usize> StridedSliceArrayMut<'a, Item, NDIM>
where
Item: Copy + Default,
{
pub fn from_shape_and_stride(
slice: &'a mut [Item],
shape: [usize; NDIM],
stride: [usize; NDIM],
) -> Self {
Array::new(StridedBaseArray::new(
SliceContainerMut::new(slice),
shape,
stride,
))
}
}
impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> std::fmt::Debug for Array<ArrayImpl, NDIM> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let shape = self.shape();
write!(f, "Array[").unwrap();
for item in shape {
write!(f, "{item},").unwrap();
}
write!(f, "]")
}
}
impl<Item, const NDIM: usize> DynArray<Item, NDIM>
where
Item: Copy + Default,
{
pub fn new_from<ArrayImpl>(arr: &Array<ArrayImpl, NDIM>) -> Self
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
{
let mut output = empty_array::<Item, NDIM>();
output.fill_from_resize(arr);
output
}
}
impl<Item, const NDIM: usize> StridedDynArray<Item, NDIM>
where
Item: Copy + Default,
{
pub fn new_from<ArrayImpl>(stride: [usize; NDIM], arr: &Array<ArrayImpl, NDIM>) -> Self
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
{
let mut output = StridedDynArray::from_shape_and_stride(arr.shape(), stride);
output.fill_from(arr);
output
}
pub fn row_major_from<ArrayImpl>(arr: &Array<ArrayImpl, NDIM>) -> Self
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
{
StridedDynArray::new_from(row_major_stride_from_shape(arr.shape()), arr)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandom1DAccessByValue + Shape<NDIM>,
{
#[inline(always)]
pub fn iter_value(&self) -> ArrayDefaultIteratorByValue<'_, ArrayImpl, NDIM> {
ArrayDefaultIteratorByValue::new(self)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandom1DAccessByRef + Shape<NDIM>,
{
#[inline(always)]
pub fn iter_ref(&self) -> ArrayDefaultIteratorByRef<'_, ArrayImpl, NDIM> {
ArrayDefaultIteratorByRef::new(self)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + Shape<NDIM>,
{
#[inline(always)]
pub fn iter_mut(&mut self) -> ArrayDefaultIteratorMut<'_, ArrayImpl, NDIM> {
ArrayDefaultIteratorMut::new(self)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessByValue<NDIM, Item = Item> + Shape<NDIM>,
{
#[inline(always)]
pub fn diag_iter_value(&self) -> ArrayDiagIteratorByValue<'_, ArrayImpl, NDIM> {
ArrayDiagIteratorByValue::new(self)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessByRef<NDIM, Item = Item> + Shape<NDIM>,
{
#[inline(always)]
pub fn diag_iter_ref(&self) -> ArrayDiagIteratorByRef<'_, ArrayImpl, NDIM> {
ArrayDiagIteratorByRef::new(self)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item> + Shape<NDIM>,
{
#[inline(always)]
pub fn diag_iter_mut(&mut self) -> ArrayDiagIteratorMut<'_, ArrayImpl, NDIM> {
ArrayDiagIteratorMut::new(self)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut<Item = Item> + Shape<NDIM>,
{
#[inline(always)]
pub fn fill_from<ArrayImplOther>(&mut self, other: &Array<ArrayImplOther, NDIM>)
where
ArrayImplOther: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
{
assert_eq!(self.shape(), other.shape());
for (item, other_item) in self.iter_mut().zip(other.iter_value()) {
*item = other_item;
}
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
Const<NDIM>: ToUInt,
<Const<NDIM> as ToUInt>::Output: Unsigned + NonZero,
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item>
+ UnsafeRandomAccessByValue<NDIM, Item = Item>
+ Shape<NDIM>,
{
pub fn multislice_iter<const SDIM: usize, const OUTDIM: usize>(
&self,
axes: [usize; SDIM],
) -> MultisliceIterator<'_, Item, ArrayImpl, NDIM, SDIM, OUTDIM>
where
Const<SDIM>: ToUInt,
Const<OUTDIM>: ToUInt,
<Const<SDIM> as ToUInt>::Output: Unsigned + NonZero,
<Const<OUTDIM> as ToUInt>::Output: Unsigned + NonZero,
NumberType<SDIM>: IsSmallerThan<NDIM>,
PInt<U<OUTDIM>>: Add<PInt<U<SDIM>>, Output = PInt<U<NDIM>>>,
PInt<U<NDIM>>: Sub<PInt<U<SDIM>>>,
{
MultisliceIterator::new(self, axes)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
Const<NDIM>: ToUInt,
<Const<NDIM> as ToUInt>::Output: Unsigned + NonZero,
ArrayImpl: UnsafeDisjointRandom1DAccessMut<Item = Item>
+ UnsafeRandomAccessMut<NDIM, Item = Item>
+ Shape<NDIM>,
{
pub fn multislice_iter_mut<const SDIM: usize, const OUTDIM: usize>(
&mut self,
axes: [usize; SDIM],
) -> MultisliceIteratorMut<'_, Item, ArrayImpl, NDIM, SDIM, OUTDIM>
where
Const<SDIM>: ToUInt,
Const<OUTDIM>: ToUInt,
<Const<SDIM> as ToUInt>::Output: Unsigned + NonZero,
<Const<OUTDIM> as ToUInt>::Output: Unsigned + NonZero,
NumberType<SDIM>: IsSmallerThan<NDIM>,
PInt<U<OUTDIM>>: Add<PInt<U<SDIM>>, Output = PInt<U<NDIM>>>,
PInt<U<NDIM>>: Sub<PInt<U<SDIM>>>,
{
MultisliceIteratorMut::new(self, axes)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + Shape<NDIM>,
{
#[inline(always)]
pub fn fill_from_iter<Iter>(&mut self, iter: Iter)
where
Iter: Iterator<Item = ArrayImpl::Item>,
{
for (item, other_item) in izip!(self.iter_mut(), iter) {
*item = other_item;
}
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: ResizeInPlace<NDIM> + Shape<NDIM> + UnsafeDisjointRandom1DAccessMut,
{
#[inline(always)]
pub fn fill_from_resize<ArrayImplOther>(&mut self, other: &Array<ArrayImplOther, NDIM>)
where
ArrayImplOther: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = ArrayImpl::Item>,
{
self.resize_in_place(other.shape());
self.fill_from(other);
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + Shape<NDIM>,
{
#[inline(always)]
pub fn fill_with_value(&mut self, value: ArrayImpl::Item) {
for item in self.iter_mut() {
*item = value;
}
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + Shape<NDIM>,
{
pub fn set_zero(&mut self) {
for item in self.iter_mut() {
*item = Default::default();
}
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + UnsafeRandomAccessMut<NDIM> + Shape<NDIM>,
ArrayImpl::Item: One,
{
pub fn set_identity(&mut self) {
self.set_zero();
self.diag_iter_mut().for_each(|elem| *elem = One::one());
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessByValue<NDIM> + Shape<NDIM>,
ArrayImpl::Item: std::iter::Sum,
{
#[inline(always)]
pub fn trace(&self) -> ArrayImpl::Item {
self.diag_iter_value().sum()
}
}
impl<Item, ArrayImpl> Array<ArrayImpl, 1>
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
ArrayImpl::Item: Conj<Output = Item>
+ std::ops::Mul<Output = ArrayImpl::Item>
+ std::ops::Add<Output = Item>,
{
#[inline(always)]
pub fn inner<ArrayImplOther>(&self, other: &Array<ArrayImplOther, 1>) -> Option<Item>
where
ArrayImplOther: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
{
assert_eq!(self.shape(), other.shape());
izip!(self.iter_value(), other.iter_value())
.map(|(x, y)| x * y.conj())
.reduce(std::ops::Add::add)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
Item: Abs,
<Item as Abs>::Output: Max<Output = <Item as Abs>::Output>,
{
#[inline(always)]
pub fn max_abs(&self) -> Option<<Item as Abs>::Output> {
self.iter_value().map(|elem| elem.abs()).reduce(Max::max)
}
}
impl<Item, ArrayImpl> Array<ArrayImpl, 1>
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
Item: Abs,
<Item as Abs>::Output: std::ops::Add<Output = <Item as Abs>::Output>,
{
#[inline(always)]
pub fn norm_1(&self) -> Option<<Item as Abs>::Output> {
self.iter_value()
.map(|elem| elem.abs())
.reduce(std::ops::Add::add)
}
}
impl<Item, ArrayImpl> Array<ArrayImpl, 1>
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<1>,
Item: AbsSquare,
<Item as AbsSquare>::Output: Sqrt<Output = <Item as AbsSquare>::Output>
+ std::ops::Add<Output = <Item as AbsSquare>::Output>,
{
#[inline(always)]
pub fn norm_2(&self) -> Option<<Item as AbsSquare>::Output> {
self.iter_value()
.map(|elem| elem.abs_square())
.reduce(std::ops::Add::add)
.map(|elem| elem.sqrt())
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandom1DAccessByValue<Item = Item> + Shape<NDIM>,
Item: AbsSquare,
<Item as AbsSquare>::Output: Sqrt<Output = <Item as AbsSquare>::Output>
+ std::ops::Add<Output = <Item as AbsSquare>::Output>
+ Copy
+ Default,
{
#[inline(always)]
pub fn norm_fro(&self) -> Option<<Item as AbsSquare>::Output> {
self.r()
.abs_square()
.iter_value()
.reduce(std::ops::Add::add)
.map(|elem| elem.sqrt())
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandom1DAccessByValue + Shape<NDIM>,
{
#[allow(clippy::type_complexity)]
pub fn into_type<T>(
self,
) -> Array<
ArrayUnaryOperator<ArrayImpl::Item, T, ArrayImpl, fn(ArrayImpl::Item) -> T, NDIM>,
NDIM,
>
where
ArrayImpl::Item: Into<T>,
{
Array::new(ArrayUnaryOperator::new(self, |item| item.into()))
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
{
pub fn col_iter(&self) -> ColIterator<'_, ArrayImpl, 2> {
ColIterator::new(self)
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: Shape<2> + UnsafeRandomAccessMut<2>,
{
pub fn col_iter_mut(&mut self) -> ColIteratorMut<'_, ArrayImpl, 2> {
ColIteratorMut::new(self)
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
{
pub fn row_iter(&self) -> RowIterator<'_, ArrayImpl, 2> {
RowIterator::new(self)
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: Shape<2> + UnsafeRandomAccessByValue<2>,
{
pub fn row_iter_mut(&mut self) -> RowIteratorMut<'_, ArrayImpl, 2> {
RowIteratorMut::new(self)
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: UnsafeRandom1DAccessByValue + Shape<2>,
{
pub fn iter_aij_value(&self) -> impl Iterator<Item = ([usize; 2], ArrayImpl::Item)> + '_ {
let iter = ArrayDefaultIteratorByValue::new(self);
AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), self.shape())
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: UnsafeRandom1DAccessByRef + Shape<2>,
{
pub fn iter_aij_ref(&self) -> impl Iterator<Item = ([usize; 2], &ArrayImpl::Item)> + '_ {
let iter = ArrayDefaultIteratorByRef::new(self);
AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), self.shape())
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: UnsafeDisjointRandom1DAccessMut + Shape<2>,
{
pub fn iter_aij_mut(
&mut self,
) -> impl Iterator<Item = ([usize; 2], &mut ArrayImpl::Item)> + '_ {
let shape = self.shape();
let iter = ArrayDefaultIteratorMut::new(self);
AsMultiIndex::multi_index(std::iter::Iterator::enumerate(iter), shape)
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn cast<Target>(self) -> Array<ArrayCast<Target, ArrayImpl, NDIM>, NDIM> {
Array::new(ArrayCast::new(self))
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn coerce_dim<const CDIM: usize>(
self,
) -> RlstResult<Array<CoerceArray<ArrayImpl, NDIM, CDIM>, CDIM>> {
if CDIM == NDIM {
Ok(Array::new(CoerceArray::new(self)))
} else {
Err(RlstError::GeneralError(
format!("Cannot coerce array: dimensions do not match {CDIM} != {NDIM}.")
.to_string(),
))
}
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM>,
{
pub fn permute_axes(
self,
permutation: [usize; NDIM],
) -> Array<ArrayTranspose<ArrayImpl, NDIM>, NDIM> {
Array::new(ArrayTranspose::new(self, permutation))
}
pub fn transpose(self) -> Array<ArrayTranspose<ArrayImpl, NDIM>, NDIM> {
let mut permutation = [0; NDIM];
for (ind, p) in (0..NDIM).rev().enumerate() {
permutation[ind] = p;
}
self.permute_axes(permutation)
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
Item: RandScalar + RlstScalar,
ArrayImpl: UnsafeDisjointRandom1DAccessMut<Item = Item> + Shape<NDIM>,
StandardNormal: Distribution<Item::Real>,
StandardUniform: Distribution<Item::Real>,
{
pub fn fill_from_standard_normal<R: Rng>(&mut self, rng: &mut R) {
let dist = StandardNormal;
self.iter_mut()
.for_each(|val| *val = <Item>::random_scalar(rng, &dist));
}
pub fn fill_from_equally_distributed<R: Rng>(&mut self, rng: &mut R) {
let dist = StandardUniform;
self.iter_mut()
.for_each(|val| *val = <Item>::random_scalar(rng, &dist));
}
pub fn fill_from_seed_equally_distributed(&mut self, seed: usize) {
let mut rng = ChaCha8Rng::seed_from_u64(seed as u64);
let dist = StandardUniform;
self.iter_mut()
.for_each(|val| *val = <Item>::random_scalar(&mut rng, &dist));
}
pub fn fill_from_seed_normally_distributed(&mut self, seed: usize) {
let mut rng = ChaCha8Rng::seed_from_u64(seed as u64);
let dist = StandardNormal;
self.iter_mut()
.for_each(|val| *val = <Item>::random_scalar(&mut rng, &dist));
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn reverse_axis(self, axis: usize) -> Array<ReverseAxis<ArrayImpl, NDIM>, NDIM> {
assert!(axis < NDIM, "Axis out of bounds");
Array::new(ReverseAxis::new(self, axis))
}
}
impl<ArrayImpl, const ADIM: usize> Array<ArrayImpl, ADIM>
where
ArrayImpl: Shape<ADIM>,
{
pub fn slice<const OUTDIM: usize>(
self,
axis: usize,
index: usize,
) -> Array<ArraySlice<ArrayImpl, ADIM, OUTDIM>, OUTDIM>
where
NumberType<ADIM>: IsGreaterByOne<OUTDIM>,
NumberType<OUTDIM>: IsGreaterZero,
{
Array::new(ArraySlice::new(self, [axis, index]))
}
pub fn multislice<const SDIM: usize, const OUTDIM: usize>(
self,
axes: [usize; SDIM],
indices: [usize; SDIM],
) -> Array<Multislice<ArrayImpl, ADIM, SDIM, OUTDIM>, OUTDIM>
where
NumberType<SDIM>: IsSmallerThan<ADIM>,
NumberDifference<ADIM, SDIM>: SubtractEqual<OUTDIM>,
{
Array::new(Multislice::new(self, axes, indices))
}
}
impl<ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: Shape<2>,
{
pub fn row(self, row_index: usize) -> Array<ArraySlice<ArrayImpl, 2, 1>, 1> {
self.slice::<1>(0, row_index)
}
pub fn col(self, col_index: usize) -> Array<ArraySlice<ArrayImpl, 2, 1>, 1> {
self.slice::<1>(1, col_index)
}
}
impl<ArrayImpl: Shape<NDIM>, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn into_subview(
self,
offset: [usize; NDIM],
shape: [usize; NDIM],
) -> Array<ArraySubView<ArrayImpl, NDIM>, NDIM> {
Array::new(ArraySubView::new(self, offset, shape))
}
}
impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> MulAdd<Item, Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: BaseItem<Item = Item> + Shape<NDIM>,
ArrayImpl2: BaseItem<Item = Item> + Shape<NDIM>,
Item: MulAdd<Output = Item> + Copy,
{
type Output = Array<MulAddImpl<ArrayImpl1, ArrayImpl2, Item, NDIM>, NDIM>;
fn mul_add(
self,
a: Item,
b: Array<ArrayImpl2, NDIM>,
) -> Array<MulAddImpl<ArrayImpl1, ArrayImpl2, Item, NDIM>, NDIM>
where
ArrayImpl1: BaseItem<Item = Item> + Shape<NDIM>,
ArrayImpl2: BaseItem<Item = Item> + Shape<NDIM>,
Item: MulAdd<Output = Item> + Copy,
{
Array::new(MulAddImpl::new(self, b, a))
}
}
impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Add<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM>,
ArrayImpl2: Shape<NDIM>,
{
type Output = Array<ArrayAddition<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
#[inline(always)]
fn add(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
Array::new(ArrayAddition::new(self, rhs))
}
}
impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::AddAssign<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
Item: std::ops::AddAssign,
{
#[inline(always)]
fn add_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
assert_eq!(self.shape(), rhs.shape());
for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
*item1 += item2;
}
}
}
impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Sub<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM>,
ArrayImpl2: Shape<NDIM>,
{
type Output = Array<ArraySubtraction<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
fn sub(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
Array::new(ArraySubtraction::new(self, rhs))
}
}
impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::SubAssign<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
Item: std::ops::SubAssign,
{
#[inline(always)]
fn sub_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
assert_eq!(self.shape(), rhs.shape());
for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
*item1 -= item2;
}
}
}
impl<Item, ArrayImpl, const NDIM: usize> std::ops::Neg for Array<ArrayImpl, NDIM>
where
ArrayImpl: BaseItem<Item = Item>,
Item: std::ops::Neg<Output = Item>,
{
type Output = Array<ArrayNeg<ArrayImpl, NDIM>, NDIM>;
fn neg(self) -> Self::Output {
Array::new(ArrayNeg::new(self))
}
}
impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Div<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM>,
ArrayImpl2: Shape<NDIM>,
{
type Output = Array<CmpWiseDivision<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
fn div(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
Array::new(CmpWiseDivision::new(self, rhs))
}
}
impl<Item, ArrayImpl, const NDIM: usize> std::ops::Div<Item> for Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM> + BaseItem<Item = Item>,
Item: Recip<Output = Item> + std::ops::Div<Output = Item>,
{
type Output = Array<ArrayScalarMult<Item, ArrayImpl, NDIM>, NDIM>;
fn div(self, rhs: Item) -> Self::Output {
Array::new(ArrayScalarMult::new(rhs.recip(), self))
}
}
impl<Item, ArrayImpl, const NDIM: usize> std::ops::DivAssign<Item> for Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
Item: std::ops::DivAssign + Copy,
{
fn div_assign(&mut self, rhs: Item) {
for item in self.iter_mut() {
*item /= rhs;
}
}
}
impl<ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::Mul<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM>,
ArrayImpl2: Shape<NDIM>,
{
type Output = Array<CmpWiseProduct<ArrayImpl1, ArrayImpl2, NDIM>, NDIM>;
fn mul(self, rhs: Array<ArrayImpl2, NDIM>) -> Self::Output {
Array::new(CmpWiseProduct::new(self, rhs))
}
}
impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::MulAssign<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
Item: std::ops::MulAssign,
{
fn mul_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
*item1 *= item2;
}
}
}
impl<Item, ArrayImpl1, ArrayImpl2, const NDIM: usize> std::ops::DivAssign<Array<ArrayImpl2, NDIM>>
for Array<ArrayImpl1, NDIM>
where
ArrayImpl1: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
ArrayImpl2: Shape<NDIM> + UnsafeRandom1DAccessByValue<Item = Item>,
Item: std::ops::DivAssign,
{
fn div_assign(&mut self, rhs: Array<ArrayImpl2, NDIM>) {
for (item1, item2) in izip!(self.iter_mut(), rhs.iter_value()) {
*item1 /= item2;
}
}
}
macro_rules! impl_scalar_mult {
($ScalarType:ty) => {
impl<ArrayImpl, const NDIM: usize> std::ops::Mul<Array<ArrayImpl, NDIM>> for $ScalarType {
type Output = Array<ArrayScalarMult<$ScalarType, ArrayImpl, NDIM>, NDIM>;
fn mul(self, rhs: Array<ArrayImpl, NDIM>) -> Self::Output {
Array::new(ArrayScalarMult::new(self, rhs))
}
}
impl<ArrayImpl, const NDIM: usize> std::ops::Mul<$ScalarType> for Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM> + BaseItem<Item = $ScalarType>,
{
type Output = Array<ArrayScalarMult<$ScalarType, ArrayImpl, NDIM>, NDIM>;
fn mul(self, rhs: $ScalarType) -> Self::Output {
Array::new(ArrayScalarMult::new(rhs, self))
}
}
};
}
impl_scalar_mult!(f64);
impl_scalar_mult!(f32);
impl_scalar_mult!(c64);
impl_scalar_mult!(c32);
impl_scalar_mult!(usize);
impl_scalar_mult!(i8);
impl_scalar_mult!(i16);
impl_scalar_mult!(i32);
impl_scalar_mult!(i64);
impl_scalar_mult!(u8);
impl_scalar_mult!(u16);
impl_scalar_mult!(u32);
impl_scalar_mult!(u64);
impl<Item, ArrayImpl, const NDIM: usize> std::ops::MulAssign<Item> for Array<ArrayImpl, NDIM>
where
ArrayImpl: Shape<NDIM> + UnsafeDisjointRandom1DAccessMut<Item = Item>,
Item: std::ops::MulAssign + Copy,
{
fn mul_assign(&mut self, rhs: Item) {
for item in self.iter_mut() {
*item *= rhs;
}
}
}
impl<ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM> {
pub fn unary_op<OpItem, OpTarget, Op: Fn(OpItem) -> OpTarget>(
self,
op: Op,
) -> Array<ArrayUnaryOperator<OpItem, OpTarget, ArrayImpl, Op, NDIM>, NDIM> {
Array::new(ArrayUnaryOperator::new(self, op))
}
}
impl<Item: Default + Copy + std::ops::Mul<Output = Item>, ArrayImpl, const NDIM: usize>
Array<ArrayImpl, NDIM>
where
ArrayImpl: BaseItem<Item = Item>,
{
pub fn scalar_mul(self, scalar: Item) -> Array<ArrayScalarMult<Item, ArrayImpl, NDIM>, NDIM> {
Array::new(ArrayScalarMult::new(scalar, self))
}
}
macro_rules! impl_unary_op_trait {
($name:ident, $method_name:ident) => {
paste! {
use crate::traits::number_traits::$name;
impl<Item: $name, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: BaseItem<Item = Item>,
{
#[inline(always)]
#[doc = "Componentwise apply `" $method_name "` to the array."]
pub fn $method_name(self) -> Array<ArrayUnaryOperator<Item, <Item as $name>::Output, ArrayImpl, fn(Item) -> <Item as $name>::Output, NDIM>, NDIM> {
self.unary_op(|x| x.$method_name())
}
}
}
};
}
impl_unary_op_trait!(Conj, conj);
impl_unary_op_trait!(Abs, abs);
impl_unary_op_trait!(Square, square);
impl_unary_op_trait!(AbsSquare, abs_square);
impl_unary_op_trait!(Sqrt, sqrt);
impl_unary_op_trait!(Exp, exp);
impl_unary_op_trait!(Ln, ln);
impl_unary_op_trait!(Recip, recip);
impl_unary_op_trait!(Sin, sin);
impl_unary_op_trait!(Cos, cos);
impl_unary_op_trait!(Tan, tan);
impl_unary_op_trait!(Asin, asin);
impl_unary_op_trait!(Acos, acos);
impl_unary_op_trait!(Atan, atan);
impl_unary_op_trait!(Sinh, sinh);
impl_unary_op_trait!(Cosh, cosh);
impl_unary_op_trait!(Tanh, tanh);
impl_unary_op_trait!(Asinh, asinh);
impl_unary_op_trait!(Acosh, acosh);
impl_unary_op_trait!(Atanh, atanh);
impl<Item, ArrayImplX, ArrayImplY, ArrayImpl>
AsMatrixApply<Array<ArrayImplX, 2>, Array<ArrayImplY, 2>> for Array<ArrayImpl, 2>
where
Item: Gemm + Copy,
ArrayImplX: RawAccess<Item = Item> + Stride<2> + Shape<2>,
ArrayImplY: RawAccessMut<Item = Item> + Stride<2> + Shape<2>,
ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
{
fn apply(
&self,
alpha: Self::Item,
x: &Array<ArrayImplX, 2>,
beta: Self::Item,
y: &mut Array<ArrayImplY, 2>,
) {
let stride_a = self.stride();
let stride_x = x.stride();
let stride_y = y.stride();
let shape_a = self.shape();
let shape_x = x.shape();
let shape_y = y.shape();
assert_eq!(
shape_a[0], shape_y[0],
"`y` has incompatible shape {shape_y:#?} with `self` {shape_a:#?}."
);
assert_eq!(
shape_a[1], shape_x[0],
"`x` has incompatible shape {shape_x:#?} with `self` {shape_a:#?}."
);
assert_eq!(
shape_y[1], shape_x[1],
"`x` has incompatible shape {shape_x:#?} with `y` {shape_y:#?}."
);
let m = shape_a[0];
let n = shape_x[1];
let k = shape_a[1];
<Item as Gemm>::gemm(
TransMode::NoTrans,
TransMode::NoTrans,
m,
n,
k,
alpha,
self.data().unwrap(),
stride_a[0],
stride_a[1],
x.data().unwrap(),
stride_x[0],
stride_x[1],
beta,
y.data_mut().unwrap(),
stride_y[0],
stride_y[1],
);
}
}
impl<Item, ArrayImplX, ArrayImplY, ArrayImpl>
AsMatrixApply<Array<ArrayImplX, 1>, Array<ArrayImplY, 1>> for Array<ArrayImpl, 2>
where
Item: Gemm + Copy,
ArrayImplX: RawAccess<Item = Item> + Stride<1> + Shape<1>,
ArrayImplY: RawAccessMut<Item = Item> + Stride<1> + Shape<1>,
ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
{
fn apply(
&self,
alpha: Self::Item,
x: &Array<ArrayImplX, 1>,
beta: Self::Item,
y: &mut Array<ArrayImplY, 1>,
) {
self.apply(
alpha,
&x.r().insert_empty_axis::<2>(AxisPosition::Back),
beta,
&mut y.r_mut().insert_empty_axis::<2>(AxisPosition::Back),
);
}
}
impl<Item, ArrayImpl> Array<ArrayImpl, 2>
where
ArrayImpl: RawAccess<Item = Item> + Stride<2> + Shape<2>,
Item: Gemm + Copy + Default + One + Zero,
{
pub fn dot<ArrayImplOther, const NDIM: usize>(
&self,
other: &Array<ArrayImplOther, NDIM>,
) -> DynArray<Item, NDIM>
where
ArrayImplOther: RawAccess<Item = Item> + Shape<NDIM> + Stride<NDIM>,
{
let a_shape = self.shape();
let x_shape = other.shape();
let mut out = empty_array::<Item, NDIM>();
if NDIM == 1 {
let mut out = out.r_mut().coerce_dim::<1>().unwrap();
out.resize_in_place([a_shape[0]]);
self.apply(
Item::one(),
&other.r().coerce_dim::<1>().unwrap(),
Item::zero(),
&mut out,
);
} else if NDIM == 2 {
let mut out = out.r_mut().coerce_dim::<2>().unwrap();
out.resize_in_place([a_shape[0], x_shape[1]]);
self.apply(
Item::one(),
&other.r().coerce_dim::<2>().unwrap(),
Item::zero(),
&mut out,
)
} else {
panic!("NDIM = {NDIM} not supported.");
}
out
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item>,
Item: Copy,
{
#[inline(always)]
pub unsafe fn swap_unsafe(&mut self, pos1: [usize; NDIM], pos2: [usize; NDIM]) {
let elem1 = *unsafe { self.get_unchecked_mut(pos1) };
let elem2 = *unsafe { self.get_unchecked_mut(pos2) };
*unsafe { self.get_unchecked_mut(pos1) } = elem2;
*unsafe { self.get_unchecked_mut(pos2) } = elem1;
}
}
impl<Item, ArrayImpl, const NDIM: usize> Array<ArrayImpl, NDIM>
where
ArrayImpl: UnsafeRandomAccessMut<NDIM, Item = Item> + Shape<NDIM>,
Item: Copy,
{
#[inline(always)]
pub fn swap(&mut self, pos1: [usize; NDIM], pos2: [usize; NDIM]) -> RlstResult<()> {
let shape = self.shape();
if !check_multi_index_in_bounds(pos1, shape) {
Err(RlstError::GeneralError(format!(
"Index `pos1` = {:#?} out of bounds.",
pos1
)))
} else if !check_multi_index_in_bounds(pos2, shape) {
Err(RlstError::GeneralError(format!(
"Index `pos2` = {:#?} out of bounds.",
pos2
)))
} else {
unsafe {
self.swap_unsafe(pos1, pos2);
}
Ok(())
}
}
}
impl<Item> ScalarAsArray for Item
where
Item: Default + Copy,
{
fn as_array<const NDIM: usize>(&self) -> StaticArray<Self, NDIM, 1> {
let mut arr = Array::new(BaseArray::new(ArrayContainer::<Item, 1>::new(), [1; NDIM]));
*unsafe { arr.get_unchecked_mut([0; NDIM]) } = *self;
arr
}
}
impl<Item, const NDIM: usize> StaticArray<Item, NDIM, 1>
where
Item: Copy + Default,
{
pub fn one_element_array() -> Self {
Array::new(BaseArray::new(ArrayContainer::<Item, 1>::new(), [1; NDIM]))
}
}
macro_rules! impl_linspace_float {
($dtype:ty) => {
impl<ArrayImpl> Linspace<$dtype> for Array<ArrayImpl, 1>
where
ArrayImpl: MutableArrayImpl<$dtype, 1>,
{
fn linspace(&mut self, a: $dtype, b: $dtype) {
let n = self.len() - 1;
assert!(b > a, "`linspace` requires `{b} = b > a = {a} ");
assert!(n > 1, "`linspace` requires `{n} = self.len() > 1");
let delta = (b - a) / (n as $dtype);
for (ind, elem) in self.iter_mut().enumerate().take(n) {
*elem = a + delta * ind as $dtype;
}
self[[n]] = b;
}
}
};
}
impl_linspace_float!(f32);
impl_linspace_float!(f64);
#[cfg(test)]
mod test {
use approx::assert_relative_eq;
use crate::assert_array_relative_eq;
use super::*;
#[test]
pub fn test_add() {
let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![2.5, 4.3, 7.0].into();
crate::assert_array_relative_eq!((a.r() + b.r()).eval(), expected, 1E-10);
}
#[test]
pub fn test_add_assign() {
let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![2.5, 4.3, 7.0].into();
a += b;
crate::assert_array_relative_eq!(a, expected, 1E-10);
}
#[test]
pub fn test_sub() {
let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![-0.5, -0.3, -1.0].into();
crate::assert_array_relative_eq!((a.r() - b.r()).eval(), expected, 1E-10);
}
#[test]
pub fn test_sub_assign() {
let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![-0.5, -0.3, -1.0].into();
a -= b;
crate::assert_array_relative_eq!(a, expected, 1E-10);
}
#[test]
pub fn test_cmp_mul() {
let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
crate::assert_array_relative_eq!(a.r() * b.r(), expected, 1E-10);
}
#[test]
pub fn test_cmp_mul_assign() {
let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let b: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
let expected: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
a *= b;
crate::assert_array_relative_eq!(a, expected, 1E-10);
}
#[test]
pub fn test_scalar_mul() {
let a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let expected: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
crate::assert_array_relative_eq!(2.0_f64 * a.r(), expected, 1E-10);
}
#[test]
pub fn test_scalar_mul_assign() {
let mut a: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let expected: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
a *= 2.0;
crate::assert_array_relative_eq!(a.r(), expected, 1E-10);
}
#[test]
pub fn test_cmp_div() {
let a: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
let b: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let expected: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
crate::assert_array_relative_eq!(a.r() / b.r(), expected, 1E-10);
}
#[test]
pub fn test_cmp_div_assign() {
let mut a: DynArray<_, 1> = vec![1.5, 4.6, 12.0].into();
let b: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
let expected: DynArray<_, 1> = vec![1.5, 2.3, 4.0].into();
a /= b;
crate::assert_array_relative_eq!(a, expected, 1E-10);
}
#[test]
pub fn test_scalar_div() {
let a: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
let expected: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
crate::assert_array_relative_eq!(a.r() / 2.0, expected, 1E-10);
}
#[test]
pub fn test_scalar_div_assign() {
let mut a: DynArray<_, 1> = vec![2.0, 4.0, 6.0].into();
let expected: DynArray<_, 1> = vec![1.0, 2.0, 3.0].into();
a /= 2.0;
crate::assert_array_relative_eq!(a.r(), expected, 1E-10);
}
#[test]
pub fn test_into_subview() {
let mut rng = ChaCha8Rng::seed_from_u64(0);
let mut a = DynArray::<f64, 3>::from_shape([7, 5, 9]);
a.fill_from_equally_distributed(&mut rng);
let view = a.r().into_subview([1, 0, 3], [2, 3, 4]);
assert_eq!(view[[1, 2, 3]], a[[2, 2, 6]])
}
#[test]
pub fn test_apply() {
let mut rng = ChaCha8Rng::seed_from_u64(0);
let alpha = 1.5;
let beta = 3.6;
let mut a = DynArray::<f64, 2>::from_shape([3, 2]);
a.fill_from_equally_distributed(&mut rng);
let mut x = DynArray::<f64, 1>::from_shape([2]);
x.fill_from_equally_distributed(&mut rng);
let mut y = DynArray::<f64, 1>::from_shape([3]);
y.fill_from_equally_distributed(&mut rng);
let y2 = y.eval();
a.apply(alpha, &x, beta, &mut y);
for (row_index, row) in a.row_iter().enumerate() {
assert_relative_eq!(
alpha * row.inner(&x).unwrap() + beta * y2[[row_index]],
y[[row_index]],
epsilon = 1E-10
);
}
let mut x = DynArray::<f64, 2>::from_shape([2, 5]);
x.fill_from_equally_distributed(&mut rng);
let mut y = DynArray::<f64, 2>::from_shape([3, 5]);
y.fill_from_equally_distributed(&mut rng);
let y2 = y.eval();
a.apply(alpha, &x, beta, &mut y);
for (row_index, row) in a.row_iter().enumerate() {
for (col_index, col) in x.col_iter().enumerate() {
assert_relative_eq!(
alpha * row.inner(&col).unwrap() + beta * y2[[row_index, col_index]],
y[[row_index, col_index]],
epsilon = 1E-10
);
}
}
}
#[test]
pub fn test_dot() {
let mut rng = ChaCha8Rng::seed_from_u64(0);
let mut a = DynArray::<f64, 2>::from_shape([3, 2]);
a.fill_from_equally_distributed(&mut rng);
let mut x = DynArray::<f64, 1>::from_shape([2]);
x.fill_from_equally_distributed(&mut rng);
let mut y = DynArray::<f64, 1>::from_shape([3]);
a.apply(1.0, &x, 0.0, &mut y);
crate::assert_array_relative_eq!(y, a.dot(&x), 1E-10);
let mut x = DynArray::<f64, 2>::from_shape([2, 5]);
x.fill_from_equally_distributed(&mut rng);
let mut y = DynArray::<f64, 2>::from_shape([3, 5]);
a.apply(1.0, &x, 0.0, &mut y);
crate::assert_array_relative_eq!(y, a.dot(&x), 1E-10);
}
#[test]
pub fn test_swap() {
let mut a = DynArray::<usize, 1>::from_shape([2]);
a[[0]] = 1;
a[[1]] = 2;
a.swap([0], [1]).unwrap();
assert_eq!(a[[0]], 2);
assert_eq!(a[[1]], 1);
}
#[test]
pub fn test_multislice_iterator() {
let mut arr = DynArray::<f64, _>::from_shape([2, 3, 4, 5]);
arr.fill_from_seed_equally_distributed(0);
let mut count = 0;
for (indices, slice) in arr.r().multislice_iter::<_, 2>([2, 3]) {
assert_array_relative_eq!(slice, arr.r().multislice::<_, 2>([2, 3], indices), 1E-12);
count += 1;
}
assert_eq!(count, 20);
}
#[test]
pub fn test_multislice_iterator_mut() {
let mut arr = DynArray::<f64, _>::from_shape([2, 3, 4, 5]);
arr.fill_from_seed_equally_distributed(0);
let expected = arr.eval();
let mut count = 0;
for (indices, slice) in arr.r_mut().multislice_iter_mut::<_, 2>([2, 3]) {
assert_array_relative_eq!(
slice,
expected.r().multislice::<_, 2>([2, 3], indices),
1E-12
);
count += 1;
}
assert_eq!(count, 20);
}
}