Coverage for icet/core/sublattices.py: 100%

1from copy import deepcopy 

2from icet.core.structure import Structure 

3from typing import List, Iterator 

4from ase import Atoms 

5import copy 

6from itertools import product 

7from string import ascii_uppercase 

8import numpy as np 

9from icet.tools.geometry import chemical_symbols_to_numbers 

10 

11 

12class Sublattice: 

13 """ 

14 This class stores and provides information about a specific 

15 sublattice. A sublattice is always supercell specific since 

16 it contains lattice indices. 

17 

18 Note 

19 ---- 

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

21 

22 Parameters 

23 ---------- 

24 chemical_symbols 

25 The allowed species on this sublattice. 

26 indices 

27 The lattice indices the sublattice consists of. 

28 symbol 

29 String used to mark the sublattice. 

30 """ 

31 

32 def __init__(self, 

33 chemical_symbols: List[str], 

34 indices: List[int], 

35 symbol: str): 

36 self._chemical_symbols = chemical_symbols 

37 self._indices = indices 

38 self._symbol = symbol 

39 self._numbers = chemical_symbols_to_numbers(chemical_symbols) 

40 

41 @property 

42 def chemical_symbols(self): 

43 return copy.deepcopy(self._chemical_symbols) 

44 

45 @property 

46 def atomic_numbers(self): 

47 return self._numbers.copy() 

48 

49 @property 

50 def indices(self): 

51 return self._indices.copy() 

52 

53 @property 

54 def symbol(self): 

55 """Symbol representation of sublattice, i.e. A, B, C, etc.""" 

56 return self._symbol 

57 

58 

59class Sublattices: 

60 """ 

61 This class stores and provides information about the sublattices 

62 of a structure. 

63 

64 Note 

65 ---- 

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

67 

68 Parameters 

69 ---------- 

70 allowed_species 

71 List of the allowed species on each site of the primitve 

72 structure. For example this can be the chemical_symbols from 

73 a cluster space. 

74 primitive_structure 

75 The primitive structure the allowed species reference to. 

76 structure 

77 The structure that the sublattices will be based on. 

78 fractional_position_tolerance 

79 Tolerance applied when comparing positions in fractional coordinates. 

80 """ 

81 

82 def __init__(self, 

83 allowed_species: List[List[str]], 

84 primitive_structure: Atoms, 

85 structure: Atoms, 

86 fractional_position_tolerance: float): 

87 self._structure = structure 

88 # sorted unique sites, this basically decides A, B, C... sublattices 

89 active_lattices = sorted(set([tuple(sorted(symbols)) 

90 for symbols in allowed_species if len(symbols) > 1])) 

91 inactive_lattices = sorted( 

92 set([tuple(sorted(symbols)) for symbols in allowed_species if len(symbols) == 1])) 

93 self._allowed_species = active_lattices + inactive_lattices 

94 

95 n = int(np.sqrt(len(self._allowed_species))) + 1 

96 symbols = [''.join(p) for r in range(1, n+1) for p in product(ascii_uppercase, repeat=r)] 

97 

98 cpp_prim_structure = Structure.from_atoms(primitive_structure) 

99 self._sublattices = [] 

100 sublattice_to_indices = [[] for _ in range(len(self._allowed_species))] 

101 for index, position in enumerate(structure.positions): 

102 lattice_site = cpp_prim_structure.find_lattice_site_by_position( 

103 position=position, fractional_position_tolerance=fractional_position_tolerance) 

104 

105 # Get allowed species on this site 

106 species = allowed_species[lattice_site.index] 

107 

108 # Get what sublattice those species correspond to 

109 sublattice = self._allowed_species.index(tuple(sorted(species))) 

110 

111 sublattice_to_indices[sublattice].append(index) 

112 

113 for symbol, species, indices in zip(symbols, self._allowed_species, sublattice_to_indices): 

114 sublattice = Sublattice(chemical_symbols=species, indices=indices, symbol=symbol) 

115 self._sublattices.append(sublattice) 

116 

117 # Map lattice index to sublattice index 

118 self._index_to_sublattice = {} 

119 for k, sublattice in enumerate(self): 

120 for index in sublattice.indices: 

121 self._index_to_sublattice[index] = k 

122 

123 def __getitem__(self, key: int) -> Sublattice: 

124 """ Returns a sublattice according to key. """ 

125 return self._sublattices[key] 

126 

127 def __len__(self): 

128 """ Returns number of sublattices. """ 

129 return len(self._sublattices) 

130 

131 def __iter__(self) -> Iterator[Sublattice]: 

132 """ Generator over sublattices. """ 

133 yield from self._sublattices 

134 

135 def get_sublattice_index_from_site_index(self, index: int) -> int: 

136 """ Returns the index of the sublattice the 

137 index in the structure belongs to. 

138 

139 Parameters 

140 ---------- 

141 index 

142 Index of site in the structure. 

143 """ 

144 return self._index_to_sublattice[index] 

145 

146 @property 

147 def allowed_species(self) -> List[List[str]]: 

148 """Lists of the allowed species on each sublattice in order. """ 

149 return deepcopy(self._allowed_species) 

150 

151 def get_sublattice_sites(self, index: int) -> List[int]: 

152 """Returns the sites that belong to the sublattice with the 

153 corresponding index. 

154 

155 Parameters 

156 ---------- 

157 index 

158 Index of the sublattice. 

159 """ 

160 return self[index].indices 

161 

162 def get_allowed_symbols_on_site(self, index: int) -> List[str]: 

163 """Returns the allowed symbols on the site. 

164 

165 Parameters 

166 ----------- 

167 index 

168 Lattice site index. 

169 """ 

170 return self[self._index_to_sublattice[index]].chemical_symbols 

171 

172 def get_allowed_numbers_on_site(self, index: int) -> List[int]: 

173 """Returns the allowed atomic numbers on the site. 

174 

175 Parameters 

176 ----------- 

177 index 

178 Lattice site index. 

179 """ 

180 return self[self._index_to_sublattice[index]].atomic_numbers 

181 

182 @property 

183 def active_sublattices(self) -> List[Sublattice]: 

184 """ List of the active sublattices. """ 

185 return [sl for sl in self if len(sl.chemical_symbols) > 1] 

186 

187 @property 

188 def inactive_sublattices(self) -> List[Sublattice]: 

189 """ List of the active sublattices. """ 

190 return [sl for sl in self if len(sl.chemical_symbols) == 1] 

191 

192 def assert_occupation_is_allowed(self, chemical_symbols: List[str]): 

193 """ Asserts that the current occupation obeys the sublattices. """ 

194 if len(chemical_symbols) != len(self._structure): 

195 raise ValueError(f'Length of input chemical symbols ({len(chemical_symbols)}) does not' 

196 f' match length of supercell ({len(self._structure)}') 

197 for sl in self: 

198 for i in sl.indices: 

199 if not chemical_symbols[i] in sl.chemical_symbols: 

200 msg = ('Occupations of structure not compatible with the sublattice.' 

201 ' Site {} with occupation {} not allowed on sublattice {}' 

202 .format(i, chemical_symbols[i], sl.chemical_symbols)) 

203 raise ValueError(msg)