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