Learning Rust by Simulating Heat Diffusion: Unabridged

heat.rs

#[derive(Debug)]
pub struct Heat1D(pub Vec<f64>);

impl Heat1D {
    // r = k / h^2, where k is the time step and h is the space step
    pub fn step(&mut self, r: f64) -> &mut Self {
        let mut iter = self.0.iter_mut().peekable();
        // skip first
        if let Some(&mut mut prev) = iter.next() {
            // ensure
            let mut delta;
            while let (Some(cur), Some(next)) = (iter.next(), iter.peek()) {
                delta = (prev - 2. * *cur + **next) * r;
                prev = *cur;
                *cur += delta;
            }
        }
        self
    }

    pub fn copy_edges(&mut self) -> &mut Self {
        // set boundaries to neighbors to simulate no BC
        let len = self.0.len();
        self.0[0] = self.0[1];
        self.0[len - 1] = self.0[len - 2];
        self
    }
}

#[derive(Debug)]
pub struct Heat2D {
    pub values: Vec<f64>,
    pub dimensions: (usize, usize),
    // auxillary space when stepping
    auxillary: Vec<f64>,
}

impl Heat2D {
    pub fn new(dimensions: (usize, usize)) -> Self {
        Heat2D {
            values: vec![0.; dimensions.0 * dimensions.1],
            dimensions,
            auxillary: vec![0.; 2 * dimensions.0],
        }
    }

    // r = k / h^2, where k is the time step and h is the space step
    pub fn step(&mut self, r: f64) -> &mut Self {
        // prepare auxillary memory
        let (prev_row, aux) = self.auxillary.split_at_mut(self.dimensions.0);

        let mut iter_row = self.values.chunks_exact_mut(self.dimensions.0).peekable();
        // iterate by threes
        if let Some(first_row) = iter_row.next() {
            // previous row to use in the computation
            prev_row.copy_from_slice(first_row);
            while let (Some(cur_row), Some(next_row)) = (iter_row.next(), iter_row.peek()) {
                // save the current row before mutation
                aux.copy_from_slice(cur_row);
                // iterate other dimension now
                let mut iter = cur_row.iter_mut().enumerate().peekable();
                // skip first
                if let Some((_, &mut mut prev)) = iter.next() {
                    // ensure
                    let mut delta;
                    while let (Some((i, ref mut cur)), Some((_, next))) = (iter.next(), iter.peek())
                    {
                        // 2d laplacian
                        delta = (prev + **next + prev_row[i] + next_row[i] - 4. * **cur) * r;
                        // change references
                        prev = **cur;
                        **cur += delta;
                    }
                }
                core::mem::swap(&mut prev_row, &mut aux);
            }
        }
        self
    }

    pub fn copy_edges(&mut self) -> &mut Self {
        // two opposite consecutive sides
        let m = self.dimensions.0;
        let n = self.values.len();
        for i in 1..(m - 1) {
            self.values[i] = self.values[m + i];
            self.values[n - m + i] = self.values[n - 2 * m + i];
        }
        // two non local sides, include corners for the plot
        for i in 0..self.dimensions.1 {
            self.values[m * i] = self.values[m * i + 1];
            self.values[m * (i + 1) - 1] = self.values[m * (i + 1) - 2];
        }
        self
    }
}

main.rs

mod heat;

use heat::*;

fn main() {
    let args: Vec<_> = std::env::args().collect();

    use gnuplot::{Color, Figure};
    let mut fg = Figure::new();

    // diffeq parameters
    // time
    let t0 = 0.;
    let tf = args.get(1).map_or(4., |s| s.parse().unwrap());
    let steps = 2usize.pow(8) * (tf - t0) as usize;
    // position
    let x0 = 0.;
    let xf = 16.;
    let res = args.get(2).map_or(64, |s| s.parse().unwrap());
    // calculated params
    // time step
    let h = (xf - x0) / res as f64;
    // position step
    let k = (tf - t0) / steps as f64;

    // convergence
    let r = k / (h * h);
    assert!(r <= 0.5, "divergent, k={} h={}, r=k/h^2={} > 0.5", k, h, r);
    println!("r={}", r);

    // 2D
    if false {
        let axes = fg.axes2d();
        // values
        let mut u = Heat1D(vec![0.; res]);
        let x_values = (0..).map(|i| f64::from(i) * h + x0);
        u.0[res / 2] = 16.;

        println!("init: {:?}", u);
        axes.lines(x_values.clone(), &u.0, &[Color("black")]);

        // heat equation: laplacian u wrt x = du/dt
        // iterate through time
        for _ in 0..steps {
            // iterate over each discrete point in space, keeping endpoints for boundary conditions
            // simulate no boundary conditions (closed system)
            u.step(r).copy_edges();
        }

        println!("done: {:?}", u);
        axes.lines(x_values, &u.0, &[Color("red")]);

        use gnuplot::{AutoOption::Fix, AxesCommon};
        axes.set_x_range(Fix(x0), Fix(xf));
        axes.set_y_range(Fix(0.), Fix(1.));
    }

    // 3D
    if true {
        let mut u = Heat2D::new((res, res));
        u.values[(res + 1) * res / 2] = 256.;

        for _ in 0..steps {
            u.step(r).copy_edges();
        }

        use gnuplot::{AutoOption::*, AxesCommon, ContourStyle};
        let axes = fg.axes3d();
        axes.surface(&u.values, res, res, Some((x0, x0, xf, xf)), &[]);
        axes.set_aspect_ratio(Fix(1.));
        axes.set_x_range(Fix(x0), Fix(xf));
        axes.set_y_range(Fix(x0), Fix(xf));
        axes.show_contours(false, true, ContourStyle::Linear, Fix(""), Auto);
        // top down or ortho camera
        //axes.set_view_map();
    }

    fg.set_terminal("wxt size 768,768", "");
    fg.show();
}