Coverage for /builds/kinetik161/ase/ase/io/mustem.py: 95.41%

1"""Module to read and write atoms in xtl file format for the muSTEM software. 

2 

3See http://tcmp.ph.unimelb.edu.au/mustem/muSTEM.html for a few examples of 

4this format and the documentation of muSTEM. 

5 

6See https://github.com/HamishGBrown/MuSTEM for the source code of muSTEM. 

7""" 

8 

9import numpy as np 

10 

11from ase.atoms import Atoms, symbols2numbers 

12from ase.data import chemical_symbols 

13from ase.utils import reader, writer 

14 

15from .utils import verify_cell_for_export, verify_dictionary 

16 

17 

18@reader 

19def read_mustem(fd): 

20 r"""Import muSTEM input file. 

21 

22 Reads cell, atom positions, etc. from muSTEM xtl file. 

23 The mustem xtl save the root mean square (RMS) displacement, which is 

24 convert to Debye-Waller (in Ų) factor by: 

25 

26 .. math:: 

27 

28 B = RMS * 8\pi^2 

29 

30 """ 

31 from ase.geometry import cellpar_to_cell 

32 

33 # Read comment: 

34 fd.readline() 

35 

36 # Parse unit cell parameter 

37 cellpar = [float(i) for i in fd.readline().split()[:3]] 

38 cell = cellpar_to_cell(cellpar) 

39 

40 # beam energy 

41 fd.readline() 

42 

43 # Number of different type of atoms 

44 element_number = int(fd.readline().strip()) 

45 

46 # List of numpy arrays: 

47 # length of the list = number of different atom type (element_number) 

48 # length of each array = number of atoms for each atom type 

49 atomic_numbers = [] 

50 positions = [] 

51 debye_waller_factors = [] 

52 occupancies = [] 

53 

54 for i in range(element_number): 

55 # Read the element 

56 _ = fd.readline() 

57 line = fd.readline().split() 

58 atoms_number = int(line[0]) 

59 atomic_number = int(line[1]) 

60 occupancy = float(line[2]) 

61 DW = float(line[3]) * 8 * np.pi**2 

62 # read all the position for each element 

63 positions.append(np.genfromtxt(fname=fd, max_rows=atoms_number)) 

64 atomic_numbers.append(np.ones(atoms_number, dtype=int) * atomic_number) 

65 occupancies.append(np.ones(atoms_number) * occupancy) 

66 debye_waller_factors.append(np.ones(atoms_number) * DW) 

67 

68 positions = np.vstack(positions) 

69 

70 atoms = Atoms(cell=cell, scaled_positions=positions) 

71 atoms.set_atomic_numbers(np.hstack(atomic_numbers)) 

72 atoms.set_array('occupancies', np.hstack(occupancies)) 

73 atoms.set_array('debye_waller_factors', np.hstack(debye_waller_factors)) 

74 

75 return atoms 

76 

77 

78class XtlmuSTEMWriter: 

79 """See the docstring of the `write_mustem` function. 

80 """ 

81 

82 def __init__(self, atoms, keV, debye_waller_factors=None, 

83 comment=None, occupancies=None, fit_cell_to_atoms=False): 

84 verify_cell_for_export(atoms.get_cell()) 

85 

86 self.atoms = atoms.copy() 

87 self.atom_types = sorted(set(atoms.symbols)) 

88 self.keV = keV 

89 self.comment = comment 

90 self.occupancies = self._get_occupancies(occupancies) 

91 self.RMS = self._get_RMS(debye_waller_factors) 

92 

93 self.numbers = symbols2numbers(self.atom_types) 

94 if fit_cell_to_atoms: 

95 self.atoms.translate(-self.atoms.positions.min(axis=0)) 

96 self.atoms.set_cell(self.atoms.positions.max(axis=0)) 

97 

98 def _get_occupancies(self, occupancies): 

99 if occupancies is None: 

100 if 'occupancies' in self.atoms.arrays: 

101 occupancies = {element: 

102 self._parse_array_from_atoms( 

103 'occupancies', element, True) 

104 for element in self.atom_types} 

105 else: 

106 occupancies = 1.0 

107 if np.isscalar(occupancies): 

108 occupancies = {atom: occupancies for atom in self.atom_types} 

109 elif isinstance(occupancies, dict): 

110 verify_dictionary(self.atoms, occupancies, 'occupancies') 

111 

112 return occupancies 

113 

114 def _get_RMS(self, DW): 

115 if DW is None: 

116 if 'debye_waller_factors' in self.atoms.arrays: 

117 DW = {element: self._parse_array_from_atoms( 

118 'debye_waller_factors', element, True) / (8 * np.pi**2) 

119 for element in self.atom_types} 

120 elif np.isscalar(DW): 

121 if len(self.atom_types) > 1: 

122 raise ValueError('This cell contains more then one type of ' 

123 'atoms and the Debye-Waller factor needs to ' 

124 'be provided for each atom using a ' 

125 'dictionary.') 

126 DW = {self.atom_types[0]: DW / (8 * np.pi**2)} 

127 elif isinstance(DW, dict): 

128 verify_dictionary(self.atoms, DW, 'debye_waller_factors') 

129 for key, value in DW.items(): 

130 DW[key] = value / (8 * np.pi**2) 

131 

132 if DW is None: 

133 raise ValueError('Missing Debye-Waller factors. It can be ' 

134 'provided as a dictionary with symbols as key or ' 

135 'if the cell contains only a single type of ' 

136 'element, the Debye-Waller factor can also be ' 

137 'provided as float.') 

138 

139 return DW 

140 

141 def _parse_array_from_atoms(self, name, element, check_same_value): 

142 """ 

143 Return the array "name" for the given element. 

144 

145 Parameters 

146 ---------- 

147 name : str 

148 The name of the arrays. Can be any key of `atoms.arrays` 

149 element : str, int 

150 The element to be considered. 

151 check_same_value : bool 

152 Check if all values are the same in the array. Necessary for 

153 'occupancies' and 'debye_waller_factors' arrays. 

154 

155 Returns 

156 ------- 

157 array containing the values corresponding defined by "name" for the 

158 given element. If check_same_value, return a single element. 

159 

160 """ 

161 if isinstance(element, str): 

162 element = symbols2numbers(element)[0] 

163 sliced_array = self.atoms.arrays[name][self.atoms.numbers == element] 

164 

165 if check_same_value: 

166 # to write the occupancies and the Debye Waller factor of xtl file 

167 # all the value must be equal 

168 if np.unique(sliced_array).size > 1: 

169 raise ValueError( 

170 "All the '{}' values for element '{}' must be " 

171 "equal.".format(name, chemical_symbols[element]) 

172 ) 

173 sliced_array = sliced_array[0] 

174 

175 return sliced_array 

176 

177 def _get_position_array_single_atom_type(self, number): 

178 # Get the scaled (reduced) position for a single atom type 

179 return self.atoms.get_scaled_positions()[ 

180 self.atoms.numbers == number] 

181 

182 def _get_file_header(self): 

183 # 1st line: comment line 

184 if self.comment is None: 

185 s = "{} atoms with chemical formula: {}\n".format( 

186 len(self.atoms), 

187 self.atoms.get_chemical_formula()) 

188 else: 

189 s = self.comment 

190 # 2nd line: lattice parameter 

191 s += "{} {} {} {} {} {}\n".format( 

192 *self.atoms.cell.cellpar().tolist()) 

193 # 3td line: acceleration voltage 

194 s += f"{self.keV}\n" 

195 # 4th line: number of different atom 

196 s += f"{len(self.atom_types)}\n" 

197 return s 

198 

199 def _get_element_header(self, atom_type, number, atom_type_number, 

200 occupancy, RMS): 

201 return "{}\n{} {} {} {:.3g}\n".format(atom_type, 

202 number, 

203 atom_type_number, 

204 occupancy, 

205 RMS) 

206 

207 def _get_file_end(self): 

208 return "Orientation\n 1 0 0\n 0 1 0\n 0 0 1\n" 

209 

210 def write_to_file(self, fd): 

211 if isinstance(fd, str): 

212 fd = open(fd, 'w') 

213 

214 fd.write(self._get_file_header()) 

215 for atom_type, number, occupancy in zip(self.atom_types, 

216 self.numbers, 

217 self.occupancies): 

218 positions = self._get_position_array_single_atom_type(number) 

219 atom_type_number = positions.shape[0] 

220 fd.write(self._get_element_header(atom_type, atom_type_number, 

221 number, 

222 self.occupancies[atom_type], 

223 self.RMS[atom_type])) 

224 np.savetxt(fname=fd, X=positions, fmt='%.6g', newline='\n') 

225 

226 fd.write(self._get_file_end()) 

227 

228 

229@writer 

230def write_mustem(fd, *args, **kwargs): 

231 r"""Write muSTEM input file. 

232 

233 Parameters: 

234 

235 atoms: Atoms object 

236 

237 keV: float 

238 Energy of the electron beam in keV required for the image simulation. 

239 

240 debye_waller_factors: float or dictionary of float with atom type as key 

241 Debye-Waller factor of each atoms. Since the prismatic/computem 

242 software use root means square RMS) displacements, the Debye-Waller 

243 factors (B) needs to be provided in Ų and these values are converted 

244 to RMS displacement by: 

245 

246 .. math:: 

247 

248 RMS = \frac{B}{8\pi^2} 

249 

250 occupancies: float or dictionary of float with atom type as key (optional) 

251 Occupancy of each atoms. Default value is `1.0`. 

252 

253 comment: str (optional) 

254 Comments to be written in the first line of the file. If not 

255 provided, write the total number of atoms and the chemical formula. 

256 

257 fit_cell_to_atoms: bool (optional) 

258 If `True`, fit the cell to the atoms positions. If negative coordinates 

259 are present in the cell, the atoms are translated, so that all 

260 positions are positive. If `False` (default), the atoms positions and 

261 the cell are unchanged. 

262 """ 

263 writer = XtlmuSTEMWriter(*args, **kwargs) 

264 writer.write_to_file(fd)