Source code for hiphive.calculators.ase_calculator

"""Contains a calculator which given an arbitrary list of clusters and
associated force constants can calculate the energy and forces of a displaced
system
"""
import math
from collections import defaultdict as dd
from typing import List, Tuple

import numpy as np
from ase import Atoms
from ase.calculators.calculator import Calculator, all_changes
from hiphive.force_constants import SortedForceConstants
from hiphive.utilities import get_displacements
from .numba_calc import _get_forces


[docs]class ForceConstantCalculator(Calculator): """This class provides an ASE calculator that can be used in conjunction with integrators and optimizers with the `atomic simulation environment (ASE) <https://wiki.fysik.dtu.dk/ase/index.html>`_. To initialize an object of this class one must provide the ideal atomic configuration along with a compatible force constant model. Parameters ----------- fcs : ForceConstants the force constants instance must include the atomic configuration max_disp : float maximum allowed displacement before calculator raises ValueError """ implemented_properties = ['energy', 'forces'] def __init__(self, fcs: SortedForceConstants, max_disp: float = 3.0): Calculator.__init__(self) if not isinstance(fcs, SortedForceConstants): raise TypeError('The FC calculator requires sorted FCs.') self.atoms_ideal = fcs.supercell.copy() # Nearest neighbor distance used as maximum displacement allowed, # stops exploding MD simulations. self.max_allowed_disp = max_disp self.clusters = dd(list) self.force_constants = dd(list) self.atom_indices = dd(list) self.atom_positions = dd(list) self.atom_counts = dd(list) self.prefactors = dd(list) # The main idea is to precompute the prefactor and multiplicities of # belonging to each cluster for cluster, fc in fcs.get_fc_dict().items(): argsort = np.argsort(cluster) # TODO: is already True? cluster = np.array(sorted(cluster)) nbody = len(set(cluster)) order = len(cluster) key = (nbody, order) self.clusters[key].append(cluster) assert fc.shape == (3,) * order self.force_constants[key].append(fc.transpose(argsort)) unique = np.unique(cluster, return_index=True, return_counts=True) self.atom_indices[key].append(unique[0]) self.atom_positions[key].append(unique[1]) self.atom_counts[key].append(unique[2]) prefac = [-count/np.prod(list(map(math.factorial, unique[2]))) for count in unique[2]] self.prefactors[key].append(prefac) for d in [self.clusters, self.force_constants, self.atom_indices, self.atom_positions, self.atom_counts, self.prefactors, ]: for k, v in d.items(): d[k] = np.array(v)
[docs] def calculate(self, atoms: Atoms = None, properties: List[str] = ['energy'], system_changes: List[str] = all_changes) -> None: Calculator.calculate(self, atoms, properties, system_changes) self._check_atoms() self._compute_displacements() if 'forces' in properties or 'energy' in properties: E, forces = self.compute_energy_and_forces() self.results['forces'] = forces self.results['energy'] = E
def _check_atoms(self) -> None: """Checks that the atomic configuration, with which the calculator is associated, is compatible with the ideal configuration provided during initialization.""" if len(self.atoms) != len(self.atoms_ideal): raise ValueError('Length of atoms does not match reference structure') if not all(self.atoms.numbers == self.atoms_ideal.numbers): raise ValueError('Atomic numbers do not match reference structure') def _compute_displacements(self) -> None: """Evaluates the atomic displacements between the current and the ideal (reference) configuration.""" self.displacements = get_displacements(self.atoms, self.atoms_ideal) # sanity check that displacements are not too large max_disp = np.max(np.linalg.norm(self.displacements, axis=1)) if max_disp > self.max_allowed_disp: msg = 'Displacement {:.5f} larger than maximum allowed displacement {:.5f}' raise ValueError(msg.format(max_disp, self.max_allowed_disp))
[docs] def compute_energy_and_forces(self) -> Tuple[float, list]: """Computes energy and forces. Returns ------- float, list(list(float)) energy and forces """ E = np.zeros(1) forces = np.zeros((len(self.atoms), 3)) for key in self.clusters.keys(): nbody = key[0] order = key[1] _get_forces(self.clusters[key], self.force_constants[key], self.atom_indices[key], self.atom_positions[key], self.atom_counts[key], self.prefactors[key], nbody, order, forces, E, self.displacements) return float(E), forces
def __repr__(self) -> str: fc_dict_str = '{{{}: {}, ...}}'.format(self.clusters[0], self.force_constants[0]) fcs_str = 'ForceConstants(fc_dict={}, atoms={!r})'.format(fc_dict_str, self.atoms_ideal) return 'ForceConstantCalculator({})'.format(fcs_str)