Coverage for /builds/ase/ase/ase/io/res.py: 94.77%
153 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +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 Res object.
94 """
95 with open(filename) as fd:
96 return Res.from_string(fd.read())
98 @staticmethod
99 def parse_title(line):
100 info = {}
102 tokens = line.split()
103 num_tokens = len(tokens)
104 # 1 = Name
105 if num_tokens <= 1:
106 return info
107 info['name'] = tokens[1]
108 # 2 = Pressure
109 if num_tokens <= 2:
110 return info
111 info['pressure'] = float(tokens[2])
112 # 3 = Volume
113 # 4 = Internal energy
114 if num_tokens <= 4:
115 return info
116 info['energy'] = float(tokens[4])
117 # 5 = Spin density, 6 - Abs spin density
118 # 7 = Space group OR num atoms (new format ONLY)
119 idx = 7
120 if tokens[idx][0] != '(':
121 idx += 1
123 if num_tokens <= idx:
124 return info
125 info['spacegroup'] = tokens[idx][1:len(tokens[idx]) - 1]
126 # idx + 1 = n, idx + 2 = -
127 # idx + 3 = times found
128 if num_tokens <= idx + 3:
129 return info
130 info['times_found'] = int(tokens[idx + 3])
132 return info
134 @staticmethod
135 def from_string(data):
136 """
137 Reads a Res from a string.
139 Args:
140 data (str): string containing Res data.
142 Returns:
143 Res object.
144 """
145 abc = []
146 ang = []
147 sp = []
148 coords = []
149 info = {}
150 coord_patt = re.compile(r"""(\w+)\s+
151 ([0-9]+)\s+
152 ([0-9\-\.]+)\s+
153 ([0-9\-\.]+)\s+
154 ([0-9\-\.]+)\s+
155 ([0-9\-\.]+)""", re.VERBOSE)
156 lines = data.splitlines()
157 line_no = 0
158 while line_no < len(lines):
159 line = lines[line_no]
160 tokens = line.split()
161 if tokens:
162 if tokens[0] == 'TITL':
163 try:
164 info = Res.parse_title(line)
165 except (ValueError, IndexError):
166 info = {}
167 elif tokens[0] == 'CELL' and len(tokens) == 8:
168 abc = [float(tok) for tok in tokens[2:5]]
169 ang = [float(tok) for tok in tokens[5:8]]
170 elif tokens[0] == 'SFAC':
171 for atom_line in lines[line_no:]:
172 if line.strip() == 'END':
173 break
174 else:
175 match = coord_patt.search(atom_line)
176 if match:
177 sp.append(match.group(1)) # 1-indexed
178 cs = match.groups()[2:5]
179 coords.append([float(c) for c in cs])
180 line_no += 1 # Make sure the global is updated
181 line_no += 1
183 return Res(Atoms(symbols=sp,
184 scaled_positions=coords,
185 cell=cellpar_to_cell(list(abc) + list(ang)),
186 pbc=True, info=info),
187 info.get('name'),
188 info.get('pressure'),
189 info.get('energy'),
190 info.get('spacegroup'),
191 info.get('times_found'))
193 def get_string(self, significant_figures=6, write_info=False):
194 """
195 Returns a string to be written as a Res file.
197 Args:
198 significant_figures (int): No. of significant figures to
199 output all quantities. Defaults to 6.
201 write_info (bool): if True, format TITL line using key-value pairs
202 from atoms.info in addition to attributes stored in Res object
204 Returns:
205 String representation of Res.
206 """
208 # Title line
209 if write_info:
210 info = self.atoms.info.copy()
211 for attribute in ['name', 'pressure', 'energy',
212 'spacegroup', 'times_found']:
213 if getattr(self, attribute) and attribute not in info:
214 info[attribute] = getattr(self, attribute)
215 lines = ['TITL ' + ' '.join([f'{k}={v}'
216 for (k, v) in info.items()])]
217 else:
218 lines = ['TITL ' + self.print_title()]
220 # Cell
221 abc_ang = cell_to_cellpar(self.atoms.get_cell())
222 fmt = f'{{0:.{significant_figures}f}}'
223 cell = ' '.join([fmt.format(a) for a in abc_ang])
224 lines.append('CELL 1.0 ' + cell)
226 # Latt
227 lines.append('LATT -1')
229 # Atoms
230 symbols = self.atoms.get_chemical_symbols()
231 species_types = []
232 for symbol in symbols:
233 if symbol not in species_types:
234 species_types.append(symbol)
235 lines.append('SFAC ' + ' '.join(species_types))
237 fmt = '{{0}} {{1}} {{2:.{0}f}} {{3:.{0}f}} {{4:.{0}f}} 1.0'
238 fmtstr = fmt.format(significant_figures)
239 for symbol, coords in zip(symbols,
240 self.atoms_.get_scaled_positions()):
241 lines.append(
242 fmtstr.format(symbol,
243 species_types.index(symbol) + 1,
244 coords[0],
245 coords[1],
246 coords[2]))
247 lines.append('END')
248 return '\n'.join(lines)
250 def __str__(self):
251 """
252 String representation of Res file.
253 """
254 return self.get_string()
256 def write_file(self, filename, **kwargs):
257 """
258 Writes Res to a file. The supported kwargs are the same as those for
259 the Res.get_string method and are passed through directly.
260 """
261 with open(filename, 'w') as fd:
262 fd.write(self.get_string(**kwargs) + '\n')
264 def print_title(self):
265 tokens = [self.name, self.pressure, self.atoms.get_volume(),
266 self.energy, 0.0, 0.0, len(self.atoms)]
267 if self.spacegroup:
268 tokens.append('(' + self.spacegroup + ')')
269 else:
270 tokens.append('(P1)')
271 if self.times_found:
272 tokens.append('n - ' + str(self.times_found))
273 else:
274 tokens.append('n - 1')
276 return ' '.join([str(tok) for tok in tokens])
279def read_res(filename, index=-1):
280 """
281 Read input in SHELX (.res) format
283 Multiple frames are read if `filename` contains a wildcard character,
284 e.g. `file_*.res`. `index` specifes which frames to retun: default is
285 last frame only (index=-1).
286 """
287 images = []
288 for fn in sorted(glob.glob(filename)):
289 res = Res.from_file(fn)
290 if res.energy:
291 calc = SinglePointCalculator(res.atoms,
292 energy=res.energy)
293 res.atoms.calc = calc
294 images.append(res.atoms)
295 return images[index]
298def write_res(filename, images, write_info=True,
299 write_results=True, significant_figures=6):
300 r"""
301 Write output in SHELX (.res) format
303 To write multiple images, include a % format string in filename,
304 e.g. `file\_\%03d.res`.
306 Optionally include contents of Atoms.info dictionary if `write_info`
307 is True, and/or results from attached calculator if `write_results`
308 is True (only energy results are supported).
309 """
311 if not isinstance(images, (list, tuple)):
312 images = [images]
314 if len(images) > 1 and '%' not in filename:
315 raise RuntimeError('More than one Atoms provided but no %' +
316 ' format string found in filename')
318 for i, atoms in enumerate(images):
319 fn = filename
320 if '%' in filename:
321 fn = filename % i
322 res = Res(atoms)
323 if write_results:
324 calculator = atoms.calc
325 if (calculator is not None and
326 isinstance(calculator, Calculator)):
327 energy = calculator.results.get('energy')
328 if energy is not None:
329 res.energy = energy
330 res.write_file(fn, write_info=write_info,
331 significant_figures=significant_figures)