def read_series(trr_filename, n_frames, readx=True, readv=False, readf=False):
"""Reads a trr trajectory. Returns a dictionary with positions,
velocities, boxes, steps, times, lambdas; all as series (last
array dimension is the frame number)."""
# read first frame
with TRRFile(trr_filename, "r") as trr_file:
frame = next(trr_file)
n_atoms = len(frame.x)
# prepare output series
boxes = np.zeros((3, 3, n_frames))
times = np.zeros(n_frames)
steps = np.zeros(n_frames)
lambdas = np.zeros(n_frames)
if readx:
positions = np.zeros((n_atoms * 3, n_frames))
else:
positions = None
if readv:
velocities = np.zeros((n_atoms * 3, n_frames))
else:
velocities = None
if readf:
forces = np.zeros((n_atoms * 3, n_frames))
else:
forces = None
# read n_frames frames
with TRRFile(trr_filename, "r") as trr_file:
for frame_nr, frame in enumerate(trr_file):
boxes[:, :, frame_nr] = frame.box
times[frame_nr] = frame.time
steps[frame_nr] = frame.step
lambdas[frame_nr] = frame.lmbda
if frame.hasx:
positions[:, frame_nr] = frame.x.flat
if frame.hasv:
velocities[:, frame_nr] = frame.v.flat
if frame.hasf:
forces[:, frame_nr] = frame.f.flat
if frame_nr == n_frames - 1:
break
# check if more frames read than available
else:
raise ValueError(
f"n_frames ({n_frames}) is higher than the number of frames in "
f"{trr_filename} ({frame_nr+1})")
series_dict = {
"positions": positions,
"velocities": velocities,
"forces": forces,
"boxes": boxes,
"steps": steps,
"times": times,
"lambdas": lambdas,
}
return series_dict
def written_trr(tmpdir, trr):
fname = str(tmpdir.join("foo.trr"))
with TRRFile(fname, 'w') as f:
for frame in trr:
natoms = frame.x.shape[0]
f.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
with TRRFile(fname) as f:
yield f
def test_write_different_natoms():
with TRRFile(TRR_multi_frame) as f_in, TRRFile('foo.trr', 'w') as f_out:
assert_equal(f_out.n_atoms, 0)
frame = f_in.read()
natoms = frame.x.shape[0]
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms - 1)
def test_trr():
f = TRRFile(TRR_single_frame)
assert_equal(f.n_atoms, 10)
frame = f.read()
# trr only saves with 3 decimal places precision
assert_array_almost_equal(frame.x.flat, np.arange(30))
assert_array_almost_equal(frame.v.flat, np.arange(30))
assert_array_almost_equal(frame.f.flat, np.arange(30))
assert_array_almost_equal(frame.box, np.eye(3) * 20)
assert_equal(frame.step, 10)
assert_equal(frame.time, 5.0)
assert_equal(frame.lmbda, 0.5)
def test_trr():
f = TRRFile(TRR_single_frame)
assert_equal(f.n_atoms, 10)
frame = f.read()
# trr only saves with 3 decimal places precision
assert_array_almost_equal(frame.x.flat, np.arange(30))
assert_array_almost_equal(frame.v.flat, np.arange(30))
assert_array_almost_equal(frame.f.flat, np.arange(30))
assert_array_almost_equal(frame.box, np.eye(3) * 20)
assert_equal(frame.step, 10)
assert_equal(frame.time, 5.0)
assert_equal(frame.lmbda, 0.5)
def trr_writer(tmpdir, trr):
outname = str(tmpdir.join('foo.trr'))
with TRRFile(outname, 'w') as f_out:
assert f_out.n_atoms == 0
frame = trr.read()
natoms = frame.x.shape[0]
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
yield frame, f_out
def test_write_trr_array_like(tmpdir, array_like, trr):
fname = str(tmpdir.join("foo.trr"))
with TRRFile(fname, 'w') as fout:
for frame in trr:
natoms = frame.x.shape[0]
x = array_like(frame.x)
v = array_like(frame.v)
f = array_like(frame.f)
box = array_like(frame.box)
fout.write(x, v, f, box, frame.step, frame.time, frame.lmbda,
natoms)
def test_write_trr_dtype(tmpdir, dtype, trr):
fname = str(tmpdir.join("foo.trr"))
with TRRFile(fname, 'w') as fout:
for frame in trr:
natoms = frame.x.shape[0]
x = frame.x.astype(dtype)
v = frame.v.astype(dtype)
f = frame.f.astype(dtype)
box = frame.box.astype(dtype)
fout.write(x, v, f, box, frame.step, frame.time, frame.lmbda,
natoms)
def test_write_trr():
with TRRFile(TRR_multi_frame) as f_in, TRRFile('foo.trr', 'w') as f_out:
assert_equal(f_out.n_atoms, 0)
for frame in f_in:
natoms = frame.x.shape[0]
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
with TRRFile('foo.trr') as f:
assert_equal(len(f), 10)
ones = np.ones(30).reshape(10, 3)
box_compare = np.eye(3) * 20
for i, frame in enumerate(f):
assert_array_almost_equal(frame.x, ones * i)
assert_array_almost_equal(frame.v, ones * i + 10)
assert_array_almost_equal(frame.f, ones * i + 20)
assert_array_almost_equal(frame.box, box_compare)
assert_equal(frame.step, i)
assert_equal(frame.time, i * 0.5)
assert_almost_equal(frame.lmbda, .01 * i)
def test_write_different_box_trr(tmpdir, trr):
"""test if we can write different box-sizes for different frames.
"""
orig_box = None
fname = str(tmpdir.join('foo.trr'))
with TRRFile(fname, 'w') as f_out:
assert f_out.n_atoms == 0
frame = trr.read()
natoms = frame.x.shape[0]
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
orig_box = frame.box.copy()
box = frame.box.copy() + 1
f_out.write(frame.x, frame.v, frame.f, box, frame.step, frame.time,
frame.lmbda, natoms)
with TRRFile(fname) as trr:
assert len(trr) == 2
frame_1 = trr.read()
frame_2 = trr.read()
assert_array_almost_equal(frame_1.box, orig_box)
assert_array_almost_equal(frame_1.box + 1, frame_2.box)
def test_write_different_box_trr():
"""test if we can write different box-sizes for different frames.
"""
fname = 'foo.trr'
orig_box = None
with TRRFile(TRR_multi_frame) as f_in, TRRFile(fname, 'w') as f_out:
assert_equal(f_out.n_atoms, 0)
frame = f_in.read()
natoms = frame.x.shape[0]
f_out.write(frame.x, frame.v, frame.f, frame.box, frame.step,
frame.time, frame.lmbda, natoms)
orig_box = frame.box.copy()
box = frame.box.copy() + 1
f_out.write(frame.x, frame.v, frame.f, box, frame.step, frame.time,
frame.lmbda, natoms)
with TRRFile(fname) as trr:
assert_equal(len(trr), 2)
frame_1 = trr.read()
frame_2 = trr.read()
assert_array_almost_equal(frame_1.box, orig_box)
assert_array_almost_equal(frame_1.box + 1, frame_2.box)
def test_read_multi_frame_trr():
with TRRFile(TRR_multi_frame) as f:
f.seek(5)
assert_equal(f.tell(), 5)
assert_equal(len(f), 10)
assert_equal(f.tell(), 5)
f.seek(0)
assert_equal(f.tell(), 0)
ones = np.ones(30).reshape(10, 3)
box_compare = np.eye(3) * 20
for i, frame in enumerate(f):
assert_array_almost_equal(frame.x, ones * i)
assert_array_almost_equal(frame.v, ones * i + 10)
assert_array_almost_equal(frame.f, ones * i + 20)
assert_array_almost_equal(frame.box, box_compare)
assert_equal(frame.step, i)
assert_equal(frame.time, i * 0.5)
assert_almost_equal(frame.lmbda, .01 * i)
def write_series(trr_filename, series_dict):
"""Writes a trr trajectory. Takes a series dictionary
similar to what read_series returns."""
def reshape_if_not_none(array, frame_nr):
if array is None:
return None
else:
return array[:, frame_nr].reshape((-1, 3))
# write frames
with TRRFile(trr_filename, "w") as trr_file:
for frame_nr in range(len(series_dict["steps"])):
positions = reshape_if_not_none(series_dict["positions"], frame_nr)
velocities = reshape_if_not_none(series_dict["velocities"],
frame_nr)
forces = reshape_if_not_none(series_dict["forces"], frame_nr)
box = series_dict["boxes"][:, :, frame_nr]
step = series_dict["steps"][frame_nr]
time = series_dict["times"][frame_nr]
lmbda = series_dict["lambdas"][frame_nr]
n_atoms = len(positions)
trr_file.write(positions, velocities, forces, box, step, time,
lmbda, n_atoms)
def trr():
with TRRFile(TRR_multi_frame) as f:
yield f