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