1from ase.calculators.calculator import BaseCalculator, all_changes 

2from ase.calculators.calculator import (PropertyNotImplementedError, 

3 CalculatorSetupError) 

4 

5 

6class Mixer: 

7 def __init__(self, calcs, weights): 

8 self.check_input(calcs, weights) 

9 common_properties = set.intersection(*(set(calc.implemented_properties) 

10 for calc in calcs)) 

11 self.implemented_properties = list(common_properties) 

12 if not self.implemented_properties: 

13 raise PropertyNotImplementedError('The provided Calculators have no' 

14 ' properties in common!') 

15 self.calcs = calcs 

16 self.weights = weights 

17 

18 @staticmethod 

19 def check_input(calcs, weights): 

20 if len(calcs) == 0: 

21 raise CalculatorSetupError('Please provide a list of Calculators') 

22 for calc in calcs: 

23 if not isinstance(calc, BaseCalculator): 

24 raise CalculatorSetupError('All Calculators should be inherited' 

25 ' form the BaseCalculator class') 

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

27 raise ValueError('The length of the weights must be the same as the' 

28 ' number of Calculators!') 

29 

30 def get_properties(self, properties, atoms): 

31 results = {} 

32 

33 def get_property(prop): 

34 contributs = [calc.get_property(prop, atoms) for calc in self.calcs] 

35 results[f'{prop}_contributions'] = contributs 

36 results[prop] = sum(weight * value for weight, value 

37 in zip(self.weights, contributs)) 

38 

39 for prop in properties: # get requested properties 

40 get_property(prop) 

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

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

43 get_property(prop) 

44 return results 

45 

46 

47class LinearCombinationCalculator(BaseCalculator): 

48 """Weighted summation of multiple calculators.""" 

49 

50 def __init__(self, calcs, weights): 

51 """Implementation of sum of calculators. 

52 

53 calcs: list 

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

55 weights: list of float 

56 Weights for each calculator in the list. 

57 """ 

58 super().__init__() 

59 self.mixer = Mixer(calcs, weights) 

60 self.implemented_properties = self.mixer.implemented_properties 

61 

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

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

64 returns with the summed value. 

65 

66 """ 

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

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

69 

70 def __str__(self): 

71 calculators = ', '.join( 

72 calc.__class__.__name__ for calc in self.mixer.calcs) 

73 return '{}({})'.format(self.__class__.__name__, calculators) 

74 

75 

76class MixedCalculator(LinearCombinationCalculator): 

77 """ 

78 Mixing of two calculators with different weights 

79 

80 H = weight1 * H1 + weight2 * H2 

81 

82 Has functionality to get the energy contributions from each calculator 

83 

84 Parameters 

85 ---------- 

86 calc1 : ASE-calculator 

87 calc2 : ASE-calculator 

88 weight1 : float 

89 weight for calculator 1 

90 weight2 : float 

91 weight for calculator 2 

92 """ 

93 

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

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

96 

97 def set_weights(self, w1, w2): 

98 self.mixer.weights[0] = w1 

99 self.mixer.weights[1] = w2 

100 

101 def get_energy_contributions(self, atoms=None): 

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

103 self.calculate( 

104 properties=['energy'], 

105 atoms=atoms, 

106 system_changes=all_changes) 

107 return self.results['energy_contributions'] 

108 

109 

110class SumCalculator(LinearCombinationCalculator): 

111 """SumCalculator for combining multiple calculators. 

112 

113 This calculator can be used when there are different calculators 

114 for the different chemical environment or for example during delta 

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

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

117 properties are the intersection of the implemented properties in 

118 each calculator. 

119 

120 """ 

121 

122 def __init__(self, calcs): 

123 """Implementation of sum of calculators. 

124 

125 calcs: list 

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

127 """ 

128 

129 weights = [1.] * len(calcs) 

130 super().__init__(calcs, weights) 

131 

132 

133class AverageCalculator(LinearCombinationCalculator): 

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

135 thermodynamic purposes).""" 

136 

137 def __init__(self, calcs): 

138 """Implementation of average of calculators. 

139 

140 calcs: list 

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

142 """ 

143 n = len(calcs) 

144 

145 if n == 0: 

146 raise CalculatorSetupError( 

147 'The value of the calcs must be a list of Calculators') 

148 

149 weights = [1 / n] * n 

150 super().__init__(calcs, weights)