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

remove some code smell

Browse Source
This commit is contained in:
Daniel Bauer
2020-03-09 10:05:36 +01:00
parent 133234706e
commit 5390b089ac
5 changed files with 92 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/error_analysis.rs
View File

@@ -20,7 +20,7 @@ fn generate_random_weights(num_windows: usize, rng: &mut StdRng) -> Vec<f64> {
for i in 0..num_windows { for i in 0..num_windows {
weights[i] = rnds[i+1] - rnds[i] weights[i] = rnds[i+1] - rnds[i]
} }
return weights weights
} }
// Generate a random weighted dataset from the given dataset by changing the weights // Generate a random weighted dataset from the given dataset by changing the weights

35
src/histogram.rs
View File

@@ -61,7 +61,9 @@ pub struct Dataset {
impl Dataset { impl Dataset {
pub fn new(num_bins: usize, dimens_lengths: Vec<usize>, bin_width: Vec<f64>, hist_min: Vec<f64>, hist_max: Vec<f64>, bias_pos: Vec<f64>, bias_fc: Vec<f64>, kT: f64, histograms: Vec<Histogram>, cyclic: bool) -> Dataset { pub fn new(num_bins: usize, dimens_lengths: Vec<usize>, bin_width: Vec<f64>,
hist_min: Vec<f64>, hist_max: Vec<f64>, bias_pos: Vec<f64>,
bias_fc: Vec<f64>, kT: f64, histograms: Vec<Histogram>, cyclic: bool) -> Dataset {
let num_windows = histograms.len(); let num_windows = histograms.len();
let bias: Vec<f64> = vec![0.0; num_bins*num_windows]; let bias: Vec<f64> = vec![0.0; num_bins*num_windows];
let weights = vec![1.0; num_windows]; let weights = vec![1.0; num_windows];
@@ -92,7 +94,7 @@ impl Dataset {
pub fn new_weighted(ds: Dataset, weights: Vec<f64>) -> Dataset { pub fn new_weighted(ds: Dataset, weights: Vec<f64>) -> Dataset {
Dataset { Dataset {
weights: weights, weights,
..ds ..ds
} }
} }
@@ -107,7 +109,7 @@ impl Dataset {
for dimen in (1..lengths.len()).rev() { for dimen in (1..lengths.len()).rev() {
let denom = lengths.iter().take(dimen).fold(1, |s,&x| s*x); let denom = lengths.iter().take(dimen).fold(1, |s,&x| s*x);
idx[dimen] = tmp / denom; idx[dimen] = tmp / denom;
tmp = tmp % denom; tmp %= denom;
} }
idx[0] = tmp; idx[0] = tmp;
idx idx
@@ -151,8 +153,7 @@ impl Dataset {
// store exp(U/kT) for better performance // store exp(U/kT) for better performance
bias_sum += 0.5 * bias_fc[i] * dist * dist bias_sum += 0.5 * bias_fc[i] * dist * dist
} }
let bias_sum = (-bias_sum/self.kT).exp(); (-bias_sum/self.kT).exp()
bias_sum
} }
} }
@@ -205,13 +206,13 @@ mod tests {
let ds = build_hist_set(); // k = 10 let ds = build_hist_set(); // k = 10
// 3th element -> x=3.5, x0=3.5 // 3th element -> x=3.5, x0=3.5
assert_delta!(0.134722337796, ds.calc_bias(3, 0), 0.00000001); assert_delta!(0.134_722_337_796, ds.calc_bias(3, 0), 0.000_000_01);
// 8th element -> x=8.5, x0=3.5 // 8th element -> x=8.5, x0=3.5
assert_delta!(1.0, ds.calc_bias(4,0), 0.00000001); assert_delta!(1.0, ds.calc_bias(4,0), 0.000_000_01);
// 1st element -> x=0.5, x0=3.5. non-cyclic! // 1st element -> x=0.5, x0=3.5. non-cyclic!
assert_delta!(0.0, ds.calc_bias(0,0), 0.0000001); assert_delta!(0.0, ds.calc_bias(0,0), 0.000_000_1);
} }
#[test] #[test]
@@ -220,18 +221,18 @@ mod tests {
ds.cyclic = true; ds.cyclic = true;
// 7th element -> x=3.5, x0=3.5 // 7th element -> x=3.5, x0=3.5
assert_delta!(0.134722337796, ds.calc_bias(3, 0), 0.00000001); assert_delta!(0.134_722_337_796, ds.calc_bias(3, 0), 0.000_000_01);
// 8th element -> x=4.5, x0=3.5 // 8th element -> x=4.5, x0=3.5
assert_delta!(1.0, ds.calc_bias(4, 0), 0.00000001); assert_delta!(1.0, ds.calc_bias(4, 0), 0.000_000_01);
// 1th element -> x=0.5, x0=3.5 // 1th element -> x=0.5, x0=3.5
// cyclic flag makes bin 0 neighboring bin 9, so the distance is actually 2 // cyclic flag makes bin 0 neighboring bin 9, so the distance is actually 2
assert_delta!(0.0000000000000117769, ds.calc_bias(0, 0), 0.00000001); assert_delta!(0.000_000_000_000_011_776_9, ds.calc_bias(0, 0), 0.000_000_01);
// 2nd element -> x=1.5, x0=3.5 // 2nd element -> x=1.5, x0=3.5
assert_delta!(0.00000001, ds.calc_bias(1, 0), 0.00000001); assert_delta!(0.000_000_01, ds.calc_bias(1, 0), 0.000_000_01);
} }
#[test] #[test]
@@ -258,10 +259,10 @@ mod tests {
vec![build_hist(), build_hist()], // hists vec![build_hist(), build_hist()], // hists
false // cyclic false // cyclic
); );
assert_delta!(2.0, ds.get_weighted_bin_count(0), 0.0000000001); assert_delta!(2.0, ds.get_weighted_bin_count(0), 0.000_000_000_1);
assert_delta!(2.0, ds.get_weighted_bin_count(1), 0.0000000001); assert_delta!(2.0, ds.get_weighted_bin_count(1), 0.000_000_000_1);
assert_delta!(6.0, ds.get_weighted_bin_count(2), 0.0000000001); assert_delta!(6.0, ds.get_weighted_bin_count(2), 0.000_000_000_1);
assert_delta!(10.0, ds.get_weighted_bin_count(3), 0.0000000001); assert_delta!(10.0, ds.get_weighted_bin_count(3), 0.000_000_000_1);
assert_delta!(24.0, ds.get_weighted_bin_count(4), 0.0000000001); assert_delta!(24.0, ds.get_weighted_bin_count(4), 0.000_000_000_1);
} }
} }

11
src/io.rs
View File

@@ -46,7 +46,7 @@ pub fn read_data(cfg: &Config) -> Result<Dataset> {
let line = l.chain_err(|| "Failed to read line")?; let line = l.chain_err(|| "Failed to read line")?;
// skip comments and empty lines // skip comments and empty lines
if line.starts_with("#") || line.len() == 0 { if line.starts_with('#') || line.is_empty() {
continue; continue;
} }
@@ -79,7 +79,7 @@ pub fn read_data(cfg: &Config) -> Result<Dataset> {
} }
} }
if histograms.len() > 0 { if !histograms.is_empty() {
Ok(Dataset::new(num_bins, dimens_length, bin_width, cfg.hist_min.clone(), cfg.hist_max.clone(), bias_pos, bias_fc, kT, histograms, cfg.cyclic)) Ok(Dataset::new(num_bins, dimens_length, bin_width, cfg.hist_min.clone(), cfg.hist_max.clone(), bias_pos, bias_fc, kT, histograms, cfg.cyclic))
} else { } else {
bail!("Histogram has no datapoints.") bail!("Histogram has no datapoints.")
@@ -174,14 +174,16 @@ fn read_window_file(window_file: &str, cfg: &Config) -> Result<Histogram> {
} }
// Write WHAM calculation results to out_file. // Write WHAM calculation results to out_file.
pub fn write_results(out_file: &str, ds: &Dataset, free: &Vec<f64>, free_std: &Vec<f64>, prob: &Vec<f64>, prob_std: &Vec<f64>) -> Result<()> { pub fn write_results(out_file: &str, ds: &Dataset, free: &[f64],
free_std: &[f64], prob: &[f64], prob_std: &[f64]) -> Result<()> {
let output = File::create(out_file) let output = File::create(out_file)
.chain_err(|| format!("Failed to create file with path {}", out_file))?; .chain_err(|| format!("Failed to create file with path {}", out_file))?;
let mut buf = BufWriter::new(output); let mut buf = BufWriter::new(output);
let header: String = (0..ds.dimens_lengths.len()).map(|d| format!("coord{}", d+1)) let header: String = (0..ds.dimens_lengths.len()).map(|d| format!("coord{}", d+1))
.collect::<Vec<String>>().join(" "); .collect::<Vec<String>>().join(" ");
writeln!(buf, "#{} {} {} {} {}", header, "Free Energy", "+/-", "Probability", "+/-"); writeln!(buf, "#{} {} {} {} {}",
header, "Free Energy", "+/-", "Probability", "+/-").unwrap();
for bin in 0..free.len() { for bin in 0..free.len() {
let coords = ds.get_coords_for_bin(bin); let coords = ds.get_coords_for_bin(bin);
@@ -200,6 +202,7 @@ mod tests {
fn cfg() -> Config { fn cfg() -> Config {
Config { Config {
metadata_file: "example/1d_cyclic/metadata.dat".to_string(), metadata_file: "example/1d_cyclic/metadata.dat".to_string(),
hist_min: vec![-3.14], hist_min: vec![-3.14],
hist_max: vec![3.14], hist_max: vec![3.14],
num_bins: vec![10], num_bins: vec![10],

102
src/lib.rs
View File

@@ -45,7 +45,9 @@ 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={:?} verbose={}, tolerance={}, iterations={}, temperature={}, cyclic={:?}, bootstrap={:?}, seed={:?}", write!(f, "Metadata={}, hist_min={:?}, hist_max={:?}, bins={:?},
verbose={}, tolerance={}, iterations={}, temperature={},
cyclic={:?}, bootstrap={:?}, seed={:?}",
self.metadata_file, self.hist_min, self.hist_max, self.num_bins, self.metadata_file, self.hist_min, self.hist_max, self.num_bins,
self.verbose, self.tolerance, self.max_iterations, self.temperature, self.verbose, self.tolerance, self.max_iterations, self.temperature,
self.cyclic, self.bootstrap, self.bootstrap_seed) self.cyclic, self.bootstrap, self.bootstrap_seed)
@@ -56,25 +58,34 @@ impl fmt::Display for Config {
// converged if the maximal difference for the calculated bias offsets is // converged if the maximal difference for the calculated bias offsets is
// smaller then a tolerance value. // 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()) // calculates abs diff between every old and new F and checks if any
// is larger than tolerance
!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 given bias offsets (F) // estimate the probability of a bin of the histogram set based on given bias
// This evaluates the first WHAM equation for each bin. // offsets (F). This evaluates the first WHAM equation for each bin:
// P(x) = \frac {\sum_{i=1}^N{n_i(x)}}
// {\sum_{i=1}^N{ N_i exp(\beta [F_i - U_{bias,i}(x)])}}
fn calc_bin_probability(bin: usize, dataset: &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;
let bin_count: f64 = dataset.get_weighted_bin_count(bin); let bin_count: f64 = dataset.get_weighted_bin_count(bin);
for (window, h) in dataset.histograms.iter().enumerate() { for (window, h) in dataset.histograms.iter().enumerate() {
let bias = dataset.get_bias(bin, window); let bias = dataset.get_bias(bin, window);
denom_sum += (dataset.weights[window] * h.num_points as f64) * bias * F[window]; denom_sum += (dataset.weights[window] * 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. // 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).
// exp(F/kT) is not required in intermediate steps so we save some time by not
// calculating it for every iteration.
// F_i = - 1/\beta ln[\sum_{X_{bins}}{P(x)exp(-\beta U_{bias,i}(x))}]
fn calc_window_F(window: usize, dataset: &Dataset, P: &[f64]) -> f64 { 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 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)| {
@@ -84,16 +95,18 @@ fn calc_window_F(window: usize, dataset: &Dataset, P: &[f64]) -> f64 {
1.0/f 1.0/f
} }
// One full WHAM iteration includes calculation of new probabilities P and // One full WHAM iteration: calculation of new probabilities P and new bias
// new bias offsets F based on previous bias offsets F_prev. This updates // offsets F based on previous bias offsets F_prev. This updates the values in
// the values in vectors F and P // vectors F and P.
fn perform_wham_iteration(dataset: &Dataset, F_prev: &[f64], F: &mut Vec<f64>, P: &mut Vec<f64>) { fn perform_wham_iteration(dataset: &Dataset, F_prev: &[f64], F: &mut Vec<f64>, P: &mut Vec<f64>) {
// Update P
// evaluate first WHAM equation for each bin to // evaluate first WHAM equation for each bin to
// estimage probabilities based on previous offsets (F_prev)) // estimate probabilities based on previous offsets (F_prev))
(0..dataset.num_bins).into_par_iter() (0..dataset.num_bins).into_par_iter()
.map(|bin| { calc_bin_probability(bin, dataset, F_prev) }) .map(|bin| { calc_bin_probability(bin, dataset, F_prev) })
.collect_into_vec(P); .collect_into_vec(P);
// Update F
// 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
(0..dataset.num_windows).into_par_iter() (0..dataset.num_windows).into_par_iter()
@@ -101,12 +114,20 @@ fn perform_wham_iteration(dataset: &Dataset, F_prev: &[f64], F: &mut Vec<f64>, P
.collect_into_vec(F); .collect_into_vec(F);
} }
pub fn perform_wham(cfg: &Config, dataset: &Dataset) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>)> { // Full WHAM calculation. Calls `perform_wham_iteration` until convergence
// criteria are met or max iterations reached.
pub fn perform_wham(cfg: &Config, dataset: &Dataset)
-> Result<(Vec<f64>, Vec<f64>, Vec<f64>)> {
// allocate required vectors. // allocate required vectors.
let mut P: Vec<f64> = vec![f64::NAN; dataset.num_bins]; // bin probability
let mut F: Vec<f64> = vec![1.0; dataset.num_windows]; // bias offset exp(F/kT) // bin probability
let mut F_prev: Vec<f64> = vec![f64::NAN; dataset.num_windows]; // previous bias offset let mut P: Vec<f64> = vec![f64::NAN; dataset.num_bins];
let mut F_tmp: Vec<f64> = vec![f64::NAN; dataset.num_windows]; // temp storage for F // bias offset exp(F/kT)
let mut F: Vec<f64> = vec![1.0; dataset.num_windows];
// previous bias offset
let mut F_prev: Vec<f64> = vec![f64::NAN; dataset.num_windows];
// temp storage for F
let mut F_tmp: Vec<f64> = vec![f64::NAN; dataset.num_windows];
let mut iteration = 0; let mut iteration = 0;
let mut converged = false; let mut converged = false;
@@ -118,14 +139,15 @@ pub fn perform_wham(cfg: &Config, dataset: &Dataset) -> Result<(Vec<f64>, Vec<f6
// store F values before the next iteration // store F values before the next iteration
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(&dataset, &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 {
// This backups exp(F/kT) in a temporary vector and calculates true F and F_prev for // This backups exp(F/kT) in a temporary vector and calculates
// convergence. Finally, F is restored. F_prev does not need to be restored because // true F and F_prev for convergence. Finally, F is restored.
// its overwritten for the next iteration. // F_prev does not need to be restored because 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 = -dataset.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() } for f in F_prev.iter_mut() { *f = -dataset.kT * f.ln() }
@@ -134,17 +156,13 @@ pub fn perform_wham(cfg: &Config, dataset: &Dataset) -> Result<(Vec<f64>, Vec<f6
println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F)); println!("Iteration {}: dF={}", &iteration, &diff_avg(&F_prev, &F));
F.copy_from_slice(&F_tmp); 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 // Normalize P to sum(P) = 1.0
let P_sum: f64 = P.iter().sum(); let P_sum: f64 = P.iter().sum();
P.iter_mut().map(|p| *p /= P_sum).count(); for p in P.iter_mut() {
*p /= P_sum;
}
if iteration == cfg.max_iterations { if iteration == cfg.max_iterations {
bail!("WHAM not converged! (max iterations reached)"); bail!("WHAM not converged! (max iterations reached)");
@@ -214,19 +232,23 @@ fn calc_free_energy(dataset: &Dataset, P: &[f64]) -> Vec<f64> {
free_energy free_energy
} }
fn dump_state(dataset: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64], P_std: &[f64], A: &[f64], A_std: &[f64]) { // Print the current WHAM iteration state. Dumps the PMF and associated vectors
fn dump_state(dataset: &Dataset, F: &[f64], F_prev: &[f64], P: &[f64],
P_std: &[f64], A: &[f64], A_std: &[f64]) {
// TODO fix output of F/F_prev // TODO fix output of F/F_prev
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").unwrap();
writeln!(lock, "#bin\t\tFree Energy\t\t+/-\t\tP(x)\t\t+/-"); writeln!(lock, "#bin\t\tFree Energy\t\t+/-\t\tP(x)\t\t+/-").unwrap();
for bin in 0..dataset.num_bins { for bin in 0..dataset.num_bins {
writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}\t{:9.5}\t{:9.5}", bin, A[bin], A_std[bin], P[bin], P_std[bin]); writeln!(lock, "{:9.5}\t{:9.5}\t{:9.5}\t{:9.5}\t{:9.5}",
bin, A[bin], A_std[bin], P[bin], P_std[bin]).unwrap();
} }
writeln!(lock, "# Bias offsets"); writeln!(lock, "# Bias offsets").unwrap();
writeln!(lock, "#Window\t\tF\t\tF_prev"); writeln!(lock, "#Window\t\tF\t\tF_prev").unwrap();
for window in 0..dataset.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()).unwrap();
} }
} }
@@ -268,11 +290,11 @@ mod tests {
fn calc_bin_probability() { fn calc_bin_probability() {
let dataset = create_test_dataset(); let dataset = create_test_dataset();
let F = vec![1.0; dataset.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.082_529_668_703_131_6, 40.923_558_470_974_93,
124226.70003377, 2308526035.5283747); 124_226.700_033_77, 2_308_526_035.528_374_7);
for b in 0..dataset.num_bins { for b in 0..dataset.num_bins {
let p = super::calc_bin_probability(b, &dataset, &F); let p = super::calc_bin_probability(b, &dataset, &F);
assert_delta!(expected[b], p, 0.0000001); assert_delta!(expected[b], p, 0.000_000_1);
} }
} }
@@ -280,10 +302,10 @@ mod tests {
fn calc_bias_offset() { fn calc_bias_offset() {
let dataset = create_test_dataset(); 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.927_477_169_990_633, 15.927_477_169_990_633);
for window in 0..dataset.num_windows { for window in 0..dataset.num_windows {
let F = super::calc_window_F(window, &dataset, &probability); let F = super::calc_window_F(window, &dataset, &probability);
assert_delta!(expected[window], F, 0.0000001); assert_delta!(expected[window], F, 0.000_000_1);
} }
} }
@@ -295,8 +317,8 @@ mod tests {
let mut P = vec![f64::NAN; dataset.num_bins]; let mut P = vec![f64::NAN; dataset.num_bins];
super::perform_wham_iteration(&dataset, &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.082_529_668_703_131_6, 40.923_558_470_974_93,
124226.70003377, 2308526035.5283747); 124_226.700_033_77, 2_308_526_035.528_374_7);
for bin in 0..dataset.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)
} }

2
tests/cli.rs
View File

@@ -2,8 +2,6 @@ mod command;
#[cfg(test)] #[cfg(test)]
mod integration { mod integration {
use std::process::Command;
use super::command::get_command; use super::command::get_command;
#[test] #[test]