Skip to content

Add Split-Operator in Rust #686

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 17, 2020
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 199 additions & 0 deletions contents/split-operator_method/code/rust/split_op.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,199 @@
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;

#[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 {
// this should use an iterator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure how you turn this into an iterator as you need the index for accessing a different array than the one you are iterating over. Not saying it is impossible, just harder. I think this is perfectly fine. This also makes the code a bit clearer.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just realized I could've used the enumerate() method, but (as you point out) that may come at the cost of clarity.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Honestly, I don't mind using that either. It just depends on what Rust users are used to reading.

for j in 0..par.res {
opr.wfc[j] *= opr.pe[j];
}

fft(&mut opr.wfc, false);

// this should use an iterator
for j in 0..par.res {
opr.wfc[j] *= opr.ke[j];
}

fft(&mut opr.wfc, true);

// this should use an iterator
for j in 0..par.res {
opr.wfc[j] *= opr.pe[j];
}

// this is a good example (c.f. the cpp implementation)
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));
}
8 changes: 8 additions & 0 deletions contents/split-operator_method/split-operator_method.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:11-49, 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`.
Expand All @@ -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:51-90, 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.
Expand All @@ -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:103-157, lang:"rust"](code/rust/split_op.rs)
{% endmethod %}

And that's it.
Expand Down Expand Up @@ -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>
Expand Down