Coverage for ase / io / lammpsdata.py: 94.75%
381 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +0000
1"""IO for LAMMPS data files."""
3from __future__ import annotations
5import re
6import warnings
7from dataclasses import dataclass, field
9import numpy as np
11from ase.atoms import Atoms
12from ase.calculators.lammps import Prism, convert
13from ase.data import atomic_masses, atomic_numbers
14from ase.utils import reader, writer
17def _make_cell(box):
18 cell = np.zeros((3, 3))
19 celldisp = np.zeros(3)
20 if 'avec' in box:
21 cell[0] = box['avec']
22 cell[1] = box['bvec']
23 cell[2] = box['cvec']
24 celldisp = box['abc origin']
25 else:
26 cell[0, 0] = box['xhi'] - box['xlo']
27 cell[1, 1] = box['yhi'] - box['ylo']
28 cell[2, 2] = box['zhi'] - box['zlo']
29 cell[1, 0] = box['xy']
30 cell[2, 0] = box['xz']
31 cell[2, 1] = box['yz']
32 return cell, celldisp
35@reader
36def read_lammps_data(
37 fileobj,
38 Z_of_type: dict = None,
39 sort_by_id: bool = True,
40 read_image_flags: bool = True,
41 units: str = 'metal',
42 atom_style: str = None,
43 style: str = None,
44):
45 """Method which reads a LAMMPS data file.
47 Parameters
48 ----------
49 fileobj : file | str
50 File from which data should be read.
51 Z_of_type : dict[int, int], optional
52 Mapping from LAMMPS atom types (typically starting from 1) to atomic
53 numbers. If None, if there is the "Masses" section, atomic numbers are
54 guessed from the atomic masses. Otherwise, atomic numbers of 1 (H), 2
55 (He), etc. are assigned to atom types of 1, 2, etc. Default is None.
56 sort_by_id : bool, optional
57 Order the particles according to their id. Might be faster to set it
58 False. Default is True.
59 read_image_flags: bool, default True
60 If True, the lattice translation vectors derived from image flags are
61 added to atomic positions.
62 units : str, optional
63 `LAMMPS units <https://docs.lammps.org/units.html>`__. Default is
64 'metal'.
65 atom_style : {'atomic', 'charge', 'full'} etc., optional
66 `LAMMPS atom style <https://docs.lammps.org/atom_style.html>`__.
67 If None, `atom_style` is guessed in the following priority (1) comment
68 after `Atoms` (2) length of fields (valid only `atomic` and `full`).
69 Default is None.
70 """
71 if style is not None:
72 warnings.warn(
73 FutureWarning('"style" is deprecated; please use "atom_style".'),
74 )
75 atom_style = style
76 # begin read_lammps_data
77 file_comment = next(fileobj).rstrip()
79 # default values (https://docs.lammps.org/read_data.html)
80 # in most cases these will be updated below
81 natoms = 0
82 # N_types = 0
83 box: dict[str, float | list[float]] = {
84 'xlo': -0.5,
85 'xhi': +0.5,
86 'ylo': -0.5,
87 'yhi': +0.5,
88 'zlo': -0.5,
89 'zhi': +0.5,
90 'xy': 0.0,
91 'xz': 0.0,
92 'yz': 0.0,
93 }
95 mass_in = {}
96 vel_in = {}
97 atom_type_labels = {}
98 bonds_in = []
99 angles_in = []
100 dihedrals_in = []
102 sections = [
103 'Atoms',
104 'Velocities',
105 'Masses',
106 'Charges',
107 'Ellipsoids',
108 'Lines',
109 'Triangles',
110 'Bodies',
111 'Bonds',
112 'Angles',
113 'Dihedrals',
114 'Impropers',
115 'Impropers Pair Coeffs',
116 'PairIJ Coeffs',
117 'Pair Coeffs',
118 'Bond Coeffs',
119 'Angle Coeffs',
120 'Dihedral Coeffs',
121 'Improper Coeffs',
122 'BondBond Coeffs',
123 'BondAngle Coeffs',
124 'MiddleBondTorsion Coeffs',
125 'EndBondTorsion Coeffs',
126 'AngleTorsion Coeffs',
127 'AngleAngleTorsion Coeffs',
128 'BondBond13 Coeffs',
129 'AngleAngle Coeffs',
130 'Atom Type Labels',
131 'Bond Type Labels',
132 'Angle Type Labels',
133 'Dihedral Type Labels',
134 'Improper Type Labels',
135 ]
136 header_fields = [
137 'atoms',
138 'bonds',
139 'angles',
140 'dihedrals',
141 'impropers',
142 'atom types',
143 'bond types',
144 'angle types',
145 'dihedral types',
146 'improper types',
147 'extra bond per atom',
148 'extra angle per atom',
149 'extra dihedral per atom',
150 'extra improper per atom',
151 'extra special per atom',
152 'ellipsoids',
153 'lines',
154 'triangles',
155 'bodies',
156 ]
157 header_fields_restricted_box = [
158 'xlo xhi',
159 'ylo yhi',
160 'zlo zhi',
161 'xy xz yz',
162 ]
163 header_fields_general_box = [
164 'avec',
165 'bvec',
166 'cvec',
167 'abc origin',
168 ]
169 header_fields += header_fields_restricted_box + header_fields_general_box
170 sections_re = '(' + '|'.join(sections).replace(' ', '\\s+') + ')'
171 header_fields_re = '(' + '|'.join(header_fields).replace(' ', '\\s+') + ')'
173 section = None
174 header = True
175 for line in fileobj:
176 # get string after #; if # does not exist, return ''
177 line_comment = re.sub(r'^.*#|^.*$', '', line).strip()
178 line = re.sub('#.*', '', line).rstrip().lstrip()
179 if re.match('^\\s*$', line): # skip blank lines
180 continue
182 # check for known section names
183 match = re.match(sections_re, line)
184 if match is not None:
185 section = match.group(0).rstrip().lstrip()
186 header = False
187 if section == 'Atoms': # id *
188 # guess `atom_style` from the comment after `Atoms` if exists
189 if atom_style is None and line_comment != '':
190 atom_style = line_comment
191 atoms_section = _read_atoms_section(fileobj, natoms, atom_style)
192 continue
194 if header:
195 field = None
196 val = None
197 match = re.match('(.*)\\s+' + header_fields_re, line)
198 if match is not None:
199 field = match.group(2).lstrip().rstrip()
200 val = match.group(1).lstrip().rstrip()
201 if field is not None and val is not None:
202 if field == 'atoms':
203 natoms = int(val)
204 elif field in header_fields_restricted_box:
205 keys = field.split()
206 values = (float(x) for x in val.split())
207 box.update(dict(zip(keys, values)))
208 elif field in header_fields_general_box:
209 box[field] = [float(x) for x in val.split()]
211 if section is not None:
212 fields = line.split()
213 if section == 'Velocities': # id vx vy vz
214 vel_in[int(fields[0])] = [float(fields[_]) for _ in (1, 2, 3)]
215 elif section == 'Masses':
216 mass_in[int(fields[0])] = float(fields[1])
217 elif section == 'Atom Type Labels':
218 atom_type_labels[int(fields[0])] = fields[1]
219 elif section == 'Bonds': # id type atom1 atom2
220 bonds_in.append([int(fields[_]) for _ in (1, 2, 3)])
221 elif section == 'Angles': # id type atom1 atom2 atom3
222 angles_in.append([int(fields[_]) for _ in (1, 2, 3, 4)])
223 elif section == 'Dihedrals': # id type atom1 atom2 atom3 atom4
224 dihedrals_in.append([int(fields[_]) for _ in (1, 2, 3, 4, 5)])
226 # set cell
227 cell, celldisp = _make_cell(box)
229 if sort_by_id:
230 atoms_section.sort()
232 ids = atoms_section.ids
234 if np.all(atoms_section.types != 0): # numeric
235 types = atoms_section.types
236 else: # labels
237 labels2types = {v: k for k, v in atom_type_labels.items()}
238 types = np.array([labels2types[_] for _ in atoms_section.labels])
240 if Z_of_type:
241 # The user-specified `Z_of_type` has the highest priority.
242 numbers = np.array([Z_of_type[_] for _ in types])
243 elif atom_type_labels and all(
244 v in atomic_numbers for v in atom_type_labels.values()
245 ):
246 # if all the labels in the `Atom Type Labels` section are element names
247 numbers = np.array([atomic_numbers[atom_type_labels[_]] for _ in types])
248 else:
249 numbers = types
251 masses = np.array([mass_in[_] for _ in types]) if mass_in else None
252 velocities = np.array([vel_in[_] for _ in ids]) if vel_in else None
254 # convert units
255 positions = convert(atoms_section.positions, 'distance', units, 'ASE')
256 cell = convert(cell, 'distance', units, 'ASE')
257 if masses is not None:
258 masses = convert(masses, 'mass', units, 'ASE')
259 if velocities is not None:
260 velocities = convert(velocities, 'velocity', units, 'ASE')
262 # guess atomic numbers from atomic masses
263 # this must be after the above mass-unit conversion
264 if Z_of_type is None and masses is not None:
265 numbers = _masses2numbers(masses)
267 # create ase.Atoms
268 atoms = Atoms(
269 positions=positions,
270 numbers=numbers,
271 masses=masses,
272 cell=cell,
273 pbc=[True, True, True],
274 celldisp=celldisp,
275 )
277 # add lattice translation vectors
278 if read_image_flags:
279 scaled_positions = atoms.get_scaled_positions(wrap=False)
280 atoms.set_scaled_positions(scaled_positions + atoms_section.cell_ids)
282 # set velocities (can't do it via constructor)
283 if velocities is not None:
284 atoms.set_velocities(velocities)
286 atoms.arrays['id'] = atoms_section.ids
287 atoms.arrays['type'] = atoms_section.types
288 if not np.all(atoms_section.mol_ids == 0):
289 atoms.arrays['mol-id'] = atoms_section.mol_ids
290 if not np.all(np.isnan(atoms_section.charges)):
291 atoms.arrays['initial_charges'] = atoms_section.charges
292 atoms.arrays['mmcharges'] = atoms_section.charges.copy()
294 # mapping from LAMMPS atom-IDs to ASE Atoms IDs
295 mapping = {atom_id: i for i, atom_id in enumerate(atoms_section.ids)}
297 if bonds_in:
298 key = 'bonds'
299 atoms.arrays[key] = _parse_bonds(bonds_in, natoms, mapping)
301 if angles_in:
302 key = 'angles'
303 atoms.arrays[key] = _parse_angles(angles_in, natoms, mapping)
305 if dihedrals_in:
306 key = 'dihedrals'
307 atoms.arrays[key] = _parse_dihedrals(dihedrals_in, natoms, mapping)
309 atoms.info['comment'] = file_comment
311 return atoms
314@dataclass
315class _AtomsSection:
316 natoms: int
317 ids: np.ndarray = field(init=False)
318 types: np.ndarray = field(init=False)
319 labels: np.ndarray = field(init=False)
320 positions: np.ndarray = field(init=False)
321 mol_ids: np.ndarray = field(init=False)
322 charges: np.ndarray = field(init=False)
323 cell_ids: np.ndarray = field(init=False)
325 def __post_init__(self):
326 self.ids = np.zeros(self.natoms, int)
327 self.types = np.zeros(self.natoms, int)
328 self.labels = np.zeros(self.natoms, object)
329 self.positions = np.zeros((self.natoms, 3), float)
330 self.mol_ids = np.zeros(self.natoms, int)
331 self.charges = np.full(self.natoms, np.nan, float)
332 self.cell_ids = np.zeros((self.natoms, 3), int)
334 def sort(self):
335 """Sort IDs."""
336 args = np.argsort(self.ids)
337 self.ids = self.ids[args]
338 self.types = self.types[args]
339 self.positions = self.positions[args]
340 self.mol_ids = self.mol_ids[args]
341 self.charges = self.charges[args]
342 self.cell_ids = self.cell_ids[args]
345def _read_atoms_section(fileobj, natoms: int, style: str = None):
346 data = _AtomsSection(natoms)
347 next(fileobj) # skip blank line just after `Atoms`
348 for i in range(natoms):
349 line = next(fileobj)
350 line = re.sub(r'#.*', '', line).rstrip().lstrip()
351 fields = line.split()
352 if style is None:
353 style = _guess_atom_style(fields)
354 data.ids[i] = int(fields[0])
355 if style == 'full' and len(fields) in (7, 10):
356 # id mol-id type q x y z [tx ty tz]
357 data.labels[i] = fields[2]
358 data.positions[i] = tuple(float(fields[_]) for _ in (4, 5, 6))
359 data.mol_ids[i] = int(fields[1])
360 data.charges[i] = float(fields[3])
361 if len(fields) == 10:
362 data.cell_ids[i] = tuple(int(fields[_]) for _ in (7, 8, 9))
363 elif style == 'atomic' and len(fields) in (5, 8):
364 # id type x y z [tx ty tz]
365 data.labels[i] = fields[1]
366 data.positions[i] = tuple(float(fields[_]) for _ in (2, 3, 4))
367 if len(fields) == 8:
368 data.cell_ids[i] = tuple(int(fields[_]) for _ in (5, 6, 7))
369 elif style in ('angle', 'bond', 'molecular') and len(fields) in (6, 9):
370 # id mol-id type x y z [tx ty tz]
371 data.labels[i] = fields[2]
372 data.positions[i] = tuple(float(fields[_]) for _ in (3, 4, 5))
373 data.mol_ids[i] = int(fields[1])
374 if len(fields) == 9:
375 data.cell_ids[i] = tuple(int(fields[_]) for _ in (6, 7, 8))
376 elif style == 'charge' and len(fields) in (6, 9):
377 # id type q x y z [tx ty tz]
378 data.labels[i] = fields[1]
379 data.positions[i] = tuple(float(fields[_]) for _ in (3, 4, 5))
380 data.charges[i] = float(fields[2])
381 if len(fields) == 9:
382 data.cell_ids[i] = tuple(int(fields[_]) for _ in (6, 7, 8))
383 else:
384 raise RuntimeError(
385 f'Style "{style}" not supported or invalid. '
386 f'Number of fields: {len(fields)}'
387 )
388 if all(_.isdigit() for _ in data.labels):
389 data.types = data.labels.astype(int)
390 return data
393def _guess_atom_style(fields):
394 """Guess `atom_sytle` from the length of fields."""
395 if len(fields) in (5, 8):
396 return 'atomic'
397 if len(fields) in (7, 10):
398 return 'full'
399 raise ValueError('atom_style cannot be guessed from len(fields)')
402def _masses2numbers(masses):
403 """Guess atomic numbers from atomic masses."""
404 return np.argmin(np.abs(atomic_masses - masses[:, None]), axis=1)
407def _parse_bonds(bonds_in, natoms: int, mapping: dict):
408 bonds = [''] * natoms
409 for bond_type, at1, at2 in bonds_in:
410 i_a1 = mapping[at1]
411 i_a2 = mapping[at2]
412 if len(bonds[i_a1]) > 0:
413 bonds[i_a1] += ','
414 bonds[i_a1] += f'{i_a2:d}({bond_type:d})'
415 for i, bond in enumerate(bonds):
416 if len(bond) == 0:
417 bonds[i] = '_'
418 return np.array(bonds)
421def _parse_angles(angles_in, natoms: int, mapping: dict):
422 angles = [''] * natoms
423 for angle_type, at1, at2, at3 in angles_in:
424 i_a1 = mapping[at1]
425 i_a2 = mapping[at2]
426 i_a3 = mapping[at3]
427 if len(angles[i_a2]) > 0:
428 angles[i_a2] += ','
429 angles[i_a2] += f'{i_a1:d}-{i_a3:d}({angle_type:d})'
430 for i, angle in enumerate(angles):
431 if len(angle) == 0:
432 angles[i] = '_'
433 return np.array(angles)
436def _parse_dihedrals(dihedrals_in, natoms: int, mapping: dict):
437 dihedrals = [''] * natoms
438 for dihedral_type, at1, at2, at3, at4 in dihedrals_in:
439 i_a1 = mapping[at1]
440 i_a2 = mapping[at2]
441 i_a3 = mapping[at3]
442 i_a4 = mapping[at4]
443 if len(dihedrals[i_a1]) > 0:
444 dihedrals[i_a1] += ','
445 dihedrals[i_a1] += f'{i_a2:d}-{i_a3:d}-{i_a4:d}({dihedral_type:d})'
446 for i, dihedral in enumerate(dihedrals):
447 if len(dihedral) == 0:
448 dihedrals[i] = '_'
449 return np.array(dihedrals)
452@writer
453def write_lammps_data(
454 fd,
455 atoms: Atoms,
456 *,
457 specorder: list = None,
458 reduce_cell: bool = False,
459 force_skew: bool = False,
460 prismobj: Prism = None,
461 write_image_flags: bool = False,
462 masses: bool = False,
463 velocities: bool = False,
464 atom_type_labels: bool = False,
465 units: str = 'metal',
466 bonds: bool = True,
467 atom_style: str = 'atomic',
468):
469 """Write atomic structure data to a LAMMPS data file.
471 Parameters
472 ----------
473 fd : file|str
474 File to which the output will be written.
475 atoms : Atoms
476 Atoms to be written.
477 specorder : list[str], optional
478 Chemical symbols in the order of LAMMPS atom types, by default None
479 force_skew : bool, optional
480 Force to write the cell as a
481 `triclinic <https://docs.lammps.org/Howto_triclinic.html>`__ box,
482 by default False
483 reduce_cell : bool, optional
484 Whether the cell shape is reduced or not, by default False
485 prismobj : Prism|None, optional
486 Prism, by default None
487 write_image_flags : bool, default False
488 If True, the image flags, i.e., in which images of the periodic
489 simulation box the atoms are, are written.
490 masses : bool, optional
491 Whether the atomic masses are written or not, by default False
492 velocities : bool, optional
493 Whether the atomic velocities are written or not, by default False
494 atom_type_labels : bool, optional
495 Whether the atom type labels are written or not, by default False
496 units : str, optional
497 `LAMMPS units <https://docs.lammps.org/units.html>`__,
498 by default 'metal'
499 bonds : bool, optional
500 Whether the bonds are written or not. Bonds can only be written
501 for atom_style='full', by default True
502 atom_style : {'atomic', 'charge', 'full'}, optional
503 `LAMMPS atom style <https://docs.lammps.org/atom_style.html>`__,
504 by default 'atomic'
505 """
506 # FIXME: We should add a check here that the encoding of the file object
507 # is actually ascii once the 'encoding' attribute of IOFormat objects
508 # starts functioning in implementation (currently it doesn't do
509 # anything).
511 if isinstance(atoms, list):
512 if len(atoms) > 1:
513 raise ValueError(
514 'Can only write one configuration to a lammps data file!'
515 )
516 atoms = atoms[0]
518 fd.write('(written by ASE)\n\n')
520 symbols = atoms.get_chemical_symbols()
521 n_atoms = len(symbols)
522 fd.write(f'{n_atoms} atoms\n')
524 if specorder is not None:
525 # To index elements in the LAMMPS data file
526 # (indices must correspond to order in the potential file)
527 species = specorder
528 elif 'type' in atoms.arrays:
529 species = _get_symbols_by_types(atoms)
530 else:
531 # This way it is assured that LAMMPS atom types are always
532 # assigned predictably according to the alphabetic order
533 species = sorted(set(symbols))
535 n_atom_types = len(species)
536 fd.write(f'{n_atom_types} atom types\n\n')
538 bonds_in = []
539 if (
540 bonds
541 and (atom_style == 'full')
542 and (atoms.arrays.get('bonds') is not None)
543 ):
544 n_bonds = 0
545 n_bond_types = 1
546 for i, bondsi in enumerate(atoms.arrays['bonds']):
547 if bondsi != '_':
548 for bond in bondsi.split(','):
549 dummy1, dummy2 = bond.split('(')
550 bond_type = int(dummy2.split(')')[0])
551 at1 = int(i) + 1
552 at2 = int(dummy1) + 1
553 bonds_in.append((bond_type, at1, at2))
554 n_bonds = n_bonds + 1
555 if bond_type > n_bond_types:
556 n_bond_types = bond_type
557 fd.write(f'{n_bonds} bonds\n')
558 fd.write(f'{n_bond_types} bond types\n\n')
560 if prismobj is None:
561 prismobj = Prism(atoms.get_cell(), reduce_cell=reduce_cell)
563 # Get cell parameters and convert from ASE units to LAMMPS units
564 xhi, yhi, zhi, xy, xz, yz = convert(
565 prismobj.get_lammps_prism(), 'distance', 'ASE', units
566 )
568 fd.write(f'0.0 {xhi:23.17g} xlo xhi\n')
569 fd.write(f'0.0 {yhi:23.17g} ylo yhi\n')
570 fd.write(f'0.0 {zhi:23.17g} zlo zhi\n')
572 if force_skew or prismobj.is_skewed():
573 fd.write(f'{xy:23.17g} {xz:23.17g} {yz:23.17g} xy xz yz\n')
575 if atom_type_labels:
576 _write_atom_type_labels(fd, species)
578 if masses:
579 _write_masses(fd, atoms, species, units)
581 # Write (unwrapped) atomic positions. If wrapping of atoms back into the
582 # cell along periodic directions is desired, this should be done manually
583 # on the Atoms object itself beforehand.
584 fd.write(f'\nAtoms # {atom_style}\n\n')
586 if write_image_flags:
587 scaled_positions = atoms.get_scaled_positions(wrap=False)
588 image_flags = np.floor(scaled_positions).astype(int)
590 # when `write_image_flags` is True, the positions are wrapped while the
591 # unwrapped positions can be recovered from the image flags
592 pos = prismobj.vector_to_lammps(
593 atoms.get_positions(),
594 wrap=write_image_flags,
595 )
597 types = _get_types(atoms, species)
599 if atom_style == 'atomic':
600 # Convert position from ASE units to LAMMPS units
601 pos = convert(pos, 'distance', 'ASE', units)
602 for i, r in enumerate(pos):
603 s = types[i]
604 line = (
605 f'{i + 1:>6} {s:>3} {r[0]:23.17g} {r[1]:23.17g} {r[2]:23.17g}'
606 )
607 if write_image_flags:
608 img = image_flags[i]
609 line += f' {img[0]:6d} {img[1]:6d} {img[2]:6d}'
610 line += '\n'
611 fd.write(line)
612 elif atom_style == 'charge':
613 charges = atoms.get_initial_charges()
614 # Convert position and charge from ASE units to LAMMPS units
615 pos = convert(pos, 'distance', 'ASE', units)
616 charges = convert(charges, 'charge', 'ASE', units)
617 for i, (q, r) in enumerate(zip(charges, pos)):
618 s = types[i]
619 line = (
620 f'{i + 1:>6} {s:>3} {q:>5}'
621 f' {r[0]:23.17g} {r[1]:23.17g} {r[2]:23.17g}'
622 )
623 if write_image_flags:
624 img = image_flags[i]
625 line += f' {img[0]:6d} {img[1]:6d} {img[2]:6d}'
626 line += '\n'
627 fd.write(line)
628 elif atom_style == 'full':
629 charges = atoms.get_initial_charges()
630 # The label 'mol-id' has apparenlty been introduced in read earlier,
631 # but so far not implemented here. Wouldn't a 'underscored' label
632 # be better, i.e. 'mol_id' or 'molecule_id'?
633 if atoms.has('mol-id'):
634 molecules = atoms.get_array('mol-id')
635 if not np.issubdtype(molecules.dtype, np.integer):
636 raise TypeError(
637 f'If "atoms" object has "mol-id" array, then '
638 f'mol-id dtype must be subtype of np.integer, and '
639 f'not {molecules.dtype!s:s}.'
640 )
641 if (len(molecules) != len(atoms)) or (molecules.ndim != 1):
642 raise TypeError(
643 'If "atoms" object has "mol-id" array, then '
644 'each atom must have exactly one mol-id.'
645 )
646 else:
647 # Assigning each atom to a distinct molecule id would seem
648 # preferableabove assigning all atoms to a single molecule
649 # id per default, as done within ase <= v 3.19.1. I.e.,
650 # molecules = np.arange(start=1, stop=len(atoms)+1,
651 # step=1, dtype=int) However, according to LAMMPS default
652 # behavior,
653 molecules = np.zeros(len(atoms), dtype=int)
654 # which is what happens if one creates new atoms within LAMMPS
655 # without explicitly taking care of the molecule id.
656 # Quote from docs at https://lammps.sandia.gov/doc/read_data.html:
657 # The molecule ID is a 2nd identifier attached to an atom.
658 # Normally, it is a number from 1 to N, identifying which
659 # molecule the atom belongs to. It can be 0 if it is a
660 # non-bonded atom or if you don't care to keep track of molecule
661 # assignments.
663 # Convert position and charge from ASE units to LAMMPS units
664 pos = convert(pos, 'distance', 'ASE', units)
665 charges = convert(charges, 'charge', 'ASE', units)
666 for i, (m, q, r) in enumerate(zip(molecules, charges, pos)):
667 s = types[i]
668 line = (
669 f'{i + 1:>6} {m:>3} {s:>3} {q:>5}'
670 f' {r[0]:23.17g} {r[1]:23.17g} {r[2]:23.17g}'
671 )
672 if write_image_flags:
673 img = image_flags[i]
674 line += f' {img[0]:6d} {img[1]:6d} {img[2]:6d}'
675 line += '\n'
676 fd.write(line)
677 if bonds and (atoms.arrays.get('bonds') is not None):
678 fd.write('\nBonds\n\n')
679 for i in range(n_bonds):
680 bond_type = bonds_in[i][0]
681 at1 = bonds_in[i][1]
682 at2 = bonds_in[i][2]
683 fd.write(f'{i + 1:>3} {bond_type:>3} {at1:>3} {at2:>3}\n')
684 else:
685 raise ValueError(atom_style)
687 if velocities and atoms.get_velocities() is not None:
688 fd.write('\nVelocities\n\n')
689 vel = prismobj.vector_to_lammps(atoms.get_velocities())
690 # Convert velocity from ASE units to LAMMPS units
691 vel = convert(vel, 'velocity', 'ASE', units)
692 for i, v in enumerate(vel):
693 fd.write(f'{i + 1:>6} {v[0]:23.17g} {v[1]:23.17g} {v[2]:23.17g}\n')
695 fd.flush()
698def _write_masses(fd, atoms: Atoms, species: list, units: str):
699 symbols_indices = atoms.symbols.indices()
700 fd.write('\nMasses\n\n')
701 for i, s in enumerate(species):
702 if s in symbols_indices:
703 # Find the first atom of the element `s` and extract its mass
704 # Cover by `float` to make a new object for safety
705 mass = float(atoms[symbols_indices[s][0]].mass)
706 else:
707 # Fetch from ASE data if the element `s` is not in the system
708 mass = atomic_masses[atomic_numbers[s]]
709 # Convert mass from ASE units to LAMMPS units
710 mass = convert(mass, 'mass', 'ASE', units)
711 atom_type = i + 1
712 fd.write(f'{atom_type} {mass:23.17g} # {s}\n')
715def _write_atom_type_labels(fd, species: list[str]):
716 fd.write('\nAtom Type Labels\n\n')
717 for i, s in enumerate(species):
718 atom_type = i + 1
719 fd.write(f'{atom_type} {s}\n')
722def _get_types(atoms: Atoms, species: list):
723 if 'type' in atoms.arrays:
724 return atoms.arrays['type']
725 symbols = atoms.get_chemical_symbols()
726 return [species.index(symbols[i]) + 1 for i in range(len(symbols))]
729def _get_symbols_by_types(atoms: Atoms) -> list[str]:
730 _, first_idx = np.unique(atoms.arrays['type'], return_index=True)
731 return [atoms.symbols[i] for i in first_idx]