Skip to main content
`Ctrl`+`K`

#

# Source code for ase.io.qbox

    # fmt: off

    """This module contains functions to read from QBox output files"""

    import re
    import xml.etree.ElementTree as ET

    from ase import Atom, Atoms
    from ase.calculators.singlepoint import SinglePointCalculator
    from ase.utils import reader

    # Compile regexs for fixing XML
    re_find_bad_xml = re.compile(r'<(/?)([A-z]+) expectation ([a-z]+)')

    @reader
    def read_qbox(f, index=-1):
        """Read data from QBox output file

        Inputs:
            f - str or fileobj, path to file or file object to read from
            index - int or slice, which frames to return

        Returns:
            list of Atoms or atoms, requested frame(s)
        """

        # Check whether this is a QB@all output
        version = None
        for line in f:
            if '<release>' in line:
                version = ET.fromstring(line)
                break
        if version is None:
            raise Exception('Parse Error: Version not found')
        is_qball = 'qb@LL' in version.text or 'qball' in version.text

        # Load in atomic species
        species = {}
        if is_qball:
            # Read all of the lines between release and the first call to `run`
            species_data = []
            for line in f:
                if '<run' in line:
                    break
                species_data.append(line)
            species_data = '\n'.join(species_data)

            # Read out the species information with regular expressions
            symbols = re.findall('symbol_ = ([A-Z][a-z]?)', species_data)
            masses = re.findall('mass_ = ([0-9.]+)', species_data)
            names = re.findall('name_ = ([a-z]+)', species_data)
            numbers = re.findall('atomic_number_ = ([0-9]+)', species_data)

            # Compile them into a dictionary
            for name, symbol, mass, number in zip(names, symbols, masses, numbers):
                spec_data = dict(
                    symbol=symbol,
                    mass=float(mass),
                    number=float(number)
                )
                species[name] = spec_data
        else:
            # Find all species
            species_blocks = _find_blocks(f, 'species', '<cmd>run')

            for spec in species_blocks:
                name = spec.get('name')
                spec_data = dict(
                    symbol=spec.find('symbol').text,
                    mass=float(spec.find('mass').text),
                    number=int(spec.find('atomic_number').text))
                species[name] = spec_data

        # Find all of the frames
        frames = _find_blocks(f, 'iteration', None)

        # If index is an int, return one frame
        if isinstance(index, int):
            return _parse_frame(frames[index], species)
        else:
            return [_parse_frame(frame, species) for frame in frames[index]]

    def _find_blocks(fp, tag, stopwords='[qbox]'):
        """Find and parse a certain block of the file.

        Reads a file sequentially and stops when it either encounters the
        end of the file, or until the it encounters a line that contains a
        user-defined string *after it has already found at least one
        desired block*. Use the stopwords ``[qbox]`` to read until the next
        command is issued.

        Groups the text between the first line that contains <tag> and the
        next line that contains </tag>, inclusively. The function then
        parses the XML and returns the Element object.

        Inputs:
            fp        - file-like object, file to be read from
            tag       - str, tag to search for (e.g., 'iteration').
                        `None` if you want to read until the end of the file
            stopwords - str, halt parsing if a line containing this string
                        is encountered

        Returns:
            list of xml.ElementTree, parsed XML blocks found by this class
        """

        start_tag = f'<{tag}'
        end_tag = f'</{tag}>'

        blocks = []  # Stores all blocks
        cur_block = []  # Block being filled
        in_block = False  # Whether we are currently parsing
        for line in fp:

            # Check if the block has started
            if start_tag in line:
                if in_block:
                    raise Exception('Parsing failed: Encountered nested block')
                else:
                    in_block = True

            # Add data to block
            if in_block:
                cur_block.append(line)

            # Check for stopping conditions
            if stopwords is not None:
                if stopwords in line and len(blocks) > 0:
                    break

            if end_tag in line:
                if in_block:
                    blocks.append(cur_block)
                    cur_block = []
                    in_block = False
                else:
                    raise Exception('Parsing failed: End tag found before start '
                                    'tag')

        # Join strings in a block into a single string
        blocks = [''.join(b) for b in blocks]

        # Ensure XML compatibility. There are two specific tags in QBall that are
        # not valid XML, so we need to run a
        blocks = [re_find_bad_xml.sub(r'<\1\2_expectation_\3', b) for b in blocks]

        # Parse the blocks
        return [ET.fromstring(b) for b in blocks]

    def _parse_frame(tree, species):
        """Parse a certain frame from QBOX output

        Inputs:
            tree - ElementTree, <iteration> block from output file
            species - dict, data about species. Key is name of atom type,
                value is data about that type
        Return:
            Atoms object describing this iteration"""

        # Load in data about the system
        energy = float(tree.find("etotal").text)

        # Load in data about the cell
        unitcell = tree.find('atomset').find('unit_cell')
        cell = []
        for d in ['a', 'b', 'c']:
            cell.append([float(x) for x in unitcell.get(d).split()])

        stress_tree = tree.find('stress_tensor')
        if stress_tree is None:
            stresses = None
        else:
            stresses = [float(stress_tree.find(f'sigma_{x}').text)
                        for x in ['xx', 'yy', 'zz', 'yz', 'xz', 'xy']]

        # Create the Atoms object
        atoms = Atoms(pbc=True, cell=cell)

        # Load in the atom information
        forces = []
        for atom in tree.find('atomset').findall('atom'):
            # Load data about the atom type
            spec = atom.get('species')
            symbol = species[spec]['symbol']
            mass = species[spec]['mass']

            # Get data about position / velocity / force
            pos = [float(x) for x in atom.find('position').text.split()]
            force = [float(x) for x in atom.find('force').text.split()]
            momentum = [float(x) * mass
                        for x in atom.find('velocity').text.split()]

            # Create the objects
            atom = Atom(symbol=symbol, mass=mass, position=pos, momentum=momentum)
            atoms += atom
            forces.append(force)

        # Create the calculator object that holds energy/forces
        calc = SinglePointCalculator(atoms,
                                     energy=energy, forces=forces, stress=stresses)
        atoms.calc = calc

        return atoms