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