From 05c795b83cff0589dca032cba81429ea3a348fa6 Mon Sep 17 00:00:00 2001 From: Daniel Bauer Date: Sat, 6 Oct 2018 12:30:19 +0200 Subject: [PATCH] moved bias calculation methods to histogram and implemented cyclic conditions --- src/cli.yml | 6 +-- src/histogram.rs | 123 +++++++++++++++++++++++++++++++++++++++++------ src/io.rs | 5 +- src/lib.rs | 42 ++-------------- src/main.rs | 4 +- 5 files changed, 122 insertions(+), 58 deletions(-) diff --git a/src/cli.yml b/src/cli.yml index f19766f..7e241cb 100644 --- a/src/cli.yml +++ b/src/cli.yml @@ -40,9 +40,9 @@ args: value_name: ITERATIONS takes_value: true required: false - - pbc: - long: pbc - help: Apply periodic conditions (Not implemented yet.) + - cyclic: + long: cyclic + help: For periodic reaction coordinates. If this is set, the first and last coordinate bin will be assumed to be neighbors for the bias calculation. - verbose: short: v long: verbose diff --git a/src/histogram.rs b/src/histogram.rs index a070d0d..80d1fe7 100644 --- a/src/histogram.rs +++ b/src/histogram.rs @@ -18,6 +18,7 @@ pub struct Histogram { impl Histogram { pub fn new(first: usize, last: usize, num_points: u32, bins: Vec) -> Histogram { + assert_eq!(last-first+1, bins.len(), "histogram length does not match first/last."); Histogram {first, last, num_points, bins} } @@ -61,15 +62,40 @@ pub struct HistogramSet { pub kT: f32, // histogram for each window - pub histograms: Vec + pub histograms: Vec, + + // flag for cyclic reaction coordinates + pub cyclic: bool, } impl HistogramSet { - pub fn new(num_bins: usize, bin_width: f32, hist_min: f32, hist_max: f32, bias_x0: Vec, bias_fc: Vec, kT: f32, histograms: Vec) -> HistogramSet { + pub fn new(num_bins: usize, bin_width: f32, hist_min: f32, hist_max: f32, bias_x0: Vec, bias_fc: Vec, kT: f32, histograms: Vec, cyclic: bool) -> HistogramSet { let num_windows = histograms.len(); - HistogramSet{num_windows, num_bins, bin_width, hist_min, hist_max, bias_x0, bias_fc, kT, histograms} + HistogramSet{num_windows, num_bins, bin_width, hist_min, hist_max, bias_x0, bias_fc, kT, histograms, cyclic} } + + + // Harmonic bias calculation: bias = 0.5*k(dx)^2 + // if cyclic is true, lowest and highest bins are assumed to be + // neighbors + pub fn calc_bias(&self, bin: usize, window: usize) -> f32 { + let x = self.get_x_for_bin(bin); + let mut dx = (x-self.bias_x0[window]).abs(); + if self.cyclic { + let hist_len = self.hist_max-self.hist_min; + if dx > 0.5*hist_len { + dx -= hist_len; + } + } + 0.5*self.bias_fc[window]*dx*dx + } + + // get center x value for a bin + pub fn get_x_for_bin(&self, bin: usize) -> f32 { + self.hist_min + self.bin_width * ((bin as f32) + 0.5) + } + } impl fmt::Display for HistogramSet { @@ -85,23 +111,92 @@ impl fmt::Display for HistogramSet { #[cfg(test)] mod tests { use super::*; + use super::super::k_B; fn build_hist() -> Histogram { - Histogram{ - first: 5, - last: 7, - num_points: 5, - bins: vec![1.0,1.0,3.0] - } + Histogram::new( + 5, // first + 9, // last + 22, // num_points + vec![1.0, 1.0, 3.0, 5.0, 12.0] // bins + ) + } + + fn build_hist_set() -> HistogramSet { + let h = build_hist(); + HistogramSet::new( + 7, // num bins + 1.0, // bin width + 0.0, // hist min + 9.0, // hist max + vec![7.5], // x0 + vec![10.0], // fc + 300.0*k_B, // kT + vec![h], // hists + false // cyclic + ) } #[test] fn get_bin_count() { let h = build_hist(); - assert_eq!(3.0, h.get_bin_count(7).unwrap()); - assert_eq!(3.0, h.get_bin_count(7).unwrap()); // twice for borrow - assert_eq!(1.0, h.get_bin_count(5).unwrap()); - assert_eq!(None, h.get_bin_count(4)); - assert_eq!(None, h.get_bin_count(8)); + let expected = vec![None, Some(1.0), Some(1.0), Some(3.0), Some(5.0), Some(12.0), None]; + let test_offset = 4; + for i in 4..10 { + match expected[i-test_offset] { + Some(x) => assert_eq!(x, h.get_bin_count(i).unwrap()), + None => assert!(h.get_bin_count(i) == None) + } + } + } + + #[test] + fn calc_bias() { + let hs = build_hist_set(); + + // 7th element -> x=7.5, x0=7.5 + assert_eq!(0.0, hs.calc_bias(7, 0)); + + // 8th element -> x=8.5, x0=7.5 + assert_eq!(5.0, hs.calc_bias(8, 0)); + + + // 9th element -> x=9.5, x0=7.5 + assert_eq!(20.0, hs.calc_bias(9, 0)); + + + // 1st element -> x=0.5, x0=7.5. non-cyclic! + assert_eq!(245.0, hs.calc_bias(0, 0)); + } + + #[test] + fn calc_bias_offset_cyclic() { + let mut hs = build_hist_set(); + hs.cyclic = true; + + // 7th element -> x=7.5, x0=7.5 + assert_eq!(0.0, hs.calc_bias(7, 0)); + + // 8th element -> x=8.5, x0=7.5 + assert_eq!(5.0, hs.calc_bias(8, 0)); + + + // 9th element -> x=9.5, x0=7.5 + assert_eq!(20.0, hs.calc_bias(9, 0)); + + + // 1st element -> x=0.5, x0=7.5 + // cyclic flag makes bin 0 neighboring bin 9, so the distance is actually 2 + assert_eq!(20.0, hs.calc_bias(0, 0)); + } + + #[test] + fn get_x_for_bin() { + let hs = build_hist_set(); + let expected: Vec = vec![0,1,2,3,4,5,6,7,8].iter() + .map(|x| *x as f32 + 0.5).collect(); + for i in 0..9 { + assert_eq!(expected[i], hs.get_x_for_bin(i)); + } } } \ No newline at end of file diff --git a/src/io.rs b/src/io.rs index d673ed7..4e6eb67 100644 --- a/src/io.rs +++ b/src/io.rs @@ -71,7 +71,7 @@ pub fn read_data(cfg: &Config) -> Option { if histograms.len() > 0 { let bin_width = (cfg.hist_max - cfg.hist_min)/(cfg.num_bins as f32); - Some(HistogramSet::new(cfg.num_bins, bin_width, cfg.hist_min, cfg.hist_max, bias_x0, bias_fc, kT, histograms)) + Some(HistogramSet::new(cfg.num_bins, bin_width, cfg.hist_min, cfg.hist_max, bias_x0, bias_fc, kT, histograms, cfg.cyclic)) } else { None } @@ -146,7 +146,8 @@ mod tests { verbose: false, tolerance: 0.0, max_iterations: 0, - temperature: 300.0 + temperature: 300.0, + cyclic: false } } diff --git a/src/lib.rs b/src/lib.rs index c55d776..2f17862 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,7 @@ pub struct Config { pub tolerance: f32, pub max_iterations: usize, pub temperature: f32, + pub cyclic: bool, } impl fmt::Display for Config { @@ -49,29 +50,17 @@ fn is_converged(old_F: &Vec, new_F: &Vec, tolerance: f32) -> bool { true } -// Harmonic bias calculation: bias = 0.5*k(dx)^2 -fn calc_bias(k: f32, x0: f32, x: f32) -> f32 { - let dx = (x-x0).abs(); - 0.5*k*dx*dx -} - -// get center x value for a bin -fn get_x_for_bin(bin: usize, min: f32, width: f32) -> f32 { - min + width * ((bin as f32) + 0.5) -} - // estimate the probability of a bin of the histogram set based on F values // This evaluates the first WHAM equation for each bin fn calc_bin_probability(bin: usize, hs: &HistogramSet, F: &Vec) -> f32 { let mut denom_sum = 0.0; let mut bin_count = 0.0; - let x = get_x_for_bin(bin, hs.hist_min, hs.bin_width); for window in 0..hs.num_windows { let h: &Histogram = &hs.histograms[window]; if let Some(count) = h.get_bin_count(bin) { bin_count += count; } - let bias = calc_bias(hs.bias_fc[window], hs.bias_x0[window], x); + let bias = hs.calc_bias(bin, window); let bias_offset = ((F[window] - bias) / hs.kT).exp(); denom_sum += (h.num_points as f32) * bias_offset; } @@ -83,8 +72,7 @@ fn calc_bin_probability(bin: usize, hs: &HistogramSet, F: &Vec) -> f32 { fn calc_window_F(window: usize, hs: &HistogramSet, P: &Vec) -> f32 { let mut ln_sum = 0.0; for bin in 0..hs.num_bins { - let x = get_x_for_bin(bin, hs.hist_min, hs.bin_width); - let bias = calc_bias(hs.bias_fc[window], hs.bias_x0[window], x); + let bias = hs.calc_bias(bin, window); ln_sum += P[bin] * (-bias/hs.kT).exp() } -hs.kT * ln_sum.ln() @@ -247,7 +235,7 @@ fn dump_state(hs: &HistogramSet, F: &Vec, F_prev: &Vec, P: &Vec, println!("# PMF"); println!("#x\t\tFree Energy\t\tP(x)"); for bin in 0..hs.num_bins { - let x = get_x_for_bin(bin, hs.hist_min, hs.bin_width); + let x = hs.get_x_for_bin(bin); println!("{:9.5}\t{:9.5}\t{:9.5}", x, A[bin], P[bin]); } println!("# Bias offsets"); @@ -275,18 +263,10 @@ mod tests { assert!(!converged); } - #[test] - fn calc_bias() { - let x0 = 10.0; - let x = 5.0; - let k = 500.0; - assert_eq!(6250.0, super::calc_bias(k, x0, x)); - } - fn create_test_hs() -> HistogramSet { let h1 = Histogram::new(0, 2, 10, vec![3.0, 4.0, 3.0]); let h2 = Histogram::new(0, 3, 20, vec![3.0, 2.0, 5.0, 10.0]); - HistogramSet::new(4, 1.0, 0.0, 4.0, vec![1.0, 2.0], vec![10.0, 10.0], 2.479, vec![h1, h2]) + HistogramSet::new(4, 1.0, 0.0, 4.0, vec![1.0, 2.0], vec![10.0, 10.0], 2.479, vec![h1, h2], false) } fn assert_near(a: f32, b: f32, tolerance: f32) { @@ -303,7 +283,6 @@ mod tests { let F = super::calc_window_F(window, &hs, &probability); assert_near(expected[window], F, 0.001); } - } #[test] @@ -324,17 +303,6 @@ mod tests { } } - #[test] - fn get_x_for_bin() { - let min = 0.0; - let width = 1.0; - let expected = vec!(0.5, 1.5, 2.5, 3.5, 4.5); - for i in 0..5 { - let x = super::get_x_for_bin(i, min, width); - assert_eq!(expected[i], x); - } - } - #[test] fn perform_wham_iteration() { let hs = create_test_hs(); diff --git a/src/main.rs b/src/main.rs index e655f85..cf6250e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,12 +18,12 @@ fn cli() -> Result> { let num_bins: usize = matches.value_of("bins").unwrap().parse()?; let verbose: bool = matches.is_present("verbose"); let temperature: f32 = matches.value_of("temperature").unwrap().parse()?; - let tolerance: f32 = matches.value_of("tolerance").unwrap_or("0.000001").parse()?; let max_iterations: usize = matches.value_of("iterations").unwrap_or("100000").parse()?; + let cyclic: bool = matches.is_present("cyclic"); Ok(wham::Config{metadata_file, hist_min, hist_max, num_bins, - verbose, tolerance, max_iterations, temperature}) + verbose, tolerance, max_iterations, temperature, cyclic}) } fn main() {