tools/i-pi/ipi/engine/thermostats.py

   1 """Contains the classes that deal with constant temperature dynamics.
   2
   3 Copyright (C) 2013, Joshua More and Michele Ceriotti
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, either version 3 of the License, or
   8 (at your option) any later version.
   9
  10 This program is distributed in the hope that it will be useful,
  11 but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13 GNU General Public License for more details.
  14
  15 You should have received a copy of the GNU General Public License
  16 along with this program. If not, see <http.//www.gnu.org/licenses/>.
  17
  18
  19 Contains the algorithms which propagate the thermostatting steps in the constant
  20 temperature ensembles. Includes the new GLE thermostat, which can be used to
  21 run PI+GLE dynamics, reducing the number of path integral beads required.
  22
  23 Classes:
  24    Thermostat: Base thermostat class with the generic methods and attributes.
  25    ThermoLangevin: Holds the algorithms for a langevin thermostat.
  26    ThermoPILE_L: Holds the algorithms for a path-integral langevin equation
  27       thermostat, with a thermostat coupled directly to the
  28       centroid coordinate of each bead.
  29    ThermoPILE_G: Holds the algorithms for a path-integral langevin equation
  30       thermostat, with a thermostat coupled to the kinetic energy for
  31       the entire system.
  32    ThermoSVR: Holds the algorithms for a stochastic velocity rescaling
  33       thermostat.
  34    ThermoGLE: Holds the algorithms for a generalised langevin equation
  35       thermostat.
  36    ThermoNMGLE: Holds the algorithms for a generalised langevin equation
  37       thermostat in the normal mode representation.
  38    ThermoNMGLEG: Holds the algorithms for a generalised langevin equation
  39       thermostat in the normal mode representation, with kinetic energy as
  40       well as potential energy sampling optimization.
  41 """
  42
  43 __all__ = ['Thermostat', 'ThermoLangevin', 'ThermoPILE_L', 'ThermoPILE_G',
  44            'ThermoSVR', 'ThermoGLE', 'ThermoNMGLE', 'ThermoNMGLEG']
  45
  46 import numpy as np
  47 from ipi.utils.depend   import *
  48 from ipi.utils.units    import *
  49 from ipi.utils.mathtools import matrix_exp, stab_cholesky, root_herm
  50 from ipi.utils.prng import Random
  51 from ipi.utils.messages import verbosity, warning, info
  52 from ipi.engine.beads import Beads
  53 from ipi.engine.normalmodes import NormalModes
  54
  55 class Thermostat(dobject):
  56    """Base thermostat class.
  57
  58    Gives the standard methods and attributes needed in all the thermostat
  59    classes.
  60
  61    Attributes:
  62       prng: A pseudo random number generator object.
  63       ndof: The number of degrees of freedom that the thermostat will be
  64          attached to.
  65
  66    Depend objects:
  67       dt: The time step used in the algorithms. Depends on the simulation dt.
  68       temp: The simulation temperature. Higher than the system temperature by
  69          a factor of the number of beads. Depends on the simulation temp.
  70       ethermo: The total energy exchanged with the bath due to the thermostat.
  71       p: The momentum vector that the thermostat is coupled to. Depends on the
  72          beads p object.
  73       m: The mass vector associated with p. Depends on the beads m object.
  74       sm: The square root of the mass vector.
  75    """
  76
  77    def __init__(self, temp = 1.0, dt = 1.0, ethermo=0.0):
  78       """Initialises Thermostat.
  79
  80       Args:
  81          temp: The simulation temperature. Defaults to 1.0.
  82          dt: The simulation time step. Defaults to 1.0.
  83          ethermo: The initial heat energy transferred to the bath.
  84             Defaults to 0.0. Will be non-zero if the thermostat is
  85             initialised from a checkpoint file.
  86       """
  87
  88       dset(self,"temp",   depend_value(name='temp', value=temp))
  89       dset(self,"dt",     depend_value(name='dt', value=dt))
  90       dset(self,"ethermo",depend_value(name='ethermo',value=ethermo))
  91
  92    def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
  93       """Binds the appropriate degrees of freedom to the thermostat.
  94
  95       This takes an object with degrees of freedom, and makes their momentum
  96       and mass vectors members of the thermostat. It also then creates the
  97       objects that will hold the data needed in the thermostat algorithms
  98       and the dependency network.
  99
 100       Args:
 101          beads: An optional beads object to take the mass and momentum vectors
 102             from.
 103          atoms: An optional atoms object to take the mass and momentum vectors
 104             from.
 105          pm: An optional tuple containing a single momentum value and its
 106             conjugate mass.
 107          prng: An optional pseudo random number generator object. Defaults to
 108             Random().
 109          fixdof: An optional integer which can specify the number of constraints
 110             applied to the system. Defaults to zero.
 111
 112       Raises:
 113          TypeError: Raised if no appropriate degree of freedom or object
 114             containing a momentum vector is specified for
 115             the thermostat to couple to.
 116       """
 117
 118       if prng is None:
 119          warning("Initializing thermostat from standard random PRNG", verbosity.medium)
 120          self.prng = Random()
 121       else:
 122          self.prng = prng
 123
 124       if not beads is None:
 125          dset(self,"p",beads.p.flatten())
 126          dset(self,"m",beads.m3.flatten())
 127       elif not atoms is None:
 128          dset(self,"p",dget(atoms, "p"))
 129          dset(self,"m",dget(atoms, "m3"))
 130       elif not pm is None:
 131          dset(self,"p",pm[0])
 132          dset(self,"m",pm[1])
 133       else:
 134          raise TypeError("Thermostat.bind expects either Beads, Atoms, NormalModes, or a (p,m) tuple to bind to")
 135
 136       if fixdof is None:
 137          self.ndof = len(self.p)
 138       else:
 139          self.ndof = float(len(self.p) - fixdof)
 140
 141       dset(self, "sm",
 142          depend_array(name="sm", value=np.zeros(len(dget(self,"m"))),
 143             func=self.get_sm, dependencies=[dget(self,"m")]))
 144
 145    def get_sm(self):
 146       """Retrieves the square root of the mass matrix.
 147
 148       Returns:
 149          A vector of the square root of the mass matrix with one value for
 150          each degree of freedom.
 151       """
 152
 153       return np.sqrt(self.m)
 154
 155    def step(self):
 156       """Dummy thermostat step."""
 157
 158       pass
 159
 160
 161 class ThermoLangevin(Thermostat):
 162    """Represents a langevin thermostat.
 163
 164    Depend objects:
 165       tau: Thermostat damping time scale. Larger values give a less strongly
 166          coupled thermostat.
 167       T: Coefficient of the diffusive contribution of the thermostat, i.e. the
 168          drift back towards equilibrium. Depends on tau and the time step.
 169       S: Coefficient of the stochastic contribution of the thermostat, i.e.
 170          the uncorrelated Gaussian noise. Depends on T and the temperature.
 171    """
 172
 173    def get_T(self):
 174       """Calculates the coefficient of the overall drift of the velocities."""
 175
 176       return np.exp(-0.5*self.dt/self.tau)
 177
 178    def get_S(self):
 179       """Calculates the coefficient of the white noise."""
 180
 181       return np.sqrt(Constants.kb*self.temp*(1 - self.T**2))
 182
 183    def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
 184       """Initialises ThermoLangevin.
 185
 186       Args:
 187          temp: The simulation temperature. Defaults to 1.0.
 188          dt: The simulation time step. Defaults to 1.0.
 189          tau: The thermostat damping timescale. Defaults to 1.0.
 190          ethermo: The initial heat energy transferred to the bath.
 191             Defaults to 0.0. Will be non-zero if the thermostat is
 192             initialised from a checkpoint file.
 193       """
 194
 195       super(ThermoLangevin,self).__init__(temp, dt, ethermo)
 196
 197       dset(self,"tau",depend_value(value=tau,name='tau'))
 198       dset(self,"T",
 199          depend_value(name="T",func=self.get_T,
 200             dependencies=[dget(self,"tau"), dget(self,"dt")]))
 201       dset(self,"S",
 202          depend_value(name="S",func=self.get_S,
 203             dependencies=[dget(self,"temp"), dget(self,"T")]))
 204
 205    def step(self):
 206       """Updates the bound momentum vector with a langevin thermostat."""
 207
 208       p = depstrip(self.p).copy()
 209       sm = depstrip(self.sm)
 210
 211       p /= sm
 212
 213       self.ethermo += np.dot(p,p)*0.5
 214       p *= self.T
 215       p += self.S*self.prng.gvec(len(p))
 216       self.ethermo -= np.dot(p,p)*0.5
 217
 218       p *= sm
 219
 220       self.p = p
 221
 222
 223 class ThermoPILE_L(Thermostat):
 224    """Represents a PILE thermostat with a local centroid thermostat.
 225
 226    Attributes:
 227       _thermos: The list of the different thermostats for all the ring polymer
 228          normal modes.
 229       nm: A normal modes object to attach the thermostat to.
 230       prng: Random number generator used in the stochastic integration
 231          algorithms.
 232
 233    Depend objects:
 234       tau: Centroid thermostat damping time scale. Larger values give a
 235          less strongly coupled centroid thermostat.
 236       tauk: Thermostat damping time scale for the non-centroid normal modes.
 237          Depends on the ring polymer spring constant, and thus the simulation
 238          temperature.
 239       pilescale: A float used to reduce the intensity of the PILE thermostat if
 240          required.
 241    """
 242
 243    def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale=1.0):
 244       """Initialises ThermoPILE_L.
 245
 246       Args:
 247          temp: The simulation temperature. Defaults to 1.0.
 248          dt: The simulation time step. Defaults to 1.0.
 249          tau: The centroid thermostat damping timescale. Defaults to 1.0.
 250          ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
 251             be non-zero if the thermostat is initialised from a checkpoint file.
 252          scale: A float used to reduce the intensity of the PILE thermostat if
 253             required.
 254
 255       Raises:
 256          TypeError: Raised if the thermostat is used with any object other than
 257             a beads object, so that we make sure that the objects needed for the
 258             normal mode transformation exist.
 259       """
 260
 261       super(ThermoPILE_L,self).__init__(temp,dt,ethermo)
 262       dset(self,"tau",depend_value(value=tau,name='tau'))
 263       dset(self,"pilescale",depend_value(value=scale,name='pilescale'))
 264
 265    def bind(self, nm=None, prng=None, bindcentroid=True, fixdof=None):
 266       """Binds the appropriate degrees of freedom to the thermostat.
 267
 268       This takes a beads object with degrees of freedom, and makes its momentum
 269       and mass vectors members of the thermostat. It also then creates the
 270       objects that will hold the data needed in the thermostat algorithms
 271       and the dependency network.
 272
 273       Gives the interface for both the PILE_L and PILE_G thermostats, which
 274       only differ in their treatment of the centroid coordinate momenta.
 275
 276       Args:
 277          nm: An optional normal mode object to take the mass and momentum
 278             vectors from.
 279          prng: An optional pseudo random number generator object. Defaults to
 280             Random().
 281          bindcentroid: An optional boolean which decides whether a Langevin
 282             thermostat is attached to the centroid mode of each atom
 283             separately, or the total kinetic energy. Defaults to True, which
 284             gives a thermostat bound to each centroid momentum.
 285          fixdof: An optional integer which can specify the number of constraints
 286             applied to the system. Defaults to zero.
 287
 288       Raises:
 289          TypeError: Raised if no appropriate degree of freedom or object
 290             containing a momentum vector is specified for
 291             the thermostat to couple to.
 292       """
 293
 294       if nm is None or not type(nm) is NormalModes:
 295          raise TypeError("ThermoPILE_L.bind expects a NormalModes argument to bind to")
 296       if prng is None:
 297          self.prng = Random()
 298       else:
 299          self.prng = prng
 300
 301       prev_ethermo = self.ethermo
 302
 303       # creates a set of thermostats to be applied to individual normal modes
 304       self._thermos = [ ThermoLangevin(temp=1, dt=1, tau=1) for b in range(nm.nbeads) ]
 305       # optionally does not bind the centroid, so we can re-use all of this
 306       # in the PILE_G case
 307       if not bindcentroid:
 308          self._thermos[0] = None
 309
 310       self.nm = nm
 311
 312       dset(self,"tauk",
 313          depend_array(name="tauk", value=np.zeros(nm.nbeads-1,float),
 314             func=self.get_tauk, dependencies=[dget(self,"pilescale"), dget(nm,"dynomegak")] ) )
 315
 316       # must pipe all the dependencies in such a way that values for the nm thermostats
 317       # are automatically updated based on the "master" thermostat
 318       def make_taugetter(k):
 319          return lambda: self.tauk[k-1]
 320       it = 0
 321       for t in self._thermos:
 322          if t is None:
 323             it += 1
 324             continue
 325          if it > 0:
 326             fixdof = None # only the centroid thermostat may have constraints
 327
 328          # bind thermostat t to the it-th bead
 329          t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]),prng=self.prng, fixdof=fixdof)
 330          # pipes temp and dt
 331          deppipe(self,"temp", t, "temp")
 332          deppipe(self,"dt", t, "dt")
 333
 334          # for tau it is slightly more complex
 335          if it == 0:
 336             deppipe(self,"tau", t, "tau")
 337          else:
 338             # Here we manually connect _thermos[i].tau to tauk[i].
 339             # Simple and clear.
 340             dget(t,"tau").add_dependency(dget(self,"tauk"))
 341             dget(t,"tau")._func = make_taugetter(it)
 342          dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
 343          it += 1
 344
 345       # since the ethermo will be "delegated" to the normal modes thermostats,
 346       # one has to split
 347       # any previously-stored value between the sub-thermostats
 348       if bindcentroid:
 349          for t in self._thermos:
 350             t.ethermo = prev_ethermo/nm.nbeads
 351          dget(self,"ethermo")._func = self.get_ethermo;
 352          # if we are not binding the centroid just yet, this bit of the piping
 353          # is delegated to the function which is actually calling this
 354
 355    def get_tauk(self):
 356       """Computes the thermostat damping time scale for the non-centroid
 357       normal modes.
 358
 359       Returns:
 360          An array with the damping time scales for the non-centroid modes.
 361       """
 362
 363       # Also include an optional scaling factor to reduce the intensity of NM thermostats
 364       return  np.array([ self.pilescale/(2*self.nm.dynomegak[k])  for k in range(1,len(self._thermos)) ])
 365
 366    def get_ethermo(self):
 367       """Computes the total energy transferred to the heat bath for all the
 368       thermostats.
 369       """
 370
 371       et = 0.0;
 372       for t in self._thermos:
 373          et += t.ethermo
 374       return et
 375
 376    def step(self):
 377       """Updates the bound momentum vector with a PILE thermostat."""
 378
 379       # super-cool! just loop over the thermostats! it's as easy as that!
 380       for t in self._thermos:
 381          t.step()
 382
 383
 384 class ThermoSVR(Thermostat):
 385    """Represents a stochastic velocity rescaling thermostat.
 386
 387    Depend objects:
 388       tau: Centroid thermostat damping time scale. Larger values give a
 389          less strongly coupled centroid thermostat.
 390       K: Scaling factor for the total kinetic energy. Depends on the
 391          temperature.
 392       et: Parameter determining the strength of the thermostat coupling.
 393          Depends on tau and the time step.
 394    """
 395
 396    def get_et(self):
 397       """Calculates the damping term in the propagator."""
 398
 399       return np.exp(-0.5*self.dt/self.tau)
 400
 401    def get_K(self):
 402       """Calculates the average kinetic energy per degree of freedom."""
 403
 404       return Constants.kb*self.temp*0.5
 405
 406    def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0):
 407       """Initialises ThermoSVR.
 408
 409       Args:
 410          temp: The simulation temperature. Defaults to 1.0.
 411          dt: The simulation time step. Defaults to 1.0.
 412          tau: The thermostat damping timescale. Defaults to 1.0.
 413          ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
 414             be non-zero if the thermostat is initialised from a checkpoint file.
 415       """
 416
 417       super(ThermoSVR,self).__init__(temp,dt,ethermo)
 418
 419       dset(self,"tau",depend_value(value=tau,name='tau'))
 420       dset(self,"et",
 421          depend_value(name="et",func=self.get_et,
 422             dependencies=[dget(self,"tau"), dget(self,"dt")]))
 423       dset(self,"K",
 424          depend_value(name="K",func=self.get_K, dependencies=[dget(self,"temp")]))
 425
 426    def step(self):
 427       """Updates the bound momentum vector with a stochastic velocity rescaling
 428       thermostat. See G Bussi, D Donadio, M Parrinello,
 429       Journal of Chemical Physics 126, 014101 (2007)
 430       """
 431
 432       K = np.dot(depstrip(self.p),depstrip(self.p)/depstrip(self.m))*0.5
 433
 434       # rescaling is un-defined if the KE is zero
 435       if K == 0.0:
 436          return
 437
 438       # gets the stochastic term (basically a Gamma distribution for the kinetic energy)
 439       r1 = self.prng.g
 440       if (self.ndof-1)%2 == 0:
 441          rg = 2.0*self.prng.gamma((self.ndof-1)/2)
 442       else:
 443          rg = 2.0*self.prng.gamma((self.ndof-2)/2) + self.prng.g**2
 444
 445       alpha2 = self.et + self.K/K*(1 - self.et)*(r1**2 + rg) + 2.0*r1*np.sqrt(self.K/K*self.et*(1 - self.et))
 446       alpha = np.sqrt(alpha2)
 447       if (r1 + np.sqrt(2*K/self.K*self.et/(1 - self.et))) < 0:
 448          alpha *= -1
 449
 450       self.ethermo += K*(1 - alpha2)
 451       self.p *= alpha
 452
 453
 454 class ThermoPILE_G(ThermoPILE_L):
 455    """Represents a PILE thermostat with a global centroid thermostat.
 456
 457    Simply replaces the Langevin thermostat for the centroid normal mode with
 458    a global velocity rescaling thermostat.
 459    """
 460
 461    def __init__(self, temp = 1.0, dt = 1.0, tau = 1.0, ethermo=0.0, scale = 1.0):
 462       """Initialises ThermoPILE_G.
 463
 464       Args:
 465          temp: The simulation temperature. Defaults to 1.0.
 466          dt: The simulation time step. Defaults to 1.0.
 467          tau: The centroid thermostat damping timescale. Defaults to 1.0.
 468          ethermo: The initial conserved energy quantity. Defaults to 0.0. Will
 469             be non-zero if the thermostat is initialised from a checkpoint file.
 470          scale: A float used to reduce the intensity of the PILE thermostat if
 471             required.
 472       """
 473
 474       super(ThermoPILE_G,self).__init__(temp,dt,tau,ethermo)
 475       dset(self,"pilescale",depend_value(value=scale,name='pilescale'))
 476
 477    def bind(self, nm=None, prng=None, fixdof=None):
 478       """Binds the appropriate degrees of freedom to the thermostat.
 479
 480       This takes a beads object with degrees of freedom, and makes its momentum
 481       and mass vectors members of the thermostat. It also then creates the
 482       objects that will hold the data needed in the thermostat algorithms
 483       and the dependency network.
 484
 485       Uses the PILE_L bind interface, with bindcentroid set to false so we can
 486       specify that thermostat separately, by binding a global
 487       thermostat to the centroid mode.
 488
 489       Args:
 490          beads: An optional beads object to take the mass and momentum vectors
 491             from.
 492          prng: An optional pseudo random number generator object. Defaults to
 493             Random().
 494          fixdof: An optional integer which can specify the number of constraints
 495             applied to the system. Defaults to zero.
 496
 497       """
 498
 499       # first binds as a local PILE, then substitutes the thermostat on the centroid
 500       prev_ethermo = self.ethermo
 501       super(ThermoPILE_G,self).bind(nm=nm,prng=prng,bindcentroid=False, fixdof=fixdof)
 502
 503       #centroid thermostat
 504       self._thermos[0] = ThermoSVR(temp=1, dt=1, tau=1)
 505
 506       t = self._thermos[0]
 507       t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]),prng=self.prng, fixdof=fixdof)
 508       deppipe(self,"temp", t, "temp")
 509       deppipe(self,"dt", t, "dt")
 510       deppipe(self,"tau", t, "tau")
 511       dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
 512
 513       # splits any previous ethermo between the thermostats, and finishes to bind ethermo to the sum function
 514       for t in self._thermos:
 515          t.ethermo = prev_ethermo/nm.nbeads
 516       dget(self,"ethermo")._func = self.get_ethermo;
 517
 518
 519 class ThermoGLE(Thermostat):
 520    """Represents a GLE thermostat.
 521
 522    This is similar to a langevin thermostat, in that it uses Gaussian random
 523    numbers to simulate a heat bath acting on the system, but simulates a
 524    non-Markovian system by using a Markovian formulation in an extended phase
 525    space. This allows for a much greater degree of flexibility, and this
 526    thermostat, properly fitted, can give the an approximation to the correct
 527    quantum ensemble even for a classical, 1-bead simulation. More reasonably,
 528    using this thermostat allows for a far smaller number of replicas of the
 529    system to be used, as the convergence of the properties
 530    of the system is accelerated with respect to number of beads when PI+GLE
 531    are used in combination. (See M. Ceriotti, D. E. Manolopoulos, M. Parinello,
 532    J. Chem. Phys. 134, 084104 (2011)).
 533
 534    Attributes:
 535       ns: The number of auxilliary degrees of freedom.
 536       s: An array holding all the momenta, including the ones for the
 537          auxilliary degrees of freedom.
 538
 539    Depend objects:
 540       A: Drift matrix giving the damping time scales for all the different
 541          degrees of freedom.
 542       C: Static covariance matrix.
 543          Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
 544          diffusion matrix, giving the strength of the coupling of the system
 545          with the heat bath, and thus the size of the stochastic
 546          contribution of the thermostat.
 547       T: Matrix for the diffusive contribution of the thermostat, i.e. the
 548          drift back towards equilibrium. Depends on A and the time step.
 549       S: Matrix for the stochastic contribution of the thermostat, i.e.
 550          the uncorrelated Gaussian noise. Depends on C and T.
 551    """
 552
 553    def get_T(self):
 554       """Calculates the matrix for the overall drift of the velocities."""
 555
 556       return matrix_exp(-0.5*self.dt*self.A)
 557
 558    def get_S(self):
 559       """Calculates the matrix for the coloured noise."""
 560
 561       SST = Constants.kb*(self.C - np.dot(self.T,np.dot(self.C,self.T.T)))
 562
 563       # Uses a symetric decomposition rather than Cholesky, since it is more stable
 564       return root_herm(SST)
 565
 566    def get_C(self):
 567       """Calculates C from temp (if C is not set explicitly)"""
 568
 569       rC = np.identity(self.ns + 1,float)*self.temp
 570       return rC[:]
 571
 572    def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
 573       """Initialises ThermoGLE.
 574
 575       Args:
 576          temp: The simulation temperature. Defaults to 1.0.
 577          dt: The simulation time step. Defaults to 1.0.
 578          A: An optional matrix giving the drift matrix. Defaults to a single
 579             value of 1.0.
 580          C: An optional matrix giving the covariance matrix. Defaults to an
 581             identity matrix times temperature with the same dimensions as the
 582             total number of degrees of freedom in the system.
 583          ethermo: The initial heat energy transferred to the bath.
 584             Defaults to 0.0. Will be non-zero if the thermostat is
 585             initialised from a checkpoint file.
 586       """
 587
 588       super(ThermoGLE,self).__init__(temp,dt,ethermo)
 589
 590       if A is None:
 591          A = np.identity(1,float)
 592       dset(self,"A",depend_value(value=A.copy(),name='A'))
 593
 594       self.ns = len(self.A) - 1;
 595
 596       # now, this is tricky. if C is taken from temp, then we want it to be updated
 597       # as a depend of temp. Otherwise, we want it to be an independent beast.
 598       if C is None:
 599          C = np.identity(self.ns+1,float)*self.temp
 600          dset(self,"C",
 601             depend_value(name='C', func=self.get_C,
 602                dependencies=[dget(self,"temp")]))
 603       else:
 604          dset(self,"C",depend_value(value=C.copy(),name='C'))
 605
 606       dset(self,"T",
 607          depend_value(name="T",func=self.get_T,
 608             dependencies=[dget(self,"A"), dget(self,"dt")]))
 609       dset(self,"S",
 610          depend_value(name="S",func=self.get_S,
 611             dependencies=[dget(self,"C"), dget(self,"T")]))
 612
 613       self.s = np.zeros(0)
 614
 615    def bind(self, beads=None, atoms=None, pm=None, prng=None, fixdof=None):
 616       """Binds the appropriate degrees of freedom to the thermostat.
 617
 618       This takes an object with degrees of freedom, and makes their momentum
 619       and mass vectors members of the thermostat. It also then creates the
 620       objects that will hold the data needed in the thermostat algorithms
 621       and the dependency network.
 622
 623       Args:
 624          beads: An optional beads object to take the mass and momentum vectors
 625             from.
 626          atoms: An optional atoms object to take the mass and momentum vectors
 627             from.
 628          pm: An optional tuple containing a single momentum value and its
 629             conjugate mass.
 630          prng: An optional pseudo random number generator object. Defaults to
 631             Random().
 632          fixdof: An optional integer which can specify the number of constraints
 633             applied to the system. Defaults to zero.
 634
 635       Raises:
 636          TypeError: Raised if no appropriate degree of freedom or object
 637             containing a momentum vector is specified for
 638             the thermostat to couple to.
 639       """
 640
 641       super(ThermoGLE,self).bind(beads,atoms,pm,prng,fixdof)
 642
 643       # allocates, initializes or restarts an array of s's
 644       if self.s.shape != (self.ns + 1, len(dget(self,"m"))):
 645          if len(self.s) > 0:
 646             warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
 647          self.s = np.zeros((self.ns + 1, len(dget(self,"m"))))
 648
 649          # Initializes the s vector in the free-particle limit
 650          info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
 651          SC = stab_cholesky(self.C*Constants.kb)
 652          self.s[:] = np.dot(SC, self.prng.gvec(self.s.shape))
 653       else:
 654          info("GLE additional DOFs initialised from input.", verbosity.medium)
 655
 656    def step(self):
 657       """Updates the bound momentum vector with a GLE thermostat"""
 658
 659       p = depstrip(self.p).copy()
 660
 661       self.s[0,:] = self.p/self.sm
 662
 663       self.ethermo += np.dot(self.s[0],self.s[0])*0.5
 664       self.s[:] = np.dot(self.T,self.s) + np.dot(self.S,self.prng.gvec(self.s.shape))
 665       self.ethermo -= np.dot(self.s[0],self.s[0])*0.5
 666
 667       self.p = self.s[0]*self.sm
 668
 669
 670 class ThermoNMGLE(Thermostat):
 671    """Represents a 'normal-modes' GLE thermostat.
 672
 673    An extension to the GLE thermostat which is applied in the
 674    normal modes representation, and which allows to use a different
 675    GLE for each normal mode
 676
 677    Attributes:
 678       ns: The number of auxilliary degrees of freedom.
 679       nb: The number of beads.
 680       s: An array holding all the momenta, including the ones for the
 681          auxilliary degrees of freedom.
 682
 683    Depend objects:
 684       A: Drift matrix giving the damping time scales for all the different
 685          degrees of freedom (must contain nb terms).
 686       C: Static covariance matrix.
 687          Satisfies A.C + C.transpose(A) = B.transpose(B), where B is the
 688          diffusion matrix, giving the strength of the coupling of the system
 689          with the heat bath, and thus the size of the stochastic
 690          contribution of the thermostat.
 691    """
 692
 693    def get_C(self):
 694       """Calculates C from temp (if C is not set explicitely)."""
 695
 696       rv = np.ndarray((self.nb, self.ns+1, self.ns+1), float)
 697       for b in range(0,self.nb):
 698          rv[b] = np.identity(self.ns + 1,float)*self.temp
 699       return rv[:]
 700
 701    def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, ethermo=0.0):
 702       """Initialises ThermoGLE.
 703
 704       Args:
 705          temp: The simulation temperature. Defaults to 1.0.
 706          dt: The simulation time step. Defaults to 1.0.
 707          A: An optional matrix giving the drift matrix. Defaults to a single
 708             value of 1.0.
 709          C: An optional matrix giving the covariance matrix. Defaults to an
 710             identity matrix times temperature with the same dimensions as the
 711             total number of degrees of freedom in the system.
 712          ethermo: The initial heat energy transferred to the bath.
 713             Defaults to 0.0. Will be non-zero if the thermostat is
 714             initialised from a checkpoint file.
 715       """
 716
 717       super(ThermoNMGLE,self).__init__(temp,dt,ethermo)
 718
 719       if A is None:
 720          A = np.identity(1,float)
 721       dset(self,"A",depend_value(value=A.copy(),name='A'))
 722
 723       self.nb = len(self.A)
 724       self.ns = len(self.A[0]) - 1;
 725
 726       # now, this is tricky. if C is taken from temp, then we want it to be
 727       # updated as a depend of temp.
 728       # Otherwise, we want it to be an independent beast.
 729       if C is None:
 730          dset(self,"C",depend_value(name='C', func=self.get_C, dependencies=[dget(self,"temp")]))
 731       else:
 732          dset(self,"C",depend_value(value=C.copy(),name='C'))
 733
 734    def bind(self, nm=None, prng=None, fixdof=None):
 735       """Binds the appropriate degrees of freedom to the thermostat.
 736
 737       This takes an object with degrees of freedom, and makes their momentum
 738       and mass vectors members of the thermostat. It also then creates the
 739       objects that will hold the data needed in the thermostat algorithms
 740       and the dependency network. Actually, this specific thermostat requires
 741       being called on a beads object.
 742
 743       Args:
 744          nm: An optional normal modes object to take the mass and momentum
 745             vectors from.
 746          prng: An optional pseudo random number generator object. Defaults to
 747             Random().
 748          fixdof: An optional integer which can specify the number of constraints
 749             applied to the system. Defaults to zero.
 750
 751       Raises:
 752          TypeError: Raised if no beads object is specified for
 753             the thermostat to couple to.
 754       """
 755
 756       if nm is None or not type(nm) is NormalModes:
 757          raise TypeError("ThermoNMGLE.bind expects a NormalModes argument to bind to")
 758
 759       if prng is None:
 760          self.prng = Random()
 761       else:
 762          self.prng = prng
 763
 764       if (nm.nbeads != self.nb):
 765          raise IndexError("The parameters in nm_gle options correspond to a bead number "+str(self.nb)+ " which does not match the number of beads in the path" + str(nm.nbeads) )
 766
 767       # allocates, initializes or restarts an array of s's
 768       if self.s.shape != (self.nb, self.ns + 1, nm.natoms *3) :
 769          if len(self.s) > 0:
 770             warning("Mismatch in GLE s array size on restart, will reinitialise to free particle.", verbosity.low)
 771          self.s = np.zeros((self.nb, self.ns + 1, nm.natoms*3))
 772
 773          # Initializes the s vector in the free-particle limit
 774          info(" GLE additional DOFs initialised to the free-particle limit.", verbosity.low)
 775          for b in range(self.nb):
 776             SC = stab_cholesky(self.C[b]*Constants.kb)
 777             self.s[b] = np.dot(SC, self.prng.gvec(self.s[b].shape))
 778       else:
 779          info("GLE additional DOFs initialised from input.", verbosity.medium)
 780
 781       prev_ethermo = self.ethermo
 782
 783       # creates a set of thermostats to be applied to individual normal modes
 784       self._thermos = [ThermoGLE(temp=1, dt=1, A=self.A[b], C=self.C[b]) for b in range(self.nb)]
 785
 786       # must pipe all the dependencies in such a way that values for the nm
 787       # thermostats are automatically updated based on the "master" thermostat
 788       def make_Agetter(k):
 789          return lambda: self.A[k]
 790       def make_Cgetter(k):
 791          return lambda: self.C[k]
 792
 793       it = 0
 794       for t in self._thermos:
 795          t.s = self.s[it]  # gets the s's as a slice of self.s
 796          t.bind(pm=(nm.pnm[it,:],nm.dynm3[it,:]), prng=self.prng) # bind thermostat t to the it-th normal mode
 797
 798          # pipes temp and dt
 799          deppipe(self,"temp", t, "temp")
 800          deppipe(self,"dt", t, "dt")
 801
 802          # here we pipe the A and C of individual NM to the "master" arrays
 803          dget(t,"A").add_dependency(dget(self,"A"))
 804          dget(t,"A")._func = make_Agetter(it)
 805          dget(t,"C").add_dependency(dget(self,"C"))
 806          dget(t,"C")._func = make_Cgetter(it)
 807          dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
 808          it += 1
 809
 810       # since the ethermo will be "delegated" to the normal modes thermostats,
 811       # one has to split
 812       # any previously-stored value between the sub-thermostats
 813       for t in self._thermos:
 814          t.ethermo = prev_ethermo/self.nb
 815
 816       dget(self,"ethermo")._func = self.get_ethermo;
 817
 818    def step(self):
 819       """Updates the thermostat in NM representation by looping over the
 820       individual DOFs.
 821       """
 822
 823       for t in self._thermos:
 824          t.step()
 825
 826    def get_ethermo(self):
 827       """Computes the total energy transferred to the heat bath for all the nm
 828       thermostats.
 829       """
 830
 831       et = 0.0;
 832       for t in self._thermos:
 833          et += t.ethermo
 834       return et
 835
 836
 837 class ThermoNMGLEG(ThermoNMGLE):
 838    """Represents a 'normal-modes' GLE thermostat + SVR.
 839
 840    An extension to the above NMGLE thermostat which also adds a stochastic velocity
 841    rescaling to the centroid.
 842
 843    Depend objects:
 844       tau: Thermostat damping time scale. Larger values give a less strongly
 845          coupled thermostat.
 846    """
 847
 848    def __init__(self, temp = 1.0, dt = 1.0, A = None, C = None, tau=1.0, ethermo=0.0):
 849
 850       super(ThermoNMGLEG,self).__init__(temp, dt, A, C, ethermo)
 851       dset(self,"tau",depend_value(value=tau,name='tau'))
 852
 853    def bind(self, nm=None, prng=None, fixdof=None):
 854       """Binds the appropriate degrees of freedom to the thermostat.
 855
 856       This takes an object with degrees of freedom, and makes their momentum
 857       and mass vectors members of the thermostat. It also then creates the
 858       objects that will hold the data needed in the thermostat algorithms
 859       and the dependency network. Actually, this specific thermostat requires
 860       being called on a beads object.
 861
 862       Args:
 863          nm: An optional normal modes object to take the mass and momentum
 864             vectors from.
 865          prng: An optional pseudo random number generator object. Defaults to
 866             Random().
 867          fixdof: An optional integer which can specify the number of constraints
 868             applied to the system. Defaults to zero.
 869       """
 870
 871       super(ThermoNMGLEG,self).bind(nm, prng, fixdof)
 872
 873       t = ThermoSVR(self.temp, self.dt, self.tau)
 874
 875       t.bind(pm=(nm.pnm[0,:],nm.dynm3[0,:]), prng=self.prng) # bind global thermostat to centroid
 876
 877       # pipes temp and dt
 878       deppipe(self,"temp", t, "temp")
 879       deppipe(self,"dt", t, "dt")
 880       deppipe(self,"tau", t, "tau")
 881
 882       dget(self,"ethermo").add_dependency(dget(t,"ethermo"))
 883       self._thermos.append(t)
 884