1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
4pub struct CsrMatrix {
5 pub row_offsets: Vec<usize>,
6 pub column_indices: Vec<usize>,
7 pub values: Vec<f64>,
8}
9
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct OperatorSystem {
12 pub dof_count: usize,
13 pub constrained: Vec<bool>,
14 #[serde(default)]
15 pub stiffness_dense: Option<Vec<f64>>,
16 #[serde(default)]
17 pub stiffness_csr: Option<CsrMatrix>,
18 pub stiffness_diag: Vec<f64>,
19 pub stiffness_upper: Vec<f64>,
20 pub mass_diag: Vec<f64>,
21 pub damping_diag: Vec<f64>,
22 pub rhs: Vec<f64>,
23}
24
25pub fn apply_k(system: &OperatorSystem, x: &[f64]) -> Vec<f64> {
26 if let Some(dense) = dense_stiffness(system) {
27 let mut y = vec![0.0; x.len()];
28 for i in 0..x.len() {
29 if system.constrained[i] {
30 y[i] = x[i];
31 continue;
32 }
33 y[i] = (0..x.len())
34 .filter(|j| !system.constrained[*j])
35 .map(|j| dense[i * system.dof_count + j] * x[j])
36 .sum();
37 }
38 return y;
39 }
40 if let Some(csr) = csr_stiffness(system) {
41 return apply_csr_with_constraints(csr, &system.constrained, x);
42 }
43
44 let mut y = vec![0.0; x.len()];
45 for i in 0..x.len() {
46 if system.constrained[i] {
47 y[i] = x[i];
48 continue;
49 }
50
51 let mut value = system.stiffness_diag[i] * x[i];
52
53 if i > 0 && !system.constrained[i - 1] {
54 value -= system.stiffness_upper[i - 1] * x[i - 1];
55 }
56 if i + 1 < x.len() && !system.constrained[i + 1] {
57 value -= system.stiffness_upper[i] * x[i + 1];
58 }
59
60 y[i] = value;
61 }
62 y
63}
64
65pub fn apply_k_unconstrained(system: &OperatorSystem, x: &[f64]) -> Vec<f64> {
66 if let Some(dense) = dense_stiffness(system) {
67 let mut y = vec![0.0; x.len()];
68 for i in 0..x.len() {
69 y[i] = (0..x.len())
70 .map(|j| dense[i * system.dof_count + j] * x[j])
71 .sum();
72 }
73 return y;
74 }
75 if let Some(csr) = csr_stiffness(system) {
76 return apply_csr_unconstrained(csr, x);
77 }
78
79 let mut y = vec![0.0; x.len()];
80 for i in 0..x.len() {
81 let mut value = system.stiffness_diag[i] * x[i];
82
83 if i > 0 {
84 value -= system.stiffness_upper[i - 1] * x[i - 1];
85 }
86 if i + 1 < x.len() {
87 value -= system.stiffness_upper[i] * x[i + 1];
88 }
89
90 y[i] = value;
91 }
92 y
93}
94
95pub fn dense_stiffness(system: &OperatorSystem) -> Option<&[f64]> {
96 system
97 .stiffness_dense
98 .as_deref()
99 .filter(|dense| dense.len() == system.dof_count.saturating_mul(system.dof_count))
100}
101
102pub fn csr_stiffness(system: &OperatorSystem) -> Option<&CsrMatrix> {
103 let csr = system.stiffness_csr.as_ref()?;
104 if csr.row_offsets.len() != system.dof_count.saturating_add(1) {
105 return None;
106 }
107 if csr.column_indices.len() != csr.values.len() {
108 return None;
109 }
110 if csr.row_offsets.last().copied()? != csr.values.len() {
111 return None;
112 }
113 Some(csr)
114}
115
116pub fn apply_m(system: &OperatorSystem, x: &[f64]) -> Vec<f64> {
117 apply_diag_with_constraints(&system.mass_diag, &system.constrained, x)
118}
119
120pub fn apply_c(system: &OperatorSystem, x: &[f64]) -> Vec<f64> {
121 apply_diag_with_constraints(&system.damping_diag, &system.constrained, x)
122}
123
124fn apply_diag_with_constraints(diag: &[f64], constrained: &[bool], x: &[f64]) -> Vec<f64> {
125 diag.iter()
126 .zip(constrained.iter())
127 .zip(x.iter())
128 .map(
129 |((&d, &is_constrained), &value)| {
130 if is_constrained {
131 value
132 } else {
133 d * value
134 }
135 },
136 )
137 .collect()
138}
139
140fn apply_csr_with_constraints(csr: &CsrMatrix, constrained: &[bool], x: &[f64]) -> Vec<f64> {
141 let mut y = vec![0.0; x.len()];
142 for row in 0..x.len() {
143 if constrained[row] {
144 y[row] = x[row];
145 continue;
146 }
147 let start = csr.row_offsets[row];
148 let end = csr.row_offsets[row + 1];
149 y[row] = (start..end)
150 .filter_map(|entry| {
151 let col = csr.column_indices[entry];
152 (!constrained[col]).then_some(csr.values[entry] * x[col])
153 })
154 .sum();
155 }
156 y
157}
158
159fn apply_csr_unconstrained(csr: &CsrMatrix, x: &[f64]) -> Vec<f64> {
160 let mut y = vec![0.0; x.len()];
161 for (row, value) in y.iter_mut().enumerate().take(x.len()) {
162 let start = csr.row_offsets[row];
163 let end = csr.row_offsets[row + 1];
164 *value = (start..end)
165 .map(|entry| csr.values[entry] * x[csr.column_indices[entry]])
166 .sum();
167 }
168 y
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174
175 #[test]
176 fn operator_apply_respects_constraint_identity() {
177 let system = OperatorSystem {
178 dof_count: 3,
179 constrained: vec![true, false, false],
180 stiffness_dense: None,
181 stiffness_csr: None,
182 stiffness_diag: vec![100.0, 10.0, 5.0],
183 stiffness_upper: vec![1.0, 0.5],
184 mass_diag: vec![1.0, 2.0, 3.0],
185 damping_diag: vec![0.1, 0.2, 0.3],
186 rhs: vec![0.0, -100.0, 0.0],
187 };
188 let x = vec![4.0, 5.0, 6.0];
189
190 assert_eq!(apply_k(&system, &x), vec![4.0, 47.0, 27.5]);
191 assert_eq!(apply_k_unconstrained(&system, &x), vec![395.0, 43.0, 27.5]);
192 assert_eq!(apply_m(&system, &x), vec![4.0, 10.0, 18.0]);
193 let c = apply_c(&system, &x);
194 assert!((c[0] - 4.0).abs() <= 1.0e-12);
195 assert!((c[1] - 1.0).abs() <= 1.0e-12);
196 assert!((c[2] - 1.8).abs() <= 1.0e-12);
197 }
198
199 #[test]
200 fn operator_apply_uses_dense_stiffness_when_present() {
201 let system = OperatorSystem {
202 dof_count: 3,
203 constrained: vec![true, false, false],
204 stiffness_dense: Some(vec![100.0, 2.0, 3.0, 2.0, 10.0, 4.0, 3.0, 4.0, 5.0]),
205 stiffness_csr: None,
206 stiffness_diag: vec![100.0, 10.0, 5.0],
207 stiffness_upper: vec![1.0, 0.5],
208 mass_diag: vec![1.0, 2.0, 3.0],
209 damping_diag: vec![0.1, 0.2, 0.3],
210 rhs: vec![0.0, -100.0, 0.0],
211 };
212 let x = vec![4.0, 5.0, 6.0];
213
214 assert_eq!(apply_k(&system, &x), vec![4.0, 74.0, 50.0]);
215 assert_eq!(apply_k_unconstrained(&system, &x), vec![428.0, 82.0, 62.0]);
216 }
217
218 #[test]
219 fn operator_apply_uses_csr_stiffness_when_present() {
220 let system = OperatorSystem {
221 dof_count: 3,
222 constrained: vec![true, false, false],
223 stiffness_dense: None,
224 stiffness_csr: Some(CsrMatrix {
225 row_offsets: vec![0, 2, 5, 7],
226 column_indices: vec![0, 1, 0, 1, 2, 1, 2],
227 values: vec![4.0, -1.0, -1.0, 3.0, -2.0, -2.0, 5.0],
228 }),
229 stiffness_diag: vec![4.0, 3.0, 5.0],
230 stiffness_upper: vec![0.0, 0.0],
231 mass_diag: vec![1.0; 3],
232 damping_diag: vec![0.0; 3],
233 rhs: vec![0.0; 3],
234 };
235 let x = vec![10.0, 2.0, 4.0];
236
237 assert_eq!(apply_k(&system, &x), vec![10.0, -2.0, 16.0]);
238 assert_eq!(apply_k_unconstrained(&system, &x), vec![38.0, -12.0, 16.0]);
239 }
240}