Coverage for /builds/ase/ase/ase/calculators/mixing.py: 90.67%

1# fmt: off 

2 

3from ase.calculators.calculator import ( 

4 BaseCalculator, 

5 CalculatorSetupError, 

6 all_changes, 

7) 

8from ase.stress import full_3x3_to_voigt_6_stress 

9 

10 

11class Mixer: 

12 def __init__(self, calcs, weights): 

13 self.check_input(calcs, weights) 

14 common_properties = set.intersection( 

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

16 ) 

17 self.implemented_properties = list(common_properties) 

18 self.calcs = calcs 

19 self.weights = weights 

20 

21 @staticmethod 

22 def check_input(calcs, weights): 

23 if len(calcs) == 0: 

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

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

26 raise ValueError( 

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

28 " the number of Calculators!" 

29 ) 

30 

31 def get_properties(self, properties, atoms): 

32 results = {} 

33 

34 def get_property(prop): 

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

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

37 if prop == "stress": 

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

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

40 if prop == "stress": 

41 contribs = self.make_stress_voigt(contribs) 

42 else: 

43 raise ValueError( 

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

45 " are not the same from all" 

46 " calculators" 

47 ) 

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

49 results[prop] = sum( 

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

51 ) 

52 

53 for prop in properties: # get requested properties 

54 get_property(prop) 

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

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

57 get_property(prop) 

58 return results 

59 

60 @staticmethod 

61 def make_stress_voigt(stresses): 

62 new_contribs = [] 

63 for contrib in stresses: 

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

65 new_contribs.append(contrib) 

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

67 new_cont = full_3x3_to_voigt_6_stress(contrib) 

68 new_contribs.append(new_cont) 

69 else: 

70 raise ValueError( 

71 "The shapes of the stress" 

72 " property are not the same" 

73 " from all calculators" 

74 ) 

75 return new_contribs 

76 

77 

78class LinearCombinationCalculator(BaseCalculator): 

79 """Weighted summation of multiple calculators.""" 

80 

81 def __init__(self, calcs, weights): 

82 """Implementation of sum of calculators. 

83 

84 calcs: list 

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

86 weights: list of float 

87 Weights for each calculator in the list. 

88 """ 

89 super().__init__() 

90 self.mixer = Mixer(calcs, weights) 

91 self.implemented_properties = self.mixer.implemented_properties 

92 

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

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

95 returns with the summed value. 

96 

97 """ 

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

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

100 

101 def __str__(self): 

102 calculators = ", ".join( 

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

104 ) 

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

106 

107 

108class MixedCalculator(LinearCombinationCalculator): 

109 """ 

110 Mixing of two calculators with different weights 

111 

112 H = weight1 * H1 + weight2 * H2 

113 

114 Has functionality to get the energy contributions from each calculator 

115 

116 Parameters 

117 ---------- 

118 calc1 : ASE-calculator 

119 calc2 : ASE-calculator 

120 weight1 : float 

121 weight for calculator 1 

122 weight2 : float 

123 weight for calculator 2 

124 """ 

125 

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

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

128 

129 def set_weights(self, w1, w2): 

130 self.mixer.weights[0] = w1 

131 self.mixer.weights[1] = w2 

132 

133 def get_energy_contributions(self, atoms=None): 

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

135 self.calculate( 

136 properties=["energy"], 

137 atoms=atoms, 

138 system_changes=all_changes 

139 ) 

140 return self.results["energy_contributions"] 

141 

142 

143class SumCalculator(LinearCombinationCalculator): 

144 """SumCalculator for combining multiple calculators. 

145 

146 This calculator can be used when there are different calculators 

147 for the different chemical environment or for example during delta 

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

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

150 properties are the intersection of the implemented properties in 

151 each calculator. 

152 

153 """ 

154 

155 def __init__(self, calcs): 

156 """Implementation of sum of calculators. 

157 

158 calcs: list 

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

160 """ 

161 

162 weights = [1.0] * len(calcs) 

163 super().__init__(calcs, weights) 

164 

165 

166class AverageCalculator(LinearCombinationCalculator): 

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

168 thermodynamic purposes).""" 

169 

170 def __init__(self, calcs): 

171 """Implementation of average of calculators. 

172 

173 calcs: list 

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

175 """ 

176 n = len(calcs) 

177 

178 if n == 0: 

179 raise CalculatorSetupError( 

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

181 ) 

182 

183 weights = [1 / n] * n 

184 super().__init__(calcs, weights)