import re
import numpy as np
from ase.atoms import Atoms
from ase.calculators.lammps import Prism, convert
from ase.utils import reader, writer
[docs]@reader
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.
"""
# load everything into memory
lines = fileobj.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]),
)
elif (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
[docs]@writer
def write_lammps_data(fd, atoms, specorder=None, force_skew=False,
prismobj=None, velocities=False, units="metal",
atom_style='atomic'):
"""Write atomic structure data to a LAMMPS data file."""
# 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]
if hasattr(fd, "name"):
fd.write("{0} (written by ASE) \n\n".format(fd.name))
else:
fd.write("(written by ASE) \n\n")
symbols = atoms.get_chemical_symbols()
n_atoms = len(symbols)
fd.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)
fd.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)
fd.write("0.0 {0:23.17g} xlo xhi\n".format(xhi))
fd.write("0.0 {0:23.17g} ylo yhi\n".format(yhi))
fd.write("0.0 {0:23.17g} zlo zhi\n".format(zhi))
if force_skew or p.is_skewed():
fd.write(
"{0:23.17g} {1:23.17g} {2:23.17g} xy xz yz\n".format(
xy, xz, yz
)
)
fd.write("\n\n")
# Write (unwrapped) atomic positions. If wrapping of atoms back into the
# cell along periodic directions is desired, this should be done manually
# on the Atoms object itself beforehand.
fd.write("Atoms \n\n")
pos = p.vector_to_lammps(atoms.get_positions(), wrap=False)
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
fd.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
fd.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()
# The label 'mol-id' has apparenlty been introduced in read earlier,
# but so far not implemented here. Wouldn't a 'underscored' label
# be better, i.e. 'mol_id' or 'molecule_id'?
if atoms.has('mol-id'):
molecules = atoms.get_array('mol-id')
if not np.issubdtype(molecules.dtype, np.integer):
raise TypeError((
"If 'atoms' object has 'mol-id' array, then"
" mol-id dtype must be subtype of np.integer, and"
" not {:s}.").format(str(molecules.dtype)))
if (len(molecules) != len(atoms)) or (molecules.ndim != 1):
raise TypeError((
"If 'atoms' object has 'mol-id' array, then"
" each atom must have exactly one mol-id."))
else:
# Assigning each atom to a distinct molecule id would seem
# preferableabove assigning all atoms to a single molecule id per
# default, as done within ase <= v 3.19.1. I.e.,
# molecules = np.arange(start=1, stop=len(atoms)+1, step=1, dtype=int)
# However, according to LAMMPS default behavior,
molecules = np.zeros(len(atoms), dtype=int)
# which is what happens if one creates new atoms within LAMMPS
# without explicitly taking care of the molecule id.
# Quote from docs at https://lammps.sandia.gov/doc/read_data.html:
# The molecule ID is a 2nd identifier attached to an atom.
# Normally, it is a number from 1 to N, identifying which
# molecule the atom belongs to. It can be 0 if it is a
# non-bonded atom or if you don't care to keep track of molecule
# assignments.
for i, (m, q, r) in enumerate(zip(molecules, 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
fd.write("{0:>6} {1:>3} {2:>3} {3:>5} {4:23.17g} {5:23.17g} "
"{6:23.17g}\n".format(*(i + 1, m, s, q) + tuple(r)))
else:
raise NotImplementedError
if velocities and atoms.get_velocities() is not None:
fd.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)
fd.write(
"{0:>6} {1:23.17g} {2:23.17g} {3:23.17g}\n".format(
*(i + 1,) + tuple(v)
)
)
fd.flush()