Coverage for ase / io / magres.py: 79.65%
226 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-21 15:52 +0000
1# fmt: off
3"""This module provides I/O functions for the MAGRES file format, introduced
4by CASTEP as an output format to store structural data and ab-initio
5calculated NMR parameters.
6Authors: Simone Sturniolo (ase implementation), Tim Green (original magres
7 parser code)
8"""
10from collections import OrderedDict
12import numpy as np
14import ase.io._magres as _magres
15import ase.units
16from ase.atoms import Atoms
17from ase.calculators.singlepoint import SinglePointDFTCalculator
18from ase.spacegroup import Spacegroup
19from ase.spacegroup.spacegroup import SpacegroupNotFoundError
21_mprops = {
22 'ms': ('sigma', 1),
23 'sus': ('S', 0),
24 'efg': ('V', 1),
25 'isc': ('K', 2)}
26# (matrix name, number of atoms in interaction) for various magres quantities
29def read_magres(fd, include_unrecognised=False):
30 """Read magres file."""
32 block_parsers = {'magres': _magres.parse_magres_block,
33 'atoms': _magres.parse_atoms_block,
34 'calculation': _magres.parse_generic_block, }
36 file_contents = fd.read()
38 # This works as a validity check
39 version = _magres.get_version(file_contents)
40 if version is None:
41 # This isn't even a .magres file!
42 raise RuntimeError('File is not in standard Magres format')
43 blocks = _magres.parse_blocks(file_contents)
45 data_dict = {}
47 for block_data in blocks:
48 block = _magres.parse_block(block_data)
50 if block[0] in block_parsers:
51 block_dict = block_parsers[block[0]](block)
52 data_dict[block[0]] = block_dict
53 else:
54 # Throw in the text content of blocks we don't recognise
55 if include_unrecognised:
56 data_dict[block[0]] = block_data[1]
58 # Now the loaded data must be turned into an ASE Atoms object
60 # First check if the file is even viable
61 if 'atoms' not in data_dict:
62 raise RuntimeError('Magres file does not contain structure data')
64 # Allowed units handling. This is redundant for now but
65 # could turn out useful in the future
67 magres_units = {'Angstrom': ase.units.Ang}
69 # Lattice parameters?
70 if 'lattice' in data_dict['atoms']:
71 try:
72 u = dict(data_dict['atoms']['units'])['lattice']
73 except KeyError:
74 raise RuntimeError('No units detected in file for lattice')
75 u = magres_units[u]
76 cell = np.array(data_dict['atoms']['lattice'][0]) * u
77 pbc = True
78 else:
79 cell = None
80 pbc = False
82 # Now the atoms
83 symbols = []
84 positions = []
85 indices = []
86 labels = []
88 if 'atom' in data_dict['atoms']:
89 try:
90 u = dict(data_dict['atoms']['units'])['atom']
91 except KeyError:
92 raise RuntimeError('No units detected in file for atom positions')
93 u = magres_units[u]
94 # Now we have to account for the possibility of there being CASTEP
95 # 'custom' species amongst the symbols
96 custom_species = None
97 for a in data_dict['atoms']['atom']:
98 spec_custom = a['species'].split(':', 1)
99 if len(spec_custom) > 1 and custom_species is None:
100 # Add it to the custom info!
101 custom_species = list(symbols)
102 symbols.append(spec_custom[0])
103 positions.append(a['position'])
104 indices.append(a['index'])
105 labels.append(a['label'])
106 if custom_species is not None:
107 custom_species.append(a['species'])
109 atoms = Atoms(cell=cell,
110 pbc=pbc,
111 symbols=symbols,
112 positions=positions)
114 # Add custom species if present
115 if custom_species is not None:
116 atoms.new_array('castep_custom_species', np.array(custom_species))
118 # Add the spacegroup, if present and recognizable
119 if 'symmetry' in data_dict['atoms']:
120 try:
121 spg = Spacegroup(data_dict['atoms']['symmetry'][0])
122 except SpacegroupNotFoundError:
123 # Not found
124 spg = Spacegroup(1) # Most generic one
125 atoms.info['spacegroup'] = spg
127 # Set up the rest of the properties as arrays
128 atoms.new_array('indices', np.array(indices))
129 atoms.new_array('labels', np.array(labels))
131 # Now for the magres specific stuff
132 li_list = list(zip(labels, indices))
134 def create_magres_array(name, order, block):
136 if order == 1:
137 u_arr = [None] * len(li_list)
138 elif order == 2:
139 u_arr = [[None] * (i + 1) for i in range(len(li_list))]
140 else:
141 raise ValueError(
142 'Invalid order value passed to create_magres_array')
144 for s in block:
145 # Find the atom index/indices
146 if order == 1:
147 # First find out which atom this is
148 at = (s['atom']['label'], s['atom']['index'])
149 try:
150 ai = li_list.index(at)
151 except ValueError:
152 raise RuntimeError('Invalid data in magres block')
153 # Then add the relevant quantity
154 u_arr[ai] = s[mn]
155 else:
156 at1 = (s['atom1']['label'], s['atom1']['index'])
157 at2 = (s['atom2']['label'], s['atom2']['index'])
158 ai1 = li_list.index(at1)
159 ai2 = li_list.index(at2)
160 # Sort them
161 ai1, ai2 = sorted((ai1, ai2), reverse=True)
162 u_arr[ai1][ai2] = s[mn]
164 if order == 1:
165 return np.array(u_arr)
166 else:
167 return np.array(u_arr, dtype=object)
169 if 'magres' in data_dict:
170 if 'units' in data_dict['magres']:
171 atoms.info['magres_units'] = dict(data_dict['magres']['units'])
172 for u in atoms.info['magres_units']:
173 # This bit to keep track of tags
174 u0 = u.split('_')[0]
176 if u0 not in _mprops:
177 raise RuntimeError('Invalid data in magres block')
179 mn, order = _mprops[u0]
181 if order > 0:
182 u_arr = create_magres_array(mn, order,
183 data_dict['magres'][u])
184 atoms.new_array(u, u_arr)
185 else:
186 # atoms.info['magres_data'] = atoms.info.get('magres_data',
187 # {})
188 # # We only take element 0 because for this sort of data
189 # # there should be only that
190 # atoms.info['magres_data'][u] = \
191 # data_dict['magres'][u][0][mn]
192 if atoms.calc is None:
193 calc = SinglePointDFTCalculator(atoms)
194 atoms.calc = calc
195 atoms.calc.results[u] = data_dict['magres'][u][0][mn]
197 if 'calculation' in data_dict:
198 atoms.info['magresblock_calculation'] = data_dict['calculation']
200 if include_unrecognised:
201 for b in data_dict:
202 if b not in block_parsers:
203 atoms.info['magresblock_' + b] = data_dict[b]
205 return atoms
208def tensor_string(tensor):
209 return ' '.join(' '.join(str(x) for x in xs) for xs in tensor)
212def write_magres(fd, image):
213 """
214 A writing function for magres files. Two steps: first data are arranged
215 into structures, then dumped to the actual file
216 """
218 image_data = {}
219 image_data['atoms'] = {'units': []}
220 # Contains units, lattice and each individual atom
221 if np.all(image.get_pbc()):
222 # Has lattice!
223 image_data['atoms']['units'].append(['lattice', 'Angstrom'])
224 image_data['atoms']['lattice'] = [image.get_cell()]
226 # Now for the atoms
227 if image.has('labels'):
228 labels = image.get_array('labels')
229 else:
230 labels = image.get_chemical_symbols()
232 if image.has('indices'):
233 indices = image.get_array('indices')
234 else:
235 indices = [labels[:i + 1].count(labels[i]) for i in range(len(labels))]
237 # Iterate over atoms
238 symbols = (image.get_array('castep_custom_species')
239 if image.has('castep_custom_species')
240 else image.get_chemical_symbols())
242 atom_info = list(zip(symbols,
243 image.get_positions()))
244 if len(atom_info) > 0:
245 image_data['atoms']['units'].append(['atom', 'Angstrom'])
246 image_data['atoms']['atom'] = []
248 for i, a in enumerate(atom_info):
249 image_data['atoms']['atom'].append({
250 'index': indices[i],
251 'position': a[1],
252 'species': a[0],
253 'label': labels[i]})
255 # Spacegroup, if present
256 if 'spacegroup' in image.info:
257 image_data['atoms']['symmetry'] = [image.info['spacegroup']
258 .symbol.replace(' ', '')]
260 # Now go on to do the same for magres information
261 if 'magres_units' in image.info:
263 image_data['magres'] = {'units': []}
265 for u in image.info['magres_units']:
266 # Get the type
267 p = u.split('_')[0]
268 if p in _mprops:
269 image_data['magres']['units'].append(
270 [u, image.info['magres_units'][u]])
271 image_data['magres'][u] = []
272 mn, order = _mprops[p]
274 if order == 0:
275 # The case of susceptibility
276 tens = {
277 mn: image.calc.results[u]
278 }
279 image_data['magres'][u] = tens
280 else:
281 arr = image.get_array(u)
282 li_tab = zip(labels, indices)
283 for i, (lab, ind) in enumerate(li_tab):
284 if order == 2:
285 for j, (lab2, ind2) in enumerate(li_tab[:i + 1]):
286 if arr[i][j] is not None:
287 tens = {mn: arr[i][j],
288 'atom1': {'label': lab,
289 'index': ind},
290 'atom2': {'label': lab2,
291 'index': ind2}}
292 image_data['magres'][u].append(tens)
293 elif order == 1:
294 if arr[i] is not None:
295 tens = {mn: arr[i],
296 'atom': {'label': lab,
297 'index': ind}}
298 image_data['magres'][u].append(tens)
300 # Calculation block, if present
301 if 'magresblock_calculation' in image.info:
302 image_data['calculation'] = image.info['magresblock_calculation']
304 def write_units(data, out):
305 if 'units' in data:
306 for tag, units in data['units']:
307 out.append(f' units {tag} {units}')
309 def write_magres_block(data):
310 """
311 Write out a <magres> block from its dictionary representation
312 """
314 out = []
316 def nout(tag, tensor_name):
317 if tag in data:
318 out.append(' '.join([' ', tag,
319 tensor_string(data[tag][tensor_name])]))
321 def siout(tag, tensor_name):
322 if tag in data:
323 for atom_si in data[tag]:
324 out.append((' %s %s %d '
325 '%s') % (tag,
326 atom_si['atom']['label'],
327 atom_si['atom']['index'],
328 tensor_string(atom_si[tensor_name])))
330 write_units(data, out)
332 nout('sus', 'S')
333 siout('ms', 'sigma')
335 siout('efg_local', 'V')
336 siout('efg_nonlocal', 'V')
337 siout('efg', 'V')
339 def sisiout(tag, tensor_name):
340 if tag in data:
341 for isc in data[tag]:
342 out.append((' %s %s %d %s %d '
343 '%s') % (tag,
344 isc['atom1']['label'],
345 isc['atom1']['index'],
346 isc['atom2']['label'],
347 isc['atom2']['index'],
348 tensor_string(isc[tensor_name])))
350 sisiout('isc_fc', 'K')
351 sisiout('isc_orbital_p', 'K')
352 sisiout('isc_orbital_d', 'K')
353 sisiout('isc_spin', 'K')
354 sisiout('isc', 'K')
356 return '\n'.join(out)
358 def write_atoms_block(data):
359 out = []
361 write_units(data, out)
363 if 'lattice' in data:
364 for lat in data['lattice']:
365 out.append(f" lattice {tensor_string(lat)}")
367 if 'symmetry' in data:
368 for sym in data['symmetry']:
369 out.append(f' symmetry {sym}')
371 if 'atom' in data:
372 for a in data['atom']:
373 out.append((' atom %s %s %s '
374 '%s') % (a['species'],
375 a['label'],
376 a['index'],
377 ' '.join(str(x) for x in a['position'])))
379 return '\n'.join(out)
381 def write_generic_block(data):
382 out = []
384 for tag, data in data.items():
385 for value in data:
386 out.append('{} {}'.format(tag, ' '.join(str(x) for x in value)))
388 return '\n'.join(out)
390 # Using this to preserve order
391 block_writers = OrderedDict([('calculation', write_generic_block),
392 ('atoms', write_atoms_block),
393 ('magres', write_magres_block)])
395 # First, write the header
396 fd.write('#$magres-abinitio-v1.0\n')
397 fd.write('# Generated by the Atomic Simulation Environment library\n')
399 for b in block_writers:
400 if b in image_data:
401 fd.write(f'[{b}]\n')
402 fd.write(block_writers[b](image_data[b]))
403 fd.write(f'\n[/{b}]\n')
405 # Now on to check for any non-standard blocks...
406 for i in image.info:
407 if '_' in i:
408 ismag, b = i.split('_', 1)
409 if ismag == 'magresblock' and b not in block_writers:
410 fd.write(f'[{b}]\n')
411 fd.write(image.info[i])
412 fd.write(f'[/{b}]\n')