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("./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_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 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)
