Coverage for /builds/ase/ase/ase/io/castep/geom_md_ts.py: 93.63%

1"""Parsers for CASTEP .geom, .md, .ts files""" 

2 

3from math import sqrt 

4from typing import Callable, Dict, List, Optional, Sequence, TextIO, Union 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.io.formats import string2index 

10from ase.stress import full_3x3_to_voigt_6_stress, voigt_6_to_full_3x3_stress 

11from ase.utils import reader, writer 

12 

13 

14class Parser: 

15 """Parser for <-- `key` in .geom, .md, .ts files""" 

16 

17 def __init__(self, units: Optional[Dict[str, float]] = None): 

18 if units is None: 

19 from ase.io.castep import units_CODATA2002 

20 

21 self.units = units_CODATA2002 

22 else: 

23 self.units = units 

24 

25 def parse(self, lines: List[str], key: str, method: Callable): 

26 """Parse <-- `key` in `lines` using `method`""" 

27 relevant_lines = [line for line in lines if line.strip().endswith(key)] 

28 if relevant_lines: 

29 return method(relevant_lines, self.units) 

30 return None 

31 

32 

33@reader 

34def _read_images( 

35 fd: TextIO, 

36 index: Union[int, slice, str] = -1, 

37 units: Optional[Dict[str, float]] = None, 

38): 

39 """Read a .geom or a .md file written by CASTEP. 

40 

41 - .geom: written by the CASTEP GeometryOptimization task 

42 - .md: written by the CATSTEP MolecularDynamics task 

43 

44 Original contribution by Wei-Bing Zhang. Thanks! 

45 

46 Parameters 

47 ---------- 

48 fd : str | TextIO 

49 File name or object (possibly compressed with .gz and .bz2) to be read. 

50 index : int | slice | str, default: -1 

51 Index of image to be read. 

52 units : dict[str, float], default: None 

53 Dictionary with conversion factors from atomic units to ASE units. 

54 

55 - ``Eh``: Hartree energy in eV 

56 - ``a0``: Bohr radius in Å 

57 - ``me``: electron mass in Da 

58 - ``kB``: Boltzmann constant in eV/K 

59 

60 If None, values based on CODATA2002 are used. 

61 

62 Returns 

63 ------- 

64 Atoms | list[Atoms] 

65 ASE Atoms object or list of them. 

66 

67 Notes 

68 ----- 

69 The force-consistent energy, forces, stress are stored in ``atoms.calc``. 

70 

71 Everything in the .geom or the .md file is in atomic units. 

72 They are converted in ASE units, i.e., Å (length), eV (energy), Da (mass). 

73 

74 Stress in the .geom or the .md file includes kinetic contribution, which is 

75 subtracted in ``atoms.calc``. 

76 

77 """ 

78 if isinstance(index, str): 

79 index = string2index(index) 

80 if isinstance(index, str): 

81 raise ValueError(index) 

82 return list(_iread_images(fd, units))[index] 

83 

84 

85read_castep_geom = _read_images 

86read_castep_md = _read_images 

87 

88 

89def _iread_images(fd: TextIO, units: Optional[Dict[str, float]] = None): 

90 """Read a .geom or .md file of CASTEP MolecularDynamics as a generator.""" 

91 parser = Parser(units) 

92 _read_header(fd) 

93 lines = [] 

94 for line in fd: 

95 if line.strip(): 

96 lines.append(line) 

97 else: 

98 yield _read_atoms(lines, parser) 

99 lines = [] 

100 

101 

102iread_castep_geom = _iread_images 

103iread_castep_md = _iread_images 

104 

105 

106def _read_atoms(lines: List[str], parser: Parser) -> Atoms: 

107 from ase.calculators.singlepoint import SinglePointCalculator 

108 

109 energy = parser.parse(lines, '<-- E', _read_energies) 

110 cell = parser.parse(lines, '<-- h', _read_cell) 

111 stress = parser.parse(lines, '<-- S', _read_stress) 

112 symbols, positions = parser.parse(lines, '<-- R', _read_positions) 

113 velocities = parser.parse(lines, '<-- V', _read_velocities) 

114 forces = parser.parse(lines, '<-- F', _read_forces) 

115 

116 # Currently unused tags: 

117 # 

118 # temperature = parser.parse(lines, '<-- T', _read_temperature) 

119 # pressure = parser.parse(lines, '<-- P', _read_pressure) 

120 # cell_velocities = parser.extract(lines, '<-- hv', _read_cell_velocities) 

121 

122 atoms = Atoms(symbols, positions, cell=cell, pbc=True) 

123 

124 if velocities is not None: # MolecularDynamics 

125 units = parser.units 

126 factor = units['a0'] * sqrt(units['me'] / units['Eh']) 

127 atoms.info['time'] = float(lines[0].split()[0]) * factor # au -> ASE 

128 atoms.set_velocities(velocities) 

129 

130 if stress is not None: 

131 stress -= atoms.get_kinetic_stress(voigt=True) 

132 

133 # The energy in .geom or .md file is the force-consistent one 

134 # (possibly with the the finite-basis-set correction when, e.g., 

135 # finite_basis_corr!=0 in GeometryOptimisation). 

136 # It should therefore be reasonable to assign it to `free_energy`. 

137 # Be also aware that the energy in .geom file not 0K extrapolated. 

138 atoms.calc = SinglePointCalculator( 

139 atoms=atoms, 

140 free_energy=energy, 

141 forces=forces, 

142 stress=stress, 

143 ) 

144 

145 return atoms 

146 

147 

148def _read_header(fd: TextIO): 

149 for line in fd: 

150 if 'END header' in line: 

151 next(fd) # read blank line below 'END header' 

152 break 

153 

154 

155def _read_energies(lines: List[str], units: Dict[str, float]) -> float: 

156 """Read force-consistent energy 

157 

158 Notes 

159 ----- 

160 Enthalpy and kinetic energy (in .md) are also written in the same line. 

161 They are however not parsed because they can be computed using stress and 

162 atomic velocties, respsectively. 

163 """ 

164 return float(lines[0].split()[0]) * units['Eh'] 

165 

166 

167def _read_temperature(lines: List[str], units: Dict[str, float]) -> float: 

168 """Read temperature 

169 

170 Notes 

171 ----- 

172 Temperature can be computed from kinetic energy and hence not necessary. 

173 """ 

174 factor = units['Eh'] / units['kB'] # hartree -> K 

175 return float(lines[0].split()[0]) * factor 

176 

177 

178def _read_pressure(lines: List[str], units: Dict[str, float]) -> float: 

179 """Read pressure 

180 

181 Notes 

182 ----- 

183 Pressure can be computed from stress and hence not necessary. 

184 """ 

185 factor = units['Eh'] / units['a0'] ** 3 # au -> eV/A3 

186 return float(lines[0].split()[0]) * factor 

187 

188 

189def _read_cell(lines: List[str], units: Dict[str, float]) -> np.ndarray: 

190 bohr = units['a0'] 

191 cell = np.array([line.split()[0:3] for line in lines], dtype=float) 

192 return cell * bohr 

193 

194 

195# def _read_cell_velocities(lines: List[str], units: Dict[str, float]): 

196# hartree = units['Eh'] 

197# me = units['me'] 

198# cell_velocities = np.array([_.split()[0:3] for _ in lines], dtype=float) 

199# return cell_velocities * np.sqrt(hartree / me) 

200 

201 

202def _read_stress(lines: List[str], units: Dict[str, float]) -> np.ndarray: 

203 hartree = units['Eh'] 

204 bohr = units['a0'] 

205 stress = np.array([line.split()[0:3] for line in lines], dtype=float) 

206 return full_3x3_to_voigt_6_stress(stress) * (hartree / bohr**3) 

207 

208 

209def _read_positions( 

210 lines: List[str], units: Dict[str, float] 

211) -> tuple[list[str], np.ndarray]: 

212 bohr = units['a0'] 

213 symbols = [line.split()[0] for line in lines] 

214 positions = np.array([line.split()[2:5] for line in lines], dtype=float) 

215 return symbols, positions * bohr 

216 

217 

218def _read_velocities(lines: List[str], units: Dict[str, float]) -> np.ndarray: 

219 hartree = units['Eh'] 

220 me = units['me'] 

221 velocities = np.array([line.split()[2:5] for line in lines], dtype=float) 

222 return velocities * np.sqrt(hartree / me) 

223 

224 

225def _read_forces(lines: List[str], units: Dict[str, float]) -> np.ndarray: 

226 hartree = units['Eh'] 

227 bohr = units['a0'] 

228 forces = np.array([line.split()[2:5] for line in lines], dtype=float) 

229 return forces * (hartree / bohr) 

230 

231 

232@writer 

233def write_castep_geom( 

234 fd: TextIO, 

235 images: Union[Atoms, Sequence[Atoms]], 

236 units: Optional[Dict[str, float]] = None, 

237 *, 

238 pressure: float = 0.0, 

239 sort: bool = False, 

240): 

241 """Write a CASTEP .geom file. 

242 

243 .. versionadded:: 3.25.0 

244 

245 Parameters 

246 ---------- 

247 fd : str | TextIO 

248 File name or object (possibly compressed with .gz and .bz2) to be read. 

249 images : Atoms | Sequenece[Atoms] 

250 ASE Atoms object(s) to be written. 

251 units : dict[str, float], default: None 

252 Dictionary with conversion factors from atomic units to ASE units. 

253 

254 - ``Eh``: Hartree energy in eV 

255 - ``a0``: Bohr radius in Å 

256 - ``me``: electron mass in Da 

257 - ``kB``: Boltzmann constant in eV/K 

258 

259 If None, values based on CODATA2002 are used. 

260 pressure : float, default: 0.0 

261 External pressure in eV/Å\\ :sup:`3`. 

262 sort : bool, default: False 

263 If True, atoms are sorted in ascending order of atomic number. 

264 

265 Notes 

266 ----- 

267 - Values in the .geom file are in atomic units. 

268 - Stress is printed including kinetic contribution. 

269 

270 """ 

271 if isinstance(images, Atoms): 

272 images = [images] 

273 

274 if units is None: 

275 from ase.io.castep import units_CODATA2002 

276 

277 units = units_CODATA2002 

278 

279 _write_header(fd) 

280 

281 for index, atoms in enumerate(images): 

282 if sort: 

283 atoms = atoms[atoms.numbers.argsort()] 

284 _write_convergence_status(fd, index) 

285 _write_energies_geom(fd, atoms, units, pressure) 

286 _write_cell(fd, atoms, units) 

287 _write_stress(fd, atoms, units) 

288 _write_positions(fd, atoms, units) 

289 _write_forces(fd, atoms, units) 

290 fd.write(' \n') 

291 

292 

293@writer 

294def write_castep_md( 

295 fd: TextIO, 

296 images: Union[Atoms, Sequence[Atoms]], 

297 units: Optional[Dict[str, float]] = None, 

298 *, 

299 pressure: float = 0.0, 

300 sort: bool = False, 

301): 

302 """Write a CASTEP .md file. 

303 

304 .. versionadded:: 3.25.0 

305 

306 Parameters 

307 ---------- 

308 fd : str | TextIO 

309 File name or object (possibly compressed with .gz and .bz2) to be read. 

310 images : Atoms | Sequenece[Atoms] 

311 ASE Atoms object(s) to be written. 

312 units : dict[str, float], default: None 

313 Dictionary with conversion factors from atomic units to ASE units. 

314 

315 - ``Eh``: Hartree energy in eV 

316 - ``a0``: Bohr radius in Å 

317 - ``me``: electron mass in Da 

318 - ``kB``: Boltzmann constant in eV/K 

319 

320 If None, values based on CODATA2002 are used. 

321 pressure : float, default: 0.0 

322 External pressure in eV/Å\\ :sup:`3`. 

323 sort : bool, default: False 

324 If True, atoms are sorted in ascending order of atomic number. 

325 

326 Notes 

327 ----- 

328 - Values in the .md file are in atomic units. 

329 - Stress is printed including kinetic contribution. 

330 

331 """ 

332 if isinstance(images, Atoms): 

333 images = [images] 

334 

335 if units is None: 

336 from ase.io.castep import units_CODATA2002 

337 

338 units = units_CODATA2002 

339 

340 _write_header(fd) 

341 

342 for index, atoms in enumerate(images): 

343 if sort: 

344 atoms = atoms[atoms.numbers.argsort()] 

345 _write_time(fd, index) 

346 _write_energies_md(fd, atoms, units, pressure) 

347 _write_temperature(fd, atoms, units) 

348 _write_cell(fd, atoms, units) 

349 _write_cell_velocities(fd, atoms, units) 

350 _write_stress(fd, atoms, units) 

351 _write_positions(fd, atoms, units) 

352 _write_velocities(fd, atoms, units) 

353 _write_forces(fd, atoms, units) 

354 fd.write(' \n') 

355 

356 

357def _format_float(x: float) -> str: 

358 """Format a floating number for .geom and .md files""" 

359 return np.format_float_scientific( 

360 x, 

361 precision=16, 

362 unique=False, 

363 pad_left=2, 

364 exp_digits=3, 

365 ).replace('e', 'E') 

366 

367 

368def _write_header(fd: TextIO): 

369 fd.write(' BEGIN header\n') 

370 fd.write(' \n') 

371 fd.write(' END header\n') 

372 fd.write(' \n') 

373 

374 

375def _write_convergence_status(fd: TextIO, index: int): 

376 fd.write(21 * ' ') 

377 fd.write(f'{index:18d}') 

378 fd.write(34 * ' ') 

379 for _ in range(4): 

380 fd.write(' F') # Convergence status. So far F for all. 

381 fd.write(10 * ' ') 

382 fd.write(' <-- c\n') 

383 

384 

385def _write_time(fd: TextIO, index: int): 

386 fd.write(18 * ' ' + f' {_format_float(index)}\n') # So far index. 

387 

388 

389def _write_energies_geom( 

390 fd: TextIO, 

391 atoms: Atoms, 

392 units: Dict[str, float], 

393 pressure: float = 0.0, 

394): 

395 """Write energies (in hartree) in a CASTEP .geom file. 

396 

397 The energy and the enthalpy are printed. 

398 """ 

399 hartree = units['Eh'] 

400 if atoms.calc is None: 

401 return 

402 if atoms.calc.results.get('free_energy') is None: 

403 return 

404 energy = atoms.calc.results.get('free_energy') / hartree 

405 pv = pressure * atoms.get_volume() / hartree 

406 fd.write(18 * ' ') 

407 fd.write(f' {_format_float(energy)}') 

408 fd.write(f' {_format_float(energy + pv)}') 

409 fd.write(27 * ' ') 

410 fd.write(' <-- E\n') 

411 

412 

413def _write_energies_md( 

414 fd: TextIO, 

415 atoms: Atoms, 

416 units: Dict[str, float], 

417 pressure: float = 0.0, 

418): 

419 """Write energies (in hartree) in a CASTEP .md file. 

420 

421 The potential energy, the total energy or enthalpy, and the kinetic energy 

422 are printed. 

423 

424 Notes 

425 ----- 

426 For the second item, CASTEP prints the total energy for the NVE and the NVT 

427 ensembles and the total enthalpy for the NPH and NPT ensembles. 

428 For the Nosé–Hoover (chain) thermostat, furthermore, the energies of the 

429 thermostats are also added. 

430 

431 """ 

432 hartree = units['Eh'] 

433 if atoms.calc is None: 

434 return 

435 if atoms.calc.results.get('free_energy') is None: 

436 return 

437 potential = atoms.calc.results.get('free_energy') / hartree 

438 kinetic = atoms.get_kinetic_energy() / hartree 

439 pv = pressure * atoms.get_volume() / hartree 

440 fd.write(18 * ' ') 

441 fd.write(f' {_format_float(potential)}') 

442 fd.write(f' {_format_float(potential + kinetic + pv)}') 

443 fd.write(f' {_format_float(kinetic)}') 

444 fd.write(' <-- E\n') 

445 

446 

447def _write_temperature(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

448 """Write temperature (in hartree) in a CASTEP .md file. 

449 

450 CASTEP writes the temperature in a .md file with 3`N` degrees of freedom 

451 regardless of the given constraints including the fixed center of mass 

452 (``FIX_COM`` which is by default ``TRUE`` in many cases). 

453 To get the consistent result with the above behavior of CASTEP using 

454 this method, we must set the corresponding constraints (e.g. ``FixCom``) 

455 to the ``Atoms`` object beforehand. 

456 """ 

457 hartree = units['Eh'] # eV 

458 boltzmann = units['kB'] # eV/K 

459 temperature = atoms.get_temperature() * boltzmann / hartree 

460 fd.write(18 * ' ') 

461 fd.write(f' {_format_float(temperature)}') 

462 fd.write(54 * ' ') 

463 fd.write(' <-- T\n') 

464 

465 

466def _write_cell(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

467 bohr = units['a0'] 

468 cell = atoms.cell / bohr # in bohr 

469 for i in range(3): 

470 fd.write(18 * ' ') 

471 for j in range(3): 

472 fd.write(f' {_format_float(cell[i, j])}') 

473 fd.write(' <-- h\n') 

474 

475 

476def _write_cell_velocities(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

477 pass # TODO: to be implemented 

478 

479 

480def _write_stress(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

481 if atoms.calc is None: 

482 return 

483 

484 stress = atoms.calc.results.get('stress') 

485 if stress is None: 

486 return 

487 

488 if stress.shape != (3, 3): 

489 stress = voigt_6_to_full_3x3_stress(stress) 

490 

491 stress += atoms.get_kinetic_stress(voigt=False) 

492 

493 hartree = units['Eh'] 

494 bohr = units['a0'] 

495 stress = stress / (hartree / bohr**3) 

496 

497 for i in range(3): 

498 fd.write(18 * ' ') 

499 for j in range(3): 

500 fd.write(f' {_format_float(stress[i, j])}') 

501 fd.write(' <-- S\n') 

502 

503 

504def _write_positions(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

505 bohr = units['a0'] 

506 positions = atoms.positions / bohr 

507 symbols = atoms.symbols 

508 indices = symbols.species_indices() 

509 for i, symbol, position in zip(indices, symbols, positions): 

510 fd.write(f' {symbol:8s}') 

511 fd.write(f' {i + 1:8d}') 

512 for j in range(3): 

513 fd.write(f' {_format_float(position[j])}') 

514 fd.write(' <-- R\n') 

515 

516 

517def _write_forces(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

518 if atoms.calc is None: 

519 return 

520 

521 forces = atoms.calc.results.get('forces') 

522 if forces is None: 

523 return 

524 

525 hartree = units['Eh'] 

526 bohr = units['a0'] 

527 forces = forces / (hartree / bohr) 

528 

529 symbols = atoms.symbols 

530 indices = symbols.species_indices() 

531 for i, symbol, force in zip(indices, symbols, forces): 

532 fd.write(f' {symbol:8s}') 

533 fd.write(f' {i + 1:8d}') 

534 for j in range(3): 

535 fd.write(f' {_format_float(force[j])}') 

536 fd.write(' <-- F\n') 

537 

538 

539def _write_velocities(fd: TextIO, atoms: Atoms, units: Dict[str, float]): 

540 hartree = units['Eh'] 

541 me = units['me'] 

542 velocities = atoms.get_velocities() / np.sqrt(hartree / me) 

543 symbols = atoms.symbols 

544 indices = symbols.species_indices() 

545 for i, symbol, velocity in zip(indices, symbols, velocities): 

546 fd.write(f' {symbol:8s}') 

547 fd.write(f' {i + 1:8d}') 

548 for j in range(3): 

549 fd.write(f' {_format_float(velocity[j])}') 

550 fd.write(' <-- V\n')