1""" 

2SHELX (.res) input/output 

3 

4Read/write files in SHELX (.res) file format. 

5 

6Format documented at http://shelx.uni-ac.gwdg.de/SHELX/ 

7 

8Written by Martin Uhren and Georg Schusteritsch. 

9Adapted for ASE by James Kermode. 

10""" 

11 

12 

13import glob 

14import re 

15 

16from ase.atoms import Atoms 

17from ase.geometry import cellpar_to_cell, cell_to_cellpar 

18from ase.calculators.calculator import Calculator 

19from ase.calculators.singlepoint import SinglePointCalculator 

20 

21__all__ = ['Res', 'read_res', 'write_res'] 

22 

23 

24class Res: 

25 

26 """ 

27 Object for representing the data in a Res file. 

28 Most attributes can be set directly. 

29 

30 Args: 

31 atoms (Atoms): Atoms object. 

32 

33 .. attribute:: atoms 

34 

35 Associated Atoms object. 

36 

37 .. attribute:: name 

38 

39 The name of the structure. 

40 

41 .. attribute:: pressure 

42 

43 The external pressure. 

44 

45 .. attribute:: energy 

46 

47 The internal energy of the structure. 

48 

49 .. attribute:: spacegroup 

50 

51 The space group of the structure. 

52 

53 .. attribute:: times_found 

54 

55 The number of times the structure was found. 

56 """ 

57 

58 def __init__(self, atoms, name=None, pressure=None, 

59 energy=None, spacegroup=None, times_found=None): 

60 self.atoms_ = atoms 

61 if name is None: 

62 name = atoms.info.get('name') 

63 if pressure is None: 

64 pressure = atoms.info.get('pressure') 

65 if spacegroup is None: 

66 spacegroup = atoms.info.get('spacegroup') 

67 if times_found is None: 

68 times_found = atoms.info.get('times_found') 

69 self.name = name 

70 self.pressure = pressure 

71 self.energy = energy 

72 self.spacegroup = spacegroup 

73 self.times_found = times_found 

74 

75 @property 

76 def atoms(self): 

77 """ 

78 Returns Atoms object associated with this Res. 

79 """ 

80 return self.atoms_ 

81 

82 @staticmethod 

83 def from_file(filename): 

84 """ 

85 Reads a Res from a file. 

86 

87 Args: 

88 filename (str): File name containing Res data. 

89 

90 Returns: 

91 Res object. 

92 """ 

93 with open(filename, 'r') as fd: 

94 return Res.from_string(fd.read()) 

95 

96 @staticmethod 

97 def parse_title(line): 

98 info = dict() 

99 

100 tokens = line.split() 

101 num_tokens = len(tokens) 

102 # 1 = Name 

103 if num_tokens <= 1: 

104 return info 

105 info['name'] = tokens[1] 

106 # 2 = Pressure 

107 if num_tokens <= 2: 

108 return info 

109 info['pressure'] = float(tokens[2]) 

110 # 3 = Volume 

111 # 4 = Internal energy 

112 if num_tokens <= 4: 

113 return info 

114 info['energy'] = float(tokens[4]) 

115 # 5 = Spin density, 6 - Abs spin density 

116 # 7 = Space group OR num atoms (new format ONLY) 

117 idx = 7 

118 if tokens[idx][0] != '(': 

119 idx += 1 

120 

121 if num_tokens <= idx: 

122 return info 

123 info['spacegroup'] = tokens[idx][1:len(tokens[idx]) - 1] 

124 # idx + 1 = n, idx + 2 = - 

125 # idx + 3 = times found 

126 if num_tokens <= idx + 3: 

127 return info 

128 info['times_found'] = int(tokens[idx + 3]) 

129 

130 return info 

131 

132 @staticmethod 

133 def from_string(data): 

134 """ 

135 Reads a Res from a string. 

136 

137 Args: 

138 data (str): string containing Res data. 

139 

140 Returns: 

141 Res object. 

142 """ 

143 abc = [] 

144 ang = [] 

145 sp = [] 

146 coords = [] 

147 info = dict() 

148 coord_patt = re.compile(r"""(\w+)\s+ 

149 ([0-9]+)\s+ 

150 ([0-9\-\.]+)\s+ 

151 ([0-9\-\.]+)\s+ 

152 ([0-9\-\.]+)\s+ 

153 ([0-9\-\.]+)""", re.VERBOSE) 

154 lines = data.splitlines() 

155 line_no = 0 

156 while line_no < len(lines): 

157 line = lines[line_no] 

158 tokens = line.split() 

159 if tokens: 

160 if tokens[0] == 'TITL': 

161 try: 

162 info = Res.parse_title(line) 

163 except (ValueError, IndexError): 

164 info = dict() 

165 elif tokens[0] == 'CELL' and len(tokens) == 8: 

166 abc = [float(tok) for tok in tokens[2:5]] 

167 ang = [float(tok) for tok in tokens[5:8]] 

168 elif tokens[0] == 'SFAC': 

169 for atom_line in lines[line_no:]: 

170 if line.strip() == 'END': 

171 break 

172 else: 

173 match = coord_patt.search(atom_line) 

174 if match: 

175 sp.append(match.group(1)) # 1-indexed 

176 cs = match.groups()[2:5] 

177 coords.append([float(c) for c in cs]) 

178 line_no += 1 # Make sure the global is updated 

179 line_no += 1 

180 

181 return Res(Atoms(symbols=sp, 

182 scaled_positions=coords, 

183 cell=cellpar_to_cell(list(abc) + list(ang)), 

184 pbc=True, info=info), 

185 info.get('name'), 

186 info.get('pressure'), 

187 info.get('energy'), 

188 info.get('spacegroup'), 

189 info.get('times_found')) 

190 

191 def get_string(self, significant_figures=6, write_info=False): 

192 """ 

193 Returns a string to be written as a Res file. 

194 

195 Args: 

196 significant_figures (int): No. of significant figures to 

197 output all quantities. Defaults to 6. 

198 

199 write_info (bool): if True, format TITL line using key-value pairs 

200 from atoms.info in addition to attributes stored in Res object 

201 

202 Returns: 

203 String representation of Res. 

204 """ 

205 

206 # Title line 

207 if write_info: 

208 info = self.atoms.info.copy() 

209 for attribute in ['name', 'pressure', 'energy', 

210 'spacegroup', 'times_found']: 

211 if getattr(self, attribute) and attribute not in info: 

212 info[attribute] = getattr(self, attribute) 

213 lines = ['TITL ' + ' '.join(['{0}={1}'.format(k, v) 

214 for (k, v) in info.items()])] 

215 else: 

216 lines = ['TITL ' + self.print_title()] 

217 

218 # Cell 

219 abc_ang = cell_to_cellpar(self.atoms.get_cell()) 

220 fmt = '{{0:.{0}f}}'.format(significant_figures) 

221 cell = ' '.join([fmt.format(a) for a in abc_ang]) 

222 lines.append('CELL 1.0 ' + cell) 

223 

224 # Latt 

225 lines.append('LATT -1') 

226 

227 # Atoms 

228 symbols = self.atoms.get_chemical_symbols() 

229 species_types = [] 

230 for symbol in symbols: 

231 if symbol not in species_types: 

232 species_types.append(symbol) 

233 lines.append('SFAC ' + ' '.join(species_types)) 

234 

235 fmt = '{{0}} {{1}} {{2:.{0}f}} {{3:.{0}f}} {{4:.{0}f}} 1.0' 

236 fmtstr = fmt.format(significant_figures) 

237 for symbol, coords in zip(symbols, 

238 self.atoms_.get_scaled_positions()): 

239 lines.append( 

240 fmtstr.format(symbol, 

241 species_types.index(symbol) + 1, 

242 coords[0], 

243 coords[1], 

244 coords[2])) 

245 lines.append('END') 

246 return '\n'.join(lines) 

247 

248 def __str__(self): 

249 """ 

250 String representation of Res file. 

251 """ 

252 return self.get_string() 

253 

254 def write_file(self, filename, **kwargs): 

255 """ 

256 Writes Res to a file. The supported kwargs are the same as those for 

257 the Res.get_string method and are passed through directly. 

258 """ 

259 with open(filename, 'w') as fd: 

260 fd.write(self.get_string(**kwargs) + '\n') 

261 

262 def print_title(self): 

263 tokens = [self.name, self.pressure, self.atoms.get_volume(), 

264 self.energy, 0.0, 0.0, len(self.atoms)] 

265 if self.spacegroup: 

266 tokens.append('(' + self.spacegroup + ')') 

267 else: 

268 tokens.append('(P1)') 

269 if self.times_found: 

270 tokens.append('n - ' + str(self.times_found)) 

271 else: 

272 tokens.append('n - 1') 

273 

274 return ' '.join([str(tok) for tok in tokens]) 

275 

276 

277def read_res(filename, index=-1): 

278 """ 

279 Read input in SHELX (.res) format 

280 

281 Multiple frames are read if `filename` contains a wildcard character, 

282 e.g. `file_*.res`. `index` specifes which frames to retun: default is 

283 last frame only (index=-1). 

284 """ 

285 images = [] 

286 for fn in sorted(glob.glob(filename)): 

287 res = Res.from_file(fn) 

288 if res.energy: 

289 calc = SinglePointCalculator(res.atoms, 

290 energy=res.energy) 

291 res.atoms.calc = calc 

292 images.append(res.atoms) 

293 return images[index] 

294 

295 

296def write_res(filename, images, write_info=True, 

297 write_results=True, significant_figures=6): 

298 """ 

299 Write output in SHELX (.res) format 

300 

301 To write multiple images, include a % format string in filename, 

302 e.g. `file_%03d.res`. 

303 

304 Optionally include contents of Atoms.info dictionary if `write_info` 

305 is True, and/or results from attached calculator if `write_results` 

306 is True (only energy results are supported). 

307 """ 

308 

309 if not isinstance(images, (list, tuple)): 

310 images = [images] 

311 

312 if len(images) > 1 and '%' not in filename: 

313 raise RuntimeError('More than one Atoms provided but no %' + 

314 ' format string found in filename') 

315 

316 for i, atoms in enumerate(images): 

317 fn = filename 

318 if '%' in filename: 

319 fn = filename % i 

320 res = Res(atoms) 

321 if write_results: 

322 calculator = atoms.calc 

323 if (calculator is not None and 

324 isinstance(calculator, Calculator)): 

325 energy = calculator.results.get('energy') 

326 if energy is not None: 

327 res.energy = energy 

328 res.write_file(fn, write_info=write_info, 

329 significant_figures=significant_figures)