Source code for CDTK.Dynamics.Thermostat

#*  **************************************************************************
#*
#*  CDTK, Chemical Dynamics Toolkit
#*  A modular system for chemical dynamics applications and more
#*
#*  Copyright (C) 2011, 2012, 2013, 2014, 2015, 2016
#*  Oriol Vendrell, DESY, <oriol.vendrell@desy.de>
#*
#*  Copyright (C) 2020, 2021, 2022, 2023
#*  Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
#*  This file is part of CDTK.
#*
#*  CDTK is free software: you can redistribute it and/or modify
#*  it under the terms of the GNU General Public License as published by
#*  the Free Software Foundation, either version 3 of the License, or
#*  (at your option) any later version.
#*
#*  This program is distributed in the hope that it will be useful,
#*  but WITHOUT ANY WARRANTY; without even the implied warranty of
#*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#*  GNU General Public License for more details.
#*
#*  You should have received a copy of the GNU General Public License
#*  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#*
#*  **************************************************************************
#! /usr/bin/env python

import os
import sys
import datetime
import uuid
import math as mt

import numpy as np

import CDTK.Tools.Conversion as conv




class Thermostat(object):
    """
    This class describes different thermostats that can be attached to Trajectories.
    Currently implemented:
    massive-andersen: Resample all momenta every DT timesteps
    """
    def __init__(self,**opts):

        # name of thermostat
        self.type = opts.get('type', 'massive-andersen') 

        # temperature in Kelvin and beta in au
        self.temp = float(opts.get('temp', '300')) 
        self.beta = 1.0 / (conv.FPC_K_B_EV * conv.ev2au * self.temp) 

        # masses
        self.M = None

        # timescale for thermostat
        self.dt = float(opts.get('dt', 0.0)) 



    def init_input_options(self,iopts):
        """
        This functions sets input options.

        Input: 
        iopts  --- Input options as a dictonary, where the key defines the variable 
                   and the value is a list with the value of the variable as first entry.
        """
        for k in iopts.keys():
            setattr(self,k,iopts[k][0])





    def apply(self, v, t_old, t, n, nb=None):
        """
        This function applies the thermostat to a given set of 
        at a given time interval (current time and last time it was applied).

        Input:
        v     --- current velocities
        t_old --- previous time in au
        t     --- current time in au
        n     --- step in the velocity verlet cycle
        nb    --- number of beads per DOF

        Output:
        new velocities.
        """

#        print "Appling thermostatt", v, t_old, t, n, nb

        if self.type == 'massive-andersen':
            # mid way in the  Velocity Verlet cycle
            if n < 2:
                return v

            # end of the velocity verlet cycle -> resample
            else:
                old_sample = np.floor(t_old / self.dt)
                new_sample = np.floor(t / self.dt)

                # we stepped over a full resampling step
                if old_sample < new_sample:

                    # for older versions
                    if nb is None:
                        sigma = np.sqrt(self.M / self.beta)
                        v_new = np.random.normal(np.zeros(len(v)), sigma) / self.M
                    # draw for each sample
                    else:
                        v_new = []
                        for i in range(len(nb)):
                            m = self.M[i]
                            nf = float(nb[i])
                            sigma = np.sqrt(nf * m / self.beta)
                            v_new.append(np.random.normal(np.zeros(nb[i]), sigma, nb[i]) / m)

                    return v_new
                else:
                    return v

        elif self.type == 'andersen':
            if nb is not None:
                print("Andersen thermostat not implemented for RPMD. Please use massive-andersen.")
                exit(-1)
                
            # mid way in the  Velocity Verlet cycle
            if n < 2:
                return v

            # end of the velocity verlet cycle -> resample
            else:
                timestep = t - t_old
                comp = timestep / self.dt

                v_new = []
                for i, vs in enumerate(v.reshape((-1,3))):

                    r = np.random.rand()
                    if r < comp:
                        m =  self.M[i*3]
                        sigma = np.sqrt(m / self.beta)
                        v_new.append(np.random.normal(np.zeros(3), sigma) / m)
                    else: 
                        v_new.append(vs)

                return np.array(v_new).flatten()
            

        elif self.type == 'nose-hoover':
            

            print("nose-hoover needs to be fully implemented")
            exit(-1)

        else:
            print("Given thermostat not implemented: ", self.type)
            exit(-1)