Coverage for /builds/debichem-team/python-ase/ase/io/lammpsdata.py : 82.73%
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=False, 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 ind_of_id[id] = i
309 if sort_by_id:
310 ind = id - 1
311 else:
312 ind = i
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(fd, atoms, specorder=None, force_skew=False,
407 prismobj=None, velocities=False, units="metal",
408 atom_style='atomic'):
409 """Write atomic structure data to a LAMMPS data file."""
411 # FIXME: We should add a check here that the encoding of the file object
412 # is actually ascii once the 'encoding' attribute of IOFormat objects
413 # starts functioning in implementation (currently it doesn't do
414 # anything).
416 if isinstance(atoms, list):
417 if len(atoms) > 1:
418 raise ValueError(
419 "Can only write one configuration to a lammps data file!"
420 )
421 atoms = atoms[0]
423 if hasattr(fd, "name"):
424 fd.write("{0} (written by ASE) \n\n".format(fd.name))
425 else:
426 fd.write("(written by ASE) \n\n")
428 symbols = atoms.get_chemical_symbols()
429 n_atoms = len(symbols)
430 fd.write("{0} \t atoms \n".format(n_atoms))
432 if specorder is None:
433 # This way it is assured that LAMMPS atom types are always
434 # assigned predictably according to the alphabetic order
435 species = sorted(set(symbols))
436 else:
437 # To index elements in the LAMMPS data file
438 # (indices must correspond to order in the potential file)
439 species = specorder
440 n_atom_types = len(species)
441 fd.write("{0} atom types\n".format(n_atom_types))
443 if prismobj is None:
444 p = Prism(atoms.get_cell())
445 else:
446 p = prismobj
448 # Get cell parameters and convert from ASE units to LAMMPS units
449 xhi, yhi, zhi, xy, xz, yz = convert(p.get_lammps_prism(), "distance",
450 "ASE", units)
452 fd.write("0.0 {0:23.17g} xlo xhi\n".format(xhi))
453 fd.write("0.0 {0:23.17g} ylo yhi\n".format(yhi))
454 fd.write("0.0 {0:23.17g} zlo zhi\n".format(zhi))
456 if force_skew or p.is_skewed():
457 fd.write(
458 "{0:23.17g} {1:23.17g} {2:23.17g} xy xz yz\n".format(
459 xy, xz, yz
460 )
461 )
462 fd.write("\n\n")
464 # Write (unwrapped) atomic positions. If wrapping of atoms back into the
465 # cell along periodic directions is desired, this should be done manually
466 # on the Atoms object itself beforehand.
467 fd.write("Atoms \n\n")
468 pos = p.vector_to_lammps(atoms.get_positions(), wrap=False)
470 if atom_style == 'atomic':
471 for i, r in enumerate(pos):
472 # Convert position from ASE units to LAMMPS units
473 r = convert(r, "distance", "ASE", units)
474 s = species.index(symbols[i]) + 1
475 fd.write(
476 "{0:>6} {1:>3} {2:23.17g} {3:23.17g} {4:23.17g}\n".format(
477 *(i + 1, s) + tuple(r)
478 )
479 )
480 elif atom_style == 'charge':
481 charges = atoms.get_initial_charges()
482 for i, (q, r) in enumerate(zip(charges, pos)):
483 # Convert position and charge from ASE units to LAMMPS units
484 r = convert(r, "distance", "ASE", units)
485 q = convert(q, "charge", "ASE", units)
486 s = species.index(symbols[i]) + 1
487 fd.write("{0:>6} {1:>3} {2:>5} {3:23.17g} {4:23.17g} {5:23.17g}\n"
488 .format(*(i + 1, s, q) + tuple(r)))
489 elif atom_style == 'full':
490 charges = atoms.get_initial_charges()
491 # The label 'mol-id' has apparenlty been introduced in read earlier,
492 # but so far not implemented here. Wouldn't a 'underscored' label
493 # be better, i.e. 'mol_id' or 'molecule_id'?
494 if atoms.has('mol-id'):
495 molecules = atoms.get_array('mol-id')
496 if not np.issubdtype(molecules.dtype, np.integer):
497 raise TypeError((
498 "If 'atoms' object has 'mol-id' array, then"
499 " mol-id dtype must be subtype of np.integer, and"
500 " not {:s}.").format(str(molecules.dtype)))
501 if (len(molecules) != len(atoms)) or (molecules.ndim != 1):
502 raise TypeError((
503 "If 'atoms' object has 'mol-id' array, then"
504 " each atom must have exactly one mol-id."))
505 else:
506 # Assigning each atom to a distinct molecule id would seem
507 # preferableabove assigning all atoms to a single molecule id per
508 # default, as done within ase <= v 3.19.1. I.e.,
509 # molecules = np.arange(start=1, stop=len(atoms)+1, step=1, dtype=int)
510 # However, according to LAMMPS default behavior,
511 molecules = np.zeros(len(atoms), dtype=int)
512 # which is what happens if one creates new atoms within LAMMPS
513 # without explicitly taking care of the molecule id.
514 # Quote from docs at https://lammps.sandia.gov/doc/read_data.html:
515 # The molecule ID is a 2nd identifier attached to an atom.
516 # Normally, it is a number from 1 to N, identifying which
517 # molecule the atom belongs to. It can be 0 if it is a
518 # non-bonded atom or if you don't care to keep track of molecule
519 # assignments.
521 for i, (m, q, r) in enumerate(zip(molecules, charges, pos)):
522 # Convert position and charge from ASE units to LAMMPS units
523 r = convert(r, "distance", "ASE", units)
524 q = convert(q, "charge", "ASE", units)
525 s = species.index(symbols[i]) + 1
526 fd.write("{0:>6} {1:>3} {2:>3} {3:>5} {4:23.17g} {5:23.17g} "
527 "{6:23.17g}\n".format(*(i + 1, m, s, q) + tuple(r)))
528 else:
529 raise NotImplementedError
531 if velocities and atoms.get_velocities() is not None:
532 fd.write("\n\nVelocities \n\n")
533 vel = p.vector_to_lammps(atoms.get_velocities())
534 for i, v in enumerate(vel):
535 # Convert velocity from ASE units to LAMMPS units
536 v = convert(v, "velocity", "ASE", units)
537 fd.write(
538 "{0:>6} {1:23.17g} {2:23.17g} {3:23.17g}\n".format(
539 *(i + 1,) + tuple(v)
540 )
541 )
543 fd.flush()