Coverage for /builds/ase/ase/ase/calculators/kim/kimmodel.py : 95.60%
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
1"""
2ASE Calculator for interatomic models compatible with the Knowledgebase
3of Interatomic Models (KIM) application programming interface (API).
4Written by:
6Mingjian Wen
7Daniel S. Karls
8University of Minnesota
9"""
10import numpy as np
12from ase.calculators.calculator import Calculator
13from ase.calculators.calculator import compare_atoms
15from . import kimpy_wrappers
16from . import neighborlist
19class KIMModelData:
20 """Initializes and subsequently stores the KIM API Portable Model
21 object, KIM API ComputeArguments object, and the neighbor list
22 object used by instances of KIMModelCalculator. Also stores the
23 arrays which are registered in the KIM API and which are used to
24 communicate with the model.
25 """
27 def __init__(self, model_name, ase_neigh, neigh_skin_ratio, debug=False):
28 self.model_name = model_name
29 self.ase_neigh = ase_neigh
30 self.debug = debug
32 # Initialize KIM API Portable Model object and ComputeArguments object
33 self.kim_model, self.compute_args = self._init_kim()
35 self.species_map = self._create_species_map()
37 # Ask model to provide information relevant for neighbor list
38 # construction
39 (
40 model_influence_dist,
41 model_cutoffs,
42 padding_not_require_neigh,
43 ) = self.get_model_neighbor_list_parameters()
45 # Initialize neighbor list object
46 self.neigh = self._init_neigh(
47 neigh_skin_ratio,
48 model_influence_dist,
49 model_cutoffs,
50 padding_not_require_neigh,
51 )
53 def _init_kim(self):
54 """Create the KIM API Portable Model object and KIM API ComputeArguments
55 object
56 """
57 if self.kim_initialized:
58 return
60 kim_model = kimpy_wrappers.PortableModel(self.model_name, self.debug)
62 # KIM API model object is what actually creates/destroys the
63 # ComputeArguments object, so we must pass it as a parameter
64 compute_args = kim_model.compute_arguments_create()
66 return kim_model, compute_args
68 def _init_neigh(
69 self,
70 neigh_skin_ratio,
71 model_influence_dist,
72 model_cutoffs,
73 padding_not_require_neigh,
74 ):
75 """Initialize neighbor list, either an ASE-native neighborlist
76 or one created using the neighlist module in kimpy
77 """
78 neigh_list_object_type = (
79 neighborlist.ASENeighborList
80 if self.ase_neigh
81 else neighborlist.KimpyNeighborList
82 )
83 return neigh_list_object_type(
84 self.compute_args,
85 neigh_skin_ratio,
86 model_influence_dist,
87 model_cutoffs,
88 padding_not_require_neigh,
89 self.debug,
90 )
92 def get_model_neighbor_list_parameters(self):
93 model_influence_dist = self.kim_model.get_influence_distance()
94 (
95 model_cutoffs,
96 padding_not_require_neigh,
97 ) = self.kim_model.get_neighbor_list_cutoffs_and_hints()
99 return model_influence_dist, model_cutoffs, padding_not_require_neigh
101 def update_compute_args_pointers(self, energy, forces):
102 self.compute_args.update(
103 self.num_particles,
104 self.species_code,
105 self._particle_contributing,
106 self.coords,
107 energy,
108 forces,
109 )
111 def _create_species_map(self):
112 """Get all the supported species of the KIM model and the
113 corresponding integer codes used by the model
115 Returns
116 -------
117 species_map : dict
118 key : str
119 chemical symbols (e.g. "Ar")
120 value : int
121 species integer code (e.g. 1)
122 """
123 supported_species, codes = self._get_model_supported_species_and_codes()
124 species_map = dict()
125 for i, spec in enumerate(supported_species):
126 species_map[spec] = codes[i]
127 if self.debug:
128 print(
129 "Species {} is supported and its code is: {}".format(
130 spec, codes[i])
131 )
133 return species_map
135 @property
136 def padding_image_of(self):
137 return self.neigh.padding_image_of
139 @property
140 def num_particles(self):
141 return self.neigh.num_particles
143 @property
144 def coords(self):
145 return self.neigh.coords
147 @property
148 def _particle_contributing(self):
149 return self.neigh.particle_contributing
151 @property
152 def species_code(self):
153 return self.neigh.species_code
155 @property
156 def kim_initialized(self):
157 return hasattr(self, "kim_model")
159 @property
160 def _neigh_initialized(self):
161 return hasattr(self, "neigh")
163 @property
164 def _get_model_supported_species_and_codes(self):
165 return self.kim_model.get_model_supported_species_and_codes
168class KIMModelCalculator(Calculator):
169 """Calculator that works with KIM Portable Models (PMs).
171 Calculator that carries out direct communication between ASE and a
172 KIM Portable Model (PM) through the kimpy library (which provides a
173 set of python bindings to the KIM API).
175 Parameters
176 ----------
177 model_name : str
178 The unique identifier assigned to the interatomic model (for
179 details, see https://openkim.org/doc/schema/kim-ids)
181 ase_neigh : bool, optional
182 False (default): Use kimpy's neighbor list library
184 True: Use ASE's internal neighbor list mechanism (usually slower
185 than the kimpy neighlist library)
187 neigh_skin_ratio : float, optional
188 Used to determine the neighbor list cutoff distance, r_neigh,
189 through the relation r_neigh = (1 + neigh_skin_ratio) * rcut,
190 where rcut is the model's influence distance. (Default: 0.2)
192 release_GIL : bool, optional
193 Whether to release python GIL. Releasing the GIL allows a KIM
194 model to run with multiple concurrent threads. (Default: False)
196 debug : bool, optional
197 If True, detailed information is printed to stdout. (Default:
198 False)
199 """
201 implemented_properties = ["energy", "free_energy", "forces", "stress"]
203 ignored_changes = {"initial_charges", "initial_magmoms"}
205 def __init__(
206 self,
207 model_name,
208 ase_neigh=False,
209 neigh_skin_ratio=0.2,
210 release_GIL=False,
211 debug=False,
212 *args,
213 **kwargs
214 ):
215 super().__init__(*args, **kwargs)
217 self.model_name = model_name
218 self.release_GIL = release_GIL
219 self.debug = debug
221 if neigh_skin_ratio < 0:
222 raise ValueError('Argument "neigh_skin_ratio" must be non-negative')
223 self.neigh_skin_ratio = neigh_skin_ratio
225 # Model output
226 self.energy = None
227 self.forces = None
229 # Create KIMModelData object. This will take care of creating
230 # and storing the KIM API Portable Model object, KIM API
231 # ComputeArguments object, and the neighbor list object that
232 # our calculator needs
233 self._kimmodeldata = KIMModelData(
234 self.model_name, ase_neigh, self.neigh_skin_ratio, self.debug
235 )
237 self._parameters_changed = False
239 def __enter__(self):
240 return self
242 def __exit__(self, exc_type, value, traceback):
243 pass
245 def __repr__(self):
246 return "KIMModelCalculator(model_name={})".format(self.model_name)
248 def calculate(
249 self,
250 atoms=None,
251 properties=["energy", "forces", "stress"],
252 system_changes=["positions", "numbers", "cell", "pbc"],
253 ):
254 """
255 Inherited method from the ase Calculator class that is called by
256 get_property()
258 Parameters
259 ----------
260 atoms : Atoms
261 Atoms object whose properties are desired
263 properties : list of str
264 List of what needs to be calculated. Can be any combination
265 of 'energy', 'forces' and 'stress'.
267 system_changes : list of str
268 List of what has changed since last calculation. Can be any
269 combination of these six: 'positions', 'numbers', 'cell',
270 and 'pbc'.
271 """
273 super().calculate(atoms, properties, system_changes)
275 if self._parameters_changed:
276 self._parameters_changed = False
278 if system_changes:
280 # Ask model to update all of its parameters and the parameters
281 # related to the neighbor list(s). This update is necessary to do
282 # here since the user will generally have made changes the model
283 # parameters since the last time an update was performed and we
284 # need to ensure that any properties calculated here are made using
285 # the up-to-date model and neighbor list parameters.
286 self._model_refresh_and_update_neighbor_list_parameters()
288 if self._need_neigh_update(atoms, system_changes):
289 self._update_neigh(atoms, self._species_map)
290 self.energy = np.array([0.0], dtype=kimpy_wrappers.c_double)
291 self.forces = np.zeros(
292 [self._num_particles[0], 3], dtype=kimpy_wrappers.c_double
293 )
294 self._update_compute_args_pointers(self.energy, self.forces)
295 else:
296 self._update_kim_coords(atoms)
298 self._kim_model.compute(self._compute_args, self.release_GIL)
300 energy = self.energy[0]
301 forces = self._assemble_padding_forces()
303 try:
304 volume = atoms.get_volume()
305 stress = self._compute_virial_stress(
306 self.forces, self._coords, volume)
307 except ValueError: # Volume cannot be computed
308 stress = None
310 # Quantities passed back to ASE
311 self.results["energy"] = energy
312 self.results["free_energy"] = energy
313 self.results["forces"] = forces
314 self.results["stress"] = stress
316 def check_state(self, atoms, tol=1e-15):
317 # Check for change in atomic configuration (positions or pbc)
318 system_changes = compare_atoms(
319 self.atoms, atoms, excluded_properties=self.ignored_changes
320 )
322 # Check if model parameters were changed
323 if self._parameters_changed:
324 system_changes.append("calculator")
326 return system_changes
328 def _assemble_padding_forces(self):
329 """
330 Assemble forces on padding atoms back to contributing atoms.
332 Parameters
333 ----------
334 forces : 2D array of doubles
335 Forces on both contributing and padding atoms
337 num_contrib: int
338 Number of contributing atoms
340 padding_image_of : 1D array of int
341 Atom number, of which the padding atom is an image
344 Returns
345 -------
346 Total forces on contributing atoms.
347 """
349 total_forces = np.array(self.forces[:self._num_contributing_particles])
351 if self._padding_image_of.size != 0:
352 pad_forces = self.forces[self._num_contributing_particles:]
353 for f, org_index in zip(pad_forces, self._padding_image_of):
354 total_forces[org_index] += f
356 return total_forces
358 @staticmethod
359 def _compute_virial_stress(forces, coords, volume):
360 """Compute the virial stress in Voigt notation.
362 Parameters
363 ----------
364 forces : 2D array
365 Partial forces on all atoms (padding included)
367 coords : 2D array
368 Coordinates of all atoms (padding included)
370 volume : float
371 Volume of cell
373 Returns
374 -------
375 stress : 1D array
376 stress in Voigt order (xx, yy, zz, yz, xz, xy)
377 """
378 stress = np.zeros(6)
379 stress[0] = -np.dot(forces[:, 0], coords[:, 0]) / volume
380 stress[1] = -np.dot(forces[:, 1], coords[:, 1]) / volume
381 stress[2] = -np.dot(forces[:, 2], coords[:, 2]) / volume
382 stress[3] = -np.dot(forces[:, 1], coords[:, 2]) / volume
383 stress[4] = -np.dot(forces[:, 0], coords[:, 2]) / volume
384 stress[5] = -np.dot(forces[:, 0], coords[:, 1]) / volume
386 return stress
388 @property
389 def _update_compute_args_pointers(self):
390 return self._kimmodeldata.update_compute_args_pointers
392 @property
393 def _kim_model(self):
394 return self._kimmodeldata.kim_model
396 @property
397 def _compute_args(self):
398 return self._kimmodeldata.compute_args
400 @property
401 def _num_particles(self):
402 return self._kimmodeldata.num_particles
404 @property
405 def _coords(self):
406 return self._kimmodeldata.coords
408 @property
409 def _padding_image_of(self):
410 return self._kimmodeldata.padding_image_of
412 @property
413 def _species_map(self):
414 return self._kimmodeldata.species_map
416 @property
417 def _neigh(self):
418 # WARNING: This property is underscored for a reason! The
419 # neighborlist(s) itself (themselves) may not be up to date with
420 # respect to changes that have been made to the model's parameters, or
421 # even since the positions in the Atoms object may have changed.
422 # Neighbor lists are only potentially updated inside the ``calculate``
423 # method.
424 return self._kimmodeldata.neigh
426 @property
427 def _num_contributing_particles(self):
428 return self._neigh.num_contributing_particles
430 @property
431 def _update_kim_coords(self):
432 return self._neigh.update_kim_coords
434 @property
435 def _need_neigh_update(self):
436 return self._neigh.need_neigh_update
438 @property
439 def _update_neigh(self):
440 return self._neigh.update
442 @property
443 def parameters_metadata(self):
444 return self._kim_model.parameters_metadata
446 @property
447 def parameter_names(self):
448 return self._kim_model.parameter_names
450 @property
451 def get_parameters(self):
452 # Ask model to update all of its parameters and the parameters related
453 # to the neighbor list(s). This update is necessary to do here since
454 # the user will generally have made changes the model parameters since
455 # the last time an update was performed and we need to ensure the
456 # parameters returned by this method are fully up to date.
457 self._model_refresh_and_update_neighbor_list_parameters()
459 return self._kim_model.get_parameters
461 def set_parameters(self, **kwargs):
462 parameters = self._kim_model.set_parameters(**kwargs)
463 self._parameters_changed = True
465 return parameters
467 def _model_refresh_and_update_neighbor_list_parameters(self):
468 """
469 Call the model's refresh routine and update the neighbor list object
470 for any necessary changes arising from changes to the model parameters,
471 e.g. a change in one of its cutoffs. After a model's parameters have
472 been changed, this method *must* be called before calling the model's
473 compute routine.
474 """
475 self._kim_model.clear_then_refresh()
477 # Update neighbor list parameters
478 (
479 model_influence_dist,
480 model_cutoffs,
481 padding_not_require_neigh,
482 ) = self._kimmodeldata.get_model_neighbor_list_parameters()
484 self._neigh.set_neigh_parameters(
485 self.neigh_skin_ratio,
486 model_influence_dist,
487 model_cutoffs,
488 padding_not_require_neigh,
489 )