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

moved wham evaluation out of run function

Browse Source
This commit is contained in:
Daniel Bauer
2018-11-05 10:22:26 +01:00
parent d5df049151
commit 93cd0d1a71
1 changed files with 61 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
src/lib.rs
View File

@@ -52,21 +52,21 @@ fn is_converged(old_F: &[f64], new_F: &[f64], tolerance: f64) -> bool {
// estimate the probability of a bin of the histogram set based on given bias offsets (F) // 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, dataset: &Dataset, F: &[f64]) -> f64 {
let mut denom_sum: f64 = 0.0; let mut denom_sum: f64 = 0.0;
for (window, h) in ds.histograms.iter().enumerate() { for (window, h) in dataset.histograms.iter().enumerate() {
let bias = ds.calc_bias(bin, window); let bias = dataset.calc_bias(bin, window);
denom_sum += (h.num_points as f64) * bias * F[window]; denom_sum += (h.num_points as f64) * bias * F[window];
} }
ds.bin_count[bin] / denom_sum dataset.bin_count[bin] / 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 and returns exp(F/kT) // 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, dataset: &Dataset, P: &[f64]) -> f64 {
let f: f64 = (0..ds.num_bins).zip(P.iter()) // zip bins and P let f: f64 = (0..dataset.num_bins).zip(P.iter()) // zip bins and P
.map(|bin_and_prob: (usize, &f64)| { .map(|bin_and_prob: (usize, &f64)| {
let bias = ds.calc_bias(bin_and_prob.0, window); let bias = dataset.calc_bias(bin_and_prob.0, window);
bin_and_prob.1 * bias bin_and_prob.1 * bias
}).sum(); }).sum();
1.0/f 1.0/f
@@ -75,33 +75,26 @@ fn calc_window_F(window: usize, ds: &Dataset, P: &[f64]) -> f64 {
// One full WHAM iteration includes calculation of new probabilities P and // One full WHAM iteration includes calculation of new probabilities P and
// 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(dataset: &Dataset, F_prev: &[f64], F: &mut [f64], P: &mut [f64]) {
// 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..dataset.num_bins {
P[bin] = calc_bin_probability(bin, ds, F_prev); P[bin] = calc_bin_probability(bin, dataset, F_prev);
} }
// evaluate second WHAM equation for each window to // evaluate second WHAM equation for each window to
// estimate new bias offsets from propabilities // estimate new bias offsets from propabilities
for window in 0..ds.num_windows { for window in 0..dataset.num_windows {
F[window] = calc_window_F(window, ds, P); F[window] = calc_window_F(window, dataset, P);
} }
} }
pub fn run(cfg: &Config) -> Result<()>{ pub fn perform_wham(cfg: &Config, dataset: &Dataset) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>)> {
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).chain_err(|| "Failed to create histogram.")?;
println!("{}",&histograms);
// allocate required vectors. // allocate required vectors.
let mut P: Vec<f64> = vec![f64::NAN; histograms.num_bins]; // bin probability let mut P: Vec<f64> = vec![f64::NAN; dataset.num_bins]; // bin probability
let mut F: Vec<f64> = vec![1.0; histograms.num_windows]; // bias offset exp(F/kT) let mut F: Vec<f64> = vec![1.0; dataset.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_prev: Vec<f64> = vec![f64::NAN; dataset.num_windows]; // previous bias offset
let mut F_tmp: Vec<f64> = vec![f64::NAN; histograms.num_windows]; // temp storage for F let mut F_tmp: Vec<f64> = vec![f64::NAN; dataset.num_windows]; // temp storage for F
let mut iteration = 0; let mut iteration = 0;
let mut converged = false; let mut converged = false;
@@ -114,7 +107,7 @@ pub fn run(cfg: &Config) -> Result<()>{
F_prev.copy_from_slice(&F); F_prev.copy_from_slice(&F);
// perform wham iteration (this updates F and P) // perform wham iteration (this updates F and P)
perform_wham_iteration(&histograms, &F_prev, &mut F, &mut P); perform_wham_iteration(&dataset, &F_prev, &mut F, &mut P);
// convergence check // convergence check
if iteration % 10 == 0 { if iteration % 10 == 0 {
@@ -122,8 +115,8 @@ pub fn run(cfg: &Config) -> Result<()>{
// convergence. Finally, F is restored. F_prev does not need to be restored because // convergence. Finally, F is restored. F_prev does not need to be restored because
// its overwritten for the next iteration. // its overwritten for the next iteration.
F_tmp.copy_from_slice(&F); F_tmp.copy_from_slice(&F);
for f in F.iter_mut() { *f = -histograms.kT * f.ln() } for f in F.iter_mut() { *f = -dataset.kT * f.ln() }
for f in F_prev.iter_mut() { *f = -histograms.kT * f.ln() } for f in F_prev.iter_mut() { *f = -dataset.kT * f.ln() }
converged = is_converged(&F_prev, &F, cfg.tolerance); converged = is_converged(&F_prev, &F, cfg.tolerance);
println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F)); println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F));
@@ -141,16 +134,29 @@ pub fn run(cfg: &Config) -> Result<()>{
let P_sum: f64 = P.iter().sum(); let P_sum: f64 = P.iter().sum();
P.iter_mut().map(|p| *p /= P_sum).count(); 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 { if iteration == cfg.max_iterations {
println!("!!!!! WHAM not converged! (max iterations reached) !!!!!"); bail!("WHAM not converged! (max iterations reached)");
} }
io::write_results(&cfg.output, &histograms, &free_energy, &P) Ok((P, F, F_prev))
}
pub fn run(cfg: &Config) -> Result<()>{
println!("Supplied WHAM options: {}", &cfg);
println!("Reading input files.");
// TODO Better error handling with nice error messages instead of a panic!
let dataset = io::read_data(&cfg).chain_err(|| "Failed to create histogram.")?;
println!("{}", &dataset);
let (P, F, F_prev) = perform_wham(&cfg, &dataset)?;
// calculate free energy and dump state
println!("Finished. Dumping final PMF");
let free_energy = calc_free_energy(&dataset, &P);
dump_state(&dataset, &F, &F_prev, &P, &free_energy);
io::write_results(&cfg.output, &dataset, &free_energy, &P)
.chain_err(|| "Could not write results to output file")?; .chain_err(|| "Could not write results to output file")?;
Ok(()) Ok(())
@@ -167,11 +173,11 @@ fn diff_avg(F: &[f64], F_prev: &[f64]) -> f64 {
} }
// calculate the normalized free energy from probability values // calculate the normalized free energy from probability values
fn calc_free_energy(ds: &Dataset, P: &[f64]) -> Vec<f64> { fn calc_free_energy(dataset: &Dataset, P: &[f64]) -> Vec<f64> {
let mut minimum = f64::MAX; let mut minimum = f64::MAX;
let mut free_energy: Vec<f64> = P.iter() let mut free_energy: Vec<f64> = P.iter()
.map(|p| { .map(|p| {
-ds.kT * p.ln() -dataset.kT * p.ln()
}) })
.inspect(|free_e| { .inspect(|free_e| {
if free_e < &minimum { if free_e < &minimum {
@@ -186,17 +192,17 @@ fn calc_free_energy(ds: &Dataset, P: &[f64]) -> Vec<f64> {
free_energy free_energy
} }
fn dump_state(ds: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64], A: &[f64]) { fn dump_state(dataset: &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, "#bin\t\tFree Energy\t\tP(x)"); writeln!(lock, "#bin\t\tFree Energy\t\tP(x)");
for bin in 0..ds.num_bins { for bin in 0..dataset.num_bins {
writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}", bin, A[bin], P[bin]); writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}", bin, A[bin], P[bin]);
} }
writeln!(lock, "# Bias offsets"); writeln!(lock, "# Bias offsets");
writeln!(lock, "#Window\t\tF\t\tdF"); writeln!(lock, "#Window\t\tF\t\tdF");
for window in 0..ds.num_windows { for window in 0..dataset.num_windows {
writeln!(lock, "{}\t{:9.5}\t{:8.8}", window, F[window], (F[window]-F_prev[window]).abs()); writeln!(lock, "{}\t{:9.5}\t{:8.8}", window, F[window], (F[window]-F_prev[window]).abs());
} }
} }
@@ -215,7 +221,7 @@ mod tests {
} }
fn create_test_ds() -> Dataset { fn create_test_dataset() -> Dataset {
let h1 = Histogram::new(10, vec![0.0, 1.0, 1.0, 8.0, 0.0]); 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]); 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], Dataset::new(5, vec![5], vec![1.0], vec![0.0], vec![4.0],
@@ -237,41 +243,41 @@ mod tests {
#[test] #[test]
fn calc_bin_probability() { fn calc_bin_probability() {
let ds = create_test_ds(); let dataset = create_test_dataset();
let F = vec![1.0; ds.num_bins] ; let F = vec![1.0; dataset.num_bins] ;
let expected = vec!(0.0, 0.0825296687031316, 40.92355847097493, let expected = vec!(0.0, 0.0825296687031316, 40.92355847097493,
124226.70003377, 2308526035.5283747); 124226.70003377, 2308526035.5283747);
for b in 0..ds.num_bins { for b in 0..dataset.num_bins {
let p = super::calc_bin_probability(b, &ds, &F); let p = super::calc_bin_probability(b, &dataset, &F);
assert_delta!(expected[b], p, 0.0000001); assert_delta!(expected[b], p, 0.0000001);
} }
} }
#[test] #[test]
fn calc_bias_offset() { fn calc_bias_offset() {
let ds = create_test_ds(); let dataset = create_test_dataset();
let probability = vec!(0.0, 0.1, 0.2, 0.3, 0.4); let probability = vec!(0.0, 0.1, 0.2, 0.3, 0.4);
let expected = vec!(15.927477169990633, 15.927477169990633); let expected = vec!(15.927477169990633, 15.927477169990633);
for window in 0..ds.num_windows { for window in 0..dataset.num_windows {
let F = super::calc_window_F(window, &ds, &probability); let F = super::calc_window_F(window, &dataset, &probability);
assert_delta!(expected[window], F, 0.0000001); assert_delta!(expected[window], F, 0.0000001);
} }
} }
#[test] #[test]
fn perform_wham_iteration() { fn perform_wham_iteration() {
let ds = create_test_ds(); let dataset = create_test_dataset();
let prev_F = vec![1.0; ds.num_windows]; let prev_F = vec![1.0; dataset.num_windows];
let mut F = vec![f64::NAN; ds.num_windows]; let mut F = vec![f64::NAN; dataset.num_windows];
let mut P = vec![f64::NAN; ds.num_bins]; let mut P = vec![f64::NAN; dataset.num_bins];
super::perform_wham_iteration(&ds, &prev_F, &mut F, &mut P); super::perform_wham_iteration(&dataset, &prev_F, &mut F, &mut P);
let expected_F = vec!(1.0, 1.0); let expected_F = vec!(1.0, 1.0);
let expected_P = vec!(0.0, 0.0825296687031316, 40.92355847097493, let expected_P = vec!(0.0, 0.0825296687031316, 40.92355847097493,
124226.70003377, 2308526035.5283747); 124226.70003377, 2308526035.5283747);
for bin in 0..ds.num_bins { for bin in 0..dataset.num_bins {
assert_delta!(expected_P[bin], P[bin], 0.01) assert_delta!(expected_P[bin], P[bin], 0.01)
} }
for window in 0..ds.num_windows { for window in 0..dataset.num_windows {
assert_delta!(expected_F[window], F[window], 0.01) assert_delta!(expected_F[window], F[window], 0.01)
} }