Coverage for /builds/ase/ase/ase/io/dlp4.py: 89.12%
147 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3""" Read/Write DL_POLY_4 CONFIG files """
4import itertools
5import re
6from typing import IO, Generator, Iterable, List, Optional, Tuple, Union
8import numpy as np
10from ase.atoms import Atoms
11from ase.calculators.singlepoint import SinglePointCalculator
12from ase.data import chemical_symbols
13from ase.units import _amu, _auf, _auv
14from ase.utils import reader, writer
16__all__ = ["read_dlp4", "write_dlp4"]
18# dlp4 labels will be registered in atoms.arrays[DLP4_LABELS_KEY]
19DLP4_LABELS_KEY = "dlp4_labels"
20DLP4_DISP_KEY = "dlp4_displacements"
21DLP_F_SI = 1.0e-10 * _amu / (1.0e-12 * 1.0e-12) # Å*Da/ps^2
22DLP_F_ASE = DLP_F_SI / _auf
23DLP_V_SI = 1.0e-10 / 1.0e-12 # Å/ps
24DLP_V_ASE = DLP_V_SI / _auv
27def _get_frame_positions(fd: IO) -> Tuple[int, int, int, List[int]]:
28 """Get positions of frames in HISTORY file."""
29 init_pos = fd.tell()
31 fd.seek(0)
33 record_len = len(fd.readline()) # system name, and record size
35 items = fd.readline().strip().split()
37 if len(items) not in (3, 5):
38 raise RuntimeError("Cannot determine version of HISTORY file format.")
40 levcfg, imcon, natoms = map(int, items[0:3])
42 if len(items) == 3: # DLPoly classic
43 # we have to iterate over the entire file
44 pos = [fd.tell() for line in fd if "timestep" in line]
45 else:
46 nframes = int(items[3])
47 pos = [((natoms * (levcfg + 2) + 4) * i + 3) * record_len
48 for i in range(nframes)]
50 fd.seek(init_pos)
51 return levcfg, imcon, natoms, pos
54@reader
55def read_dlp_history(fd: IO,
56 index: Optional[Union[int, slice]] = None,
57 symbols: Optional[List[str]] = None) -> List[Atoms]:
58 """Read a HISTORY file.
60 Compatible with DLP4 and DLP_CLASSIC.
62 *Index* can be integer or a slice.
64 Provide a list of element strings to symbols to ignore naming
65 from the HISTORY file.
66 """
67 return list(iread_dlp_history(fd, index, symbols))
70@reader
71def iread_dlp_history(fd: IO,
72 index: Optional[Union[int, slice]] = None,
73 symbols: Optional[List[str]] = None
74 ) -> Generator[Atoms, None, None]:
75 """Generator version of read_dlp_history
77 Compatible with DLP4 and DLP_CLASSIC.
79 *Index* can be integer or a slice.
81 Provide a list of element strings to symbols to ignore naming
82 from the HISTORY file.
83 """
84 levcfg, imcon, natoms, pos = _get_frame_positions(fd)
86 positions: Iterable[int] = ()
87 if index is None:
88 positions = pos
89 elif isinstance(index, int):
90 positions = (pos[index],)
91 elif isinstance(index, slice):
92 positions = itertools.islice(pos, index.start, index.stop, index.step)
94 for pos_i in positions:
95 fd.seek(pos_i + 1)
96 yield read_single_image(fd, levcfg, imcon, natoms, is_trajectory=True,
97 symbols=symbols)
100@reader
101def read_dlp4(fd: IO,
102 symbols: Optional[List[str]] = None) -> Atoms:
103 """Read a DL_POLY_4 config/revcon file.
105 Typically used indirectly through read("filename", atoms, format="dlp4").
107 Can be unforgiving with custom chemical element names.
108 Please complain to alin@elena.space for bugs."""
110 fd.readline()
111 levcfg, imcon = map(int, fd.readline().split()[0:2])
113 position = fd.tell()
114 is_trajectory = fd.readline().split()[0] == "timestep"
116 if not is_trajectory:
117 # Difficult to distinguish between trajectory and non-trajectory
118 # formats without reading one line too many.
119 fd.seek(position)
121 natoms = (int(fd.readline().split()[2]) if is_trajectory else None)
122 return read_single_image(fd, levcfg, imcon, natoms, is_trajectory,
123 symbols)
126def read_single_image(fd: IO, levcfg: int, imcon: int,
127 natoms: Optional[int], is_trajectory: bool,
128 symbols: Optional[List[str]] = None) -> Atoms:
129 """ Read a DLP frame """
130 sym = symbols if symbols else []
131 positions = []
132 velocities = []
133 forces = []
134 charges = []
135 masses = []
136 disps = []
137 labels = []
139 is_pbc = imcon > 0
141 cell = np.zeros((3, 3), dtype=np.float64)
142 if is_pbc or is_trajectory:
143 for j, line in enumerate(itertools.islice(fd, 3)):
144 cell[j, :] = np.array(line.split(), dtype=np.float64)
146 i = 0
147 for i, line in enumerate(itertools.islice(fd, natoms), 1):
148 match = re.match(r"\s*([A-Za-z][a-z]?)(\S*)", line)
149 if not match:
150 raise OSError(f"Line {line} does not match valid format.")
152 symbol, label = match.group(1, 2)
153 symbol = symbol.capitalize()
155 if is_trajectory:
156 mass, charge, *disp = map(float, line.split()[2:])
157 charges.append(charge)
158 masses.append(mass)
159 disps.extend(disp if disp else [None]) # type: ignore[list-item]
161 if not symbols:
162 if symbol not in chemical_symbols:
163 raise ValueError(
164 f"Symbol '{symbol}' not found in chemical symbols table"
165 )
166 sym.append(symbol)
168 # make sure label is not empty
169 labels.append(label if label else "")
171 positions.append([float(x) for x in next(fd).split()[:3]])
172 if levcfg > 0:
173 velocities.append([float(x) * DLP_V_ASE
174 for x in next(fd).split()[:3]])
175 if levcfg > 1:
176 forces.append([float(x) * DLP_F_ASE
177 for x in next(fd).split()[:3]])
179 if symbols and (i != len(symbols)):
180 raise ValueError(
181 f"Counter is at {i} but {len(symbols)} symbols provided."
182 )
184 if imcon == 0:
185 pbc = (False, False, False)
186 elif imcon in (1, 2, 3):
187 pbc = (True, True, True)
188 elif imcon == 6:
189 pbc = (True, True, False)
190 elif imcon in (4, 5):
191 raise ValueError(f"Unsupported imcon: {imcon}")
192 else:
193 raise ValueError(f"Invalid imcon: {imcon}")
195 atoms = Atoms(positions=positions,
196 symbols=sym,
197 cell=cell,
198 pbc=pbc,
199 # Cell is centered around (0, 0, 0) in dlp4:
200 celldisp=-cell.sum(axis=0) / 2)
202 if is_trajectory:
203 atoms.set_masses(masses)
204 atoms.set_array(DLP4_DISP_KEY, disps, float)
205 atoms.set_initial_charges(charges)
207 atoms.set_array(DLP4_LABELS_KEY, labels, str)
208 if levcfg > 0:
209 atoms.set_velocities(velocities)
210 if levcfg > 1:
211 atoms.calc = SinglePointCalculator(atoms, forces=forces)
213 return atoms
216@writer
217def write_dlp4(fd: IO, atoms: Atoms,
218 levcfg: int = 0,
219 title: str = "CONFIG generated by ASE"):
220 """Write a DL_POLY_4 config file.
222 Typically used indirectly through write("filename", atoms, format="dlp4").
224 Can be unforgiven with custom chemical element names.
225 Please complain to alin@elena.space in case of bugs"""
227 def float_format(flt):
228 return format(flt, "20.10f")
230 natoms = len(atoms)
232 if tuple(atoms.pbc) == (False, False, False):
233 imcon = 0
234 elif tuple(atoms.pbc) == (True, True, True):
235 imcon = 3
236 elif tuple(atoms.pbc) == (True, True, False):
237 imcon = 6
238 else:
239 raise ValueError(f"Unsupported pbc: {atoms.pbc}. "
240 "Supported pbc are 000, 110, and 111.")
242 print(f"{title:72s}", file=fd)
243 print(f"{levcfg:10d}{imcon:10d} {natoms}", file=fd)
245 if imcon > 0:
246 for row in atoms.get_cell():
247 print("".join(map(float_format, row)), file=fd)
249 vels = atoms.get_velocities() / DLP_V_ASE if levcfg > 0 else None
250 forces = atoms.get_forces() / DLP_F_ASE if levcfg > 1 else None
252 labels = atoms.arrays.get(DLP4_LABELS_KEY)
254 for i, atom in enumerate(atoms):
256 sym = atom.symbol
257 if labels is not None:
258 sym += labels[i]
260 print(f"{sym:8s}{atom.index + 1:10d}", file=fd)
262 to_write = (atom.x, atom.y, atom.z)
263 print("".join(map(float_format, to_write)), file=fd)
265 if levcfg > 0:
266 to_write = (vels[atom.index, :]
267 if vels is not None else (0.0, 0.0, 0.0))
268 print("".join(map(float_format, to_write)), file=fd)
270 if levcfg > 1:
271 to_write = (forces[atom.index, :]
272 if forces is not None else (0.0, 0.0, 0.0))
273 print("".join(map(float_format, to_write)), file=fd)