1from collections import defaultdict 

2 

3import numpy as np 

4import kimpy 

5from kimpy import neighlist 

6from ase.neighborlist import neighbor_list 

7from ase import Atom 

8 

9from .kimpy_wrappers import check_call_wrapper, c_int, c_double 

10 

11 

12class NeighborList: 

13 

14 kimpy_arrays = { 

15 "num_particles": c_int, 

16 "coords": c_double, 

17 "particle_contributing": c_int, 

18 "species_code": c_int, 

19 "cutoffs": c_double, 

20 "padding_image_of": c_int, 

21 "need_neigh": c_int, 

22 } 

23 

24 def __setattr__(self, name, value): 

25 """ 

26 Override assignment to any of the attributes listed in 

27 kimpy_arrays to automatically cast the object to a numpy array. 

28 This is done to avoid a ton of explicit numpy.array() calls (and 

29 the possibility that we forget to do the cast). It is important 

30 to use np.asarray() here instead of np.array() because using the 

31 latter will mean that incrementation (+=) will create a new 

32 object that the reference is bound to, which becomes a problem 

33 if update_compute_args isn't called to reregister the 

34 corresponding address with the KIM API. 

35 """ 

36 if name in self.kimpy_arrays and value is not None: 

37 value = np.asarray(value, dtype=self.kimpy_arrays[name]) 

38 self.__dict__[name] = value 

39 

40 def __init__( 

41 self, 

42 neigh_skin_ratio, 

43 model_influence_dist, 

44 model_cutoffs, 

45 padding_not_require_neigh, 

46 debug, 

47 ): 

48 

49 self.set_neigh_parameters( 

50 neigh_skin_ratio, 

51 model_influence_dist, 

52 model_cutoffs, 

53 padding_not_require_neigh, 

54 ) 

55 

56 self.debug = debug 

57 

58 if self.debug: 

59 print() 

60 print("Calculator skin: {}".format(self.skin)) 

61 print(f"Model influence distance: {model_influence_dist}") 

62 print( 

63 "Calculator influence distance (including skin distance): {}" 

64 "".format(self.influence_dist) 

65 ) 

66 print("Number of cutoffs: {}".format(model_cutoffs.size)) 

67 print("Model cutoffs: {}".format(model_cutoffs)) 

68 print( 

69 "Calculator cutoffs (including skin distance): {}" 

70 "".format(self.cutoffs) 

71 ) 

72 print( 

73 "Model needs neighbors of padding atoms: {}" 

74 "".format(self.padding_need_neigh) 

75 ) 

76 print() 

77 

78 # Attributes to be set by subclasses 

79 self.neigh = None 

80 self.num_contributing_particles = None 

81 self.padding_image_of = None 

82 self.num_particles = None 

83 self.coords = None 

84 self.particle_contributing = None 

85 self.species_code = None 

86 self.need_neigh = None 

87 self.last_update_positions = None 

88 

89 def set_neigh_parameters( 

90 self, 

91 neigh_skin_ratio, 

92 model_influence_dist, 

93 model_cutoffs, 

94 padding_not_require_neigh, 

95 ): 

96 self.skin = neigh_skin_ratio * model_influence_dist 

97 self.influence_dist = model_influence_dist + self.skin 

98 self.cutoffs = model_cutoffs + self.skin 

99 self.padding_need_neigh = not padding_not_require_neigh.all() 

100 

101 def update_kim_coords(self, atoms): 

102 """Update atomic positions in self.coords, which is where the KIM 

103 API will look to find them in order to pass them to the model. 

104 """ 

105 if self.padding_image_of.size != 0: 

106 disp_contrib = atoms.positions - self.coords[: len(atoms)] 

107 disp_pad = disp_contrib[self.padding_image_of] 

108 self.coords += np.concatenate((disp_contrib, disp_pad)) 

109 else: 

110 np.copyto(self.coords, atoms.positions) 

111 

112 if self.debug: 

113 print("Debug: called update_kim_coords") 

114 print() 

115 

116 def need_neigh_update(self, atoms, system_changes): 

117 need_neigh_update = True 

118 if len(system_changes) == 1 and "positions" in system_changes: 

119 # Only position changes 

120 if self.last_update_positions is not None: 

121 a = self.last_update_positions 

122 b = atoms.positions 

123 if a.shape == b.shape: 

124 delta = np.linalg.norm(a - b, axis=1) 

125 # Indices of the two largest elements 

126 ind = np.argpartition(delta, -2)[-2:] 

127 if sum(delta[ind]) <= self.skin: 

128 need_neigh_update = False 

129 

130 return need_neigh_update 

131 

132 

133class ASENeighborList(NeighborList): 

134 def __init__( 

135 self, 

136 compute_args, 

137 neigh_skin_ratio, 

138 model_influence_dist, 

139 model_cutoffs, 

140 padding_not_require_neigh, 

141 debug, 

142 ): 

143 super().__init__( 

144 neigh_skin_ratio, 

145 model_influence_dist, 

146 model_cutoffs, 

147 padding_not_require_neigh, 

148 debug, 

149 ) 

150 

151 self.neigh = {} 

152 compute_args.set_callback( 

153 kimpy.compute_callback_name.GetNeighborList, self.get_neigh, 

154 self.neigh 

155 ) 

156 

157 @staticmethod 

158 def get_neigh(data, cutoffs, neighbor_list_index, particle_number): 

159 """Retrieves the neighbors of each atom using ASE's native neighbor 

160 list library 

161 """ 

162 # We can only return neighbors of particles that were stored 

163 number_of_particles = data["num_particles"] 

164 if particle_number >= number_of_particles or particle_number < 0: 

165 return (np.array([]), 1) 

166 

167 neighbors = data["neighbors"][neighbor_list_index][particle_number] 

168 return (neighbors, 0) 

169 

170 def build(self, orig_atoms): 

171 """Build the ASE neighbor list and return an Atoms object with 

172 all of the neighbors added. First a neighbor list is created 

173 from ase.neighbor_list, having only information about the 

174 neighbors of the original atoms. If neighbors of padding atoms 

175 are required, they are calculated using information from the 

176 first neighbor list. 

177 """ 

178 syms = orig_atoms.get_chemical_symbols() 

179 orig_num_atoms = len(orig_atoms) 

180 orig_pos = orig_atoms.get_positions() 

181 

182 # New atoms object that will contain the contributing atoms plus 

183 # the padding atoms 

184 new_atoms = orig_atoms.copy() 

185 

186 neigh_list = defaultdict(list) 

187 neigh_dists = defaultdict(list) 

188 

189 # Information for padding atoms 

190 padding_image_of = [] 

191 padding_shifts = [] 

192 

193 # Ask ASE to build the neighbor list using the influence distance. 

194 # This is not a neighbor list for each atom, but rather a listing 

195 # of all neighbor pairs that exist. Atoms with no neighbors will 

196 # not show up. 

197 ( 

198 neigh_indices_i, 

199 neigh_indices_j, 

200 relative_pos, 

201 neigh_cell_offsets, 

202 dists, 

203 ) = neighbor_list("ijDSd", orig_atoms, self.influence_dist) 

204 

205 # Loop over all neighbor pairs. Because this loop will generally 

206 # include image atoms (for periodic systems), we keep track of 

207 # which atoms/images we've accounted for in the `used` dictionary. 

208 used = dict() 

209 for neigh_i, neigh_j, rel_pos, offset, dist in zip( 

210 neigh_indices_i, neigh_indices_j, 

211 relative_pos, neigh_cell_offsets, dists 

212 ): 

213 # Get neighbor position of neighbor 

214 # (mapped back into unit cell, so this may overlap with other atoms) 

215 wrapped_pos = orig_pos[neigh_i] + rel_pos 

216 

217 shift = tuple(offset) 

218 uniq_index = (neigh_j,) + shift 

219 if shift == (0, 0, 0): 

220 # This atom is in the unit cell, i.e. it is contributing 

221 neigh_list[neigh_i].append(neigh_j) 

222 neigh_dists[neigh_i].append(dist) 

223 if uniq_index not in used: 

224 used[uniq_index] = neigh_j 

225 else: 

226 # This atom is not in the unit cell, i.e. it is padding 

227 if uniq_index not in used: 

228 # Add the neighbor as a padding atom 

229 used[uniq_index] = len(new_atoms) 

230 new_atoms.append(Atom(syms[neigh_j], position=wrapped_pos)) 

231 padding_image_of.append(neigh_j) 

232 padding_shifts.append(offset) 

233 neigh_list[neigh_i].append(used[uniq_index]) 

234 neigh_dists[neigh_i].append(dist) 

235 neighbor_list_size = orig_num_atoms 

236 

237 # Add neighbors of padding atoms if the potential requires them 

238 if self.padding_need_neigh: 

239 neighbor_list_size = len(new_atoms) 

240 inv_used = dict((v, k) for k, v in used.items()) 

241 # Loop over all the neighbors (k) and the image of that neighbor 

242 # in the cell (neigh) 

243 for k, neigh in enumerate(padding_image_of): 

244 # Shift from original atom in cell to neighbor 

245 shift = padding_shifts[k] 

246 for orig_neigh, orig_dist in zip( 

247 neigh_list[neigh], neigh_dists[neigh]): 

248 # Get the shift of the neighbor of the original atom 

249 orig_shift = inv_used[orig_neigh][1:] 

250 

251 # Apply sum of original shift and current shift to 

252 # neighbors of original atom 

253 total_shift = orig_shift + shift 

254 

255 # Get the image in the cell of the original neighbor 

256 if orig_neigh <= orig_num_atoms - 1: 

257 orig_neigh_image = orig_neigh 

258 else: 

259 orig_neigh_image = padding_image_of[orig_neigh - 

260 orig_num_atoms] 

261 

262 # If the original image with the total shift has been 

263 # used before then it is also a neighbor of this atom 

264 uniq_index = (orig_neigh_image,) + tuple(total_shift) 

265 if uniq_index in used: 

266 neigh_list[k + orig_num_atoms].append(used[uniq_index]) 

267 neigh_dists[k + orig_num_atoms].append(orig_dist) 

268 

269 # If the model has multiple cutoffs, we need to return neighbor lists 

270 # corresponding to each of them 

271 neigh_lists = [] 

272 for cut in self.cutoffs: 

273 neigh_list = [ 

274 np.array(neigh_list[k], dtype=c_int)[neigh_dists[k] <= cut] 

275 for k in range(neighbor_list_size) 

276 ] 

277 neigh_lists.append(neigh_list) 

278 

279 self.padding_image_of = padding_image_of 

280 

281 self.neigh["neighbors"] = neigh_lists 

282 self.neigh["num_particles"] = neighbor_list_size 

283 

284 return new_atoms 

285 

286 def update(self, orig_atoms, species_map): 

287 """Create the neighbor list along with the other required 

288 parameters (which are stored as instance attributes). The 

289 required parameters are: 

290 

291 - num_particles 

292 - coords 

293 - particle_contributing 

294 - species_code 

295 

296 Note that the KIM API requires a neighbor list that has indices 

297 corresponding to each atom. 

298 """ 

299 

300 # Information of original atoms 

301 self.num_contributing_particles = len(orig_atoms) 

302 

303 new_atoms = self.build(orig_atoms) 

304 

305 # Save the number of atoms and all their neighbors and positions 

306 num_atoms = len(new_atoms) 

307 num_padding = num_atoms - self.num_contributing_particles 

308 self.num_particles = [num_atoms] 

309 self.coords = new_atoms.get_positions() 

310 

311 # Save which coordinates are from original atoms and which are from 

312 # neighbors using a mask 

313 indices_mask = [1] * self.num_contributing_particles + [0] * num_padding 

314 self.particle_contributing = indices_mask 

315 

316 # Species support and code 

317 try: 

318 self.species_code = [ 

319 species_map[s] for s in new_atoms.get_chemical_symbols() 

320 ] 

321 except KeyError as e: 

322 raise RuntimeError( 

323 "Species not supported by KIM model; {}".format( 

324 str(e))) 

325 

326 self.last_update_positions = orig_atoms.get_positions() 

327 

328 if self.debug: 

329 print("Debug: called update_ase_neigh") 

330 print() 

331 

332 

333class KimpyNeighborList(NeighborList): 

334 def __init__( 

335 self, 

336 compute_args, 

337 neigh_skin_ratio, 

338 model_influence_dist, 

339 model_cutoffs, 

340 padding_not_require_neigh, 

341 debug, 

342 ): 

343 super().__init__( 

344 neigh_skin_ratio, 

345 model_influence_dist, 

346 model_cutoffs, 

347 padding_not_require_neigh, 

348 debug, 

349 ) 

350 

351 self.neigh = neighlist.NeighList() 

352 

353 compute_args.set_callback_pointer( 

354 kimpy.compute_callback_name.GetNeighborList, 

355 neighlist.get_neigh_kim(), 

356 self.neigh, 

357 ) 

358 

359 @check_call_wrapper 

360 def build(self): 

361 return self.neigh.build( 

362 self.coords, self.influence_dist, self.cutoffs, self.need_neigh 

363 ) 

364 

365 @check_call_wrapper 

366 def create_paddings( 

367 self, cell, pbc, contributing_coords, contributing_species_code 

368 ): 

369 # Cast things passed through kimpy to numpy arrays 

370 cell = np.asarray(cell, dtype=c_double) 

371 pbc = np.asarray(pbc, dtype=c_int) 

372 contributing_coords = np.asarray(contributing_coords, dtype=c_double) 

373 

374 return neighlist.create_paddings( 

375 self.influence_dist, 

376 cell, 

377 pbc, 

378 contributing_coords, 

379 contributing_species_code, 

380 ) 

381 

382 def update(self, atoms, species_map): 

383 """Create the neighbor list along with the other required 

384 parameters (which are stored as instance attributes). The 

385 required parameters are: 

386 

387 - num_particles 

388 - coords 

389 - particle_contributing 

390 - species_code 

391 

392 Note that the KIM API requires a neighbor list that has indices 

393 corresponding to each atom. 

394 """ 

395 

396 # Get info from Atoms object 

397 cell = np.asarray(atoms.get_cell(), dtype=c_double) 

398 pbc = np.asarray(atoms.get_pbc(), dtype=c_int) 

399 contributing_coords = np.asarray(atoms.get_positions(), dtype=c_double) 

400 self.num_contributing_particles = atoms.get_global_number_of_atoms() 

401 num_contributing = self.num_contributing_particles 

402 

403 # Species support and code 

404 try: 

405 contributing_species_code = np.array( 

406 [species_map[s] for s in atoms.get_chemical_symbols()], 

407 dtype=c_int 

408 ) 

409 except KeyError as e: 

410 raise RuntimeError( 

411 "Species not supported by KIM model; {}".format( 

412 str(e))) 

413 

414 if pbc.any(): # Need padding atoms 

415 # Create padding atoms 

416 

417 ( 

418 padding_coords, 

419 padding_species_code, 

420 self.padding_image_of, 

421 ) = self.create_paddings( 

422 cell, pbc, contributing_coords, contributing_species_code 

423 ) 

424 num_padding = padding_species_code.size 

425 

426 self.num_particles = [num_contributing + num_padding] 

427 self.coords = np.concatenate((contributing_coords, padding_coords)) 

428 self.species_code = np.concatenate( 

429 (contributing_species_code, padding_species_code) 

430 ) 

431 self.particle_contributing = [1] * \ 

432 num_contributing + [0] * num_padding 

433 self.need_neigh = [1] * self.num_particles[0] 

434 if not self.padding_need_neigh: 

435 self.need_neigh[num_contributing:] = 0 

436 

437 else: # Do not need padding atoms 

438 self.padding_image_of = [] 

439 self.num_particles = [num_contributing] 

440 self.coords = contributing_coords 

441 self.species_code = contributing_species_code 

442 self.particle_contributing = [1] * num_contributing 

443 self.need_neigh = self.particle_contributing 

444 

445 # Create neighborlist 

446 self.build() 

447 

448 self.last_update_positions = atoms.get_positions() 

449 

450 if self.debug: 

451 print("Debug: called update_kimpy_neigh") 

452 print()