# Coverage for hiphive/core/structures.py: 95%

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

## 108 statements

1import numpy as np

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]

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)

16def pos_to_spos(pos, cell):

17 """ Inverse of sps_to_pos"""

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

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))

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)

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)

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)

56 @property

57 def site(self):

58 return self._site

60 @property

61 def offset(self):

62 return self._offset

64 def astype(self, dtype):

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

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

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)

75 @property

76 def pos(self):

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

78 self._structure.cell,

79 self._structure.spos)

81 @property

82 def number(self):

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

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)

94 @property

95 def index(self):

96 return self._index

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

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

112 def __len__(self):

113 return len(self._spos)

115 @property

116 def spos(self):

117 return self._spos

119 @property

120 def cell(self):

121 return self._cell

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)

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)

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()

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))

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]

161 def index(self, site, offset):

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

163 return atom.index

165 def __getitem__(self, index):

166 tup = self._map[index]

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

168 index=index)

170 def __len__(self):

171 return len(self._supercell)