index | modules | gitlab

# Source code for ase.io.lammpsdata

    import re
    import numpy as np

    from ase.atoms import Atoms
    from ase.parallel import paropen
    from ase.utils import basestring
    from ase.calculators.lammps import Prism, convert

    def read_lammps_data(fileobj, Z_of_type=None, style="full",
                         sort_by_id=False, units="metal"):
        """Method which reads a LAMMPS data file.

        sort_by_id: Order the particles according to their id. Might be faster to
        switch it off.
        Units are set by default to the style=metal setting in LAMMPS.
        """
        if isinstance(fileobj, basestring):
            f = paropen(fileobj)
        else:
            f = fileobj

        # load everything into memory
        lines = f.readlines()

        # begin read_lammps_data
        comment = None
        N = None
        # N_types = None
        xlo = None
        xhi = None
        ylo = None
        yhi = None
        zlo = None
        zhi = None
        xy = None
        xz = None
        yz = None
        pos_in = {}
        travel_in = {}
        mol_id_in = {}
        charge_in = {}
        mass_in = {}
        vel_in = {}
        bonds_in = []
        angles_in = []
        dihedrals_in = []

        sections = [
            "Atoms",
            "Velocities",
            "Masses",
            "Charges",
            "Ellipsoids",
            "Lines",
            "Triangles",
            "Bodies",
            "Bonds",
            "Angles",
            "Dihedrals",
            "Impropers",
            "Impropers Pair Coeffs",
            "PairIJ Coeffs",
            "Pair Coeffs",
            "Bond Coeffs",
            "Angle Coeffs",
            "Dihedral Coeffs",
            "Improper Coeffs",
            "BondBond Coeffs",
            "BondAngle Coeffs",
            "MiddleBondTorsion Coeffs",
            "EndBondTorsion Coeffs",
            "AngleTorsion Coeffs",
            "AngleAngleTorsion Coeffs",
            "BondBond13 Coeffs",
            "AngleAngle Coeffs",
        ]
        header_fields = [
            "atoms",
            "bonds",
            "angles",
            "dihedrals",
            "impropers",
            "atom types",
            "bond types",
            "angle types",
            "dihedral types",
            "improper types",
            "extra bond per atom",
            "extra angle per atom",
            "extra dihedral per atom",
            "extra improper per atom",
            "extra special per atom",
            "ellipsoids",
            "lines",
            "triangles",
            "bodies",
            "xlo xhi",
            "ylo yhi",
            "zlo zhi",
            "xy xz yz",
        ]
        sections_re = "(" + "|".join(sections).replace(" ", "\\s+") + ")"
        header_fields_re = "(" + "|".join(header_fields).replace(" ", "\\s+") + ")"

        section = None
        header = True
        for line in lines:
            if comment is None:
                comment = line.rstrip()
            else:
                line = re.sub("#.*", "", line).rstrip().lstrip()
                if re.match("^\\s*$", line):  # skip blank lines
                    continue

            # check for known section names
            m = re.match(sections_re, line)
            if m is not None:
                section = m.group(0).rstrip().lstrip()
                header = False
                continue

            if header:
                field = None
                val = None
                # m = re.match(header_fields_re+"\s+=\s*(.*)", line)
                # if m is not None: # got a header line
                #   field=m.group(1).lstrip().rstrip()
                #   val=m.group(2).lstrip().rstrip()
                # else: # try other format
                #   m = re.match("(.*)\s+"+header_fields_re, line)
                #   if m is not None:
                #       field = m.group(2).lstrip().rstrip()
                #       val = m.group(1).lstrip().rstrip()
                m = re.match("(.*)\\s+" + header_fields_re, line)
                if m is not None:
                    field = m.group(2).lstrip().rstrip()
                    val = m.group(1).lstrip().rstrip()
                if field is not None and val is not None:
                    if field == "atoms":
                        N = int(val)
                    # elif field == "atom types":
                    #     N_types = int(val)
                    elif field == "xlo xhi":
                        (xlo, xhi) = [float(x) for x in val.split()]
                    elif field == "ylo yhi":
                        (ylo, yhi) = [float(x) for x in val.split()]
                    elif field == "zlo zhi":
                        (zlo, zhi) = [float(x) for x in val.split()]
                    elif field == "xy xz yz":
                        (xy, xz, yz) = [float(x) for x in val.split()]

            if section is not None:
                fields = line.split()
                if section == "Atoms":  # id *
                    id = int(fields[0])
                    if style == "full" and (len(fields) == 7 or len(fields) == 10):
                        # id mol-id type q x y z [tx ty tz]
                        pos_in[id] = (
                            int(fields[2]),
                            float(fields[4]),
                            float(fields[5]),
                            float(fields[6]),
                        )
                        mol_id_in[id] = int(fields[1])
                        charge_in[id] = float(fields[3])
                        if len(fields) == 10:
                            travel_in[id] = (
                                int(fields[7]),
                                int(fields[8]),
                                int(fields[9]),
                            )
                    elif style == "atomic" and (
                            len(fields) == 5 or len(fields) == 8
                    ):
                        # id type x y z [tx ty tz]
                        pos_in[id] = (
                            int(fields[1]),
                            float(fields[2]),
                            float(fields[3]),
                            float(fields[4]),
                        )
                        if len(fields) == 8:
                            travel_in[id] = (
                                int(fields[5]),
                                int(fields[6]),
                                int(fields[7]),
                            )
                    elif (style in ("angle", "bond", "molecular")
                          ) and (len(fields) == 6 or len(fields) == 9):
                        # id mol-id type x y z [tx ty tz]
                        pos_in[id] = (
                            int(fields[2]),
                            float(fields[3]),
                            float(fields[4]),
                            float(fields[5]),
                        )
                        mol_id_in[id] = int(fields[1])
                        if len(fields) == 9:
                            travel_in[id] = (
                                int(fields[6]),
                                int(fields[7]),
                                int(fields[8]),
                            )
                    if style == "charge" and (len(fields) == 6 or len(fields) == 9):
                        # id type q x y z [tx ty tz]
                        pos_in[id] = (
                            int(fields[1]),
                            float(fields[3]),
                            float(fields[4]),
                            float(fields[5]),
                        )
                        charge_in[id] = float(fields[2])
                        if len(fields) == 9:
                            travel_in[id] = (
                                int(fields[6]),
                                int(fields[7]),
                                int(fields[8]),
                            )
                    else:
                        raise RuntimeError(
                            "Style '{}' not supported or invalid "
                            "number of fields {}"
                            "".format(style, len(fields))
                        )
                elif section == "Velocities":  # id vx vy vz
                    vel_in[int(fields[0])] = (
                        float(fields[1]),
                        float(fields[2]),
                        float(fields[3]),
                    )
                elif section == "Masses":
                    mass_in[int(fields[0])] = float(fields[1])
                elif section == "Bonds":  # id type atom1 atom2
                    bonds_in.append(
                        (int(fields[1]), int(fields[2]), int(fields[3]))
                    )
                elif section == "Angles":  # id type atom1 atom2 atom3
                    angles_in.append(
                        (
                            int(fields[1]),
                            int(fields[2]),
                            int(fields[3]),
                            int(fields[4]),
                        )
                    )
                elif section == "Dihedrals":  # id type atom1 atom2 atom3 atom4
                    dihedrals_in.append(
                        (
                            int(fields[1]),
                            int(fields[2]),
                            int(fields[3]),
                            int(fields[4]),
                            int(fields[5]),
                        )
                    )

        # set cell
        cell = np.zeros((3, 3))
        cell[0, 0] = xhi - xlo
        cell[1, 1] = yhi - ylo
        cell[2, 2] = zhi - zlo
        if xy is not None:
            cell[1, 0] = xy
        if xz is not None:
            cell[2, 0] = xz
        if yz is not None:
            cell[2, 1] = yz

        # initialize arrays for per-atom quantities
        positions = np.zeros((N, 3))
        numbers = np.zeros((N), int)
        ids = np.zeros((N), int)
        types = np.zeros((N), int)
        if len(vel_in) > 0:
            velocities = np.zeros((N, 3))
        else:
            velocities = None
        if len(mass_in) > 0:
            masses = np.zeros((N))
        else:
            masses = None
        if len(mol_id_in) > 0:
            mol_id = np.zeros((N), int)
        else:
            mol_id = None
        if len(charge_in) > 0:
            charge = np.zeros((N), float)
        else:
            charge = None
        if len(travel_in) > 0:
            travel = np.zeros((N, 3), int)
        else:
            travel = None
        if len(bonds_in) > 0:
            bonds = [""] * N
        else:
            bonds = None
        if len(angles_in) > 0:
            angles = [""] * N
        else:
            angles = None
        if len(dihedrals_in) > 0:
            dihedrals = [""] * N
        else:
            dihedrals = None

        ind_of_id = {}
        # copy per-atom quantities from read-in values
        for (i, id) in enumerate(pos_in.keys()):
            # by id
            ind_of_id[id] = i
            if sort_by_id:
                ind = id - 1
            else:
                ind = i
            type = pos_in[id][0]
            positions[ind, :] = [pos_in[id][1], pos_in[id][2], pos_in[id][3]]
            if velocities is not None:
                velocities[ind, :] = [vel_in[id][0], vel_in[id][1], vel_in[id][2]]
            if travel is not None:
                travel[ind] = travel_in[id]
            if mol_id is not None:
                mol_id[ind] = mol_id_in[id]
            if charge is not None:
                charge[ind] = charge_in[id]
            ids[ind] = id
            # by type
            types[ind] = type
            if Z_of_type is None:
                numbers[ind] = type
            else:
                numbers[ind] = Z_of_type[type]
            if masses is not None:
                masses[ind] = mass_in[type]
        # convert units
        positions = convert(positions, "distance", units, "ASE")
        cell = convert(cell, "distance", units, "ASE")
        if masses is not None:
            masses = convert(masses, "mass", units, "ASE")
        if velocities is not None:
            velocities = convert(velocities, "velocity", units, "ASE")

        # create ase.Atoms
        at = Atoms(
            positions=positions,
            numbers=numbers,
            masses=masses,
            cell=cell,
            pbc=[True, True, True],
        )
        # set velocities (can't do it via constructor)
        if velocities is not None:
            at.set_velocities(velocities)
        at.arrays["id"] = ids
        at.arrays["type"] = types
        if travel is not None:
            at.arrays["travel"] = travel
        if mol_id is not None:
            at.arrays["mol-id"] = mol_id
        if charge is not None:
            at.arrays["initial_charges"] = charge
            at.arrays["mmcharges"] = charge.copy()

        if bonds is not None:
            for (type, a1, a2) in bonds_in:
                i_a1 = ind_of_id[a1]
                i_a2 = ind_of_id[a2]
                if len(bonds[i_a1]) > 0:
                    bonds[i_a1] += ","
                bonds[i_a1] += "%d(%d)" % (i_a2, type)
            for i in range(len(bonds)):
                if len(bonds[i]) == 0:
                    bonds[i] = "_"
            at.arrays["bonds"] = np.array(bonds)

        if angles is not None:
            for (type, a1, a2, a3) in angles_in:
                i_a1 = ind_of_id[a1]
                i_a2 = ind_of_id[a2]
                i_a3 = ind_of_id[a3]
                if len(angles[i_a2]) > 0:
                    angles[i_a2] += ","
                angles[i_a2] += "%d-%d(%d)" % (i_a1, i_a3, type)
            for i in range(len(angles)):
                if len(angles[i]) == 0:
                    angles[i] = "_"
            at.arrays["angles"] = np.array(angles)

        if dihedrals is not None:
            for (type, a1, a2, a3, a4) in dihedrals_in:
                i_a1 = ind_of_id[a1]
                i_a2 = ind_of_id[a2]
                i_a3 = ind_of_id[a3]
                i_a4 = ind_of_id[a4]
                if len(dihedrals[i_a1]) > 0:
                    dihedrals[i_a1] += ","
                dihedrals[i_a1] += "%d-%d-%d(%d)" % (i_a2, i_a3, i_a4, type)
            for i in range(len(dihedrals)):
                if len(dihedrals[i]) == 0:
                    dihedrals[i] = "_"
            at.arrays["dihedrals"] = np.array(dihedrals)

        at.info["comment"] = comment

        return at

    def write_lammps_data(fileobj, atoms, specorder=None, force_skew=False,
                          prismobj=None, velocities=False, units="metal",
                          atom_style='atomic'):
        """Write atomic structure data to a LAMMPS data file."""
        if isinstance(fileobj, basestring):
            f = paropen(fileobj, "w", encoding="ascii")
            close_file = True
        else:
            # Presume fileobj acts like a fileobj
            f = fileobj
            close_file = False

        # FIXME: We should add a check here that the encoding of the file object
        #        is actually ascii once the 'encoding' attribute of IOFormat objects
        #        starts functioning in implementation (currently it doesn't do
        #         anything).

        if isinstance(atoms, list):
            if len(atoms) > 1:
                raise ValueError(
                    "Can only write one configuration to a lammps data file!"
                )
            atoms = atoms[0]

        f.write("{0} (written by ASE) \n\n".format(f.name))

        symbols = atoms.get_chemical_symbols()
        n_atoms = len(symbols)
        f.write("{0} \t atoms \n".format(n_atoms))

        if specorder is None:
            # This way it is assured that LAMMPS atom types are always
            # assigned predictably according to the alphabetic order
            species = sorted(set(symbols))
        else:
            # To index elements in the LAMMPS data file
            # (indices must correspond to order in the potential file)
            species = specorder
        n_atom_types = len(species)
        f.write("{0}  atom types\n".format(n_atom_types))

        if prismobj is None:
            p = Prism(atoms.get_cell())
        else:
            p = prismobj

        # Get cell parameters and convert from ASE units to LAMMPS units
        xhi, yhi, zhi, xy, xz, yz = convert(p.get_lammps_prism(), "distance",
                "ASE", units)

        f.write("0.0 {0:23.17g}  xlo xhi\n".format(xhi))
        f.write("0.0 {0:23.17g}  ylo yhi\n".format(yhi))
        f.write("0.0 {0:23.17g}  zlo zhi\n".format(zhi))

        if force_skew or p.is_skewed():
            f.write(
                "{0:23.17g} {1:23.17g} {2:23.17g}  xy xz yz\n".format(
                    xy, xz, yz
                )
            )
        f.write("\n\n")

        f.write("Atoms \n\n")
        pos = p.vector_to_lammps(atoms.get_positions(), wrap=True)

        if atom_style == 'atomic':
            for i, r in enumerate(pos):
                # Convert position from ASE units to LAMMPS units
                r = convert(r, "distance", "ASE", units)
                s = species.index(symbols[i]) + 1
                f.write(
                    "{0:>6} {1:>3} {2:23.17g} {3:23.17g} {4:23.17g}\n".format(
                        *(i + 1, s) + tuple(r)
                    )
                )
        elif atom_style == 'charge':
            charges = atoms.get_initial_charges()
            for i, (q, r) in enumerate(zip(charges, pos)):
                # Convert position and charge from ASE units to LAMMPS units
                r = convert(r, "distance", "ASE", units)
                q = convert(q, "charge", "ASE", units)
                s = species.index(symbols[i]) + 1
                f.write(
                    "{0:>6} {1:>3} {2:>5} {3:23.17g} {4:23.17g} {5:23.17g}\n".format(
                        *(i + 1, s, q) + tuple(r)
                    )
                )
        elif atom_style == 'full':
            charges = atoms.get_initial_charges()
            molecule = 1 # Assign all atoms to a single molecule
            for i, (q, r) in enumerate(zip(charges, pos)):
                # Convert position and charge from ASE units to LAMMPS units
                r = convert(r, "distance", "ASE", units)
                q = convert(q, "charge", "ASE", units)
                s = species.index(symbols[i]) + 1
                f.write(
                    "{0:>6} {1:>3} {2:>3} {3:>5} {4:23.17g} {5:23.17g} {6:23.17g}\n".format(
                        *(i + 1, molecule, s, q) + tuple(r)
                    )
                )
        else:
            raise NotImplementedError

        if velocities and atoms.get_velocities() is not None:
            f.write("\n\nVelocities \n\n")
            vel = p.vector_to_lammps(atoms.get_velocities())
            for i, v in enumerate(vel):
                # Convert velocity from ASE units to LAMMPS units
                v = convert(v, "velocity", "ASE", units)
                f.write(
                    "{0:>6} {1:23.17g} {2:23.17g} {3:23.17g}\n".format(
                        *(i + 1,) + tuple(v)
                    )
                )

        f.flush()
        if close_file:
            f.close()

Documentation v: 3.19.0

Versions
  3.20.0b1 (development)
  3.19.0 (latest stable)