Coverage for /builds/ase/ase/ase/calculators/tip3p.py: 98.37%

1# fmt: off 

2 

3"""TIP3P potential.""" 

4 

5import numpy as np 

6 

7import ase.units as units 

8from ase.calculators.calculator import Calculator, all_changes 

9 

10qH = 0.417 

11sigma0 = 3.15061 

12epsilon0 = 0.1521 * units.kcal / units.mol 

13rOH = 0.9572 

14angleHOH = 104.52 

15thetaHOH = 104.52 / 180 * np.pi # we keep this for backwards compatibility 

16 

17 

18class TIP3P(Calculator): 

19 implemented_properties = ['energy', 'forces'] 

20 nolabel = True 

21 pcpot = None 

22 

23 def __init__(self, rc=5.0, width=1.0): 

24 """TIP3P potential. 

25 

26 rc: float 

27 Cutoff radius for Coulomb part. 

28 width: float 

29 Width for cutoff function for Coulomb part. 

30 """ 

31 self.rc = rc 

32 self.width = width 

33 Calculator.__init__(self) 

34 self.sites_per_mol = 3 

35 

36 def calculate(self, atoms=None, 

37 properties=['energy'], 

38 system_changes=all_changes): 

39 Calculator.calculate(self, atoms, properties, system_changes) 

40 

41 R = self.atoms.positions.reshape((-1, 3, 3)) 

42 Z = self.atoms.numbers 

43 pbc = self.atoms.pbc 

44 cell = self.atoms.cell.diagonal() 

45 nh2o = len(R) 

46 

47 assert (self.atoms.cell == np.diag(cell)).all(), 'not orthorhombic' 

48 assert ((cell >= 2 * self.rc) | ~pbc).all(), 'cutoff too large' # ??? 

49 if Z[0] == 8: 

50 o = 0 

51 else: 

52 o = 2 

53 assert (Z[o::3] == 8).all() 

54 assert (Z[(o + 1) % 3::3] == 1).all() 

55 assert (Z[(o + 2) % 3::3] == 1).all() 

56 

57 charges = np.array([qH, qH, qH]) 

58 charges[o] *= -2 

59 

60 energy = 0.0 

61 forces = np.zeros((3 * nh2o, 3)) 

62 

63 for m in range(nh2o - 1): 

64 DOO = R[m + 1:, o] - R[m, o] 

65 shift = np.zeros_like(DOO) 

66 for i, periodic in enumerate(pbc): 

67 if periodic: 

68 L = cell[i] 

69 shift[:, i] = (DOO[:, i] + L / 2) % L - L / 2 - DOO[:, i] 

70 DOO += shift 

71 d2 = (DOO**2).sum(1) 

72 d = d2**0.5 

73 x1 = d > self.rc - self.width 

74 x2 = d < self.rc 

75 x12 = np.logical_and(x1, x2) 

76 y = (d[x12] - self.rc + self.width) / self.width 

77 t = np.zeros(len(d)) # cutoff function 

78 t[x2] = 1.0 

79 t[x12] -= y**2 * (3.0 - 2.0 * y) 

80 dtdd = np.zeros(len(d)) 

81 dtdd[x12] -= 6.0 / self.width * y * (1.0 - y) 

82 c6 = (sigma0**2 / d2)**3 

83 c12 = c6**2 

84 e = 4 * epsilon0 * (c12 - c6) 

85 energy += np.dot(t, e) 

86 F = (24 * epsilon0 * (2 * c12 - c6) / d2 * t - 

87 e * dtdd / d)[:, np.newaxis] * DOO 

88 forces[m * 3 + o] -= F.sum(0) 

89 forces[m * 3 + 3 + o::3] += F 

90 

91 for j in range(3): 

92 D = R[m + 1:] - R[m, j] + shift[:, np.newaxis] 

93 r2 = (D**2).sum(axis=2) 

94 r = r2**0.5 

95 e = charges[j] * charges / r * units.Hartree * units.Bohr 

96 energy += np.dot(t, e).sum() 

97 F = (e / r2 * t[:, np.newaxis])[:, :, np.newaxis] * D 

98 FOO = -(e.sum(1) * dtdd / d)[:, np.newaxis] * DOO 

99 forces[(m + 1) * 3 + o::3] += FOO 

100 forces[m * 3 + o] -= FOO.sum(0) 

101 forces[(m + 1) * 3:] += F.reshape((-1, 3)) 

102 forces[m * 3 + j] -= F.sum(axis=0).sum(axis=0) 

103 

104 if self.pcpot: 

105 e, f = self.pcpot.calculate(np.tile(charges, nh2o), 

106 self.atoms.positions) 

107 energy += e 

108 forces += f 

109 

110 self.results['energy'] = energy 

111 self.results['forces'] = forces 

112 

113 def embed(self, charges): 

114 """Embed atoms in point-charges.""" 

115 self.pcpot = PointChargePotential(charges) 

116 return self.pcpot 

117 

118 def check_state(self, atoms, tol=1e-15): 

119 system_changes = Calculator.check_state(self, atoms, tol) 

120 if self.pcpot and self.pcpot.mmpositions is not None: 

121 system_changes.append('positions') 

122 return system_changes 

123 

124 def add_virtual_sites(self, positions): 

125 return positions # no virtual sites 

126 

127 def redistribute_forces(self, forces): 

128 return forces 

129 

130 def get_virtual_charges(self, atoms): 

131 charges = np.empty(len(atoms)) 

132 charges[:] = qH 

133 if atoms.numbers[0] == 8: 

134 charges[::3] = -2 * qH 

135 else: 

136 charges[2::3] = -2 * qH 

137 return charges 

138 

139 

140class PointChargePotential: 

141 def __init__(self, mmcharges): 

142 """Point-charge potential for TIP3P. 

143 

144 Only used for testing QMMM. 

145 """ 

146 self.mmcharges = mmcharges 

147 self.mmpositions = None 

148 self.mmforces = None 

149 

150 def set_positions(self, mmpositions, com_pv=None): 

151 self.mmpositions = mmpositions 

152 

153 def calculate(self, qmcharges, qmpositions): 

154 energy = 0.0 

155 self.mmforces = np.zeros_like(self.mmpositions) 

156 qmforces = np.zeros_like(qmpositions) 

157 for C, R, F in zip(self.mmcharges, self.mmpositions, self.mmforces): 

158 d = qmpositions - R 

159 r2 = (d**2).sum(1) 

160 e = units.Hartree * units.Bohr * C * r2**-0.5 * qmcharges 

161 energy += e.sum() 

162 f = (e / r2)[:, np.newaxis] * d 

163 qmforces += f 

164 F -= f.sum(0) 

165 self.mmpositions = None 

166 return energy, qmforces 

167 

168 def get_forces(self, calc): 

169 return self.mmforces