def test_get_velocities_and_forces():
"""Write data with velocities and forces, and read it back"""
# NOTE: this is a test of a hidden API
xyz = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
vel = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
forces = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
box = np.array(np.random.randn(500, 3, 3), dtype=np.float32)
time = np.array(np.random.randn(500), dtype=np.float32)
step = np.array(np.arange(500), dtype=np.int32)
lambd = np.array(np.random.randn(500), dtype=np.float32)
with TRRTrajectoryFile(temp, 'w') as f:
f._write(xyz=xyz,
time=time,
step=step,
box=box,
lambd=lambd,
vel=vel,
forces=forces)
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2, vel2, forces2 = f._read(
n_frames=500,
atom_indices=None,
get_velocities=True,
get_forces=True)
eq(xyz, xyz2)
eq(vel, vel2)
eq(forces, forces2)
eq(time, time2)
eq(step, step2)
eq(lambd, lambd2)
def test_get_velocities_forces_atom_indices_2():
# NOTE: this is a test of a hidden API
xyz = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
vel = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
forces = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
box = np.array(np.random.randn(500, 3, 3), dtype=np.float32)
time = np.array(np.random.randn(500), dtype=np.float32)
step = np.array(np.arange(500), dtype=np.int32)
lambd = np.array(np.random.randn(500), dtype=np.float32)
with TRRTrajectoryFile(temp, 'w') as f:
f._write(xyz=xyz,
time=time,
step=step,
box=box,
lambd=lambd,
vel=vel,
forces=forces)
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2, vel2, forces2 = f._read(
n_frames=500,
atom_indices=slice(None, None, 2),
get_velocities=True,
get_forces=True)
eq(xyz[:, ::2], xyz2)
eq(vel[:, ::2], vel2)
eq(forces[:, ::2], forces2)
eq(time, time2)
eq(step, step2)
eq(lambd, lambd2)
pass
def test_seek(get_fn):
reference = TRRTrajectoryFile(get_fn('frame0.trr')).read()[0]
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
eq(len(f), len(reference))
eq(len(f.offsets), len(reference))
eq(f.tell(), 0)
eq(f.read(1)[0][0], reference[0])
eq(f.tell(), 1)
xyz = f.read(1)[0][0]
eq(xyz, reference[1])
eq(f.tell(), 2)
f.seek(0)
eq(f.tell(), 0)
xyz = f.read(1)[0][0]
eq(f.tell(), 1)
eq(xyz, reference[0])
f.seek(5)
eq(f.read(1)[0][0], reference[5])
eq(f.tell(), 6)
f.seek(-5, 1)
eq(f.tell(), 1)
eq(f.read(1)[0][0], reference[1])
def read_frame_data(self, filename, frame_num):
"""
Returns pos, vel, box or raises error
"""
# if self._last_filename != filename:
# try:
# self._file.close()
# except AttributeError:
# pass # first time thru, self._file is None
# self._file = TRRTrajectoryFile(filename)
# f = self._file
# do we need to reopen the TRR each time to avoid problems with the
# fiel length changing?
trr = TRRTrajectoryFile(filename)
f = trr
logger.debug("Reading file %s frame %d (of %d)", filename, frame_num,
len(f))
# logger.debug("File length: %d", len(f))
try:
f.seek(offset=frame_num)
data = f._read(n_frames=1, atom_indices=None, get_velocities=True)
finally:
trr.close()
return data[0][0], data[5][0], data[3][0]
def test_read_stride(get_fn):
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz, time, step, box, lambd = f.read()
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz3, time3, step3, box3, lambd3 = f.read(stride=3)
assert eq(xyz[::3], xyz3)
assert eq(step[::3], step3)
assert eq(box[::3], box3)
assert eq(time[::3], time3)
def test_read_stride_n_frames(get_fn):
# trr read stride when n_frames is supplied (different path)
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz, time, step, box, lambd = f.read()
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz3, time3, step3, box3, lambd3 = f.read(n_frames=1000, stride=3)
assert eq(xyz[::3], xyz3)
assert eq(step[::3], step3)
assert eq(box[::3], box3)
assert eq(time[::3], time3)
def test_read_atomindices_1():
with TRRTrajectoryFile(temp) as f:
xyz, time, step, box, lambd = f.read()
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2 = f.read(atom_indices=[0, 1, 2])
assert eq(xyz[:, [0, 1, 2]], xyz2)
assert eq(step, step2)
assert eq(box, box2)
assert eq(lambd, lambd2)
assert eq(time, time2)
def test_read_atomindices_2():
with TRRTrajectoryFile(temp) as f:
xyz, time, step, box, lambd = f.read()
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2 = f.read(atom_indices=slice(None, None, 2))
assert eq(xyz[:, ::2], xyz2)
assert eq(step, step2)
assert eq(box, box2)
assert eq(lambd, lambd2)
assert eq(time, time2)
def test_read_stride_offsets(get_fn):
# read xtc with stride and offsets
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz, time, step, box, lambd = f.read()
for s in (1, 2, 3, 4, 5):
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
f.offsets # pre-compute byte offsets between frames
xyz_s, time_s, step_s, box_s, lamb_s = f.read(stride=s)
assert eq(xyz_s, xyz[::s])
assert eq(step_s, step[::s])
assert eq(box_s, box[::s])
assert eq(time_s, time[::s])
def test_read_forces_do_not_exist():
"""Requesting forces from a file that does not have them"""
# NOTE: this is a test of a hidden API
xyz = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
box = np.array(np.random.randn(500, 3, 3), dtype=np.float32)
time = np.array(np.random.randn(500), dtype=np.float32)
step = np.array(np.arange(500), dtype=np.int32)
lambd = np.array(np.random.randn(500), dtype=np.float32)
with TRRTrajectoryFile(temp, 'w') as f:
f.write(xyz=xyz, time=time, step=step, box=box, lambd=lambd)
with TRRTrajectoryFile(temp) as f:
with pytest.raises(RuntimeError):
f._read(n_frames=500, atom_indices=None, get_velocities=False, get_forces=True)
def test_1():
# Write data and read it back
xyz = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
time = np.random.randn(500)
step = np.arange(500)
lambd = np.random.randn(500)
with TRRTrajectoryFile(temp, 'w') as f:
f.write(xyz=xyz, time=time, step=step, lambd=lambd)
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2 = f.read()
assert eq(xyz, xyz2)
assert eq(time, time2)
assert eq(step, step2)
assert eq(lambd, lambd2)
def test_tell(get_fn):
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
eq(f.tell(), 0)
f.read(101)
eq(f.tell(), 101)
f.read(3)
eq(f.tell(), 104)
def test_deficient_shape():
# Write data one frame at a time. This checks how the shape is dealt with,
# because each frame is deficient in shape.
xyz = np.array(np.random.randn(500, 50, 3), dtype=np.float32)
time = np.random.randn(500)
step = np.arange(500)
lambd = np.random.randn(500)
with TRRTrajectoryFile(temp, 'w') as f:
for i in range(len(xyz)):
f.write(xyz=xyz[i], time=time[i], step=step[i], lambd=lambd[i])
with TRRTrajectoryFile(temp) as f:
xyz2, time2, step2, box2, lambd2 = f.read()
assert eq(xyz, xyz2)
assert eq(time, time2)
assert eq(step, step2)
assert eq(lambd, lambd2)
def test_ragged_2():
# try first writing no box vectors, and then adding some
xyz = np.random.randn(100, 5, 3)
time = np.random.randn(100)
box = np.random.randn(100, 3, 3)
with TRRTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time=time, box=box)
assert_raises(ValueError, lambda: f.write(xyz))
def test_malformed():
with open(temp, 'w') as tmpf:
tmpf.write("foo") # very badly malformed TRR
with pytest.raises(IOError):
TRRTrajectoryFile(temp)
psutil = pytest.importorskip("psutil")
open_files = psutil.Process().open_files()
paths = [os.path.realpath(f.path) for f in open_files]
assert os.path.realpath(temp) not in paths
def test_read_stride_switching(get_fn):
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
xyz, time, step, box, lambd = f.read()
with TRRTrajectoryFile(get_fn('frame0.trr')) as f:
f.offsets # pre-compute byte offsets between frames
# read the first 10 frames with stride of 2
s = 2
n_frames = 10
xyz_s, time_s, step_s, box_s, lamb_s = f.read(n_frames=n_frames, stride=s)
assert eq(xyz_s, xyz[:n_frames*s:s])
assert eq(step_s, step[:n_frames*s:s])
assert eq(box_s, box[:n_frames*s:s])
assert eq(time_s, time[:n_frames*s:s])
# now read the rest with stride 3, should start from frame index 8.
# eg. np.arange(0, n_frames*s + 1, 2)[-1] == 18
offset = f.tell()
assert offset == 20
s = 3
xyz_s, time_s, step_s, box_s, lamb_s = f.read(n_frames=None, stride=s)
assert eq(xyz_s, xyz[offset::s])
assert eq(step_s, step[offset::s])
assert eq(box_s, box[offset::s])
assert eq(time_s, time[offset::s])
def write_frame_to_file(self, filename, snapshot, mode='w'):
if os.path.isfile(filename):
# stop if we already have this file; could also happen because
# of a weird behavior in a mover. You must remove the files if
# you don't want them.
raise RuntimeError("File " + str(filename) + " exists. " +
"Preventing overwrite.")
# type control before passing things to Cython code
xyz = np.asarray([snapshot.xyz], dtype=np.float32)
time = np.asarray([0.0], dtype=np.float32)
step = np.asarray([0], dtype=np.int32)
box = np.asarray([snapshot.box_vectors], dtype=np.float32)
lambd = np.asarray([0.0], dtype=np.float32)
vel = np.asarray([snapshot.velocities], dtype=np.float32)
try:
trr = TRRTrajectoryFile(filename, mode)
trr._write(xyz, time, step, box, lambd, vel)
finally:
trr.close()
close_file_descriptors(filename)
import numpy as np
import sys
import math
from mdtraj.formats import XTCTrajectoryFile
from mdtraj.formats import TRRTrajectoryFile
trrFileName = sys.argv[1]
oFileName = sys.argv[2]
with TRRTrajectoryFile(trrFileName) as trrFile:
coordtMat, tVec, step, boxMat, lambdaVec = trrFile.read()
trrFile.close()
nFrame, mMol, xyz = np.shape(coordtMat)
print 'nFrame = {}'.format(nFrame)
apm = 1
mMol /= apm
Lx = boxMat[0][0][0]
Ly = boxMat[0][1][1]
Lz = boxMat[0][2][2]
Lx2 = Lx * 0.5
Ly2 = Ly * 0.5
Lz2 = Lz * 0.5
# Setting PMat
dz = 0.02
zMin = -8. + dz * 0.5
zMax = 8. - dz * 0.5
zVec = np.zeros(int((zMax - zMin) / dz) + 1)
denVec = np.zeros(len(zVec))
async def _launch_traj(self, projname: str) -> dict:
"""Launch a trajectory with the current state to completion.
Launch a trajectory using the current state with the given command in
a new process. Runs in the given working directory. Waits for its
completion with async, then checks for failures or warnings.
Parameters
----------
projname
The unique project name. No other project should have this name
Returns
-------
A dictionary with the keys:
"commit": basin integer the trajectory committed to or None if it
did not commit
"frames": np.array with the +delta_t and +2delta_t xyz frames. Has
the shape (2, n_atoms, 3)
Raises
------
RuntimeError
If CP2K fails to run.
"""
# We are saving the state of the class before calling a method with
# async.sleep in a local variable so it is not changed out from underneath
# us. Any call to async.sleep gives an opportunity for another async method
# to modify this class. All other variables are safe, but the pin_offset
# is in contention between the forwards and reverse, so we save it here.
pin_offset = str(self.pin_offset)
tpr_path = await self._run_grompp(projname)
# Set the name for the committor output and write the unique plumed file
plumed_out_name = f"{projname}_plumed.out"
plumed_in_path = os.path.join(self.working_dir,
f"{projname}_plumed.dat")
self.plumed_handler.write_plumed(plumed_in_path, plumed_out_name)
command_list = ["-s", tpr_path, "-plumed", plumed_in_path, "-deffnm", projname]
if self.should_pin:
# total_instances * 2 because each has a forward and reverse mdrun
command_list.extend(["-pinoffset", pin_offset, "-pinstride",
str(self.total_instances * 2), "-pin", "on"])
# run
proc = await self._open_md_and_wait(command_list, projname)
plumed_out_path = os.path.join(self.working_dir, plumed_out_name)
# Check if there was a fatal error that wasn't caused by a committing
# basin
if proc.returncode != 0:
stdout, stderr = proc.communicate()
stdout_msg = stdout.decode('ascii')
stderror_msg = stderr.decode('ascii')
# Copy the output file to a place we can see it
failed_log = os.path.join(self.working_dir, f"{projname}.log")
copied_log = f"{projname}_FATAL.log"
with open(copied_log, "a") as out:
with open(failed_log, "r") as f:
out.write(f.read())
out.write("\nFAILURE \n")
out.write("STDOUT: \n")
out.write(stdout_msg)
out.write("\nSTDERR: \n")
out.write(stderror_msg)
self.logger.warning("Trajectory %s exited fatally:\n stdout: %s\n stderr: %s",
projname, stdout_msg, stderror_msg)
raise RuntimeError(f"Trajectory {projname} failed")
# TODO: check warnings in gromacs log file
parser = PlumedOutputHandler(plumed_out_path)
basin = parser.check_basin()
if basin is not None:
self.logger.info("Trajectory %s committed to basin %s", projname,
basin)
else:
self.logger.info("Trajectory %s did not commit before simulation ended",
projname)
try:
traj_path = os.path.join(self.working_dir, f"{projname}.trr")
with TRRTrajectoryFile(traj_path, "r") as file:
xyz, _, _, box, _ = file.read(3, stride=1)
# Convert from nm read to A
xyz *= 10
box *= 10
# return last two frames of the three read
return {"commit": basin,
"frames": xyz[1:, :, :]}
except EOFError:
self.logger.warning("Required frames could not be be read from the"
" output trajectory. This may be cased by a delta_t"
" that is too large where the traj committed to a"
" basin before 2*delta_t fs or a simulation wall time"
" that is too short and exited before reaching 2*delta_t fs")
return None
