Coverage for /builds/debichem-team/python-ase/ase/calculators/mixing.py: 90.67%

1from ase.calculators.calculator import ( 

2 BaseCalculator, 

3 CalculatorSetupError, 

4 all_changes, 

5) 

6from ase.stress import full_3x3_to_voigt_6_stress 

7 

8 

9class Mixer: 

10 def __init__(self, calcs, weights): 

11 self.check_input(calcs, weights) 

12 common_properties = set.intersection( 

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

14 ) 

15 self.implemented_properties = list(common_properties) 

16 self.calcs = calcs 

17 self.weights = weights 

18 

19 @staticmethod 

20 def check_input(calcs, weights): 

21 if len(calcs) == 0: 

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

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

24 raise ValueError( 

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

26 " the number of Calculators!" 

27 ) 

28 

29 def get_properties(self, properties, atoms): 

30 results = {} 

31 

32 def get_property(prop): 

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

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

35 if prop == "stress": 

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

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

38 if prop == "stress": 

39 contribs = self.make_stress_voigt(contribs) 

40 else: 

41 raise ValueError( 

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

43 " are not the same from all" 

44 " calculators" 

45 ) 

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

47 results[prop] = sum( 

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

49 ) 

50 

51 for prop in properties: # get requested properties 

52 get_property(prop) 

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

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

55 get_property(prop) 

56 return results 

57 

58 @staticmethod 

59 def make_stress_voigt(stresses): 

60 new_contribs = [] 

61 for contrib in stresses: 

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

63 new_contribs.append(contrib) 

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

65 new_cont = full_3x3_to_voigt_6_stress(contrib) 

66 new_contribs.append(new_cont) 

67 else: 

68 raise ValueError( 

69 "The shapes of the stress" 

70 " property are not the same" 

71 " from all calculators" 

72 ) 

73 return new_contribs 

74 

75 

76class LinearCombinationCalculator(BaseCalculator): 

77 """Weighted summation of multiple calculators.""" 

78 

79 def __init__(self, calcs, weights): 

80 """Implementation of sum of calculators. 

81 

82 calcs: list 

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

84 weights: list of float 

85 Weights for each calculator in the list. 

86 """ 

87 super().__init__() 

88 self.mixer = Mixer(calcs, weights) 

89 self.implemented_properties = self.mixer.implemented_properties 

90 

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

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

93 returns with the summed value. 

94 

95 """ 

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

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

98 

99 def __str__(self): 

100 calculators = ", ".join( 

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

102 ) 

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

104 

105 

106class MixedCalculator(LinearCombinationCalculator): 

107 """ 

108 Mixing of two calculators with different weights 

109 

110 H = weight1 * H1 + weight2 * H2 

111 

112 Has functionality to get the energy contributions from each calculator 

113 

114 Parameters 

115 ---------- 

116 calc1 : ASE-calculator 

117 calc2 : ASE-calculator 

118 weight1 : float 

119 weight for calculator 1 

120 weight2 : float 

121 weight for calculator 2 

122 """ 

123 

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

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

126 

127 def set_weights(self, w1, w2): 

128 self.mixer.weights[0] = w1 

129 self.mixer.weights[1] = w2 

130 

131 def get_energy_contributions(self, atoms=None): 

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

133 self.calculate( 

134 properties=["energy"], 

135 atoms=atoms, 

136 system_changes=all_changes 

137 ) 

138 return self.results["energy_contributions"] 

139 

140 

141class SumCalculator(LinearCombinationCalculator): 

142 """SumCalculator for combining multiple calculators. 

143 

144 This calculator can be used when there are different calculators 

145 for the different chemical environment or for example during delta 

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

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

148 properties are the intersection of the implemented properties in 

149 each calculator. 

150 

151 """ 

152 

153 def __init__(self, calcs): 

154 """Implementation of sum of calculators. 

155 

156 calcs: list 

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

158 """ 

159 

160 weights = [1.0] * len(calcs) 

161 super().__init__(calcs, weights) 

162 

163 

164class AverageCalculator(LinearCombinationCalculator): 

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

166 thermodynamic purposes).""" 

167 

168 def __init__(self, calcs): 

169 """Implementation of average of calculators. 

170 

171 calcs: list 

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

173 """ 

174 n = len(calcs) 

175 

176 if n == 0: 

177 raise CalculatorSetupError( 

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

179 ) 

180 

181 weights = [1 / n] * n 

182 super().__init__(calcs, weights)