Coverage for /builds/kinetik161/ase/ase/calculators/kim/neighborlist.py: 52.88%
191 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1from collections import defaultdict
3import kimpy
4import numpy as np
5from kimpy import neighlist
7from ase import Atom
8from ase.neighborlist import neighbor_list
10from .kimpy_wrappers import c_double, c_int, check_call_wrapper
13class NeighborList:
15 kimpy_arrays = {
16 "num_particles": c_int,
17 "coords": c_double,
18 "particle_contributing": c_int,
19 "species_code": c_int,
20 "cutoffs": c_double,
21 "padding_image_of": c_int,
22 "need_neigh": c_int,
23 }
25 def __setattr__(self, name, value):
26 """
27 Override assignment to any of the attributes listed in
28 kimpy_arrays to automatically cast the object to a numpy array.
29 This is done to avoid a ton of explicit numpy.array() calls (and
30 the possibility that we forget to do the cast). It is important
31 to use np.asarray() here instead of np.array() because using the
32 latter will mean that incrementation (+=) will create a new
33 object that the reference is bound to, which becomes a problem
34 if update_compute_args isn't called to reregister the
35 corresponding address with the KIM API.
36 """
37 if name in self.kimpy_arrays and value is not None:
38 value = np.asarray(value, dtype=self.kimpy_arrays[name])
39 self.__dict__[name] = value
41 def __init__(
42 self,
43 neigh_skin_ratio,
44 model_influence_dist,
45 model_cutoffs,
46 padding_not_require_neigh,
47 debug,
48 ):
50 self.set_neigh_parameters(
51 neigh_skin_ratio,
52 model_influence_dist,
53 model_cutoffs,
54 padding_not_require_neigh,
55 )
57 self.debug = debug
59 if self.debug:
60 print()
61 print(f"Calculator skin: {self.skin}")
62 print(f"Model influence distance: {model_influence_dist}")
63 print(
64 "Calculator influence distance (including skin distance): {}"
65 "".format(self.influence_dist)
66 )
67 print(f"Number of cutoffs: {model_cutoffs.size}")
68 print(f"Model cutoffs: {model_cutoffs}")
69 print(
70 "Calculator cutoffs (including skin distance): {}"
71 "".format(self.cutoffs)
72 )
73 print(
74 "Model needs neighbors of padding atoms: {}"
75 "".format(self.padding_need_neigh)
76 )
77 print()
79 # Attributes to be set by subclasses
80 self.neigh = None
81 self.num_contributing_particles = None
82 self.padding_image_of = None
83 self.num_particles = None
84 self.coords = None
85 self.particle_contributing = None
86 self.species_code = None
87 self.need_neigh = None
88 self.last_update_positions = None
90 def set_neigh_parameters(
91 self,
92 neigh_skin_ratio,
93 model_influence_dist,
94 model_cutoffs,
95 padding_not_require_neigh,
96 ):
97 self.skin = neigh_skin_ratio * model_influence_dist
98 self.influence_dist = model_influence_dist + self.skin
99 self.cutoffs = model_cutoffs + self.skin
100 self.padding_need_neigh = not padding_not_require_neigh.all()
102 def update_kim_coords(self, atoms):
103 """Update atomic positions in self.coords, which is where the KIM
104 API will look to find them in order to pass them to the model.
105 """
106 if self.padding_image_of.size != 0:
107 disp_contrib = atoms.positions - self.coords[: len(atoms)]
108 disp_pad = disp_contrib[self.padding_image_of]
109 self.coords += np.concatenate((disp_contrib, disp_pad))
110 else:
111 np.copyto(self.coords, atoms.positions)
113 if self.debug:
114 print("Debug: called update_kim_coords")
115 print()
117 def need_neigh_update(self, atoms, system_changes):
118 need_neigh_update = True
119 if len(system_changes) == 1 and "positions" in system_changes:
120 # Only position changes
121 if self.last_update_positions is not None:
122 a = self.last_update_positions
123 b = atoms.positions
124 if a.shape == b.shape:
125 delta = np.linalg.norm(a - b, axis=1)
126 # Indices of the two largest elements
127 try:
128 ind = np.argpartition(delta, -2)[-2:]
130 if sum(delta[ind]) <= self.skin:
131 need_neigh_update = False
132 except ValueError as error:
133 # if there is only a single atom that gets displaced
134 # np.argpartition(delta, -2) will fail with a
135 # ValueError, a single atom has no neighbors to update
136 if atoms.positions.shape[0] != 1:
137 raise error
138 need_neigh_update = False
140 return need_neigh_update
143class ASENeighborList(NeighborList):
144 def __init__(
145 self,
146 compute_args,
147 neigh_skin_ratio,
148 model_influence_dist,
149 model_cutoffs,
150 padding_not_require_neigh,
151 debug,
152 ):
153 super().__init__(
154 neigh_skin_ratio,
155 model_influence_dist,
156 model_cutoffs,
157 padding_not_require_neigh,
158 debug,
159 )
161 self.neigh = {}
162 compute_args.set_callback(
163 kimpy.compute_callback_name.GetNeighborList, self.get_neigh,
164 self.neigh
165 )
167 @staticmethod
168 def get_neigh(data, cutoffs, neighbor_list_index, particle_number):
169 """Retrieves the neighbors of each atom using ASE's native neighbor
170 list library
171 """
172 # We can only return neighbors of particles that were stored
173 number_of_particles = data["num_particles"]
174 if particle_number >= number_of_particles or particle_number < 0:
175 return (np.array([]), 1)
177 neighbors = data["neighbors"][neighbor_list_index][particle_number]
178 return (neighbors, 0)
180 def build(self, orig_atoms):
181 """Build the ASE neighbor list and return an Atoms object with
182 all of the neighbors added. First a neighbor list is created
183 from ase.neighbor_list, having only information about the
184 neighbors of the original atoms. If neighbors of padding atoms
185 are required, they are calculated using information from the
186 first neighbor list.
187 """
188 syms = orig_atoms.get_chemical_symbols()
189 orig_num_atoms = len(orig_atoms)
190 orig_pos = orig_atoms.get_positions()
192 # New atoms object that will contain the contributing atoms plus
193 # the padding atoms
194 new_atoms = orig_atoms.copy()
196 neigh_list = defaultdict(list)
197 neigh_dists = defaultdict(list)
199 # Information for padding atoms
200 padding_image_of = []
201 padding_shifts = []
203 # Ask ASE to build the neighbor list using the influence distance.
204 # This is not a neighbor list for each atom, but rather a listing
205 # of all neighbor pairs that exist. Atoms with no neighbors will
206 # not show up.
207 (
208 neigh_indices_i,
209 neigh_indices_j,
210 relative_pos,
211 neigh_cell_offsets,
212 dists,
213 ) = neighbor_list("ijDSd", orig_atoms, self.influence_dist)
215 # Loop over all neighbor pairs. Because this loop will generally
216 # include image atoms (for periodic systems), we keep track of
217 # which atoms/images we've accounted for in the `used` dictionary.
218 used = {}
219 for neigh_i, neigh_j, rel_pos, offset, dist in zip(
220 neigh_indices_i, neigh_indices_j,
221 relative_pos, neigh_cell_offsets, dists
222 ):
223 # Get neighbor position of neighbor
224 # (mapped back into unit cell, so this may overlap with other atoms)
225 wrapped_pos = orig_pos[neigh_i] + rel_pos
227 shift = tuple(offset)
228 uniq_index = (neigh_j,) + shift
229 if shift == (0, 0, 0):
230 # This atom is in the unit cell, i.e. it is contributing
231 neigh_list[neigh_i].append(neigh_j)
232 neigh_dists[neigh_i].append(dist)
233 if uniq_index not in used:
234 used[uniq_index] = neigh_j
235 else:
236 # This atom is not in the unit cell, i.e. it is padding
237 if uniq_index not in used:
238 # Add the neighbor as a padding atom
239 used[uniq_index] = len(new_atoms)
240 new_atoms.append(Atom(syms[neigh_j], position=wrapped_pos))
241 padding_image_of.append(neigh_j)
242 padding_shifts.append(offset)
243 neigh_list[neigh_i].append(used[uniq_index])
244 neigh_dists[neigh_i].append(dist)
245 neighbor_list_size = orig_num_atoms
247 # Add neighbors of padding atoms if the potential requires them
248 if self.padding_need_neigh:
249 neighbor_list_size = len(new_atoms)
250 inv_used = {v: k for k, v in used.items()}
251 # Loop over all the neighbors (k) and the image of that neighbor
252 # in the cell (neigh)
253 for k, neigh in enumerate(padding_image_of):
254 # Shift from original atom in cell to neighbor
255 shift = padding_shifts[k]
256 for orig_neigh, orig_dist in zip(
257 neigh_list[neigh], neigh_dists[neigh]):
258 # Get the shift of the neighbor of the original atom
259 orig_shift = inv_used[orig_neigh][1:]
261 # Apply sum of original shift and current shift to
262 # neighbors of original atom
263 total_shift = orig_shift + shift
265 # Get the image in the cell of the original neighbor
266 if orig_neigh <= orig_num_atoms - 1:
267 orig_neigh_image = orig_neigh
268 else:
269 orig_neigh_image = padding_image_of[orig_neigh -
270 orig_num_atoms]
272 # If the original image with the total shift has been
273 # used before then it is also a neighbor of this atom
274 uniq_index = (orig_neigh_image,) + tuple(total_shift)
275 if uniq_index in used:
276 neigh_list[k + orig_num_atoms].append(used[uniq_index])
277 neigh_dists[k + orig_num_atoms].append(orig_dist)
279 # If the model has multiple cutoffs, we need to return neighbor lists
280 # corresponding to each of them
281 neigh_lists = []
282 for cut in self.cutoffs:
283 neigh_list = [
284 np.array(neigh_list[k], dtype=c_int)[neigh_dists[k] <= cut]
285 for k in range(neighbor_list_size)
286 ]
287 neigh_lists.append(neigh_list)
289 self.padding_image_of = padding_image_of
291 self.neigh["neighbors"] = neigh_lists
292 self.neigh["num_particles"] = neighbor_list_size
294 return new_atoms
296 def update(self, orig_atoms, species_map):
297 """Create the neighbor list along with the other required
298 parameters (which are stored as instance attributes). The
299 required parameters are:
301 - num_particles
302 - coords
303 - particle_contributing
304 - species_code
306 Note that the KIM API requires a neighbor list that has indices
307 corresponding to each atom.
308 """
310 # Information of original atoms
311 self.num_contributing_particles = len(orig_atoms)
313 new_atoms = self.build(orig_atoms)
315 # Save the number of atoms and all their neighbors and positions
316 num_atoms = len(new_atoms)
317 num_padding = num_atoms - self.num_contributing_particles
318 self.num_particles = [num_atoms]
319 self.coords = new_atoms.get_positions()
321 # Save which coordinates are from original atoms and which are from
322 # neighbors using a mask
323 indices_mask = [1] * self.num_contributing_particles + [0] * num_padding
324 self.particle_contributing = indices_mask
326 # Species support and code
327 try:
328 self.species_code = [
329 species_map[s] for s in new_atoms.get_chemical_symbols()
330 ]
331 except KeyError as e:
332 raise RuntimeError(
333 "Species not supported by KIM model; {}".format(
334 str(e)))
336 self.last_update_positions = orig_atoms.get_positions()
338 if self.debug:
339 print("Debug: called update_ase_neigh")
340 print()
343class KimpyNeighborList(NeighborList):
344 def __init__(
345 self,
346 compute_args,
347 neigh_skin_ratio,
348 model_influence_dist,
349 model_cutoffs,
350 padding_not_require_neigh,
351 debug,
352 ):
353 super().__init__(
354 neigh_skin_ratio,
355 model_influence_dist,
356 model_cutoffs,
357 padding_not_require_neigh,
358 debug,
359 )
361 self.neigh = neighlist.NeighList()
363 compute_args.set_callback_pointer(
364 kimpy.compute_callback_name.GetNeighborList,
365 neighlist.get_neigh_kim(),
366 self.neigh,
367 )
369 @check_call_wrapper
370 def build(self):
371 return self.neigh.build(
372 self.coords, self.influence_dist, self.cutoffs, self.need_neigh
373 )
375 @check_call_wrapper
376 def create_paddings(
377 self, cell, pbc, contributing_coords, contributing_species_code
378 ):
379 # Cast things passed through kimpy to numpy arrays
380 cell = np.asarray(cell, dtype=c_double)
381 pbc = np.asarray(pbc, dtype=c_int)
382 contributing_coords = np.asarray(contributing_coords, dtype=c_double)
384 return neighlist.create_paddings(
385 self.influence_dist,
386 cell,
387 pbc,
388 contributing_coords,
389 contributing_species_code,
390 )
392 def update(self, atoms, species_map):
393 """Create the neighbor list along with the other required
394 parameters (which are stored as instance attributes). The
395 required parameters are:
397 - num_particles
398 - coords
399 - particle_contributing
400 - species_code
402 Note that the KIM API requires a neighbor list that has indices
403 corresponding to each atom.
404 """
406 # Get info from Atoms object
407 cell = np.asarray(atoms.get_cell(), dtype=c_double)
408 pbc = np.asarray(atoms.get_pbc(), dtype=c_int)
409 contributing_coords = np.asarray(atoms.get_positions(), dtype=c_double)
410 self.num_contributing_particles = atoms.get_global_number_of_atoms()
411 num_contributing = self.num_contributing_particles
413 # Species support and code
414 try:
415 contributing_species_code = np.array(
416 [species_map[s] for s in atoms.get_chemical_symbols()],
417 dtype=c_int
418 )
419 except KeyError as e:
420 raise RuntimeError(
421 "Species not supported by KIM model; {}".format(
422 str(e)))
424 if pbc.any(): # Need padding atoms
425 # Create padding atoms
427 (
428 padding_coords,
429 padding_species_code,
430 self.padding_image_of,
431 ) = self.create_paddings(
432 cell, pbc, contributing_coords, contributing_species_code
433 )
434 num_padding = padding_species_code.size
436 self.num_particles = [num_contributing + num_padding]
437 self.coords = np.concatenate((contributing_coords, padding_coords))
438 self.species_code = np.concatenate(
439 (contributing_species_code, padding_species_code)
440 )
441 self.particle_contributing = [1] * \
442 num_contributing + [0] * num_padding
443 self.need_neigh = [1] * self.num_particles[0]
444 if not self.padding_need_neigh:
445 self.need_neigh[num_contributing:] = 0
447 else: # Do not need padding atoms
448 self.padding_image_of = []
449 self.num_particles = [num_contributing]
450 self.coords = contributing_coords
451 self.species_code = contributing_species_code
452 self.particle_contributing = [1] * num_contributing
453 self.need_neigh = self.particle_contributing
455 # Create neighborlist
456 self.build()
458 self.last_update_positions = atoms.get_positions()
460 if self.debug:
461 print("Debug: called update_kimpy_neigh")
462 print()