Coverage for /builds/alexhroom/ase/ase/md/bussi.py: 100.00%

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

2 

3import math 

4 

5import numpy as np 

6from ase import units 

7from ase.md.md import MolecularDynamics 

8from ase.parallel import world 

9 

10 

11class Bussi(MolecularDynamics): 

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__( 

42 atoms, 

43 timestep, 

44 **md_kwargs, 

45 ) 

46 

47 self.temp = temperature_K * units.kB 

48 self.taut = taut 

49 self.communicator = world 

50 self.rng = rng 

51 

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

53 

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

55 

56 if np.isclose( 

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

58 ): 

59 raise ValueError( 

60 "Initial kinetic energy is zero. " 

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

62 ) 

63 

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

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

66 

67 self.transferred_energy = 0.0 

68 

69 def scale_velocities(self): 

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

71 kinetic_energy = self.atoms.get_kinetic_energy() 

72 alpha = self.calculate_alpha(kinetic_energy) 

73 

74 momenta = self.atoms.get_momenta() 

75 self.atoms.set_momenta(alpha * momenta) 

76 

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

78 

79 def calculate_alpha(self, kinetic_energy): 

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

81 from the Bussi paper.""" 

82 

83 energy_scaling_term = ( 

84 (1 - self._exp_term) 

85 * self.target_kinetic_energy 

86 / kinetic_energy 

87 / self.ndof 

88 ) 

89 

90 # R1 in Eq. (A7) 

91 normal_noise = self.rng.standard_normal() 

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

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

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

95 

96 return math.sqrt( 

97 self._exp_term 

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

99 + 2 

100 * normal_noise 

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

102 ) 

103 

104 def step(self, forces=None): 

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

106 if forces is None: 

107 forces = self.atoms.get_forces(md=True) 

108 

109 self.scale_velocities() 

110 

111 self.atoms.set_momenta( 

112 self.atoms.get_momenta() + 0.5 * self.dt * forces 

113 ) 

114 momenta = self.atoms.get_momenta() 

115 

116 self.atoms.set_positions( 

117 self.atoms.positions + self.dt * momenta / self._masses 

118 ) 

119 

120 forces = self.atoms.get_forces(md=True) 

121 

122 self.atoms.set_momenta(momenta + 0.5 * self.dt * forces) 

123 

124 return forces