Package csb :: Package statmech :: Module wham
[frames] | no frames]

Source Code for Module csb.statmech.wham

  1  """ 
  2  Estimate the free energy and density of states from tempered ensembles using 
  3  histogram re-weighting. 
  4  """ 
  5   
  6  import numpy 
  7   
  8  from csb.numeric import log, log_sum_exp 
  9  from csb.statistics import histogram_nd 
 10   
 11  from abc import abstractmethod, ABCMeta 
12 13 14 -class AbstractWHAM(object):
15 """ 16 Abstract base class 17 """ 18 __metaclass__ = ABCMeta 19
20 - def __init__(self, ensembles, raw_energies, n):
21 22 self._f = numpy.zeros(len(ensembles)) 23 self._e = raw_energies 24 self._n = n 25 self._L = [] 26 self._log_g = None 27 self._ensembles = ensembles
28
29 - def log_g(self, normalize=True):
30 """ 31 Return the Density of states (DOS). 32 33 @param normalize: Ensure that the density of states sums to one 34 @rtype: float 35 """ 36 if normalize: 37 return self._log_g - log_sum_exp(self._log_g) 38 else: 39 return self._log_g
40 41 @property
42 - def free_energies(self):
43 """ 44 Free energies 45 """ 46 return self._f
47
48 - def _stop_criterium(self, tol=1e-10):
49 """ 50 general stop criterium; if the relative difference between 51 sequential negative log likelihoods is less than a predefined 52 tolerance 53 54 @param tol: tolerance 55 @type tol: float 56 57 @rtype: boolean 58 """ 59 L = self._L 60 return tol is not None and len(L) > 1 and \ 61 abs((L[-2] - L[-1]) / (L[-2] + L[-1])) < tol
62 63 64 @abstractmethod
65 - def estimate(self, *params):
66 """ 67 Estimate the density of states 68 """ 69 pass
70 71 @abstractmethod
72 - def log_z(self, beta=1., ensembles=None):
73 """ 74 Compute the partition function for an ensemble at inverse temperature 75 beta or for a defined ensemble 76 77 @param beta: Inverse Temperature 78 @type beta: float or list 79 80 @param ensembles: List of ensembles for which the partition function should be evaluated 81 @type ensembles: List of ensembles 82 83 @rtype: float or array 84 """ 85 pass
86
87 88 -class WHAM(AbstractWHAM):
89 """ 90 Implementation of the original WHAM methods based on histograms. 91 """ 92
93 - def __init__(self, ensembles, raw_energies, n):
94 super(WHAM, self).__init__(ensembles, raw_energies, n) 95 96 self._ex = None 97 self._h = None
98
99 - def estimate(self, n_bins=100, n_iter=10000, tol=1e-10):
100 101 self._L = [] 102 h, e = histogram_nd(self._e, nbins=n_bins, normalize=False) 103 self._ex = e = numpy.array(e) 104 self._h = h 105 f = self._f 106 107 log_h = log(h) 108 log_g = h * 0.0 109 log_g -= log_sum_exp(log_g) 110 log_n = log(self._n) 111 112 e_ij = -numpy.squeeze(numpy.array([ensemble.energy(e) 113 for ensemble in self._ensembles])).T 114 115 for _i in range(n_iter): 116 117 ## update density of states 118 y = log_sum_exp(numpy.reshape((e_ij - f + log_n).T, 119 (len(f), -1)), 0) 120 log_g = log_h - numpy.reshape(y, log_g.shape) 121 log_g -= log_sum_exp(log_g) 122 123 ## update free energies 124 f = log_sum_exp(numpy.reshape(e_ij.T + log_g.flatten(), 125 (len(f), -1)).T, 0) 126 self._L.append((self._n * f).sum() - (h * log_g).sum()) 127 128 self._log_g = log_g 129 self._f = f 130 131 if self._stop_criterium(tol): 132 break 133 134 return f, log_g
135
136 - def log_z(self, beta=1., ensembles=None):
137 """ 138 Use trapezoidal rule to evaluate the partition function. 139 """ 140 from numpy import array, multiply, reshape 141 142 is_float = False 143 144 if type(beta) == float: 145 beta = reshape(array(beta), (-1,)) 146 is_float = True 147 148 x = self._ex[0, 1:] - self._ex[0, :-1] 149 y = self._ex[0] 150 151 for i in range(1, self._ex.shape[0]): 152 x = multiply.outer(x, self._ex[i, 1:] - self._ex[i, :-1]) 153 y = multiply.outer(y, self._ex[i]) 154 155 y = -multiply.outer(beta, y) + self._log_g 156 y = reshape(array([y.T[1:], y.T[:-1]]), (2, -1)) 157 y = log_sum_exp(y, 0) - log(2) 158 y = reshape(y, (-1, len(beta))).T + log(x) 159 160 log_z = log_sum_exp(y.T, 0) 161 162 if is_float: 163 return float(log_z) 164 else: 165 return log_z
166
167 168 -class NonparametricWHAM(AbstractWHAM):
169 """ 170 Implementation of the nonparametric WHAM outlined in Habeck 2012, in which histograms 171 are reduced to delta peaks, this allows to use energies samples at different orders 172 of magnitude, improving the accuracy of the DOS estimates. 173 """ 174
175 - def estimate(self, n_iter=10000, tol=1e-10):
176 177 e_ij = numpy.array([ensemble.energy(self._e) 178 for ensemble in self._ensembles]).T 179 180 f = self._f 181 log_n = log(self._n) 182 self._L = [] 183 for _i in range(n_iter): 184 185 ## update density of states 186 log_g = -log_sum_exp((-e_ij - f + log_n).T, 0) 187 log_g -= log_sum_exp(log_g) 188 189 ## update free energies 190 f = log_sum_exp((-e_ij.T + log_g).T, 0) 191 self._L.append((self._n * f).sum() - log_g.sum()) 192 193 self._f = f 194 self._log_g = log_g 195 196 if self._stop_criterium(tol): 197 break 198 199 return f, log_g
200
201 - def log_g(self, normalize=True):
202 203 e_ij = numpy.array([ensemble.energy(self._e) 204 for ensemble in self._ensembles]).T 205 206 log_g = -log_sum_exp((-e_ij - self._f + log(self._n)).T, 0) 207 208 if normalize: 209 log_g -= log_sum_exp(log_g) 210 211 return log_g
212
213 - def log_z(self, beta=1., ensembles=None):
214 215 from numpy import multiply 216 217 if ensembles is not None: 218 e_ij_prime = numpy.array([ensemble.energy(self._e) 219 for ensemble in ensembles]) 220 else: 221 e_ij_prime = multiply.outer(beta, self._e) 222 223 224 log_z = log_sum_exp((-e_ij_prime + self.log_g()).T, 0) 225 226 return log_z
227