def test_combine_cpptraj_iterating_with_pytraj(self): traj = pt.iterload(fn("tz2.ortho.nc"), fn("tz2.ortho.parm7")) commands = [ 'autoimage', 'rms', ] dslist = CpptrajDatasetList() actlist = ActionList(commands, top=traj.top, dslist=dslist) def get_fi(actlist, traj): '''create a frame iterator with pre-processed by cpptraj ''' for frame in traj: actlist.compute(frame) yield frame ref = traj[3] pt.autoimage(ref, top=traj.top) fi = get_fi(actlist, traj) rmsd_nofit_after_fitting = pt.rmsd_nofit(fi, ref=ref, top=traj.top) t0 = traj[:].autoimage().superpose() saved_rmsd_ = pt.rmsd_nofit(t0, ref=ref) aa_eq(rmsd_nofit_after_fitting, saved_rmsd_)
def test_get_coordinates_trajectory(): '''mutable pytraj.Trajectory ''' traj = pt.Trajectory(xyz=traj_tz2_ortho.xyz, top=traj_tz2_ortho.top) # make a different copy since ``traj`` is mutable traj2 = traj.copy() # all coordinates xyz = pt.get_coordinates(traj) aa_eq(traj.xyz, xyz) # given frames xyz = pt.get_coordinates(traj, frame_indices=[0, 5]) aa_eq(traj[[0, 5]].xyz, xyz) # given frames, autoimage=True xyz = pt.get_coordinates(traj, frame_indices=[0, 5], autoimage=True) aa_eq(traj2[[0, 5]].autoimage().xyz, xyz) # given frames, autoimage=True, rmsfit=ref ref = traj[-3] pt.autoimage(ref, top=traj.top) xyz = pt.get_coordinates(traj, frame_indices=[0, 5], autoimage=True, rmsfit=ref) aa_eq(traj2[[0, 5]].autoimage().superpose(ref=ref).xyz, xyz)
def test_center_autoimagesuperpose(self): traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7") t0 = traj[:] t1 = traj[:] pt.center(t0) pt.autoimage(t0) pt.superpose(t0) aa_eq(pt.center(t1).autoimage().superpose().xyz, t0.xyz)
def test_center_autoimagesuperpose(self): traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) t0 = traj[:] t1 = traj[:] pt.center(t0) pt.autoimage(t0) pt.superpose(t0) aa_eq(pt.center(t1).autoimage().superpose().xyz, t0.xyz)
def test_autoimage_for_tightly_packed_systems(): trajin_fn = os.path.join(cpptraj_test_dir, 'Test_AutoImage', 'G3_3A.rst7') prmtop_fn = os.path.join(cpptraj_test_dir, 'Test_AutoImage', 'nowat.G3_3A.parm7') saved_rst7 = os.path.join(cpptraj_test_dir, 'Test_AutoImage', 'image.G3_3A.rst7.save') saved_traj = pt.iterload(saved_rst7, prmtop_fn) traj = pt.load(trajin_fn, prmtop_fn) pt.autoimage(traj, 'anchor :96 origin') aa_eq(traj.xyz, saved_traj.xyz)
def test_autoimage_with_rmsfit(self): traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) t0 = traj[:] pt.autoimage(t0).superpose() avg_0 = pt.mean_structure(t0, '@CA') avg_1 = pt.mean_structure(traj(autoimage=True, rmsfit=0), '@CA') aa_eq(avg_0.xyz, avg_1.xyz) # 3rd frame # assign traj again t0 = traj[:] pt.autoimage(t0).superpose(ref=3) avg_0 = pt.mean_structure(t0, '@CA') avg_1 = pt.mean_structure(traj(autoimage=True, rmsfit=3), '@CA') avg_2 = pt.mean_structure(traj, autoimage=True, rmsfit=3, mask='@CA') aa_eq(avg_0.xyz, avg_1.xyz) aa_eq(avg_0.xyz, avg_2.xyz) # 3rd frame, frame_indices # assign traj again frame_indices = [0, 8, 5] t0 = traj[frame_indices] t1 = traj[frame_indices] t0.autoimage().superpose(ref=-1) avg_0 = pt.mean_structure(t0, '@CA') # use ref=5 which correspond to original index # try with pytraj.TrajectoryIterator avg_1 = pt.mean_structure( traj, autoimage=True, rmsfit=5, mask='@CA', frame_indices=frame_indices) # try with pytraj.Trajectory avg_2 = pt.mean_structure(t1, autoimage=True, rmsfit=-1, mask='@CA') avg_3 = pt.mean_structure( traj[:], autoimage=True, rmsfit=5, mask='@CA', frame_indices=frame_indices) aa_eq(avg_0.xyz, avg_1.xyz) aa_eq(avg_0.xyz, avg_2.xyz) aa_eq(avg_0.xyz, avg_3.xyz)
def view_structure(self): """Visualize a restart file with NGL""" log.debug( f"Visualizing top {self.topology.value} and restart {self.rst_traj.value}" ) # load rst7 with pytraj (NGL cannot read it directly) traj = pt.load(os.path.join(self.md_dir.value, self.rst_traj.value), os.path.join(self.md_dir.value, self.topology.value)) traj = pt.autoimage(traj) ## pass to parmed to write the pdb data in a StringIO # struct = pmd.load_file(os.path.join(self.md_dir.value, self.topology.value), xyz=traj.xyz) # f = StringIO() # struct.write_pdb(f) # pdb_data = f.getvalue().encode("ascii") # f.close() # write as pdb to temporary file (much faster than parmed + StringIO) with NamedTemporaryFile() as f: pt.write_traj(f.name, traj, format="pdb", overwrite=True) pdb_data = f.read() # create javascript code with open(os.path.join(self.src_dir, "static", "js", "nglviewer.js")) as f: JS_TEMPLATE = f.read() self.js_view_structure.code = JS_TEMPLATE % (pdb_data) # trigger javascript callback by adding an invisible character to the button label self.view_button.label += " "
def test_autoimage_with_rmsfit(self): traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7") t0 = traj[:] pt.autoimage(t0).superpose() avg_0 = pt.mean_structure(t0, '@CA') avg_1 = pt.mean_structure(traj(autoimage=True, rmsfit=0), '@CA') aa_eq(avg_0.xyz, avg_1.xyz) # 3rd frame # assign traj again t0 = traj[:] pt.autoimage(t0).superpose(ref=3) avg_0 = pt.mean_structure(t0, '@CA') avg_1 = pt.mean_structure(traj(autoimage=True, rmsfit=3), '@CA') avg_2 = pt.mean_structure(traj, autoimage=True, rmsfit=3, mask='@CA') aa_eq(avg_0.xyz, avg_1.xyz) aa_eq(avg_0.xyz, avg_2.xyz) # 3rd frame, frame_indices # assign traj again frame_indices = [0, 8, 5] t0 = traj[frame_indices] t1 = traj[frame_indices] t0.autoimage().superpose(ref=-1) avg_0 = pt.mean_structure(t0, '@CA') # use ref=5 which correspond to original index # try with pytraj.TrajectoryIterator avg_1 = pt.mean_structure(traj, autoimage=True, rmsfit=5, mask='@CA', frame_indices=frame_indices) # try with pytraj.Trajectory avg_2 = pt.mean_structure(t1, autoimage=True, rmsfit=-1, mask='@CA') avg_3 = pt.mean_structure(traj[:], autoimage=True, rmsfit=5, mask='@CA', frame_indices=frame_indices) aa_eq(avg_0.xyz, avg_1.xyz) aa_eq(avg_0.xyz, avg_2.xyz) aa_eq(avg_0.xyz, avg_3.xyz)
def test_pytraj_do(): traj = pt.iterload(fn("tz2.nc"), fn("tz2.parm7")) ref0 = pt.autoimage(traj[0], top=traj.top) ref1 = pt.autoimage(traj[1], top=traj.top) data = pt.tools.dict_to_ndarray(pt.compute( ['autoimage', 'radgyr nomax @CA', 'rms refindex 0', 'rms refindex 1'], traj, ref=[ref0, ref1])) t0 = pt.autoimage(traj[:]) aa_eq(pt.radgyr(t0, '@CA'), data[0]) aa_eq(pt.rmsd(t0, ref=ref0), data[1]) aa_eq(pt.rmsd(t0, ref=ref1), data[2]) aa_eq(pt.compute(pt.radgyr, t0, '@CA'), data[0])
def test_rmsfit_with_autoimage(self): traj = self.traj.copy() # test rmsfit and autoimage ref0 = traj[0] # need to autoimage reference frame first pt.autoimage(ref0, top=traj.top) for chunk in traj.iterchunk(chunksize=2, rmsfit=(ref0, self.mask), autoimage=True): pass fa0 = traj[-2:] fa0.autoimage() ref01 = traj[0] pt.autoimage(ref01, top=traj.top) fa0.rmsfit(ref=ref01, mask=self.mask) aa_eq(chunk.xyz, fa0.xyz[-2:])
def test_action_bounds(self): # creat mutable trajectory traj = pt.load(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) pt.autoimage(traj) pt.superpose(traj, ref=0, mask=':1-13&!@H=', mass=True) grid_data = pt._grid(traj, mask=':1-13', grid_spacing=[0.5, 0., 0.]) text = ''' parm {} trajin {} autoimage rms first :1-13&!@H= mass bounds :1-13 dx .5 name MyGrid '''.format(fn('tz2.ortho.parm7'), fn('tz2.ortho.nc')) state = pt.load_cpptraj_state(text) state.run() cpp_grid = state.data['MyGrid'].values aa_eq(cpp_grid, grid_data)
def test_get_coordinates_trajecotoryiterator(self): '''immutable pytraj.TrajectoryIterator ''' traj = self.traj_tz2_ortho.copy() # all coordinates xyz = pt.get_coordinates(traj) aa_eq(traj.xyz, xyz) # given frames xyz = pt.get_coordinates(traj, frame_indices=[0, 5]) aa_eq(traj[[0, 5]].xyz, xyz) # given frames, autoimage=True xyz = pt.get_coordinates(traj, frame_indices=[0, 5], autoimage=True) aa_eq(traj[[0, 5]].autoimage().xyz, xyz) # given frames, autoimage=True, rmsfit=ref ref = traj[-3] pt.autoimage(ref, top=traj.top) xyz = pt.get_coordinates(traj, frame_indices=[0, 5], autoimage=True, rmsfit=ref) aa_eq(traj[[0, 5]].autoimage().superpose(ref=ref).xyz, xyz) xyz = pt.get_coordinates(traj, frame_indices=range(3), autoimage=True, rmsfit=ref) # with mask xyz = pt.get_coordinates(traj, mask='@CA') aa_eq(xyz, traj['@CA'].xyz) # slow fi = pt.pipe(traj, ['autoimage']) aa_eq(pt.get_coordinates(fi), traj[:].autoimage().xyz) # raise self.assertRaises( ValueError, lambda: pt.get_coordinates(traj(), frame_indices=[0, 2]))
def test_pytraj_do(self): traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7") ref0 = pt.autoimage(traj[0], top=traj.top) ref1 = pt.autoimage(traj[1], top=traj.top) data = pt.tools.dict_to_ndarray(pt.compute( ['autoimage', 'radgyr nomax @CA', 'rms refindex 0', 'rms refindex 1'], traj, ref=[ref0, ref1])) t0 = pt.autoimage(traj[:]) aa_eq(pt.radgyr(t0, '@CA'), data[0]) aa_eq(pt.rmsd(t0, ref=ref0), data[1]) aa_eq(pt.rmsd(t0, ref=ref1), data[2]) # pytraj's method aa_eq(pt.compute(pt.radgyr, t0, '@CA'), data[0])
def test_combine_cpptraj_iterating_with_pytraj(self): traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7") commands = ['autoimage', 'rms', ] dslist = CpptrajDatasetList() actlist = ActionList(commands, top=traj.top, dslist=dslist) def get_fi(actlist, traj): '''create a frame iterator with pre-processed by cpptraj ''' for frame in traj: actlist.compute(frame) yield frame ref = traj[3] pt.autoimage(ref, top=traj.top) fi = get_fi(actlist, traj) rmsd_nofit_after_fitting = pt.rmsd_nofit(fi, ref=ref, top=traj.top) t0 = traj[:].autoimage().superpose() saved_rmsd_ = pt.rmsd_nofit(t0, ref=ref) aa_eq(rmsd_nofit_after_fitting, saved_rmsd_)
def test_combination_of_differnt_transformations(self): traj_on_disk = pt.datafiles.load_tz2_ortho() traj_on_disk_2 = pt.datafiles.load_tz2_ortho() traj_on_mem = pt.datafiles.load_tz2_ortho()[:] ref = pt.autoimage(traj_on_disk[:1]) # note: if using autoimage, must provide pre-processed reference (traj_on_mem.autoimage().superpose(ref=ref).scale('x 1.2')) (traj_on_disk.autoimage().superpose(ref=ref).scale('x 1.2')) aa_eq(traj_on_disk.xyz, traj_on_mem.xyz) # remove traj_on_disk._remove_transformations() aa_eq(traj_on_disk.xyz, traj_on_disk_2.xyz) assert len(traj_on_disk._cdslist) == 0
O_idx = pt.select_atoms(traj.top, ':WAT@O') Solute_idx = pt.select_atoms(traj.top, sol_mask) N_frames = traj.n_frames ### Key : distance ### Item: [water_O_idx, frame_idx, count] crd_dict = dict() for d in range(int((stop-start)/step)): crd_dict[start+d*step] = list() for frame in traj.iterframe(): pt.autoimage(frame, top=traj.topology) solute = frame[Solute_idx] solvent = frame[O_idx] dists = scipy_spatial.distance.cdist(solvent, solute) for dist, l in crd_dict.items(): valids = np.where(dists < dist)[0] valids = np.unique(valids) N_valids = valids.shape[0] if N_valids > 0: for N_i in range(N_valids):
def test_autoimage_regular(): # OK for frame in traj(stop=400): pt.autoimage(frame, top=traj.top)
def compute_rmsd( itraj: str, itop: Optional[str] = None, iref: Optional[str] = None, lig_mask: str = default_mask, reimage: bool = False, verbose: bool = False, ) -> np.ndarray: """ Compute RMSD for a trajectory with respect to a reference structure. Args: itraj (str): Trajectory file name itop (str, optional): Topology file name iref (str, optional): Reference file name lig_mask (str, optional): Selection mask (in `pytraj` format) for the ligand reimage (bool): Re-image coordinates according to PBC Returns: Returns a `np.ndarray` containing the frame number, an the RMSD (in angstrom) with respect to the reference structure `iref`. """ if verbose: print("Loading trajectory...", file=sys.stdout, end="") lig_traj = load_traj(itraj, itop, mask=lig_mask) if verbose: print("done", file=sys.stdout) if iref is not None: if verbose: print("Loading reference...", file=sys.stdout, end="") lig_ref = load_ref(iref, itop, mask=lig_mask) if verbose: print("done", file=sys.stdout) else: lig_ref = 0 # Autoimage (for PBC) if reimage: if verbose: print("Reimaging...", file=sys.stdout, end="") lig_traj = pt.autoimage(lig_traj) if iref is not None: lig_ref = pt.autoimage(lig_ref) if verbose: print("done", file=sys.stdout) # TODO: Align trajectory with reference structure # (needs to load the whole trajectory) # Compute RMSD (symmetrized) if verbose: print("Computing RMSD...", file=sys.stdout, end="") rmsd = pt.analysis.rmsd.symmrmsd( lig_traj, mask=lig_mask, ref=lig_ref, ref_mask=lig_mask, fit=False ) if verbose: print("done", file=sys.stdout) # TODO: Add time return np.stack((np.arange(0, lig_traj.n_frames), rmsd), axis=1)