Skip to main content
`Ctrl`+`K`

#

# Source code for ase.io.res

    # fmt: off

    """
    SHELX (.res) input/output

    Read/write files in SHELX (.res) file format.

    Format documented at http://shelx.uni-ac.gwdg.de/SHELX/

    Written by Martin Uhren and Georg Schusteritsch.
    Adapted for ASE by James Kermode.
    """

    import glob
    import re

    from ase.atoms import Atoms
    from ase.calculators.calculator import Calculator
    from ase.calculators.singlepoint import SinglePointCalculator
    from ase.geometry import cell_to_cellpar, cellpar_to_cell

    __all__ = ['Res', 'read_res', 'write_res']

    class Res:

        """
        Object for representing the data in a Res file.
        Most attributes can be set directly.

        Args:
            atoms (Atoms):  Atoms object.

        .. attribute:: atoms

            Associated Atoms object.

        .. attribute:: name

            The name of the structure.

        .. attribute:: pressure

            The external pressure.

        .. attribute:: energy

            The internal energy of the structure.

        .. attribute:: spacegroup

            The space group of the structure.

        .. attribute:: times_found

            The number of times the structure was found.
        """

        def __init__(self, atoms, name=None, pressure=None,
                     energy=None, spacegroup=None, times_found=None):
            self.atoms_ = atoms
            if name is None:
                name = atoms.info.get('name')
            if pressure is None:
                pressure = atoms.info.get('pressure')
            if spacegroup is None:
                spacegroup = atoms.info.get('spacegroup')
            if times_found is None:
                times_found = atoms.info.get('times_found')
            self.name = name
            self.pressure = pressure
            self.energy = energy
            self.spacegroup = spacegroup
            self.times_found = times_found

        @property
        def atoms(self):
            """
            Returns Atoms object associated with this Res.
            """
            return self.atoms_

        @staticmethod
        def from_file(filename):
            """
            Reads a Res from a file.

            Args:
                filename (str): File name containing Res data.

            Returns:
                Res object.
            """
            with open(filename) as fd:
                return Res.from_string(fd.read())

        @staticmethod
        def parse_title(line):
            info = {}

            tokens = line.split()
            num_tokens = len(tokens)
            # 1 = Name
            if num_tokens <= 1:
                return info
            info['name'] = tokens[1]
            # 2 = Pressure
            if num_tokens <= 2:
                return info
            info['pressure'] = float(tokens[2])
            # 3 = Volume
            # 4 = Internal energy
            if num_tokens <= 4:
                return info
            info['energy'] = float(tokens[4])
            # 5 = Spin density, 6 - Abs spin density
            # 7 = Space group OR num atoms (new format ONLY)
            idx = 7
            if tokens[idx][0] != '(':
                idx += 1

            if num_tokens <= idx:
                return info
            info['spacegroup'] = tokens[idx][1:len(tokens[idx]) - 1]
            # idx + 1 = n, idx + 2 = -
            # idx + 3 = times found
            if num_tokens <= idx + 3:
                return info
            info['times_found'] = int(tokens[idx + 3])

            return info

        @staticmethod
        def from_string(data):
            """
            Reads a Res from a string.

            Args:
                data (str): string containing Res data.

            Returns:
                Res object.
            """
            abc = []
            ang = []
            sp = []
            coords = []
            info = {}
            coord_patt = re.compile(r"""(\w+)\s+
                                        ([0-9]+)\s+
                                        ([0-9\-\.]+)\s+
                                        ([0-9\-\.]+)\s+
                                        ([0-9\-\.]+)\s+
                                        ([0-9\-\.]+)""", re.VERBOSE)
            lines = data.splitlines()
            line_no = 0
            while line_no < len(lines):
                line = lines[line_no]
                tokens = line.split()
                if tokens:
                    if tokens[0] == 'TITL':
                        try:
                            info = Res.parse_title(line)
                        except (ValueError, IndexError):
                            info = {}
                    elif tokens[0] == 'CELL' and len(tokens) == 8:
                        abc = [float(tok) for tok in tokens[2:5]]
                        ang = [float(tok) for tok in tokens[5:8]]
                    elif tokens[0] == 'SFAC':
                        for atom_line in lines[line_no:]:
                            if line.strip() == 'END':
                                break
                            else:
                                match = coord_patt.search(atom_line)
                                if match:
                                    sp.append(match.group(1))  # 1-indexed
                                    cs = match.groups()[2:5]
                                    coords.append([float(c) for c in cs])
                            line_no += 1  # Make sure the global is updated
                line_no += 1

            return Res(Atoms(symbols=sp,
                             scaled_positions=coords,
                             cell=cellpar_to_cell(list(abc) + list(ang)),
                             pbc=True, info=info),
                       info.get('name'),
                       info.get('pressure'),
                       info.get('energy'),
                       info.get('spacegroup'),
                       info.get('times_found'))

        def get_string(self, significant_figures=6, write_info=False):
            """
            Returns a string to be written as a Res file.

            Args:
                significant_figures (int): No. of significant figures to
                    output all quantities. Defaults to 6.

                write_info (bool): if True, format TITL line using key-value pairs
                   from atoms.info in addition to attributes stored in Res object

            Returns:
                String representation of Res.
            """

            # Title line
            if write_info:
                info = self.atoms.info.copy()
                for attribute in ['name', 'pressure', 'energy',
                                  'spacegroup', 'times_found']:
                    if getattr(self, attribute) and attribute not in info:
                        info[attribute] = getattr(self, attribute)
                lines = ['TITL ' + ' '.join([f'{k}={v}'
                                             for (k, v) in info.items()])]
            else:
                lines = ['TITL ' + self.print_title()]

            # Cell
            abc_ang = cell_to_cellpar(self.atoms.get_cell())
            fmt = f'{{0:.{significant_figures}f}}'
            cell = ' '.join([fmt.format(a) for a in abc_ang])
            lines.append('CELL 1.0 ' + cell)

            # Latt
            lines.append('LATT -1')

            # Atoms
            symbols = self.atoms.get_chemical_symbols()
            species_types = []
            for symbol in symbols:
                if symbol not in species_types:
                    species_types.append(symbol)
            lines.append('SFAC ' + ' '.join(species_types))

            fmt = '{{0}} {{1}} {{2:.{0}f}} {{3:.{0}f}} {{4:.{0}f}} 1.0'
            fmtstr = fmt.format(significant_figures)
            for symbol, coords in zip(symbols,
                                      self.atoms_.get_scaled_positions()):
                lines.append(
                    fmtstr.format(symbol,
                                  species_types.index(symbol) + 1,
                                  coords[0],
                                  coords[1],
                                  coords[2]))
            lines.append('END')
            return '\n'.join(lines)

        def __str__(self):
            """
            String representation of Res file.
            """
            return self.get_string()

        def write_file(self, filename, **kwargs):
            """
            Writes Res to a file. The supported kwargs are the same as those for
            the Res.get_string method and are passed through directly.
            """
            with open(filename, 'w') as fd:
                fd.write(self.get_string(**kwargs) + '\n')

        def print_title(self):
            tokens = [self.name, self.pressure, self.atoms.get_volume(),
                      self.energy, 0.0, 0.0, len(self.atoms)]
            if self.spacegroup:
                tokens.append('(' + self.spacegroup + ')')
            else:
                tokens.append('(P1)')
            if self.times_found:
                tokens.append('n - ' + str(self.times_found))
            else:
                tokens.append('n - 1')

            return ' '.join([str(tok) for tok in tokens])

    def read_res(filename, index=-1):
        """
        Read input in SHELX (.res) format

        Multiple frames are read if `filename` contains a wildcard character,
        e.g. `file_*.res`. `index` specifes which frames to retun: default is
        last frame only (index=-1).
        """
        images = []
        for fn in sorted(glob.glob(filename)):
            res = Res.from_file(fn)
            if res.energy:
                calc = SinglePointCalculator(res.atoms,
                                             energy=res.energy)
                res.atoms.calc = calc
            images.append(res.atoms)
        return images[index]

    def write_res(filename, images, write_info=True,
                  write_results=True, significant_figures=6):
        r"""
        Write output in SHELX (.res) format

        To write multiple images, include a % format string in filename,
        e.g. `file\_\%03d.res`.

        Optionally include contents of Atoms.info dictionary if `write_info`
        is True, and/or results from attached calculator if `write_results`
        is True (only energy results are supported).
        """

        if not isinstance(images, (list, tuple)):
            images = [images]

        if len(images) > 1 and '%' not in filename:
            raise RuntimeError('More than one Atoms provided but no %' +
                               ' format string found in filename')

        for i, atoms in enumerate(images):
            fn = filename
            if '%' in filename:
                fn = filename % i
            res = Res(atoms)
            if write_results:
                calculator = atoms.calc
                if (calculator is not None and
                        isinstance(calculator, Calculator)):
                    energy = calculator.results.get('energy')
                    if energy is not None:
                        res.energy = energy
            res.write_file(fn, write_info=write_info,
                           significant_figures=significant_figures)