Add Split-Operator in Rust (algorithm-archivists#686)

* Add Split-Operator in Rust

* Remove comments

The small for-loops are fine as they are.

Author: fanninpm

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<f64>,
+    dk: f64,
+    k: Vec<f64>,
+    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<f64> = Vec::with_capacity(res);
+        let dk = PI / xmax;
+        let mut k: Vec<f64> = 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<Complex<f64>>,
+    pe: Vec<Complex<f64>>,
+    ke: Vec<Complex<f64>>,
+    wfc: Vec<Complex<f64>>,
+}
+
+impl Operators {
+    pub fn new(par: &Parameters, v_offset: f64, wfc_offset: f64) -> Operators {
+        let mut v: Vec<Complex<f64>> = Vec::with_capacity(par.res);
+        let mut pe: Vec<Complex<f64>> = Vec::with_capacity(par.res);
+        let mut ke: Vec<Complex<f64>> = Vec::with_capacity(par.res);
+        let mut wfc: Vec<Complex<f64>> = 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<Complex<f64>>, 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<f64>;
+
+    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::<f64>() * 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));
+}

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 %}
 
 <script>