Coverage for gamdpy/interactions/dihedrals.py: 66%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda 

5from .make_fixed_interactions import make_fixed_interactions # bonds is an example of 'fixed' interactions 

6from .interaction import Interaction 

7 

8# Abstract Base Class and type annotation 

9from .interaction import Interaction 

10from gamdpy import Configuration 

11 

12class Dihedrals(Interaction): 

13 

14 def __init__(self, potential, indices, parameters): 

15 self.potential = potential 

16 self.indices = np.array(indices, dtype=np.int32) 

17 self.params = np.array(parameters, dtype=np.float32) 

18 

19 def get_params(self, configuration: Configuration, compute_plan: dict, verbose=False) -> tuple: 

20 self.d_indices = cuda.to_device(self.indices) 

21 self.d_params = cuda.to_device(self.params) 

22 

23 return (self.d_indices, self.d_params) 

24 

25 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags: dict[str,bool], verbose=False): 

26 # Unpack parameters from configuration and compute_plan 

27 D, N = configuration.D, configuration.N 

28 pb, tp, gridsync, UtilizeNIII = [compute_plan[key] for key in ['pb', 'tp', 'gridsync', 'UtilizeNIII']] 

29 num_blocks = (N - 1) // pb + 1 

30 

31 if D != 3: 

32 raise ValueError("Dihedral interactions only support D=3 (three dimensional) systems") 

33 

34 compute_u = compute_flags['U'] 

35 compute_w = compute_flags['W'] 

36 compute_lap = compute_flags['lapU'] 

37 compute_stresses = compute_flags['stresses'] 

38 if compute_stresses: 

39 print('WARNING: computation of stresses is not implemented yet for dihedrals') 

40 

41 if compute_u: 

42 u_id = configuration.sid['U'] 

43 if compute_w: 

44 w_id = configuration.sid['W'] 

45 if compute_lap: 

46 lap_id = configuration.sid['lapU'] 

47 

48 if compute_stresses: 

49 sx_id = configuration.vectors.indices['sx'] 

50 if D > 1: 

51 sy_id = configuration.vectors.indices['sy'] 

52 if D > 2: 

53 sz_id = configuration.vectors.indices['sz'] 

54 if D > 3: 

55 sw_id = configuration.vectors.indices['sw'] 

56 

57 # Unpack indices for vectors and scalars to be compiled into kernel 

58 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']] 

59 #u_id = configuration.sid['U'] 

60 

61 dist_sq_dr_function = numba.njit(configuration.simbox.get_dist_sq_dr_function()) 

62 potential_function = numba.njit(self.potential) 

63 

64 def dihedral_calculator(vectors, scalars, ptype, sim_box, indices, values): 

65 ''' 

66 Algorithm is based on D.C. Rapaport, The Art of Molecular Dynamics Simulations,  

67 Cambridge University Press (1995). 

68 ''' 

69 

70 nparams = 6 

71 p = cuda.local.array(shape=nparams,dtype=numba.float32) 

72 for n in range(nparams): 

73 p[n] = values[indices[4]][n] 

74 

75 one = numba.float32(1.0) 

76 

77 dr_1 = cuda.local.array(shape=D,dtype=numba.float32) 

78 dr_2 = cuda.local.array(shape=D,dtype=numba.float32) 

79 dr_3 = cuda.local.array(shape=D,dtype=numba.float32) 

80 

81 dist_sq_dr_function(vectors[r_id][indices[1]], vectors[r_id][indices[0]], sim_box, dr_1) 

82 dist_sq_dr_function(vectors[r_id][indices[2]], vectors[r_id][indices[1]], sim_box, dr_2) 

83 dist_sq_dr_function(vectors[r_id][indices[3]], vectors[r_id][indices[2]], sim_box, dr_3) 

84 

85 c11 = dr_1[0]*dr_1[0] + dr_1[1]*dr_1[1] + dr_1[2]*dr_1[2] 

86 c12 = dr_1[0]*dr_2[0] + dr_1[1]*dr_2[1] + dr_1[2]*dr_2[2] 

87 c13 = dr_1[0]*dr_3[0] + dr_1[1]*dr_3[1] + dr_1[2]*dr_3[2] 

88 c22 = dr_2[0]*dr_2[0] + dr_2[1]*dr_2[1] + dr_2[2]*dr_2[2] 

89 c23 = dr_2[0]*dr_3[0] + dr_2[1]*dr_3[1] + dr_2[2]*dr_3[2] 

90 c33 = dr_3[0]*dr_3[0] + dr_3[1]*dr_3[1] + dr_3[2]*dr_3[2] 

91 

92 cA = c13*c22 - c12*c23 

93 cB1 = c11*c22 - c12*c12 

94 cB2 = c22*c33 - c23*c23 

95 cD = math.sqrt(cB1*cB2) 

96 cd = cA/cD 

97 

98 t1 = cA 

99 t2 = c11*c23 - c12*c13 

100 t3 = -cB1 

101 t4 = cB2 

102 t5 = c13*c23 - c12*c33 

103 t6 = -cA 

104 cR1 = c12/c22 

105 cR2 = c23/c22 

106 

107 if cd > one: 

108 cd = one 

109 elif cd < -one: 

110 cd = -one 

111 

112 dihedral = math.pi - math.acos(cd) 

113 u, f = potential_function(dihedral, p) 

114 

115 for k in range(3): 

116 f1 = f*c22*(t1*dr_1[k] + t2*dr_2[k] + t3*dr_3[k])/(cD*cB1) 

117 f2 = f*c22*(t4*dr_1[k] + t5*dr_2[k] + t6*dr_3[k])/(cD*cB2) 

118 

119 cuda.atomic.add(vectors, (f_id, indices[0], k), f1) # Force 

120 cuda.atomic.add(vectors, (f_id, indices[1], k), -(one + cR1)*f1 + cR2*f2) 

121 cuda.atomic.add(vectors, (f_id, indices[2], k), cR1*f1 - (one + cR2)*f2) 

122 cuda.atomic.add(vectors, (f_id, indices[3], k), f2) 

123 

124 u_per_part = numba.float32(0.25)*u 

125 

126 if compute_u: 

127 for n in range(4): 

128 cuda.atomic.add(scalars, (indices[n], u_id), u_per_part) 

129 

130 return 

131 

132 return make_fixed_interactions(configuration, dihedral_calculator, compute_plan, verbose=False) 

133 

134 

135 def get_exclusions(self, configuration, max_number_exclusions=20): 

136 

137 exclusions = np.zeros( (configuration.N, max_number_exclusions+1), dtype=np.int32 ) 

138 

139 nangles = len(self.indices) 

140 for n in range(nangles): 

141 pidx = self.indices[n][:4] 

142 for k in range(4): 

143 

144 offset = exclusions[pidx[k]][-1] 

145 if offset > max_number_exclusions-2: 

146 raise ValueError("Number of max. exclusion breached") 

147 

148 if k==0: 

149 idx = [1, 2, 3] 

150 elif k==1: 

151 idx = [0, 2, 3] 

152 elif k==2: 

153 idx = [0, 1, 3] 

154 else: 

155 idx = [0, 1, 2] 

156 

157 for kk in idx: 

158 if dihedrals_entry_not_exists(pidx[kk], exclusions[pidx[k]], offset): 

159 exclusions[pidx[k]][offset] = pidx[kk] 

160 offset += 1 

161 

162 exclusions[pidx[k]][-1] = offset 

163 

164 return exclusions 

165 

166 

167 def get_dihedral(self, dihedral_idx, configuration): 

168 

169 pidx = self.indices[dihedral_idx][:4] 

170 

171 r0 = configuration['r'][pidx[0]] 

172 r1 = configuration['r'][pidx[1]] 

173 r2 = configuration['r'][pidx[2]] 

174 r3 = configuration['r'][pidx[3]] 

175 

176 dr_1 = dihedrals_get_dist_vector(r1, r0, configuration.simbox) 

177 dr_2 = dihedrals_get_dist_vector(r2, r1, configuration.simbox) 

178 dr_3 = dihedrals_get_dist_vector(r3, r2, configuration.simbox) 

179 

180 c11 = dr_1[0]*dr_1[0] + dr_1[1]*dr_1[1] + dr_1[2]*dr_1[2] 

181 c12 = dr_1[0]*dr_2[0] + dr_1[1]*dr_2[1] + dr_1[2]*dr_2[2] 

182 c13 = dr_1[0]*dr_3[0] + dr_1[1]*dr_3[1] + dr_1[2]*dr_3[2] 

183 c22 = dr_2[0]*dr_2[0] + dr_2[1]*dr_2[1] + dr_2[2]*dr_2[2] 

184 c23 = dr_2[0]*dr_3[0] + dr_2[1]*dr_3[1] + dr_2[2]*dr_3[2] 

185 c33 = dr_3[0]*dr_3[0] + dr_3[1]*dr_3[1] + dr_3[2]*dr_3[2] 

186 

187 cA = c13*c22 - c12*c23 

188 cB1 = c11*c22 - c12*c12 

189 cB2 = c22*c33 - c23*c23 

190 cD = math.sqrt(cB1*cB2) 

191 cc = cA/cD 

192 

193 dihedral = math.pi - math.acos(cc) 

194 

195 return dihedral 

196 

197# Helpers (copies of angles helpers: This should be centralized somehow) 

198def dihedrals_get_dist_vector(ri, rj, simbox): 

199 """ 

200 Function for checking that angles are correctly calculated. Assumes Orthorhombic box 

201 """ 

202 dr = np.zeros(3) 

203 lengths = simbox.get_lengths() 

204 for k in range(simbox.D): 

205 dr[k] = ri[k] - rj[k] 

206 box_k = lengths[k] 

207 #PP 

208 dr[k] += (-box_k if 2.0*dr[k] > +box_k else (+box_k if 2.0*dr[k] < -box_k else 0.0)) 

209 

210 return dr 

211 

212def dihedrals_entry_not_exists(idx, exclusion_list, nentries): 

213 

214 for n in range(nentries): 

215 if exclusion_list[n]==idx: 

216 return False 

217 

218 return True 

219