runmat-analysis-fea 0.5.6

Finite element assembly/solve/post scaffolding for RunMat
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use crate::{
    assembly::AssemblySummary,
    operator::{apply_k, OperatorSystem},
    solve::{
        preconditioner::SpdPreconditionerKind,
        runtime_tensor_solver::{
            prepare_runtime_tensor_linear_system, solve_prepared_linear_system_runtime_tensor,
            RuntimeTensorPreparedLinearSystem,
        },
    },
};

use super::math::dot;

pub(super) struct LinearSolveAttempt {
    pub(super) vector: Vec<f64>,
    pub(super) runtime_tensor: Option<crate::solve::linear::LinearSolveResult>,
}

pub(super) struct CgSolveOptions<'a> {
    pub(super) max_iters: usize,
    pub(super) tol: f64,
    pub(super) use_runtime_tensor: bool,
    pub(super) prepared_runtime_system: Option<&'a RuntimeTensorPreparedLinearSystem>,
    pub(super) initial_guess: Option<&'a [f64]>,
}

pub(super) fn solve_k_system_cg(
    summary: &AssemblySummary,
    system: &OperatorSystem,
    rhs: &[f64],
    options: CgSolveOptions<'_>,
) -> LinearSolveAttempt {
    if options.use_runtime_tensor {
        if let Some(prepared) = options.prepared_runtime_system {
            if let Some(result) = solve_prepared_linear_system_runtime_tensor(
                summary,
                prepared,
                rhs,
                SpdPreconditionerKind::Jacobi,
                None,
            ) {
                return LinearSolveAttempt {
                    vector: result.solution.clone(),
                    runtime_tensor: Some(result),
                };
            }
        } else if let Some(fallback_prepared) = prepare_runtime_tensor_linear_system(summary) {
            if let Some(result) = solve_prepared_linear_system_runtime_tensor(
                summary,
                &fallback_prepared,
                rhs,
                SpdPreconditionerKind::Jacobi,
                None,
            ) {
                return LinearSolveAttempt {
                    vector: result.solution.clone(),
                    runtime_tensor: Some(result),
                };
            }
        }
    }

    let mut x = match options.initial_guess {
        Some(values) if values.len() == rhs.len() => values.to_vec(),
        _ => vec![0.0; rhs.len()],
    };
    let mut r = {
        let kx = apply_k(system, &x);
        rhs.iter()
            .zip(kx.iter())
            .map(|(rhs_i, kx_i)| rhs_i - kx_i)
            .collect::<Vec<f64>>()
    };
    let mut p = r.clone();
    let mut rr_old = dot(&r, &r);
    if rr_old.sqrt() <= options.tol {
        return LinearSolveAttempt {
            vector: x,
            runtime_tensor: None,
        };
    }

    for _ in 0..options.max_iters {
        let ap = apply_k(system, &p);
        let denom = dot(&p, &ap).abs().max(1.0e-12);
        let alpha = rr_old / denom;
        for i in 0..x.len() {
            x[i] += alpha * p[i];
            r[i] -= alpha * ap[i];
        }
        let rr_new = dot(&r, &r);
        if rr_new.sqrt() <= options.tol {
            break;
        }
        let beta = rr_new / rr_old.max(1.0e-12);
        for i in 0..p.len() {
            p[i] = r[i] + beta * p[i];
        }
        rr_old = rr_new;
    }

    LinearSolveAttempt {
        vector: x,
        runtime_tensor: None,
    }
}