Coverage for ase / io / dlp4.py: 89.19%
148 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3""" Read/Write DL_POLY_4 CONFIG files """
4import itertools
5import re
6from collections.abc import Generator, Iterable
7from typing import IO
9import numpy as np
11from ase.atoms import Atoms
12from ase.calculators.singlepoint import SinglePointCalculator
13from ase.data import chemical_symbols
14from ase.units import _amu, _auf, _auv
15from ase.utils import reader, writer
17__all__ = ["read_dlp4", "write_dlp4"]
19# dlp4 labels will be registered in atoms.arrays[DLP4_LABELS_KEY]
20DLP4_LABELS_KEY = "dlp4_labels"
21DLP4_DISP_KEY = "dlp4_displacements"
22DLP_F_SI = 1.0e-10 * _amu / (1.0e-12 * 1.0e-12) # Å*Da/ps^2
23DLP_F_ASE = DLP_F_SI / _auf
24DLP_V_SI = 1.0e-10 / 1.0e-12 # Å/ps
25DLP_V_ASE = DLP_V_SI / _auv
28def _get_frame_positions(fd: IO) -> tuple[int, int, int, list[int]]:
29 """Get positions of frames in HISTORY file."""
30 init_pos = fd.tell()
32 fd.seek(0)
34 record_len = len(fd.readline()) # system name, and record size
36 items = fd.readline().strip().split()
38 if len(items) not in (3, 5):
39 raise RuntimeError("Cannot determine version of HISTORY file format.")
41 levcfg, imcon, natoms = map(int, items[0:3])
43 if len(items) == 3: # DLPoly classic
44 # we have to iterate over the entire file
45 pos = [fd.tell() for line in fd if "timestep" in line]
46 else:
47 nframes = int(items[3])
48 pos = [((natoms * (levcfg + 2) + 4) * i + 3) * record_len
49 for i in range(nframes)]
51 fd.seek(init_pos)
52 return levcfg, imcon, natoms, pos
55@reader
56def read_dlp_history(fd: IO,
57 index: int | slice | None = None,
58 symbols: list[str] | None = None) -> list[Atoms]:
59 """Read a HISTORY file.
61 Compatible with DLP4 and DLP_CLASSIC.
63 *Index* can be integer or a slice.
65 Provide a list of element strings to symbols to ignore naming
66 from the HISTORY file.
67 """
68 return list(iread_dlp_history(fd, index, symbols))
71@reader
72def iread_dlp_history(fd: IO,
73 index: int | slice | None = None,
74 symbols: list[str] | None = None
75 ) -> Generator[Atoms, None, None]:
76 """Generator version of read_dlp_history
78 Compatible with DLP4 and DLP_CLASSIC.
80 *Index* can be integer or a slice.
82 Provide a list of element strings to symbols to ignore naming
83 from the HISTORY file.
84 """
85 levcfg, imcon, natoms, pos = _get_frame_positions(fd)
87 positions: Iterable[int] = ()
88 if index is None:
89 positions = pos
90 elif isinstance(index, int):
91 positions = (pos[index],)
92 elif isinstance(index, slice):
93 positions = itertools.islice(pos, index.start, index.stop, index.step)
95 for pos_i in positions:
96 fd.seek(pos_i + 1)
97 yield read_single_image(fd, levcfg, imcon, natoms, is_trajectory=True,
98 symbols=symbols)
101@reader
102def read_dlp4(fd: IO,
103 symbols: list[str] | None = None) -> Atoms:
104 """Read a DL_POLY_4 config/revcon file.
106 Typically used indirectly through read("filename", atoms, format="dlp4").
108 Can be unforgiving with custom chemical element names.
109 Please complain to alin@elena.space for bugs."""
111 fd.readline()
112 levcfg, imcon = map(int, fd.readline().split()[0:2])
114 position = fd.tell()
115 is_trajectory = fd.readline().split()[0] == "timestep"
117 if not is_trajectory:
118 # Difficult to distinguish between trajectory and non-trajectory
119 # formats without reading one line too many.
120 fd.seek(position)
122 natoms = (int(fd.readline().split()[2]) if is_trajectory else None)
123 return read_single_image(fd, levcfg, imcon, natoms, is_trajectory,
124 symbols)
127def read_single_image(fd: IO, levcfg: int, imcon: int,
128 natoms: int | None, is_trajectory: bool,
129 symbols: list[str] | None = None) -> Atoms:
130 """ Read a DLP frame """
131 sym = symbols if symbols else []
132 positions = []
133 velocities = []
134 forces = []
135 charges = []
136 masses = []
137 disps = []
138 labels = []
140 is_pbc = imcon > 0
142 cell = np.zeros((3, 3), dtype=np.float64)
143 if is_pbc or is_trajectory:
144 for j, line in enumerate(itertools.islice(fd, 3)):
145 cell[j, :] = np.array(line.split(), dtype=np.float64)
147 i = 0
148 for i, line in enumerate(itertools.islice(fd, natoms), 1):
149 match = re.match(r"\s*([A-Za-z][a-z]?)(\S*)", line)
150 if not match:
151 raise OSError(f"Line {line} does not match valid format.")
153 symbol, label = match.group(1, 2)
154 symbol = symbol.capitalize()
156 if is_trajectory:
157 mass, charge, *disp = map(float, line.split()[2:])
158 charges.append(charge)
159 masses.append(mass)
160 disps.extend(disp if disp else [None]) # type: ignore[list-item]
162 if not symbols:
163 if symbol not in chemical_symbols:
164 raise ValueError(
165 f"Symbol '{symbol}' not found in chemical symbols table"
166 )
167 sym.append(symbol)
169 # make sure label is not empty
170 labels.append(label if label else "")
172 positions.append([float(x) for x in next(fd).split()[:3]])
173 if levcfg > 0:
174 velocities.append([float(x) * DLP_V_ASE
175 for x in next(fd).split()[:3]])
176 if levcfg > 1:
177 forces.append([float(x) * DLP_F_ASE
178 for x in next(fd).split()[:3]])
180 if symbols and (i != len(symbols)):
181 raise ValueError(
182 f"Counter is at {i} but {len(symbols)} symbols provided."
183 )
185 if imcon == 0:
186 pbc = (False, False, False)
187 elif imcon in (1, 2, 3):
188 pbc = (True, True, True)
189 elif imcon == 6:
190 pbc = (True, True, False)
191 elif imcon in (4, 5):
192 raise ValueError(f"Unsupported imcon: {imcon}")
193 else:
194 raise ValueError(f"Invalid imcon: {imcon}")
196 atoms = Atoms(positions=positions,
197 symbols=sym,
198 cell=cell,
199 pbc=pbc,
200 # Cell is centered around (0, 0, 0) in dlp4:
201 celldisp=-cell.sum(axis=0) / 2)
203 if is_trajectory:
204 atoms.set_masses(masses)
205 atoms.set_array(DLP4_DISP_KEY, disps, float)
206 atoms.set_initial_charges(charges)
208 atoms.set_array(DLP4_LABELS_KEY, labels, str)
209 if levcfg > 0:
210 atoms.set_velocities(velocities)
211 if levcfg > 1:
212 atoms.calc = SinglePointCalculator(atoms, forces=forces)
214 return atoms
217@writer
218def write_dlp4(fd: IO, atoms: Atoms,
219 levcfg: int = 0,
220 title: str = "CONFIG generated by ASE"):
221 """Write a DL_POLY_4 config file.
223 Typically used indirectly through write("filename", atoms, format="dlp4").
225 Can be unforgiven with custom chemical element names.
226 Please complain to alin@elena.space in case of bugs"""
228 def float_format(flt):
229 return format(flt, "20.10f")
231 natoms = len(atoms)
233 if tuple(atoms.pbc) == (False, False, False):
234 imcon = 0
235 elif tuple(atoms.pbc) == (True, True, True):
236 imcon = 3
237 elif tuple(atoms.pbc) == (True, True, False):
238 imcon = 6
239 else:
240 raise ValueError(f"Unsupported pbc: {atoms.pbc}. "
241 "Supported pbc are 000, 110, and 111.")
243 print(f"{title:72s}", file=fd)
244 print(f"{levcfg:10d}{imcon:10d} {natoms}", file=fd)
246 if imcon > 0:
247 for row in atoms.get_cell():
248 print("".join(map(float_format, row)), file=fd)
250 vels = atoms.get_velocities() / DLP_V_ASE if levcfg > 0 else None
251 forces = atoms.get_forces() / DLP_F_ASE if levcfg > 1 else None
253 labels = atoms.arrays.get(DLP4_LABELS_KEY)
255 for i, atom in enumerate(atoms):
257 sym = atom.symbol
258 if labels is not None:
259 sym += labels[i]
261 print(f"{sym:8s}{atom.index + 1:10d}", file=fd)
263 to_write = (atom.x, atom.y, atom.z)
264 print("".join(map(float_format, to_write)), file=fd)
266 if levcfg > 0:
267 to_write = (vels[atom.index, :]
268 if vels is not None else (0.0, 0.0, 0.0))
269 print("".join(map(float_format, to_write)), file=fd)
271 if levcfg > 1:
272 to_write = (forces[atom.index, :]
273 if forces is not None else (0.0, 0.0, 0.0))
274 print("".join(map(float_format, to_write)), file=fd)