Coverage for ase / io / res.py: 94.77%
153 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3"""
4SHELX (.res) input/output
6Read/write files in SHELX (.res) file format.
8Format documented at http://shelx.uni-ac.gwdg.de/SHELX/
10Written by Martin Uhren and Georg Schusteritsch.
11Adapted for ASE by James Kermode.
12"""
15import glob
16import re
18from ase.atoms import Atoms
19from ase.calculators.calculator import Calculator
20from ase.calculators.singlepoint import SinglePointCalculator
21from ase.geometry import cell_to_cellpar, cellpar_to_cell
23__all__ = ['Res', 'read_res', 'write_res']
26class Res:
28 """
29 Object for representing the data in a Res file.
30 Most attributes can be set directly.
32 Args:
33 atoms (Atoms): Atoms object.
35 .. attribute:: atoms
37 Associated Atoms object.
39 .. attribute:: name
41 The name of the structure.
43 .. attribute:: pressure
45 The external pressure.
47 .. attribute:: energy
49 The internal energy of the structure.
51 .. attribute:: spacegroup
53 The space group of the structure.
55 .. attribute:: times_found
57 The number of times the structure was found.
58 """
60 def __init__(self, atoms, name=None, pressure=None,
61 energy=None, spacegroup=None, times_found=None):
62 self.atoms_ = atoms
63 if name is None:
64 name = atoms.info.get('name')
65 if pressure is None:
66 pressure = atoms.info.get('pressure')
67 if spacegroup is None:
68 spacegroup = atoms.info.get('spacegroup')
69 if times_found is None:
70 times_found = atoms.info.get('times_found')
71 self.name = name
72 self.pressure = pressure
73 self.energy = energy
74 self.spacegroup = spacegroup
75 self.times_found = times_found
77 @property
78 def atoms(self):
79 """
80 Returns Atoms object associated with this Res.
81 """
82 return self.atoms_
84 @staticmethod
85 def from_file(filename):
86 """
87 Reads a Res from a file.
89 Args:
90 filename (str): File name containing Res data.
92 Returns
93 -------
94 Res object.
95 """
96 with open(filename) as fd:
97 return Res.from_string(fd.read())
99 @staticmethod
100 def parse_title(line):
101 info = {}
103 tokens = line.split()
104 num_tokens = len(tokens)
105 # 1 = Name
106 if num_tokens <= 1:
107 return info
108 info['name'] = tokens[1]
109 # 2 = Pressure
110 if num_tokens <= 2:
111 return info
112 info['pressure'] = float(tokens[2])
113 # 3 = Volume
114 # 4 = Internal energy
115 if num_tokens <= 4:
116 return info
117 info['energy'] = float(tokens[4])
118 # 5 = Spin density, 6 - Abs spin density
119 # 7 = Space group OR num atoms (new format ONLY)
120 idx = 7
121 if tokens[idx][0] != '(':
122 idx += 1
124 if num_tokens <= idx:
125 return info
126 info['spacegroup'] = tokens[idx][1:len(tokens[idx]) - 1]
127 # idx + 1 = n, idx + 2 = -
128 # idx + 3 = times found
129 if num_tokens <= idx + 3:
130 return info
131 info['times_found'] = int(tokens[idx + 3])
133 return info
135 @staticmethod
136 def from_string(data):
137 """
138 Reads a Res from a string.
140 Args:
141 data (str): string containing Res data.
143 Returns
144 -------
145 Res object.
146 """
147 abc = []
148 ang = []
149 sp = []
150 coords = []
151 info = {}
152 coord_patt = re.compile(r"""(\w+)\s+
153 ([0-9]+)\s+
154 ([0-9\-\.]+)\s+
155 ([0-9\-\.]+)\s+
156 ([0-9\-\.]+)\s+
157 ([0-9\-\.]+)""", re.VERBOSE)
158 lines = data.splitlines()
159 line_no = 0
160 while line_no < len(lines):
161 line = lines[line_no]
162 tokens = line.split()
163 if tokens:
164 if tokens[0] == 'TITL':
165 try:
166 info = Res.parse_title(line)
167 except (ValueError, IndexError):
168 info = {}
169 elif tokens[0] == 'CELL' and len(tokens) == 8:
170 abc = [float(tok) for tok in tokens[2:5]]
171 ang = [float(tok) for tok in tokens[5:8]]
172 elif tokens[0] == 'SFAC':
173 for atom_line in lines[line_no:]:
174 if line.strip() == 'END':
175 break
176 else:
177 match = coord_patt.search(atom_line)
178 if match:
179 sp.append(match.group(1)) # 1-indexed
180 cs = match.groups()[2:5]
181 coords.append([float(c) for c in cs])
182 line_no += 1 # Make sure the global is updated
183 line_no += 1
185 return Res(Atoms(symbols=sp,
186 scaled_positions=coords,
187 cell=cellpar_to_cell(list(abc) + list(ang)),
188 pbc=True, info=info),
189 info.get('name'),
190 info.get('pressure'),
191 info.get('energy'),
192 info.get('spacegroup'),
193 info.get('times_found'))
195 def get_string(self, significant_figures=6, write_info=False):
196 """
197 Returns a string to be written as a Res file.
199 Args:
200 significant_figures (int): No. of significant figures to
201 output all quantities. Defaults to 6.
203 write_info (bool): if True, format TITL line using key-value pairs
204 from atoms.info in addition to attributes stored in Res object
206 Returns
207 -------
208 String representation of Res.
209 """
211 # Title line
212 if write_info:
213 info = self.atoms.info.copy()
214 for attribute in ['name', 'pressure', 'energy',
215 'spacegroup', 'times_found']:
216 if getattr(self, attribute) and attribute not in info:
217 info[attribute] = getattr(self, attribute)
218 lines = ['TITL ' + ' '.join([f'{k}={v}'
219 for (k, v) in info.items()])]
220 else:
221 lines = ['TITL ' + self.print_title()]
223 # Cell
224 abc_ang = cell_to_cellpar(self.atoms.get_cell())
225 fmt = f'{{0:.{significant_figures}f}}'
226 cell = ' '.join([fmt.format(a) for a in abc_ang])
227 lines.append('CELL 1.0 ' + cell)
229 # Latt
230 lines.append('LATT -1')
232 # Atoms
233 symbols = self.atoms.get_chemical_symbols()
234 species_types = []
235 for symbol in symbols:
236 if symbol not in species_types:
237 species_types.append(symbol)
238 lines.append('SFAC ' + ' '.join(species_types))
240 fmt = '{{0}} {{1}} {{2:.{0}f}} {{3:.{0}f}} {{4:.{0}f}} 1.0'
241 fmtstr = fmt.format(significant_figures)
242 for symbol, coords in zip(symbols,
243 self.atoms_.get_scaled_positions()):
244 lines.append(
245 fmtstr.format(symbol,
246 species_types.index(symbol) + 1,
247 coords[0],
248 coords[1],
249 coords[2]))
250 lines.append('END')
251 return '\n'.join(lines)
253 def __str__(self):
254 """
255 String representation of Res file.
256 """
257 return self.get_string()
259 def write_file(self, filename, **kwargs):
260 """
261 Writes Res to a file. The supported kwargs are the same as those for
262 the Res.get_string method and are passed through directly.
263 """
264 with open(filename, 'w') as fd:
265 fd.write(self.get_string(**kwargs) + '\n')
267 def print_title(self):
268 tokens = [self.name, self.pressure, self.atoms.get_volume(),
269 self.energy, 0.0, 0.0, len(self.atoms)]
270 if self.spacegroup:
271 tokens.append('(' + self.spacegroup + ')')
272 else:
273 tokens.append('(P1)')
274 if self.times_found:
275 tokens.append('n - ' + str(self.times_found))
276 else:
277 tokens.append('n - 1')
279 return ' '.join([str(tok) for tok in tokens])
282def read_res(filename, index=-1):
283 """
284 Read input in SHELX (.res) format
286 Multiple frames are read if `filename` contains a wildcard character,
287 e.g. `file_*.res`. `index` specifes which frames to retun: default is
288 last frame only (index=-1).
289 """
290 images = []
291 for fn in sorted(glob.glob(filename)):
292 res = Res.from_file(fn)
293 if res.energy:
294 calc = SinglePointCalculator(res.atoms,
295 energy=res.energy)
296 res.atoms.calc = calc
297 images.append(res.atoms)
298 return images[index]
301def write_res(filename, images, write_info=True,
302 write_results=True, significant_figures=6):
303 r"""
304 Write output in SHELX (.res) format
306 To write multiple images, include a % format string in filename,
307 e.g. `file\_\%03d.res`.
309 Optionally include contents of Atoms.info dictionary if `write_info`
310 is True, and/or results from attached calculator if `write_results`
311 is True (only energy results are supported).
312 """
314 if not isinstance(images, (list, tuple)):
315 images = [images]
317 if len(images) > 1 and '%' not in filename:
318 raise RuntimeError('More than one Atoms provided but no %' +
319 ' format string found in filename')
321 for i, atoms in enumerate(images):
322 fn = filename
323 if '%' in filename:
324 fn = filename % i
325 res = Res(atoms)
326 if write_results:
327 calculator = atoms.calc
328 if (calculator is not None and
329 isinstance(calculator, Calculator)):
330 energy = calculator.results.get('energy')
331 if energy is not None:
332 res.energy = energy
333 res.write_file(fn, write_info=write_info,
334 significant_figures=significant_figures)