Coverage for /builds/hweiske/ase/ase/calculators/mixing.py: 89.61%

77 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-04-22 11:22 +0000

1from ase.calculators.calculator import (BaseCalculator, CalculatorSetupError, 

2 PropertyNotImplementedError, 

3 all_changes) 

4from ase.stress import full_3x3_to_voigt_6_stress 

5 

6 

7class Mixer: 

8 def __init__(self, calcs, weights): 

9 self.check_input(calcs, weights) 

10 common_properties = set.intersection( 

11 *(set(calc.implemented_properties) for calc in calcs) 

12 ) 

13 self.implemented_properties = list(common_properties) 

14 if not self.implemented_properties: 

15 raise PropertyNotImplementedError( 

16 "The provided Calculators have" 

17 " no properties in common!" 

18 ) 

19 self.calcs = calcs 

20 self.weights = weights 

21 

22 @staticmethod 

23 def check_input(calcs, weights): 

24 if len(calcs) == 0: 

25 raise CalculatorSetupError("Please provide a list of Calculators") 

26 if len(weights) != len(calcs): 

27 raise ValueError( 

28 "The length of the weights must be the same as" 

29 " the number of Calculators!" 

30 ) 

31 

32 def get_properties(self, properties, atoms): 

33 results = {} 

34 

35 def get_property(prop): 

36 contribs = [calc.get_property(prop, atoms) for calc in self.calcs] 

37 # ensure that the contribution shapes are the same for stress prop 

38 if prop == "stress": 

39 shapes = [contrib.shape for contrib in contribs] 

40 if not all(shape == shapes[0] for shape in shapes): 

41 if prop == "stress": 

42 contribs = self.make_stress_voigt(contribs) 

43 else: 

44 raise ValueError( 

45 f"The shapes of the property {prop}" 

46 " are not the same from all" 

47 " calculators" 

48 ) 

49 results[f"{prop}_contributions"] = contribs 

50 results[prop] = sum( 

51 weight * value for weight, value in zip(self.weights, contribs) 

52 ) 

53 

54 for prop in properties: # get requested properties 

55 get_property(prop) 

56 for prop in self.implemented_properties: # cache all available props 

57 if all(prop in calc.results for calc in self.calcs): 

58 get_property(prop) 

59 return results 

60 

61 @staticmethod 

62 def make_stress_voigt(stresses): 

63 new_contribs = [] 

64 for contrib in stresses: 

65 if contrib.shape == (6,): 

66 new_contribs.append(contrib) 

67 elif contrib.shape == (3, 3): 

68 new_cont = full_3x3_to_voigt_6_stress(contrib) 

69 new_contribs.append(new_cont) 

70 else: 

71 raise ValueError( 

72 "The shapes of the stress" 

73 " property are not the same" 

74 " from all calculators" 

75 ) 

76 return new_contribs 

77 

78 

79class LinearCombinationCalculator(BaseCalculator): 

80 """Weighted summation of multiple calculators.""" 

81 

82 def __init__(self, calcs, weights): 

83 """Implementation of sum of calculators. 

84 

85 calcs: list 

86 List of an arbitrary number of :mod:`ase.calculators` objects. 

87 weights: list of float 

88 Weights for each calculator in the list. 

89 """ 

90 super().__init__() 

91 self.mixer = Mixer(calcs, weights) 

92 self.implemented_properties = self.mixer.implemented_properties 

93 

94 def calculate(self, atoms, properties, system_changes): 

95 """Calculates all the specific property for each calculator and 

96 returns with the summed value. 

97 

98 """ 

99 self.atoms = atoms.copy() # for caching of results 

100 self.results = self.mixer.get_properties(properties, atoms) 

101 

102 def __str__(self): 

103 calculators = ", ".join( 

104 calc.__class__.__name__ for calc in self.mixer.calcs 

105 ) 

106 return f"{self.__class__.__name__}({calculators})" 

107 

108 

109class MixedCalculator(LinearCombinationCalculator): 

110 """ 

111 Mixing of two calculators with different weights 

112 

113 H = weight1 * H1 + weight2 * H2 

114 

115 Has functionality to get the energy contributions from each calculator 

116 

117 Parameters 

118 ---------- 

119 calc1 : ASE-calculator 

120 calc2 : ASE-calculator 

121 weight1 : float 

122 weight for calculator 1 

123 weight2 : float 

124 weight for calculator 2 

125 """ 

126 

127 def __init__(self, calc1, calc2, weight1, weight2): 

128 super().__init__([calc1, calc2], [weight1, weight2]) 

129 

130 def set_weights(self, w1, w2): 

131 self.mixer.weights[0] = w1 

132 self.mixer.weights[1] = w2 

133 

134 def get_energy_contributions(self, atoms=None): 

135 """Return the potential energy from calc1 and calc2 respectively""" 

136 self.calculate( 

137 properties=["energy"], 

138 atoms=atoms, 

139 system_changes=all_changes 

140 ) 

141 return self.results["energy_contributions"] 

142 

143 

144class SumCalculator(LinearCombinationCalculator): 

145 """SumCalculator for combining multiple calculators. 

146 

147 This calculator can be used when there are different calculators 

148 for the different chemical environment or for example during delta 

149 leaning. It works with a list of arbitrary calculators and 

150 evaluates them in sequence when it is required. The supported 

151 properties are the intersection of the implemented properties in 

152 each calculator. 

153 

154 """ 

155 

156 def __init__(self, calcs): 

157 """Implementation of sum of calculators. 

158 

159 calcs: list 

160 List of an arbitrary number of :mod:`ase.calculators` objects. 

161 """ 

162 

163 weights = [1.0] * len(calcs) 

164 super().__init__(calcs, weights) 

165 

166 

167class AverageCalculator(LinearCombinationCalculator): 

168 """AverageCalculator for equal summation of multiple calculators (for 

169 thermodynamic purposes).""" 

170 

171 def __init__(self, calcs): 

172 """Implementation of average of calculators. 

173 

174 calcs: list 

175 List of an arbitrary number of :mod:`ase.calculators` objects. 

176 """ 

177 n = len(calcs) 

178 

179 if n == 0: 

180 raise CalculatorSetupError( 

181 "The value of the calcs must be a list of Calculators" 

182 ) 

183 

184 weights = [1 / n] * n 

185 super().__init__(calcs, weights)