Coverage for /builds/kinetik161/ase/ase/calculators/kim/kimmodel.py: 95.56%
180 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
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, compare_atoms
14from . import kimpy_wrappers, neighborlist
17class KIMModelData:
18 """Initializes and subsequently stores the KIM API Portable Model
19 object, KIM API ComputeArguments object, and the neighbor list
20 object used by instances of KIMModelCalculator. Also stores the
21 arrays which are registered in the KIM API and which are used to
22 communicate with the model.
23 """
25 def __init__(self, model_name, ase_neigh, neigh_skin_ratio, debug=False):
26 self.model_name = model_name
27 self.ase_neigh = ase_neigh
28 self.debug = debug
30 # Initialize KIM API Portable Model object and ComputeArguments object
31 self.kim_model, self.compute_args = self._init_kim()
33 self.species_map = self._create_species_map()
35 # Ask model to provide information relevant for neighbor list
36 # construction
37 (
38 model_influence_dist,
39 model_cutoffs,
40 padding_not_require_neigh,
41 ) = self.get_model_neighbor_list_parameters()
43 # Initialize neighbor list object
44 self.neigh = self._init_neigh(
45 neigh_skin_ratio,
46 model_influence_dist,
47 model_cutoffs,
48 padding_not_require_neigh,
49 )
51 def _init_kim(self):
52 """Create the KIM API Portable Model object and KIM API ComputeArguments
53 object
54 """
55 if self.kim_initialized:
56 return
58 kim_model = kimpy_wrappers.PortableModel(self.model_name, self.debug)
60 # KIM API model object is what actually creates/destroys the
61 # ComputeArguments object, so we must pass it as a parameter
62 compute_args = kim_model.compute_arguments_create()
64 return kim_model, compute_args
66 def _init_neigh(
67 self,
68 neigh_skin_ratio,
69 model_influence_dist,
70 model_cutoffs,
71 padding_not_require_neigh,
72 ):
73 """Initialize neighbor list, either an ASE-native neighborlist
74 or one created using the neighlist module in kimpy
75 """
76 neigh_list_object_type = (
77 neighborlist.ASENeighborList
78 if self.ase_neigh
79 else neighborlist.KimpyNeighborList
80 )
81 return neigh_list_object_type(
82 self.compute_args,
83 neigh_skin_ratio,
84 model_influence_dist,
85 model_cutoffs,
86 padding_not_require_neigh,
87 self.debug,
88 )
90 def get_model_neighbor_list_parameters(self):
91 model_influence_dist = self.kim_model.get_influence_distance()
92 (
93 model_cutoffs,
94 padding_not_require_neigh,
95 ) = self.kim_model.get_neighbor_list_cutoffs_and_hints()
97 return model_influence_dist, model_cutoffs, padding_not_require_neigh
99 def update_compute_args_pointers(self, energy, forces):
100 self.compute_args.update(
101 self.num_particles,
102 self.species_code,
103 self._particle_contributing,
104 self.coords,
105 energy,
106 forces,
107 )
109 def _create_species_map(self):
110 """Get all the supported species of the KIM model and the
111 corresponding integer codes used by the model
113 Returns
114 -------
115 species_map : dict
116 key : str
117 chemical symbols (e.g. "Ar")
118 value : int
119 species integer code (e.g. 1)
120 """
121 supported_species, codes = self._get_model_supported_species_and_codes()
122 species_map = {}
123 for i, spec in enumerate(supported_species):
124 species_map[spec] = codes[i]
125 if self.debug:
126 print(
127 "Species {} is supported and its code is: {}".format(
128 spec, codes[i])
129 )
131 return species_map
133 @property
134 def padding_image_of(self):
135 return self.neigh.padding_image_of
137 @property
138 def num_particles(self):
139 return self.neigh.num_particles
141 @property
142 def coords(self):
143 return self.neigh.coords
145 @property
146 def _particle_contributing(self):
147 return self.neigh.particle_contributing
149 @property
150 def species_code(self):
151 return self.neigh.species_code
153 @property
154 def kim_initialized(self):
155 return hasattr(self, "kim_model")
157 @property
158 def _neigh_initialized(self):
159 return hasattr(self, "neigh")
161 @property
162 def _get_model_supported_species_and_codes(self):
163 return self.kim_model.get_model_supported_species_and_codes
166class KIMModelCalculator(Calculator):
167 """Calculator that works with KIM Portable Models (PMs).
169 Calculator that carries out direct communication between ASE and a
170 KIM Portable Model (PM) through the kimpy library (which provides a
171 set of python bindings to the KIM API).
173 Parameters
174 ----------
175 model_name : str
176 The unique identifier assigned to the interatomic model (for
177 details, see https://openkim.org/doc/schema/kim-ids)
179 ase_neigh : bool, optional
180 False (default): Use kimpy's neighbor list library
182 True: Use ASE's internal neighbor list mechanism (usually slower
183 than the kimpy neighlist library)
185 neigh_skin_ratio : float, optional
186 Used to determine the neighbor list cutoff distance, r_neigh,
187 through the relation r_neigh = (1 + neigh_skin_ratio) * rcut,
188 where rcut is the model's influence distance. (Default: 0.2)
190 release_GIL : bool, optional
191 Whether to release python GIL. Releasing the GIL allows a KIM
192 model to run with multiple concurrent threads. (Default: False)
194 debug : bool, optional
195 If True, detailed information is printed to stdout. (Default:
196 False)
197 """
199 implemented_properties = ["energy", "free_energy", "forces", "stress"]
201 ignored_changes = {"initial_charges", "initial_magmoms"}
203 def __init__(
204 self,
205 model_name,
206 ase_neigh=False,
207 neigh_skin_ratio=0.2,
208 release_GIL=False,
209 debug=False,
210 *args,
211 **kwargs
212 ):
213 super().__init__(*args, **kwargs)
215 self.model_name = model_name
216 self.release_GIL = release_GIL
217 self.debug = debug
219 if neigh_skin_ratio < 0:
220 raise ValueError('Argument "neigh_skin_ratio" must be non-negative')
221 self.neigh_skin_ratio = neigh_skin_ratio
223 # Model output
224 self.energy = None
225 self.forces = None
227 # Create KIMModelData object. This will take care of creating
228 # and storing the KIM API Portable Model object, KIM API
229 # ComputeArguments object, and the neighbor list object that
230 # our calculator needs
231 self._kimmodeldata = KIMModelData(
232 self.model_name, ase_neigh, self.neigh_skin_ratio, self.debug
233 )
235 self._parameters_changed = False
237 def __enter__(self):
238 return self
240 def __exit__(self, exc_type, value, traceback):
241 pass
243 def __repr__(self):
244 return f"KIMModelCalculator(model_name={self.model_name})"
246 def calculate(
247 self,
248 atoms=None,
249 properties=["energy", "forces", "stress"],
250 system_changes=["positions", "numbers", "cell", "pbc"],
251 ):
252 """
253 Inherited method from the ase Calculator class that is called by
254 get_property()
256 Parameters
257 ----------
258 atoms : Atoms
259 Atoms object whose properties are desired
261 properties : list of str
262 List of what needs to be calculated. Can be any combination
263 of 'energy', 'forces' and 'stress'.
265 system_changes : list of str
266 List of what has changed since last calculation. Can be any
267 combination of these six: 'positions', 'numbers', 'cell',
268 and 'pbc'.
269 """
271 super().calculate(atoms, properties, system_changes)
273 if self._parameters_changed:
274 self._parameters_changed = False
276 if system_changes:
278 # Ask model to update all of its parameters and the parameters
279 # related to the neighbor list(s). This update is necessary to do
280 # here since the user will generally have made changes the model
281 # parameters since the last time an update was performed and we
282 # need to ensure that any properties calculated here are made using
283 # the up-to-date model and neighbor list parameters.
284 self._model_refresh_and_update_neighbor_list_parameters()
286 if self._need_neigh_update(atoms, system_changes):
287 self._update_neigh(atoms, self._species_map)
288 self.energy = np.array([0.0], dtype=kimpy_wrappers.c_double)
289 self.forces = np.zeros(
290 [self._num_particles[0], 3], dtype=kimpy_wrappers.c_double
291 )
292 self._update_compute_args_pointers(self.energy, self.forces)
293 else:
294 self._update_kim_coords(atoms)
296 self._kim_model.compute(self._compute_args, self.release_GIL)
298 energy = self.energy[0]
299 forces = self._assemble_padding_forces()
301 try:
302 volume = atoms.get_volume()
303 stress = self._compute_virial_stress(
304 self.forces, self._coords, volume)
305 except ValueError: # Volume cannot be computed
306 stress = None
308 # Quantities passed back to ASE
309 self.results["energy"] = energy
310 self.results["free_energy"] = energy
311 self.results["forces"] = forces
312 self.results["stress"] = stress
314 def check_state(self, atoms, tol=1e-15):
315 # Check for change in atomic configuration (positions or pbc)
316 system_changes = compare_atoms(
317 self.atoms, atoms, excluded_properties=self.ignored_changes
318 )
320 # Check if model parameters were changed
321 if self._parameters_changed:
322 system_changes.append("calculator")
324 return system_changes
326 def _assemble_padding_forces(self):
327 """
328 Assemble forces on padding atoms back to contributing atoms.
330 Parameters
331 ----------
332 forces : 2D array of doubles
333 Forces on both contributing and padding atoms
335 num_contrib: int
336 Number of contributing atoms
338 padding_image_of : 1D array of int
339 Atom number, of which the padding atom is an image
342 Returns
343 -------
344 Total forces on contributing atoms.
345 """
347 total_forces = np.array(self.forces[:self._num_contributing_particles])
349 if self._padding_image_of.size != 0:
350 pad_forces = self.forces[self._num_contributing_particles:]
351 for f, org_index in zip(pad_forces, self._padding_image_of):
352 total_forces[org_index] += f
354 return total_forces
356 @staticmethod
357 def _compute_virial_stress(forces, coords, volume):
358 """Compute the virial stress in Voigt notation.
360 Parameters
361 ----------
362 forces : 2D array
363 Partial forces on all atoms (padding included)
365 coords : 2D array
366 Coordinates of all atoms (padding included)
368 volume : float
369 Volume of cell
371 Returns
372 -------
373 stress : 1D array
374 stress in Voigt order (xx, yy, zz, yz, xz, xy)
375 """
376 stress = np.zeros(6)
377 stress[0] = -np.dot(forces[:, 0], coords[:, 0]) / volume
378 stress[1] = -np.dot(forces[:, 1], coords[:, 1]) / volume
379 stress[2] = -np.dot(forces[:, 2], coords[:, 2]) / volume
380 stress[3] = -np.dot(forces[:, 1], coords[:, 2]) / volume
381 stress[4] = -np.dot(forces[:, 0], coords[:, 2]) / volume
382 stress[5] = -np.dot(forces[:, 0], coords[:, 1]) / volume
384 return stress
386 @property
387 def _update_compute_args_pointers(self):
388 return self._kimmodeldata.update_compute_args_pointers
390 @property
391 def _kim_model(self):
392 return self._kimmodeldata.kim_model
394 @property
395 def _compute_args(self):
396 return self._kimmodeldata.compute_args
398 @property
399 def _num_particles(self):
400 return self._kimmodeldata.num_particles
402 @property
403 def _coords(self):
404 return self._kimmodeldata.coords
406 @property
407 def _padding_image_of(self):
408 return self._kimmodeldata.padding_image_of
410 @property
411 def _species_map(self):
412 return self._kimmodeldata.species_map
414 @property
415 def _neigh(self):
416 # WARNING: This property is underscored for a reason! The
417 # neighborlist(s) itself (themselves) may not be up to date with
418 # respect to changes that have been made to the model's parameters, or
419 # even since the positions in the Atoms object may have changed.
420 # Neighbor lists are only potentially updated inside the ``calculate``
421 # method.
422 return self._kimmodeldata.neigh
424 @property
425 def _num_contributing_particles(self):
426 return self._neigh.num_contributing_particles
428 @property
429 def _update_kim_coords(self):
430 return self._neigh.update_kim_coords
432 @property
433 def _need_neigh_update(self):
434 return self._neigh.need_neigh_update
436 @property
437 def _update_neigh(self):
438 return self._neigh.update
440 @property
441 def parameters_metadata(self):
442 return self._kim_model.parameters_metadata
444 @property
445 def parameter_names(self):
446 return self._kim_model.parameter_names
448 @property
449 def get_parameters(self):
450 # Ask model to update all of its parameters and the parameters related
451 # to the neighbor list(s). This update is necessary to do here since
452 # the user will generally have made changes the model parameters since
453 # the last time an update was performed and we need to ensure the
454 # parameters returned by this method are fully up to date.
455 self._model_refresh_and_update_neighbor_list_parameters()
457 return self._kim_model.get_parameters
459 def set_parameters(self, **kwargs):
460 parameters = self._kim_model.set_parameters(**kwargs)
461 self._parameters_changed = True
463 return parameters
465 def _model_refresh_and_update_neighbor_list_parameters(self):
466 """
467 Call the model's refresh routine and update the neighbor list object
468 for any necessary changes arising from changes to the model parameters,
469 e.g. a change in one of its cutoffs. After a model's parameters have
470 been changed, this method *must* be called before calling the model's
471 compute routine.
472 """
473 self._kim_model.clear_then_refresh()
475 # Update neighbor list parameters
476 (
477 model_influence_dist,
478 model_cutoffs,
479 padding_not_require_neigh,
480 ) = self._kimmodeldata.get_model_neighbor_list_parameters()
482 self._neigh.set_neigh_parameters(
483 self.neigh_skin_ratio,
484 model_influence_dist,
485 model_cutoffs,
486 padding_not_require_neigh,
487 )