1use std::ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};
2
3use super::default_solver::DefaultSolver;
4use super::utils::*;
5use super::{DenseMatrix, Matrix, MatrixCommon, MatrixView, MatrixViewMut};
6use crate::error::LaError;
7use crate::scalar::{IndexType, Scalar, Scale};
8use crate::VectorIndex;
9use crate::{Dense, DenseRef, FaerContext, FaerScalar, FaerVec, Vector, VectorViewMut};
10use crate::{FaerLU, FaerVecMut, FaerVecRef};
11
12use faer::{get_global_parallelism, unzip, zip, Accum};
13use faer::{linalg::matmul::matmul, Mat, MatMut, MatRef};
14
15#[derive(Clone, Debug, PartialEq)]
16pub struct FaerMat<T: FaerScalar> {
17 pub(crate) data: Mat<T>,
18 pub(crate) context: FaerContext,
19}
20
21#[derive(Clone, Debug, PartialEq)]
22pub struct FaerMatRef<'a, T: FaerScalar> {
23 pub(crate) data: MatRef<'a, T>,
24 pub(crate) context: FaerContext,
25}
26
27#[derive(Debug, PartialEq)]
28pub struct FaerMatMut<'a, T: FaerScalar> {
29 pub(crate) data: MatMut<'a, T>,
30 pub(crate) context: FaerContext,
31}
32
33impl<T: FaerScalar> DefaultSolver for FaerMat<T> {
34 type LS = FaerLU<T>;
35}
36
37impl_matrix_common_ref!(
38 FaerMatMut<'a, T>,
39 FaerVec<T>,
40 FaerContext,
41 MatMut<'a, T>,
42 FaerScalar
43);
44impl_matrix_common_ref!(
45 FaerMatRef<'a, T>,
46 FaerVec<T>,
47 FaerContext,
48 MatRef<'a, T>,
49 FaerScalar
50);
51impl_matrix_common!(FaerMat<T>, FaerVec<T>, FaerContext, Mat<T>, FaerScalar);
52
53macro_rules! impl_mul_scalar {
54 ($mat_type:ty, $out:ty) => {
55 impl<'a, T: FaerScalar> Mul<Scale<T>> for $mat_type {
56 type Output = $out;
57
58 fn mul(self, rhs: Scale<T>) -> Self::Output {
59 let scale: faer::Scale<T> = rhs.into();
60 Self::Output {
61 data: &self.data * scale,
62 context: self.context,
63 }
64 }
65 }
66 };
67}
68
69macro_rules! impl_mul_assign_scalar {
70 ($mat_type:ty) => {
71 impl<T: FaerScalar> MulAssign<Scale<T>> for $mat_type {
72 fn mul_assign(&mut self, rhs: Scale<T>) {
73 let scale: faer::Scale<T> = rhs.into();
74 self.data *= scale;
75 }
76 }
77 };
78}
79
80impl_mul_scalar!(FaerMatRef<'_, T>, FaerMat<T>);
81impl_mul_scalar!(FaerMat<T>, FaerMat<T>);
82impl_mul_scalar!(&FaerMat<T>, FaerMat<T>);
83
84impl_mul_assign_scalar!(FaerMatMut<'_, T>);
85
86impl_add!(FaerMat<T>, &FaerMat<T>, FaerMat<T>, FaerScalar);
87impl_add!(FaerMat<T>, &FaerMatRef<'_, T>, FaerMat<T>, FaerScalar);
88impl_add!(FaerMatRef<'_, T>, &FaerMat<T>, FaerMat<T>, FaerScalar);
89
90impl_sub!(FaerMat<T>, &FaerMat<T>, FaerMat<T>, FaerScalar);
91impl_sub!(FaerMat<T>, &FaerMatRef<'_, T>, FaerMat<T>, FaerScalar);
92impl_sub!(FaerMatRef<'_, T>, &FaerMat<T>, FaerMat<T>, FaerScalar);
93
94impl_add_assign!(FaerMat<T>, &FaerMat<T>, FaerScalar);
95impl_add_assign!(FaerMat<T>, &FaerMatRef<'_, T>, FaerScalar);
96impl_add_assign!(FaerMatMut<'_, T>, &FaerMatRef<'_, T>, FaerScalar);
97impl_add_assign!(FaerMatMut<'_, T>, &FaerMatMut<'_, T>, FaerScalar);
98
99impl_sub_assign!(FaerMat<T>, &FaerMat<T>, FaerScalar);
100impl_sub_assign!(FaerMat<T>, &FaerMatRef<'_, T>, FaerScalar);
101impl_sub_assign!(FaerMatMut<'_, T>, &FaerMatRef<'_, T>, FaerScalar);
102impl_sub_assign!(FaerMatMut<'_, T>, &FaerMatMut<'_, T>, FaerScalar);
103
104impl_index!(FaerMat<T>, FaerScalar);
105impl_index!(FaerMatRef<'_, T>, FaerScalar);
106impl_index_mut!(FaerMat<T>, FaerScalar);
107
108impl<'a, T: FaerScalar> MatrixView<'a> for FaerMatRef<'a, T> {
109 type Owned = FaerMat<T>;
110
111 fn into_owned(self) -> Self::Owned {
112 Self::Owned {
113 data: self.data.to_owned(),
114 context: self.context,
115 }
116 }
117
118 fn gemv_o(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) {
119 y.mul_assign(Scale(beta));
120 matmul(
121 y.data.as_mut(),
122 Accum::Add,
123 self.data.as_ref(),
124 x.data.as_ref(),
125 alpha,
126 get_global_parallelism(),
127 );
128 }
129 fn gemv_v(
130 &self,
131 alpha: Self::T,
132 x: &<Self::V as crate::vector::Vector>::View<'_>,
133 beta: Self::T,
134 y: &mut Self::V,
135 ) {
136 y.mul_assign(Scale(beta));
137 matmul(
138 y.data.as_mut(),
139 Accum::Add,
140 self.data.as_ref(),
141 x.data.as_ref(),
142 alpha,
143 get_global_parallelism(),
144 );
145 }
146}
147
148impl<'a, T: FaerScalar> MatrixViewMut<'a> for FaerMatMut<'a, T> {
149 type Owned = FaerMat<T>;
150 type View = FaerMatRef<'a, T>;
151
152 fn into_owned(self) -> Self::Owned {
153 Self::Owned {
154 data: self.data.to_owned(),
155 context: self.context,
156 }
157 }
158
159 fn gemm_oo(&mut self, alpha: Self::T, a: &Self::Owned, b: &Self::Owned, beta: Self::T) {
160 self.mul_assign(Scale(beta));
161 matmul(
162 self.data.as_mut(),
163 Accum::Add,
164 a.data.as_ref(),
165 b.data.as_ref(),
166 alpha,
167 get_global_parallelism(),
168 )
169 }
170 fn gemm_vo(&mut self, alpha: Self::T, a: &Self::View, b: &Self::Owned, beta: Self::T) {
171 self.mul_assign(Scale(beta));
172 matmul(
173 self.data.as_mut(),
174 Accum::Add,
175 a.data.as_ref(),
176 b.data.as_ref(),
177 alpha,
178 get_global_parallelism(),
179 )
180 }
181}
182
183impl<T: FaerScalar> DenseMatrix for FaerMat<T> {
184 type View<'a> = FaerMatRef<'a, T>;
185 type ViewMut<'a> = FaerMatMut<'a, T>;
186
187 fn from_vec(nrows: IndexType, ncols: IndexType, data: Vec<Self::T>, ctx: Self::C) -> Self {
188 let data = Mat::from_fn(nrows, ncols, |i, j| data[i + j * nrows]);
189 Self { data, context: ctx }
190 }
191
192 fn resize_cols(&mut self, ncols: IndexType) {
193 if ncols == self.ncols() {
194 return;
195 }
196 let nrows = self.nrows();
197 self.data.resize_with(nrows, ncols, |_, _| T::zero());
198 }
199
200 fn get_index(&self, i: IndexType, j: IndexType) -> Self::T {
201 self.data[(i, j)]
202 }
203
204 fn gemm(&mut self, alpha: Self::T, a: &Self, b: &Self, beta: Self::T) {
205 self.data.mul_assign(faer::Scale(beta));
206 matmul(
207 self.data.as_mut(),
208 Accum::Add,
209 a.data.as_ref(),
210 b.data.as_ref(),
211 alpha,
212 get_global_parallelism(),
213 )
214 }
215 fn column_mut(&mut self, i: usize) -> <Self::V as Vector>::ViewMut<'_> {
216 let data = self.data.get_mut(0..self.nrows(), i);
217 FaerVecMut {
218 data,
219 context: self.context,
220 }
221 }
222
223 fn columns_mut(&mut self, start: usize, end: usize) -> Self::ViewMut<'_> {
224 let data = self.data.get_mut(0..self.data.nrows(), start..end);
225 FaerMatMut {
226 data,
227 context: self.context,
228 }
229 }
230
231 fn set_index(&mut self, i: IndexType, j: IndexType, value: Self::T) {
232 self.data[(i, j)] = value;
233 }
234
235 fn column(&self, i: usize) -> <Self::V as Vector>::View<'_> {
236 let data = self.data.get(0..self.data.nrows(), i);
237 FaerVecRef {
238 data,
239 context: self.context,
240 }
241 }
242 fn columns(&self, start: usize, end: usize) -> Self::View<'_> {
243 let data = self.data.get(0..self.nrows(), start..end);
244 FaerMatRef {
245 data,
246 context: self.context,
247 }
248 }
249
250 fn column_axpy(&mut self, alpha: Self::T, j: IndexType, i: IndexType) {
251 if i > self.ncols() {
252 panic!("Column index out of bounds");
253 }
254 if j > self.ncols() {
255 panic!("Column index out of bounds");
256 }
257 if i == j {
258 panic!("Column index cannot be the same");
259 }
260 for k in 0..self.nrows() {
261 let value =
262 unsafe { *self.data.get_unchecked(k, i) + alpha * *self.data.get_unchecked(k, j) };
263 unsafe { *self.data.get_mut_unchecked(k, i) = value };
264 }
265 }
266}
267
268impl<T: FaerScalar> Matrix for FaerMat<T> {
269 type Sparsity = Dense<Self>;
270 type SparsityRef<'a> = DenseRef<'a, Self>;
271
272 fn sparsity(&self) -> Option<Self::SparsityRef<'_>> {
273 None
274 }
275
276 fn context(&self) -> &Self::C {
277 &self.context
278 }
279 fn inner_mut(&mut self) -> &mut Self::Inner {
280 &mut self.data
281 }
282
283 fn gather(&mut self, other: &Self, indices: &<Self::V as Vector>::Index) {
284 assert_eq!(indices.len(), self.nrows() * self.ncols());
285 if self.nrows() == 0 || self.ncols() == 0 {
286 return;
287 }
288 let mut idx = indices.data.iter().peekable();
289 for j in 0..self.ncols() {
290 let other_col = other.data.col(*idx.peek().unwrap() / other.nrows());
291 for self_ij in self.data.col_mut(j).iter_mut() {
292 let other_i = idx.next().unwrap() % other.nrows();
293 *self_ij = other_col[other_i];
294 }
295 }
296 }
297
298 fn set_data_with_indices(
299 &mut self,
300 dst_indices: &<Self::V as Vector>::Index,
301 src_indices: &<Self::V as Vector>::Index,
302 data: &Self::V,
303 ) {
304 for (dst_i, src_i) in dst_indices.data.iter().zip(src_indices.data.iter()) {
305 let i = dst_i % self.nrows();
306 let j = dst_i / self.nrows();
307 self.data[(i, j)] = data[*src_i];
308 }
309 }
310
311 fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V) {
312 v.add_assign(&self.column(j));
313 }
314
315 fn triplet_iter(
316 &self,
317 ) -> (
318 impl Iterator<Item = (IndexType, IndexType)> + '_,
319 impl Iterator<Item = Self::T> + '_,
320 ) {
321 let indices: Vec<_> = (0..self.ncols())
322 .flat_map(move |j| (0..self.nrows()).map(move |i| (i, j)))
323 .collect();
324 let values: Vec<_> = indices.iter().map(|&(i, j)| self.data[(i, j)]).collect();
325 (indices.into_iter(), values.into_iter())
326 }
327
328 fn try_from_triplets(
329 nrows: IndexType,
330 ncols: IndexType,
331 indices: Vec<(IndexType, IndexType)>,
332 values: Vec<Self::T>,
333 ctx: Self::C,
334 ) -> Result<Self, LaError> {
335 assert_eq!(
336 values.len(),
337 indices.len(),
338 "values.len() must equal indices.len() for non-batched backend"
339 );
340 let mut m = Mat::zeros(nrows, ncols);
341 for ((i, j), v) in indices.iter().zip(values) {
342 m[(*i, *j)] = v;
343 }
344 Ok(Self {
345 data: m,
346 context: ctx,
347 })
348 }
349 fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) {
350 y.mul_assign(Scale(beta));
351 matmul(
352 y.data.as_mut(),
353 Accum::Add,
354 self.data.as_ref(),
355 x.data.as_ref(),
356 alpha,
357 get_global_parallelism(),
358 );
359 }
360 fn zeros(nrows: IndexType, ncols: IndexType, ctx: Self::C) -> Self {
361 let data = Mat::zeros(nrows, ncols);
362 Self { data, context: ctx }
363 }
364 fn copy_from(&mut self, other: &Self) {
365 self.data.copy_from(&other.data);
366 }
367 fn from_diagonal(v: &Self::V) -> Self {
368 let dim = v.len();
369 let data = Mat::from_fn(dim, dim, |i, j| if i == j { v[i] } else { T::zero() });
370 Self {
371 data,
372 context: *v.context(),
373 }
374 }
375 fn partition_indices_by_zero_diagonal(
376 &self,
377 ) -> (<Self::V as Vector>::Index, <Self::V as Vector>::Index) {
378 let diagonal = self.data.diagonal().column_vector();
379 let (zero_indices, nonzero_indices) = diagonal.iter().enumerate().fold(
380 (Vec::new(), Vec::new()),
381 |(mut zero_indices, mut nonzero_indices), (i, &v)| {
382 if v.is_zero() {
383 zero_indices.push(i);
384 } else {
385 nonzero_indices.push(i);
386 }
387 (zero_indices, nonzero_indices)
388 },
389 );
390 (
391 <Self::V as Vector>::Index::from_vec(zero_indices, self.context),
392 <Self::V as Vector>::Index::from_vec(nonzero_indices, self.context),
393 )
394 }
395 fn set_column(&mut self, j: IndexType, v: &Self::V) {
396 self.column_mut(j).copy_from(v);
397 }
398
399 fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) {
400 zip!(self.data.as_mut(), x.data.as_ref(), y.data.as_ref())
401 .for_each(|unzip!(s, x, y)| *s = *x + beta * *y);
402 }
403
404 fn new_from_sparsity(
405 nrows: IndexType,
406 ncols: IndexType,
407 _sparsity: Option<Self::Sparsity>,
408 ctx: Self::C,
409 ) -> Self {
410 Self::zeros(nrows, ncols, ctx)
411 }
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417
418 #[test]
419 fn test_column_axpy() {
420 super::super::tests::test_column_axpy::<FaerMat<f64>>();
421 }
422
423 #[test]
424 fn test_partition_indices_by_zero_diagonal() {
425 super::super::tests::test_partition_indices_by_zero_diagonal::<FaerMat<f64>>();
426 }
427
428 #[test]
429 fn test_resize_cols() {
430 super::super::tests::test_resize_cols::<FaerMat<f64>>();
431 }
432
433 super::super::generate_matrix_tests_nonbatched!(faer, FaerMat<f64>);
434 super::super::generate_dense_matrix_tests_nonbatched!(faer, FaerMat<f64>);
435}