mirror of
https://github.com/dnlbauer/WHAM.git
synced 2026-09-10 22:25:31 +00:00
WHAM works in N dimensions
This commit is contained in:
105
src/histogram.rs
105
src/histogram.rs
@@ -4,33 +4,16 @@ use std::cell::RefCell;
|
||||
// One histogram
|
||||
#[derive(Debug)]
|
||||
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
|
||||
pub num_points: u32,
|
||||
|
||||
// histogram bins
|
||||
bins: Vec<f64>
|
||||
pub bins: Vec<f64>
|
||||
}
|
||||
|
||||
impl Histogram {
|
||||
pub fn new(first: usize, last: usize, num_points: u32, bins: Vec<f64>) -> Histogram {
|
||||
assert_eq!(last-first+1, bins.len(), "histogram length does not match first/last.");
|
||||
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])
|
||||
}
|
||||
pub fn new(num_points: u32, bins: Vec<f64>) -> Histogram {
|
||||
Histogram {num_points, bins}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,17 +23,20 @@ pub struct Dataset {
|
||||
// number of histogram windows (number of simulations)
|
||||
pub num_windows: usize,
|
||||
|
||||
// number of global histogram bins
|
||||
// total number of bins
|
||||
pub num_bins: usize,
|
||||
|
||||
// min value of the histogram
|
||||
hist_min: f64,
|
||||
// number of bins in each dimension
|
||||
pub dimens_lengths: Vec<usize>,
|
||||
|
||||
// max value of the histogram
|
||||
hist_max: f64,
|
||||
// min values of the histogram in each dimension
|
||||
hist_min: Vec<f64>,
|
||||
|
||||
// width of a bin in unit of x
|
||||
bin_width: f64,
|
||||
// max values of the histogram in each dimension
|
||||
hist_max: Vec<f64>,
|
||||
|
||||
// width of a bin in unit of its dimension
|
||||
bin_width: Vec<f64>,
|
||||
|
||||
// value of kT
|
||||
pub kT: f64,
|
||||
@@ -72,16 +58,17 @@ pub struct 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 bias: RefCell<Vec<Option<f64>>> = RefCell::new(vec![None; num_bins*num_windows]);
|
||||
Dataset{
|
||||
num_windows,
|
||||
num_bins,
|
||||
dimens_lengths,
|
||||
bin_width,
|
||||
hist_min,
|
||||
hist_max,
|
||||
hist_max,
|
||||
kT,
|
||||
histograms,
|
||||
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
|
||||
// if cyclic is true, lowest and highest bins are assumed to be
|
||||
// neighbors
|
||||
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();
|
||||
match cache[ndx] {
|
||||
Some(val) => val,
|
||||
None => {
|
||||
let x = self.get_x_for_bin(bin);
|
||||
let mut dx = (x-self.bias_pos[window]).abs();
|
||||
if self.cyclic {
|
||||
let hist_len = self.hist_max-self.hist_min;
|
||||
if dx > 0.5*hist_len {
|
||||
dx -= hist_len;
|
||||
// TODO optimize this part!
|
||||
let dimens = self.hist_min.len();
|
||||
|
||||
let bias_ndx: Vec<usize> = (0..dimens)
|
||||
.map(|dimen| { window * dimens + dimen }).collect();
|
||||
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);
|
||||
bias
|
||||
cache[ndx] = Some(bias_sum);
|
||||
bias_sum
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
@@ -140,8 +151,6 @@ mod tests {
|
||||
|
||||
fn build_hist() -> Histogram {
|
||||
Histogram::new(
|
||||
5, // first
|
||||
9, // last
|
||||
22, // num_points
|
||||
vec![1.0, 1.0, 3.0, 5.0, 12.0] // bins
|
||||
)
|
||||
|
||||
126
src/io.rs
126
src/io.rs
@@ -32,15 +32,24 @@ pub fn read_data(cfg: &Config) -> Option<Dataset> {
|
||||
let mut bias_pos: Vec<f64> = Vec::new();
|
||||
let mut bias_fc: Vec<f64> = Vec::new();
|
||||
let mut histograms: Vec<Histogram> = Vec::new();
|
||||
|
||||
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| {
|
||||
eprintln!("Failed to read metadata from {}. {}", &cfg.metadata_file, x);
|
||||
process::exit(1)
|
||||
});
|
||||
});
|
||||
let buf = BufReader::new(&f);
|
||||
|
||||
// read each metadata file line and parse it
|
||||
for l in buf.lines() {
|
||||
let line = l.unwrap();
|
||||
|
||||
// skip comments and empty lines
|
||||
if line.starts_with("#") || line.len() == 0 {
|
||||
continue;
|
||||
@@ -84,15 +93,36 @@ pub fn read_data(cfg: &Config) -> Option<Dataset> {
|
||||
}
|
||||
|
||||
if histograms.len() > 0 {
|
||||
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();
|
||||
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))
|
||||
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))
|
||||
} else {
|
||||
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
|
||||
fn read_window_file(window_file: &str, cfg: &Config) -> Option<Histogram> {
|
||||
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)
|
||||
});
|
||||
let buf = BufReader::new(&f);
|
||||
|
||||
let mut global_hist = vec![0.0; cfg.num_bins[0]];
|
||||
let bin_width = (cfg.hist_max[0] - cfg.hist_min[0])/(cfg.num_bins[0] as f64);
|
||||
|
||||
// total number of bins is the product of all dimensions length
|
||||
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() {
|
||||
let line = l.unwrap();
|
||||
|
||||
// skip comments and empty lines
|
||||
if line.starts_with("#") || line.starts_with("@") || line.len() == 0 {
|
||||
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 mut min_bin: usize =(cfg.num_bins[0]-1) as usize;
|
||||
for bin in 0..global_hist.len() {
|
||||
if global_hist[bin] != 0.0 && bin > max_bin {
|
||||
max_bin = bin;
|
||||
}
|
||||
if global_hist[bin] != 0.0 && bin < min_bin {
|
||||
min_bin = bin;
|
||||
let values: Vec<f64> = (0..cfg.dimens).collect::<Vec<usize>>().iter().map(|_| {
|
||||
split.next().unwrap().parse::<f64>().unwrap()
|
||||
}).collect();
|
||||
|
||||
if is_in_hist_boundaries(&values, cfg) {
|
||||
let bin_indeces = (0..cfg.dimens).map(|dimen: usize| {
|
||||
let val = values[dimen];
|
||||
((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 {
|
||||
None // zero length histogram
|
||||
} else {
|
||||
// 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))
|
||||
let num_points: f64 = hist.iter().sum();
|
||||
if num_points == 0.0 {
|
||||
return None
|
||||
}
|
||||
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>> {
|
||||
let mut output = File::create(out_file)?;
|
||||
writeln!(output, "#{:8}\t{:8}\t{:8}", "x", "Free Energy", "Probability");
|
||||
for bin in 0..free.len() {
|
||||
let x = ds.get_x_for_bin(bin);
|
||||
writeln!(output, "{:8.6}\t{:8.6}\t{:8.6}", x, free[bin], prob[bin])?;
|
||||
}
|
||||
Ok(())
|
||||
let mut output = File::create(out_file)?;
|
||||
writeln!(output, "#{}\t{}\t{}", "x", "Free Energy", "Probability"); // TODO better format (coord1, coord2..)
|
||||
for bin in 0..free.len() {
|
||||
let coords = ds.get_coords_for_bin(bin);
|
||||
let coords_str: String = coords.iter().map(|c| {format!("{:8.6}", c)})
|
||||
.collect::<Vec<String>>().join("\t");
|
||||
writeln!(output, "{}\t{:8.6}\t{:8.6}", coords_str, free[bin], prob[bin])?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -188,9 +216,9 @@ mod tests {
|
||||
// assert_eq!(1, h.first);
|
||||
// assert_eq!(6, h.last);
|
||||
assert_eq!(11, h.num_points);
|
||||
assert_eq!(2.0, h.get_bin_count(1).unwrap());
|
||||
assert_eq!(2.0, h.get_bin_count(2).unwrap());
|
||||
assert_eq!(1.0, h.get_bin_count(6).unwrap());
|
||||
assert_eq!(2.0, h.bins[1]);
|
||||
assert_eq!(2.0, h.bins[2]);
|
||||
assert_eq!(1.0, h.bins[6]);
|
||||
}
|
||||
|
||||
|
||||
|
||||
13
src/lib.rs
13
src/lib.rs
@@ -1,8 +1,5 @@
|
||||
#![allow(non_snake_case)]
|
||||
|
||||
#[macro_use]
|
||||
extern crate scan_fmt;
|
||||
|
||||
pub mod io;
|
||||
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;
|
||||
for window in 0..ds.num_windows {
|
||||
let h: &Histogram = &ds.histograms[window];
|
||||
if let Some(count) = h.get_bin_count(bin) {
|
||||
bin_count += count;
|
||||
}
|
||||
bin_count += h.bins[bin];
|
||||
let bias = ds.calc_bias(bin, window);
|
||||
let bias_offset = ((F[window] - bias) / ds.kT).exp();
|
||||
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!("#x\t\tFree Energy\t\tP(x)");
|
||||
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!("# Bias offsets");
|
||||
@@ -274,8 +269,8 @@ mod tests {
|
||||
}
|
||||
|
||||
fn create_test_ds() -> Dataset {
|
||||
let h1 = Histogram::new(0, 2, 10, vec![3.0, 4.0, 3.0]);
|
||||
let h2 = Histogram::new(0, 3, 20, vec![3.0, 2.0, 5.0, 10.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(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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user