Coverage for /builds/kinetik161/ase/ase/io/res.py: 94.77%
153 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1"""
2SHELX (.res) input/output
4Read/write files in SHELX (.res) file format.
6Format documented at http://shelx.uni-ac.gwdg.de/SHELX/
8Written by Martin Uhren and Georg Schusteritsch.
9Adapted for ASE by James Kermode.
10"""
13import glob
14import re
16from ase.atoms import Atoms
17from ase.calculators.calculator import Calculator
18from ase.calculators.singlepoint import SinglePointCalculator
19from ase.geometry import cell_to_cellpar, cellpar_to_cell
21__all__ = ['Res', 'read_res', 'write_res']
24class Res:
26 """
27 Object for representing the data in a Res file.
28 Most attributes can be set directly.
30 Args:
31 atoms (Atoms): Atoms object.
33 .. attribute:: atoms
35 Associated Atoms object.
37 .. attribute:: name
39 The name of the structure.
41 .. attribute:: pressure
43 The external pressure.
45 .. attribute:: energy
47 The internal energy of the structure.
49 .. attribute:: spacegroup
51 The space group of the structure.
53 .. attribute:: times_found
55 The number of times the structure was found.
56 """
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
75 @property
76 def atoms(self):
77 """
78 Returns Atoms object associated with this Res.
79 """
80 return self.atoms_
82 @staticmethod
83 def from_file(filename):
84 """
85 Reads a Res from a file.
87 Args:
88 filename (str): File name containing Res data.
90 Returns:
91 Res object.
92 """
93 with open(filename) as fd:
94 return Res.from_string(fd.read())
96 @staticmethod
97 def parse_title(line):
98 info = {}
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
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])
130 return info
132 @staticmethod
133 def from_string(data):
134 """
135 Reads a Res from a string.
137 Args:
138 data (str): string containing Res data.
140 Returns:
141 Res object.
142 """
143 abc = []
144 ang = []
145 sp = []
146 coords = []
147 info = {}
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 = {}
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
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'))
191 def get_string(self, significant_figures=6, write_info=False):
192 """
193 Returns a string to be written as a Res file.
195 Args:
196 significant_figures (int): No. of significant figures to
197 output all quantities. Defaults to 6.
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
202 Returns:
203 String representation of Res.
204 """
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([f'{k}={v}'
214 for (k, v) in info.items()])]
215 else:
216 lines = ['TITL ' + self.print_title()]
218 # Cell
219 abc_ang = cell_to_cellpar(self.atoms.get_cell())
220 fmt = f'{{0:.{significant_figures}f}}'
221 cell = ' '.join([fmt.format(a) for a in abc_ang])
222 lines.append('CELL 1.0 ' + cell)
224 # Latt
225 lines.append('LATT -1')
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))
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)
248 def __str__(self):
249 """
250 String representation of Res file.
251 """
252 return self.get_string()
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')
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')
274 return ' '.join([str(tok) for tok in tokens])
277def read_res(filename, index=-1):
278 """
279 Read input in SHELX (.res) format
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]
296def write_res(filename, images, write_info=True,
297 write_results=True, significant_figures=6):
298 """
299 Write output in SHELX (.res) format
301 To write multiple images, include a % format string in filename,
302 e.g. `file_%03d.res`.
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 """
309 if not isinstance(images, (list, tuple)):
310 images = [images]
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')
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)