Source code for ase.io.magres

"""This module provides I/O functions for the MAGRES file format, introduced
by CASTEP as an output format to store structural data and ab-initio
calculated NMR parameters.
Authors: Simone Sturniolo (ase implementation), Tim Green (original magres
    parser code)
"""

import re
from collections import OrderedDict

import numpy as np

import ase.units
from ase.atoms import Atoms
from ase.calculators.singlepoint import SinglePointDFTCalculator
from ase.spacegroup import Spacegroup
from ase.spacegroup.spacegroup import SpacegroupNotFoundError

_mprops = {
    'ms': ('sigma', 1),
    'sus': ('S', 0),
    'efg': ('V', 1),
    'isc': ('K', 2)}
# (matrix name, number of atoms in interaction) for various magres quantities


[docs]def read_magres(fd, include_unrecognised=False): """ Reader function for magres files. """ blocks_re = re.compile(r'[\[<](?P<block_name>.*?)[>\]](.*?)[<\[]/' + r'(?P=block_name)[\]>]', re.M | re.S) """ Here are defined the various functions required to parse different blocks. """ def tensor33(x): return np.squeeze(np.reshape(x, (3, 3))).tolist() def tensor31(x): return np.squeeze(np.reshape(x, (3, 1))).tolist() def get_version(file_contents): """ Look for and parse the magres file format version line """ lines = file_contents.split('\n') match = re.match(r'\#\$magres-abinitio-v([0-9]+).([0-9]+)', lines[0]) if match: version = match.groups() version = tuple(vnum for vnum in version) else: version = None return version def parse_blocks(file_contents): """ Parse series of XML-like deliminated blocks into a list of (block_name, contents) tuples """ blocks = blocks_re.findall(file_contents) return blocks def parse_block(block): """ Parse block contents into a series of (tag, data) records """ def clean_line(line): # Remove comments and whitespace at start and ends of line line = re.sub('#(.*?)\n', '', line) line = line.strip() return line name, data = block lines = [clean_line(line) for line in data.split('\n')] records = [] for line in lines: xs = line.split() if len(xs) > 0: tag = xs[0] data = xs[1:] records.append((tag, data)) return (name, records) def check_units(d): """ Verify that given units for a particular tag are correct. """ allowed_units = {'lattice': 'Angstrom', 'atom': 'Angstrom', 'ms': 'ppm', 'efg': 'au', 'efg_local': 'au', 'efg_nonlocal': 'au', 'isc': '10^19.T^2.J^-1', 'isc_fc': '10^19.T^2.J^-1', 'isc_orbital_p': '10^19.T^2.J^-1', 'isc_orbital_d': '10^19.T^2.J^-1', 'isc_spin': '10^19.T^2.J^-1', 'sus': '10^-6.cm^3.mol^-1', 'calc_cutoffenergy': 'Hartree', } if d[0] in d and d[1] == allowed_units[d[0]]: pass else: raise RuntimeError(f'Unrecognized units: {d[0]} {d[1]}') return d def parse_magres_block(block): """ Parse magres block into data dictionary given list of record tuples. """ name, records = block # 3x3 tensor def ntensor33(name): return lambda d: {name: tensor33([float(x) for x in data])} # Atom label, atom index and 3x3 tensor def sitensor33(name): return lambda d: {'atom': {'label': data[0], 'index': int(data[1])}, name: tensor33([float(x) for x in data[2:]])} # 2x(Atom label, atom index) and 3x3 tensor def sisitensor33(name): return lambda d: {'atom1': {'label': data[0], 'index': int(data[1])}, 'atom2': {'label': data[2], 'index': int(data[3])}, name: tensor33([float(x) for x in data[4:]])} tags = {'ms': sitensor33('sigma'), 'sus': ntensor33('S'), 'efg': sitensor33('V'), 'efg_local': sitensor33('V'), 'efg_nonlocal': sitensor33('V'), 'isc': sisitensor33('K'), 'isc_fc': sisitensor33('K'), 'isc_spin': sisitensor33('K'), 'isc_orbital_p': sisitensor33('K'), 'isc_orbital_d': sisitensor33('K'), 'units': check_units} data_dict = {} for record in records: tag, data = record if tag not in data_dict: data_dict[tag] = [] data_dict[tag].append(tags[tag](data)) return data_dict def parse_atoms_block(block): """ Parse atoms block into data dictionary given list of record tuples. """ name, records = block # Lattice record: a1, a2 a3, b1, b2, b3, c1, c2 c3 def lattice(d): return tensor33([float(x) for x in data]) # Atom record: label, index, x, y, z def atom(d): return {'species': data[0], 'label': data[1], 'index': int(data[2]), 'position': tensor31([float(x) for x in data[3:]])} def symmetry(d): return ' '.join(data) tags = {'lattice': lattice, 'atom': atom, 'units': check_units, 'symmetry': symmetry} data_dict = {} for record in records: tag, data = record if tag not in data_dict: data_dict[tag] = [] data_dict[tag].append(tags[tag](data)) return data_dict def parse_generic_block(block): """ Parse any other block into data dictionary given list of record tuples. """ name, records = block data_dict = {} for record in records: tag, data = record if tag not in data_dict: data_dict[tag] = [] data_dict[tag].append(data) return data_dict """ Actual parser code. """ block_parsers = {'magres': parse_magres_block, 'atoms': parse_atoms_block, 'calculation': parse_generic_block, } file_contents = fd.read() # Solve incompatibility for atomic indices above 100 pattern_ms = r'(ms [a-zA-Z]{1,2})(\d+)' file_contents = re.sub(pattern_ms, r'\g<1> \g<2>', file_contents) pattern_efg = r'(efg [a-zA-Z]{1,2})(\d+)' file_contents = re.sub(pattern_efg, r'\g<1> \g<2>', file_contents) # This works as a validity check version = get_version(file_contents) if version is None: # This isn't even a .magres file! raise RuntimeError('File is not in standard Magres format') blocks = parse_blocks(file_contents) data_dict = {} for block_data in blocks: block = parse_block(block_data) if block[0] in block_parsers: block_dict = block_parsers[block[0]](block) data_dict[block[0]] = block_dict else: # Throw in the text content of blocks we don't recognise if include_unrecognised: data_dict[block[0]] = block_data[1] # Now the loaded data must be turned into an ASE Atoms object # First check if the file is even viable if 'atoms' not in data_dict: raise RuntimeError('Magres file does not contain structure data') # Allowed units handling. This is redundant for now but # could turn out useful in the future magres_units = {'Angstrom': ase.units.Ang} # Lattice parameters? if 'lattice' in data_dict['atoms']: try: u = dict(data_dict['atoms']['units'])['lattice'] except KeyError: raise RuntimeError('No units detected in file for lattice') u = magres_units[u] cell = np.array(data_dict['atoms']['lattice'][0]) * u pbc = True else: cell = None pbc = False # Now the atoms symbols = [] positions = [] indices = [] labels = [] if 'atom' in data_dict['atoms']: try: u = dict(data_dict['atoms']['units'])['atom'] except KeyError: raise RuntimeError('No units detected in file for atom positions') u = magres_units[u] # Now we have to account for the possibility of there being CASTEP # 'custom' species amongst the symbols custom_species = None for a in data_dict['atoms']['atom']: spec_custom = a['species'].split(':', 1) if len(spec_custom) > 1 and custom_species is None: # Add it to the custom info! custom_species = list(symbols) symbols.append(spec_custom[0]) positions.append(a['position']) indices.append(a['index']) labels.append(a['label']) if custom_species is not None: custom_species.append(a['species']) atoms = Atoms(cell=cell, pbc=pbc, symbols=symbols, positions=positions) # Add custom species if present if custom_species is not None: atoms.new_array('castep_custom_species', np.array(custom_species)) # Add the spacegroup, if present and recognizable if 'symmetry' in data_dict['atoms']: try: spg = Spacegroup(data_dict['atoms']['symmetry'][0]) except SpacegroupNotFoundError: # Not found spg = Spacegroup(1) # Most generic one atoms.info['spacegroup'] = spg # Set up the rest of the properties as arrays atoms.new_array('indices', np.array(indices)) atoms.new_array('labels', np.array(labels)) # Now for the magres specific stuff li_list = list(zip(labels, indices)) def create_magres_array(name, order, block): if order == 1: u_arr = [None] * len(li_list) elif order == 2: u_arr = [[None] * (i + 1) for i in range(len(li_list))] else: raise ValueError( 'Invalid order value passed to create_magres_array') for s in block: # Find the atom index/indices if order == 1: # First find out which atom this is at = (s['atom']['label'], s['atom']['index']) try: ai = li_list.index(at) except ValueError: raise RuntimeError('Invalid data in magres block') # Then add the relevant quantity u_arr[ai] = s[mn] else: at1 = (s['atom1']['label'], s['atom1']['index']) at2 = (s['atom2']['label'], s['atom2']['index']) ai1 = li_list.index(at1) ai2 = li_list.index(at2) # Sort them ai1, ai2 = sorted((ai1, ai2), reverse=True) u_arr[ai1][ai2] = s[mn] if order == 1: return np.array(u_arr) else: return np.array(u_arr, dtype=object) if 'magres' in data_dict: if 'units' in data_dict['magres']: atoms.info['magres_units'] = dict(data_dict['magres']['units']) for u in atoms.info['magres_units']: # This bit to keep track of tags u0 = u.split('_')[0] if u0 not in _mprops: raise RuntimeError('Invalid data in magres block') mn, order = _mprops[u0] if order > 0: u_arr = create_magres_array(mn, order, data_dict['magres'][u]) atoms.new_array(u, u_arr) else: # atoms.info['magres_data'] = atoms.info.get('magres_data', # {}) # # We only take element 0 because for this sort of data # # there should be only that # atoms.info['magres_data'][u] = \ # data_dict['magres'][u][0][mn] if atoms.calc is None: calc = SinglePointDFTCalculator(atoms) atoms.calc = calc atoms.calc.results[u] = data_dict['magres'][u][0][mn] if 'calculation' in data_dict: atoms.info['magresblock_calculation'] = data_dict['calculation'] if include_unrecognised: for b in data_dict: if b not in block_parsers: atoms.info['magresblock_' + b] = data_dict[b] return atoms
def tensor_string(tensor): return ' '.join(' '.join(str(x) for x in xs) for xs in tensor)
[docs]def write_magres(fd, image): """ A writing function for magres files. Two steps: first data are arranged into structures, then dumped to the actual file """ image_data = {} image_data['atoms'] = {'units': []} # Contains units, lattice and each individual atom if np.all(image.get_pbc()): # Has lattice! image_data['atoms']['units'].append(['lattice', 'Angstrom']) image_data['atoms']['lattice'] = [image.get_cell()] # Now for the atoms if image.has('labels'): labels = image.get_array('labels') else: labels = image.get_chemical_symbols() if image.has('indices'): indices = image.get_array('indices') else: indices = [labels[:i + 1].count(labels[i]) for i in range(len(labels))] # Iterate over atoms symbols = (image.get_array('castep_custom_species') if image.has('castep_custom_species') else image.get_chemical_symbols()) atom_info = list(zip(symbols, image.get_positions())) if len(atom_info) > 0: image_data['atoms']['units'].append(['atom', 'Angstrom']) image_data['atoms']['atom'] = [] for i, a in enumerate(atom_info): image_data['atoms']['atom'].append({ 'index': indices[i], 'position': a[1], 'species': a[0], 'label': labels[i]}) # Spacegroup, if present if 'spacegroup' in image.info: image_data['atoms']['symmetry'] = [image.info['spacegroup'] .symbol.replace(' ', '')] # Now go on to do the same for magres information if 'magres_units' in image.info: image_data['magres'] = {'units': []} for u in image.info['magres_units']: # Get the type p = u.split('_')[0] if p in _mprops: image_data['magres']['units'].append( [u, image.info['magres_units'][u]]) image_data['magres'][u] = [] mn, order = _mprops[p] if order == 0: # The case of susceptibility tens = { mn: image.calc.results[u] } image_data['magres'][u] = tens else: arr = image.get_array(u) li_tab = zip(labels, indices) for i, (lab, ind) in enumerate(li_tab): if order == 2: for j, (lab2, ind2) in enumerate(li_tab[:i + 1]): if arr[i][j] is not None: tens = {mn: arr[i][j], 'atom1': {'label': lab, 'index': ind}, 'atom2': {'label': lab2, 'index': ind2}} image_data['magres'][u].append(tens) elif order == 1: if arr[i] is not None: tens = {mn: arr[i], 'atom': {'label': lab, 'index': ind}} image_data['magres'][u].append(tens) # Calculation block, if present if 'magresblock_calculation' in image.info: image_data['calculation'] = image.info['magresblock_calculation'] def write_units(data, out): if 'units' in data: for tag, units in data['units']: out.append(f' units {tag} {units}') def write_magres_block(data): """ Write out a <magres> block from its dictionary representation """ out = [] def nout(tag, tensor_name): if tag in data: out.append(' '.join([' ', tag, tensor_string(data[tag][tensor_name])])) def siout(tag, tensor_name): if tag in data: for atom_si in data[tag]: out.append((' %s %s %d ' '%s') % (tag, atom_si['atom']['label'], atom_si['atom']['index'], tensor_string(atom_si[tensor_name]))) write_units(data, out) nout('sus', 'S') siout('ms', 'sigma') siout('efg_local', 'V') siout('efg_nonlocal', 'V') siout('efg', 'V') def sisiout(tag, tensor_name): if tag in data: for isc in data[tag]: out.append((' %s %s %d %s %d ' '%s') % (tag, isc['atom1']['label'], isc['atom1']['index'], isc['atom2']['label'], isc['atom2']['index'], tensor_string(isc[tensor_name]))) sisiout('isc_fc', 'K') sisiout('isc_orbital_p', 'K') sisiout('isc_orbital_d', 'K') sisiout('isc_spin', 'K') sisiout('isc', 'K') return '\n'.join(out) def write_atoms_block(data): out = [] write_units(data, out) if 'lattice' in data: for lat in data['lattice']: out.append(f" lattice {tensor_string(lat)}") if 'symmetry' in data: for sym in data['symmetry']: out.append(f' symmetry {sym}') if 'atom' in data: for a in data['atom']: out.append((' atom %s %s %s ' '%s') % (a['species'], a['label'], a['index'], ' '.join(str(x) for x in a['position']))) return '\n'.join(out) def write_generic_block(data): out = [] for tag, data in data.items(): for value in data: out.append('{} {}'.format(tag, ' '.join(str(x) for x in value))) return '\n'.join(out) # Using this to preserve order block_writers = OrderedDict([('calculation', write_generic_block), ('atoms', write_atoms_block), ('magres', write_magres_block)]) # First, write the header fd.write('#$magres-abinitio-v1.0\n') fd.write('# Generated by the Atomic Simulation Environment library\n') for b in block_writers: if b in image_data: fd.write(f'[{b}]\n') fd.write(block_writers[b](image_data[b])) fd.write(f'\n[/{b}]\n') # Now on to check for any non-standard blocks... for i in image.info: if '_' in i: ismag, b = i.split('_', 1) if ismag == 'magresblock' and b not in block_writers: fd.write(f'[{b}]\n') fd.write(image.info[i]) fd.write(f'[/{b}]\n')