finite_difference.rs

extern crate time;
use time::{PreciseTime};

fn main() {
    let nx: usize = 100;
    let ny: usize = 100;
    let delta_x: f64 = 1.0;
    let dx = delta_x / (nx as f64);
    let mut temp_lhs = Vec::<f64>::with_capacity(nx * ny);
    let mut temp_rhs = Vec::<f64>::with_capacity(nx * ny);
    
    let idx = |i: usize, j: usize| -> usize {
        debug_assert!(i < nx);
        debug_assert!(j < ny);
        let result = i + nx * j;
        debug_assert!(result < nx * ny);
        result
    };

    fn loop_all<F: FnMut(usize, usize)> (nx: usize, ny: usize, f: &mut F) {
        for j in 0..ny {
            for i in 0..nx {
                debug_assert!(i < nx);
                debug_assert!(j < ny);
                f(i, j);
            }
        }   
    }

    fn loop_internal<F: FnMut(usize, usize)> (nx: usize, ny: usize, f: &mut F) {
        debug_assert!(ny - 1 > 0);
        debug_assert!(nx - 1 > 0);
        for j in 1..ny-1 {
            for i in 1..nx-1 {
                debug_assert!(i < nx);
                debug_assert!(j < ny);
                f(i, j);
            }
        }   
    }

    // initialize temperature to constant value
    let initial_temperature: f64 = 20.;
    loop_all(nx, ny, & mut |_,_| {
        temp_lhs.push(initial_temperature);
        temp_rhs.push(initial_temperature);
    });
    
    fn compute_time_step(dx: f64, lambda: f64, cfl: f64) -> f64 {
        dx * dx * cfl * 0.5 / lambda
    }
    
    enum Boundary { Left, Right, Top, Bottom };
    
    let impose_dirichlet_boundary_conditions = |temp: &mut Vec<f64>, boundary: Boundary, value: f64| {
      match boundary {
        Boundary::Left => {
            let i = 0;
            for j in 0..ny {
                temp[idx(i, j)] = value;
            }
        }
        Boundary::Right => {
            let i = nx - 1;
            for j in 0..ny {
                temp[idx(i, j)] = value;
            }
        }
        Boundary::Bottom => {
            let j = 0;
            for i in 0..nx {
                temp[idx(i, j)] = value;
            }
        }
        Boundary::Top => {
            let j = ny - 1;
            for i in 0..nx {
                temp[idx(i, j)] = value;
            }
        }
      }
    };
    
    let impose_boundary_conditions = |mut temp_lhs: &mut Vec<f64>| {
        impose_dirichlet_boundary_conditions(&mut temp_lhs, Boundary::Left, 19.);
        impose_dirichlet_boundary_conditions(&mut temp_lhs, Boundary::Right, 20.);
        impose_dirichlet_boundary_conditions(&mut temp_lhs, Boundary::Top, 20.);
        impose_dirichlet_boundary_conditions(&mut temp_lhs, Boundary::Bottom, 20.);
    };  
    
    let clear = |temp_rhs: &mut Vec<f64>| {
        loop_internal(nx, ny, &mut |i,j| {
            temp_rhs[idx(i, j)] = 0.
        });
    };
    
    let kernel = |temp_rhs: &mut Vec<f64>, temp_lhs: & Vec<f64>, i :usize, j: usize, dx: f64| {
        debug_assert!(i < nx);
        debug_assert!(j < ny);
        debug_assert!(dx > 0.);
        debug_assert!(temp_lhs[idx(i+1, j)] >= 0.);
        debug_assert!(temp_lhs[idx(i-1, j)] >= 0.);
        debug_assert!(temp_lhs[idx(i, j+1)] >= 0.);
        debug_assert!(temp_lhs[idx(i, j-1)] >= 0.);
        temp_rhs[idx(i, j)] = (temp_lhs[idx(i+1, j)] - temp_lhs[idx(i-1, j)]) / (2. * dx)
            + (temp_lhs[idx(i, j+1)] - temp_lhs[idx(i, j-1)]) / (2. * dx);
    };
    
    let compute_rhs = |mut temp_rhs: &mut Vec<f64>, temp_lhs: &Vec<f64>| {
        loop_internal(nx, ny, &mut |i,j| {
            kernel(&mut temp_rhs, &temp_lhs, i, j, dx);
        });
    };

    let mut time: f64 = 0.;
    let time_end: f64 = 10.;
    let mut dt = compute_time_step(1. / (nx as f64), 1., 0.1);
    let mut no_iterations = 0;
    let start = PreciseTime::now();
    
    while time != time_end {
        if time + dt >= time_end {
            dt = time_end - time;
        }
        impose_boundary_conditions(&mut temp_lhs);
        clear(&mut temp_rhs);
        compute_rhs(&mut temp_rhs, &temp_lhs);

        loop_internal(nx, ny, &mut |i,j| {
            debug_assert!(dt >= 0.);
            debug_assert!(dx >= 0.);
            temp_lhs[idx(i,j)] += dt * dx * temp_rhs[idx(i, j)];
            debug_assert!(temp_lhs[idx(i,j)] >= 0.);
        });
        time += dt;
        no_iterations += 1;
        if time == time_end {
            break
        }
    }
    let end = PreciseTime::now();
    let total_duration = start.to(end);
    println!("total_duration: {}", total_duration);
}