writing trajectories

.gitignore

@@ -1,3 +1,4 @@
 target
 .idea
-*.iml
+*.iml
+*.xyz

src/energy.rs

@@ -1,5 +1,5 @@
 
-pub fn get_total_energy(rx: &[f64], ry: &[f64], rz: &[f64], num_particles: usize, l_x: f64, l_y: f64, l_z: f64, cutoff_squared: f64, e_corr: f64) -> (f64, f64) {
+pub fn get_total_energy(rx: &[f64], ry: &[f64], rz: &[f64], num_particles: usize, l_x: f64, l_y: f64, l_z: f64, cutoff_squared: f64, e_corr: f64, e_shift: f64) -> (f64, f64) {
     let mut energy = 0.0;
     let mut virial = 0.0;
     let hl_x = l_x / 2.0;
@@ -9,7 +9,7 @@ pub fn get_total_energy(rx: &[f64], ry: &[f64], rz: &[f64], num_particles: usize
         for j in i+1..num_particles-1 {
             let dist_squared = get_particle_distance_squared(rx[i], ry[i],rz[i],rx[j],ry[j],rz[j], l_x, l_y, l_z, hl_x, hl_y, hl_z);
             if dist_squared < cutoff_squared {
-                let (e,v) = eval_pair_energy(dist_squared);
+                let (e,v) = eval_pair_energy(dist_squared, e_shift);
                 energy += e;
                 virial += v;
             }
@@ -19,7 +19,7 @@ pub fn get_total_energy(rx: &[f64], ry: &[f64], rz: &[f64], num_particles: usize
     return (energy, virial);
 }
 
-pub fn get_particle_energy(rx: &[f64], ry: &[f64], rz: &[f64], p_index: usize, num_particles: usize, l_x: f64, l_y: f64, l_z: f64, cutoff_squared: f64) -> (f64, f64) {
+pub fn get_particle_energy(rx: &[f64], ry: &[f64], rz: &[f64], p_index: usize, num_particles: usize, l_x: f64, l_y: f64, l_z: f64, cutoff_squared: f64, e_shift: f64) -> (f64, f64) {
     let mut energy = 0.0;
     let mut virial = 0.0;
     let hl_x = l_x / 2.0;
@@ -30,7 +30,7 @@ pub fn get_particle_energy(rx: &[f64], ry: &[f64], rz: &[f64], p_index: usize, n
 
         let dist_squared = get_particle_distance_squared(rx[i], ry[i], rz[i], rx[p_index], ry[p_index], rz[p_index], l_x, l_y, l_z, hl_x, hl_y, hl_z);
         if dist_squared < cutoff_squared {
-            let (e,v) = eval_pair_energy(dist_squared);
+            let (e,v) = eval_pair_energy(dist_squared, e_shift);
             energy += e;
             virial += v;
         }
@@ -65,22 +65,22 @@ fn test_get_particle_distance_squared() {
     assert!( (get_particle_distance_squared(x1,y1,z1,x2,y2,z2, 9.0,9.0,9.0, 4.5,4.5,4.5) - 16.0) < 0.00001, "{}", get_particle_distance_squared(x1,y1,z1,x2,y2,z2, 9.0,9.0,9.0, 4.5,4.5,4.5));
 }
 
-fn eval_pair_energy(dist_squared: f64) -> (f64, f64) {
+fn eval_pair_energy(dist_squared: f64, e_shift: f64) -> (f64, f64) {
     let r6 = ::LJ_SIG/(dist_squared * dist_squared * dist_squared);
     let r62 = r6*r6;
-    let energy = 4.0 * ::LJ_EPS * (r62 - r6);
+    let energy = 4.0 * ::LJ_EPS * (r62 - r6) - e_shift;
     let virial = 48.0 * ::LJ_EPS * ( r62 - 0.5 * r6 );
     return (energy, virial);
 }
 
 #[test]
 fn test_eval_pair_energy() {
-    let (e,v) = eval_pair_energy(1.0);
+    let (e,v) = eval_pair_energy(1.0, 0.0);
     assert!( (e - 0.0).abs() < 0.00001, "{}",  e);
 
-    let (e,v) = eval_pair_energy(2.0);
+    let (e,v) = eval_pair_energy(2.0, 0.0);
     assert!( (e - -0.4375).abs() < 0.00001, "{}",  e);
 
-    let (e,v) = eval_pair_energy(0.5);
+    let (e,v) = eval_pair_energy(0.5, 0.0);
     assert!( (e - 224.0).abs() < 0.00001, "{}",  e);
 }

src/main.rs

@@ -7,7 +7,10 @@ mod energy;
 use energy::*;
 use std::io::prelude::*;
 extern crate argparse;
-use argparse::{ArgumentParser, Store, StoreFalse};
+use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
+mod trajectory;
+use trajectory::*;
+
 
 const LJ_EPS : f64 = 1.0;
 const LJ_SIG : f64 = 1.0;
@@ -24,7 +27,10 @@ macro_rules! println_stderr(
     } }
 );
 
-fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize, NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64, CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool) {
+fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize,
+                  NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64,
+                  CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool,
+                  OUTPUT_PREFIX: &mut String, OUTPUT_INTERVAL: &mut i64, OUTPUT_MINIM: &mut bool) {
     let mut ap = ArgumentParser::new();
     ap.set_description("LJ MC simulation.");
     ap.refer(NUM_STEPS)
@@ -51,9 +57,15 @@ fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize, NUM_PARTIC
     ap.refer(SCALE)
         .add_option(&["--nodisplacementscale"], StoreFalse,
                     "Disable displacement scaling");
-//    ap.refer(TRAJECTORY_STEPSIZE)
+    ap.refer(OUTPUT_PREFIX)
-//        .add_option(&["--trj_steps"], Store,
+        .add_option(&["-o", "--output"], Store,
-//                    "Number of steps between writing to the trajectory file");
+                    "Output file prefix");
+    ap.refer(OUTPUT_INTERVAL)
+        .add_option(&["--osteps"], Store,
+                    "Number of steps between writing to the trajectory file. -1 only writes last frame");
+    ap.refer(OUTPUT_MINIM)
+        .add_option(&["--writeminimization"], StoreTrue,
+                    "Enables writing of minimization step to trajectory");
 //    ap.refer(VACUUM_DIMENSION)
 //        .add_option(&["--vacuum"], Store,
 //                    "Dimension of vacuum space below and above the intial system.");
@@ -80,13 +92,16 @@ fn main() {
     let mut displacement = 0.1;
 
     let mut  TAILCORR : bool = true;
-    let mut SHIFT: bool = false;
+    let mut SHIFT: bool = true;
     let mut SCALE: bool = true;
+
+    let mut output_prefix = "montecarlo".to_string();
+    let mut output_interval : i64 = 100;
+    let mut output_minim : bool = false;
     parse_cmd_args(&mut sample_steps, &mut minim_steps, &mut num_particles,
                    &mut density, &mut temperature,
-                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT);
+                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT,
-
+                   &mut output_prefix, &mut output_interval, &mut output_minim);
-
 
     println_stderr!("");
     println_stderr!("################################################################");
@@ -116,6 +131,7 @@ fn main() {
         if rx.len() == num_particles { break; }
     }
 
+    let e_shift = if SHIFT { 4.0 * LJ_EPS * ( (LJ_SIG/cutoff).powi(12) - (LJ_SIG/cutoff).powi(6) ) } else { 0.0 };
     let e_corr = if TAILCORR { 8.0/3.0*std::f64::consts::PI*density*LJ_EPS*LJ_SIG.powi(3)*((1.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
     let p_corr = if TAILCORR { 16.0/3.0*std::f64::consts::PI*density.powi(2)*LJ_EPS*LJ_SIG.powi(3)*((2.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
 
@@ -123,9 +139,9 @@ fn main() {
     println_stderr!("System volume: {:8.3}, Dimensions {:.3}/{:.3}/{:.3}", volume, l_x, l_y, l_z);
     println_stderr!("Minimization steps: {}, Sampling steps: {}", minim_steps, sample_steps);
     println_stderr!("LJ params eps: {}, sigma: {}, cutoff: {}", LJ_EPS, LJ_SIG, cutoff);
-    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, SHIFT, p_corr);
+    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, e_shift, p_corr);
 
-    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr);
+    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
     let mut energy_sum = 0.0;
     let mut virial_sum = 0.0;
     let mut step_counter = 0;
@@ -137,6 +153,10 @@ fn main() {
     println_stderr!("################################################################");
     println_stderr!("");
 
+    // prepare and write first trajectory frame
+    let mut trajectory : XYZTrajectory = XYZTrajectory::new(&format!("{}.xyz", output_prefix));
+    if output_minim { trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true); }
+
     for step in 0..minim_steps+sample_steps {
 
         // select rnd particle
@@ -148,7 +168,7 @@ fn main() {
         let oldZ = rz[rnd_index];
 
         // old particle energy
-        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared);
+        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
 
         // rnd displacement and PBC
         rx[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
@@ -162,14 +182,14 @@ fn main() {
         if rz[rnd_index] > l_z { rz[rnd_index] -= l_z }
 
         // calculate energy difference
-        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared);
+        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
         let dE = new_particle_energy - old_particle_energy;
 
         //accept move
         if rng.gen::<f64>() < (-beta * dE).exp() {
             accept_counter += 1;
             if step % 1000 == 0 { // calculate total energy every 1000 steps to account for rounding errors
-                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr);
+                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
                 energy = e;
                 virial = v;
             } else {
@@ -220,7 +240,13 @@ fn main() {
 
         if step > minim_steps && step_counter % 5000 == 0 {
             println_stderr!("Step  {:<10}Energy: {:<12.3}Virial: {:<12.3}", step_counter, energy, virial);
+        }
 
+        // write trajectory maybe
+        if step as i64 % output_interval == 0 {
+            if step > minim_steps || output_minim {
+                trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
+            }
         }
 
     }
@@ -271,8 +297,8 @@ Pressure: {}",
          minim_steps, sample_steps,
         LJ_EPS, LJ_SIG, cutoff,
         num_particles, density, temperature, volume, l_x, l_y, l_z, displacement,
-        e_corr, SHIFT, p_corr,
+        e_corr, e_shift, p_corr,
         step_counter, accept_counter, final_acceptance_rate, final_energy, particle_energy, final_virial, pressure);
 
-
+    trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
 }

src/trajectory.rs

@@ -0,0 +1,167 @@
+#![allow(dead_code)]
+#![allow(unused_must_use)] // hate.
+#![allow(unused_variables)]
+
+use std::error::Error;
+use std::io::prelude::*;
+use std::fs::File;
+use std::path::Path;
+
+use std::io::BufReader;
+
+
+pub struct XYZTrajectory {
+    file: File,
+
+}
+
+impl XYZTrajectory {
+    pub fn new(filename: &String) -> XYZTrajectory {
+        let path = Path::new(filename);
+        let display = path.display();
+        // Open a file in write-only mode, returns `io::Result<File>`
+        let traj_file = match File::create(&path) {
+            Err(why) => panic!("couldn't create {}: {}",
+                               display,
+                               why.description()),
+            Ok(file) => file,
+        };
+
+        XYZTrajectory { file: traj_file }
+
+    }
+
+
+    pub fn write(&mut self, rx: &[f64], ry: &[f64], rz: &[f64], num_particles: usize, box_x : f64, box_y : f64, box_z : f64, temp: f64 ,lj_eps : f64, lj_sig : f64, lj_cutoff : f64, flush: bool) {
+        self.file.write(format!("{} ## Box: {} {} {} Temp: {} LJ: {}/{}/{}\n", num_particles, box_x,box_y,box_z,temp, lj_eps, lj_sig, lj_cutoff).as_bytes());
+        for i  in 0..num_particles {
+            let formatted = format!("atom{} {} {} {}\n",
+                                    i+1, rx[i], ry[i], rz[i]
+            );
+            self.file.write(formatted.as_bytes());
+        }
+
+        if flush { self.file.flush(); }
+
+    }
+}
+
+pub struct Frame {
+    pub rx : Vec<f64>,
+    pub ry : Vec<f64>,
+    pub rz : Vec<f64>,
+    pub num_particles: usize,
+    pub box_x: f64,
+    pub box_y: f64,
+    pub box_z: f64,
+    pub temperature: f64,
+    pub lj_eps: f64,
+    pub lj_sig: f64,
+    pub lj_cutoff: f64,
+}
+
+pub struct TrjReader {
+    pub reader: BufReader<File>,
+}
+impl TrjReader {
+    pub fn new() -> TrjReader {
+        let filename = "trajectory.xyz";
+        let file = File::open(filename).expect("Failed to open file.");
+        let reader : BufReader<File> = BufReader::new(file);
+        return TrjReader { reader:reader };
+    }
+
+    // get next frame
+    pub fn next_frame(&mut self) -> Frame {
+        let buffer_string = &mut String::new();
+        match self.reader.read_line(buffer_string) {
+            Ok(size) => {
+                if size == 0 {
+                    panic!("no new frame!")
+                }
+            },
+            Err(e) => panic!("no new frame!")
+        }
+        let first_line_vec : Vec<&str> = buffer_string.split_whitespace().collect();
+        let num_particles = first_line_vec[0].parse::<usize>().unwrap();
+        let box_x = first_line_vec[3].parse::<f64>().unwrap();
+        let box_y = first_line_vec[4].parse::<f64>().unwrap();
+        let box_z = first_line_vec[5].parse::<f64>().unwrap();
+        let temp = first_line_vec[7].parse::<f64>().unwrap();
+        let lj : Vec<&str> = first_line_vec[9].split("/").collect();
+        let lj_eps = lj[0].parse::<f64>().unwrap();
+        let lj_sig = lj[1].parse::<f64>().unwrap();
+        let lj_cutoff = lj[2].parse::<f64>().unwrap();
+        let mut rx = Vec::new();
+        let mut ry = Vec::new();
+        let mut rz = Vec::new();
+        for i in 0..num_particles-1{
+            let atom_line = &mut String::new();
+            match self.reader.read_line(atom_line) {
+                Ok(size) => {
+                    let atom_vec : Vec<&str> = atom_line.split_whitespace().collect();
+                    rx.push(atom_vec[1].parse::<f64>().unwrap());
+                    ry.push(atom_vec[2].parse::<f64>().unwrap());
+                    rz.push(atom_vec[3].parse::<f64>().unwrap());
+                },
+                Err(e) => return panic!("no new frame!")
+            };
+        }
+        let frame = Frame {
+            rx : rx,
+            ry : ry,
+            rz : rz,
+            num_particles : num_particles,
+            box_x : box_x,
+            box_y : box_y,
+            box_z : box_z,
+            temperature : temp,
+            lj_eps : lj_eps,
+            lj_sig : lj_sig,
+            lj_cutoff : lj_cutoff,
+        };
+        return frame;
+    }
+
+    // read next frame data into the frame
+    pub fn update_with_next(&mut self, frame: &mut Frame) -> bool {
+        let buffer_string = &mut String::new();
+        match self.reader.read_line(buffer_string) {
+            Ok(size) => {
+                if size == 0 {
+                    return false;
+                }
+            },
+            Err(e) => panic!("no new frame!")
+        }
+        let first_line_vec : Vec<&str> = buffer_string.split_whitespace().collect();
+        frame.num_particles = first_line_vec[0].parse::<usize>().unwrap();
+        frame.box_x = first_line_vec[3].parse::<f64>().unwrap();
+        frame.box_y = first_line_vec[4].parse::<f64>().unwrap();
+        frame.box_z = first_line_vec[5].parse::<f64>().unwrap();
+        frame.temperature = first_line_vec[7].parse::<f64>().unwrap();
+        let lj : Vec<&str> = first_line_vec[9].split("/").collect();
+        frame.lj_eps = lj[0].parse::<f64>().unwrap();
+        frame.lj_sig = lj[1].parse::<f64>().unwrap();
+        frame.lj_cutoff = lj[2].parse::<f64>().unwrap();
+        for i in 0..frame.num_particles-1 {
+            let atom_line = &mut String::new();
+            match self.reader.read_line(atom_line) {
+                Ok(size) => {
+                    if size == 0 {
+                        return false;
+                    } else {
+                        let atom_vec : Vec<&str> = atom_line.split_whitespace().collect();
+                        frame.rx[i] = atom_vec[1].parse::<f64>().unwrap();
+                        frame.ry[i]= atom_vec[2].parse::<f64>().unwrap();
+                        frame.rz[i] = atom_vec[3].parse::<f64>().unwrap();
+                    }
+                },
+                Err(e) => {return false;}
+            }
+        }
+
+        return true;
+    }
+
+}