index | modules | gitlab

# 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
    import numpy as np
    from collections import OrderedDict

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

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

    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',
                             'isc': '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('Unrecognized units: %s %s' % (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()

        # 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)

    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('  units %s %s' % (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("  lattice %s" % tensor_string(lat))

            if 'symmetry' in data:
                for sym in data['symmetry']:
                    out.append('  symmetry %s' % 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('%s %s' % (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('[{0}]\n'.format(b))
                fd.write(block_writers[b](image_data[b]))
                fd.write('\n[/{0}]\n'.format(b))

        # 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('[{0}]\n'.format(b))
                    fd.write(image.info[i])
                    fd.write('[/{0}]\n'.format(b))