Coverage for /builds/debichem-team/python-ase/ase/calculators/kim/neighborlist.py : 51.63%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from collections import defaultdict
3import numpy as np
4import kimpy
5from kimpy import neighlist
6from ase.neighborlist import neighbor_list
7from ase import Atom
9from .kimpy_wrappers import check_call_wrapper
12class NeighborList:
14 kimpy_arrays = {
15 "num_particles": np.intc,
16 "coords": np.double,
17 "particle_contributing": np.intc,
18 "species_code": np.intc,
19 "cutoffs": np.double,
20 "padding_image_of": np.intc,
21 "need_neigh": np.intc,
22 }
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
40 def __init__(
41 self,
42 neigh_skin_ratio,
43 model_influence_dist,
44 model_cutoffs,
45 padding_not_require_neigh,
46 debug,
47 ):
49 self.skin = neigh_skin_ratio * model_influence_dist
50 self.influence_dist = model_influence_dist + self.skin
51 self.cutoffs = model_cutoffs + self.skin
52 self.padding_need_neigh = not padding_not_require_neigh.all()
53 self.debug = debug
55 if self.debug:
56 print()
57 print("Calculator skin: {}".format(self.skin))
58 print(f"Model influence distance: {model_influence_dist}")
59 print(
60 "Calculator influence distance (including skin distance): {}"
61 "".format(self.influence_dist)
62 )
63 print("Number of cutoffs: {}".format(model_cutoffs.size))
64 print("Model cutoffs: {}".format(model_cutoffs))
65 print(
66 "Calculator cutoffs (including skin distance): {}"
67 "".format(self.cutoffs)
68 )
69 print(
70 "Model needs neighbors of padding atoms: {}"
71 "".format(self.padding_need_neigh)
72 )
73 print()
75 # Attributes to be set by subclasses
76 self.neigh = None
77 self.num_contributing_particles = None
78 self.padding_image_of = None
79 self.num_particles = None
80 self.coords = None
81 self.particle_contributing = None
82 self.species_code = None
83 self.need_neigh = None
84 self.last_update_positions = None
86 def update_kim_coords(self, atoms):
87 """Update atomic positions in self.coords, which is where the KIM
88 API will look to find them in order to pass them to the model.
89 """
90 if self.padding_image_of.size != 0:
91 disp_contrib = atoms.positions - self.coords[: len(atoms)]
92 disp_pad = disp_contrib[self.padding_image_of]
93 self.coords += np.concatenate((disp_contrib, disp_pad))
94 else:
95 np.copyto(self.coords, atoms.positions)
97 if self.debug:
98 print("Debug: called update_kim_coords")
99 print()
101 def need_neigh_update(self, atoms, system_changes):
102 need_neigh_update = True
103 if len(system_changes) == 1 and "positions" in system_changes:
104 # Only position changes
105 if self.last_update_positions is not None:
106 a = self.last_update_positions
107 b = atoms.positions
108 if a.shape == b.shape:
109 delta = np.linalg.norm(a - b, axis=1)
110 # Indices of the two largest elements
111 ind = np.argpartition(delta, -2)[-2:]
112 if sum(delta[ind]) <= self.skin:
113 need_neigh_update = False
115 return need_neigh_update
118class ASENeighborList(NeighborList):
119 def __init__(
120 self,
121 compute_args,
122 neigh_skin_ratio,
123 model_influence_dist,
124 model_cutoffs,
125 padding_not_require_neigh,
126 debug,
127 ):
128 super().__init__(
129 neigh_skin_ratio,
130 model_influence_dist,
131 model_cutoffs,
132 padding_not_require_neigh,
133 debug,
134 )
136 self.neigh = {}
137 compute_args.set_callback(
138 kimpy.compute_callback_name.GetNeighborList,
139 self.get_neigh,
140 self.neigh)
142 @staticmethod
143 def get_neigh(data, cutoffs, neighbor_list_index, particle_number):
144 """Retrieves the neighbors of each atom using ASE's native neighbor
145 list library
146 """
147 # We can only return neighbors of particles that were stored
148 number_of_particles = data["num_particles"]
149 if particle_number >= number_of_particles or particle_number < 0:
150 return (np.array([]), 1)
152 neighbors = data["neighbors"][neighbor_list_index][particle_number]
153 return (neighbors, 0)
155 def build(self, orig_atoms):
156 """Build the ASE neighbor list and return an Atoms object with
157 all of the neighbors added. First a neighbor list is created
158 from ase.neighbor_list, having only information about the
159 neighbors of the original atoms. If neighbors of padding atoms
160 are required, they are calculated using information from the
161 first neighbor list.
162 """
163 syms = orig_atoms.get_chemical_symbols()
164 orig_num_atoms = len(orig_atoms)
165 orig_pos = orig_atoms.get_positions()
167 # New atoms object that will contain the contributing atoms plus
168 # the padding atoms
169 new_atoms = orig_atoms.copy()
171 neigh_list = defaultdict(list)
172 neigh_dists = defaultdict(list)
174 # Information for padding atoms
175 padding_image_of = []
176 padding_shifts = []
178 # Ask ASE to build the neighbor list using the influence distance.
179 # This is not a neighbor list for each atom, but rather a listing
180 # of all neighbor pairs that exist. Atoms with no neighbors will
181 # not show up.
182 (neigh_indices_i,
183 neigh_indices_j,
184 relative_pos,
185 neigh_cell_offsets,
186 dists) = neighbor_list(
187 "ijDSd", orig_atoms, self.influence_dist
188 )
190 # Loop over all neighbor pairs. Because this loop will generally
191 # include image atoms (for periodic systems), we keep track of
192 # which atoms/images we've accounted for in the `used` dictionary.
193 used = dict()
194 for neigh_i, neigh_j, rel_pos, offset, dist in zip(neigh_indices_i,
195 neigh_indices_j,
196 relative_pos,
197 neigh_cell_offsets,
198 dists):
199 # Get neighbor position of neighbor
200 # (mapped back into unit cell, so this may overlap with other atoms)
201 wrapped_pos = orig_pos[neigh_i] + rel_pos
203 shift = tuple(offset)
204 uniq_index = (neigh_j,) + shift
205 if shift == (0, 0, 0):
206 # This atom is in the unit cell, i.e. it is contributing
207 neigh_list[neigh_i].append(neigh_j)
208 neigh_dists[neigh_i].append(dist)
209 if uniq_index not in used:
210 used[uniq_index] = neigh_j
211 else:
212 # This atom is not in the unit cell, i.e. it is padding
213 if uniq_index not in used:
214 # Add the neighbor as a padding atom
215 used[uniq_index] = len(new_atoms)
216 new_atoms.append(Atom(syms[neigh_j], position=wrapped_pos))
217 padding_image_of.append(neigh_j)
218 padding_shifts.append(offset)
219 neigh_list[neigh_i].append(used[uniq_index])
220 neigh_dists[neigh_i].append(dist)
221 neighbor_list_size = orig_num_atoms
223 # Add neighbors of padding atoms if the potential requires them
224 if self.padding_need_neigh:
225 neighbor_list_size = len(new_atoms)
226 inv_used = dict((v, k) for k, v in used.items())
227 # Loop over all the neighbors (k) and the image of that neighbor
228 # in the cell (neigh)
229 for k, neigh in enumerate(padding_image_of):
230 # Shift from original atom in cell to neighbor
231 shift = padding_shifts[k]
232 for orig_neigh, orig_dist in zip(neigh_list[neigh],
233 neigh_dists[neigh]):
234 # Get the shift of the neighbor of the original atom
235 orig_shift = inv_used[orig_neigh][1:]
237 # Apply sum of original shift and current shift to
238 # neighbors of original atom
239 total_shift = orig_shift + shift
241 # Get the image in the cell of the original neighbor
242 if orig_neigh <= orig_num_atoms - 1:
243 orig_neigh_image = orig_neigh
244 else:
245 orig_neigh_image = padding_image_of[orig_neigh -
246 orig_num_atoms]
248 # If the original image with the total shift has been
249 # used before then it is also a neighbor of this atom
250 uniq_index = (orig_neigh_image,) + tuple(total_shift)
251 if uniq_index in used:
252 neigh_list[k + orig_num_atoms].append(used[uniq_index])
253 neigh_dists[k + orig_num_atoms].append(orig_dist)
255 # If the model has multiple cutoffs, we need to return neighbor lists
256 # corresponding to each of them
257 neigh_lists = []
258 for cut in self.cutoffs:
259 neigh_list = [
260 np.array(neigh_list[k], dtype=np.intc)[neigh_dists[k] <= cut]
261 for k in range(neighbor_list_size)
262 ]
263 neigh_lists.append(neigh_list)
265 self.padding_image_of = padding_image_of
267 self.neigh["neighbors"] = neigh_lists
268 self.neigh["num_particles"] = neighbor_list_size
270 return new_atoms
272 def update(self, orig_atoms, species_map):
273 """Create the neighbor list along with the other required
274 parameters (which are stored as instance attributes). The
275 required parameters are:
277 - num_particles
278 - coords
279 - particle_contributing
280 - species_code
282 Note that the KIM API requires a neighbor list that has indices
283 corresponding to each atom.
284 """
286 # Information of original atoms
287 self.num_contributing_particles = len(orig_atoms)
289 new_atoms = self.build(orig_atoms)
291 # Save the number of atoms and all their neighbors and positions
292 num_atoms = len(new_atoms)
293 num_padding = num_atoms - self.num_contributing_particles
294 self.num_particles = [num_atoms]
295 self.coords = new_atoms.get_positions()
297 # Save which coordinates are from original atoms and which are from
298 # neighbors using a mask
299 indices_mask = [1] * self.num_contributing_particles + [0] * num_padding
300 self.particle_contributing = indices_mask
302 # Species support and code
303 try:
304 self.species_code = [
305 species_map[s] for s in new_atoms.get_chemical_symbols()
306 ]
307 except KeyError as e:
308 raise RuntimeError("Species not supported by KIM model; {}"
309 .format(str(e)))
311 self.last_update_positions = orig_atoms.get_positions()
313 if self.debug:
314 print("Debug: called update_ase_neigh")
315 print()
318class KimpyNeighborList(NeighborList):
319 def __init__(
320 self,
321 compute_args,
322 neigh_skin_ratio,
323 model_influence_dist,
324 model_cutoffs,
325 padding_not_require_neigh,
326 debug,
327 ):
328 super().__init__(
329 neigh_skin_ratio,
330 model_influence_dist,
331 model_cutoffs,
332 padding_not_require_neigh,
333 debug,
334 )
336 self.neigh = neighlist.NeighList()
338 compute_args.set_callback_pointer(
339 kimpy.compute_callback_name.GetNeighborList,
340 neighlist.get_neigh_kim(),
341 self.neigh,
342 )
344 @check_call_wrapper
345 def build(self):
346 return self.neigh.build(
347 self.coords, self.influence_dist, self.cutoffs, self.need_neigh
348 )
350 @check_call_wrapper
351 def create_paddings(
352 self, cell, pbc, contributing_coords, contributing_species_code
353 ):
354 # Cast things passed through kimpy to numpy arrays
355 cell = np.asarray(cell, dtype=np.double)
356 pbc = np.asarray(pbc, dtype=np.intc)
357 contributing_coords = np.asarray(contributing_coords, dtype=np.double)
359 return neighlist.create_paddings(
360 self.influence_dist,
361 cell,
362 pbc,
363 contributing_coords,
364 contributing_species_code,
365 )
367 def update(self, atoms, species_map):
368 """Create the neighbor list along with the other required
369 parameters (which are stored as instance attributes). The
370 required parameters are:
372 - num_particles
373 - coords
374 - particle_contributing
375 - species_code
377 Note that the KIM API requires a neighbor list that has indices
378 corresponding to each atom.
379 """
381 # Get info from Atoms object
382 cell = np.asarray(atoms.get_cell(), dtype=np.double)
383 pbc = np.asarray(atoms.get_pbc(), dtype=np.intc)
384 contributing_coords = np.asarray(atoms.get_positions(), dtype=np.double)
385 self.num_contributing_particles = atoms.get_global_number_of_atoms()
386 num_contributing = self.num_contributing_particles
388 # Species support and code
389 try:
390 contributing_species_code = np.array(
391 [species_map[s] for s in atoms.get_chemical_symbols()],
392 dtype=np.intc)
393 except KeyError as e:
394 raise RuntimeError("Species not supported by KIM model; {}"
395 .format(str(e)))
397 if pbc.any(): # Need padding atoms
398 # Create padding atoms
400 padding_coords, padding_species_code, self.padding_image_of = \
401 self.create_paddings(
402 cell, pbc, contributing_coords, contributing_species_code)
403 num_padding = padding_species_code.size
405 self.num_particles = [num_contributing + num_padding]
406 self.coords = np.concatenate((contributing_coords, padding_coords))
407 self.species_code = np.concatenate(
408 (contributing_species_code, padding_species_code)
409 )
410 self.particle_contributing = ([1] * num_contributing +
411 [0] * num_padding)
412 self.need_neigh = [1] * self.num_particles[0]
413 if not self.padding_need_neigh:
414 self.need_neigh[num_contributing:] = 0
416 else: # Do not need padding atoms
417 self.padding_image_of = []
418 self.num_particles = [num_contributing]
419 self.coords = contributing_coords
420 self.species_code = contributing_species_code
421 self.particle_contributing = [1] * num_contributing
422 self.need_neigh = self.particle_contributing
424 # Create neighborlist
425 self.build()
427 self.last_update_positions = atoms.get_positions()
429 if self.debug:
430 print("Debug: called update_kimpy_neigh")
431 print()