del_fem_cpu/
sparse_square.rs

1//! sparse matrix class and functions
2
3pub fn set_fixed_bc<T, const N: usize, const NN: usize>(
4    val_dia: T,
5    bc_flag: &[[i32; N]],
6    row2val: &mut [[T; NN]],
7    idx2val: &mut [[T; NN]],
8    row2idx: &[usize],
9    idx2col: &[usize],
10) where
11    T: num_traits::Float,
12{
13    let num_blk = bc_flag.len();
14    assert_eq!(bc_flag.len(), row2val.len());
15    for i_blk in 0..num_blk {
16        // set diagonal
17        for i_dim in 0..N {
18            if bc_flag[i_blk][i_dim] == 0 {
19                continue;
20            };
21            for j_dim in 0..N {
22                row2val[i_blk][i_dim + N * j_dim] = T::zero();
23                row2val[i_blk][j_dim + N * i_dim] = T::zero();
24            }
25            row2val[i_blk][i_dim + N * i_dim] = val_dia;
26        }
27    }
28    //
29    assert_eq!(bc_flag.len(), num_blk);
30    for i_blk in 0..num_blk {
31        // set row
32        for idx in row2idx[i_blk]..row2idx[i_blk + 1] {
33            for i_dim in 0..N {
34                if bc_flag[i_blk][i_dim] == 0 {
35                    continue;
36                };
37                for j_dim in 0..N {
38                    idx2val[idx][i_dim + N * j_dim] = T::zero();
39                }
40            }
41        }
42    }
43    //
44    for idx in 0..idx2col.len() {
45        let j_blk1 = idx2col[idx];
46        for j_dim in 0..N {
47            if bc_flag[j_blk1][j_dim] == 0 {
48                continue;
49            };
50            for i_dim in 0..N {
51                idx2val[idx][i_dim + N * j_dim] = T::zero();
52            }
53        }
54    }
55}
56
57pub fn set_fix_dof_to_rhs_vector<T, const N: usize>(blk2rhs: &mut [[T; N]], blk2isfix: &[[i32; N]])
58where
59    T: num_traits::Float,
60{
61    let num_vtx = blk2rhs.len();
62    for i_vtx in 0..num_vtx {
63        for i_dof in 0..N {
64            if blk2isfix[i_vtx][i_dof] == 0 {
65                continue;
66            }
67            blk2rhs[i_vtx][i_dof] = T::zero();
68        }
69    }
70}
71
72/// sparse matrix class
73/// Compressed Row Storage (CRS) data structure
74/// * `num_blk` - number of row and col blocks
75pub struct Matrix<MAT> {
76    pub num_blk: usize,
77    pub row2idx: Vec<usize>,
78    pub idx2col: Vec<usize>,
79    pub idx2val: Vec<MAT>,
80    pub row2val: Vec<MAT>,
81}
82
83impl<T, const NN: usize> Matrix<[T; NN]>
84where
85    T: num_traits::Float,
86{
87    pub fn new() -> Self {
88        Matrix {
89            num_blk: 0,
90            row2idx: vec![0],
91            idx2col: Vec::<usize>::new(),
92            idx2val: Vec::<[T; NN]>::new(),
93            row2val: Vec::<[T; NN]>::new(),
94        }
95    }
96
97    pub fn as_ref_mut(&mut self) -> MatrixRefMut<T, NN> {
98        MatrixRefMut {
99            num_blk: self.num_blk,
100            row2idx: &self.row2idx,
101            idx2col: &self.idx2col,
102            idx2val: &mut self.idx2val,
103            row2val: &mut self.row2val,
104        }
105    }
106
107    pub fn as_ref(&self) -> MatrixRef<T, NN> {
108        MatrixRef {
109            num_blk: self.num_blk,
110            row2idx: &self.row2idx,
111            idx2col: &self.idx2col,
112            idx2val: &self.idx2val,
113            row2val: &self.row2val,
114        }
115    }
116
117    pub fn from_vtx2vtx(vtx2idx: &[usize], idx2vtx: &[usize]) -> Self {
118        let num_blk = vtx2idx.len() - 1;
119        let num_idx = vtx2idx[num_blk];
120        Self {
121            num_blk,
122            row2idx: vtx2idx.to_vec(),
123            idx2col: idx2vtx.to_vec(),
124            idx2val: vec![[T::zero(); NN]; num_idx],
125            row2val: vec![[T::zero(); NN]; num_blk],
126        }
127    }
128
129    pub fn set_fixed_dof<const N: usize>(&mut self, val_dia: T, blk2isfix: &[[i32; N]]) {
130        set_fixed_bc(
131            val_dia,
132            blk2isfix,
133            &mut self.row2val,
134            &mut self.idx2val,
135            &self.row2idx,
136            &self.idx2col,
137        );
138    }
139
140    /// generalized matrix-vector multiplication
141    /// where matrix is sparse (not block) matrix
142    /// `{y_vec} <- \alpha * [a_mat] * {x_vec} + \beta * {y_vec}`
143    pub fn mult_vec<const N: usize>(
144        &self,
145        y_vec: &mut [[T; N]],
146        beta: T,
147        alpha: T,
148        x_vec: &[[T; N]],
149    ) where
150        T: num_traits::Float,
151    {
152        use del_geo_core::matn_col_major;
153        use del_geo_core::vecn::VecN;
154        assert_eq!(y_vec.len(), self.num_blk);
155        for m in y_vec.iter_mut() {
156            del_geo_core::vecn::scale_in_place(m, beta);
157        }
158        for i_blk in 0..self.num_blk {
159            for idx in self.row2idx[i_blk]..self.row2idx[i_blk + 1] {
160                assert!(idx < self.idx2col.len());
161                let j_blk = self.idx2col[idx];
162                assert!(j_blk < self.num_blk);
163                let a = matn_col_major::mult_vec(&self.idx2val[idx], &x_vec[j_blk]).scale(alpha);
164                del_geo_core::vecn::add_in_place(&mut y_vec[i_blk], &a);
165            }
166            {
167                let a = matn_col_major::mult_vec(&self.row2val[i_blk], &x_vec[i_blk]).scale(alpha);
168                del_geo_core::vecn::add_in_place(&mut y_vec[i_blk], &a);
169            }
170        }
171    }
172
173    /// set zero to all the values
174    pub fn set_zero(&mut self) {
175        assert_eq!(self.idx2val.len(), self.idx2col.len());
176        self.row2val.fill([T::zero(); NN]);
177        self.idx2val.fill([T::zero(); NN]);
178    }
179
180    pub fn merge_for_array_blk<const NNODE: usize>(
181        &mut self,
182        emat: &[[[T; NN]; NNODE]; NNODE],
183        node2vtx: &[usize; NNODE],
184        col2idx: &mut Vec<usize>,
185    ) {
186        col2idx.resize(self.num_blk, usize::MAX);
187        for i_node in 0..NNODE {
188            let i_vtx = node2vtx[i_node];
189            for idx in self.row2idx[i_vtx]..self.row2idx[i_vtx + 1] {
190                let j_vtx = self.idx2col[idx];
191                col2idx[j_vtx] = idx;
192            }
193            for j_node in 0..NNODE {
194                if i_node == j_node {
195                    del_geo_core::matn_col_major::add_in_place(
196                        &mut self.row2val[i_vtx],
197                        &emat[i_node][j_node],
198                    );
199                } else {
200                    let j_vtx = node2vtx[j_node];
201                    let idx0 = col2idx[j_vtx];
202                    assert_ne!(idx0, usize::MAX);
203                    del_geo_core::matn_col_major::add_in_place(
204                        &mut self.idx2val[idx0],
205                        &emat[i_node][j_node],
206                    );
207                }
208            }
209            for idx in self.row2idx[i_vtx]..self.row2idx[i_vtx + 1] {
210                let j_vtx = self.idx2col[idx];
211                col2idx[j_vtx] = usize::MAX;
212            }
213        }
214    }
215}
216
217/// solve linear system using the Conjugate Gradient (CG) method
218pub fn conjugate_gradient<T, const N: usize, const NN: usize>(
219    r_vec: &mut [[T; N]],
220    u_vec: &mut [[T; N]],
221    ap_vec: &mut [[T; N]],
222    p_vec: &mut [[T; N]],
223    conv_ratio_tol: T,
224    max_iteration: usize,
225    mat: MatrixRef<T, NN>,
226) -> Vec<T>
227where
228    T: num_traits::Float + std::fmt::Display + std::fmt::Debug,
229{
230    let _num_dim = r_vec.len() / mat.row2val.len();
231    //
232    let mut conv_hist = Vec::<T>::new();
233    crate::slice_of_array::set_zero(u_vec);
234    let mut sqnorm_res = crate::slice_of_array::dot(r_vec, r_vec);
235    if sqnorm_res < T::epsilon() {
236        return conv_hist;
237    }
238    let inv_sqnorm_res_ini = T::one() / sqnorm_res;
239    crate::slice_of_array::copy(p_vec, r_vec); // {p} = {r}  (set initial serch direction, copy value not reference)
240    for _iitr in 0..max_iteration {
241        // alpha = (r,r) / (p,Ap)
242        mat.mult_vec::<N>(ap_vec, T::zero(), T::one(), p_vec); // {Ap_vec} = [mat]*{p_vec}
243        let pap = crate::slice_of_array::dot(p_vec, ap_vec);
244        // assert!(pap >= T::zero(), "{pap}");
245        let alpha = sqnorm_res / pap;
246        crate::slice_of_array::add_scaled_vector(u_vec, alpha, p_vec); // {u} = +alpha*{p} + {u} (update x)
247        crate::slice_of_array::add_scaled_vector(r_vec, -alpha, ap_vec); // {r} = -alpha*{Ap} + {r}
248        let sqnorm_res_new = crate::slice_of_array::dot(r_vec, r_vec);
249        let conv_ratio = (sqnorm_res_new * inv_sqnorm_res_ini).sqrt();
250        conv_hist.push(conv_ratio);
251        if conv_ratio < conv_ratio_tol {
252            return conv_hist;
253        }
254        {
255            let beta = sqnorm_res_new / sqnorm_res; // beta = (r1,r1) / (r0,r0)
256            sqnorm_res = sqnorm_res_new;
257            crate::slice_of_array::scale_and_add_vec(p_vec, beta, r_vec); // {p} = {r} + beta*{p}
258        }
259    }
260    conv_hist
261}
262
263/// solve a real-valued linear system using the conjugate gradient method with preconditioner
264#[allow(clippy::too_many_arguments)]
265pub fn preconditioned_conjugate_gradient<T, const N: usize, const NN: usize>(
266    r_vec: &mut [[T; N]],
267    x_vec: &mut Vec<[T; N]>,
268    pr_vec: &mut Vec<[T; N]>,
269    p_vec: &mut Vec<[T; N]>,
270    conv_ratio_tol: T,
271    max_nitr: usize,
272    mat: &Matrix<[T; NN]>,
273    ilu: &crate::sparse_ilu::Preconditioner<[T; NN]>,
274) -> Vec<T>
275where
276    T: num_traits::Float + std::fmt::Debug,
277{
278    use crate::slice_of_array::{add_scaled_vector, copy, dot, scale_and_add_vec, set_zero};
279    {
280        let n = r_vec.len();
281        x_vec.resize(n, [T::zero(); N]);
282        pr_vec.resize(n, [T::zero(); N]);
283        p_vec.resize(n, [T::zero(); N]);
284    }
285    assert_eq!(r_vec.len(), mat.num_blk);
286    let mut conv_hist = Vec::<T>::new();
287
288    set_zero(x_vec);
289
290    let inv_sqnorm_res0 = {
291        let sqnorm_res0 = dot(r_vec, r_vec); // DotX(r_vec, r_vec, N);
292        conv_hist.push(sqnorm_res0.sqrt());
293        if sqnorm_res0 < T::epsilon() {
294            return conv_hist;
295        }
296        T::one() / sqnorm_res0
297    };
298
299    // {Pr} = [P]{r}
300    copy(pr_vec, r_vec); // std::vector<double> Pr_vec(r_vec, r_vec + N);
301
302    crate::sparse_ilu::solve_preconditioning_vec(pr_vec, ilu); // ilu.SolvePrecond(Pr_vec.data());
303
304    // {p} = {Pr}
305    copy(p_vec, pr_vec);
306
307    // rPr = ({r},{Pr})
308    let mut rpr = dot(r_vec, pr_vec); // DotX(r_vec, Pr_vec.data(), N);
309    for _iitr in 0..max_nitr {
310        // {Ap} = [A]{p}
311        mat.mult_vec(pr_vec, T::zero(), T::one(), p_vec);
312        {
313            // alpha = ({r},{Pr})/({p},{Ap})
314            let pap = dot(p_vec, pr_vec);
315            let alpha = rpr / pap;
316            add_scaled_vector(r_vec, -alpha, pr_vec); // {r} = -alpha*{Ap} + {r}
317            add_scaled_vector(x_vec, alpha, p_vec); // {x} = +alpha*{p} + {x}
318        }
319        {
320            // Converge Judgement
321            let sqnorm_res = dot(r_vec, r_vec); // DotX(r_vec, r_vec, N);
322            conv_hist.push(sqnorm_res.sqrt());
323            let conv_ratio = (sqnorm_res * inv_sqnorm_res0).sqrt();
324            if conv_ratio < conv_ratio_tol {
325                return conv_hist;
326            }
327        }
328        {
329            // calc beta
330            copy(pr_vec, r_vec);
331            // {Pr} = [P]{r}
332            crate::sparse_ilu::solve_preconditioning_vec(pr_vec, ilu);
333            // rPr1 = ({r},{Pr})
334            let rpr1 = dot(r_vec, pr_vec);
335            // beta = rPr1/rPr
336            let beta = rpr1 / rpr;
337            rpr = rpr1;
338            // {p} = {Pr} + beta*{p}
339            scale_and_add_vec(p_vec, beta, pr_vec);
340        }
341    }
342    {
343        // Converge Judgement
344        let sq_norm_res = dot(r_vec, r_vec); // DotX(r_vec, r_vec, N);
345        conv_hist.push(sq_norm_res.sqrt());
346    }
347    conv_hist
348}
349
350pub struct MatrixRefMut<'a, T, const NN: usize> {
351    pub num_blk: usize,
352    pub row2idx: &'a [usize],
353    pub idx2col: &'a [usize],
354    pub idx2val: &'a mut [[T; NN]],
355    pub row2val: &'a mut [[T; NN]],
356}
357
358impl<'a, T, const NN: usize> MatrixRefMut<'a, T, NN>
359where
360    T: num_traits::Float,
361{
362    pub fn set_zero(&mut self) {
363        assert_eq!(self.idx2val.len(), self.idx2col.len());
364        self.row2val.fill([T::zero(); NN]);
365        self.idx2val.fill([T::zero(); NN]);
366    }
367
368    pub fn merge_for_array_blk<const NNODE: usize>(
369        &mut self,
370        emat: &[[[T; NN]; NNODE]; NNODE],
371        node2vtx: &[usize; NNODE],
372        col2idx: &mut Vec<usize>,
373    ) {
374        col2idx.resize(self.num_blk, usize::MAX);
375        for i_node in 0..NNODE {
376            let i_vtx = node2vtx[i_node];
377            for idx in self.row2idx[i_vtx]..self.row2idx[i_vtx + 1] {
378                let j_vtx = self.idx2col[idx];
379                col2idx[j_vtx] = idx;
380            }
381            for j_node in 0..NNODE {
382                if i_node == j_node {
383                    del_geo_core::matn_col_major::add_in_place(
384                        &mut self.row2val[i_vtx],
385                        &emat[i_node][j_node],
386                    );
387                } else {
388                    let j_vtx = node2vtx[j_node];
389                    let idx0 = col2idx[j_vtx];
390                    assert_ne!(idx0, usize::MAX);
391                    del_geo_core::matn_col_major::add_in_place(
392                        &mut self.idx2val[idx0],
393                        &emat[i_node][j_node],
394                    );
395                }
396            }
397            for idx in self.row2idx[i_vtx]..self.row2idx[i_vtx + 1] {
398                let j_vtx = self.idx2col[idx];
399                col2idx[j_vtx] = usize::MAX;
400            }
401        }
402    }
403}
404
405pub struct MatrixRef<'a, T, const NN: usize> {
406    pub num_blk: usize,
407    pub row2idx: &'a [usize],
408    pub idx2col: &'a [usize],
409    pub idx2val: &'a [[T; NN]],
410    pub row2val: &'a [[T; NN]],
411}
412
413impl<'a, T, const NN: usize> MatrixRef<'a, T, NN>
414where
415    T: num_traits::Float,
416{
417    /// generalized matrix-vector multiplication
418    /// where matrix is sparse (not block) matrix
419    /// `{y_vec} <- \alpha * [a_mat] * {x_vec} + \beta * {y_vec}`
420    pub fn mult_vec<const N: usize>(
421        &self,
422        y_vec: &mut [[T; N]],
423        beta: T,
424        alpha: T,
425        x_vec: &[[T; N]],
426    ) where
427        T: num_traits::Float,
428    {
429        use del_geo_core::matn_col_major;
430        use del_geo_core::vecn::VecN;
431        assert_eq!(y_vec.len(), self.num_blk);
432        for m in y_vec.iter_mut() {
433            del_geo_core::vecn::scale_in_place(m, beta);
434        }
435        for i_blk in 0..self.num_blk {
436            for idx in self.row2idx[i_blk]..self.row2idx[i_blk + 1] {
437                assert!(idx < self.idx2col.len());
438                let j_blk = self.idx2col[idx];
439                assert!(j_blk < self.num_blk);
440                let a = matn_col_major::mult_vec(&self.idx2val[idx], &x_vec[j_blk]).scale(alpha);
441                del_geo_core::vecn::add_in_place(&mut y_vec[i_blk], &a);
442            }
443            {
444                let a = matn_col_major::mult_vec(&self.row2val[i_blk], &x_vec[i_blk]).scale(alpha);
445                del_geo_core::vecn::add_in_place(&mut y_vec[i_blk], &a);
446            }
447        }
448    }
449}
del_fem_cpu/sparse_square.rs

del_fem_cpu/
sparse_square.rs
