1import numpy as np 

2 

3from ase.calculators.calculator import Calculator 

4from ase.utils import ff 

5 

6 

7class ForceField(Calculator): 

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

9 nolabel = True 

10 

11 def __init__(self, morses=None, bonds=None, angles=None, dihedrals=None, 

12 vdws=None, coulombs=None, **kwargs): 

13 Calculator.__init__(self, **kwargs) 

14 if (morses is None and 

15 bonds is None and 

16 angles is None and 

17 dihedrals is None and 

18 vdws is None and 

19 coulombs is None): 

20 raise ImportError("At least one of morses, bonds, angles, dihedrals," 

21 "vdws or coulombs lists must be defined!") 

22 if morses is None: 

23 self.morses = [] 

24 else: 

25 self.morses = morses 

26 if bonds is None: 

27 self.bonds = [] 

28 else: 

29 self.bonds = bonds 

30 if angles is None: 

31 self.angles = [] 

32 else: 

33 self.angles = angles 

34 if dihedrals is None: 

35 self.dihedrals = [] 

36 else: 

37 self.dihedrals = dihedrals 

38 if vdws is None: 

39 self.vdws = [] 

40 else: 

41 self.vdws = vdws 

42 if coulombs is None: 

43 self.coulombs = [] 

44 else: 

45 self.coulombs = coulombs 

46 

47 def calculate(self, atoms, properties, system_changes): 

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

49 if system_changes: 

50 for name in ['energy', 'forces', 'hessian']: 

51 self.results.pop(name, None) 

52 if 'energy' not in self.results: 

53 energy = 0.0 

54 for morse in self.morses: 

55 i, j, e = ff.get_morse_potential_value(atoms, morse) 

56 energy += e 

57 for bond in self.bonds: 

58 i, j, e = ff.get_bond_potential_value(atoms, bond) 

59 energy += e 

60 for angle in self.angles: 

61 i, j, k, e = ff.get_angle_potential_value(atoms, angle) 

62 energy += e 

63 for dihedral in self.dihedrals: 

64 i, j, k, l, e = ff.get_dihedral_potential_value( 

65 atoms, dihedral) 

66 energy += e 

67 for vdw in self.vdws: 

68 i, j, e = ff.get_vdw_potential_value(atoms, vdw) 

69 energy += e 

70 for coulomb in self.coulombs: 

71 i, j, e = ff.get_coulomb_potential_value(atoms, coulomb) 

72 energy += e 

73 self.results['energy'] = energy 

74 if 'forces' not in self.results: 

75 forces = np.zeros(3 * len(atoms)) 

76 for morse in self.morses: 

77 i, j, g = ff.get_morse_potential_gradient(atoms, morse) 

78 limits = get_limits([i, j]) 

79 for gb, ge, lb, le in limits: 

80 forces[gb:ge] -= g[lb:le] 

81 for bond in self.bonds: 

82 i, j, g = ff.get_bond_potential_gradient(atoms, bond) 

83 limits = get_limits([i, j]) 

84 for gb, ge, lb, le in limits: 

85 forces[gb:ge] -= g[lb:le] 

86 for angle in self.angles: 

87 i, j, k, g = ff.get_angle_potential_gradient(atoms, angle) 

88 limits = get_limits([i, j, k]) 

89 for gb, ge, lb, le in limits: 

90 forces[gb:ge] -= g[lb:le] 

91 for dihedral in self.dihedrals: 

92 i, j, k, l, g = ff.get_dihedral_potential_gradient( 

93 atoms, dihedral) 

94 limits = get_limits([i, j, k, l]) 

95 for gb, ge, lb, le in limits: 

96 forces[gb:ge] -= g[lb:le] 

97 for vdw in self.vdws: 

98 i, j, g = ff.get_vdw_potential_gradient(atoms, vdw) 

99 limits = get_limits([i, j]) 

100 for gb, ge, lb, le in limits: 

101 forces[gb:ge] -= g[lb:le] 

102 for coulomb in self.coulombs: 

103 i, j, g = ff.get_coulomb_potential_gradient(atoms, coulomb) 

104 limits = get_limits([i, j]) 

105 for gb, ge, lb, le in limits: 

106 forces[gb:ge] -= g[lb:le] 

107 self.results['forces'] = np.reshape(forces, (len(atoms), 3)) 

108 if 'hessian' not in self.results: 

109 hessian = np.zeros((3 * len(atoms), 3 * len(atoms))) 

110 for morse in self.morses: 

111 i, j, h = ff.get_morse_potential_hessian(atoms, morse) 

112 limits = get_limits([i, j]) 

113 for gb1, ge1, lb1, le1 in limits: 

114 for gb2, ge2, lb2, le2 in limits: 

115 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

116 for bond in self.bonds: 

117 i, j, h = ff.get_bond_potential_hessian(atoms, bond) 

118 limits = get_limits([i, j]) 

119 for gb1, ge1, lb1, le1 in limits: 

120 for gb2, ge2, lb2, le2 in limits: 

121 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

122 for angle in self.angles: 

123 i, j, k, h = ff.get_angle_potential_hessian(atoms, angle) 

124 limits = get_limits([i, j, k]) 

125 for gb1, ge1, lb1, le1 in limits: 

126 for gb2, ge2, lb2, le2 in limits: 

127 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

128 for dihedral in self.dihedrals: 

129 i, j, k, l, h = ff.get_dihedral_potential_hessian( 

130 atoms, dihedral) 

131 limits = get_limits([i, j, k, l]) 

132 for gb1, ge1, lb1, le1 in limits: 

133 for gb2, ge2, lb2, le2 in limits: 

134 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

135 for vdw in self.vdws: 

136 i, j, h = ff.get_vdw_potential_hessian(atoms, vdw) 

137 limits = get_limits([i, j]) 

138 for gb1, ge1, lb1, le1 in limits: 

139 for gb2, ge2, lb2, le2 in limits: 

140 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

141 for coulomb in self.coulombs: 

142 i, j, h = ff.get_coulomb_potential_hessian(atoms, coulomb) 

143 limits = get_limits([i, j]) 

144 for gb1, ge1, lb1, le1 in limits: 

145 for gb2, ge2, lb2, le2 in limits: 

146 hessian[gb1:ge1, gb2:ge2] += h[lb1:le1, lb2:le2] 

147 self.results['hessian'] = hessian 

148 

149 def get_hessian(self, atoms=None): 

150 return self.get_property('hessian', atoms) 

151 

152 

153def get_limits(indices): 

154 gstarts = [] 

155 gstops = [] 

156 lstarts = [] 

157 lstops = [] 

158 for l, g in enumerate(indices): 

159 g3, l3 = 3 * g, 3 * l 

160 gstarts.append(g3) 

161 gstops.append(g3 + 3) 

162 lstarts.append(l3) 

163 lstops.append(l3 + 3) 

164 return zip(gstarts, gstops, lstarts, lstops)