Learning Rust by Simulating Heat Diffusion

In the summer of 2018, Rust caught my attention. For most people, C++ is the primary choice for any job that requires fast performance. What prevented me from having confidence in my C++ abilities were weeks of practice avoiding memory issues and working with the language architecture. Even though I also dabbled in some front-end tech like React and the whole tangled mess of NPM Dependency Hell™ (cough leftpad), I thought it was strange that C++, the language that is most people’s first choice for writing fast, low-level code, feels like 30 years of afterthoughts carefully glued onto each other. We have learned many things about programming in 30 years.

This is what made me so excited to find Rust, a language that is not only fun to code in and (after the initial learning curve) highly intuitive, but also lightning-fast and with opinionated tools. It combines the delight of npm installing some random package dumped into the package registry and watching your program execute in a fraction of a second, thanks to the package manager cargo and an LLVM-backed compiler. I could now pull in other people’s polluted, bloated code and obsess over shaving off microseconds of execution at the same time. Unsurprisingly, I like math and physics. While learning Rust, it happened that I was curious about implementing a numerical method for the heat equation, which is a partial differential equation (PDE) that models heat diffusion in most materials:

$$ \frac{\partial u}{\partial t} = - a \nabla^2 u $$

This equation says that the rate of change of heat ($\frac{\partial u}{\partial t}$) is proportional to the Laplacian of the heat distribution ($ \nabla^2 u$) at that point. The Laplacian operator tells you how different a point is compared to its neighbors. If the point is, on average, less than its neighbors, the Laplacian will be positive, and vice versa. This equation captures your intuition that a spot that much hotter or colder than its neighbors will equalize faster.

I chose to simulate heat diffusion because it’s relatively simple compared to other analytic systems, like waves or fluid.

Math

Since we’re concerned with running a simulation, we have to use discrete steps in time and space instead of an infinitely detailed continuum, assigning a certain heat value to every point and determining what it will be at the next moment. What’s so great about the heat equation is that, for each point, I can compute the first time derivative in terms of its current state in position space. Since I know how fast each point is changing, I can model the time evolution of the system! Although I wanted to eventually create a higher dimension simulation, I had to start in 1D.

Recalling the original heat equation, the change in heat $du$ is proportional to the Laplacian times the infinitesimal time step $dt$. However, we have to convert each differential operator into a discrete equivalent. The discrete Laplacian is actually quite simple in one dimension. You simply add the neighbors and subtract two times the point to evaluate the lowest order approximation of the Laplacian at that point. The finite difference equation

$$ \frac{u^{n+1}_j-u^n_j}{k} = \frac{u^n_{j+1} - 2u^n_j + u^n_{j-1}}{h^2} $$

models our heat distribution $u$, where the subscript is the space index and superscript is the time index. This formula, shamelessly ripped from Wikipedia, replaces the derivatives in the heat equation with finite difference approximations. A derivative is defined as the slope of a tangent line, which in turn is the limit of secant lines through the function at a target point and another point arbitrarily close. The finite difference simply picks another point so that one can approximate the derivative. (To get the approximation of the Laplacian on the right-hand side, use the definition of the derivative twice, but replace the arbitrarily small step with a finite one.)

The limit of secant lines is a tangent line. Since derivatives are slopes of tangent lines, they can be approximated by slopes of secant lines using a finite difference quotient.

The limit of secant lines is a tangent line. Since derivatives are slopes of tangent lines, they can be approximated by slopes of secant lines using a finite difference quotient.

In the finite difference approximation, the $k$ corresponds to $dt$, or the time step we will use. $h$ corresponds to the position step. Wikipedia notes that the numerical method is stable for $r=\frac{k}{h^2} \leq 0.5$. Otherwise, each iteration of our simulation will gradually lead to more extreme and useless values.

Gallery

I produced these images with my final version of the heat simulation program, which includes another variant that works on 2D space. These are plotted with the library gnuplot.

These are simulations of heat distribution along one or two spatial dimensions, starting with a single spike, or impulse, of heat at the center and letting it diffuse. It’s no accident that these look like normal distributions’ bell curves — the impulse response/Green’s function of the heat PDE, is actually a normal distribution that spreads out over time.

1D heat @ 16 initial spike, 4 time units, 64  position bins, 16 position units wide. This is the plot that corresponds to the abridged program above.

1D heat @ 16 initial spike, 4 time units, 64 position bins, 16 position units wide. This is the plot that corresponds to the abridged program above.

1D heat @ 16 initial spike, 8 time units, 64 position bins, 16 position units wide.

1D heat @ 16 initial spike, 8 time units, 64 position bins, 16 position units wide.

2D heat @ 16 initial spike, 4 time units, 64² position bins, 64 position units wide. Note the z-scale.

2D heat @ 16 initial spike, 4 time units, 64² position bins, 64 position units wide. Note the z-scale.

2D heat @ 16 initial spike, 16 time units, 64² position bins, 64 position units wide. Note the z-scale.

2D heat @ 16 initial spike, 16 time units, 64² position bins, 64 position units wide. Note the z-scale.

Code

We’ll walk through creating a 1D heat simulation and plotting it. This isn’t exactly how I did it — which involved a lot of trial and error — but it’s pretty close.

Writing the simulation

Let’s create a file called heat.rs to contain our simulation’s heavy lifting. We will run the simulation with a separate file.

First, we create a struct Heat1D to encapsulate the data and functionality of a one dimensional heat simulation. The pub keyword exposes the struct and its fields outside the module. This type of struct is known as a tuple struct, since it is essentially a named tuple type. The single element in the tuple is a vector of 64-bit floats (doubles), Vec<f64>. It contains the state of the simulation at a given time. The function step in the impl block will compute a time step given the $r$ value from above. step also mutably borrows the Heat1D object it is called on, allowing it to change the values in the Vec<f64>.

In heat.rs:

1
2
3
4
5
6
7
pub struct Heat1D(pub Vec<f64>);

impl Heat1D {
    pub fn step(&mut self, r: f64) {
        // TODO
    }
}

Every time we call the step method, we want the system to advance in time according to the $r$ value. Our Laplacian isn’t defined for points without neighbors, so the method should not change the edges. To compute the Laplacian, we need to iterate over each value in the vector, making a copy of the previous value and peeking at the next value. Instead of doing this by incrementing an index (boring!), we can use Rust’s iterators.

In the body of step:

1
2
3
4
5
6
7
8
9
// Gets a mutable iterator to the wrapped Vec, making it peekable.
let mut iter = self.0.iter_mut().peekable();
// Copies the first element into a new variable, making it mutable.
if let Some(&mut mut prev) = iter.next() {
    // Gets the next element and peek the following element
    while let (Some(cur), Some(next)) = (iter.next(), iter.peek()) {
        // TODO
    }
}

On line 2, we use .0 to access the enclosed Vec<f64>, get an iterator that allows you to change elements, and allow peeking of the next value. Line 4 uses the special syntax if let, which attempts to match the value on the right to the pattern on the left. Some, along with None, are the enum variants of Option, which acts as an optional value — think Java’s nullable pointers or C++’s std::optional. Similarly, the while let loop on line 4 terminates when there is no longer an element after the current one.

Next, we need to update cur while copying it to prev for use in the following computation. We could just make a copy of the whole vector, but where’s the fun in that? prev is a plain f64, so it needs no dereferencing. cur is a mutable reference to an f64 (&mut f64), which allows us to change it, but it needs one use of the dereference operator * to access its data. Since we peek next, it’s actually a &&mut f64, so we need to dereference it twice!

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
let mut iter = self.0.iter_mut().peekable();

if let Some(&mut mut prev) = iter.next() {
    // This variable will hold the delta while we copy the value of `cur`.
    // It's scoped out here to prevent excess allocations and drops.
    let mut delta;
    while let (Some(cur), Some(next)) = (iter.next(), iter.peek()) {
        // Laplacian formula from above, multiplied by our constant `r`
        delta = (prev - 2. * *cur + **next) * r;
        // copy
        prev = *cur;
        // make a time step
        *cur += delta;
    }
}

The type f64 implements a special Rust trait called Copy, which indicates that any assignments will copy the data instead of moving it and rendering the initial value useless. All Rust primitives are Copy. However, structs and enums are not Copy by default. To finish our simulation, we will add a method called copy_edges. This method will copy the heat values next to each edge to the edge itself. If we call copy_edges before every time step, then heat will neither escape nor enter via the edges, since the neighboring cells have the same heat value, simulating a closed system.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
impl Heat1D {

    // ...

    pub fn copy_edges(&mut self) {
        // Fairly certain this is optimized by the compiler,
        // but it doesn't hurt. It also makes referring to the length easier.
        let len = self.0.len();
        // Copies the left edge...
        self.0[0] = self.0[1];
        // and the right edge.
        self.0[len - 1] = self.0[len - 2];
    }
}

Running the simulation

Now, we can create a main.rs file to actually run the simulation. We use mod to declare the existence of heat.rs so we can access its contents. Then, we utilize the use keyword to bring our Heat1D struct into the scope so we don’t have to type its fully qualified name all the time.

In main.rs:

1
2
3
mod heat;

use heat::Heat1D;

We need a few parameters for our simulation, namely the physical size of the simulation, duration we will simulate, and the resolution of space and time our program will use. We can then use these parameters to compute $r$, the ratio between our discrete Laplacian and the change in heat at each point. We’ll make SIM_SIZE the number of space bins we will split our one-dimensional interval in. Rust requires that we declare the type of compile-time constants. We use type usize because it represents a count, which in this case, is 64 evenly spaced points in space.

1
const SIM_SIZE: usize = 64;

Then, we define the width of the of the simulation as 16 “unit lengths.” Note that we use the type f64, or a 64-bit float. This is a double in other languages.

1
2
// Width, in unit lengths
const SIM_WIDTH: f64 = 16.;

We do a similar thing for the time scale, using 2048 steps and running the simulation for 8 “unit durations.”

1
2
3
4
5
// Number of "time" steps we will take. This can be scaled
// to simulate a different area, diffusivity, or time.
const SIM_STEPS: usize = 2048;
// Amount of time we will simulate
const SIM_TIME: f64 = 8.;

Now we create our main function. Rust will call this function in the main file when the program is run as a binary. We compute $h$, $k$, and $r$, and use the assert! macro to ensure that $r \leq 0.5$. If it isn’t, the simulation is unstable and will fail to converge. The assert! terminates the program if this condition is false, guaranteeing that the simulation is stable.

1
2
3
4
5
6
fn main() {
    let h = SIM_WIDTH / SIM_SIZE as f64;
    let k = SIM_TIME / SIM_STEPS as f64;
    let r = k / (h * h);
    assert!(r <= 0.5);
}

We must cast the usizes to f64s using the as keyword, since Rust purposefully avoids implicit casting in exchange for compiler type inference. In our main method, we create a Heat1D instance, utilizing the vec! macro to make SIM_SIZE number of bins, all set to 0. In the center, we set an initial spike of heat of arbitrary magnitude to let diffuse. This represents the initial conditions for our simulation.

1
2
let mut sim = Heat1D(vec![0.; SIM_SIZE]);
sim.0[SIM_SIZE / 2] = 24.;

We can use the println! macro to print the contents of the simulation. The {:?} syntax formats any object that implements the Debug trait, such as our Vec of floats.

1
println!("initial: {:?}", sim.0);

We run the simulation’s step function for as many steps as SIM_STEPS, copying the edges each time to simulate a closed system. This for loop iterates over a range, using the _ syntax is used to discard the iteration variable.

1
2
3
4
for _ in 0..SIM_STEPS {
    sim.step(r);
    sim.copy_edges();
}

After the simulation is done, we print the final array.

1
println!("final: {:?}", sim.0);

When run, this should give us something like

1
2
initial: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 24.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
final: [0.17248554593305104, 0.17248554593305104, 0.1744182236115872, 0.17826985766072606, 0.18401292722175752, 0.19160596443559655, 0.20099335546434025, 0.21210511552915368, 0.22485667024967929, 0.23914868087076047, 0.254866954563554, 0.2718824826928585, 0.2900516496246303, 0.3092166522586212, 0.32920616604353714, 0.34983628688108603, 0.37091177024599253, 0.39222757931191526, 0.4135707432154226, 0.43472251520459076, 0.4554608087397201, 0.4755628781007125, 0.49480819917555635, 0.5129814963119549, 0.5298758528329515, 0.545295836421062, 0.5590605663703764, 0.5710066479168974, 0.5809908996151519, 0.5888928030637045, 0.5946166101190433, 0.598093050898522, 0.5992805960836506, 0.5981662389334395, 0.5947657755702609, 0.5891235760206062, 0.5813118526601063, 0.5714294465956663, 0.5596001655994571, 0.5459707190066957, 0.5307083050741963, 0.5139979143160072, 0.4960394180209627, 0.47704451434849277, 0.4572336050321169, 0.4368326738391171, 0.4160702336863684, 0.3951744029365273, 0.3743701632176742, 0.35387684151114096, 0.3339058486727881, 0.3146586954570221, 0.29632529597308244, 0.279082557780338, 0.2630932479532838, 0.2485051157993189, 0.23545024581214044, 0.2240446091352919, 0.2143877784559605, 0.2065627699230424, 0.20063597637171804, 0.19665715873900064, 0.19465946689000047, 0.19465946689000047]

which is impressive! The final heat distribution seems to increase towards the center, but unfortunately a list of numbers is a pain to read. I took the time to install the gnuplot crate using cargo and the native library it hooks into so that I could plot the data. I added this code at the end of main to generate the images below.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
use gnuplot::{AutoOption::Fix, AxesCommon, Color, Figure};
let mut fg = Figure::new();
let axes = fg.axes2d();
// Scale the x data so it lines up with the width given.
axes.lines((0..).map(|i| f64::from(i) * h), &sim.0, &[Color("red")]);
// Set the window.
axes.set_x_range(Fix(0.), Fix(SIM_WIDTH));
axes.set_y_range(Fix(0.), Fix(1.));
// "wxt" is the GUI name and "768,768" is the size of it in pixels.
fg.set_terminal("wxt size 768,768", "");
fg.show();

And we’re done!

Abridged Source

The complete program can be found at this gist.

heat.rs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
pub struct Heat1D(pub Vec<f64>);

impl Heat1D {
    pub fn step(&mut self, r: f64) {
        // create a mutable iterator that allows us to immutably peek.
        let mut iter = self.0.iter_mut().peekable();

        if let Some(&mut mut prev) = iter.next() {
            // This variable will hold the delta while we copy the value of `cur`.
            // It's scoped out here to prevent excess allocations and drops.
            let mut delta;
            while let (Some(cur), Some(next)) = (iter.next(), iter.peek()) {
                // Laplacian formula from above, multiplied by our constant `r`
                delta = (prev - 2. * *cur + **next) * r;
                // copy
                prev = *cur;
                // make a time step
                *cur += delta;
            }
        }
    }

    pub fn copy_edges(&mut self) {
        // Fairly certain this is optimized by the compiler,
        // but it doesn't hurt. It also makes referring to the length easier.
        let len = self.0.len();
        // Copies the left edge...
        self.0[0] = self.0[1];
        // and the right edge.
        self.0[len - 1] = self.0[len - 2];
    }
}

main.rs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
mod heat;

use heat::Heat1D;

// Number of space steps in the simulation. This can be scaled
// to simulate a different area, diffusivity, or time.
const SIM_SIZE: usize = 64;
// Width, in unit lengths
const SIM_WIDTH: f64 = 16.;
// Number of "time" steps we will take.
const SIM_STEPS: usize = 2048;
// Amount of time we will simulate
const SIM_TIME: f64 = 8.;

fn main() {
    let h = SIM_WIDTH / SIM_SIZE as f64;
    let k = SIM_TIME / SIM_STEPS as f64;

    let r = k / (h * h);
    assert!(r <= 0.5);

    let mut sim = Heat1D(vec![0.; SIM_SIZE]);
    sim.0[SIM_SIZE / 2] = 24.;

    println!("initial: {:?}", sim.0);

    for _ in 0..SIM_STEPS {
        sim.step(r);
        sim.copy_edges();
    }

    println!("final: {:?}", sim.0);

    // use gnuplot to render graphs of the result
    use gnuplot::{AutoOption::Fix, AxesCommon, Color, Figure};
    let mut fg = Figure::new();
    let axes = fg.axes2d();
    // Scale the x data so it lines up with the width given.
    axes.lines((0..).map(|i| f64::from(i) * h), &sim.0, &[Color("red")]);
    // Set the window.
    axes.set_x_range(Fix(0.), Fix(SIM_WIDTH));
    axes.set_y_range(Fix(0.), Fix(1.));
    // "wxt" is the GUI name and "768,768" is the size of it in pixels.
    fg.set_terminal("wxt size 768,768", "");
    fg.show();
}

Conclusion

This project was pretty fun. It turns out that simulating such a simple system as heat diffusion is relatively easy. The Rust borrow checker helped keep my code bug-free and kept me focused on the logic of my program rather than the technical details. Also, the graphs look really cool, which is of course the most important part.