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

WHAM works in N dimensions

Browse Source
This commit is contained in:
Daniel Bauer
2018-10-14 11:36:10 +02:00
parent 29d2dd2b03
commit 59a6cc072d
3 changed files with 138 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
src/histogram.rs
View File

@@ -4,33 +4,16 @@ use std::cell::RefCell;
// One histogram // One histogram
#[derive(Debug)] #[derive(Debug)]
pub struct Histogram { pub struct Histogram {
// offset of this histogram bins from the global histogram
first: usize,
// offset of the last element of the histogram. TODO required?
last: usize,
// total number of data points stored in the histogram // total number of data points stored in the histogram
pub num_points: u32, pub num_points: u32,
// histogram bins // histogram bins
bins: Vec<f64> pub bins: Vec<f64>
} }
impl Histogram { impl Histogram {
pub fn new(first: usize, last: usize, num_points: u32, bins: Vec<f64>) -> Histogram { pub fn new(num_points: u32, bins: Vec<f64>) -> Histogram {
assert_eq!(last-first+1, bins.len(), "histogram length does not match first/last."); Histogram {num_points, bins}
Histogram {first, last, num_points, bins}
}
// Returns the value of a bin if the bin is present in this
// histogram
pub fn get_bin_count(&self, bin: usize) -> Option<f64> {
if bin < self.first || bin > self.last {
None
} else {
Some(self.bins[bin-self.first])
}
} }
} }
@@ -40,17 +23,20 @@ pub struct Dataset {
// number of histogram windows (number of simulations) // number of histogram windows (number of simulations)
pub num_windows: usize, pub num_windows: usize,
// number of global histogram bins // total number of bins
pub num_bins: usize, pub num_bins: usize,
// min value of the histogram // number of bins in each dimension
hist_min: f64, pub dimens_lengths: Vec<usize>,
// max value of the histogram // min values of the histogram in each dimension
hist_max: f64, hist_min: Vec<f64>,
// width of a bin in unit of x // max values of the histogram in each dimension
bin_width: f64, hist_max: Vec<f64>,
// width of a bin in unit of its dimension
bin_width: Vec<f64>,
// value of kT // value of kT
pub kT: f64, pub kT: f64,
@@ -72,16 +58,17 @@ pub struct Dataset {
} }
impl Dataset { impl Dataset {
pub fn new(num_bins: usize, bin_width: f64, hist_min: f64, hist_max: 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: RefCell<Vec<Option<f64>>> = RefCell::new(vec![None; num_bins*num_windows]); let bias: RefCell<Vec<Option<f64>>> = RefCell::new(vec![None; num_bins*num_windows]);
Dataset{ Dataset{
num_windows, num_windows,
num_bins, num_bins,
dimens_lengths,
bin_width, bin_width,
hist_min, hist_min,
hist_max, hist_max,
kT, kT,
histograms, histograms,
cyclic, cyclic,
@@ -91,36 +78,60 @@ impl Dataset {
} }
} }
fn expand_index(&self, bin: usize, lengths: &Vec<usize>) -> Vec<usize> {
let mut tmp = bin;
let mut idx = vec![0; lengths.len()];
for dimen in (1..lengths.len()).rev() {
let denom = lengths.iter().take(dimen).fold(1, |s,&x| s*x);
idx[dimen] = tmp / denom;
tmp = tmp % denom;
}
idx[0] = tmp;
idx
}
// get center x value for a bin
pub fn get_coords_for_bin(&self, bin: usize) -> Vec<f64> {
self.expand_index(bin, &self.dimens_lengths).iter().enumerate().map(|(i, dimen_bin)| {
self.hist_min[i] + self.bin_width[i]*(*dimen_bin as f64 + 0.5)
}).collect()
}
// Harmonic bias calculation: bias = 0.5*k(dx)^2 // Harmonic bias calculation: bias = 0.5*k(dx)^2
// if cyclic is true, lowest and highest bins are assumed to be // if cyclic is true, lowest and highest bins are assumed to be
// neighbors // neighbors
pub fn calc_bias(&self, bin: usize, window: usize) -> f64 { pub fn calc_bias(&self, bin: usize, window: usize) -> f64 {
let ndx = bin + (self.num_bins*window); let ndx = window * self.num_bins + bin;
let mut cache = self.bias.borrow_mut(); let mut cache = self.bias.borrow_mut();
match cache[ndx] { match cache[ndx] {
Some(val) => val, Some(val) => val,
None => { None => {
let x = self.get_x_for_bin(bin); // TODO optimize this part!
let mut dx = (x-self.bias_pos[window]).abs(); let dimens = self.hist_min.len();
if self.cyclic {
let hist_len = self.hist_max-self.hist_min; let bias_ndx: Vec<usize> = (0..dimens)
if dx > 0.5*hist_len { .map(|dimen| { window * dimens + dimen }).collect();
dx -= hist_len; let coord = self.get_coords_for_bin(bin);
let bias_fc: Vec<f64> = bias_ndx.iter().map(|ndx| { self.bias_fc[*ndx] }).collect();
let bias_pos: Vec<f64> = bias_ndx.iter().map(|ndx| { self.bias_pos[*ndx] }).collect();
let mut bias_sum = 0.0;
for i in 0..dimens {
let mut dist = (coord[i] - bias_pos[i]).abs();
if self.cyclic {
let hist_len = self.hist_max[i] - self.hist_min[i];
if dist > 0.5 * hist_len {
dist -= hist_len;
}
} }
bias_sum += 0.5 * bias_fc[i] * dist * dist
} }
let bias = 0.5*self.bias_fc[window]*dx*dx; cache[ndx] = Some(bias_sum);
cache[ndx] = Some(bias); bias_sum
bias
} }
} }
} }
// get center x value for a bin
pub fn get_x_for_bin(&self, bin: usize) -> f64 {
self.hist_min + self.bin_width * ((bin as f64) + 0.5)
}
} }
impl fmt::Display for Dataset { impl fmt::Display for Dataset {
@@ -140,8 +151,6 @@ mod tests {
fn build_hist() -> Histogram { fn build_hist() -> Histogram {
Histogram::new( Histogram::new(
5, // first
9, // last
22, // num_points 22, // num_points
vec![1.0, 1.0, 3.0, 5.0, 12.0] // bins vec![1.0, 1.0, 3.0, 5.0, 12.0] // bins
) )

126
src/io.rs
View File

@@ -32,15 +32,24 @@ pub fn read_data(cfg: &Config) -> Option<Dataset> {
let mut bias_pos: Vec<f64> = Vec::new(); let mut bias_pos: Vec<f64> = Vec::new();
let mut bias_fc: Vec<f64> = Vec::new(); let mut bias_fc: Vec<f64> = Vec::new();
let mut histograms: Vec<Histogram> = Vec::new(); let mut histograms: Vec<Histogram> = Vec::new();
let kT = cfg.temperature * k_B; let kT = cfg.temperature * k_B;
let bin_width: Vec<f64> = (0..cfg.dimens).map(|idx| {
(cfg.hist_max[idx] - cfg.hist_min[idx])/(cfg.num_bins[idx] as f64)
}).collect();
let num_bins = cfg.num_bins.iter().fold(1, |state, &bins| state*bins);
let dimens_length = cfg.num_bins.clone();
let f = File::open(&cfg.metadata_file).unwrap_or_else(|x| { let f = File::open(&cfg.metadata_file).unwrap_or_else(|x| {
eprintln!("Failed to read metadata from {}. {}", &cfg.metadata_file, x); eprintln!("Failed to read metadata from {}. {}", &cfg.metadata_file, x);
process::exit(1) process::exit(1)
}); });
let buf = BufReader::new(&f); let buf = BufReader::new(&f);
// read each metadata file line and parse it
for l in buf.lines() { for l in buf.lines() {
let line = l.unwrap(); let line = l.unwrap();
// skip comments and empty lines // skip comments and empty lines
if line.starts_with("#") || line.len() == 0 { if line.starts_with("#") || line.len() == 0 {
continue; continue;
@@ -84,15 +93,36 @@ pub fn read_data(cfg: &Config) -> Option<Dataset> {
} }
if histograms.len() > 0 { if histograms.len() > 0 {
let bin_width: Vec<f64> = (0..cfg.dimens).map(|idx| { Some(Dataset::new(num_bins, dimens_length, bin_width, cfg.hist_min.clone(), cfg.hist_max.clone(), bias_pos, bias_fc, kT, histograms, cfg.cyclic))
(cfg.hist_max[idx] - cfg.hist_min[idx])/(cfg.num_bins[idx] as f64)
}).collect();
Some(Dataset::new(cfg.num_bins[0], bin_width[0], cfg.hist_min[0], cfg.hist_max[0], bias_pos, bias_fc, kT, histograms, cfg.cyclic))
} else { } else {
None None
} }
} }
// transforms a multidimensional index into a one dimensional index
// indeces: multidimensional indeces
// lengths: length of the matrix in each dimension
// returns an index if the matrix is flattened to a one dimensional vector
// example for 3 dimensions N,M,O: idx = i_O + l_O*l_M*i_M + l_O*l_M*l_N*i_N
fn flat_index(indeces: &Vec<usize>, lengths: &Vec<usize>) -> usize {
let mut idx = 0;
for i in 0..indeces.len() {
idx += indeces[i]*lengths[0..i].iter()
.fold(1, |state, &l| { state * l });
}
idx
}
// returns true if the values are inside the histogram boundaries defined by cfg
fn is_in_hist_boundaries(values: &Vec<f64>, cfg: &Config) -> bool {
for dimen in 0..cfg.dimens {
if values[dimen] < cfg.hist_min[dimen] || values[dimen] > cfg.hist_max[dimen] {
return false
}
}
true
}
// parse a timeseries file into a histogram // parse a timeseries file into a histogram
fn read_window_file(window_file: &str, cfg: &Config) -> Option<Histogram> { fn read_window_file(window_file: &str, cfg: &Config) -> Option<Histogram> {
let f = File::open(window_file).unwrap_or_else(|x| { let f = File::open(window_file).unwrap_or_else(|x| {
@@ -100,63 +130,61 @@ fn read_window_file(window_file: &str, cfg: &Config) -> Option<Histogram> {
process::exit(1) process::exit(1)
}); });
let buf = BufReader::new(&f); let buf = BufReader::new(&f);
let mut global_hist = vec![0.0; cfg.num_bins[0]]; // total number of bins is the product of all dimensions length
let bin_width = (cfg.hist_max[0] - cfg.hist_min[0])/(cfg.num_bins[0] as f64); let total_bins = cfg.num_bins.iter().fold(1, |s, &x| { s*x });
let mut hist = vec![0.0; total_bins];
// bin width for each dimension: (max-min)/bins
let bin_width: Vec<f64> = (0..cfg.dimens).map(|idx| {
(cfg.hist_max[idx] - cfg.hist_min[idx])/(cfg.num_bins[idx] as f64)
}).collect();
// read and parse each timeseries line
for l in buf.lines() { for l in buf.lines() {
let line = l.unwrap(); let line = l.unwrap();
// skip comments and empty lines // skip comments and empty lines
if line.starts_with("#") || line.starts_with("@") || line.len() == 0 { if line.starts_with("#") || line.starts_with("@") || line.len() == 0 {
continue; continue;
} }
let (_, x) = scan_fmt!(&line, "{} {}", f64, f64); let mut split = line.split_whitespace();
split.next(); // skip time/step column
match x {
Some(x) => {
if x > cfg.hist_min[0] && x < cfg.hist_max[0] {
let bin_ndx = ((x-cfg.hist_min[0]) / bin_width) as usize;
global_hist[bin_ndx] += 1.0;
}
}
None => {
eprintln!("{}, Failed to read datapoint from line: {}", &window_file, &line);
process::exit(1);
}
}
}
let mut max_bin: usize = 0; let values: Vec<f64> = (0..cfg.dimens).collect::<Vec<usize>>().iter().map(|_| {
let mut min_bin: usize =(cfg.num_bins[0]-1) as usize; split.next().unwrap().parse::<f64>().unwrap()
for bin in 0..global_hist.len() { }).collect();
if global_hist[bin] != 0.0 && bin > max_bin {
max_bin = bin; if is_in_hist_boundaries(&values, cfg) {
} let bin_indeces = (0..cfg.dimens).map(|dimen: usize| {
if global_hist[bin] != 0.0 && bin < min_bin { let val = values[dimen];
min_bin = bin; ((val-cfg.hist_min[dimen]) / bin_width[dimen]) as usize
}).collect();
let index = flat_index(&bin_indeces, &cfg.num_bins);
hist[index] += 1.0;
} }
} }
if (max_bin == min_bin && global_hist[max_bin] == 0.0) || max_bin < min_bin { let num_points: f64 = hist.iter().sum();
None // zero length histogram if num_points == 0.0 {
} else { return None
// trim global hist to save memory
global_hist.truncate(max_bin+1);
global_hist.drain(..min_bin);
let num_points: f64 = global_hist.iter().sum();
Some(Histogram::new(min_bin, max_bin, num_points as u32, global_hist))
} }
Some(Histogram::new(num_points as u32, hist))
} }
// TODO multidimensional output
pub fn write_results(out_file: &str, ds: &Dataset, free: &Vec<f64>, prob: &Vec<f64>) -> Result<(), Box<Error>> { pub fn write_results(out_file: &str, ds: &Dataset, free: &Vec<f64>, prob: &Vec<f64>) -> Result<(), Box<Error>> {
let mut output = File::create(out_file)?; let mut output = File::create(out_file)?;
writeln!(output, "#{:8}\t{:8}\t{:8}", "x", "Free Energy", "Probability"); writeln!(output, "#{}\t{}\t{}", "x", "Free Energy", "Probability"); // TODO better format (coord1, coord2..)
for bin in 0..free.len() { for bin in 0..free.len() {
let x = ds.get_x_for_bin(bin); let coords = ds.get_coords_for_bin(bin);
writeln!(output, "{:8.6}\t{:8.6}\t{:8.6}", x, free[bin], prob[bin])?; let coords_str: String = coords.iter().map(|c| {format!("{:8.6}", c)})
} .collect::<Vec<String>>().join("\t");
Ok(()) writeln!(output, "{}\t{:8.6}\t{:8.6}", coords_str, free[bin], prob[bin])?;
}
Ok(())
} }
#[cfg(test)] #[cfg(test)]
@@ -188,9 +216,9 @@ mod tests {
// assert_eq!(1, h.first); // assert_eq!(1, h.first);
// assert_eq!(6, h.last); // assert_eq!(6, h.last);
assert_eq!(11, h.num_points); assert_eq!(11, h.num_points);
assert_eq!(2.0, h.get_bin_count(1).unwrap()); assert_eq!(2.0, h.bins[1]);
assert_eq!(2.0, h.get_bin_count(2).unwrap()); assert_eq!(2.0, h.bins[2]);
assert_eq!(1.0, h.get_bin_count(6).unwrap()); assert_eq!(1.0, h.bins[6]);
} }

13
src/lib.rs
View File

@@ -1,8 +1,5 @@
#![allow(non_snake_case)] #![allow(non_snake_case)]
#[macro_use]
extern crate scan_fmt;
pub mod io; pub mod io;
pub mod histogram; pub mod histogram;
@@ -55,9 +52,7 @@ fn calc_bin_probability(bin: usize, ds: &Dataset, F: &Vec<f64>) -> f64 {
let mut bin_count: f64 = 0.0; let mut bin_count: f64 = 0.0;
for window in 0..ds.num_windows { for window in 0..ds.num_windows {
let h: &Histogram = &ds.histograms[window]; let h: &Histogram = &ds.histograms[window];
if let Some(count) = h.get_bin_count(bin) { bin_count += h.bins[bin];
bin_count += count;
}
let bias = ds.calc_bias(bin, window); let bias = ds.calc_bias(bin, window);
let bias_offset = ((F[window] - bias) / ds.kT).exp(); let bias_offset = ((F[window] - bias) / ds.kT).exp();
denom_sum += (h.num_points as f64) * bias_offset; denom_sum += (h.num_points as f64) * bias_offset;
@@ -245,7 +240,7 @@ fn dump_state(ds: &Dataset, F: &Vec<f64>, F_prev: &Vec<f64>, P: &Vec<f64>, A: &V
println!("# PMF"); println!("# PMF");
println!("#x\t\tFree Energy\t\tP(x)"); println!("#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_x_for_bin(bin); let x = ds.get_coords_for_bin(bin)[0]; // TODO
println!("{:9.5}\t{:9.5}\t{:9.5}", x, A[bin], P[bin]); println!("{:9.5}\t{:9.5}\t{:9.5}", x, A[bin], P[bin]);
} }
println!("# Bias offsets"); println!("# Bias offsets");
@@ -274,8 +269,8 @@ mod tests {
} }
fn create_test_ds() -> Dataset { fn create_test_ds() -> Dataset {
let h1 = Histogram::new(0, 2, 10, vec![3.0, 4.0, 3.0]); 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(0, 3, 20, vec![3.0, 2.0, 5.0, 10.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, 1.0, 0.0, 4.0, vec![1.0, 2.0], vec![10.0, 10.0], 2.479, vec![h1, h2], false) Dataset::new(4, 1.0, 0.0, 4.0, vec![1.0, 2.0], vec![10.0, 10.0], 2.479, vec![h1, h2], false)
} }