Coverage for mchammer/configuration_manager.py: 95%

1import random 

2import numpy as np 

3from typing import Dict, List, Tuple 

4from ase import Atoms 

5from icet.core.sublattices import Sublattices 

6from icet.tools.geometry import atomic_number_to_chemical_symbol 

7 

8 

9class SwapNotPossibleError(Exception): 

10 pass 

11 

12 

13class ConfigurationManager(object): 

14 """ 

15 The ConfigurationManager owns and handles information pertaining to a 

16 configuration being sampled in a Monte Carlo simulation. 

17 

18 Note 

19 ---- 

20 As a user you will usually not interact directly with objects of this type. 

21 

22 Parameters 

23 ---------- 

24 structure : Atoms 

25 Configuration to be handled. 

26 sublattices : :class:`Sublattices <icet.core.sublattices.Sublattices>` 

27 Sublattices used to define allowed occupations and handle related information. 

28 """ 

29 

30 def __init__(self, structure: Atoms, sublattices: Sublattices) -> None: 

31 self._structure = structure.copy() 

32 self._occupations = self._structure.numbers 

33 self._sublattices = sublattices 

34 self._sites_by_species = self._get_sites_by_species() 

35 

36 def _get_sites_by_species(self) -> List[Dict[int, List[int]]]: 

37 """Returns the sites that are occupied for each species. Each 

38 dictionary represents one sublattice where the key is the 

39 species (by atomic number) and the value is the list of sites 

40 occupied by said species in the respective sublattice. 

41 """ 

42 sites_by_species = [] 

43 for sl in self._sublattices: 

44 species_dict = {key: [] for key in sl.atomic_numbers} 

45 for site in sl.indices: 

46 species_dict[self._occupations[site]].append(site) 

47 sites_by_species.append(species_dict) 

48 return sites_by_species 

49 

50 @property 

51 def occupations(self) -> np.ndarray: 

52 """ Occupation vector of the configuration (copy). """ 

53 return self._occupations.copy() 

54 

55 @property 

56 def sublattices(self) -> Sublattices: 

57 """ Sublattices of the configuration. """ 

58 return self._sublattices 

59 

60 @property 

61 def structure(self) -> Atoms: 

62 """ Atomic structure associated with configuration (copy). """ 

63 structure = self._structure.copy() 

64 structure.set_atomic_numbers(self.occupations) 

65 return structure 

66 

67 def get_occupations_on_sublattice(self, sublattice_index: int) -> List[int]: 

68 """ 

69 Returns the occupations on one sublattice. 

70 

71 Parameters 

72 --------- 

73 sublattice_index 

74 Sublattice by index for which the occupations should be returned. 

75 """ 

76 sl = self.sublattices[sublattice_index] 

77 return list(self.occupations[sl.indices]) 

78 

79 def is_swap_possible(self, sublattice_index: int, 

80 allowed_species: List[int] = None) -> bool: 

81 """ Checks if a swap trial move is possible on a specific sublattice. 

82 

83 Parameters 

84 ---------- 

85 sublattice_index 

86 Index of sublattice to be checked. 

87 allowed_species 

88 List of atomic numbers for allowed species. 

89 """ 

90 sl = self.sublattices[sublattice_index] 

91 if allowed_species is None: 

92 swap_symbols = set(self.occupations[sl.indices]) 

93 else: 

94 swap_symbols = set([o for o in self.occupations[sl.indices] if o in 

95 allowed_species]) 

96 return len(swap_symbols) > 1 

97 

98 def get_swapped_state(self, sublattice_index: int, 

99 allowed_species: List[int] = None 

100 ) -> Tuple[List[int], List[int]]: 

101 """Returns two random sites (first element of tuple) and their 

102 occupation after a swap (second element of tuple). The new 

103 configuration will obey the occupation constraints associated 

104 with the :class:`ConfigurationManager` object. 

105 

106 Parameters 

107 ---------- 

108 sublattice_index 

109 Sublattice by index from which to pick sites. 

110 allowed_species 

111 List of atomic numbers for allowed species. 

112 """ 

113 # pick the first site 

114 if allowed_species is None: 

115 available_sites = self.sublattices[sublattice_index].indices 

116 else: 

117 available_sites = [ 

118 s for Z in allowed_species for s in 

119 self._get_sites_by_species()[sublattice_index][Z]] 

120 

121 try: 

122 site1 = random.choice(available_sites) 

123 except IndexError: 

124 raise SwapNotPossibleError(f'Sublattice {sublattice_index} is empty.') 

125 

126 # pick the second site 

127 if allowed_species is None: 

128 possible_swap_species = \ 

129 set(self._sublattices.get_allowed_numbers_on_site(site1)) - \ 

130 set([self._occupations[site1]]) 

131 else: 

132 possible_swap_species = \ 

133 set(allowed_species) - set([self._occupations[site1]]) 

134 possible_swap_sites = [] 

135 for Z in possible_swap_species: 

136 possible_swap_sites.extend(self._sites_by_species[sublattice_index][Z]) 

137 

138 possible_swap_sites = np.array(possible_swap_sites) 

139 

140 try: 

141 site2 = random.choice(possible_swap_sites) 

142 except IndexError: 

143 raise SwapNotPossibleError( 

144 'Cannot swap on sublattice {} since it is full of {} species .' 

145 .format(sublattice_index, 

146 atomic_number_to_chemical_symbol([self._occupations[site1]])[0])) 

147 

148 return ([site1, site2], [self._occupations[site2], self._occupations[site1]]) 

149 

150 def get_flip_state( 

151 self, 

152 sublattice_index: int, 

153 allowed_species: List[int] = None, 

154 ) -> Tuple[int, int]: 

155 """ 

156 Returns a site index and a new species for the site. 

157 

158 Parameters 

159 ---------- 

160 sublattice_index 

161 Index of sublattice from which to pick a site. 

162 allowed_species 

163 List of atomic numbers for allowed species. 

164 """ 

165 if allowed_species is None: 

166 available_sites = self._sublattices[sublattice_index].indices 

167 else: 

168 available_sites = [s for Z in allowed_species for s in 

169 self._get_sites_by_species()[sublattice_index][Z]] 

170 

171 site = random.choice(available_sites) 

172 if allowed_species is not None: 

173 species = random.choice(list( 

174 set(allowed_species) - set([self._occupations[site]]))) 

175 else: 

176 species = random.choice(list( 

177 set(self._sublattices[sublattice_index].atomic_numbers) - 

178 set([self._occupations[site]]))) 

179 return site, species 

180 

181 def update_occupations(self, sites: List[int], species: List[int]) -> None: 

182 """ 

183 Updates the occupation vector of the configuration being sampled. 

184 This will change the state in both the configuration in the calculator 

185 and the configuration manager. 

186 

187 Parameters 

188 ---------- 

189 sites 

190 Indices of sites of the configuration to change. 

191 species 

192 New occupations by atomic number. 

193 """ 

194 

195 # Update sublattices 

196 for site, new_Z in zip(sites, species): 

197 if 0 > new_Z > 118: 197 ↛ 198line 197 didn't jump to line 198, because the condition on line 197 was never true

198 raise ValueError('Invalid new species {} on site {}'.format(new_Z, site)) 

199 if len(self._occupations) >= site < 0: 

200 raise ValueError('Site {} is not a valid site index'.format(site)) 

201 old_Z = self._occupations[site] 

202 sublattice_index = self.sublattices.get_sublattice_index_from_site_index(site) 

203 

204 if new_Z not in self.sublattices[sublattice_index].atomic_numbers: 

205 raise ValueError('Invalid new species {} on site {}'.format(new_Z, site)) 

206 

207 # Remove site from list of sites for old species 

208 self._sites_by_species[sublattice_index][old_Z].remove(site) 

209 # Add site to list of sites for new species 

210 try: 

211 self._sites_by_species[sublattice_index][new_Z].append(site) 

212 except KeyError: 

213 raise ValueError('Invalid new species {} on site {}'.format(new_Z, site)) 

214 

215 # Update occupation vector itself 

216 self._occupations[sites] = species