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 

14from .utils import verify_cell_for_export, verify_dictionary 

15 

16 

17@reader 

18def read_mustem(fd): 

19 r"""Import muSTEM input file. 

20 

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

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

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

24 

25 .. math:: 

26 

27 B = RMS * 8\pi^2 

28 

29 """ 

30 from ase.geometry import cellpar_to_cell 

31 

32 # Read comment: 

33 fd.readline() 

34 

35 # Parse unit cell parameter 

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

37 cell = cellpar_to_cell(cellpar) 

38 

39 # beam energy 

40 fd.readline() 

41 

42 # Number of different type of atoms 

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

44 

45 # List of numpy arrays: 

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

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

48 atomic_numbers = [] 

49 positions = [] 

50 debye_waller_factors = [] 

51 occupancies = [] 

52 

53 for i in range(element_number): 

54 # Read the element 

55 _ = fd.readline() 

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

57 atoms_number = int(line[0]) 

58 atomic_number = int(line[1]) 

59 occupancy = float(line[2]) 

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

61 # read all the position for each element 

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

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

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

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

66 

67 positions = np.vstack(positions) 

68 

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

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

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

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

73 

74 return atoms 

75 

76 

77class XtlmuSTEMWriter: 

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

79 """ 

80 

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

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

83 verify_cell_for_export(atoms.get_cell()) 

84 

85 self.atoms = atoms.copy() 

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

87 self.keV = keV 

88 self.comment = comment 

89 self.occupancies = self._get_occupancies(occupancies) 

90 self.RMS = self._get_RMS(debye_waller_factors) 

91 

92 self.numbers = symbols2numbers(self.atom_types) 

93 if fit_cell_to_atoms: 

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

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

96 

97 def _get_occupancies(self, occupancies): 

98 if occupancies is None: 

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

100 occupancies = {element: 

101 self._parse_array_from_atoms( 

102 'occupancies', element, True) 

103 for element in self.atom_types} 

104 else: 

105 occupancies = 1.0 

106 if np.isscalar(occupancies): 

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

108 elif isinstance(occupancies, dict): 

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

110 

111 return occupancies 

112 

113 def _get_RMS(self, DW): 

114 if DW is None: 

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

116 DW = {element: self._parse_array_from_atoms( 

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

118 for element in self.atom_types} 

119 elif np.isscalar(DW): 

120 if len(self.atom_types) > 1: 

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

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

123 'be provided for each atom using a ' 

124 'dictionary.') 

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

126 elif isinstance(DW, dict): 

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

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

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

130 

131 if DW is None: 

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

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

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

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

136 'provided as float.') 

137 

138 return DW 

139 

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

141 """ 

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

143 

144 Parameters 

145 ---------- 

146 name : str 

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

148 element : str, int 

149 The element to be considered. 

150 check_same_value : bool 

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

152 'occupancies' and 'debye_waller_factors' arrays. 

153 

154 Returns 

155 ------- 

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

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

158 

159 """ 

160 if isinstance(element, str): 

161 element = symbols2numbers(element)[0] 

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

163 

164 if check_same_value: 

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

166 # all the value must be equal 

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

168 raise ValueError( 

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

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

171 ) 

172 sliced_array = sliced_array[0] 

173 

174 return sliced_array 

175 

176 def _get_position_array_single_atom_type(self, number): 

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

178 return self.atoms.get_scaled_positions()[ 

179 self.atoms.numbers == number] 

180 

181 def _get_file_header(self): 

182 # 1st line: comment line 

183 if self.comment is None: 

184 s = "{0} atoms with chemical formula: {1}\n".format( 

185 len(self.atoms), 

186 self.atoms.get_chemical_formula()) 

187 else: 

188 s = self.comment 

189 # 2nd line: lattice parameter 

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

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

192 # 3td line: acceleration voltage 

193 s += "{}\n".format(self.keV) 

194 # 4th line: number of different atom 

195 s += "{}\n".format(len(self.atom_types)) 

196 return s 

197 

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

199 occupancy, RMS): 

200 return "{0}\n{1} {2} {3} {4:.3g}\n".format(atom_type, 

201 number, 

202 atom_type_number, 

203 occupancy, 

204 RMS) 

205 

206 def _get_file_end(self): 

207 return "Orientation\n 1 0 0\n 0 1 0\n 0 0 1\n" 

208 

209 def write_to_file(self, fd): 

210 if isinstance(fd, str): 

211 fd = open(fd, 'w') 

212 

213 fd.write(self._get_file_header()) 

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

215 self.numbers, 

216 self.occupancies): 

217 positions = self._get_position_array_single_atom_type(number) 

218 atom_type_number = positions.shape[0] 

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

220 number, 

221 self.occupancies[atom_type], 

222 self.RMS[atom_type])) 

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

224 

225 fd.write(self._get_file_end()) 

226 

227 

228@writer 

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

230 r"""Write muSTEM input file. 

231 

232 Parameters: 

233 

234 atoms: Atoms object 

235 

236 keV: float 

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

238 

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

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

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

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

243 to RMS displacement by: 

244 

245 .. math:: 

246 

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

248 

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

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

251 

252 comment: str (optional) 

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

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

255 

256 fit_cell_to_atoms: bool (optional) 

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

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

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

260 the cell are unchanged. 

261 """ 

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

263 writer.write_to_file(fd)