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