diff --git a/contents/split-operator_method/code/rust/split_op.rs b/contents/split-operator_method/code/rust/split_op.rs new file mode 100644 index 000000000..7269f5148 --- /dev/null +++ b/contents/split-operator_method/code/rust/split_op.rs @@ -0,0 +1,197 @@ +extern crate num; +extern crate rustfft; + +use num::complex::Complex; +use rustfft::FFTplanner; +use std::f64::consts::PI; +use std::fs::File; +use std::io::Write; +use std::path::Path; + +// This implementation is based on the C and C++ implementations. + +#[derive(Clone)] +struct Parameters { + xmax: f64, + res: usize, + dt: f64, + timesteps: usize, + dx: f64, + x: Vec, + dk: f64, + k: Vec, + im_time: bool, +} + +impl Parameters { + pub fn new(xmax: f64, res: usize, dt: f64, timesteps: usize, im_time: bool) -> Parameters { + let dx = 2.0_f64 * xmax / (res as f64); + let mut x: Vec = Vec::with_capacity(res); + let dk = PI / xmax; + let mut k: Vec = Vec::with_capacity(res); + for i in 0..res { + x.push(xmax / (res as f64) - xmax + (i as f64) * dx); + match i { + i if (i < res / 2) => k.push((i as f64) * PI / xmax), + _ => k.push(((i as f64) - (res as f64)) * PI / xmax), + } + } + Parameters { + xmax, + res, + dt, + timesteps, + im_time, + dx, + x, + dk, + k, + } + } +} + +struct Operators { + v: Vec>, + pe: Vec>, + ke: Vec>, + wfc: Vec>, +} + +impl Operators { + pub fn new(par: &Parameters, v_offset: f64, wfc_offset: f64) -> Operators { + let mut v: Vec> = Vec::with_capacity(par.res); + let mut pe: Vec> = Vec::with_capacity(par.res); + let mut ke: Vec> = Vec::with_capacity(par.res); + let mut wfc: Vec> = Vec::with_capacity(par.res); + + for i in 0..par.res { + v.push(Complex::new( + 0.5_f64 * (par.x[i] - v_offset).powi(2), + 0.0_f64, + )); + wfc.push(Complex::new( + (-((par.x[i] - wfc_offset).powi(2)) / 2.0_f64).exp(), + 0.0_f64, + )); + if par.im_time { + ke.push(Complex::new( + (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(), + 0.0_f64, + )); + pe.push(Complex::new((-0.5_f64 * par.dt * v[i].re).exp(), 0.0_f64)); + } else { + ke.push(Complex::new( + 0.0_f64, + (-0.5_f64 * par.dt * par.k[i].powi(2)).exp(), + )); + pe.push(Complex::new(0.0_f64, (-0.5_f64 * par.dt * v[i].re).exp())); + } + } + Operators { v, pe, ke, wfc } + } +} + +fn fft(x: &mut Vec>, inverse: bool) { + let mut y = vec![Complex::new(0.0_f64, 0.0_f64); x.len()]; + let mut p = FFTplanner::new(inverse); + let fft = p.plan_fft(x.len()); + fft.process(x, &mut y); + + for i in 0..x.len() { + x[i] = y[i] / (x.len() as f64).sqrt(); + } +} + +fn split_op(par: &Parameters, opr: &mut Operators) { + let mut density: Vec; + + for i in 0..par.timesteps { + for j in 0..par.res { + opr.wfc[j] *= opr.pe[j]; + } + + fft(&mut opr.wfc, false); + + for j in 0..par.res { + opr.wfc[j] *= opr.ke[j]; + } + + fft(&mut opr.wfc, true); + + for j in 0..par.res { + opr.wfc[j] *= opr.pe[j]; + } + + density = opr.wfc.iter().map(|x| x.norm().powi(2)).collect(); + + if par.im_time { + let sum = density.iter().sum::() * par.dx; + + for j in 0..par.res { + opr.wfc[j] /= sum.sqrt(); + } + } + + // Writing data into a file in the format of: + // index, density, real potential. + let path_name = format!("output{}.dat", i); + let path = Path::new(&path_name); + let display = path.display(); + + let mut file = match File::create(&path) { + Err(why) => panic!("Couldn't create {}: {}", display, why), + Ok(good) => good, + }; + + for j in 0..par.res { + if let Err(why) = writeln!(file, "{}\t{}\t{}", j, density[j], opr.v[j].re) { + panic!("Couldn't write to {}: {}", display, why) + } + if let Err(why) = file.flush() { + panic!("Couldn't flush {}: {}", display, why) + } + } + } +} + +fn calculate_energy(par: &Parameters, opr: &Operators) -> f64 { + let wfc_r = opr.wfc.clone(); + let mut wfc_k = opr.wfc.clone(); + let mut wfc_c = vec![Complex::new(0.0_f64, 0.0_f64); par.res]; + + fft(&mut wfc_k, false); + + for i in 0..par.res { + wfc_c[i] = wfc_r[i].conj(); + } + + let mut energy_k = vec![Complex::new(0.0_f64, 0.0_f64); par.res]; + let mut energy_r = vec![Complex::new(0.0_f64, 0.0_f64); par.res]; + + for i in 0..par.res { + energy_k[i] = wfc_k[i] * Complex::new(par.k[i], 0.0_f64).powi(2); + } + + fft(&mut energy_k, true); + + for i in 0..par.res { + energy_k[i] *= wfc_c[i].scale(0.5_f64); + energy_r[i] = wfc_c[i] * opr.v[i] * wfc_r[i]; + } + + let energy_final = energy_k + .into_iter() + .zip(energy_r.into_iter()) + .fold(0.0_f64, |acc, x| acc + (x.0 + x.1).re); + + energy_final * par.dx +} + +fn main() { + let par = Parameters::new(5.0, 256, 0.05, 100, true); + let mut opr = Operators::new(&par, 0.0, -1.0); + + split_op(&par, &mut opr); + + println!("The energy is {}", calculate_energy(&par, &opr)); +} diff --git a/contents/split-operator_method/split-operator_method.md b/contents/split-operator_method/split-operator_method.md index 20ae39e7b..e97226238 100644 --- a/contents/split-operator_method/split-operator_method.md +++ b/contents/split-operator_method/split-operator_method.md @@ -108,6 +108,8 @@ Regardless, we first need to set all the initial parameters, including the initi [import:11-30, lang:"python"](code/python/split_op.py) {% sample lang="hs" %} [import:17-47, lang:"haskell"](code/haskell/splitOp.hs) +{% sample lang="rs" %} +[import:14-51, lang:"rust"](code/rust/split_op.rs) {% endmethod %} As a note, when we generate our grid in momentum space `k`, we need to split the grid into two lines, one that is going from `0` to `-kmax` and is then discontinuous and goes from `kmax` to `0`. @@ -129,6 +131,8 @@ Afterwards, we turn them into operators: [import:33-54, lang:"python"](code/python/split_op.py) {% sample lang="hs" %} [import:49-66, lang:"haskell"](code/haskell/splitOp.hs) +{% sample lang="rs" %} +[import:53-92, lang:"rust"](code/rust/split_op.rs) {% endmethod %} Here, we use a standard harmonic potential for the atoms to sit in and a gaussian distribution for an initial guess for the probability distribution. @@ -150,6 +154,8 @@ The final step is to do the iteration, itself. [import:57-95, lang:"python"](code/python/split_op.py) {% sample lang="hs" %} [import:68-73, lang:"haskell"](code/haskell/splitOp.hs) +{% sample lang="rs" %} +[import:105-155, lang:"rust"](code/rust/split_op.rs) {% endmethod %} And that's it. @@ -185,6 +191,8 @@ Checking to make sure your code can output the correct energy for a harmonic tra [import:5-127, lang:"python"](code/python/split_op.py) {% sample lang="hs" %} [import, lang:"haskell"](code/haskell/splitOp.hs) +{% sample lang="rs" %} +[import, lang:"rust"](code/rust/split_op.rs) {% endmethod %}