# fmt: off """ ASE Calculator for interatomic models compatible with the Knowledgebase of Interatomic Models (KIM) application programming interface (API). Written by: Mingjian Wen Daniel S. Karls University of Minnesota """ import numpy as np from ase.calculators.calculator import Calculator, compare_atoms from . import kimpy_wrappers, neighborlist class KIMModelData: """Initializes and subsequently stores the KIM API Portable Model object, KIM API ComputeArguments object, and the neighbor list object used by instances of KIMModelCalculator. Also stores the arrays which are registered in the KIM API and which are used to communicate with the model. """ def __init__(self, model_name, ase_neigh, neigh_skin_ratio, debug=False): self.model_name = model_name self.ase_neigh = ase_neigh self.debug = debug # Initialize KIM API Portable Model object and ComputeArguments object self.kim_model, self.compute_args = self._init_kim() self.species_map = self._create_species_map() # Ask model to provide information relevant for neighbor list # construction ( model_influence_dist, model_cutoffs, padding_not_require_neigh, ) = self.get_model_neighbor_list_parameters() # Initialize neighbor list object self.neigh = self._init_neigh( neigh_skin_ratio, model_influence_dist, model_cutoffs, padding_not_require_neigh, ) def _init_kim(self): """Create the KIM API Portable Model object and KIM API ComputeArguments object """ if self.kim_initialized: return None kim_model = kimpy_wrappers.PortableModel(self.model_name, self.debug) # KIM API model object is what actually creates/destroys the # ComputeArguments object, so we must pass it as a parameter compute_args = kim_model.compute_arguments_create() return kim_model, compute_args def _init_neigh( self, neigh_skin_ratio, model_influence_dist, model_cutoffs, padding_not_require_neigh, ): """Initialize neighbor list, either an ASE-native neighborlist or one created using the neighlist module in kimpy """ neigh_list_object_type = ( neighborlist.ASENeighborList if self.ase_neigh else neighborlist.KimpyNeighborList ) return neigh_list_object_type( self.compute_args, neigh_skin_ratio, model_influence_dist, model_cutoffs, padding_not_require_neigh, self.debug, ) def get_model_neighbor_list_parameters(self): model_influence_dist = self.kim_model.get_influence_distance() ( model_cutoffs, padding_not_require_neigh, ) = self.kim_model.get_neighbor_list_cutoffs_and_hints() return model_influence_dist, model_cutoffs, padding_not_require_neigh def update_compute_args_pointers(self, energy, forces): self.compute_args.update( self.num_particles, self.species_code, self._particle_contributing, self.coords, energy, forces, ) def _create_species_map(self): """Get all the supported species of the KIM model and the corresponding integer codes used by the model Returns ------- species_map : dict key : str chemical symbols (e.g. "Ar") value : int species integer code (e.g. 1) """ supported_species, codes = self._get_model_supported_species_and_codes() species_map = {} for i, spec in enumerate(supported_species): species_map[spec] = codes[i] if self.debug: print( "Species {} is supported and its code is: {}".format( spec, codes[i]) ) return species_map @property def padding_image_of(self): return self.neigh.padding_image_of @property def num_particles(self): return self.neigh.num_particles @property def coords(self): return self.neigh.coords @property def _particle_contributing(self): return self.neigh.particle_contributing @property def species_code(self): return self.neigh.species_code @property def kim_initialized(self): return hasattr(self, "kim_model") @property def _neigh_initialized(self): return hasattr(self, "neigh") @property def _get_model_supported_species_and_codes(self): return self.kim_model.get_model_supported_species_and_codes class KIMModelCalculator(Calculator): """Calculator that works with KIM Portable Models (PMs). Calculator that carries out direct communication between ASE and a KIM Portable Model (PM) through the kimpy library (which provides a set of python bindings to the KIM API). Parameters ---------- model_name : str The unique identifier assigned to the interatomic model (for details, see https://openkim.org/doc/schema/kim-ids) ase_neigh : bool, optional False (default): Use kimpy's neighbor list library True: Use ASE's internal neighbor list mechanism (usually slower than the kimpy neighlist library) neigh_skin_ratio : float, optional Used to determine the neighbor list cutoff distance, r_neigh, through the relation r_neigh = (1 + neigh_skin_ratio) * rcut, where rcut is the model's influence distance. (Default: 0.2) release_GIL : bool, optional Whether to release python GIL. Releasing the GIL allows a KIM model to run with multiple concurrent threads. (Default: False) debug : bool, optional If True, detailed information is printed to stdout. (Default: False) """ implemented_properties = ["energy", "free_energy", "forces", "stress"] ignored_changes = {"initial_charges", "initial_magmoms"} def __init__( self, model_name, ase_neigh=False, neigh_skin_ratio=0.2, release_GIL=False, debug=False, *args, **kwargs ): super().__init__(*args, **kwargs) self.model_name = model_name self.release_GIL = release_GIL self.debug = debug if neigh_skin_ratio < 0: raise ValueError('Argument "neigh_skin_ratio" must be non-negative') self.neigh_skin_ratio = neigh_skin_ratio # Model output self.energy = None self.forces = None # Create KIMModelData object. This will take care of creating # and storing the KIM API Portable Model object, KIM API # ComputeArguments object, and the neighbor list object that # our calculator needs self._kimmodeldata = KIMModelData( self.model_name, ase_neigh, self.neigh_skin_ratio, self.debug ) self._parameters_changed = False def __enter__(self): return self def __exit__(self, exc_type, value, traceback): pass def __repr__(self): return f"KIMModelCalculator(model_name={self.model_name})" def calculate( self, atoms=None, properties=["energy", "forces", "stress"], system_changes=["positions", "numbers", "cell", "pbc"], ): """ Inherited method from the ase Calculator class that is called by get_property() Parameters ---------- atoms : Atoms Atoms object whose properties are desired properties : list of str List of what needs to be calculated. Can be any combination of 'energy', 'forces' and 'stress'. system_changes : list of str List of what has changed since last calculation. Can be any combination of these six: 'positions', 'numbers', 'cell', and 'pbc'. """ super().calculate(atoms, properties, system_changes) if self._parameters_changed: self._parameters_changed = False if system_changes: # Ask model to update all of its parameters and the parameters # related to the neighbor list(s). This update is necessary to do # here since the user will generally have made changes the model # parameters since the last time an update was performed and we # need to ensure that any properties calculated here are made using # the up-to-date model and neighbor list parameters. self._model_refresh_and_update_neighbor_list_parameters() if self._need_neigh_update(atoms, system_changes): self._update_neigh(atoms, self._species_map) self.energy = np.array([0.0], dtype=kimpy_wrappers.c_double) self.forces = np.zeros( [self._num_particles[0], 3], dtype=kimpy_wrappers.c_double ) self._update_compute_args_pointers(self.energy, self.forces) else: self._update_kim_coords(atoms) self._kim_model.compute(self._compute_args, self.release_GIL) energy = self.energy[0] forces = self._assemble_padding_forces() try: volume = atoms.get_volume() stress = self._compute_virial_stress( self.forces, self._coords, volume) except ValueError: # Volume cannot be computed stress = None # Quantities passed back to ASE self.results["energy"] = energy self.results["free_energy"] = energy self.results["forces"] = forces self.results["stress"] = stress def check_state(self, atoms, tol=1e-15): # Check for change in atomic configuration (positions or pbc) system_changes = compare_atoms( self.atoms, atoms, excluded_properties=self.ignored_changes ) # Check if model parameters were changed if self._parameters_changed: system_changes.append("calculator") return system_changes def _assemble_padding_forces(self): """ Assemble forces on padding atoms back to contributing atoms. Parameters ---------- forces : 2D array of doubles Forces on both contributing and padding atoms num_contrib: int Number of contributing atoms padding_image_of : 1D array of int Atom number, of which the padding atom is an image Returns ------- Total forces on contributing atoms. """ total_forces = np.array(self.forces[:self._num_contributing_particles]) if self._padding_image_of.size != 0: pad_forces = self.forces[self._num_contributing_particles:] for f, org_index in zip(pad_forces, self._padding_image_of): total_forces[org_index] += f return total_forces @staticmethod def _compute_virial_stress(forces, coords, volume): """Compute the virial stress in Voigt notation. Parameters ---------- forces : 2D array Partial forces on all atoms (padding included) coords : 2D array Coordinates of all atoms (padding included) volume : float Volume of cell Returns ------- stress : 1D array stress in Voigt order (xx, yy, zz, yz, xz, xy) """ stress = np.zeros(6) stress[0] = -np.dot(forces[:, 0], coords[:, 0]) / volume stress[1] = -np.dot(forces[:, 1], coords[:, 1]) / volume stress[2] = -np.dot(forces[:, 2], coords[:, 2]) / volume stress[3] = -np.dot(forces[:, 1], coords[:, 2]) / volume stress[4] = -np.dot(forces[:, 0], coords[:, 2]) / volume stress[5] = -np.dot(forces[:, 0], coords[:, 1]) / volume return stress @property def _update_compute_args_pointers(self): return self._kimmodeldata.update_compute_args_pointers @property def _kim_model(self): return self._kimmodeldata.kim_model @property def _compute_args(self): return self._kimmodeldata.compute_args @property def _num_particles(self): return self._kimmodeldata.num_particles @property def _coords(self): return self._kimmodeldata.coords @property def _padding_image_of(self): return self._kimmodeldata.padding_image_of @property def _species_map(self): return self._kimmodeldata.species_map @property def _neigh(self): # WARNING: This property is underscored for a reason! The # neighborlist(s) itself (themselves) may not be up to date with # respect to changes that have been made to the model's parameters, or # even since the positions in the Atoms object may have changed. # Neighbor lists are only potentially updated inside the ``calculate`` # method. return self._kimmodeldata.neigh @property def _num_contributing_particles(self): return self._neigh.num_contributing_particles @property def _update_kim_coords(self): return self._neigh.update_kim_coords @property def _need_neigh_update(self): return self._neigh.need_neigh_update @property def _update_neigh(self): return self._neigh.update @property def parameters_metadata(self): return self._kim_model.parameters_metadata @property def parameter_names(self): return self._kim_model.parameter_names @property def get_parameters(self): # Ask model to update all of its parameters and the parameters related # to the neighbor list(s). This update is necessary to do here since # the user will generally have made changes the model parameters since # the last time an update was performed and we need to ensure the # parameters returned by this method are fully up to date. self._model_refresh_and_update_neighbor_list_parameters() return self._kim_model.get_parameters def set_parameters(self, **kwargs): parameters = self._kim_model.set_parameters(**kwargs) self._parameters_changed = True return parameters def _model_refresh_and_update_neighbor_list_parameters(self): """ Call the model's refresh routine and update the neighbor list object for any necessary changes arising from changes to the model parameters, e.g. a change in one of its cutoffs. After a model's parameters have been changed, this method *must* be called before calling the model's compute routine. """ self._kim_model.clear_then_refresh() # Update neighbor list parameters ( model_influence_dist, model_cutoffs, padding_not_require_neigh, ) = self._kimmodeldata.get_model_neighbor_list_parameters() self._neigh.set_neigh_parameters( self.neigh_skin_ratio, model_influence_dist, model_cutoffs, padding_not_require_neigh, )