def test_iterframe_from_array(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
# no mass
fi = pt.iterframe_from_array(traj.xyz, traj.n_atoms,
range(traj.n_frames))
for f_traj, (_, f_fi) in zip(traj, fi):
aa_eq(f_traj.xyz, f_fi.xyz)
# f_fi mass must be [1.0, ...]
aa_eq(f_fi.mass, [1.0 for _ in range(traj.n_atoms)])
# mass from Topology
fi = pt.iterframe_from_array(traj.xyz, traj.n_atoms,
range(traj.n_frames), traj.top)
for f_traj, (_, f_fi) in zip(traj, fi):
aa_eq(f_traj.xyz, f_fi.xyz)
aa_eq(f_fi.mass, f_traj.mass)
aa_eq(f_fi.mass, traj.top.mass)
# mass from array
fi = pt.iterframe_from_array(traj.xyz,
traj.n_atoms,
range(traj.n_frames),
mass=traj.top.mass)
for f_traj, (_, f_fi) in zip(traj, fi):
aa_eq(f_traj.xyz, f_fi.xyz)
aa_eq(f_fi.mass, f_traj.mass)
aa_eq(f_fi.mass, traj.top.mass)
def test_ired_simple_for_coverage():
'''
'''
traj = pt.iterload(traj_dir, parm_dir)
h_indices = pt.select_atoms('@H', traj.top)
n_indices = pt.select_atoms('@H', traj.top) - 1
nh_indices = list(zip(n_indices, h_indices))
vecs_and_mat = pt.ired_vector_and_matrix(traj, nh_indices, dtype='tuple')
vecs_and_mat = pt.ired_vector_and_matrix(traj, nh_indices, dtype='tuple')
state_vecs = vecs_and_mat[0]
mat_ired = vecs_and_mat[1]
# get eigenvalues and eigenvectors
modes = pt.matrix.diagonalize(mat_ired, n_vecs=len(state_vecs))
evals, evecs = modes
data_0 = _ired(
state_vecs, modes=(evals, evecs), NHbond=True, tcorr=10000, tstep=1.)
data_1 = _ired(state_vecs, modes=modes, NHbond=True, tcorr=10000, tstep=1)
for d0, d1 in zip(data_0, data_1):
if d0.dtype not in [
'modes',
]:
aa_eq(d0.values, d1.values)
else:
# modes
# values: tuple
aa_eq(d0.values[0], d1.values[0])
aa_eq(d0.values[1], d1.values[1])
# try different dtype
out_try_new_dtype = pt.ired_vector_and_matrix(
traj, nh_indices, dtype='cpptraj_dataset')
def test_hbonds_from_pdb(self):
traj = pt.load(fn('1L2Y.pdb'))
hb = pt.search_hbonds(traj)
state = pt.load_cpptraj_state('''
parm {pdb}
trajin {pdb}
hbond series
'''.format(pdb=fn('1L2Y.pdb')))
state.run()
for data_p, data_cpp in zip(hb.data, state.data[1:]):
assert len(
data_p) == traj.n_frames == 38, 'size of dataset must be 38'
aa_eq(data_p, data_cpp)
# make sure distances are smaller than cutoff
distance_cutoff = 2.5
angle_cutoff = 135.
hb = pt.search_hbonds(traj)
distances = pt.distance(traj, hb.get_amber_mask()[0])
angles = pt.angles(traj, hb.get_amber_mask()[1])
dist_indices = np.where(distances > distance_cutoff)
angle_indices = np.where(angles < angle_cutoff)
saved_donor_acceptors = [
'ASP9_OD2-ARG16_NH1-HH12', 'ASP9_OD2-ARG16_NH2-HH22',
'ASP9_OD2-ARG16_NE-HE', 'ASP9_OD2-ARG16_NH2-HH21',
'ASP9_OD2-ARG16_NH1-HH11'
]
donor_acceptors = pt.search_hbonds(traj, ':9,16').donor_acceptor
assert saved_donor_acceptors == donor_acceptors, 'saved_donor_acceptors'
aa_eq(hb.total_solute_hbonds(), hb.data['total_solute_hbonds'])
def test_real_box(self):
traj = pt.load(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
trajiter = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
saved_box = Box(
[3.94559740E+01, 4.68215170E+01, 4.04695410E+01, 90., 90., 90.])
aa_eq(traj.top.box.values, saved_box.values)
for frame in traj:
assert frame.box.type == 'ortho'
aa_eq(
frame.box.values,
[35.2627796623, 41.8455476799, 36.168629529, 90.0, 90.0, 90.0],
decimal=1)
arr0 = traj.unitcells
arr1 = trajiter.unitcells
for b0, b1, frame in zip(arr0, arr1, trajiter):
box = frame.box
# FIXME:
# python3)
b2 = box.values
aa_eq(b0, b1)
aa_eq(b0, b2)
# test assign Box with list/tuple
b = Box(saved_box.values)
b2 = Box((t for t in saved_box.values))
aa_eq(b.values, saved_box.values, decimal=7)
aa_eq(b2.values, saved_box.values, decimal=7)
# assign frame.box with list/tuple
frame.box = b.values
b3 = frame.box
aa_eq(b3.values, saved_box.values, decimal=7)
def test_basic(self):
'''slicing, select'''
top = pt.load_topology(fn('Tc5b.top'))
#
assert isinstance(top[0], Atom)
assert isinstance(top[:2], pt.Topology)
assert isinstance(top[:1], pt.Topology)
assert isinstance(top[range(10)], pt.Topology)
assert isinstance(top[list(range(10))], pt.Topology)
assert isinstance(top[np.array(range(10))], pt.Topology)
assert top[0].name == top['@1'][0].name
# mask, AtomMask, python array, list
atm = top("@CA")
indices = atm.indices
for a1, a2, a3, a4 in zip(top['@CA'], top[atm], top[indices],
top[list(indices)]):
assert a1.name == a2.name == a3.name == a4.name == 'CA'
# check len
assert len(top[:]) == top.n_atoms
assert len(top[:10]) == 10
# API
top.bond_indices
top.angle_indices
top.dihedral_indices
def test_convert_to_dict_and_rebuild(self):
'''test_convert_to_dict_and_rebuild
'''
top = self.traj.top
d = top.to_dict()
new_top = pt.Topology()
MOLNUM = 0
for idx, (aname, atype, charge, mass, resid, resname,
mol_number) in enumerate(
zip(d['atom_name'], d['atom_type'], d['atom_charge'], d[
'atom_mass'], d['resid'], d['resname'], d[
'mol_number'])):
atom = pt.core.Atom(
name=aname, type=atype, charge=charge, mass=mass, resid=resid)
atom.set_mol(mol_number)
residue = pt.core.Residue(resname, resid)
new_top.add_atom(atom, residue)
if idx == 0:
new_top.start_new_mol()
if mol_number > MOLNUM:
new_top.start_new_mol()
MOLNUM += 1
new_top.add_bonds(d['bond_index'])
new_top.add_dihedrals(d['dihedral_index'])
new_top.box = Box(top.box.values)
assert_equal_topology(top, new_top, self.traj)
def test_0_trajiter(self):
traj = io.load_sample_data("tz2")
from pytraj.externals.six import zip
for idx, (f0, f1) in enumerate(zip(traj, traj)):
pass
assert idx > 0
# tee
ilist = it.tee(traj, 2)
for idx, (f0, f1) in enumerate(zip(ilist[0], ilist[1])):
assert f0.rmsd(f1) < 1E-5
assert idx == traj.n_frames - 1
# product
it.product(traj, repeat=2)
def test_simplifed_topology(self):
'''simplify'''
top = pt.load_topology(fn('Tc5b.top'))
sim_top = top.simplify()
for atom, sim_atom in zip(top.atoms, sim_top.atoms):
assert atom.resname == sim_atom.resname, 'equal resname'
assert atom.name == sim_atom.name, 'equal resname'
assert atom.type == sim_atom.type, 'equal resname'
assert atom.charge == sim_atom.charge, 'equal resname'
assert atom.mass == sim_atom.mass, 'equal resname'
def test_autoimage_and_slicing(self):
traj_on_disk = pt.datafiles.load_tz2_ortho()
traj_on_mem = traj_on_disk[:]
aa_eq(traj_on_disk.xyz, traj_on_mem.xyz)
aa_eq(traj_on_disk.autoimage().xyz, traj_on_mem.autoimage().xyz)
aa_eq(traj_on_mem[:1].xyz, traj_on_disk[:1].xyz)
aa_eq(traj_on_mem[:].xyz, traj_on_disk[:].xyz)
from pytraj.externals.six import zip
for f0, f1 in zip(traj_on_disk(0, 8, 2), traj_on_mem(0, 8, 2)):
aa_eq(f0.xyz, f1.xyz)
def test_0(self):
traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
fa1 = traj[:]
# do inplace-autoimage for Trajectory
fa1.autoimage()
# fa2 will store frame when iterating `traj` (TrajectoryIterator)
fa2 = Trajectory()
fa2.top = traj.top.copy()
# frame_iter
for frame in traj(autoimage=True):
fa2.append(frame)
aa_eq(fa1.xyz, fa2.xyz)
# frame_iter with mask and autoimage
fa3 = fa2.copy()
# fa3 is a copy of autoimaged fa2. then we strip all but CA atoms
# just want to make sure we can use `mask`
fa3.strip("!@CA")
fa4 = Trajectory()
fa4.top = fa3.top.copy()
for frame in traj(mask='@CA', autoimage=True):
fa4.append(frame)
aa_eq(fa3.xyz, fa4.xyz)
# frame_iter with mask and autoimage, and rmsfit
# fa3 is a copy of autoimaged fa2. then we strip all but CA atoms
# just want to make sure we can use `mask`
fa3 = traj[:]
fa3.autoimage()
fa3.rmsfit(ref=5, mask='@CB')
fa3.strip("!@CA")
fa4 = Trajectory()
fa4.top = fa3.top.copy()
ref1 = traj[5].copy()
for frame in traj(mask='@CA', autoimage=True, rmsfit=(ref1, '@CB')):
fa4.append(frame)
aa_eq(fa3.xyz, fa4.xyz)
for f0, f1 in zip(fa3, fa4):
assert f0.rmsd_nofit(f1) < 1E-7
def test_0(self):
# TrajectoryIterrator
# status: failed
from pytraj.externals.six import zip
traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
pt.write_traj("./output/tz2.autoimage_with_rmsfit.nc",
traj(autoimage=True, rmsfit=(0, '@CA,C,N')),
overwrite=True)
saved_traj = pt.load(fn('tz2.autoimage_with_rmsfit.nc'), traj.top)
p_traj = pt.load('./output/tz2.autoimage_with_rmsfit.nc', traj.top)
aa_eq(saved_traj.xyz, p_traj.xyz)
for f1, f2 in zip(p_traj, saved_traj):
pass
def test_on_disk_trajectory(self):
# TrajectoryIterrator
output = "ok_to_delete.nc"
traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
pt.write_traj(
output,
traj(autoimage=True, rmsfit=(0, '@CA,C,N')),
overwrite=True)
saved_traj = pt.load(fn('tz2.autoimage_with_rmsfit.nc'), traj.top)
p_traj = pt.load(output, traj.top)
aa_eq(saved_traj.xyz, p_traj.xyz)
# ensure iterable without memory freeing
for f1, f2 in zip(p_traj, saved_traj):
aa_eq(f1.xyz, f2.xyz)
def test_ired_need_lapack_cpptraj():
state = pt.load_cpptraj_state(txt)
state.run()
xyz = state.data['CRD1'].xyz
top = state.data['CRD1'].top
traj = pt.Trajectory(xyz=xyz, top=top)
state_vecs = state.data[1:-3].values
h_indices = pt.select_atoms('@H', traj.top)
n_indices = pt.select_atoms('@H', traj.top) - 1
nh_indices = list(zip(n_indices, h_indices))
mat_ired = pt.ired_vector_and_matrix(traj, mask=nh_indices, order=2)[-1]
mat_ired /= mat_ired[0, 0]
# matired: make sure to reproduce cpptraj output
aa_eq(mat_ired, state.data['matired'].values)
# get modes
modes = state.data[-2]
cpp_eigenvalues = modes.eigenvalues
cpp_eigenvectors = modes.eigenvectors
evals, evecs = np.linalg.eigh(mat_ired)
# need to sort a bit
evals = evals[::-1]
# cpptraj's eigvenvalues
aa_eq(evals, cpp_eigenvalues)
# cpptraj's eigvenvectors
# use absolute values to avoid flipped sign
# from Dan Roe
# In practice, the "sign" of an eigenvector depends on the math library used to calculate it.
# This is in fact why the modes command displacement test is disabled for cpptraj.
# I bet if you use a different math library (e.g. use your system BLAS/LAPACK instead of the one bundled with Amber
# or vice versa) you will get different signs.
# Bottom line is that eigenvector sign doesn't matter.
aa_eq(np.abs(evecs[:, ::-1].T), np.abs(cpp_eigenvectors), decimal=4)
data = _ired(state_vecs, modes=(cpp_eigenvectors, cpp_eigenvalues))
order_s2 = data['IRED_00127[S2]']
# load cpptraj's output and compare to pytraj' values for S2 order paramters
cpp_order_s2 = np.loadtxt(
os.path.join(cpptraj_test_dir, 'Test_IRED', 'orderparam.save')).T[-1]
aa_eq(order_s2, cpp_order_s2, decimal=5)
def merge_frame_from_trajs(trajlist):
"""
Examples
--------
>>> import numpy as np
>>> import pytraj as pt
>>> traj0 = pt.load_sample_data('tz2')[:3]
>>> traj1 = pt.load_sample_data('tz2')[3:6]
>>> traj2 = pt.load_sample_data('tz2')[6:9]
>>> print(traj0.n_atoms, traj1.n_atoms, traj2.n_atoms)
5293 5293 5293
>>> for frame in pt.tools.merge_frame_from_trajs((traj0, traj1, traj2)): print(frame)
<Frame with 15879 atoms>
<Frame with 15879 atoms>
<Frame with 15879 atoms>
"""
for iterables in zip(*trajlist):
yield merge_frames(iterables)
def test_0(self):
from pytraj.externals.six import zip
from pytraj import Trajectory, TrajectoryIterator
# TrajectoryIterator object
traj0 = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
assert isinstance(traj0, TrajectoryIterator)
# Trajectory object
traj1 = traj0[:]
assert isinstance(traj1, Trajectory)
# check same coords
ref = traj0[0]
for f0, f1 in zip(traj0(autoimage=True, rmsfit=(ref, '@CA,C,N')),
traj1(autoimage=True, rmsfit=(ref, '@CA,C,N'))):
aa_eq(f0.xyz, f1.xyz)
assert f0.rmsd_nofit(f1) == 0.
def test_simplifed_topology(self):
'''simplify'''
top = pt.load_topology(fn('Tc5b.top'))
sim_top = top.simplify()
assert sim_top.select('@CA').tolist() == top.select('@CA').tolist()
for atom, sim_atom in zip(top.atoms, sim_top.atoms):
assert atom.resname == sim_atom.resname
assert atom.name == sim_atom.name
assert atom.type == sim_atom.type
assert atom.charge == sim_atom.charge
assert atom.mass == sim_atom.mass
# API
atom = sim_top.atoms[0]
atom.residue
atom.residue.name
atom.residue.index
atom.bond_partners
def n_grams(a, n):
"""n_grams
Parameters
----------
a : sequence
n : number of elements
asarray : bool, default False
if False: return an iterator
if True: return a numpy array
Examples
--------
>>> list(n_grams([2, 3, 4 ,5], 2))
[(2, 3), (3, 4), (4, 5)]
Notes
-----
adapted from: http://sahandsaba.com/thirty-python-language-features-and-tricks-you-may-not-know.html
"""
z = (islice(a, i, None) for i in range(n))
return zip(*z)
def test_assert(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
fa = Trajectory.from_iterable(iterframe_master(traj), top=traj.top)
for f0, f1 in zip(fa, traj):
aa_eq(f0.xyz, f1.xyz)
def merge_trajs(traj1, traj2, start_new_mol=True, n_frames=None):
"""
Examples
--------
>>> import pytraj as pt
>>> import numpy as np
>>> traj1 = pt.load_sample_data('ala3')[:1]
>>> traj2 = pt.load_sample_data('tz2')[:1]
>>> traj3 = merge_trajs(traj1, traj2)
>>> # from frame_iter for saving memory
>>> traj3 = merge_trajs((traj1(0, 10, 2), traj1.top), (traj2(100, 110, 2), traj2.top), n_frames=6)
>>> # raise error if not having the same n_frames
>>> traj4 = pt.load_sample_data('tz2')[:]
>>> traj4.n_frames
10
>>> traj1.n_frames
1
>>> merge_trajs(traj1, traj4)
Traceback (most recent call last):
...
ValueError: must have the same n_frames
Notes
-----
Code might be changed
"""
from pytraj.externals.six import zip
from pytraj import Trajectory
import numpy as np
if isinstance(traj1, (list, tuple)):
n_frames_1 = n_frames
top1 = traj1[1]
_traj1 = traj1[0]
else:
n_frames_1 = traj1.n_frames
top1 = traj1.top
_traj1 = traj1
if isinstance(traj2, (list, tuple)):
n_frames_2 = n_frames
top2 = traj2[1] # example: (traj(0, 5), traj.top)
_traj2 = traj2[0]
else:
n_frames_2 = traj2.n_frames
top2 = traj2.top
_traj2 = traj2
if n_frames_1 != n_frames_2:
raise ValueError("must have the same n_frames")
traj = Trajectory()
traj._allocate(n_frames_1, top1.n_atoms + top2.n_atoms)
# merge Topology
top = top1.copy()
if start_new_mol:
top.start_new_mol()
top.join(top2)
traj.top = top
# update coords
for f1, f2, frame in zip(_traj1, _traj2, traj):
frame.xyz = np.vstack((f1.xyz, f2.xyz))
return traj
def test_4_trajiter(self):
traj = pt.load_sample_data("tz2")
for idx, (f0, f1) in enumerate(zip(traj, traj)):
f0.rmsd(f1)
def test_0(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
farray = traj[:]
for i, f0 in enumerate(traj):
for j, x in enumerate(f0.xyz):
if np.all(np.abs(x - 5.707) < 1E-3):
pass
for frame in traj.iterframe():
pass
for frame0 in farray.iterframe():
pass
i = 0
for frame0 in farray.iterframe(start=0, step=1):
i += 1
aa_eq(traj[-1].xyz, frame0.xyz)
start, stop, step = 2, 8, 4
indices = list(range(start, stop, step))
for idx, frame0, f in zip(indices, farray.iterframe(start, stop, step),
traj[indices]):
aa_eq(frame0.xyz, f.xyz)
aa_eq(traj[6].xyz, frame0.xyz)
for frame0 in farray.iterframe(start=2, step=2):
pass
aa_eq(traj[8].xyz, frame0.xyz)
traj[6][0]
for frame0 in traj.iterframe(start=2, step=4, stop=8):
pass
for frame0 in traj.iterframe(start=2, step=4, stop=8):
pass
aa_eq(traj[6].xyz, frame0.xyz)
for frame0 in traj.iterframe(start=2, step=2):
pass
aa_eq(traj[8].xyz, frame0.xyz)
count = 0
for frame0 in traj.iterframe(start=2):
count += 1
aa_eq(traj[-1].xyz, frame0.xyz)
count = 0
for frame0 in traj.iterframe(start=2, stop=7):
count += 1
aa_eq(traj[6].xyz, frame0.xyz)
for frame0 in traj.iterframe():
pass
aa_eq(traj[-1].xyz, frame0.xyz)
for frame0 in farray.iterframe():
pass
aa_eq(traj[-1].xyz, frame0.xyz)
for frame0 in traj():
pass
aa_eq(traj[-1].xyz, frame0.xyz)
for frame0 in farray():
pass
aa_eq(farray[-1].xyz, frame0.xyz)
count = 0
for frame0 in traj(start=2, stop=7):
count += 1
aa_eq(traj[6].xyz, frame0.xyz)
count = 0
for frame0 in farray(start=2, stop=7):
count += 1
aa_eq(traj[6].xyz, frame0.xyz)
count = 0
for frame0 in farray(2, 7, 1):
count += 1
aa_eq(traj[6].xyz, frame0.xyz)
count = 0
for frame0 in farray(2, 7, 2):
count += 1
aa_eq(traj[6].xyz, frame0.xyz)
