index | modules | gitlab

# Source code for ase.io.eon

    # Copyright (C) 2012-2023, Jesper Friis, SINTEF
    # Copyright (C) 2024, Rohit Goswami, UI
    # (see accompanying license files for ASE).
    """Module to read and write atoms EON reactant.con files.

    See http://theory.cm.utexas.edu/eon/index.html for a description of EON.
    """
    from dataclasses import dataclass
    from typing import List, Tuple
    from warnings import warn

    import numpy as np

    from ase.atoms import Atoms
    from ase.constraints import FixAtoms
    from ase.geometry import cell_to_cellpar, cellpar_to_cell
    from ase.utils import reader, writer

    ISOTOPE_TOL = 1e-8

    @dataclass
    class EONHeader:
        """
        A data class for storing header information of an EON file.

        Attributes
        ----------
        header_lines : List[str]
            The first two comment lines from the EON file header.
        cell_lengths : Tuple[float, float, float]
            The lengths of the cell vectors.
        cell_angles : Tuple[float, float, float]
            The angles between the cell vectors.
        Ncomponent : int
            The number of distinct atom types.
        component_counts : List[int]
            A list containing the count of atoms for each type.
        masses : List[float]
            A list containing the atomic masses for each type.
        """

        header_lines: List[str]
        # Actually these are float float float but.. mypy complains
        cell_lengths: Tuple[float, ...]
        cell_angles: Tuple[float, ...]
        Ncomponent: int
        component_counts: List[int]
        masses: List[float]

    def process_header(lines: List[str]) -> EONHeader:
        """
        Processes the header lines from an EON file and returns an EONHeader object.

        This function parses the first 9 lines of an EON file, extracting
        information about the simulation cell, the number of atom types, their
        counts, and masses, and encapsulates this information in an EONHeader
        object.

        Parameters
        ----------
        lines : List[str]
            The first 9 lines of an EON file containing header information.

        Returns
        -------
        EONHeader
            An object containing the parsed header information.
        """
        header_lines = lines[:2]

        # Parse cell lengths and angles
        cell_lengths = tuple(map(float, lines[2].split()))
        cell_angles = tuple(map(float, lines[3].split()))
        if len(cell_lengths) != 3 or len(cell_angles) != 3:
            raise ValueError(
                "Cell lengths and angles must each contain exactly three values."
            )

        Ncomponent = int(lines[6])
        component_counts = list(map(int, lines[7].split()))
        masses = list(map(float, lines[8].split()))

        return EONHeader(
            header_lines=header_lines,
            cell_lengths=cell_lengths,
            cell_angles=cell_angles,
            Ncomponent=Ncomponent,
            component_counts=component_counts,
            masses=masses,
        )

    def make_atoms(coordblock, header):
        """
        Creates an Atoms object from coordinate blocks and header information.

        This function takes a list of coordinate blocks and the associated header
        information, constructs the cell, sets the atomic positions, masses, and
        optionally applies FixAtoms constraints based on the header information, and
        returns an ASE Atoms object.

        Parameters
        ----------
        coordblock : list of str
            The lines from an EON file representing atom coordinates and types.
        header : EONHeader
            The parsed header information.

        Returns
        -------
        Atoms
            An ASE Atoms object constructed from the given coordinate blocks and
            header.
        """
        symbols = []
        coords = []
        masses = []
        fixed = []
        # Ordering in EON is different from the ASE convention
        cell_angles = (
            header.cell_angles[2],
            header.cell_angles[1],
            header.cell_angles[0]
        )
        cellpar = [x for x in header.cell_lengths + cell_angles]
        for idx, nblock in enumerate(header.component_counts):
            elem_block = coordblock[:(nblock + 2)]
            symb = elem_block[0]
            symbols.extend(nblock * [symb])
            mass = header.masses[idx]
            masses.extend(nblock * [mass])
            for eline in elem_block[2:]:
                tokens = eline.split()
                coords.append([float(x) for x in tokens[:3]])
                fixed.append(bool(int(tokens[3])))
            coordblock = coordblock[(nblock + 2):]
        return Atoms(
            symbols=symbols,
            positions=coords,
            masses=masses,
            cell=cellpar_to_cell(cellpar),
            constraint=FixAtoms(mask=fixed),
        )

    @reader
    def read_eon(fileobj, index=-1):
        """
        Reads an EON file or directory and returns one or more Atoms objects.

        This function handles single EON files, in both single image and multi-image
        variants. It returns either a single Atoms object, a list of Atoms objects,
        or a specific Atoms object indexed from the file or directory.

        Parameters
        ----------
        fileobj : str or Path or file-like object
            The path to the EON file or directory, or an open file-like object.
        index : int, optional
            The index of the Atoms object to return. If -1 (default), returns all
            objects or a single object if only one is present.

        Returns
        -------
        Atoms or List[Atoms]
            Depending on the `index` parameter and the content of the fileobj,
            returns either a single Atoms object or a list of Atoms objects.
        """
        images = []
        while True:
            # Read and process headers if they exist
            try:
                lines = [next(fileobj).strip() for _ in range(9)]
            except StopIteration:
                break  # End of file
            header = process_header(lines)
            num_blocklines = (header.Ncomponent * 2) + sum(header.component_counts)
            coordblocks = [next(fileobj).strip() for _ in range(num_blocklines)]
            atoms = make_atoms(coordblocks, header)
            images.append(atoms)

        # XXX: I really don't like this since there should be a consistent return
        if index == -1:
            return images if len(images) > 1 else images[0]
        else:
            return images[index]

    @writer
    def write_eon(fileobj, images, comment="Generated by ASE"):
        """
        Writes structures to an EON trajectory file, allowing for multiple
        snapshots.

        This function iterates over all provided images, converting each one into a
        text format compatible with EON trajectory files. It handles the conversion
        of the cell parameters, chemical symbols, atomic masses, and atomic
        constraints. Only FixAtoms constraints are supported; any other constraints
        will generate a warning and be skipped. Arbitrary masses will raise an
        error, since EON will not accept them, i.e. no Duterium and Hydrogen.

        Parameters
        ----------
        fileobj : file object
            An opened, writable file or file-like object to which the EON trajectory
            information will be written.
        images : list of Atoms
            A list of ASE Atoms objects representing the images (atomic structures)
            to be written into the EON trajectory file. Each Atoms object should
            have a cell attribute and may have a constraints attribute. If the
            constraints attribute is not a FixAtoms object, a warning will be
            issued.
        comment : str
            An additional comment statement which will be written out as the first
            header

        Raises
        ------
        Warning
            If any constraint in an Atoms object is not an instance of FixAtoms.
        RuntimeError
            If the masses for the species are not the default ones, i.e. if isotopes
            are present

        Returns
        -------
        None
            The function writes directly to the provided file object.

        See Also
        --------
        ase.io.utils.segment_list : for segmenting the list of images.

        Examples
        --------
        >>> from ase.io import Trajectory
        >>> from ase.io.utils import segment_list
        >>> traj = Trajectory("neb.traj")
        >>> n_images = 10  # Segment size, i.e. number of images in the NEB
        >>> for idx, pth in enumerate(segment_list(traj, n_images)):
        ...     with open(f"outputs/neb_path_{idx:03d}.con", "w") as fobj:
        ...         write_eon_traj(fobj, pth)
        """

        for idx, atoms in enumerate(images):
            out = []
            if idx == 0:
                out.append(comment)
            else:
                out.append(f"\n{comment}")
            out.append("preBox_header_2")  # ??

            a, b, c, alpha, beta, gamma = cell_to_cellpar(atoms.cell)
            out.append("%-10.6f  %-10.6f  %-10.6f" % (a, b, c))
            out.append("%-10.6f  %-10.6f  %-10.6f" % (gamma, beta, alpha))

            out.append("postBox_header_1")  # ??
            out.append("postBox_header_2")  # ??

            symbol_indices = atoms.symbols.indices()
            natoms = [len(x) for x in symbol_indices.values()]
            masslist = atoms.get_masses()
            species_masses = []
            for symbol, indices in symbol_indices.items():
                masses_for_symb = masslist[indices]
                if np.ptp(masses_for_symb) > ISOTOPE_TOL:
                    raise RuntimeError(
                        "Isotopes cannot be handled by EON"
                        ", ensure uniform masses for symbols"
                    )
                species_masses.append(masses_for_symb[0])

            out.append(str(len(symbol_indices)))
            out.append(" ".join([str(n) for n in natoms]))
            out.append(" ".join([str(n) for n in species_masses]))

            atom_id = 0
            for cid, (species, indices) in enumerate(symbol_indices.items()):
                fixed = np.array([False] * len(atoms))
                out.append(species)
                out.append("Coordinates of Component %d" % (cid + 1))
                atom = atoms[indices]
                coords = atom.positions
                for constraint in atom.constraints:
                    if not isinstance(constraint, FixAtoms):
                        warn(
                            "Only FixAtoms constraints are supported"
                            "by con files. Dropping %r",
                            constraint,
                        )
                        continue
                    if constraint.index.dtype.kind == "b":
                        fixed = np.array(constraint.index, dtype=int)
                    else:
                        fixed = np.zeros((natoms[cid],), dtype=int)
                        for i in constraint.index:
                            fixed[i] = 1
                for xyz, fix in zip(coords, fixed):
                    line_fmt = "{:>22.17f} {:>22.17f} {:>22.17f} {:d} {:4d}"
                    line = line_fmt.format(*xyz, int(fix), atom_id)
                    out.append(line)
                    atom_id += 1
            fileobj.write("\n".join(out))