bauer/WHAM
1
0
Fork 0
mirror of https://github.com/dnlbauer/WHAM.git synced 2026-09-11 14:45:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

comments and code cleanup

Browse Source
This commit is contained in:
Daniel Bauer
2018-10-18 01:03:29 +02:00
parent 6b9661bd13
commit c835520bf4
3 changed files with 198 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
example/1d/wham.out
View File

@@ -1,127 +1,101 @@
#Coor Free +/- Prob +/- #x Free Energy Probability
-3.108600 7.109190 -nan 0.003556 -nan -3.108600 7.118166 0.003561
-3.045800 5.311412 -nan 0.007311 -nan -3.045800 5.331290 0.007290
-2.983000 3.849095 -nan 0.013139 -nan -2.983000 3.879241 0.013048
-2.920200 2.933168 -nan 0.018968 -nan -2.920200 2.968659 0.018796
-2.857400 1.907181 -nan 0.028620 -nan -2.857400 1.942494 0.028362
-2.794600 1.385773 -nan 0.035274 -nan -2.794600 1.418408 0.034994
-2.731800 1.139312 -nan 0.038937 -nan -2.731800 1.169398 0.038668
-2.669000 0.814751 -nan 0.044348 -nan -2.669000 0.843024 0.044073
-2.606200 0.609195 -nan 0.048157 -nan -2.606200 0.636021 0.047887
-2.543400 0.694709 -nan 0.046534 -nan -2.543400 0.720123 0.046299
-2.480600 1.042492 -nan 0.040478 -nan -2.480600 1.066434 0.040297
-2.417800 1.487426 -nan 0.033865 -nan -2.417800 1.509855 0.033734
-2.355000 1.994612 -nan 0.027634 -nan -2.355000 2.015549 0.027544
-2.292200 2.207453 -nan 0.025374 -nan -2.292200 2.226942 0.025306
-2.229400 2.555138 -nan 0.022072 -nan -2.229400 2.573176 0.022026
-2.166600 2.575426 -nan 0.021894 -nan -2.166600 2.591946 0.021861
-2.103800 2.461877 -nan 0.022913 -nan -2.103800 2.476835 0.022893
-2.041000 2.500517 -nan 0.022561 -nan -2.041000 2.513953 0.022555
-1.978200 2.458051 -nan 0.022949 -nan -1.978200 2.470032 0.022956
-1.915400 2.217054 -nan 0.025276 -nan -1.915400 2.227585 0.025299
-1.852600 2.096290 -nan 0.026530 -nan -1.852600 2.105310 0.026570
-1.789800 1.771307 -nan 0.030222 -nan -1.789800 1.778779 0.030286
-1.727000 1.441175 -nan 0.034499 -nan -1.727000 1.447127 0.034593
-1.664200 0.901911 -nan 0.042825 -nan -1.664200 0.906384 0.042968
-1.601400 0.301267 -nan 0.054485 -nan -1.601400 0.304263 0.054699
-1.538600 0.245895 -nan 0.055708 -nan -1.538600 0.247397 0.055960
-1.475800 0.000000 -nan 0.061480 -nan -1.475800 0.000000 0.061795
-1.413000 0.537640 -nan 0.049559 -nan -1.413000 0.536142 0.049843
-1.350200 1.431041 -nan 0.034639 -nan -1.350200 1.428046 0.034859
-1.287400 2.402640 -nan 0.023464 -nan -1.287400 2.398149 0.023627
-1.224600 3.790635 -nan 0.013450 -nan -1.224600 3.784647 0.013552
-1.161800 5.704072 -nan 0.006246 -nan -1.161800 5.696583 0.006297
-1.099000 7.665152 -nan 0.002845 -nan -1.099000 7.656158 0.002870
-1.036200 9.965495 -nan 0.001131 -nan -1.036200 9.955000 0.001142
-0.973400 12.369366 -nan 0.000432 -nan -0.973400 12.357407 0.000436
-0.910600 14.967624 -nan 0.000152 -nan -0.910600 14.954254 0.000154
-0.847800 17.760111 -nan 0.000050 -nan -0.847800 17.745327 0.000050
-0.785000 20.575788 -nan 0.000016 -nan -0.785000 20.559461 0.000016
-0.722200 22.801620 -nan 0.000007 -nan -0.722200 22.783560 0.000007
-0.659400 25.195750 -nan 0.000003 -nan -0.659400 25.175865 0.000003
-0.596600 26.441571 -nan 0.000002 -nan -0.596600 26.419980 0.000002
-0.533800 27.909051 -nan 0.000001 -nan -0.533800 27.885999 0.000001
-0.471000 29.053664 -nan 0.000001 -nan -0.471000 29.029334 0.000001
-0.408200 30.392176 -nan 0.000000 -nan -0.408200 30.366556 0.000000
-0.345400 31.642366 -nan 0.000000 -nan -0.345400 31.615252 0.000000
-0.282600 32.810590 -nan 0.000000 -nan -0.282600 32.781716 0.000000
-0.219800 33.762125 -nan 0.000000 -nan -0.219800 33.731378 0.000000
-0.157000 34.508470 -nan 0.000000 -nan -0.157000 34.476014 0.000000
-0.094200 35.421120 -nan 0.000000 -nan -0.094200 35.387300 0.000000
-0.031400 35.614610 -nan 0.000000 -nan -0.031400 35.579756 0.000000
0.031400 35.555913 -nan 0.000000 -nan 0.031400 35.520194 0.000000
0.094200 35.381388 -nan 0.000000 -nan 0.094200 35.344763 0.000000
0.157000 34.924243 -nan 0.000000 -nan 0.157000 34.886460 0.000000
0.219800 33.673332 -nan 0.000000 -nan 0.219800 33.633998 0.000000
0.282600 32.727755 -nan 0.000000 -nan 0.282600 32.686511 0.000000
0.345400 31.255053 -nan 0.000000 -nan 0.345400 31.211777 0.000000
0.408200 29.725589 -nan 0.000000 -nan 0.408200 29.680451 0.000000
0.471000 28.078949 -nan 0.000001 -nan 0.471000 28.032242 0.000001
0.533800 26.487669 -nan 0.000002 -nan 0.533800 26.439599 0.000002
0.596600 24.481873 -nan 0.000003 -nan 0.596600 24.432467 0.000003
0.659400 22.360728 -nan 0.000008 -nan 0.659400 22.309893 0.000008
0.722200 20.238391 -nan 0.000018 -nan 0.722200 20.186040 0.000019
0.785000 18.348925 -nan 0.000039 -nan 0.785000 18.295046 0.000040
0.847800 16.276613 -nan 0.000090 -nan 0.847800 16.221232 0.000093
0.910600 14.308785 -nan 0.000198 -nan 0.910600 14.251911 0.000204
0.973400 12.620723 -nan 0.000390 -nan 0.973400 12.562353 0.000402
1.036200 11.243896 -nan 0.000678 -nan 1.036200 11.184029 0.000698
1.099000 10.102041 -nan 0.001071 -nan 1.099000 10.040677 0.001103
1.161800 9.431925 -nan 0.001401 -nan 1.161800 9.369064 0.001444
1.224600 9.159750 -nan 0.001563 -nan 1.224600 9.095392 0.001612
1.287400 9.321889 -nan 0.001464 -nan 1.287400 9.256035 0.001511
1.350200 9.881419 -nan 0.001170 -nan 1.350200 9.814068 0.001208
1.413000 11.014057 -nan 0.000743 -nan 1.413000 10.945207 0.000768
1.475800 12.589255 -nan 0.000395 -nan 1.475800 12.518898 0.000409
1.538600 14.481607 -nan 0.000185 -nan 1.538600 14.409744 0.000191
1.601400 16.525295 -nan 0.000082 -nan 1.601400 16.451965 0.000084
1.664200 18.663885 -nan 0.000035 -nan 1.664200 18.589166 0.000036
1.727000 20.745586 -nan 0.000015 -nan 1.727000 20.669504 0.000016
1.789800 22.882359 -nan 0.000006 -nan 1.789800 22.804791 0.000007
1.852600 24.611022 -nan 0.000003 -nan 1.852600 24.531690 0.000003
1.915400 26.263981 -nan 0.000002 -nan 1.915400 26.182612 0.000002
1.978200 27.362886 -nan 0.000001 -nan 1.978200 27.279439 0.000001
2.041000 28.656631 -nan 0.000001 -nan 2.041000 28.571380 0.000001
2.103800 29.415584 -nan 0.000000 -nan 2.103800 29.328969 0.000000
2.166600 29.985317 -nan 0.000000 -nan 2.166600 29.897739 0.000000
2.229400 30.409211 -nan 0.000000 -nan 2.229400 30.320908 0.000000
2.292200 30.154379 -nan 0.000000 -nan 2.292200 30.065391 0.000000
2.355000 29.898976 -nan 0.000000 -nan 2.355000 29.809137 0.000000
2.417800 29.431664 -nan 0.000000 -nan 2.417800 29.340601 0.000000
2.480600 28.566154 -nan 0.000001 -nan 2.480600 28.473346 0.000001
2.543400 27.774399 -nan 0.000001 -nan 2.543400 27.679353 0.000001
2.606200 26.550657 -nan 0.000001 -nan 2.606200 26.453160 0.000002
2.669000 24.524409 -nan 0.000003 -nan 2.669000 24.424649 0.000003
2.731800 22.387526 -nan 0.000008 -nan 2.731800 22.285933 0.000008
2.794600 20.061389 -nan 0.000020 -nan 2.794600 19.958407 0.000021
2.857400 17.710243 -nan 0.000051 -nan 2.857400 17.606291 0.000053
2.920200 15.524401 -nan 0.000122 -nan 2.920200 15.420138 0.000128
2.983000 13.172725 -nan 0.000313 -nan 2.983000 13.069599 0.000328
3.045800 11.135246 -nan 0.000708 -nan 3.045800 11.036085 0.000740
3.108600 9.106349 -nan 0.001597 -nan 3.108600 9.015657 0.001664
#Window Free +/-
#0 0.000000 -nan
#1 -4.304935 -nan
#2 -11.772646 -nan
#3 -18.486638 -nan
#4 -22.834670 -nan
#5 -24.172759 -nan
#6 -22.273996 -nan
#7 -17.239462 -nan
#8 -10.107963 -nan
#9 -5.386748 -nan
#10 -10.219397 -nan
#11 -18.470489 -nan
#12 -25.454944 -nan
#13 -3.099544 -nan
#14 -10.785962 -nan
#15 -19.937669 -nan
#16 -27.163650 -nan
#17 -31.721722 -nan
#18 -33.474006 -nan
#19 -32.952581 -nan
#20 -32.029659 -nan
#21 -32.180907 -nan
#22 -33.041061 -nan
#23 -32.813126 -nan
#24 -30.673600 -nan

7
src/io.rs
View File

@@ -236,13 +236,6 @@ mod tests {
println!("{:?}", ds); println!("{:?}", ds);
assert_eq!(2, ds.num_windows); assert_eq!(2, ds.num_windows);
assert_eq!(cfg.num_bins[0], ds.dimens_lengths[0]); assert_eq!(cfg.num_bins[0], ds.dimens_lengths[0]);
// fields are private
// assert_eq!(cfg.hist_min[0], ds.hist_min[0]);
// assert_eq!(cfg.hist_max[0], ds.hist_max[0]);
// let expected_bin_width = (cfg.hist_max[0] - cfg.hist_min[0])/cfg.num_bins[0] as f64;
// assert_eq!(expected_bin_width, ds.bin_width);
// assert_eq!(vec![0.0, 1.0], ds.bias_pos);
// assert_eq!(vec![100.0, 200.0], ds.bias_fc);
assert_eq!(cfg.temperature * k_B, ds.kT); assert_eq!(cfg.temperature * k_B, ds.kT);
assert_eq!(2, ds.histograms.len()) assert_eq!(2, ds.histograms.len())
} }

225
src/lib.rs
View File

@@ -31,23 +31,22 @@ pub struct Config {
impl fmt::Display for Config { impl fmt::Display for Config {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Metadata={}, hist_min={:?}, hist_max={:?}, bins={:?}\nverbose={}, tolerance={}, iterations={}, temperature={}, cyclic={:?}" , self.metadata_file, self.hist_min, write!(f, "Metadata={}, hist_min={:?}, hist_max={:?}, bins={:?} verbose={}, tolerance={}, iterations={}, temperature={}, cyclic={:?}", self.metadata_file, self.hist_min, self.hist_max, self.num_bins,
self.hist_max, self.num_bins, self.verbose, self.tolerance, self.verbose, self.tolerance, self.max_iterations, self.temperature, self.cyclic)
self.max_iterations, self.temperature, self.cyclic)
} }
} }
// Checks for convergence between two WHAM iterations. WHAM is considered as // Checks for convergence between two WHAM iterations. WHAM is considered as
// converged if the absolute difference for the calculated bias offset is // converged if the maximal difference for the calculated bias offsets is
// smaller then a tolerance value for every simulation window. // smaller then a tolerance value.
fn is_converged(old_F: &[f64], new_F: &[f64], tolerance: f64) -> bool { fn is_converged(old_F: &[f64], new_F: &[f64], tolerance: f64) -> bool {
!new_F.iter().zip(old_F.iter()) !new_F.iter().zip(old_F.iter())
.map(|x| { (x.0-x.1).abs() }) .map(|x| { (x.0-x.1).abs() })
.any(|diff| { diff > tolerance }) .any(|diff| { diff > tolerance })
} }
// estimate the probability of a bin of the histogram set based on F values // estimate the probability of a bin of the histogram set based on given bias offsets (F)
// This evaluates the first WHAM equation for each bin // This evaluates the first WHAM equation for each bin.
fn calc_bin_probability(bin: usize, ds: &Dataset, F: &[f64]) -> f64 { fn calc_bin_probability(bin: usize, ds: &Dataset, F: &[f64]) -> f64 {
let mut denom_sum: f64 = 0.0; let mut denom_sum: f64 = 0.0;
let mut bin_count: f64 = 0.0; let mut bin_count: f64 = 0.0;
@@ -59,8 +58,8 @@ fn calc_bin_probability(bin: usize, ds: &Dataset, F: &[f64]) -> f64 {
bin_count / denom_sum bin_count / denom_sum
} }
// estimate the bias offset F of the histogram based on given probabilities // estimate the bias offset F of the histogram based on given probabilities.
// This evaluates the second WHAM equation for each window // This evaluates the second WHAM equation for each window and returns exp(F/kT)
fn calc_window_F(window: usize, ds: &Dataset, P: &[f64]) -> f64 { fn calc_window_F(window: usize, ds: &Dataset, P: &[f64]) -> f64 {
let f: 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)| { .filter_map(|bin_and_prob: (usize, &f64)| {
@@ -78,44 +77,6 @@ fn calc_window_F(window: usize, ds: &Dataset, P: &[f64]) -> f64 {
// new bias offsets F based on previous bias offsets F_prev. This updates // new bias offsets F based on previous bias offsets F_prev. This updates
// the values in vectors F and P // the values in vectors F and P
fn perform_wham_iteration(ds: &Dataset, F_prev: &[f64], F: &mut [f64], P: &mut [f64]) { fn perform_wham_iteration(ds: &Dataset, F_prev: &[f64], F: &mut [f64], P: &mut [f64]) {
// reset bias offsets
for window in 0..ds.num_windows {
F[window] = 0.0;
}
// for bin in 0..ds.num_bins {
// let x = get_x_for_bin(bin, ds.hist_min, ds.bin_width);
// let mut num = 0.0;
// let mut denom = 0.0;
// for window in 0..ds.num_windows {
// match ds.histograms[window].get_bin_count(bin) {
// Some(c) => num += c,
// _ => {}
// }
// let bias = calc_bias(
// ds.bias_fc[window],
// ds.bias_pos[window],
// x);
// let bf = ((F_prev[window]-bias) / ds.kT).exp();
// denom += ds.histograms[window].num_points as f64* bf
// }
// P[bin] = num / denom;
// for window in 0..ds.num_windows {
// let bias = calc_bias(
// ds.bias_fc[window],
// ds.bias_pos[window],
// x);
// let bf = (-bias/ds.kT).exp() * P[bin];
// F[window] += bf;
// }
// }
// for window in 0..ds.num_windows {
// F[window] = -ds.kT * F[window].ln();
// }
// evaluate first WHAM equation for each bin to // evaluate first WHAM equation for each bin to
// estimage probabilities based on previous offsets (F_prev) // estimage probabilities based on previous offsets (F_prev)
for bin in 0..ds.num_bins { for bin in 0..ds.num_bins {
@@ -127,9 +88,78 @@ fn perform_wham_iteration(ds: &Dataset, F_prev: &[f64], F: &mut [f64], P: &mut [
for window in 0..ds.num_windows { for window in 0..ds.num_windows {
F[window] = calc_window_F(window, ds, P); F[window] = calc_window_F(window, ds, P);
} }
} }
pub fn run(cfg: &Config) -> Result<(), Box<Error>>{
println!("Supplied WHAM options: {}", &cfg);
println!("Reading input files.");
// TODO Better error handling with nice error messages instead of a panic!
let histograms = io::read_data(&cfg)
.expect("No datapoints in histogram boundaries.");
println!("{}",&histograms);
// allocate required vectors.
let mut P: Vec<f64> = vec![f64::NAN; histograms.num_bins]; // bin probability
let mut F: Vec<f64> = vec![1.0; histograms.num_windows]; // bias offset exp(F/kT)
let mut F_prev: Vec<f64> = vec![f64::NAN; histograms.num_windows]; // previous bias offset
let mut F_tmp: Vec<f64> = vec![f64::NAN; histograms.num_windows]; // temp storage for F
let mut iteration = 0;
let mut converged = false;
// perform WHAM until convergence
while !converged && iteration < cfg.max_iterations {
iteration += 1;
// store F values before the next iteration
F_prev.copy_from_slice(&F);
// perform wham iteration (this updates F and P)
perform_wham_iteration(&histograms, &F_prev, &mut F, &mut P);
// convergence check
if iteration % 10 == 0 {
// This backups exp(F/kT) in a temporary vector and calculates true F and F_prev for
// convergence. Finally, F is restored. F_prev does not need to be restored because
// its overwritten for the next iteration.
F_tmp.copy_from_slice(&F);
for window in 0..histograms.num_windows {
F[window] = -histograms.kT * F[window].ln();
F_prev[window] = -histograms.kT * F_prev[window].ln();
}
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
//if iteration % 100 == 0 {
// free_energy(&histograms, &mut P, &mut A);
// dump_state(&histograms, &F, &F_prev, &P, &A);
//}
}
// Normalize P to sum(P) = 1.0
let P_sum: f64 = P.iter().sum();
P.iter_mut().map(|p| *p /= P_sum).count();
// calculate free energy and dump state
println!("Finished. Dumping final PMF");
let free_energy = calc_free_energy(&histograms, &P);
dump_state(&histograms, &F, &F_prev, &P, &free_energy);
if iteration == cfg.max_iterations {
println!("!!!!! WHAM not converged! (max iterations reached) !!!!!");
}
io::write_results(&cfg.output, &histograms, &free_energy, &P)?;
Ok(())
}
// get average difference between two bias offset sets // get average difference between two bias offset sets
fn diff_avg(F: &[f64], F_prev: &[f64]) -> f64 { fn diff_avg(F: &[f64], F_prev: &[f64]) -> f64 {
let mut F_sum: f64 = 0.0; let mut F_sum: f64 = 0.0;
@@ -139,95 +169,34 @@ fn diff_avg(F: &[f64], F_prev: &[f64]) -> f64 {
F_sum / F.len() as f64 F_sum / F.len() as f64
} }
// calculate the normalized free energy from probability values
// calculate the normalized free energy from normalized probability values fn calc_free_energy(ds: &Dataset, P: &[f64]) -> Vec<f64> {
fn free_energy(ds: &Dataset, P: &[f64], A: &mut [f64]) { let mut minimum = f64::MAX;
let mut bin_min = f64::MAX; let mut free_energy: Vec<f64> = P.iter()
.map(|p| {
// Free energy calculation -ds.kT * p.ln()
for bin in 0..ds.num_bins { })
A[bin] = -ds.kT*P[bin].ln(); .inspect(|free_e| {
if A[bin] < bin_min { if free_e < &minimum {
bin_min = A[bin]; minimum = *free_e;
} }
})
.collect();
for e in free_energy.iter_mut() {
*e -= minimum
}
free_energy
} }
// Make A relative to minimum // TODO print nice headers for N dimensions
for bin in 0..ds.num_bins {
A[bin] -= bin_min;
}
}
pub fn run(cfg: &Config) -> Result<(), Box<Error>>{
println!("Supplied WHAM options: {}", &cfg);
// read input data into the histograms object
println!("Reading input files.");
let histograms = io::read_data(&cfg) // TODO nicer error handling for this
.expect("No datapoints in histogram boundaries.");
println!("{}",&histograms);
// allocate only once for better performance
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;
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);
// 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 {
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
if iteration % 100 == 0 {
free_energy(&histograms, &mut P, &mut A);
dump_state(&histograms, &F, &F_prev, &P, &A);
}
}
// Normalize P
let P_sum: f64 = P.iter().sum();
P.iter_mut().map(|p| *p /= P_sum).count();
// final free energy calculation and state dump
println!("Finished. Dumping final PMF");
free_energy(&histograms, &mut P, &mut A);
dump_state(&histograms, &F, &F_prev, &P, &A);
if iteration == cfg.max_iterations {
println!("!!!!! WHAM not converged! (max iterations reached) !!!!!");
}
io::write_results(&cfg.output, &histograms, &A, &P)?;
Ok(())
}
fn dump_state(ds: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64], A: &[f64]) { fn dump_state(ds: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64], A: &[f64]) {
let out = std::io::stdout(); let out = std::io::stdout();
let mut lock = out.lock(); let mut lock = out.lock();
writeln!(lock, "# PMF"); writeln!(lock, "# PMF");
writeln!(lock, "#x\t\tFree Energy\t\tP(x)"); writeln!(lock, "#x\t\tFree Energy\t\tP(x)");
for bin in 0..ds.num_bins { for bin in 0..ds.num_bins {
let x = ds.get_coords_for_bin(bin)[0]; // TODO let x = ds.get_coords_for_bin(bin)[0];
writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}", x, A[bin], P[bin]); writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}", x, A[bin], P[bin]);
} }
writeln!(lock, "# Bias offsets"); writeln!(lock, "# Bias offsets");