Coverage for /builds/ase/ase/ase/io/lammpsdata.py : 91.35%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import re
2import numpy as np
4from ase.atoms import Atoms
5from ase.calculators.lammps import Prism, convert
6from ase.utils import reader, writer
9@reader
10def read_lammps_data(fileobj, Z_of_type=None, style="full",
11 sort_by_id=True, units="metal"):
12 """Method which reads a LAMMPS data file.
14 sort_by_id: Order the particles according to their id. Might be faster to
15 switch it off.
16 Units are set by default to the style=metal setting in LAMMPS.
17 """
18 # load everything into memory
19 lines = fileobj.readlines()
21 # begin read_lammps_data
22 comment = None
23 N = None
24 # N_types = None
25 xlo = None
26 xhi = None
27 ylo = None
28 yhi = None
29 zlo = None
30 zhi = None
31 xy = None
32 xz = None
33 yz = None
34 pos_in = {}
35 travel_in = {}
36 mol_id_in = {}
37 charge_in = {}
38 mass_in = {}
39 vel_in = {}
40 bonds_in = []
41 angles_in = []
42 dihedrals_in = []
44 sections = [
45 "Atoms",
46 "Velocities",
47 "Masses",
48 "Charges",
49 "Ellipsoids",
50 "Lines",
51 "Triangles",
52 "Bodies",
53 "Bonds",
54 "Angles",
55 "Dihedrals",
56 "Impropers",
57 "Impropers Pair Coeffs",
58 "PairIJ Coeffs",
59 "Pair Coeffs",
60 "Bond Coeffs",
61 "Angle Coeffs",
62 "Dihedral Coeffs",
63 "Improper Coeffs",
64 "BondBond Coeffs",
65 "BondAngle Coeffs",
66 "MiddleBondTorsion Coeffs",
67 "EndBondTorsion Coeffs",
68 "AngleTorsion Coeffs",
69 "AngleAngleTorsion Coeffs",
70 "BondBond13 Coeffs",
71 "AngleAngle Coeffs",
72 ]
73 header_fields = [
74 "atoms",
75 "bonds",
76 "angles",
77 "dihedrals",
78 "impropers",
79 "atom types",
80 "bond types",
81 "angle types",
82 "dihedral types",
83 "improper types",
84 "extra bond per atom",
85 "extra angle per atom",
86 "extra dihedral per atom",
87 "extra improper per atom",
88 "extra special per atom",
89 "ellipsoids",
90 "lines",
91 "triangles",
92 "bodies",
93 "xlo xhi",
94 "ylo yhi",
95 "zlo zhi",
96 "xy xz yz",
97 ]
98 sections_re = "(" + "|".join(sections).replace(" ", "\\s+") + ")"
99 header_fields_re = "(" + "|".join(header_fields).replace(" ", "\\s+") + ")"
101 section = None
102 header = True
103 for line in lines:
104 if comment is None:
105 comment = line.rstrip()
106 else:
107 line = re.sub("#.*", "", line).rstrip().lstrip()
108 if re.match("^\\s*$", line): # skip blank lines
109 continue
111 # check for known section names
112 m = re.match(sections_re, line)
113 if m is not None:
114 section = m.group(0).rstrip().lstrip()
115 header = False
116 continue
118 if header:
119 field = None
120 val = None
121 # m = re.match(header_fields_re+"\s+=\s*(.*)", line)
122 # if m is not None: # got a header line
123 # field=m.group(1).lstrip().rstrip()
124 # val=m.group(2).lstrip().rstrip()
125 # else: # try other format
126 # m = re.match("(.*)\s+"+header_fields_re, line)
127 # if m is not None:
128 # field = m.group(2).lstrip().rstrip()
129 # val = m.group(1).lstrip().rstrip()
130 m = re.match("(.*)\\s+" + header_fields_re, line)
131 if m is not None:
132 field = m.group(2).lstrip().rstrip()
133 val = m.group(1).lstrip().rstrip()
134 if field is not None and val is not None:
135 if field == "atoms":
136 N = int(val)
137 # elif field == "atom types":
138 # N_types = int(val)
139 elif field == "xlo xhi":
140 (xlo, xhi) = [float(x) for x in val.split()]
141 elif field == "ylo yhi":
142 (ylo, yhi) = [float(x) for x in val.split()]
143 elif field == "zlo zhi":
144 (zlo, zhi) = [float(x) for x in val.split()]
145 elif field == "xy xz yz":
146 (xy, xz, yz) = [float(x) for x in val.split()]
148 if section is not None:
149 fields = line.split()
150 if section == "Atoms": # id *
151 id = int(fields[0])
152 if style == "full" and (len(fields) == 7 or len(fields) == 10):
153 # id mol-id type q x y z [tx ty tz]
154 pos_in[id] = (
155 int(fields[2]),
156 float(fields[4]),
157 float(fields[5]),
158 float(fields[6]),
159 )
160 mol_id_in[id] = int(fields[1])
161 charge_in[id] = float(fields[3])
162 if len(fields) == 10:
163 travel_in[id] = (
164 int(fields[7]),
165 int(fields[8]),
166 int(fields[9]),
167 )
168 elif style == "atomic" and (
169 len(fields) == 5 or len(fields) == 8
170 ):
171 # id type x y z [tx ty tz]
172 pos_in[id] = (
173 int(fields[1]),
174 float(fields[2]),
175 float(fields[3]),
176 float(fields[4]),
177 )
178 if len(fields) == 8:
179 travel_in[id] = (
180 int(fields[5]),
181 int(fields[6]),
182 int(fields[7]),
183 )
184 elif (style in ("angle", "bond", "molecular")
185 ) and (len(fields) == 6 or len(fields) == 9):
186 # id mol-id type x y z [tx ty tz]
187 pos_in[id] = (
188 int(fields[2]),
189 float(fields[3]),
190 float(fields[4]),
191 float(fields[5]),
192 )
193 mol_id_in[id] = int(fields[1])
194 if len(fields) == 9:
195 travel_in[id] = (
196 int(fields[6]),
197 int(fields[7]),
198 int(fields[8]),
199 )
200 elif (style == "charge"
201 and (len(fields) == 6 or len(fields) == 9)):
202 # id type q x y z [tx ty tz]
203 pos_in[id] = (
204 int(fields[1]),
205 float(fields[3]),
206 float(fields[4]),
207 float(fields[5]),
208 )
209 charge_in[id] = float(fields[2])
210 if len(fields) == 9:
211 travel_in[id] = (
212 int(fields[6]),
213 int(fields[7]),
214 int(fields[8]),
215 )
216 else:
217 raise RuntimeError(
218 "Style '{}' not supported or invalid "
219 "number of fields {}"
220 "".format(style, len(fields))
221 )
222 elif section == "Velocities": # id vx vy vz
223 vel_in[int(fields[0])] = (
224 float(fields[1]),
225 float(fields[2]),
226 float(fields[3]),
227 )
228 elif section == "Masses":
229 mass_in[int(fields[0])] = float(fields[1])
230 elif section == "Bonds": # id type atom1 atom2
231 bonds_in.append(
232 (int(fields[1]), int(fields[2]), int(fields[3]))
233 )
234 elif section == "Angles": # id type atom1 atom2 atom3
235 angles_in.append(
236 (
237 int(fields[1]),
238 int(fields[2]),
239 int(fields[3]),
240 int(fields[4]),
241 )
242 )
243 elif section == "Dihedrals": # id type atom1 atom2 atom3 atom4
244 dihedrals_in.append(
245 (
246 int(fields[1]),
247 int(fields[2]),
248 int(fields[3]),
249 int(fields[4]),
250 int(fields[5]),
251 )
252 )
254 # set cell
255 cell = np.zeros((3, 3))
256 cell[0, 0] = xhi - xlo
257 cell[1, 1] = yhi - ylo
258 cell[2, 2] = zhi - zlo
259 if xy is not None:
260 cell[1, 0] = xy
261 if xz is not None:
262 cell[2, 0] = xz
263 if yz is not None:
264 cell[2, 1] = yz
266 # initialize arrays for per-atom quantities
267 positions = np.zeros((N, 3))
268 numbers = np.zeros((N), int)
269 ids = np.zeros((N), int)
270 types = np.zeros((N), int)
271 if len(vel_in) > 0:
272 velocities = np.zeros((N, 3))
273 else:
274 velocities = None
275 if len(mass_in) > 0:
276 masses = np.zeros((N))
277 else:
278 masses = None
279 if len(mol_id_in) > 0:
280 mol_id = np.zeros((N), int)
281 else:
282 mol_id = None
283 if len(charge_in) > 0:
284 charge = np.zeros((N), float)
285 else:
286 charge = None
287 if len(travel_in) > 0:
288 travel = np.zeros((N, 3), int)
289 else:
290 travel = None
291 if len(bonds_in) > 0:
292 bonds = [""] * N
293 else:
294 bonds = None
295 if len(angles_in) > 0:
296 angles = [""] * N
297 else:
298 angles = None
299 if len(dihedrals_in) > 0:
300 dihedrals = [""] * N
301 else:
302 dihedrals = None
304 ind_of_id = {}
305 # copy per-atom quantities from read-in values
306 for (i, id) in enumerate(pos_in.keys()):
307 # by id
308 if sort_by_id:
309 ind = id - 1
310 else:
311 ind = i
312 ind_of_id[id] = ind
313 type = pos_in[id][0]
314 positions[ind, :] = [pos_in[id][1], pos_in[id][2], pos_in[id][3]]
315 if velocities is not None:
316 velocities[ind, :] = [vel_in[id][0], vel_in[id][1], vel_in[id][2]]
317 if travel is not None:
318 travel[ind] = travel_in[id]
319 if mol_id is not None:
320 mol_id[ind] = mol_id_in[id]
321 if charge is not None:
322 charge[ind] = charge_in[id]
323 ids[ind] = id
324 # by type
325 types[ind] = type
326 if Z_of_type is None:
327 numbers[ind] = type
328 else:
329 numbers[ind] = Z_of_type[type]
330 if masses is not None:
331 masses[ind] = mass_in[type]
332 # convert units
333 positions = convert(positions, "distance", units, "ASE")
334 cell = convert(cell, "distance", units, "ASE")
335 if masses is not None:
336 masses = convert(masses, "mass", units, "ASE")
337 if velocities is not None:
338 velocities = convert(velocities, "velocity", units, "ASE")
340 # create ase.Atoms
341 at = Atoms(
342 positions=positions,
343 numbers=numbers,
344 masses=masses,
345 cell=cell,
346 pbc=[True, True, True],
347 )
348 # set velocities (can't do it via constructor)
349 if velocities is not None:
350 at.set_velocities(velocities)
351 at.arrays["id"] = ids
352 at.arrays["type"] = types
353 if travel is not None:
354 at.arrays["travel"] = travel
355 if mol_id is not None:
356 at.arrays["mol-id"] = mol_id
357 if charge is not None:
358 at.arrays["initial_charges"] = charge
359 at.arrays["mmcharges"] = charge.copy()
361 if bonds is not None:
362 for (type, a1, a2) in bonds_in:
363 i_a1 = ind_of_id[a1]
364 i_a2 = ind_of_id[a2]
365 if len(bonds[i_a1]) > 0:
366 bonds[i_a1] += ","
367 bonds[i_a1] += "%d(%d)" % (i_a2, type)
368 for i in range(len(bonds)):
369 if len(bonds[i]) == 0:
370 bonds[i] = "_"
371 at.arrays["bonds"] = np.array(bonds)
373 if angles is not None:
374 for (type, a1, a2, a3) in angles_in:
375 i_a1 = ind_of_id[a1]
376 i_a2 = ind_of_id[a2]
377 i_a3 = ind_of_id[a3]
378 if len(angles[i_a2]) > 0:
379 angles[i_a2] += ","
380 angles[i_a2] += "%d-%d(%d)" % (i_a1, i_a3, type)
381 for i in range(len(angles)):
382 if len(angles[i]) == 0:
383 angles[i] = "_"
384 at.arrays["angles"] = np.array(angles)
386 if dihedrals is not None:
387 for (type, a1, a2, a3, a4) in dihedrals_in:
388 i_a1 = ind_of_id[a1]
389 i_a2 = ind_of_id[a2]
390 i_a3 = ind_of_id[a3]
391 i_a4 = ind_of_id[a4]
392 if len(dihedrals[i_a1]) > 0:
393 dihedrals[i_a1] += ","
394 dihedrals[i_a1] += "%d-%d-%d(%d)" % (i_a2, i_a3, i_a4, type)
395 for i in range(len(dihedrals)):
396 if len(dihedrals[i]) == 0:
397 dihedrals[i] = "_"
398 at.arrays["dihedrals"] = np.array(dihedrals)
400 at.info["comment"] = comment
402 return at
405@writer
406def write_lammps_data(
407 fd,
408 atoms: Atoms,
409 *,
410 specorder: list = None,
411 force_skew: bool = False,
412 prismobj: Prism = None,
413 masses: bool = False,
414 velocities: bool = False,
415 units: str = "metal",
416 atom_style: str = "atomic",
417):
418 """Write atomic structure data to a LAMMPS data file.
420 Parameters
421 ----------
422 fd : file|str
423 File to which the output will be written.
424 atoms : Atoms
425 Atoms to be written.
426 specorder : list[str], optional
427 Chemical symbols in the order of LAMMPS atom types, by default None
428 force_skew : bool, optional
429 Force to write the cell as a
430 `triclinic <https://docs.lammps.org/Howto_triclinic.html>`__ box,
431 by default False
432 prismobj : Prism|None, optional
433 Prism, by default None
434 masses : bool, optional
435 Whether the atomic masses are written or not, by default False
436 velocities : bool, optional
437 Whether the atomic velocities are written or not, by default False
438 units : str, optional
439 `LAMMPS units <https://docs.lammps.org/units.html>`__,
440 by default "metal"
441 atom_style : {"atomic", "charge", "full"}, optional
442 `LAMMPS atom style <https://docs.lammps.org/atom_style.html>`__,
443 by default "atomic"
444 """
446 # FIXME: We should add a check here that the encoding of the file object
447 # is actually ascii once the 'encoding' attribute of IOFormat objects
448 # starts functioning in implementation (currently it doesn't do
449 # anything).
451 if isinstance(atoms, list):
452 if len(atoms) > 1:
453 raise ValueError(
454 "Can only write one configuration to a lammps data file!"
455 )
456 atoms = atoms[0]
458 fd.write("(written by ASE)\n\n")
460 symbols = atoms.get_chemical_symbols()
461 n_atoms = len(symbols)
462 fd.write(f"{n_atoms} atoms\n")
464 if specorder is None:
465 # This way it is assured that LAMMPS atom types are always
466 # assigned predictably according to the alphabetic order
467 species = sorted(set(symbols))
468 else:
469 # To index elements in the LAMMPS data file
470 # (indices must correspond to order in the potential file)
471 species = specorder
472 n_atom_types = len(species)
473 fd.write(f"{n_atom_types} atom types\n\n")
475 if prismobj is None:
476 p = Prism(atoms.get_cell())
477 else:
478 p = prismobj
480 # Get cell parameters and convert from ASE units to LAMMPS units
481 xhi, yhi, zhi, xy, xz, yz = convert(p.get_lammps_prism(), "distance",
482 "ASE", units)
484 fd.write(f"0.0 {xhi:23.17g} xlo xhi\n")
485 fd.write(f"0.0 {yhi:23.17g} ylo yhi\n")
486 fd.write(f"0.0 {zhi:23.17g} zlo zhi\n")
488 if force_skew or p.is_skewed():
489 fd.write(f"{xy:23.17g} {xz:23.17g} {yz:23.17g} xy xz yz\n")
490 fd.write("\n")
492 if masses:
493 _write_masses(fd, atoms, species, units)
495 # Write (unwrapped) atomic positions. If wrapping of atoms back into the
496 # cell along periodic directions is desired, this should be done manually
497 # on the Atoms object itself beforehand.
498 fd.write(f"Atoms # {atom_style}\n\n")
499 pos = p.vector_to_lammps(atoms.get_positions(), wrap=False)
501 if atom_style == 'atomic':
502 for i, r in enumerate(pos):
503 # Convert position from ASE units to LAMMPS units
504 r = convert(r, "distance", "ASE", units)
505 s = species.index(symbols[i]) + 1
506 fd.write(
507 "{0:>6} {1:>3} {2:23.17g} {3:23.17g} {4:23.17g}\n".format(
508 *(i + 1, s) + tuple(r)
509 )
510 )
511 elif atom_style == 'charge':
512 charges = atoms.get_initial_charges()
513 for i, (q, r) in enumerate(zip(charges, pos)):
514 # Convert position and charge from ASE units to LAMMPS units
515 r = convert(r, "distance", "ASE", units)
516 q = convert(q, "charge", "ASE", units)
517 s = species.index(symbols[i]) + 1
518 fd.write("{0:>6} {1:>3} {2:>5} {3:23.17g} {4:23.17g} {5:23.17g}\n"
519 .format(*(i + 1, s, q) + tuple(r)))
520 elif atom_style == 'full':
521 charges = atoms.get_initial_charges()
522 # The label 'mol-id' has apparenlty been introduced in read earlier,
523 # but so far not implemented here. Wouldn't a 'underscored' label
524 # be better, i.e. 'mol_id' or 'molecule_id'?
525 if atoms.has('mol-id'):
526 molecules = atoms.get_array('mol-id')
527 if not np.issubdtype(molecules.dtype, np.integer):
528 raise TypeError((
529 "If 'atoms' object has 'mol-id' array, then"
530 " mol-id dtype must be subtype of np.integer, and"
531 " not {:s}.").format(str(molecules.dtype)))
532 if (len(molecules) != len(atoms)) or (molecules.ndim != 1):
533 raise TypeError((
534 "If 'atoms' object has 'mol-id' array, then"
535 " each atom must have exactly one mol-id."))
536 else:
537 # Assigning each atom to a distinct molecule id would seem
538 # preferableabove assigning all atoms to a single molecule
539 # id per default, as done within ase <= v 3.19.1. I.e.,
540 # molecules = np.arange(start=1, stop=len(atoms)+1,
541 # step=1, dtype=int) However, according to LAMMPS default
542 # behavior,
543 molecules = np.zeros(len(atoms), dtype=int)
544 # which is what happens if one creates new atoms within LAMMPS
545 # without explicitly taking care of the molecule id.
546 # Quote from docs at https://lammps.sandia.gov/doc/read_data.html:
547 # The molecule ID is a 2nd identifier attached to an atom.
548 # Normally, it is a number from 1 to N, identifying which
549 # molecule the atom belongs to. It can be 0 if it is a
550 # non-bonded atom or if you don't care to keep track of molecule
551 # assignments.
553 for i, (m, q, r) in enumerate(zip(molecules, charges, pos)):
554 # Convert position and charge from ASE units to LAMMPS units
555 r = convert(r, "distance", "ASE", units)
556 q = convert(q, "charge", "ASE", units)
557 s = species.index(symbols[i]) + 1
558 fd.write("{0:>6} {1:>3} {2:>3} {3:>5} {4:23.17g} {5:23.17g} "
559 "{6:23.17g}\n".format(*(i + 1, m, s, q) + tuple(r)))
560 else:
561 raise NotImplementedError
563 if velocities and atoms.get_velocities() is not None:
564 fd.write("\n\nVelocities\n\n")
565 vel = p.vector_to_lammps(atoms.get_velocities())
566 for i, v in enumerate(vel):
567 # Convert velocity from ASE units to LAMMPS units
568 v = convert(v, "velocity", "ASE", units)
569 fd.write(
570 "{0:>6} {1:23.17g} {2:23.17g} {3:23.17g}\n".format(
571 *(i + 1,) + tuple(v)
572 )
573 )
575 fd.flush()
578def _write_masses(fd, atoms: Atoms, species: list, units: str):
579 symbols_indices = atoms.symbols.indices()
580 fd.write("Masses\n\n")
581 for i, s in enumerate(species):
582 # Skip if the system does not contain the element `s`.
583 if s not in symbols_indices:
584 continue
585 # Find the first atom of the element `s` and extract its mass
586 # Cover by `float` to make a new object for safety
587 mass = float(atoms[symbols_indices[s][0]].mass)
588 # Convert mass from ASE units to LAMMPS units
589 mass = convert(mass, "mass", "ASE", units)
590 atom_type = i + 1
591 fd.write(f"{atom_type:>6} {mass:23.17g} # {s}\n")
592 fd.write("\n")