Coverage for hiphive/calculators/ase_calculator.py: 88%

1"""Contains a calculator which given an arbitrary list of clusters and 

2associated force constants can calculate the energy and forces of a displaced 

3system 

4""" 

5import math 

6from collections import defaultdict as dd 

7from typing import List, Tuple 

8 

9import numpy as np 

10from ase import Atoms 

11from ase.calculators.calculator import Calculator, all_changes 

12from hiphive.force_constants import SortedForceConstants 

13from hiphive.utilities import get_displacements 

14from .numba_calc import _get_forces 

15 

16 

17class ForceConstantCalculator(Calculator): 

18 """This class provides an ASE calculator that can be used in conjunction 

19 with integrators and optimizers with the `atomic simulation environment 

20 (ASE) <https://wiki.fysik.dtu.dk/ase/index.html>`_. To initialize an object 

21 of this class one must provide the ideal atomic configuration along with a 

22 compatible force constant model. 

23 

24 Parameters 

25 ----------- 

26 fcs : ForceConstants 

27 the force constants instance must include the atomic configuration 

28 max_disp : float 

29 maximum allowed displacement before calculator raises ValueError 

30 """ 

31 

32 implemented_properties = ['energy', 'forces'] 

33 

34 def __init__(self, fcs: SortedForceConstants, max_disp: float = 3.0): 

35 Calculator.__init__(self) 

36 

37 if not isinstance(fcs, SortedForceConstants): 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 raise TypeError('The FC calculator requires sorted FCs.') 

39 self.atoms_ideal = fcs.supercell.copy() 

40 

41 # Nearest neighbor distance used as maximum displacement allowed, 

42 # stops exploding MD simulations. 

43 self.max_allowed_disp = max_disp 

44 

45 self.clusters = dd(list) 

46 self.force_constants = dd(list) 

47 self.atom_indices = dd(list) 

48 self.atom_positions = dd(list) 

49 self.atom_counts = dd(list) 

50 self.prefactors = dd(list) 

51 # The main idea is to precompute the prefactor and multiplicities of 

52 # belonging to each cluster 

53 for cluster, fc in fcs.get_fc_dict().items(): 

54 argsort = np.argsort(cluster) # TODO: is already True? 

55 cluster = np.array(sorted(cluster)) 

56 nbody = len(set(cluster)) 

57 order = len(cluster) 

58 key = (nbody, order) 

59 self.clusters[key].append(cluster) 

60 assert fc.shape == (3,) * order 

61 self.force_constants[key].append(fc.transpose(argsort)) 

62 unique = np.unique(cluster, return_index=True, return_counts=True) 

63 self.atom_indices[key].append(unique[0]) 

64 self.atom_positions[key].append(unique[1]) 

65 self.atom_counts[key].append(unique[2]) 

66 prefac = [-count/np.prod(list(map(math.factorial, unique[2]))) for count in unique[2]] 

67 self.prefactors[key].append(prefac) 

68 for d in [self.clusters, 

69 self.force_constants, 

70 self.atom_indices, 

71 self.atom_positions, 

72 self.atom_counts, 

73 self.prefactors, 

74 ]: 

75 for k, v in d.items(): 

76 d[k] = np.array(v) 

77 

78 def calculate(self, atoms: Atoms = None, 

79 properties: List[str] = ['energy'], 

80 system_changes: List[str] = all_changes) -> None: 

81 Calculator.calculate(self, atoms, properties, system_changes) 

82 self._check_atoms() 

83 self._compute_displacements() 

84 

85 if 'forces' in properties or 'energy' in properties: 85 ↛ exitline 85 didn't return from function 'calculate' because the condition on line 85 was always true

86 E, forces = self.compute_energy_and_forces() 

87 self.results['forces'] = forces 

88 self.results['energy'] = E 

89 

90 def _check_atoms(self) -> None: 

91 """Checks that the atomic configuration, with which the calculator is 

92 associated, is compatible with the ideal configuration provided during 

93 initialization.""" 

94 if len(self.atoms) != len(self.atoms_ideal): 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true

95 raise ValueError('Length of atoms does not match reference structure') 

96 if not all(self.atoms.numbers == self.atoms_ideal.numbers): 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true

97 raise ValueError('Atomic numbers do not match reference structure') 

98 

99 def _compute_displacements(self) -> None: 

100 """Evaluates the atomic displacements between the current and the ideal 

101 (reference) configuration.""" 

102 self.displacements = get_displacements(self.atoms, self.atoms_ideal) 

103 

104 # sanity check that displacements are not too large 

105 max_disp = np.max(np.linalg.norm(self.displacements, axis=1)) 

106 if max_disp > self.max_allowed_disp: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 msg = 'Displacement {:.5f} larger than maximum allowed displacement {:.5f}' 

108 raise ValueError(msg.format(max_disp, self.max_allowed_disp)) 

109 

110 def compute_energy_and_forces(self) -> Tuple[float, list]: 

111 """Computes energy and forces. 

112 

113 Returns 

114 ------- 

115 float, list(list(float)) 

116 energy and forces 

117 """ 

118 E = np.zeros(1) 

119 forces = np.zeros((len(self.atoms), 3)) 

120 

121 for key in self.clusters.keys(): 

122 nbody = key[0] 

123 order = key[1] 

124 _get_forces(self.clusters[key], 

125 self.force_constants[key], 

126 self.atom_indices[key], 

127 self.atom_positions[key], 

128 self.atom_counts[key], 

129 self.prefactors[key], 

130 nbody, order, 

131 forces, E, self.displacements) 

132 return float(E[0]), forces 

133 

134 def __repr__(self) -> str: 

135 fc_dict_str = '{{{}: {}, ...}}'.format(self.clusters[0], self.force_constants[0]) 

136 fcs_str = 'ForceConstants(fc_dict={}, atoms={!r})'.format(fc_dict_str, self.atoms_ideal) 

137 return 'ForceConstantCalculator({})'.format(fcs_str)