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

1# fmt: off 

2 

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

4 

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

6this format and the documentation of muSTEM. 

7 

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

9""" 

10 

11import numpy as np 

12 

13from ase.atoms import Atoms, symbols2numbers 

14from ase.data import chemical_symbols 

15from ase.utils import reader, writer 

16 

17from .utils import verify_cell_for_export, verify_dictionary 

18 

19 

20@reader 

21def read_mustem(fd): 

22 r"""Import muSTEM input file. 

23 

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

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

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

27 

28 .. math:: 

29 

30 B = RMS * 8\pi^2 

31 

32 """ 

33 from ase.geometry import cellpar_to_cell 

34 

35 # Read comment: 

36 fd.readline() 

37 

38 # Parse unit cell parameter 

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

40 cell = cellpar_to_cell(cellpar) 

41 

42 # beam energy 

43 fd.readline() 

44 

45 # Number of different type of atoms 

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

47 

48 # List of numpy arrays: 

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

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

51 atomic_numbers = [] 

52 positions = [] 

53 debye_waller_factors = [] 

54 occupancies = [] 

55 

56 for _ in range(element_number): 

57 # Read the element 

58 _ = fd.readline() 

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

60 atoms_number = int(line[0]) 

61 atomic_number = int(line[1]) 

62 occupancy = float(line[2]) 

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

64 # read all the position for each element 

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

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

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

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

69 

70 positions = np.vstack(positions) 

71 

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

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

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

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

76 

77 return atoms 

78 

79 

80class XtlmuSTEMWriter: 

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

82 """ 

83 

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

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

86 verify_cell_for_export(atoms.get_cell()) 

87 

88 self.atoms = atoms.copy() 

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

90 self.keV = keV 

91 self.comment = comment 

92 self.occupancies = self._get_occupancies(occupancies) 

93 self.RMS = self._get_RMS(debye_waller_factors) 

94 

95 self.numbers = symbols2numbers(self.atom_types) 

96 if fit_cell_to_atoms: 

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

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

99 

100 def _get_occupancies(self, occupancies): 

101 if occupancies is None: 

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

103 occupancies = {element: 

104 self._parse_array_from_atoms( 

105 'occupancies', element, True) 

106 for element in self.atom_types} 

107 else: 

108 occupancies = 1.0 

109 if np.isscalar(occupancies): 

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

111 elif isinstance(occupancies, dict): 

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

113 

114 return occupancies 

115 

116 def _get_RMS(self, DW): 

117 if DW is None: 

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

119 DW = {element: self._parse_array_from_atoms( 

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

121 for element in self.atom_types} 

122 elif np.isscalar(DW): 

123 if len(self.atom_types) > 1: 

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

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

126 'be provided for each atom using a ' 

127 'dictionary.') 

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

129 elif isinstance(DW, dict): 

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

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

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

133 

134 if DW is None: 

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

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

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

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

139 'provided as float.') 

140 

141 return DW 

142 

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

144 """ 

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

146 

147 Parameters 

148 ---------- 

149 name : str 

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

151 element : str, int 

152 The element to be considered. 

153 check_same_value : bool 

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

155 'occupancies' and 'debye_waller_factors' arrays. 

156 

157 Returns 

158 ------- 

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

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

161 

162 """ 

163 if isinstance(element, str): 

164 element = symbols2numbers(element)[0] 

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

166 

167 if check_same_value: 

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

169 # all the value must be equal 

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

171 raise ValueError( 

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

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

174 ) 

175 sliced_array = sliced_array[0] 

176 

177 return sliced_array 

178 

179 def _get_position_array_single_atom_type(self, number): 

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

181 return self.atoms.get_scaled_positions()[ 

182 self.atoms.numbers == number] 

183 

184 def _get_file_header(self): 

185 # 1st line: comment line 

186 if self.comment is None: 

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

188 len(self.atoms), 

189 self.atoms.get_chemical_formula()) 

190 else: 

191 s = self.comment 

192 # 2nd line: lattice parameter 

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

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

195 # 3td line: acceleration voltage 

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

197 # 4th line: number of different atom 

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

199 return s 

200 

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

202 occupancy, RMS): 

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

204 number, 

205 atom_type_number, 

206 occupancy, 

207 RMS) 

208 

209 def _get_file_end(self): 

210 return "Orientation\n 1 0 0\n 0 1 0\n 0 0 1\n" 

211 

212 def write_to_file(self, fd): 

213 if isinstance(fd, str): 

214 fd = open(fd, 'w') 

215 

216 fd.write(self._get_file_header()) 

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

218 self.numbers, 

219 self.occupancies): 

220 positions = self._get_position_array_single_atom_type(number) 

221 atom_type_number = positions.shape[0] 

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

223 number, 

224 self.occupancies[atom_type], 

225 self.RMS[atom_type])) 

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

227 

228 fd.write(self._get_file_end()) 

229 

230 

231@writer 

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

233 r"""Write muSTEM input file. 

234 

235 Parameters: 

236 

237 atoms: Atoms object 

238 

239 keV: float 

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

241 

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

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

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

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

246 to RMS displacement by: 

247 

248 .. math:: 

249 

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

251 

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

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

254 

255 comment: str (optional) 

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

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

258 

259 fit_cell_to_atoms: bool (optional) 

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

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

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

263 the cell are unchanged. 

264 """ 

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

266 writer.write_to_file(fd)