File size: 2,164 Bytes
aac5fad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
"""Functions for CP2K-specific input/output."""
__all__ = ['add_energy_cp2k_comment', 'read_frames_cp2k']
from itertools import repeat
from .utilities import Frame, merge_frames, read_frames
def add_energy_cp2k_comment(frames):
"""Parse CP2K energy and inject it into frames.
For each frame in `frames`, try to extract a CP2K-formatted potential energy
from the comment string and inject it back into the frame. Energy from CP2K is
in Hartree, so no conversion is needed.
"""
for frame in frames:
if frame.energy is not None:
raise ValueError('Energy already present.')
try:
for pair in frame.comment.split(','):
items = pair.split('=')
if items[0].strip() == 'E':
frame.energy = float(items[1])
break
except (IndexError, ValueError):
raise ValueError('No CP2K energy found in comment line.')
yield frame
def read_frames_cp2k(fn_positions, cell=None, fn_forces=None, read_energy: bool = True, force_unit=1.0):
"""Read data specifically produced by CP2K.
Arguments:
fn_positions: position trajectory file name, XYZ format
cell: a constant cell to use in all frames, optional
fn_forces: forces file name, XYZ format, optional
read_energy: whether to read energies from comments in `fn_positions`
Returns:
a `Frame` object
"""
# positions from XYZ, energies from comment if requested
# we expect units of angstrom for positions from CP2K
frames_pos = read_frames(fn_positions, fformat='xyz')
if read_energy:
frames_pos = add_energy_cp2k_comment(frames_pos)
frames = [frames_pos]
# add a constant cell if provided
if cell is not None:
frames.append(repeat(Frame(cell=cell)))
# add forces from XYZ if filename was provided
# we expect atomic units for forces from CP2K per default
if fn_forces is not None:
frames.append(read_frames(fn_forces, fformat='xyz', name_data='forces', unit=force_unit))
# iterate over merged frames
yield from merge_frames(*frames)
|