Add matrix approach to heat equation

heat-equation/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2018"
 
 
 [dependencies]
-sbp = { path = "../sbp" }
+sbp = { path = "../sbp", features = ["sparse"] }
 ndarray = "0.13.1"
 plotters = { version = "0.3.0", default-features = false, features = ["bitmap_gif", "bitmap_backend", "line_series"] }
+sprs = "0.7.0"

heat-equation/src/main.rs

@@ -2,7 +2,7 @@ use ndarray::{Array1, ArrayView1};
 use plotters::prelude::*;
 use sbp::{
     integrate::{integrate, Rk4},
-    operators::{SbpOperator1d2, SBP4},
+    operators::{SbpOperator1d, SbpOperator1d2, SBP4},
     Float,
 };
 
@@ -13,6 +13,9 @@ fn main() {
     dual_dirichlet(v0.view(), 1.0, 1.0);
     // Neumann boundary is introducing energy into the system
     neumann_dirichlet(v0.view(), -0.2, 1.0);
+
+    // The sparse formulation
+    dual_dirichlet_sparse(v0.view(), 1.0, 1.0);
 }
 
 fn dual_dirichlet(v: ArrayView1<Float>, v0: Float, vn: Float) {
@@ -134,3 +137,90 @@ fn neumann_dirichlet(v: ArrayView1<Float>, v0: Float, vn: Float) {
         std::mem::swap(&mut v1, &mut v2);
     }
 }
+
+fn dual_dirichlet_sparse(v: ArrayView1<Float>, v0: Float, vn: Float) {
+    let drawing_area = BitMapBackend::gif("dual_dirichlet_sparse.gif", (300, 300), 100)
+        .unwrap()
+        .into_drawing_area();
+    let mut chart = ChartBuilder::on(&drawing_area)
+        .x_label_area_size(40)
+        .y_label_area_size(40)
+        .build_cartesian_2d(0.0..1.01, -1.0..2.0)
+        .unwrap();
+
+    let op = SBP4;
+    let nx = v.len();
+    let system = op.diff2_matrix(nx);
+    let e0 = {
+        let mut e0 = sprs::CsMat::zero((nx, 1));
+        e0.insert(0, 0, 1.0);
+        e0
+    };
+    let en = {
+        let mut en = sprs::CsMat::zero((nx, 1));
+        en.insert(nx - 1, 0, 1.0);
+        en
+    };
+
+    let mut hi = op.h_matrix(nx);
+    hi.map_inplace(|v| 1.0 / v);
+
+    let tau = (1.0, -1.0);
+    let sat0 = {
+        let d1 = op.d1_vec(nx, true);
+        let mut mat = &hi * &d1.transpose_view();
+        mat.map_inplace(|v| v * tau.0);
+        (&mat * &e0.transpose_view(), mat)
+    };
+    let satn = {
+        let dn = op.d1_vec(nx, false);
+        let mut mat = &hi * &dn.transpose_view();
+        mat.map_inplace(|v| v * tau.1);
+        (&mat * &en.transpose_view(), mat)
+    };
+
+    let system = &system + &(&sat0.0 + &satn.0);
+    // Stack the two matrices to allow easy book-keeping
+    // of boundary conditions
+    let mut bc = sprs::hstack(&[sat0.1.view(), satn.1.view()]).to_csr();
+    bc.map_inplace(|v| -v);
+    let system = &system;
+    let bc = &bc;
+
+    let dt = 0.2 / nx.pow(2) as Float / 3.0;
+    let x = Array1::from_shape_fn((nx,), |i| i as Float / (nx - 1) as Float);
+
+    let mut k = [v.to_owned(), v.to_owned(), v.to_owned(), v.to_owned()];
+    let rhs = move |fut: &mut Array1<Float>, prev: &Array1<Float>, _t: Float| {
+        fut.fill(0.0);
+        let prev = prev.as_slice().unwrap();
+        let fut = fut.as_slice_mut().unwrap();
+
+        sprs::prod::mul_acc_mat_vec_csr(system.view(), prev, fut);
+
+        sprs::prod::mul_acc_mat_vec_csr(bc.view(), &[v0, vn][..], fut);
+    };
+
+    let mut v1 = v.to_owned();
+    let mut v2 = v.to_owned();
+    for i in 0..90 {
+        if i % 3 == 0 {
+            drawing_area.fill(&WHITE).unwrap();
+            chart
+                .configure_mesh()
+                .x_desc("x")
+                .y_desc("y")
+                .draw()
+                .unwrap();
+            chart
+                .draw_series(LineSeries::new(
+                    x.iter().zip(v1.iter()).map(|(&x, &y)| (x, y)),
+                    &BLACK,
+                ))
+                .unwrap();
+            drawing_area.present().unwrap();
+        }
+        integrate::<Rk4, _, _, _>(rhs, &v1, &mut v2, &mut 0.0, dt, &mut k);
+        std::mem::swap(&mut v1, &mut v2);
+    }
+}

sbp/src/operators.rs

@@ -19,8 +19,13 @@ pub trait SbpOperator1d: Send + Sync {
 
 pub trait SbpOperator1d2: SbpOperator1d {
     fn diff2(&self, prev: ArrayView1<Float>, fut: ArrayViewMut1<Float>);
-    /// Lacks a scaling of 1/h^2
+    /// Result of H^-1 * d1, without the 1/h scaling
     fn d1(&self) -> &[Float];
+
+    #[cfg(feature = "sparse")]
+    fn diff2_matrix(&self, n: usize) -> sprs::CsMat<Float>;
+    #[cfg(feature = "sparse")]
+    fn d1_vec(&self, n: usize, front: bool) -> sprs::CsMat<Float>;
 }
 
 pub trait SbpOperator2d: Send + Sync {

sbp/src/operators/traditional4.rs

@@ -109,6 +109,41 @@ impl super::SbpOperator1d2 for SBP4 {
     fn d1(&self) -> &[Float] {
         &[-88.0 / 17.0, 144.0 / 17.0, -72.0 / 17.0, 16.0 / 17.0]
     }
+    #[cfg(feature = "sparse")]
+    fn diff2_matrix(&self, n: usize) -> sprs::CsMat<Float> {
+        let mut m = super::sparse_from_block(
+            Self::D2BLOCK,
+            Self::D2DIAG,
+            super::Symmetry::Symmetric,
+            super::OperatorType::Normal,
+            n,
+        );
+        let hi = (n - 1) as Float;
+        m.map_inplace(|v| v * hi);
+        m
+    }
+    #[cfg(feature = "sparse")]
+    fn d1_vec(&self, n: usize, front: bool) -> sprs::CsMat<Float> {
+        let d1 = &[-11.0 / 6.0, 3.0, -3.0 / 2.0, 1.0 / 3.0];
+        assert!(n >= d1.len());
+        let mut d1 = d1.to_vec();
+        let hi = (n - 1) as Float;
+        d1.iter_mut().for_each(|v| *v *= hi);
+
+        if front {
+            sprs::CsMat::new((1, n), vec![0, d1.len()], (0..d1.len()).collect(), d1)
+        } else {
+            // d_x => -d_x when reversed
+            d1.iter_mut().for_each(|v| *v *= -1.0);
+            // Indices are now in reverse order
+            sprs::CsMat::new(
+                (1, n),
+                vec![0, d1.len()],
+                (0..n).rev().take(d1.len()).collect(),
+                d1,
+            )
+        }
+    }
 }
 
 #[test]