Coverage for /builds/ase/ase/ase/io/eon.py: 95.41%
109 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +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 in EON is different from the ASE convention
124 cell_angles = (
125 header.cell_angles[2],
126 header.cell_angles[1],
127 header.cell_angles[0]
128 )
129 cellpar = [x for x in header.cell_lengths + cell_angles]
130 for idx, nblock in enumerate(header.component_counts):
131 elem_block = coordblock[:(nblock + 2)]
132 symb = elem_block[0]
133 symbols.extend(nblock * [symb])
134 mass = header.masses[idx]
135 masses.extend(nblock * [mass])
136 for eline in elem_block[2:]:
137 tokens = eline.split()
138 coords.append([float(x) for x in tokens[:3]])
139 fixed.append(bool(int(tokens[3])))
140 coordblock = coordblock[(nblock + 2):]
141 return Atoms(
142 symbols=symbols,
143 positions=coords,
144 masses=masses,
145 cell=cellpar_to_cell(cellpar),
146 constraint=FixAtoms(mask=fixed),
147 )
150@reader
151def read_eon(fileobj, index=-1):
152 """
153 Reads an EON file or directory and returns one or more Atoms objects.
155 This function handles single EON files, in both single image and multi-image
156 variants. It returns either a single Atoms object, a list of Atoms objects,
157 or a specific Atoms object indexed from the file or directory.
159 Parameters
160 ----------
161 fileobj : str or Path or file-like object
162 The path to the EON file or directory, or an open file-like object.
163 index : int, optional
164 The index of the Atoms object to return. If -1 (default), returns all
165 objects or a single object if only one is present.
167 Returns
168 -------
169 Atoms or List[Atoms]
170 Depending on the `index` parameter and the content of the fileobj,
171 returns either a single Atoms object or a list of Atoms objects.
172 """
173 images = []
174 while True:
175 # Read and process headers if they exist
176 try:
177 lines = [next(fileobj).strip() for _ in range(9)]
178 except StopIteration:
179 break # End of file
180 header = process_header(lines)
181 num_blocklines = (header.Ncomponent * 2) + sum(header.component_counts)
182 coordblocks = [next(fileobj).strip() for _ in range(num_blocklines)]
183 atoms = make_atoms(coordblocks, header)
184 images.append(atoms)
186 # XXX: I really don't like this since there should be a consistent return
187 if index == -1:
188 return images if len(images) > 1 else images[0]
189 else:
190 return images[index]
193@writer
194def write_eon(fileobj, images, comment="Generated by ASE"):
195 """
196 Writes structures to an EON trajectory file, allowing for multiple
197 snapshots.
199 This function iterates over all provided images, converting each one into a
200 text format compatible with EON trajectory files. It handles the conversion
201 of the cell parameters, chemical symbols, atomic masses, and atomic
202 constraints. Only FixAtoms constraints are supported; any other constraints
203 will generate a warning and be skipped. Arbitrary masses will raise an
204 error, since EON will not accept them, i.e. no Duterium and Hydrogen.
206 Parameters
207 ----------
208 fileobj : file object
209 An opened, writable file or file-like object to which the EON trajectory
210 information will be written.
211 images : list of Atoms
212 A list of ASE Atoms objects representing the images (atomic structures)
213 to be written into the EON trajectory file. Each Atoms object should
214 have a cell attribute and may have a constraints attribute. If the
215 constraints attribute is not a FixAtoms object, a warning will be
216 issued.
217 comment : str
218 An additional comment statement which will be written out as the first
219 header
221 Raises
222 ------
223 Warning
224 If any constraint in an Atoms object is not an instance of FixAtoms.
225 RuntimeError
226 If the masses for the species are not the default ones, i.e. if isotopes
227 are present
229 Returns
230 -------
231 None
232 The function writes directly to the provided file object.
234 See Also
235 --------
236 ase.io.utils.segment_list : for segmenting the list of images.
238 Examples
239 --------
240 >>> from ase.io import Trajectory
241 >>> from ase.io.utils import segment_list
242 >>> traj = Trajectory("neb.traj")
243 >>> n_images = 10 # Segment size, i.e. number of images in the NEB
244 >>> for idx, pth in enumerate(segment_list(traj, n_images)):
245 ... with open(f"outputs/neb_path_{idx:03d}.con", "w") as fobj:
246 ... write_eon_traj(fobj, pth)
247 """
249 for idx, atoms in enumerate(images):
250 out = []
251 if idx == 0:
252 out.append(comment)
253 else:
254 out.append(f"\n{comment}")
255 out.append("preBox_header_2") # ??
257 a, b, c, alpha, beta, gamma = cell_to_cellpar(atoms.cell)
258 out.append("%-10.6f %-10.6f %-10.6f" % (a, b, c))
259 out.append("%-10.6f %-10.6f %-10.6f" % (gamma, beta, alpha))
261 out.append("postBox_header_1") # ??
262 out.append("postBox_header_2") # ??
264 symbol_indices = atoms.symbols.indices()
265 natoms = [len(x) for x in symbol_indices.values()]
266 masslist = atoms.get_masses()
267 species_masses = []
268 for symbol, indices in symbol_indices.items():
269 masses_for_symb = masslist[indices]
270 if np.ptp(masses_for_symb) > ISOTOPE_TOL:
271 raise RuntimeError(
272 "Isotopes cannot be handled by EON"
273 ", ensure uniform masses for symbols"
274 )
275 species_masses.append(masses_for_symb[0])
277 out.append(str(len(symbol_indices)))
278 out.append(" ".join([str(n) for n in natoms]))
279 out.append(" ".join([str(n) for n in species_masses]))
281 atom_id = 0
282 for cid, (species, indices) in enumerate(symbol_indices.items()):
283 fixed = np.array([False] * len(atoms))
284 out.append(species)
285 out.append("Coordinates of Component %d" % (cid + 1))
286 atom = atoms[indices]
287 coords = atom.positions
288 for constraint in atom.constraints:
289 if not isinstance(constraint, FixAtoms):
290 warn(
291 "Only FixAtoms constraints are supported"
292 "by con files. Dropping %r",
293 constraint,
294 )
295 continue
296 if constraint.index.dtype.kind == "b":
297 fixed = np.array(constraint.index, dtype=int)
298 else:
299 fixed = np.zeros((natoms[cid],), dtype=int)
300 for i in constraint.index:
301 fixed[i] = 1
302 for xyz, fix in zip(coords, fixed):
303 line_fmt = "{:>22.17f} {:>22.17f} {:>22.17f} {:d} {:4d}"
304 line = line_fmt.format(*xyz, int(fix), atom_id)
305 out.append(line)
306 atom_id += 1
307 fileobj.write("\n".join(out))