use num::{One, Zero}; use std::iter; use std::ops::{ Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign, }; use simba::scalar::{ClosedAdd, ClosedDiv, ClosedMul, ClosedNeg, ClosedSub}; use crate::base::allocator::{Allocator, SameShapeAllocator, SameShapeC, SameShapeR}; use crate::base::constraint::{ AreMultipliable, DimEq, SameNumberOfColumns, SameNumberOfRows, ShapeConstraint, }; use crate::base::dimension::{Dim, DimMul, DimName, DimProd, Dynamic}; use crate::base::storage::{ContiguousStorageMut, Storage, StorageMut}; use crate::base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, MatrixSum, Scalar, VectorSliceN}; use crate::SimdComplexField; /* * * Indexing. * */ impl> Index for Matrix { type Output = N; #[inline] fn index(&self, i: usize) -> &Self::Output { let ij = self.vector_to_matrix_index(i); &self[ij] } } impl Index<(usize, usize)> for Matrix where N: Scalar, S: Storage, { type Output = N; #[inline] fn index(&self, ij: (usize, usize)) -> &Self::Output { let shape = self.shape(); assert!( ij.0 < shape.0 && ij.1 < shape.1, "Matrix index out of bounds." ); unsafe { self.get_unchecked((ij.0, ij.1)) } } } // Mutable versions. impl> IndexMut for Matrix { #[inline] fn index_mut(&mut self, i: usize) -> &mut N { let ij = self.vector_to_matrix_index(i); &mut self[ij] } } impl IndexMut<(usize, usize)> for Matrix where N: Scalar, S: StorageMut, { #[inline] fn index_mut(&mut self, ij: (usize, usize)) -> &mut N { let shape = self.shape(); assert!( ij.0 < shape.0 && ij.1 < shape.1, "Matrix index out of bounds." ); unsafe { self.get_unchecked_mut((ij.0, ij.1)) } } } /* * * Neg * */ impl Neg for Matrix where N: Scalar + ClosedNeg, S: Storage, DefaultAllocator: Allocator, { type Output = MatrixMN; #[inline] fn neg(self) -> Self::Output { let mut res = self.into_owned(); res.neg_mut(); res } } impl<'a, N, R: Dim, C: Dim, S> Neg for &'a Matrix where N: Scalar + ClosedNeg, S: Storage, DefaultAllocator: Allocator, { type Output = MatrixMN; #[inline] fn neg(self) -> Self::Output { -self.clone_owned() } } impl Matrix where N: Scalar + ClosedNeg, S: StorageMut, { /// Negates `self` in-place. #[inline] pub fn neg_mut(&mut self) { for e in self.iter_mut() { *e = -e.inlined_clone() } } } /* * * Addition & Subtraction * */ macro_rules! componentwise_binop_impl( ($Trait: ident, $method: ident, $bound: ident; $TraitAssign: ident, $method_assign: ident, $method_assign_statically_unchecked: ident, $method_assign_statically_unchecked_rhs: ident; $method_to: ident, $method_to_statically_unchecked: ident) => { impl> Matrix where N: Scalar + $bound { /* * * Methods without dimension checking at compile-time. * This is useful for code reuse because the sum representative system does not plays * easily with static checks. * */ #[inline] fn $method_to_statically_unchecked(&self, rhs: &Matrix, out: &mut Matrix) where SB: Storage, SC: StorageMut { assert_eq!(self.shape(), rhs.shape(), "Matrix addition/subtraction dimensions mismatch."); assert_eq!(self.shape(), out.shape(), "Matrix addition/subtraction output dimensions mismatch."); // This is the most common case and should be deduced at compile-time. // TODO: use specialization instead? if self.data.is_contiguous() && rhs.data.is_contiguous() && out.data.is_contiguous() { let arr1 = self.data.as_slice(); let arr2 = rhs.data.as_slice(); let out = out.data.as_mut_slice(); for i in 0 .. arr1.len() { unsafe { *out.get_unchecked_mut(i) = arr1.get_unchecked(i).inlined_clone().$method(arr2.get_unchecked(i).inlined_clone()); } } } else { for j in 0 .. self.ncols() { for i in 0 .. self.nrows() { unsafe { let val = self.get_unchecked((i, j)).inlined_clone().$method(rhs.get_unchecked((i, j)).inlined_clone()); *out.get_unchecked_mut((i, j)) = val; } } } } } #[inline] fn $method_assign_statically_unchecked(&mut self, rhs: &Matrix) where R2: Dim, C2: Dim, SA: StorageMut, SB: Storage { assert_eq!(self.shape(), rhs.shape(), "Matrix addition/subtraction dimensions mismatch."); // This is the most common case and should be deduced at compile-time. // TODO: use specialization instead? if self.data.is_contiguous() && rhs.data.is_contiguous() { let arr1 = self.data.as_mut_slice(); let arr2 = rhs.data.as_slice(); for i in 0 .. arr2.len() { unsafe { arr1.get_unchecked_mut(i).$method_assign(arr2.get_unchecked(i).inlined_clone()); } } } else { for j in 0 .. rhs.ncols() { for i in 0 .. rhs.nrows() { unsafe { self.get_unchecked_mut((i, j)).$method_assign(rhs.get_unchecked((i, j)).inlined_clone()) } } } } } #[inline] fn $method_assign_statically_unchecked_rhs(&self, rhs: &mut Matrix) where R2: Dim, C2: Dim, SB: StorageMut { assert_eq!(self.shape(), rhs.shape(), "Matrix addition/subtraction dimensions mismatch."); // This is the most common case and should be deduced at compile-time. // TODO: use specialization instead? if self.data.is_contiguous() && rhs.data.is_contiguous() { let arr1 = self.data.as_slice(); let arr2 = rhs.data.as_mut_slice(); for i in 0 .. arr1.len() { unsafe { let res = arr1.get_unchecked(i).inlined_clone().$method(arr2.get_unchecked(i).inlined_clone()); *arr2.get_unchecked_mut(i) = res; } } } else { for j in 0 .. self.ncols() { for i in 0 .. self.nrows() { unsafe { let r = rhs.get_unchecked_mut((i, j)); *r = self.get_unchecked((i, j)).inlined_clone().$method(r.inlined_clone()) } } } } } /* * * Methods without dimension checking at compile-time. * This is useful for code reuse because the sum representative system does not plays * easily with static checks. * */ /// Equivalent to `self + rhs` but stores the result into `out` to avoid allocations. #[inline] pub fn $method_to(&self, rhs: &Matrix, out: &mut Matrix) where SB: Storage, SC: StorageMut, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns + SameNumberOfRows + SameNumberOfColumns { self.$method_to_statically_unchecked(rhs, out) } } impl<'b, N, R1, C1, R2, C2, SA, SB> $Trait<&'b Matrix> for Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: Storage, SB: Storage, DefaultAllocator: SameShapeAllocator, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { type Output = MatrixSum; #[inline] fn $method(self, rhs: &'b Matrix) -> Self::Output { assert_eq!(self.shape(), rhs.shape(), "Matrix addition/subtraction dimensions mismatch."); let mut res = self.into_owned_sum::(); res.$method_assign_statically_unchecked(rhs); res } } impl<'a, N, R1, C1, R2, C2, SA, SB> $Trait> for &'a Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: Storage, SB: Storage, DefaultAllocator: SameShapeAllocator, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { type Output = MatrixSum; #[inline] fn $method(self, rhs: Matrix) -> Self::Output { let mut rhs = rhs.into_owned_sum::(); assert_eq!(self.shape(), rhs.shape(), "Matrix addition/subtraction dimensions mismatch."); self.$method_assign_statically_unchecked_rhs(&mut rhs); rhs } } impl $Trait> for Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: Storage, SB: Storage, DefaultAllocator: SameShapeAllocator, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { type Output = MatrixSum; #[inline] fn $method(self, rhs: Matrix) -> Self::Output { self.$method(&rhs) } } impl<'a, 'b, N, R1, C1, R2, C2, SA, SB> $Trait<&'b Matrix> for &'a Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: Storage, SB: Storage, DefaultAllocator: SameShapeAllocator, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { type Output = MatrixSum; #[inline] fn $method(self, rhs: &'b Matrix) -> Self::Output { let mut res = unsafe { let (nrows, ncols) = self.shape(); let nrows: SameShapeR = Dim::from_usize(nrows); let ncols: SameShapeC = Dim::from_usize(ncols); #[cfg(feature="no_unsound_assume_init")] { unimplemented!() } #[cfg(not(feature="no_unsound_assume_init"))] { Matrix::new_uninitialized_generic(nrows, ncols).assume_init() } }; self.$method_to_statically_unchecked(rhs, &mut res); res } } impl<'b, N, R1, C1, R2, C2, SA, SB> $TraitAssign<&'b Matrix> for Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: StorageMut, SB: Storage, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { #[inline] fn $method_assign(&mut self, rhs: &'b Matrix) { self.$method_assign_statically_unchecked(rhs) } } impl $TraitAssign> for Matrix where R1: Dim, C1: Dim, R2: Dim, C2: Dim, N: Scalar + $bound, SA: StorageMut, SB: Storage, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns { #[inline] fn $method_assign(&mut self, rhs: Matrix) { self.$method_assign(&rhs) } } } ); componentwise_binop_impl!(Add, add, ClosedAdd; AddAssign, add_assign, add_assign_statically_unchecked, add_assign_statically_unchecked_mut; add_to, add_to_statically_unchecked); componentwise_binop_impl!(Sub, sub, ClosedSub; SubAssign, sub_assign, sub_assign_statically_unchecked, sub_assign_statically_unchecked_mut; sub_to, sub_to_statically_unchecked); impl iter::Sum for MatrixMN where N: Scalar + ClosedAdd + Zero, DefaultAllocator: Allocator, { fn sum>>(iter: I) -> MatrixMN { iter.fold(Matrix::zero(), |acc, x| acc + x) } } impl iter::Sum for MatrixMN where N: Scalar + ClosedAdd + Zero, DefaultAllocator: Allocator, { /// # Example /// ``` /// # use nalgebra::DVector; /// assert_eq!(vec![DVector::repeat(3, 1.0f64), /// DVector::repeat(3, 1.0f64), /// DVector::repeat(3, 1.0f64)].into_iter().sum::>(), /// DVector::repeat(3, 1.0f64) + DVector::repeat(3, 1.0f64) + DVector::repeat(3, 1.0f64)); /// ``` /// /// # Panics /// Panics if the iterator is empty: /// ```should_panic /// # use std::iter; /// # use nalgebra::DMatrix; /// iter::empty::>().sum::>(); // panics! /// ``` fn sum>>(mut iter: I) -> MatrixMN { if let Some(first) = iter.next() { iter.fold(first, |acc, x| acc + x) } else { panic!("Cannot compute `sum` of empty iterator.") } } } impl<'a, N, R: DimName, C: DimName> iter::Sum<&'a MatrixMN> for MatrixMN where N: Scalar + ClosedAdd + Zero, DefaultAllocator: Allocator, { fn sum>>(iter: I) -> MatrixMN { iter.fold(Matrix::zero(), |acc, x| acc + x) } } impl<'a, N, C: Dim> iter::Sum<&'a MatrixMN> for MatrixMN where N: Scalar + ClosedAdd + Zero, DefaultAllocator: Allocator, { /// # Example /// ``` /// # use nalgebra::DVector; /// let v = &DVector::repeat(3, 1.0f64); /// /// assert_eq!(vec![v, v, v].into_iter().sum::>(), /// v + v + v); /// ``` /// /// # Panics /// Panics if the iterator is empty: /// ```should_panic /// # use std::iter; /// # use nalgebra::DMatrix; /// iter::empty::<&DMatrix>().sum::>(); // panics! /// ``` fn sum>>( mut iter: I, ) -> MatrixMN { if let Some(first) = iter.next() { iter.fold(first.clone(), |acc, x| acc + x) } else { panic!("Cannot compute `sum` of empty iterator.") } } } /* * * Multiplication * */ // Matrix × Scalar // Matrix / Scalar macro_rules! componentwise_scalarop_impl( ($Trait: ident, $method: ident, $bound: ident; $TraitAssign: ident, $method_assign: ident) => { impl $Trait for Matrix where N: Scalar + $bound, S: Storage, DefaultAllocator: Allocator { type Output = MatrixMN; #[inline] fn $method(self, rhs: N) -> Self::Output { let mut res = self.into_owned(); // XXX: optimize our iterator! // // Using our own iterator prevents loop unrolling, which breaks some optimization // (like SIMD). On the other hand, using the slice iterator is 4x faster. // for left in res.iter_mut() { for left in res.as_mut_slice().iter_mut() { *left = left.inlined_clone().$method(rhs.inlined_clone()) } res } } impl<'a, N, R: Dim, C: Dim, S> $Trait for &'a Matrix where N: Scalar + $bound, S: Storage, DefaultAllocator: Allocator { type Output = MatrixMN; #[inline] fn $method(self, rhs: N) -> Self::Output { self.clone_owned().$method(rhs) } } impl $TraitAssign for Matrix where N: Scalar + $bound, S: StorageMut { #[inline] fn $method_assign(&mut self, rhs: N) { for j in 0 .. self.ncols() { for i in 0 .. self.nrows() { unsafe { self.get_unchecked_mut((i, j)).$method_assign(rhs.inlined_clone()) }; } } } } } ); componentwise_scalarop_impl!(Mul, mul, ClosedMul; MulAssign, mul_assign); componentwise_scalarop_impl!(Div, div, ClosedDiv; DivAssign, div_assign); macro_rules! left_scalar_mul_impl( ($($T: ty),* $(,)*) => {$( impl> Mul> for $T where DefaultAllocator: Allocator<$T, R, C> { type Output = MatrixMN<$T, R, C>; #[inline] fn mul(self, rhs: Matrix<$T, R, C, S>) -> Self::Output { let mut res = rhs.into_owned(); // XXX: optimize our iterator! // // Using our own iterator prevents loop unrolling, which breaks some optimization // (like SIMD). On the other hand, using the slice iterator is 4x faster. // for rhs in res.iter_mut() { for rhs in res.as_mut_slice().iter_mut() { *rhs = self * *rhs } res } } impl<'b, R: Dim, C: Dim, S: Storage<$T, R, C>> Mul<&'b Matrix<$T, R, C, S>> for $T where DefaultAllocator: Allocator<$T, R, C> { type Output = MatrixMN<$T, R, C>; #[inline] fn mul(self, rhs: &'b Matrix<$T, R, C, S>) -> Self::Output { self * rhs.clone_owned() } } )*} ); left_scalar_mul_impl!(u8, u16, u32, u64, usize, i8, i16, i32, i64, isize, f32, f64); // Matrix × Matrix impl<'a, 'b, N, R1: Dim, C1: Dim, R2: Dim, C2: Dim, SA, SB> Mul<&'b Matrix> for &'a Matrix where N: Scalar + Zero + One + ClosedAdd + ClosedMul, SA: Storage, SB: Storage, DefaultAllocator: Allocator, ShapeConstraint: AreMultipliable, { type Output = MatrixMN; #[inline] fn mul(self, rhs: &'b Matrix) -> Self::Output { let mut res = unsafe { crate::zero_or_uninitialized_generic!(self.data.shape().0, rhs.data.shape().1) }; self.mul_to(rhs, &mut res); res } } impl<'a, N, R1: Dim, C1: Dim, R2: Dim, C2: Dim, SA, SB> Mul> for &'a Matrix where N: Scalar + Zero + One + ClosedAdd + ClosedMul, SB: Storage, SA: Storage, DefaultAllocator: Allocator, ShapeConstraint: AreMultipliable, { type Output = MatrixMN; #[inline] fn mul(self, rhs: Matrix) -> Self::Output { self * &rhs } } impl<'b, N, R1: Dim, C1: Dim, R2: Dim, C2: Dim, SA, SB> Mul<&'b Matrix> for Matrix where N: Scalar + Zero + One + ClosedAdd + ClosedMul, SB: Storage, SA: Storage, DefaultAllocator: Allocator, ShapeConstraint: AreMultipliable, { type Output = MatrixMN; #[inline] fn mul(self, rhs: &'b Matrix) -> Self::Output { &self * rhs } } impl Mul> for Matrix where N: Scalar + Zero + One + ClosedAdd + ClosedMul, SB: Storage, SA: Storage, DefaultAllocator: Allocator, ShapeConstraint: AreMultipliable, { type Output = MatrixMN; #[inline] fn mul(self, rhs: Matrix) -> Self::Output { &self * &rhs } } // TODO: this is too restrictive: // − we can't use `a *= b` when `a` is a mutable slice. // − we can't use `a *= b` when C2 is not equal to C1. impl MulAssign> for Matrix where R1: Dim, C1: Dim, R2: Dim, N: Scalar + Zero + One + ClosedAdd + ClosedMul, SB: Storage, SA: ContiguousStorageMut + Clone, ShapeConstraint: AreMultipliable, DefaultAllocator: Allocator, { #[inline] fn mul_assign(&mut self, rhs: Matrix) { *self = &*self * rhs } } impl<'b, N, R1, C1, R2, SA, SB> MulAssign<&'b Matrix> for Matrix where R1: Dim, C1: Dim, R2: Dim, N: Scalar + Zero + One + ClosedAdd + ClosedMul, SB: Storage, SA: ContiguousStorageMut + Clone, ShapeConstraint: AreMultipliable, // TODO: this is too restrictive. See comments for the non-ref version. DefaultAllocator: Allocator, { #[inline] fn mul_assign(&mut self, rhs: &'b Matrix) { *self = &*self * rhs } } /// # Special multiplications. impl Matrix where N: Scalar + Zero + One + ClosedAdd + ClosedMul, SA: Storage, { /// Equivalent to `self.transpose() * rhs`. #[inline] pub fn tr_mul(&self, rhs: &Matrix) -> MatrixMN where SB: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { let mut res = unsafe { crate::zero_or_uninitialized_generic!(self.data.shape().1, rhs.data.shape().1) }; self.tr_mul_to(rhs, &mut res); res } /// Equivalent to `self.adjoint() * rhs`. #[inline] pub fn ad_mul(&self, rhs: &Matrix) -> MatrixMN where N: SimdComplexField, SB: Storage, DefaultAllocator: Allocator, ShapeConstraint: SameNumberOfRows, { let mut res = unsafe { crate::zero_or_uninitialized_generic!(self.data.shape().1, rhs.data.shape().1) }; self.ad_mul_to(rhs, &mut res); res } #[inline(always)] fn xx_mul_to( &self, rhs: &Matrix, out: &mut Matrix, dot: impl Fn( &VectorSliceN, &VectorSliceN, ) -> N, ) where SB: Storage, SC: StorageMut, ShapeConstraint: SameNumberOfRows + DimEq + DimEq, { let (nrows1, ncols1) = self.shape(); let (nrows2, ncols2) = rhs.shape(); let (nrows3, ncols3) = out.shape(); assert!( nrows1 == nrows2, "Matrix multiplication dimensions mismatch {:?} and {:?}: left rows != right rows.", self.shape(), rhs.shape() ); assert!( ncols1 == nrows3, "Matrix multiplication output dimensions mismatch {:?} and {:?}: left cols != right rows.", self.shape(), out.shape() ); assert!( ncols2 == ncols3, "Matrix multiplication output dimensions mismatch {:?} and {:?}: left cols != right cols", rhs.shape(), out.shape() ); for i in 0..ncols1 { for j in 0..ncols2 { let dot = dot(&self.column(i), &rhs.column(j)); unsafe { *out.get_unchecked_mut((i, j)) = dot }; } } } /// Equivalent to `self.transpose() * rhs` but stores the result into `out` to avoid /// allocations. #[inline] pub fn tr_mul_to( &self, rhs: &Matrix, out: &mut Matrix, ) where SB: Storage, SC: StorageMut, ShapeConstraint: SameNumberOfRows + DimEq + DimEq, { self.xx_mul_to(rhs, out, |a, b| a.dot(b)) } /// Equivalent to `self.adjoint() * rhs` but stores the result into `out` to avoid /// allocations. #[inline] pub fn ad_mul_to( &self, rhs: &Matrix, out: &mut Matrix, ) where N: SimdComplexField, SB: Storage, SC: StorageMut, ShapeConstraint: SameNumberOfRows + DimEq + DimEq, { self.xx_mul_to(rhs, out, |a, b| a.dotc(b)) } /// Equivalent to `self * rhs` but stores the result into `out` to avoid allocations. #[inline] pub fn mul_to( &self, rhs: &Matrix, out: &mut Matrix, ) where SB: Storage, SC: StorageMut, ShapeConstraint: SameNumberOfRows + SameNumberOfColumns + AreMultipliable, { out.gemm(N::one(), self, rhs, N::zero()); } /// The kronecker product of two matrices (aka. tensor product of the corresponding linear /// maps). pub fn kronecker( &self, rhs: &Matrix, ) -> MatrixMN, DimProd> where N: ClosedMul, R1: DimMul, C1: DimMul, SB: Storage, DefaultAllocator: Allocator, DimProd>, { let (nrows1, ncols1) = self.data.shape(); let (nrows2, ncols2) = rhs.data.shape(); let mut res = unsafe { crate::zero_or_uninitialized_generic!(nrows1.mul(nrows2), ncols1.mul(ncols2)) }; { let mut data_res = res.data.ptr_mut(); for j1 in 0..ncols1.value() { for j2 in 0..ncols2.value() { for i1 in 0..nrows1.value() { unsafe { let coeff = self.get_unchecked((i1, j1)).inlined_clone(); for i2 in 0..nrows2.value() { *data_res = coeff.inlined_clone() * rhs.get_unchecked((i2, j2)).inlined_clone(); data_res = data_res.offset(1); } } } } } } res } } impl iter::Product for MatrixN where N: Scalar + Zero + One + ClosedMul + ClosedAdd, DefaultAllocator: Allocator, { fn product>>(iter: I) -> MatrixN { iter.fold(Matrix::one(), |acc, x| acc * x) } } impl<'a, N, D: DimName> iter::Product<&'a MatrixN> for MatrixN where N: Scalar + Zero + One + ClosedMul + ClosedAdd, DefaultAllocator: Allocator, { fn product>>(iter: I) -> MatrixN { iter.fold(Matrix::one(), |acc, x| acc * x) } }