Coverage for /builds/debichem-team/python-ase/ase/calculators/kim/kimpy_wrappers.py : 67.17%
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"""
2Wrappers that provide a minimal interface to kimpy methods and objects
4Daniel S. Karls
5University of Minnesota
6"""
8import functools
10import kimpy
12from .exceptions import KIMModelNotFound, KIMModelInitializationError, KimpyError
15def check_call(f, *args, **kwargs):
16 """
17 Call a kimpy function using its arguments and, if a RuntimeError is raised,
18 catch it and raise a KimpyError with the exception's message.
20 (Starting with kimpy 2.0.0, a RuntimeError is the only exception type raised
21 when something goes wrong.)
22 """
23 try:
24 return f(*args, **kwargs)
25 except RuntimeError as e:
26 raise KimpyError(f'Calling kimpy function "{f.__name__}" failed:\n {str(e)}')
29def check_call_wrapper(func):
30 @functools.wraps(func)
31 def myfunc(*args, **kwargs):
32 return check_call(func, *args, **kwargs)
34 return myfunc
37# kimpy methods
38collections_create = functools.partial(check_call, kimpy.collections.create)
39model_create = functools.partial(check_call, kimpy.model.create)
40simulator_model_create = functools.partial(check_call, kimpy.simulator_model.create)
41get_species_name = functools.partial(check_call, kimpy.species_name.get_species_name)
43# kimpy attributes (here to avoid importing kimpy in higher-level modules)
44collection_item_type_portableModel = kimpy.collection_item_type.portableModel
47class ModelCollections:
48 """
49 KIM Portable Models and Simulator Models are installed/managed into
50 different "collections". In order to search through the different
51 KIM API model collections on the system, a corresponding object must
52 be instantiated. For more on model collections, see the KIM API's
53 install file:
54 https://github.com/openkim/kim-api/blob/master/INSTALL
55 """
57 def __init__(self):
58 self.collection = collections_create()
60 def __enter__(self):
61 return self
63 def __exit__(self, exc_type, value, traceback):
64 pass
66 def get_item_type(self, model_name):
67 try:
68 model_type = check_call(self.collection.get_item_type, model_name)
69 except KimpyError:
70 msg = (
71 "Could not find model {} installed in any of the KIM API model "
72 "collections on this system. See "
73 "https://openkim.org/doc/usage/obtaining-models/ for instructions on "
74 "installing models.".format(model_name)
75 )
76 raise KIMModelNotFound(msg)
78 return model_type
80 @property
81 def initialized(self):
82 return hasattr(self, "collection")
85class PortableModel:
86 """Creates a KIM API Portable Model object and provides a minimal interface to it"""
88 def __init__(self, model_name, debug):
89 self.model_name = model_name
90 self.debug = debug
92 # Create KIM API Model object
93 units_accepted, self.kim_model = model_create(
94 kimpy.numbering.zeroBased,
95 kimpy.length_unit.A,
96 kimpy.energy_unit.eV,
97 kimpy.charge_unit.e,
98 kimpy.temperature_unit.K,
99 kimpy.time_unit.ps,
100 self.model_name,
101 )
103 if not units_accepted:
104 raise KIMModelInitializationError(
105 "Requested units not accepted in kimpy.model.create"
106 )
108 if self.debug:
109 l_unit, e_unit, c_unit, te_unit, ti_unit = check_call(
110 self.kim_model.get_units
111 )
112 print("Length unit is: {}".format(l_unit))
113 print("Energy unit is: {}".format(e_unit))
114 print("Charge unit is: {}".format(c_unit))
115 print("Temperature unit is: {}".format(te_unit))
116 print("Time unit is: {}".format(ti_unit))
117 print()
119 def __enter__(self):
120 return self
122 def __exit__(self, exc_type, value, traceback):
123 pass
125 def get_model_supported_species_and_codes(self):
126 """Get all of the supported species for this model and their
127 corresponding integer codes that are defined in the KIM API
129 Returns
130 -------
131 species : list of str
132 Abbreviated chemical symbols of all species the mmodel
133 supports (e.g. ["Mo", "S"])
135 codes : list of int
136 Integer codes used by the model for each species (order
137 corresponds to the order of ``species``)
138 """
139 species = []
140 codes = []
141 num_kim_species = kimpy.species_name.get_number_of_species_names()
143 for i in range(num_kim_species):
144 species_name = get_species_name(i)
145 species_support, code = self.get_species_support_and_code(species_name)
147 if species_support:
148 species.append(str(species_name))
149 codes.append(code)
151 return species, codes
153 @check_call_wrapper
154 def compute(self, compute_args_wrapped, release_GIL):
155 return self.kim_model.compute(compute_args_wrapped.compute_args, release_GIL)
157 @check_call_wrapper
158 def get_species_support_and_code(self, species_name):
159 return self.kim_model.get_species_support_and_code(species_name)
161 @check_call_wrapper
162 def get_influence_distance(self):
163 return self.kim_model.get_influence_distance()
165 @check_call_wrapper
166 def get_neighbor_list_cutoffs_and_hints(self):
167 return self.kim_model.get_neighbor_list_cutoffs_and_hints()
169 def compute_arguments_create(self):
170 return ComputeArguments(self, self.debug)
172 @property
173 def initialized(self):
174 return hasattr(self, "kim_model")
177class ComputeArguments:
178 """
179 Creates a KIM API ComputeArguments object from a KIM Portable Model object and
180 configures it for ASE. A ComputeArguments object is associated with a KIM Portable
181 Model and is used to inform the KIM API of what the model can compute. It is also
182 used to register the data arrays that allow the KIM API to pass the atomic
183 coordinates to the model and retrieve the corresponding energy and forces, etc.
184 """
186 def __init__(self, kim_model_wrapped, debug):
187 self.kim_model_wrapped = kim_model_wrapped
188 self.debug = debug
190 # Create KIM API ComputeArguments object
191 self.compute_args = check_call(
192 self.kim_model_wrapped.kim_model.compute_arguments_create
193 )
195 # Check compute arguments
196 kimpy_arg_name = kimpy.compute_argument_name
197 num_arguments = kimpy_arg_name.get_number_of_compute_argument_names()
198 if self.debug:
199 print("Number of compute_args: {}".format(num_arguments))
201 for i in range(num_arguments):
202 name = check_call(kimpy_arg_name.get_compute_argument_name, i)
203 dtype = check_call(kimpy_arg_name.get_compute_argument_data_type, name)
205 arg_support = self.get_argument_support_status(name)
207 if self.debug:
208 print(
209 "Compute Argument name {:21} is of type {:7} and has support "
210 "status {}".format(*[str(x) for x in [name, dtype, arg_support]])
211 )
213 # See if the model demands that we ask it for anything other than energy and
214 # forces. If so, raise an exception.
215 if arg_support == kimpy.support_status.required:
216 if (
217 name != kimpy.compute_argument_name.partialEnergy
218 and name != kimpy.compute_argument_name.partialForces
219 ):
220 raise KIMModelInitializationError(
221 "Unsupported required ComputeArgument {}".format(name)
222 )
224 # Check compute callbacks
225 callback_name = kimpy.compute_callback_name
226 num_callbacks = callback_name.get_number_of_compute_callback_names()
227 if self.debug:
228 print()
229 print("Number of callbacks: {}".format(num_callbacks))
231 for i in range(num_callbacks):
232 name = check_call(callback_name.get_compute_callback_name, i)
234 support_status = self.get_callback_support_status(name)
236 if self.debug:
237 print(
238 "Compute callback {:17} has support status {}".format(
239 str(name), support_status
240 )
241 )
243 # Cannot handle any "required" callbacks
244 if support_status == kimpy.support_status.required:
245 raise KIMModelInitializationError(
246 "Unsupported required ComputeCallback: {}".format(name)
247 )
249 @check_call_wrapper
250 def set_argument_pointer(self, compute_arg_name, data_object):
251 return self.compute_args.set_argument_pointer(compute_arg_name, data_object)
253 @check_call_wrapper
254 def get_argument_support_status(self, name):
255 return self.compute_args.get_argument_support_status(name)
257 @check_call_wrapper
258 def get_callback_support_status(self, name):
259 return self.compute_args.get_callback_support_status(name)
261 @check_call_wrapper
262 def set_callback(self, compute_callback_name, callback_function, data_object):
263 return self.compute_args.set_callback(
264 compute_callback_name, callback_function, data_object
265 )
267 @check_call_wrapper
268 def set_callback_pointer(self, compute_callback_name, callback, data_object):
269 return self.compute_args.set_callback_pointer(
270 compute_callback_name, callback, data_object
271 )
273 def update(
274 self, num_particles, species_code, particle_contributing, coords, energy, forces
275 ):
276 """Register model input and output in the kim_model object."""
277 compute_arg_name = kimpy.compute_argument_name
278 set_argument_pointer = self.set_argument_pointer
280 set_argument_pointer(compute_arg_name.numberOfParticles, num_particles)
281 set_argument_pointer(compute_arg_name.particleSpeciesCodes, species_code)
282 set_argument_pointer(
283 compute_arg_name.particleContributing, particle_contributing
284 )
285 set_argument_pointer(compute_arg_name.coordinates, coords)
286 set_argument_pointer(compute_arg_name.partialEnergy, energy)
287 set_argument_pointer(compute_arg_name.partialForces, forces)
289 if self.debug:
290 print("Debug: called update_kim")
291 print()
294class SimulatorModel:
295 """Creates a KIM API Simulator Model object and provides a minimal
296 interface to it. This is only necessary in this package in order to
297 extract any information about a given simulator model because it is
298 generally embedded in a shared object.
299 """
301 def __init__(self, model_name):
302 # Create a KIM API Simulator Model object for this model
303 self.model_name = model_name
304 self.simulator_model = simulator_model_create(self.model_name)
306 # Need to close template map in order to access simulator model metadata
307 self.simulator_model.close_template_map()
309 def __enter__(self):
310 return self
312 def __exit__(self, exc_type, value, traceback):
313 pass
315 @property
316 def simulator_name(self):
317 simulator_name, _ = self.simulator_model.get_simulator_name_and_version()
318 return simulator_name
320 @property
321 def num_supported_species(self):
322 num_supported_species = self.simulator_model.get_number_of_supported_species()
323 if num_supported_species == 0:
324 raise KIMModelInitializationError(
325 "Unable to determine supported species of simulator model {}.".format(
326 self.model_name
327 )
328 )
329 else:
330 return num_supported_species
332 @property
333 def supported_species(self):
334 supported_species = []
335 for spec_code in range(self.num_supported_species):
336 species = check_call(self.simulator_model.get_supported_species, spec_code)
337 supported_species.append(species)
339 return tuple(supported_species)
341 @property
342 def num_metadata_fields(self):
343 return self.simulator_model.get_number_of_simulator_fields()
345 @property
346 def metadata(self):
347 sm_metadata_fields = {}
348 for field in range(self.num_metadata_fields):
349 extent, field_name = check_call(
350 self.simulator_model.get_simulator_field_metadata, field
351 )
352 sm_metadata_fields[field_name] = []
353 for ln in range(extent):
354 field_line = check_call(
355 self.simulator_model.get_simulator_field_line, field, ln
356 )
357 sm_metadata_fields[field_name].append(field_line)
359 return sm_metadata_fields
361 @property
362 def supported_units(self):
363 try:
364 supported_units = self.metadata["units"][0]
365 except (KeyError, IndexError):
366 raise KIMModelInitializationError(
367 "Unable to determine supported units of simulator model {}.".format(
368 self.model_name
369 )
370 )
372 return supported_units
374 @property
375 def atom_style(self):
376 """
377 See if a 'model-init' field exists in the SM metadata and, if
378 so, whether it contains any entries including an "atom_style"
379 command. This is specific to LAMMPS SMs and is only required
380 for using the LAMMPSrun calculator because it uses
381 lammps.inputwriter to create a data file. All other content in
382 'model-init', if it exists, is ignored.
383 """
384 atom_style = None
385 for ln in self.metadata.get("model-init", []):
386 if ln.find("atom_style") != -1:
387 atom_style = ln.split()[1]
389 return atom_style
391 @property
392 def model_defn(self):
393 return self.metadata["model-defn"]
395 @property
396 def initialized(self):
397 return hasattr(self, "simulator_model")