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