Coverage for /builds/debichem-team/python-ase/ase/md/bussi.py: 100.00%

1"""Bussi NVT dynamics class.""" 

2 

3import math 

4 

5import numpy as np 

6 

7from ase import units 

8from ase.md.verlet import VelocityVerlet 

9 

10 

11class Bussi(VelocityVerlet): 

12 """Bussi stochastic velocity rescaling (NVT) molecular dynamics. 

13 Based on the paper from Bussi et al. (https://arxiv.org/abs/0803.4060) 

14 

15 Parameters 

16 ---------- 

17 atoms : Atoms 

18 The atoms object. 

19 timestep : float 

20 The time step in ASE time units. 

21 temperature_K : float 

22 The desired temperature, in Kelvin. 

23 taut : float 

24 Time constant for Bussi temperature coupling in ASE time units. 

25 rng : numpy.random, optional 

26 Random number generator. 

27 **md_kwargs : dict, optional 

28 Additional arguments passed to :class:~ase.md.md.MolecularDynamics 

29 base class. 

30 """ 

31 

32 def __init__( 

33 self, 

34 atoms, 

35 timestep, 

36 temperature_K, 

37 taut, 

38 rng=np.random, 

39 **md_kwargs, 

40 ): 

41 super().__init__(atoms, timestep, **md_kwargs) 

42 

43 self.temp = temperature_K * units.kB 

44 self.taut = taut 

45 self.rng = rng 

46 

47 self.ndof = self.atoms.get_number_of_degrees_of_freedom() 

48 

49 self.target_kinetic_energy = 0.5 * self.temp * self.ndof 

50 

51 if np.isclose( 

52 self.atoms.get_kinetic_energy(), 0.0, rtol=0, atol=1e-12 

53 ): 

54 raise ValueError( 

55 "Initial kinetic energy is zero. " 

56 "Please set the initial velocities before running Bussi NVT." 

57 ) 

58 

59 self._exp_term = math.exp(-self.dt / self.taut) 

60 self._masses = self.atoms.get_masses()[:, np.newaxis] 

61 

62 self.transferred_energy = 0.0 

63 

64 def scale_velocities(self): 

65 """Do the NVT Bussi stochastic velocity scaling.""" 

66 kinetic_energy = self.atoms.get_kinetic_energy() 

67 alpha = self.calculate_alpha(kinetic_energy) 

68 

69 momenta = self.atoms.get_momenta() 

70 self.atoms.set_momenta(alpha * momenta) 

71 

72 self.transferred_energy += (alpha**2 - 1.0) * kinetic_energy 

73 

74 def calculate_alpha(self, kinetic_energy): 

75 """Calculate the scaling factor alpha using equation (A7) 

76 from the Bussi paper.""" 

77 

78 energy_scaling_term = ( 

79 (1 - self._exp_term) 

80 * self.target_kinetic_energy 

81 / kinetic_energy 

82 / self.ndof 

83 ) 

84 

85 # R1 in Eq. (A7) 

86 normal_noise = self.rng.standard_normal() 

87 # \sum_{i=2}^{Nf} R_i^2 in Eq. (A7) 

88 # 2 * standard_gamma(n / 2) is equal to chisquare(n) 

89 sum_of_noises = 2.0 * self.rng.standard_gamma(0.5 * (self.ndof - 1)) 

90 

91 return math.sqrt( 

92 self._exp_term 

93 + energy_scaling_term * (sum_of_noises + normal_noise**2) 

94 + 2 

95 * normal_noise 

96 * math.sqrt(self._exp_term * energy_scaling_term) 

97 ) 

98 

99 def step(self, forces=None): 

100 """Move one timestep forward using Bussi NVT molecular dynamics.""" 

101 self.scale_velocities() 

102 return super().step(forces)