Coverage for ase / io / eon.py: 95.41%
109 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +0000
1# fmt: off
3# Copyright (C) 2012-2023, Jesper Friis, SINTEF
4# Copyright (C) 2024, Rohit Goswami, UI
5# (see accompanying license files for ASE).
6"""Module to read and write atoms EON reactant.con files.
8See http://theory.cm.utexas.edu/eon/index.html for a description of EON.
9"""
10from dataclasses import dataclass
11from typing import List, Tuple
12from warnings import warn
14import numpy as np
16from ase.atoms import Atoms
17from ase.constraints import FixAtoms
18from ase.geometry import cell_to_cellpar, cellpar_to_cell
19from ase.utils import reader, writer
21ISOTOPE_TOL = 1e-8
24@dataclass
25class EONHeader:
26 """
27 A data class for storing header information of an EON file.
29 Attributes
30 ----------
31 header_lines : List[str]
32 The first two comment lines from the EON file header.
33 cell_lengths : Tuple[float, float, float]
34 The lengths of the cell vectors.
35 cell_angles : Tuple[float, float, float]
36 The angles between the cell vectors.
37 Ncomponent : int
38 The number of distinct atom types.
39 component_counts : List[int]
40 A list containing the count of atoms for each type.
41 masses : List[float]
42 A list containing the atomic masses for each type.
43 """
45 header_lines: List[str]
46 # Actually these are float float float but.. mypy complains
47 cell_lengths: Tuple[float, ...]
48 cell_angles: Tuple[float, ...]
49 Ncomponent: int
50 component_counts: List[int]
51 masses: List[float]
54def process_header(lines: List[str]) -> EONHeader:
55 """
56 Processes the header lines from an EON file and returns an EONHeader object.
58 This function parses the first 9 lines of an EON file, extracting
59 information about the simulation cell, the number of atom types, their
60 counts, and masses, and encapsulates this information in an EONHeader
61 object.
63 Parameters
64 ----------
65 lines : List[str]
66 The first 9 lines of an EON file containing header information.
68 Returns
69 -------
70 EONHeader
71 An object containing the parsed header information.
72 """
73 header_lines = lines[:2]
75 # Parse cell lengths and angles
76 cell_lengths = tuple(map(float, lines[2].split()))
77 cell_angles = tuple(map(float, lines[3].split()))
78 if len(cell_lengths) != 3 or len(cell_angles) != 3:
79 raise ValueError(
80 "Cell lengths and angles must each contain exactly three values."
81 )
83 Ncomponent = int(lines[6])
84 component_counts = list(map(int, lines[7].split()))
85 masses = list(map(float, lines[8].split()))
87 return EONHeader(
88 header_lines=header_lines,
89 cell_lengths=cell_lengths,
90 cell_angles=cell_angles,
91 Ncomponent=Ncomponent,
92 component_counts=component_counts,
93 masses=masses,
94 )
97def make_atoms(coordblock, header):
98 """
99 Creates an Atoms object from coordinate blocks and header information.
101 This function takes a list of coordinate blocks and the associated header
102 information, constructs the cell, sets the atomic positions, masses, and
103 optionally applies FixAtoms constraints based on the header information, and
104 returns an ASE Atoms object.
106 Parameters
107 ----------
108 coordblock : list of str
109 The lines from an EON file representing atom coordinates and types.
110 header : EONHeader
111 The parsed header information.
113 Returns
114 -------
115 Atoms
116 An ASE Atoms object constructed from the given coordinate blocks and
117 header.
118 """
119 symbols = []
120 coords = []
121 masses = []
122 fixed = []
123 # Ordering follows the EON convention
124 # (alpha, beta, gamma)
125 cell_angles = (
126 header.cell_angles[0],
127 header.cell_angles[1],
128 header.cell_angles[2]
129 )
130 cellpar = [x for x in header.cell_lengths + cell_angles]
131 for idx, nblock in enumerate(header.component_counts):
132 elem_block = coordblock[:(nblock + 2)]
133 symb = elem_block[0]
134 symbols.extend(nblock * [symb])
135 mass = header.masses[idx]
136 masses.extend(nblock * [mass])
137 for eline in elem_block[2:]:
138 tokens = eline.split()
139 coords.append([float(x) for x in tokens[:3]])
140 fixed.append(bool(int(tokens[3])))
141 coordblock = coordblock[(nblock + 2):]
142 return Atoms(
143 symbols=symbols,
144 positions=coords,
145 masses=masses,
146 cell=cellpar_to_cell(cellpar),
147 constraint=FixAtoms(mask=fixed),
148 )
151@reader
152def read_eon(fileobj, index=-1):
153 """
154 Reads an EON file or directory and returns one or more Atoms objects.
156 This function handles single EON files, in both single image and multi-image
157 variants. It returns either a single Atoms object, a list of Atoms objects,
158 or a specific Atoms object indexed from the file or directory.
160 Parameters
161 ----------
162 fileobj : str or Path or file-like object
163 The path to the EON file or directory, or an open file-like object.
164 index : int, optional
165 The index of the Atoms object to return. If -1 (default), returns all
166 objects or a single object if only one is present.
168 Returns
169 -------
170 Atoms or List[Atoms]
171 Depending on the `index` parameter and the content of the fileobj,
172 returns either a single Atoms object or a list of Atoms objects.
173 """
174 images = []
175 while True:
176 # Read and process headers if they exist
177 try:
178 lines = [next(fileobj).strip() for _ in range(9)]
179 except StopIteration:
180 break # End of file
181 header = process_header(lines)
182 num_blocklines = (header.Ncomponent * 2) + sum(header.component_counts)
183 coordblocks = [next(fileobj).strip() for _ in range(num_blocklines)]
184 atoms = make_atoms(coordblocks, header)
185 images.append(atoms)
187 # XXX: I really don't like this since there should be a consistent return
188 if index == -1:
189 return images if len(images) > 1 else images[0]
190 else:
191 return images[index]
194@writer
195def write_eon(fileobj, images, comment="Generated by ASE"):
196 """
197 Writes structures to an EON trajectory file, allowing for multiple
198 snapshots.
200 This function iterates over all provided images, converting each one into a
201 text format compatible with EON trajectory files. It handles the conversion
202 of the cell parameters, chemical symbols, atomic masses, and atomic
203 constraints. Only FixAtoms constraints are supported; any other constraints
204 will generate a warning and be skipped. Arbitrary masses will raise an
205 error, since EON will not accept them, i.e. no Duterium and Hydrogen.
207 Parameters
208 ----------
209 fileobj : file object
210 An opened, writable file or file-like object to which the EON trajectory
211 information will be written.
212 images : list of Atoms
213 A list of ASE Atoms objects representing the images (atomic structures)
214 to be written into the EON trajectory file. Each Atoms object should
215 have a cell attribute and may have a constraints attribute. If the
216 constraints attribute is not a FixAtoms object, a warning will be
217 issued.
218 comment : str
219 An additional comment statement which will be written out as the first
220 header
222 Raises
223 ------
224 Warning
225 If any constraint in an Atoms object is not an instance of FixAtoms.
226 RuntimeError
227 If the masses for the species are not the default ones, i.e. if isotopes
228 are present
230 Returns
231 -------
232 None
233 The function writes directly to the provided file object.
235 See Also
236 --------
237 ase.io.utils.segment_list : for segmenting the list of images.
239 Examples
240 --------
241 >>> from ase.io import Trajectory
242 >>> from ase.io.utils import segment_list
243 >>> traj = Trajectory("neb.traj")
244 >>> n_images = 10 # Segment size, i.e. number of images in the NEB
245 >>> for idx, pth in enumerate(segment_list(traj, n_images)):
246 ... with open(f"outputs/neb_path_{idx:03d}.con", "w") as fobj:
247 ... write_eon_traj(fobj, pth)
248 """
250 for idx, atoms in enumerate(images):
251 out = []
252 if idx == 0:
253 out.append(comment)
254 else:
255 out.append(f"\n{comment}")
256 out.append("preBox_header_2") # ??
258 a, b, c, alpha, beta, gamma = cell_to_cellpar(atoms.cell)
259 out.append("%-10.6f %-10.6f %-10.6f" % (a, b, c))
260 out.append("%-10.6f %-10.6f %-10.6f" % (alpha, beta, gamma))
262 out.append("postBox_header_1") # ??
263 out.append("postBox_header_2") # ??
265 symbol_indices = atoms.symbols.indices()
266 natoms = [len(x) for x in symbol_indices.values()]
267 masslist = atoms.get_masses()
268 species_masses = []
269 for symbol, indices in symbol_indices.items():
270 masses_for_symb = masslist[indices]
271 if np.ptp(masses_for_symb) > ISOTOPE_TOL:
272 raise RuntimeError(
273 "Isotopes cannot be handled by EON"
274 ", ensure uniform masses for symbols"
275 )
276 species_masses.append(masses_for_symb[0])
278 out.append(str(len(symbol_indices)))
279 out.append(" ".join([str(n) for n in natoms]))
280 out.append(" ".join([str(n) for n in species_masses]))
282 atom_id = 0
283 for cid, (species, indices) in enumerate(symbol_indices.items()):
284 fixed = np.array([False] * len(atoms))
285 out.append(species)
286 out.append("Coordinates of Component %d" % (cid + 1))
287 atom = atoms[indices]
288 coords = atom.positions
289 for constraint in atom.constraints:
290 if not isinstance(constraint, FixAtoms):
291 warn(
292 "Only FixAtoms constraints are supported"
293 "by con files. Dropping %r",
294 constraint,
295 )
296 continue
297 if constraint.index.dtype.kind == "b":
298 fixed = np.array(constraint.index, dtype=int)
299 else:
300 fixed = np.zeros((natoms[cid],), dtype=int)
301 for i in constraint.index:
302 fixed[i] = 1
303 for xyz, fix in zip(coords, fixed):
304 line_fmt = "{:>22.17f} {:>22.17f} {:>22.17f} {:d} {:4d}"
305 line = line_fmt.format(*xyz, int(fix), atom_id)
306 out.append(line)
307 atom_id += 1
308 fileobj.write("\n".join(out))