bauer/WHAM
1
0
Fork 0
mirror of https://github.com/dnlbauer/WHAM.git synced 2026-09-10 22:25:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

use exp(F) and exp(U) for wham iterations - huge performance increase!

Browse Source
This commit is contained in:
Daniel Bauer
2018-10-18 00:13:54 +02:00
parent b064ecea6d
commit 6b9661bd13
3 changed files with 10176 additions and 10153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
src/lib.rs
View File

@@ -53,25 +53,25 @@ fn calc_bin_probability(bin: usize, ds: &Dataset, F: &[f64]) -> f64 {
let mut bin_count: f64 = 0.0;
for (window, h) in ds.histograms.iter().enumerate() {
bin_count += h.bins[bin];
let bias = -ds.kT*ds.calc_bias(bin, window).ln();
let bias_offset = ((F[window] - bias) / ds.kT).exp();
denom_sum += (h.num_points as f64) * bias_offset;
let bias = ds.calc_bias(bin, window);
denom_sum += (h.num_points as f64) * bias * F[window];
}
bin_count / denom_sum
bin_count / denom_sum
}
// estimate the bias offset F of the histogram based on given probabilities
// This evaluates the second WHAM equation for each window
fn calc_window_F(window: usize, ds: &Dataset, P: &[f64]) -> f64 {
(0..ds.num_bins).zip(P.iter()) // zip bins and P
let f: f64 = (0..ds.num_bins).zip(P.iter()) // zip bins and P
.filter_map(|bin_and_prob: (usize, &f64)| {
if bin_and_prob.1 == &0.0 { // skip zeros for speed
None
} else {
Some(bin_and_prob.1 * ds.calc_bias(bin_and_prob.0, window))
let bias = ds.calc_bias(bin_and_prob.0, window);
Some(bin_and_prob.1 * bias)
}
}).sum()
}).sum();
1.0/f
}
// One full WHAM iteration includes calculation of new probabilities P and
@@ -169,24 +169,31 @@ pub fn run(cfg: &Config) -> Result<(), Box<Error>>{
println!("{}",&histograms);
// allocate only once for better performance
let mut F_prev: Vec<f64> = vec![f64::INFINITY; histograms.num_windows];
let mut F: Vec<f64> = vec![0.0; histograms.num_windows];
let mut F_prev: Vec<f64> = vec![f64::NAN; histograms.num_windows];
let mut F: Vec<f64> = vec![1.0; histograms.num_windows];
let mut F_tmp: Vec<f64> = vec![f64::NAN; histograms.num_windows];
let mut P: Vec<f64> = vec![f64::NAN; histograms.num_bins];
let mut A: Vec<f64> = vec![f64::NAN; histograms.num_bins];
// perform WHAM until convergence
let mut iteration = 0;
while !is_converged(&F_prev, &F, cfg.tolerance) && iteration < cfg.max_iterations {
let mut converged = false;
while !converged && iteration < cfg.max_iterations {
iteration += 1;
// store F values before the next iteration
F_prev.copy_from_slice(&F[..]);
F_prev.copy_from_slice(&F);
// perform wham iteration and update F
perform_wham_iteration(&histograms, &F_prev, &mut F, &mut P);
// output some stats during calculation
if iteration % 10 == 0 {
println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F));
F_tmp.copy_from_slice(&F);
F.iter_mut().map(|f| *f=-histograms.kT*f.ln()).count();
F_prev.iter_mut().map(|f| *f=-histograms.kT*f.ln()).count();
converged = is_converged(&F_prev, &F, cfg.tolerance);
println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F));
F.copy_from_slice(&F_tmp);
}
// Dump free energy and bias offsets
@@ -235,6 +242,7 @@ fn dump_state(ds: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64], A: &[f64]) {
mod tests {
use super::histogram::{Dataset,Histogram};
use std::f64;
use super::k_B;
macro_rules! assert_delta {
($x:expr, $y:expr, $d:expr) => {
@@ -242,6 +250,14 @@ mod tests {
}
}
fn create_test_ds() -> Dataset {
let h1 = Histogram::new(10, vec![0.0, 1.0, 1.0, 8.0, 0.0]);
let h2 = Histogram::new(10, vec![0.0, 0.0, 8.0, 1.0, 1.0]);
Dataset::new(5, vec![5], vec![1.0], vec![0.0], vec![4.0],
vec![1.0, 1.0], vec![10.0, 10.0], 300.0*k_B, vec![h1, h2], false)
}
#[test]
fn is_converged() {
let new = vec![1.0,1.0];
@@ -255,23 +271,23 @@ mod tests {
assert!(!converged);
}
fn create_test_ds() -> Dataset {
let h1 = Histogram::new(10, vec![0.0, 0.0, 3.0, 4.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
let h2 = Histogram::new(20, vec![0.0, 0.0, 0.0, 3.0, 2.0, 5.0, 10.0, 0.0, 0.0, 0.0, 0.0]);
Dataset::new(4, vec![1], vec![1.0], vec![0.0], vec![4.0],
vec![1.0, 2.0], vec![10.0, 10.0], 2.479, vec![h1, h2], false)
#[test]
fn calc_bin_probability() {
let ds = create_test_ds();
let F = vec![1.0; ds.num_bins] ;
let expected = vec!(0.0, 0.0825296687031316, 40.92355847097493,
124226.70003377, 2308526035.5283747);
for b in 0..ds.num_bins {
let p = super::calc_bin_probability(b, &ds, &F);
assert_delta!(expected[b], p, 0.0000001);
}
}
fn assert_near(a: f64, b: f64, tolerance: f64) {
let d = (a-b).abs();
assert!(d <= tolerance, "Values are not close: {}, {}, d={}", &a, &b, &d);
}
#[test]
#[test]
fn calc_bias_offset() {
let ds = create_test_ds();
let probability = vec!(0.959, 0.331, 0.656, 46.750);
let expected = vec!(0.786289183, 1.10629119);
let probability = vec!(0.0, 0.1, 0.2, 0.3, 0.4);
let expected = vec!(15.927477169990633, 15.927477169990633);
for window in 0..ds.num_windows {
let F = super::calc_window_F(window, &ds, &probability);
assert_delta!(expected[window], F, 0.0000001);
@@ -279,39 +295,20 @@ mod tests {
}
#[test]
#[ignore] // TODO
fn calc_bin_probability() {
let ds = create_test_ds();
let F = vec!(0.0, 0.0);
let expected = vec!(0.959, 0.331, 0.656, 46.750);
for b in 0..4 {
let p = super::calc_bin_probability(b, &ds, &F);
assert_near(expected[b], p, 0.001);
}
let F = vec!(1.0, 1.0);
let expected = vec!(0.641, 0.221, 0.439, 31.232);
for b in 0..4 {
let p = super::calc_bin_probability(b, &ds, &F);
assert_near(expected[b], p, 0.001);
}
}
#[test]
#[ignore] // TODO
fn perform_wham_iteration() {
let ds = create_test_ds();
let prev_F = vec![0.0; ds.num_windows];
let mut F = vec![0.0; ds.num_windows];
let prev_F = vec![1.0; ds.num_windows];
let mut F = vec![f64::NAN; ds.num_windows];
let mut P = vec![f64::NAN; ds.num_bins];
super::perform_wham_iteration(&ds, &prev_F, &mut F, &mut P);
let expected_F = vec!(0.5948, -0.2513);
let expected_P = vec!(0.959, 0.331, 0.656, 46.750);
let expected_F = vec!(1.0, 1.0);
let expected_P = vec!(0.0, 0.0825296687031316, 40.92355847097493,
124226.70003377, 2308526035.5283747);
for bin in 0..ds.num_bins {
assert_near(expected_P[bin], P[bin], 0.01)
assert_delta!(expected_P[bin], P[bin], 0.01)
}
for window in 0..ds.num_windows {
assert_near(expected_F[window], F[window], 0.01)
assert_delta!(expected_F[window], F[window], 0.01)
}
}