diff --git a/src/lib.rs b/src/lib.rs index 1a82b21..333bbaa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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) // 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; - for (window, h) in ds.histograms.iter().enumerate() { - let bias = ds.calc_bias(bin, window); + for (window, h) in dataset.histograms.iter().enumerate() { + let bias = dataset.calc_bias(bin, 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. // 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 { - let f: f64 = (0..ds.num_bins).zip(P.iter()) // zip bins and P +fn calc_window_F(window: usize, dataset: &Dataset, P: &[f64]) -> f64 { + let f: f64 = (0..dataset.num_bins).zip(P.iter()) // zip bins and P .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 }).sum(); 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 // new bias offsets F based on previous bias offsets F_prev. This updates // 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 // estimage probabilities based on previous offsets (F_prev) - for bin in 0..ds.num_bins { - P[bin] = calc_bin_probability(bin, ds, F_prev); + for bin in 0..dataset.num_bins { + P[bin] = calc_bin_probability(bin, dataset, F_prev); } // evaluate second WHAM equation for each window to // estimate new bias offsets from propabilities - for window in 0..ds.num_windows { - F[window] = calc_window_F(window, ds, P); + for window in 0..dataset.num_windows { + F[window] = calc_window_F(window, dataset, P); } } -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 histograms = io::read_data(&cfg).chain_err(|| "Failed to create histogram.")?; - println!("{}",&histograms); - - // allocate required vectors. - let mut P: Vec = vec![f64::NAN; histograms.num_bins]; // bin probability - let mut F: Vec = vec![1.0; histograms.num_windows]; // bias offset exp(F/kT) - let mut F_prev: Vec = vec![f64::NAN; histograms.num_windows]; // previous bias offset - let mut F_tmp: Vec = vec![f64::NAN; histograms.num_windows]; // temp storage for F +pub fn perform_wham(cfg: &Config, dataset: &Dataset) -> Result<(Vec, Vec, Vec)> { + // allocate required vectors. + let mut P: Vec = vec![f64::NAN; dataset.num_bins]; // bin probability + let mut F: Vec = vec![1.0; dataset.num_windows]; // bias offset exp(F/kT) + let mut F_prev: Vec = vec![f64::NAN; dataset.num_windows]; // previous bias offset + let mut F_tmp: Vec = vec![f64::NAN; dataset.num_windows]; // temp storage for F let mut iteration = 0; let mut converged = false; @@ -114,7 +107,7 @@ pub fn run(cfg: &Config) -> Result<()>{ F_prev.copy_from_slice(&F); // 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 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 // its overwritten for the next iteration. F_tmp.copy_from_slice(&F); - for f in F.iter_mut() { *f = -histograms.kT * f.ln() } - for f in F_prev.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 = -dataset.kT * f.ln() } converged = is_converged(&F_prev, &F, cfg.tolerance); 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(); 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) !!!!!"); + 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")?; Ok(()) @@ -167,11 +173,11 @@ fn diff_avg(F: &[f64], F_prev: &[f64]) -> f64 { } // calculate the normalized free energy from probability values -fn calc_free_energy(ds: &Dataset, P: &[f64]) -> Vec { +fn calc_free_energy(dataset: &Dataset, P: &[f64]) -> Vec { let mut minimum = f64::MAX; let mut free_energy: Vec = P.iter() .map(|p| { - -ds.kT * p.ln() + -dataset.kT * p.ln() }) .inspect(|free_e| { if free_e < &minimum { @@ -186,17 +192,17 @@ fn calc_free_energy(ds: &Dataset, P: &[f64]) -> Vec { 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 mut lock = out.lock(); writeln!(lock, "# PMF"); 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, "# Bias offsets"); 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()); } } @@ -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 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], @@ -237,41 +243,41 @@ mod tests { #[test] fn calc_bin_probability() { - let ds = create_test_ds(); - let F = vec![1.0; ds.num_bins] ; + let dataset = create_test_dataset(); + let F = vec![1.0; dataset.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); + for b in 0..dataset.num_bins { + let p = super::calc_bin_probability(b, &dataset, &F); assert_delta!(expected[b], p, 0.0000001); } } #[test] 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 expected = vec!(15.927477169990633, 15.927477169990633); - for window in 0..ds.num_windows { - let F = super::calc_window_F(window, &ds, &probability); + for window in 0..dataset.num_windows { + let F = super::calc_window_F(window, &dataset, &probability); assert_delta!(expected[window], F, 0.0000001); } } #[test] fn perform_wham_iteration() { - let ds = create_test_ds(); - 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 dataset = create_test_dataset(); + let prev_F = vec![1.0; dataset.num_windows]; + let mut F = vec![f64::NAN; dataset.num_windows]; + let mut P = vec![f64::NAN; dataset.num_bins]; + super::perform_wham_iteration(&dataset, &prev_F, &mut F, &mut P); 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 { + for bin in 0..dataset.num_bins { 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) }