Coverage for hiphive/core/structures.py: 94%

1import numpy as np 

2 

3 

4def site_offset_to_spos(site, offset, basis_spos): 

5 """ Returns the scaled position of an atom at specified site and offset 

6 relative to the basis in scaled coordinates""" 

7 return offset + basis_spos[site] 

8 

9 

10def spos_to_pos(spos, cell): 

11 """ Returns the Cartesian coordinate given the scaled coordinate and cell 

12 metric (cell vectors as rows)""" 

13 return np.dot(spos, cell) 

14 

15 

16def pos_to_spos(pos, cell): 

17 """ Inverse of sps_to_pos""" 

18 return np.linalg.solve(cell.T, pos) 

19 

20 

21def spos_to_site_offset(spos, basis_spos, symprec): 

22 """ Returns the site and offset of the atom at the specified scaled 

23 coordinate given the scaled positions of the basis atoms""" 

24 for site, sp in enumerate(basis_spos): 24 ↛ 30line 24 didn't jump to line 30 because the loop on line 24 didn't complete

25 offset = spos - sp 

26 rounded_offset = offset.round(0).astype(np.int64) 

27 # TODO: fix tolerance (symprec should be w.r.t. cart. coord.) 

28 if np.allclose(rounded_offset, offset, rtol=0, atol=symprec): 

29 return site, rounded_offset 

30 raise Exception('spos {} not compatible with basis {} using symprec {}' 

31 .format(spos, basis_spos, symprec)) 

32 

33 

34def pos_to_site_offset(pos, cell, basis_spos, symprec): 

35 """ helper to map pos -> spos -> site/offset""" 

36 spos = pos_to_spos(pos, cell) 

37 return spos_to_site_offset(spos, basis_spos, symprec) 

38 

39 

40def site_offset_to_pos(site, offset, cell, basis_spos): 

41 """ helper to map site/offset -> spos -> pos""" 

42 spos = site_offset_to_spos(site, offset, basis_spos) 

43 return spos_to_pos(spos, cell) 

44 

45 

46class BaseAtom: 

47 """ This class represents an atom placed in an infinite crustal""" 

48 def __init__(self, site, offset): 

49 assert type(site) is int, type(site) 

50 assert len(offset) == 3, len(offset) 

51 assert (all(type(i) is int for i in offset) or 

52 all(type(i) is np.int64 for i in offset)), type(offset[0]) 

53 self._site = site 

54 self._offset = np.array(offset) 

55 

56 @property 

57 def site(self): 

58 return self._site 

59 

60 @property 

61 def offset(self): 

62 return self._offset 

63 

64 def astype(self, dtype): 

65 """ Useful arguments: list, tuple, np.int64""" 

66 return dtype((self._site, *self._offset)) 

67 

68 

69class Atom(BaseAtom): 

70 """ This class represents a crystal atom in a given structure""" 

71 def __init__(self, *args, **kwargs): 

72 self._structure = kwargs.pop('structure', None) 

73 super().__init__(*args, **kwargs) 

74 

75 @property 

76 def pos(self): 

77 return site_offset_to_pos(self._site, self._offset, 

78 self._structure.cell, 

79 self._structure.spos) 

80 

81 @property 

82 def number(self): 

83 return self._structure.numbers[self._site] 

84 

85 

86class SupercellAtom(Atom): 

87 """ Represents an atom in a supercell but site and offset given by an 

88 underlying primitve cell""" 

89 def __init__(self, *args, **kwargs): 

90 self._index = kwargs.pop('index') 

91 assert type(self._index) is int 

92 super().__init__(*args, **kwargs) 

93 

94 @property 

95 def index(self): 

96 return self._index 

97 

98 

99class Structure: 

100 """ This class essentially wraps the ase.Atoms class but is a bit more 

101 carefull about pbc and scaled coordinates. It also returns hiphive.Atom 

102 objects instead""" 

103 def __init__(self, atoms, symprec=1e-6): 

104 spos = atoms.get_scaled_positions(wrap=False) 

105 for sp in spos.flat: 

106 if not (-symprec < sp < (1 - symprec)): 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 raise ValueError('bad spos {}'.format(sp)) 

108 self._spos = spos 

109 self._cell = atoms.cell 

110 self._numbers = atoms.numbers 

111 

112 def __len__(self): 

113 return len(self._spos) 

114 

115 @property 

116 def spos(self): 

117 return self._spos 

118 

119 @property 

120 def cell(self): 

121 return self._cell 

122 

123 def __getitem__(self, index): 

124 if index >= len(self): 

125 raise IndexError('Structure contains {} atoms'.format(len(self))) 

126 return Atom(index, (0, 0, 0), structure=self) 

127 

128 def atom_from_pos(self, pos, symprec=None): 

129 if symprec is None: 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true

130 symprec = self._symprec 

131 site, offset = pos_to_site_offset(pos, self._cell, self._spos, symprec) 

132 return Atom(site, offset, structure=self) 

133 

134 

135class Supercell: 

136 """ This class tries to represent atoms in a supercell as positioned on the 

137 primitve lattice""" 

138 def __init__(self, supercell, prim, symprec): 

139 self._supercell = Structure(supercell) 

140 self._prim = Structure(prim) 

141 self._symprec = symprec 

142 self._map = list() 

143 self._inverse_map_lookup = dict() 

144 self._create_map() 

145 

146 def _create_map(self): 

147 for atom in self._supercell: 

148 atom = self._prim.atom_from_pos(atom.pos, self._symprec) 

149 self._map.append(atom.astype(tuple)) 

150 

151 def wrap_atom(self, atom): 

152 atom = Atom(atom.site, atom.offset, structure=self._prim) 

153 tup = atom.astype(tuple) 

154 index = self._inverse_map_lookup.get(tup, None) 

155 if index is None: 

156 atom = self._supercell.atom_from_pos(atom.pos, self._symprec) 

157 index = atom.site 

158 self._inverse_map_lookup[tup] = index 

159 return self[index] 

160 

161 def index(self, site, offset): 

162 atom = self.wrap_atom(BaseAtom(site, offset)) 

163 return atom.index 

164 

165 def __getitem__(self, index): 

166 tup = self._map[index] 

167 return SupercellAtom(tup[0], tup[1:], structure=self._prim, 

168 index=index) 

169 

170 def __len__(self): 

171 return len(self._supercell)