1pub mod higher_order;
14pub mod petrov_galerkin;
15
16#[cfg(test)]
17mod higher_order_tests;
18
19use scirs2_core::ndarray::{Array1, Array2};
20use std::collections::HashMap;
21use std::time::Instant;
22
23use crate::pde::{
24 BoundaryCondition, BoundaryConditionType, BoundaryLocation, PDEError, PDEResult, PDESolution,
25 PDESolverInfo,
26};
27
28pub use higher_order::{
30 HigherOrderMeshGenerator, HigherOrderTriangle, ShapeFunctions, TriangularQuadrature,
31};
32
33pub use petrov_galerkin::{PetrovGalerkinSolver, PetrovGalerkinType, StabilizedFormulations};
35
36#[derive(Debug, Clone, Copy, PartialEq)]
38pub struct Point {
39 pub x: f64,
41
42 pub y: f64,
44}
45
46impl Point {
47 pub fn new(x: f64, y: f64) -> Self {
49 Point { x, y }
50 }
51
52 pub fn distance(&self, other: &Point) -> f64 {
54 ((self.x - other.x).powi(2) + (self.y - other.y).powi(2)).sqrt()
55 }
56}
57
58#[derive(Debug, Clone)]
60pub struct Triangle {
61 pub nodes: [usize; 3],
63
64 pub marker: Option<i32>,
66}
67
68impl Triangle {
69 pub fn new(nodes: [usize; 3], marker: Option<i32>) -> Self {
71 Triangle { nodes, marker }
72 }
73}
74
75#[derive(Debug, Clone)]
77pub struct TriangularMesh {
78 pub points: Vec<Point>,
80
81 pub elements: Vec<Triangle>,
83
84 pub boundary_edges: Vec<(usize, usize, Option<i32>)>,
86
87 pub boundary_nodes: HashMap<usize, BoundaryNodeInfo>,
89}
90
91#[derive(Debug, Clone)]
93pub struct BoundaryNodeInfo {
94 pub bc_type: BoundaryConditionType,
96
97 pub value: f64,
99
100 pub coefficients: Option<[f64; 3]>,
102
103 pub marker: Option<i32>,
105}
106
107impl Default for TriangularMesh {
108 fn default() -> Self {
109 Self::new()
110 }
111}
112
113impl TriangularMesh {
114 pub fn new() -> Self {
116 TriangularMesh {
117 points: Vec::new(),
118 elements: Vec::new(),
119 boundary_edges: Vec::new(),
120 boundary_nodes: HashMap::new(),
121 }
122 }
123
124 pub fn generate_rectangular(
126 x_range: (f64, f64),
127 y_range: (f64, f64),
128 nx: usize,
129 ny: usize,
130 ) -> Self {
131 let mut mesh = TriangularMesh::new();
132
133 let dx = (x_range.1 - x_range.0) / (nx as f64);
135 let dy = (y_range.1 - y_range.0) / (ny as f64);
136
137 for j in 0..=ny {
139 for i in 0..=nx {
140 let x = x_range.0 + i as f64 * dx;
141 let y = y_range.0 + j as f64 * dy;
142 mesh.points.push(Point::new(x, y));
143 }
144 }
145
146 for j in 0..ny {
148 for i in 0..nx {
149 let n00 = j * (nx + 1) + i; let n10 = j * (nx + 1) + (i + 1); let n01 = (j + 1) * (nx + 1) + i; let n11 = (j + 1) * (nx + 1) + (i + 1); mesh.elements.push(Triangle::new([n00, n10, n01], None));
158
159 mesh.elements.push(Triangle::new([n11, n01, n10], None));
161 }
162 }
163
164 for i in 0..nx {
168 let n1 = i;
169 let n2 = i + 1;
170 mesh.boundary_edges.push((n1, n2, Some(1))); }
172
173 for j in 0..ny {
175 let n1 = (j + 1) * (nx + 1) - 1;
176 let n2 = (j + 2) * (nx + 1) - 1;
177 mesh.boundary_edges.push((n1, n2, Some(2))); }
179
180 for i in 0..nx {
182 let n1 = (ny + 1) * (nx + 1) - i - 1;
183 let n2 = (ny + 1) * (nx + 1) - i - 2;
184 mesh.boundary_edges.push((n1, n2, Some(3))); }
186
187 for j in 0..ny {
189 let n1 = (ny - j) * (nx + 1);
190 let n2 = (ny - j - 1) * (nx + 1);
191 mesh.boundary_edges.push((n1, n2, Some(4))); }
193
194 mesh
195 }
196
197 pub fn set_boundary_conditions(
199 &mut self,
200 boundary_conditions: &[BoundaryCondition<f64>],
201 ) -> PDEResult<()> {
202 self.boundary_nodes.clear();
204
205 for &(n1, n2, marker) in &self.boundary_edges {
207 for bc in boundary_conditions {
209 let bc_marker = match (bc.dimension, bc.location) {
211 (1, BoundaryLocation::Lower) => Some(1), (0, BoundaryLocation::Upper) => Some(2), (1, BoundaryLocation::Upper) => Some(3), (0, BoundaryLocation::Lower) => Some(4), _ => None,
216 };
217
218 if bc_marker == marker {
220 let bc_info = BoundaryNodeInfo {
222 bc_type: bc.bc_type,
223 value: bc.value,
224 coefficients: bc.coefficients,
225 marker,
226 };
227
228 self.boundary_nodes.insert(n1, bc_info.clone());
229 self.boundary_nodes.insert(n2, bc_info);
230 }
231 }
232 }
233
234 Ok(())
235 }
236
237 pub fn triangle_area(&self, element: &Triangle) -> f64 {
239 let [i, j, k] = element.nodes;
240 let pi = &self.points[i];
241 let pj = &self.points[j];
242 let pk = &self.points[k];
243
244 0.5 * ((pj.x - pi.x) * (pk.y - pi.y) - (pk.x - pi.x) * (pj.y - pi.y)).abs()
246 }
247
248 pub fn shape_function_gradients(&self, element: &Triangle) -> PDEResult<[Point; 3]> {
250 let [i, j, k] = element.nodes;
251 let pi = &self.points[i];
252 let pj = &self.points[j];
253 let pk = &self.points[k];
254
255 let area = self.triangle_area(element);
256 if area < 1e-10 {
257 return Err(PDEError::FiniteElementError(format!(
258 "Element has nearly zero area: {area}"
259 )));
260 }
261
262 let gradients = [
264 Point::new((pj.y - pk.y) / (2.0 * area), (pk.x - pj.x) / (2.0 * area)),
265 Point::new((pk.y - pi.y) / (2.0 * area), (pi.x - pk.x) / (2.0 * area)),
266 Point::new((pi.y - pj.y) / (2.0 * area), (pj.x - pi.x) / (2.0 * area)),
267 ];
268
269 Ok(gradients)
270 }
271}
272
273#[derive(Debug, Clone, Copy, PartialEq)]
275pub enum ElementType {
276 Linear,
278
279 Quadratic,
281
282 Cubic,
284}
285
286#[derive(Debug, Clone)]
288pub struct FEMOptions {
289 pub element_type: ElementType,
291
292 pub quadrature_order: usize,
294
295 pub max_iterations: usize,
297
298 pub tolerance: f64,
300
301 pub save_convergence_history: bool,
303
304 pub verbose: bool,
306}
307
308impl Default for FEMOptions {
309 fn default() -> Self {
310 FEMOptions {
311 element_type: ElementType::Linear,
312 quadrature_order: 3, max_iterations: 1000,
314 tolerance: 1e-6,
315 save_convergence_history: false,
316 verbose: false,
317 }
318 }
319}
320
321#[derive(Debug, Clone)]
323pub struct FEMResult {
324 pub u: Array1<f64>,
326
327 pub mesh: TriangularMesh,
329
330 pub residual_norm: f64,
332
333 pub num_iterations: usize,
335
336 pub computation_time: f64,
338
339 pub convergence_history: Option<Vec<f64>>,
341}
342
343pub struct FEMPoissonSolver {
345 mesh: TriangularMesh,
347
348 higher_order_elements: Option<Vec<HigherOrderTriangle>>,
350
351 higher_order_points: Option<Vec<Point>>,
353
354 source_term: Box<dyn Fn(f64, f64) -> f64 + Send + Sync>,
356
357 boundary_conditions: Vec<BoundaryCondition<f64>>,
359
360 options: FEMOptions,
362}
363
364impl FEMPoissonSolver {
365 pub fn new(
367 mesh: TriangularMesh,
368 source_term: impl Fn(f64, f64) -> f64 + Send + Sync + 'static,
369 boundary_conditions: Vec<BoundaryCondition<f64>>,
370 options: Option<FEMOptions>,
371 ) -> PDEResult<Self> {
372 if boundary_conditions.is_empty() {
374 return Err(PDEError::BoundaryConditions(
375 "At least one boundary condition is required".to_string(),
376 ));
377 }
378
379 let opts = options.unwrap_or_default();
380
381 let (higher_order_elements, higher_order_points) = match opts.element_type {
383 ElementType::Linear => (None, None),
384 ElementType::Quadratic => {
385 let (points, elements) = HigherOrderMeshGenerator::linear_to_quadratic(&mesh)?;
386 (Some(elements), Some(points))
387 }
388 ElementType::Cubic => {
389 let (points, elements) = HigherOrderMeshGenerator::linear_to_cubic(&mesh)?;
390 (Some(elements), Some(points))
391 }
392 };
393
394 Ok(FEMPoissonSolver {
395 mesh,
396 higher_order_elements,
397 higher_order_points,
398 source_term: Box::new(source_term),
399 boundary_conditions,
400 options: opts,
401 })
402 }
403
404 pub fn solve(&mut self) -> PDEResult<FEMResult> {
406 let start_time = Instant::now();
407
408 self.mesh
410 .set_boundary_conditions(&self.boundary_conditions)?;
411
412 let _n = if let Some(ref higher_order_points) = self.higher_order_points {
414 higher_order_points.len()
415 } else {
416 self.mesh.points.len()
417 };
418
419 let (mut a, mut b) = self.assemble_system()?;
421
422 self.apply_dirichlet_boundary_conditions(&mut a, &mut b)?;
424
425 let u = FEMPoissonSolver::solve_linear_system(&a, &b)?;
427
428 let residual_norm = FEMPoissonSolver::compute_residual(&a, &b, &u);
430
431 let computation_time = start_time.elapsed().as_secs_f64();
432
433 Ok(FEMResult {
434 u,
435 mesh: self.mesh.clone(),
436 residual_norm,
437 num_iterations: 1, computation_time,
439 convergence_history: None,
440 })
441 }
442
443 fn assemble_system(&self) -> PDEResult<(Array2<f64>, Array1<f64>)> {
445 let n = if let Some(ref higher_order_points) = self.higher_order_points {
446 higher_order_points.len()
447 } else {
448 self.mesh.points.len()
449 };
450
451 let mut a = Array2::zeros((n, n));
453 let mut b = Array1::zeros(n);
454
455 match self.options.element_type {
456 ElementType::Linear => {
457 for element in &self.mesh.elements {
459 let (a_e, b_e) = self.element_matrices_linear(element)?;
460
461 let [i, j, k] = element.nodes;
463
464 a[[i, i]] += a_e[0][0];
466 a[[j, j]] += a_e[1][1];
467 a[[k, k]] += a_e[2][2];
468
469 a[[i, j]] += a_e[0][1];
471 a[[i, k]] += a_e[0][2];
472 a[[j, i]] += a_e[1][0];
473 a[[j, k]] += a_e[1][2];
474 a[[k, i]] += a_e[2][0];
475 a[[k, j]] += a_e[2][1];
476
477 b[i] += b_e[0];
479 b[j] += b_e[1];
480 b[k] += b_e[2];
481 }
482 }
483 _ => {
484 if let Some(ref higher_order_elements) = self.higher_order_elements {
486 for element in higher_order_elements {
487 let (a_e, b_e) = self.element_matrices_higher_order(element)?;
488
489 for (i, node_i) in element.nodes.iter().enumerate() {
491 b[*node_i] += b_e[i];
492 for (j, node_j) in element.nodes.iter().enumerate() {
493 a[[*node_i, *node_j]] += a_e[[i, j]];
494 }
495 }
496 }
497 }
498 }
499 }
500
501 Ok((a, b))
502 }
503
504 fn element_matrices_linear(&self, element: &Triangle) -> PDEResult<([[f64; 3]; 3], [f64; 3])> {
506 let [i, j, k] = element.nodes;
508 let pi = &self.mesh.points[i];
509 let pj = &self.mesh.points[j];
510 let pk = &self.mesh.points[k];
511
512 let area = self.mesh.triangle_area(element);
514
515 let gradients = self.mesh.shape_function_gradients(element)?;
517
518 let mut a_e = [[0.0; 3]; 3];
520
521 for m in 0..3 {
522 for n in 0..3 {
523 a_e[m][n] =
525 (gradients[m].x * gradients[n].x + gradients[m].y * gradients[n].y) * area;
526 }
527 }
528
529 let mut b_e = [0.0; 3];
531
532 let centroid_x = (pi.x + pj.x + pk.x) / 3.0;
534 let centroid_y = (pi.y + pj.y + pk.y) / 3.0;
535 let f_centroid = (self.source_term)(centroid_x, centroid_y);
536
537 b_e.iter_mut().for_each(|value| {
539 *value = f_centroid * (area / 3.0);
540 });
541
542 Ok((a_e, b_e))
543 }
544
545 fn element_matrices_higher_order(
547 &self,
548 element: &HigherOrderTriangle,
549 ) -> PDEResult<(Array2<f64>, Array1<f64>)> {
550 let num_nodes = element.nodes.len();
551 let mut a_e = Array2::zeros((num_nodes, num_nodes));
552 let mut b_e = Array1::zeros(num_nodes);
553
554 let points = if let Some(ref ho_points) = self.higher_order_points {
556 ho_points
557 } else {
558 return Err(PDEError::FiniteElementError(
559 "Higher-order points not available".to_string(),
560 ));
561 };
562
563 let corner_nodes = element.corner_nodes();
565 let p1 = &points[corner_nodes[0]];
566 let p2 = &points[corner_nodes[1]];
567 let p3 = &points[corner_nodes[2]];
568
569 let jacobian = Array2::from_shape_vec(
571 (2, 2),
572 vec![p2.x - p1.x, p3.x - p1.x, p2.y - p1.y, p3.y - p1.y],
573 )
574 .expect("Operation failed");
575
576 let det_j = jacobian[[0, 0]] * jacobian[[1, 1]] - jacobian[[0, 1]] * jacobian[[1, 0]];
577 if det_j.abs() < 1e-12 {
578 return Err(PDEError::FiniteElementError(
579 "Degenerate element with zero Jacobian determinant".to_string(),
580 ));
581 }
582
583 let inv_j = Array2::from_shape_vec(
585 (2, 2),
586 vec![
587 jacobian[[1, 1]] / det_j,
588 -jacobian[[0, 1]] / det_j,
589 -jacobian[[1, 0]] / det_j,
590 jacobian[[0, 0]] / det_j,
591 ],
592 )
593 .expect("Operation failed");
594
595 let (xi_coords, eta_coords, weights) =
597 TriangularQuadrature::get_rule(self.options.quadrature_order)?;
598
599 for q in 0..xi_coords.len() {
601 let xi = xi_coords[q];
602 let eta = eta_coords[q];
603 let weight = weights[q];
604
605 let shape_funcs = ShapeFunctions::evaluate(element.element_type, xi, eta)?;
607 let (d_n_dxi, d_n_deta) =
608 ShapeFunctions::evaluate_derivatives(element.element_type, xi, eta)?;
609
610 let mut d_n_dx = Array1::zeros(num_nodes);
612 let mut d_n_dy = Array1::zeros(num_nodes);
613
614 for i in 0..num_nodes {
615 d_n_dx[i] = inv_j[[0, 0]] * d_n_dxi[i] + inv_j[[0, 1]] * d_n_deta[i];
616 d_n_dy[i] = inv_j[[1, 0]] * d_n_dxi[i] + inv_j[[1, 1]] * d_n_deta[i];
617 }
618
619 let mut x_phys = 0.0;
621 let mut y_phys = 0.0;
622 for i in 0..num_nodes {
623 x_phys += shape_funcs[i] * points[element.nodes[i]].x;
624 y_phys += shape_funcs[i] * points[element.nodes[i]].y;
625 }
626
627 let f_val = (self.source_term)(x_phys, y_phys);
629
630 for i in 0..num_nodes {
632 b_e[i] += f_val * shape_funcs[i] * weight * det_j.abs();
634
635 for j in 0..num_nodes {
636 a_e[[i, j]] +=
638 (d_n_dx[i] * d_n_dx[j] + d_n_dy[i] * d_n_dy[j]) * weight * det_j.abs();
639 }
640 }
641 }
642
643 Ok((a_e, b_e))
644 }
645
646 fn apply_dirichlet_boundary_conditions(
648 &self,
649 a: &mut Array2<f64>,
650 b: &mut Array1<f64>,
651 ) -> PDEResult<()> {
652 let n = self.mesh.points.len();
653
654 for (&node_idx, bc_info) in &self.mesh.boundary_nodes {
656 if bc_info.bc_type == BoundaryConditionType::Dirichlet {
657 for j in 0..n {
659 a[[node_idx, j]] = 0.0;
660 }
661 a[[node_idx, node_idx]] = 1.0;
662
663 b[node_idx] = bc_info.value;
665 } else if bc_info.bc_type == BoundaryConditionType::Neumann {
666 let boundary_edges: Vec<_> = self
672 .mesh
673 .boundary_edges
674 .iter()
675 .filter(|&&(n1, n2, _)| n1 == node_idx || n2 == node_idx)
676 .collect();
677
678 for &(n1, n2, _) in &boundary_edges {
680 let other_node = if *n1 == node_idx { *n2 } else { *n1 };
681
682 let p1 = &self.mesh.points[node_idx];
684 let p2 = &self.mesh.points[other_node];
685
686 let edge_length = p1.distance(p2);
688
689 b[node_idx] += bc_info.value * (edge_length / 2.0);
692 }
693 } else if bc_info.bc_type == BoundaryConditionType::Robin {
694 if let Some([a_coef, b_coef, c_coef]) = bc_info.coefficients {
696 let boundary_edges: Vec<_> = self
698 .mesh
699 .boundary_edges
700 .iter()
701 .filter(|&&(n1, n2, _)| n1 == node_idx || n2 == node_idx)
702 .collect();
703
704 for &(n1, n2, _) in &boundary_edges {
705 let other_node = if *n1 == node_idx { *n2 } else { *n1 };
706
707 let p1 = &self.mesh.points[node_idx];
709 let p2 = &self.mesh.points[other_node];
710
711 let edge_length = p1.distance(p2);
713
714 a[[node_idx, node_idx]] += a_coef * edge_length / 2.0;
718
719 b[node_idx] += c_coef * edge_length / 2.0;
721 }
722 }
723 }
724 }
725
726 Ok(())
727 }
728
729 fn solve_linear_system(a: &Array2<f64>, b: &Array1<f64>) -> PDEResult<Array1<f64>> {
731 let n = b.len();
732
733 let mut a_copy = a.clone();
738 let mut b_copy = b.clone();
739
740 for i in 0..n {
742 let mut max_val = a_copy[[i, i]].abs();
744 let mut max_row = i;
745
746 for k in i + 1..n {
747 if a_copy[[k, i]].abs() > max_val {
748 max_val = a_copy[[k, i]].abs();
749 max_row = k;
750 }
751 }
752
753 if max_val < 1e-10 {
755 return Err(PDEError::Other(
756 "Matrix is singular or nearly singular".to_string(),
757 ));
758 }
759
760 if max_row != i {
762 for j in i..n {
763 let temp = a_copy[[i, j]];
764 a_copy[[i, j]] = a_copy[[max_row, j]];
765 a_copy[[max_row, j]] = temp;
766 }
767
768 let temp = b_copy[i];
769 b_copy[i] = b_copy[max_row];
770 b_copy[max_row] = temp;
771 }
772
773 for k in i + 1..n {
775 let factor = a_copy[[k, i]] / a_copy[[i, i]];
776
777 for j in i..n {
778 a_copy[[k, j]] -= factor * a_copy[[i, j]];
779 }
780
781 b_copy[k] -= factor * b_copy[i];
782 }
783 }
784
785 let mut x = Array1::zeros(n);
787 for i in (0..n).rev() {
788 let mut sum = 0.0;
789 for j in i + 1..n {
790 sum += a_copy[[i, j]] * x[j];
791 }
792
793 x[i] = (b_copy[i] - sum) / a_copy[[i, i]];
794 }
795
796 Ok(x)
797 }
798
799 fn compute_residual(a: &Array2<f64>, b: &Array1<f64>, x: &Array1<f64>) -> f64 {
801 let n = b.len();
802 let mut residual = 0.0;
803
804 for i in 0..n {
805 let mut row_sum = 0.0;
806 for j in 0..n {
807 row_sum += a[[i, j]] * x[j];
808 }
809
810 let diff = row_sum - b[i];
811 residual += diff * diff;
812 }
813
814 residual.sqrt()
815 }
816}
817
818impl From<FEMResult> for PDESolution<f64> {
820 fn from(result: FEMResult) -> Self {
821 let mut grids = Vec::new();
822 let n = result.mesh.points.len();
823
824 let mut x_coords = Array1::zeros(n);
826 let mut y_coords = Array1::zeros(n);
827
828 for (i, point) in result.mesh.points.iter().enumerate() {
829 x_coords[i] = point.x;
830 y_coords[i] = point.y;
831 }
832
833 grids.push(x_coords);
834 grids.push(y_coords);
835
836 let mut values = Vec::new();
838 let u_reshaped = result
839 .u
840 .into_shape_with_order((n, 1))
841 .expect("Operation failed");
842 values.push(u_reshaped);
843
844 let info = PDESolverInfo {
846 num_iterations: result.num_iterations,
847 computation_time: result.computation_time,
848 residual_norm: Some(result.residual_norm),
849 convergence_history: result.convergence_history,
850 method: "Finite Element Method".to_string(),
851 };
852
853 PDESolution {
854 grids,
855 values,
856 error_estimate: None,
857 info,
858 }
859 }
860}
861
862impl PDEError {
864 pub fn finite_element_error(msg: String) -> Self {
866 PDEError::Other(format!("Finite element error: {msg}"))
867 }
868}