Coverage for gamdpy/interactions/angles.py: 55%

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 

6 

7# Abstract Base Class and type annotation 

8from .interaction import Interaction 

9from gamdpy import Configuration 

10 

11class Angles(Interaction): 

12 

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

14 self.potential = potential 

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

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

17 

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 return (self.d_indices, self.d_params) 

23 

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

25 # Unpack parameters from configuration and compute_plan 

26 D, N = configuration.D, configuration.N 

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

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

29 

30 if D != 3: 

31 raise ValueError("Angle interactions only support D=3 (three dimensional) simulations") 

32 

33 compute_u = compute_flags['U'] 

34 compute_w = compute_flags['W'] 

35 compute_lap = compute_flags['lapU'] 

36 compute_stresses = compute_flags['stresses'] 

37 if compute_stresses: 

38 print('WARNING: computation of stresses is not implemented yet for angles') 

39 

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

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

42 u_id = configuration.sid['U'] 

43 

44 if compute_u: 

45 u_id = configuration.sid['U'] 

46 if compute_w: 

47 w_id = configuration.sid['W'] 

48 if compute_lap: 

49 lap_id = configuration.sid['lapU'] 

50 

51 if compute_stresses: 

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

53 if D > 1: 

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

55 if D > 2: 

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

57 if D > 3: 

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

59 

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

61 potential_function = numba.njit(self.potential) 

62 

63 def angle_calculator(vectors, scalars, ptype, sim_box, indices, values): 

64 ''' 

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

66 Cambridge University Press (1995). 

67 ''' 

68 

69 nparams = 2 #numba.int32(values.shape[1]) 

70 params = cuda.local.array(shape=nparams, dtype=numba.float32) 

71 

72 for n in range(nparams): 

73 params[n] = values[indices[3]][n] 

74 

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

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

77 

78 one = numba.float32(1.0) 

79 

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

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

82 

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

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

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

86 

87 cD = math.sqrt(c11*c22) 

88 ca = c12/cD 

89 

90 if ca > one: 

91 ca = one 

92 elif ca < -one: 

93 ca = -one 

94 angle = math.acos(ca) 

95 

96 u, f = potential_function(angle, params) 

97 

98 for k in range(D): 

99 f_1 = f*( (c12/c11)*dr_1[k] - dr_2[k] )/cD 

100 f_2 = f*( dr_1[k] - (c12/c22)*dr_2[k] )/cD 

101 

102 cuda.atomic.add(vectors, (f_id, indices[0], k), f_1) 

103 cuda.atomic.add(vectors, (f_id, indices[1], k), -f_1-f_2) 

104 cuda.atomic.add(vectors, (f_id, indices[2], k), f_2) 

105 

106 onethird = numba.float32(1.0/3.0)*u; 

107 cuda.atomic.add(scalars, (indices[0], u_id), onethird) 

108 cuda.atomic.add(scalars, (indices[1], u_id), onethird) 

109 cuda.atomic.add(scalars, (indices[2], u_id), onethird) 

110 

111 return 

112 

113 return make_fixed_interactions(configuration, angle_calculator, compute_plan, verbose=False) 

114 

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

116 

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

118 

119 nangles = len(self.indices) 

120 for n in range(nangles): 

121 pidx = self.indices[n][:3] 

122 for k in range(3): 

123 

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

125 if offset > max_number_exclusions-2: 

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

127 

128 if k==0: 

129 idx = [1, 2] 

130 elif k==1: 

131 idx = [0, 2] 

132 else: 

133 idx = [0, 1] 

134 

135 for kk in idx: 

136 if angles_entry_not_exists(pidx[1], exclusions[pidx[k]],offset): 

137 exclusions[pidx[k]][offset] = pidx[1] 

138 offset += 1 

139 

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

141 

142 return exclusions 

143 

144 def get_angle(self, angle_idx, configuration): 

145 

146 pidx = self.indices[angle_idx][:3] 

147 

148 r1 = configuration['r'][pidx[0]] 

149 r2 = configuration['r'][pidx[1]] 

150 r3 = configuration['r'][pidx[2]] 

151 

152 dr_1 = angles_get_dist_vector(r1, r2, configuration.simbox) 

153 dr_2 = angles_get_dist_vector(r3, r2, configuration.simbox) 

154 

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

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

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

158 

159 cD = math.sqrt(c11*c22) 

160 cc = c12/cD 

161 

162 angle = math.acos(cc) 

163 

164 return angle 

165 

166# Helpers  

167def angles_get_dist_vector(ri, rj, simbox): 

168 """ 

169 Assumes Orthorhombic box 

170 """ 

171 dr = np.zeros(3) 

172 lengths = simbox.get_lengths() 

173 for k in range(simbox.D): 

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

175 box_k = lengths[k] 

176 #PP 

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

178 

179 return dr 

180 

181def angles_entry_not_exists(idx, exclusion_list, nentries): 

182 

183 for n in range(nentries): 

184 if exclusion_list[n]==idx: 

185 return False 

186 

187 return True