Skip to main content

diffsol_la/matrix/
dense_faer_serial.rs

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}