Coverage for /builds/kinetik161/ase/ase/io/dlp4.py: 89.12%

1""" Read/Write DL_POLY_4 CONFIG files """ 

2import itertools 

3import re 

4from typing import IO, Generator, Iterable, List, Optional, Tuple, Union 

5 

6import numpy as np 

7 

8from ase.atoms import Atoms 

9from ase.calculators.singlepoint import SinglePointCalculator 

10from ase.data import chemical_symbols 

11from ase.units import _amu, _auf, _auv 

12from ase.utils import reader, writer 

13 

14__all__ = ["read_dlp4", "write_dlp4"] 

15 

16# dlp4 labels will be registered in atoms.arrays[DLP4_LABELS_KEY] 

17DLP4_LABELS_KEY = "dlp4_labels" 

18DLP4_DISP_KEY = "dlp4_displacements" 

19DLP_F_SI = 1.0e-10 * _amu / (1.0e-12 * 1.0e-12) # Å*Da/ps^2 

20DLP_F_ASE = DLP_F_SI / _auf 

21DLP_V_SI = 1.0e-10 / 1.0e-12 # Å/ps 

22DLP_V_ASE = DLP_V_SI / _auv 

23 

24 

25def _get_frame_positions(fd: IO) -> Tuple[int, int, int, List[int]]: 

26 """Get positions of frames in HISTORY file.""" 

27 init_pos = fd.tell() 

28 

29 fd.seek(0) 

30 

31 record_len = len(fd.readline()) # system name, and record size 

32 

33 items = fd.readline().strip().split() 

34 

35 if len(items) not in (3, 5): 

36 raise RuntimeError("Cannot determine version of HISTORY file format.") 

37 

38 levcfg, imcon, natoms = map(int, items[0:3]) 

39 

40 if len(items) == 3: # DLPoly classic 

41 # we have to iterate over the entire file 

42 pos = [fd.tell() for line in fd if "timestep" in line] 

43 else: 

44 nframes = int(items[3]) 

45 pos = [((natoms * (levcfg + 2) + 4) * i + 3) * record_len 

46 for i in range(nframes)] 

47 

48 fd.seek(init_pos) 

49 return levcfg, imcon, natoms, pos 

50 

51 

52@reader 

53def read_dlp_history(fd: IO, 

54 index: Optional[Union[int, slice]] = None, 

55 symbols: Optional[List[str]] = None) -> List[Atoms]: 

56 """Read a HISTORY file. 

57 

58 Compatible with DLP4 and DLP_CLASSIC. 

59 

60 *Index* can be integer or a slice. 

61 

62 Provide a list of element strings to symbols to ignore naming 

63 from the HISTORY file. 

64 """ 

65 return list(iread_dlp_history(fd, index, symbols)) 

66 

67 

68@reader 

69def iread_dlp_history(fd: IO, 

70 index: Optional[Union[int, slice]] = None, 

71 symbols: Optional[List[str]] = None 

72 ) -> Generator[Atoms, None, None]: 

73 """Generator version of read_dlp_history 

74 

75 Compatible with DLP4 and DLP_CLASSIC. 

76 

77 *Index* can be integer or a slice. 

78 

79 Provide a list of element strings to symbols to ignore naming 

80 from the HISTORY file. 

81 """ 

82 levcfg, imcon, natoms, pos = _get_frame_positions(fd) 

83 

84 positions: Iterable[int] = () 

85 if index is None: 

86 positions = pos 

87 elif isinstance(index, int): 

88 positions = (pos[index],) 

89 elif isinstance(index, slice): 

90 positions = itertools.islice(pos, index.start, index.stop, index.step) 

91 

92 for pos_i in positions: 

93 fd.seek(pos_i + 1) 

94 yield read_single_image(fd, levcfg, imcon, natoms, is_trajectory=True, 

95 symbols=symbols) 

96 

97 

98@reader 

99def read_dlp4(fd: IO, 

100 symbols: Optional[List[str]] = None) -> Atoms: 

101 """Read a DL_POLY_4 config/revcon file. 

102 

103 Typically used indirectly through read("filename", atoms, format="dlp4"). 

104 

105 Can be unforgiving with custom chemical element names. 

106 Please complain to alin@elena.space for bugs.""" 

107 

108 fd.readline() 

109 levcfg, imcon = map(int, fd.readline().split()[0:2]) 

110 

111 position = fd.tell() 

112 is_trajectory = fd.readline().split()[0] == "timestep" 

113 

114 if not is_trajectory: 

115 # Difficult to distinguish between trajectory and non-trajectory 

116 # formats without reading one line too many. 

117 fd.seek(position) 

118 

119 natoms = (int(fd.readline().split()[2]) if is_trajectory else None) 

120 return read_single_image(fd, levcfg, imcon, natoms, is_trajectory, 

121 symbols) 

122 

123 

124def read_single_image(fd: IO, levcfg: int, imcon: int, 

125 natoms: Optional[int], is_trajectory: bool, 

126 symbols: Optional[List[str]] = None) -> Atoms: 

127 """ Read a DLP frame """ 

128 sym = symbols if symbols else [] 

129 positions = [] 

130 velocities = [] 

131 forces = [] 

132 charges = [] 

133 masses = [] 

134 disps = [] 

135 labels = [] 

136 

137 is_pbc = imcon > 0 

138 

139 cell = np.zeros((3, 3), dtype=np.float64) 

140 if is_pbc or is_trajectory: 

141 for j, line in enumerate(itertools.islice(fd, 3)): 

142 cell[j, :] = np.array(line.split(), dtype=np.float64) 

143 

144 i = 0 

145 for i, line in enumerate(itertools.islice(fd, natoms), 1): 

146 match = re.match(r"\s*([A-Za-z][a-z]?)(\S*)", line) 

147 if not match: 

148 raise OSError(f"Line {line} does not match valid format.") 

149 

150 symbol, label = match.group(1, 2) 

151 symbol = symbol.capitalize() 

152 

153 if is_trajectory: 

154 mass, charge, *disp = map(float, line.split()[2:]) 

155 charges.append(charge) 

156 masses.append(mass) 

157 disps.extend(disp if disp else [None]) # type: ignore[list-item] 

158 

159 if not symbols: 

160 if symbol not in chemical_symbols: 

161 raise ValueError( 

162 f"Symbol '{symbol}' not found in chemical symbols table" 

163 ) 

164 sym.append(symbol) 

165 

166 # make sure label is not empty 

167 labels.append(label if label else "") 

168 

169 positions.append([float(x) for x in next(fd).split()[:3]]) 

170 if levcfg > 0: 

171 velocities.append([float(x) * DLP_V_ASE 

172 for x in next(fd).split()[:3]]) 

173 if levcfg > 1: 

174 forces.append([float(x) * DLP_F_ASE 

175 for x in next(fd).split()[:3]]) 

176 

177 if symbols and (i != len(symbols)): 

178 raise ValueError( 

179 f"Counter is at {i} but {len(symbols)} symbols provided." 

180 ) 

181 

182 if imcon == 0: 

183 pbc = (False, False, False) 

184 elif imcon in (1, 2, 3): 

185 pbc = (True, True, True) 

186 elif imcon == 6: 

187 pbc = (True, True, False) 

188 elif imcon in (4, 5): 

189 raise ValueError(f"Unsupported imcon: {imcon}") 

190 else: 

191 raise ValueError(f"Invalid imcon: {imcon}") 

192 

193 atoms = Atoms(positions=positions, 

194 symbols=sym, 

195 cell=cell, 

196 pbc=pbc, 

197 # Cell is centered around (0, 0, 0) in dlp4: 

198 celldisp=-cell.sum(axis=0) / 2) 

199 

200 if is_trajectory: 

201 atoms.set_masses(masses) 

202 atoms.set_array(DLP4_DISP_KEY, disps, float) 

203 atoms.set_initial_charges(charges) 

204 

205 atoms.set_array(DLP4_LABELS_KEY, labels, str) 

206 if levcfg > 0: 

207 atoms.set_velocities(velocities) 

208 if levcfg > 1: 

209 atoms.calc = SinglePointCalculator(atoms, forces=forces) 

210 

211 return atoms 

212 

213 

214@writer 

215def write_dlp4(fd: IO, atoms: Atoms, 

216 levcfg: int = 0, 

217 title: str = "CONFIG generated by ASE"): 

218 """Write a DL_POLY_4 config file. 

219 

220 Typically used indirectly through write("filename", atoms, format="dlp4"). 

221 

222 Can be unforgiven with custom chemical element names. 

223 Please complain to alin@elena.space in case of bugs""" 

224 

225 def float_format(flt): 

226 return format(flt, "20.10f") 

227 

228 natoms = len(atoms) 

229 

230 if tuple(atoms.pbc) == (False, False, False): 

231 imcon = 0 

232 elif tuple(atoms.pbc) == (True, True, True): 

233 imcon = 3 

234 elif tuple(atoms.pbc) == (True, True, False): 

235 imcon = 6 

236 else: 

237 raise ValueError(f"Unsupported pbc: {atoms.pbc}. " 

238 "Supported pbc are 000, 110, and 111.") 

239 

240 print(f"{title:72s}", file=fd) 

241 print(f"{levcfg:10d}{imcon:10d}{natoms:10d}", file=fd) 

242 

243 if imcon > 0: 

244 for row in atoms.get_cell(): 

245 print("".join(map(float_format, row)), file=fd) 

246 

247 vels = atoms.get_velocities() / DLP_V_ASE if levcfg > 0 else [] 

248 forces = atoms.get_forces() / DLP_F_ASE if levcfg > 1 else [] 

249 

250 labels = atoms.arrays.get(DLP4_LABELS_KEY) 

251 

252 for i, atom in enumerate(atoms): 

253 

254 sym = atom.symbol 

255 if labels is not None: 

256 sym += labels[i] 

257 

258 print(f"{sym:8s}{atom.index + 1:10d}", file=fd) 

259 

260 to_write = (atom.x, atom.y, atom.z) 

261 print("".join(map(float_format, to_write)), file=fd) 

262 

263 if levcfg > 0: 

264 to_write = (vels[atom.index, :] 

265 if vels is not None else (0.0, 0.0, 0.0)) 

266 print("".join(map(float_format, to_write)), file=fd) 

267 

268 if levcfg > 1: 

269 to_write = (forces[atom.index, :] 

270 if forces is not None else (0.0, 0.0, 0.0)) 

271 print("".join(map(float_format, to_write)), file=fd)