Coverage for hiphive/core/atoms.py: 90%

1""" 

2Collection of functions and classes for handling information concerning atoms 

3and structures, including the relationship between primitive cell and 

4supercells that are derived thereof. 

5""" 

6 

7import pickle 

8from ase import Atoms as aseAtoms 

9import numpy as np 

10from ..input_output.logging_tools import logger 

11 

12# TODO: Rename logger 

13logger = logger.getChild('atoms') 

14 

15 

16class Atom: 

17 # TODO: This class should inherit some immutable to make it clear that 

18 # there is no reference to any other obj 

19 """Unique representation of an atom in a lattice with a basis 

20 

21 Class for storing information about the position of an atom in a supercell 

22 relative to the origin of the underlying primitive cell. This class is used 

23 for handling the relationship between a primitive cell and supercells 

24 derived thereof. 

25 

26 Parameters 

27 ---------- 

28 site : int 

29 site index 

30 offset : list(float) or numpy.ndarray 

31 must contain three elements, offset_x, offset_y, offset_z 

32 """ 

33 def __init__(self, site, offset): 

34 offset = tuple(offset) 

35 self._site = site 

36 self._offset = offset 

37 

38 @property 

39 def site(self): 

40 """int : index of corresponding site in the primitive basis""" 

41 return self._site 

42 

43 @property 

44 def offset(self): 

45 """list(int) : translational offset of the supercell site relative 

46 to the origin of the primitive cell in units of primitive lattice 

47 vectors""" 

48 return self._offset 

49 

50 def __repr__(self): 

51 return 'Atom({}, {})'.format(self.site, self.offset) 

52 

53 def spos(atom, basis): 

54 return np.add(basis[atom.site], atom.offset) 

55 

56 def pos(atom, basis, cell): 

57 spos = atom.spos(basis) 

58 return np.dot(spos, cell) 

59 

60 @staticmethod 

61 def spos_to_atom(spos, basis, tol=None): 

62 # TODO: Why is this simply duplicated spos_to_atom from helper below? 

63 if not tol: 63 ↛ 65line 63 didn't jump to line 65, because the condition on line 63 was never true

64 # TODO: Link to config file 

65 tol = 1e-4 

66 for site, base in enumerate(basis): 66 ↛ 76line 66 didn't jump to line 76, because the loop on line 66 didn't complete

67 offset = np.subtract(spos, base) 

68 diff = offset - np.round(offset, 0) 

69 if np.linalg.norm(diff) < tol: 

70 offset = np.round(offset, 0).astype(int) 

71 atom = Atom(site, offset) 

72 assert np.linalg.norm(spos - atom.spos(basis)) < tol, ( 

73 '{} with basis {} != {}'.format(atom, basis, spos)) 

74 return atom 

75 

76 s = '{} not compatible with {} and tolerance {}' 

77 raise Exception(s.format(spos, basis, tol)) 

78 

79 def __hash__(self): 

80 return hash((self._site, *self.offset)) 

81 

82 def __eq__(self, other): 

83 if not isinstance(other, Atom): 83 ↛ 84line 83 didn't jump to line 84, because the condition on line 83 was never true

84 return False 

85 return self.site == other.site and self.offset == other.offset 

86 

87 

88class Atoms(aseAtoms): 

89 """Minimally augmented version of the ASE Atoms class suitable for handling 

90 primitive cell information. 

91 

92 Saves and loads by pickle. 

93 """ 

94 @property 

95 def basis(self): 

96 """numpy.ndarray : scaled coordinates of the sites in the primitive basis 

97 """ 

98 return self.get_scaled_positions().round(12) % 1 

99 

100 def write(self, f): 

101 """ Writes the object to file. 

102 

103 Note: Only the cell, basis and numbers are stored! 

104 

105 Parameters 

106 ---------- 

107 f : str or file object 

108 name of input file (str) or stream to write to (file object) 

109 """ 

110 data = {} 

111 data['cell'] = self.cell 

112 data['basis'] = self.basis 

113 data['numbers'] = self.numbers 

114 

115 pickle.dump(data, f) 

116 

117 @staticmethod 

118 def read(f): 

119 """ Load an hiPhive Atoms object from file. 

120 

121 Parameters 

122 ---------- 

123 f : str or file object 

124 name of input file (str) or stream to load from (file object) 

125 

126 Returns 

127 ------- 

128 hiPhive Atoms object 

129 """ 

130 data = pickle.load(f) 

131 atoms = aseAtoms(numbers=data['numbers'], 

132 scaled_positions=data['basis'], 

133 cell=data['cell'], 

134 pbc=True) 

135 return Atoms(atoms) 

136 

137 

138def atom_to_spos(atom, basis): 

139 """Helper function for obtaining the position of a supercell atom in scaled 

140 coordinates. 

141 

142 Parameters 

143 ---------- 

144 atom : hiPhive.Atom 

145 supercell atom 

146 basis : list(list(float)) or numpy.ndarray 

147 positions of sites in the primitive basis 

148 

149 Returns 

150 ------- 

151 numpy.ndarray 

152 scaled coordinates of an atom in a supercell 

153 """ 

154 return np.add(atom.offset, basis[atom.site]) 

155 

156 

157def spos_to_atom(spos, basis, tol=1e-4): 

158 """Helper function for transforming a supercell position to the primitive 

159 basis. 

160 

161 Parameters 

162 ---------- 

163 spos : list(list(float)) or numpy.ndarray 

164 scaled coordinates of an atom in a supercell 

165 basis : list(list(float)) or numpy.ndarray 

166 positions of sites in the primitive basis 

167 tol : float 

168 a general tolerance 

169 

170 Returns 

171 ------- 

172 hiphive.Atom 

173 supercell atom 

174 """ 

175 # TODO: Fix tolerance 

176 # If needed, convert inputs to arrays to make use of numpy vectorization 

177 spos = np.asarray(spos) 

178 basis = np.asarray(basis) 

179 # If the scaled position belongs to this site, the offset is the 

180 # difference in scaled coordinates and should be integer 

181 offsets = spos - basis 

182 # The diff is the difference between the offset vector and the nearest 

183 # integer vector. 

184 diffs = offsets - np.round(offsets, 0) 

185 # It should be close to the null vector if this is the correct site. 

186 match_indices = np.nonzero(np.linalg.norm(diffs, axis=1) < tol)[0] 

187 # If no atom was found or more than one atoms were found we throw an error 

188 if len(match_indices) != 1: 188 ↛ 189line 188 didn't jump to line 189, because the condition on line 188 was never true

189 raise ValueError(f'{spos} not compatible with {basis} and tolerance {tol}') 

190 

191 # This should be the correct atom 

192 site = match_indices[0] 

193 # If the difference is less than the tol make the offset integers 

194 offset = np.rint(offsets[site]) 

195 atom = Atom(site, offset) 

196 # Just to be sure we check that the atom actually produces the 

197 # input spos given the input basis 

198 s = ('Atom=[{},{}] with basis {} != {}' 

199 .format(atom.site, atom.offset, basis, spos)) 

200 assert np.linalg.norm(spos - atom_to_spos(atom, basis)) < tol, s 

201 return atom