Skip to content

Instantly share code, notes, and snippets.

@tqn

tqn/heat.rs Secret

Last active Feb 25, 2019
Embed
What would you like to do?
Learning Rust by Simulating Heat Diffusion: Unabridged
#[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
}
}
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();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.