File size: 2,676 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
"""Functions to read and write XYZ files."""
__all__ = [
'write_frame_xyz',
'read_frame_xyz',
]
import numpy as np
from ..constants import angstrom
from .utilities import Frame, register_io
@register_io('xyz', 'read', 'xyz')
def read_frame_xyz(f_in, name_data='positions', unit=angstrom):
"""Read one frame of XYZ format from an open file.
Arguments:
f_in: open file in XYZ format
name_data: what quantity to take the XYZ data as
unit: unit to scale data by, multiplicative factor in atomic units
Returns:
`Frame` object or `None` if there is no more data
"""
# read first line to examine it
line_begin = f_in.readline()
# no more data in the file
if not line_begin:
return None
# there is some data, frame should begin with natoms
natoms = int(line_begin)
# read comment line
comment = f_in.readline().rstrip()
names = []
data = []
for _ in range(natoms):
line = f_in.readline()
if line.strip() == '':
raise ValueError('Unexpected data in file.')
items = line.split()
names.append(items[0])
data.append([float(item) for item in items[1:4]])
data = np.array(data) * unit
# so unless the code fails, this will not trigger.
if len(names) != natoms:
raise ValueError('Inconsistent number of atoms in XYZ file.')
# prepare data
if name_data == 'positions':
positions = data
forces = None
elif name_data == 'forces':
positions = None
forces = data
else:
raise ValueError(f'Unsupported `name_data`: {name_data}. Expected "positions" or "forces".')
return Frame(names=names, positions=positions, comment=comment, energy=None, forces=forces)
@register_io('xyz', 'write', 'xyz')
def write_frame_xyz(f_out, frame, unit=angstrom):
"""Print a single frame into an open XYZ file.
This is currently hard-coded to write positions, if we ever need to write forces
or something else, it needs generalizing.
"""
# Check that required things are in frame:
if (frame.positions is None) or (frame.names is None):
raise ValueError('Frame does not contain required properties.')
fmt_one = '{:13.6f}'
fmt_prop = '{:6s} ' + 3*fmt_one + '\n'
# write number of atoms and comment line
f_out.write(f'{len(frame.names):d}\n')
if frame.comment is not None:
f_out.write(f'{frame.comment:s}\n')
else:
f_out.write('\n')
data = frame.positions / unit
# write atomic lines
for i, name in enumerate(frame.names):
f_out.write(fmt_prop.format(name, *data[i]))
|