Optimize the Jacobi method and add documentation

Moved the stencil to the outer loop. Added documentation to the library and
readme.

.gitignore

@@ -1,2 +1,3 @@
 target
 Cargo.lock
+examples/fluid-frame*.ppm

Cargo.toml

@@ -5,10 +5,8 @@ authors = ["Ben"]
 
 [dependencies]
 nalgebra = "*"
+pbr = "*"
 
 [lib]
 name = "jacobi"
 path = "src/lib.rs"
-
-[profile.release]
-debug = true

examples/fluid.rs

@@ -1,37 +1,27 @@
-#![feature(alloc_system)]
-
 extern crate nalgebra;
 extern crate jacobi;
-extern crate alloc_system;
+extern crate pbr;
 
-use jacobi::{jacobi_iter, StencilElement};
-use nalgebra::{DVector, Vector2};
 use std::io::Write;
 use std::fs::File;
 use std::f32;
+use std::thread;
+use pbr::ProgressBar;
+
+use jacobi::{jacobi_iter, StencilElement};
+use nalgebra::{DVector, Vector2};
 
 pub fn write_ppm<W: Write>(w: &mut W,
                         c: &DVector<f32>,
-                        v: &DVector<Vector2<f32>>,
                         width: usize, height: usize) {
-    write!(w, "P3\n{} {}\n255\n", width, height).unwrap();
+    write!(w, "P2\n{} {}\n255\n", width, height).unwrap();
 
     let f = |v| -> u8 {(v * 255f32) as u8};
-    let n = |v: f32| -> u8 {(v * 128f32 + 128f32) as u8};
 
     for i in 0..height {
         for j in 0..width {
-            if cfg!(debug) {
-                write!(w, "{} {} {} ",
-                       f(c[j + 256*i]),
-                       n(v[j + 256*i].x),
-                       n(v[j + 256*i].y)).unwrap();
-            } else { 
-                write!(w, "{} {} {} ",
-                       f(c[j + 256*i]),
-                       f(c[j + 256*i]),
-                       f(c[j + 256*i])).unwrap();
-            }
+            write!(w, "{} ",
+                   f(c[j + 256*i])).unwrap();
         }
         write!(w, "\n").unwrap();
     }
@@ -42,20 +32,26 @@ fn main() {
     let size = 256*256;
     let stencil = [
         StencilElement{offset:1, value:1f32},
-        StencilElement{offset:256, value:1f32}
+        StencilElement{offset:256, value:1f32},
+        StencilElement{offset:-1, value:1f32},
+        StencilElement{offset:-256, value:1f32}
     ];
-    let indexer = |x: usize, y: usize| -> usize { x + 256*y };
 
+    let indexer = |x: usize, y: usize| -> usize { x + 256*y };
 
     let mut color_grid = DVector::<f32>::new_zeros(size);
     let mut velocity_grid = DVector::<Vector2<f32>>::new_zeros(size);
     let mut next_color_grid = DVector::<f32>::new_zeros(size);
     let mut next_velocity_grid = DVector::<Vector2<f32>>::new_zeros(size);
     let mut divergence = DVector::<f32>::new_zeros(size);
     let mut pressure_grid = DVector::<f32>::new_zeros(size);
+    let mut next_pressure_grid = DVector::<f32>::new_zeros(size);
+
+    println!("Simulating 1000 frames.");
+    let mut progress_bar = ProgressBar::new(1000);
+    for frame in 0..1000 {
+        progress_bar.inc();
 
-    // Simulate 200 timesteps.
-    for _ in 0..200 {
         velocity_grid[indexer(127,0)] = Vector2::new(0.0f32, 3.0f32);
         velocity_grid[indexer(128,0)] = Vector2::new(0.0f32, 3.0f32);
         velocity_grid[indexer(126,0)] = Vector2::new(0.0f32, 3.0f32);
@@ -133,8 +129,11 @@ fn main() {
 
         // 3. Approximate pressure jacobi iterations.
         for _ in 0..50 { 
-            jacobi_iter(-0.25f32, &stencil, &stencil,
-                        &mut pressure_grid, &divergence);
+            jacobi_iter(-0.25f32, stencil.iter(), 
+                        &pressure_grid, 
+                        &mut next_pressure_grid, 
+                        &divergence);
+            std::mem::swap(&mut pressure_grid, &mut next_pressure_grid);
         }
 
         // 4. Apply pressure.
@@ -159,8 +158,12 @@ fn main() {
                 velocity_grid[indexer(x, y)] -= Vector2::new(vx, vy);
             }
         }
+        let buffer = color_grid.clone();
+        let frame = frame;
+        thread::spawn(move || {
+            let mut file = File::create(
+                &format!("images/fluid-frame{}.ppm", frame)).unwrap();
+            let _ = write_ppm(&mut file, &buffer, 256, 256);
+        });
     }
-
-    let mut file = File::create("out.ppm").unwrap();
-    let _ = write_ppm(&mut file, &color_grid, &velocity_grid, 256, 256);
 }

readme.md

@@ -1,2 +1,19 @@
+# Jacobi solver
+
+This small library computes the [jacobi
+iteration](https://en.wikipedia.org/wiki/Jacobi_method) using a given
+[stencil](https://en.wikipedia.org/wiki/Stencil_code). 
+
+# Fluid example
+
+An example of using the jacobi method to solve the discrete pressure Poisson
+equation that arrises in fluid simulation is given. The fluid simulation example
+can be run through Cargo, and the resulting .ppm frames can be assembled with
+ffmpeg:
+
+~~~bash
 cargo run --release --example fluid
-sips -s format png out.ppm --out png; open png
+
+ffmpeg -y -r 60 -f image2 -i examples/fluid-frame%d.ppm -vcodec libx264 -pix_fmt yuv420p -crf 1 -threads 0 -bf 0 images/fluid.mp4
+~~~
+

src/lib.rs

@@ -1,5 +1,3 @@
-#![feature(test)]
-
 extern crate nalgebra;
 #[cfg(test)]
 extern crate test;
@@ -8,81 +6,78 @@ use nalgebra::{BaseFloat, Indexable};
 
 #[derive(Clone)]
 pub struct StencilElement<N: BaseFloat> {
-    pub offset: usize,
+    pub offset: isize,
     pub value: N,
 }
 
-pub fn jacobi_iter<N, V>(id_inverse: N,
-                         pos_stencil: &[StencilElement<N>],
-                         neg_stencil: &[StencilElement<N>],
-                         x: &mut V, b: &V)
-    where N: BaseFloat,
-          V: Indexable<usize, N> {
+/// Computes a [jacobi iteration](https://en.wikipedia.org/wiki/Jacobi_method)
+/// using a [stencil](https://en.wikipedia.org/wiki/Stencil_code). Reapted
+/// applications of this method solve Ax=b, where the matrix A can be
+/// represented by a stencil. `id_inverse` is the inverse of the main diagonal
+/// element A_ii. `stencil` is some iterable representing the other stencil
+/// elements of A.
+/// 
+/// # Examples
+///
+/// Using the jacobi method to approximate pressure from the divergence of a
+/// velocity field in a 256x256 grid:
+///
+/// ```
+/// let size = 256*256;
+/// let stencil = [ 
+///     StencilElement{offset:1, value:1f32}, 
+///     StencilElement{offset:256, value:1f32}, 
+///     StencilElement{offset:-1, value:1f32}, 
+///     StencilElement{offset:-256, value:1f32} 
+/// ];
+///
+/// let divergence = DVector::<f32>::new_zeros(size); 
+/// let mut pressure_grid = DVector::<f32>::new_zeros(size); 
+/// let mut refined_pressure_grid = DVector::<f32>::new_zeros(size);
+///
+/// for _ in 0..50 { 
+///     jacobi_iter(-0.25f32, stencil.iter(), 
+///         &pressure_grid, 
+///         &mut refined_pressure_grid, 
+///         &divergence);
+///     std::mem::swap(&mut pressure_grid, &mut refined_pressure_grid);
+/// }
+///
+/// let _ = pressure_grid;
+///
+/// ```
+///
+pub fn jacobi_iter<'a, Float, Vector, Stencil>(id_inverse: Float,
+                                               stencil: Stencil,
+                                               x: &Vector,
+                                               x_n: &mut Vector,
+                                               b: &Vector)
+    where Float: 'static + BaseFloat,
+          Vector: Indexable<usize, Float>,
+          Stencil: Iterator<Item = &'a StencilElement<Float>>
+{
     //
+    use std::cmp::{max, min};
     let len = x.shape();
 
     for i in 0..len {
-        let mut sigma: N = N::zero();
-
-        for s in pos_stencil.iter() {
-            let idx = i + s.offset;
-            if idx >= len { continue; }
-            sigma = sigma + s.value * x[idx];
+        unsafe {
+            x_n.unsafe_set(i, b.unsafe_at(i) * id_inverse);
         }
+    }
 
-        for s in neg_stencil.iter() {
-            if i <= s.offset { continue; }
-            let idx = i - s.offset;
-            sigma = sigma + s.value * x[idx];
+    for s in stencil {
+        let start = max(-s.offset, 0) as usize;
+        let stop = min(len as isize - s.offset, len as isize) as usize;
+        let multiplier = s.value * id_inverse;
+
+        for i in start..stop {
+            let sample_offset = (i as isize + s.offset) as usize;
+            let current = unsafe { x_n.unsafe_at(i) };
+            unsafe {
+                x_n.unsafe_set(i, current - x.unsafe_at(sample_offset) * multiplier);
+            };
         }
 
-        let x_i = (b[i] - sigma) * id_inverse;
-        x[i] = x_i;
     }
 }
-
-#[test]
-fn jacobi_test() {
-    use nalgebra::DVector;
-    let stencil = [
-        StencilElement{offset:1, value:-1f32},
-        StencilElement{offset:10, value:-1f32}
-    ];
-    let b = DVector::<f32>::new_random(100);
-    let mut x = DVector::<f32>::new_zeros(100);
-
-    for _ in 0..1 {
-        jacobi_iter(4f32, &stencil, &stencil, &mut x, &b);
-    } 
-
-    let x_1 = x.clone();
-
-    for _ in 0..10 {
-        jacobi_iter(4f32, &stencil, &stencil, &mut x, &b);
-    } 
-
-    for _ in 0..1000 {
-        jacobi_iter(4f32, &stencil, &stencil, &mut x, &b);
-    } 
-
-    let x_1011 = x.clone();
-
-    for _ in 0..1000 {
-        jacobi_iter(4f32, &stencil, &stencil, &mut x, &b);
-    } 
-
-    assert!(x_1 != x);
-    assert!(x_1011 == x); // Converged.
-}
-
-#[bench]
-fn bench_safe (bench: &mut test::Bencher) { 
-    use nalgebra::DVector;
-    let stencil = [
-        StencilElement{offset:1, value:-1f32},
-        StencilElement{offset:10, value:-1f32}
-    ];
-    let b = DVector::<f32>::new_random(100);
-    let mut x = DVector::<f32>::new_zeros(100);
-    bench.iter(|| jacobi_iter(4f32, &stencil, &stencil, &mut x, &b));
-}