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

1# fmt: off 

2 

3"""Bussi NVT dynamics class.""" 

4 

5import math 

6 

7import numpy as np 

8 

9from ase import units 

10from ase.md.verlet import VelocityVerlet 

11 

12 

13class Bussi(VelocityVerlet): 

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

15 

16 Based on the paper from Bussi et al., J. Chem. Phys. 126, 014101 (2007) 

17 (also available from https://arxiv.org/abs/0803.4060). 

18 """ 

19 

20 def __init__( 

21 self, 

22 atoms, 

23 timestep, 

24 temperature_K, 

25 taut, 

26 rng=None, 

27 **kwargs, 

28 ): 

29 """ 

30 Parameters 

31 ---------- 

32 atoms : Atoms 

33 The atoms object. 

34 timestep : float 

35 The time step in ASE time units. 

36 temperature_K : float 

37 The desired temperature, in Kelvin. 

38 taut : float 

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

40 rng : RNG object, optional 

41 Random number generator, by default numpy.random. 

42 **kwargs : dict, optional 

43 Additional arguments are passed to 

44 :class:~ase.md.md.MolecularDynamics base class. 

45 """ 

46 super().__init__(atoms, timestep, **kwargs) 

47 

48 self.temp = temperature_K * units.kB 

49 self.taut = taut 

50 if rng is None: 

51 self.rng = np.random 

52 else: 

53 self.rng = rng 

54 

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

56 

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

58 

59 if np.isclose( 

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

61 ): 

62 raise ValueError( 

63 "Initial kinetic energy is zero. " 

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

65 ) 

66 

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

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

69 

70 self.transferred_energy = 0.0 

71 

72 def scale_velocities(self): 

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

74 kinetic_energy = self.atoms.get_kinetic_energy() 

75 alpha = self.calculate_alpha(kinetic_energy) 

76 

77 momenta = self.atoms.get_momenta() 

78 self.atoms.set_momenta(alpha * momenta) 

79 

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

81 

82 def calculate_alpha(self, kinetic_energy): 

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

84 from the Bussi paper.""" 

85 

86 energy_scaling_term = ( 

87 (1 - self._exp_term) 

88 * self.target_kinetic_energy 

89 / kinetic_energy 

90 / self.ndof 

91 ) 

92 

93 # R1 in Eq. (A7) 

94 noisearray = self.rng.standard_normal(size=(1,)) 

95 # ASE mpi interfaces can only broadcast arrays, not scalars 

96 self.comm.broadcast(noisearray, 0) 

97 normal_noise = noisearray[0] 

98 

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

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

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

102 

103 return math.sqrt( 

104 self._exp_term 

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

106 + 2 

107 * normal_noise 

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

109 ) 

110 

111 def step(self, forces=None): 

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

113 self.scale_velocities() 

114 return super().step(forces)