def test_matrix_module(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
for n_cores in [2, 3]:
for func in [matrix.dist, matrix.idea]:
x = pt.pmap(func, traj, '@CA', n_cores=n_cores)
aa_eq(x, func(traj, '@CA'))
def test_xyz(self):
traj = pt.datafiles.load_tz2_ortho()
t0 = traj[:]
def set_xyz_not_c_contiguous():
t0.xyz = np.asfortranarray(traj.xyz)
def append_2d():
pt.load_sample_data('ala3')
def set_xyz_not_same_n_atoms():
traj1 = pt.load_sample_data('ala3')
t0.xyz = traj1.xyz
def append_2d():
pt.load_sample_data('ala3')
t0.append_xyz(pt.tools.as_2darray(traj))
# fortran order, autoconvert
# make sure there is no TypeError
set_xyz_not_c_contiguous()
self.assertRaises(ValueError, lambda: set_xyz_not_same_n_atoms())
self.assertRaises(ValueError, lambda: append_2d())
# make sure to autoconvert from f4 to f8
xyz_f4 = np.array(traj.xyz, dtype='f4')
assert xyz_f4.itemsize == 4, 'must be f4'
t0.xyz = xyz_f4
aa_eq(t0.xyz, xyz_f4)
assert t0.xyz.itemsize == 8, 'must be converted from f4 to f8'
def test_0(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
f0 = traj[0]
pt.io.to_pickle(f0, './output/frame_pk.pk')
f1 = pt.io.read_pickle('./output/frame_pk.pk')
aa_eq(f0.xyz, f1.xyz)
def test_antechamber(self):
fn = 'data/cro.ac'
traj = pt.io.load_antechamber(fn)
assert traj.n_atoms == 40
assert traj.top.n_residues == 1
aa_eq([atom.charge for atom in traj.top.atoms][:3], [-0.895800, 0.113200, 0.11110])
aa_eq([atom.charge for atom in traj.top.atoms][-3:], [0.146500, 0.146500, 0.42300])
def test_0(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
f0 = traj[0]
pt.io.to_pickle(f0, './output/frame_pk.pk')
f1 = pt.io.read_pickle('./output/frame_pk.pk')
aa_eq(f0.xyz, f1.xyz)
def test_read_gaussian_output():
filename = "./data/gaussian/GF2.log"
gau = cclib.parser.Gaussian(filename)
go = gau.parse()
traj = pt.tools.read_gaussian_output(filename, "./data/gaussian/GF2.pdb")
aa_eq(traj.xyz, go.atomcoords)
def test_makestructure_with_reference(self):
tz2_parm7 = cpptraj_test_dir + '/tz2.parm7'
ref_rst7 = cpptraj_test_dir + '/tz2.rst7'
trajin_rst7 = cpptraj_test_dir + '/Test_MakeStructure/pp2.rst7.save'
# cpptraj
command = """
parm {}
reference {}
trajin {}
makestructure "ref:1-13:tz2.rst7"
rmsd reference
createcrd mycrd
""".format(tz2_parm7, ref_rst7, trajin_rst7)
state = pt.load_cpptraj_state(command)
state.run()
cpp_traj = state.data['mycrd']
# pytraj
traj = pt.load(trajin_rst7, top=tz2_parm7)
ref = pt.load(ref_rst7, top=tz2_parm7)
pt.make_structure(traj, "ref:1-13:1-13", ref=ref)
pt.rmsd(traj, ref=ref)
aa_eq(traj.xyz, cpp_traj.xyz)
def test_crdframes(self):
'''test crdframes in cpptraj
'''
max_frames = 50
traj = pt.iterload('data/tz2.nc',
'data/tz2.parm7',
frame_slice=(0, max_frames, 2))
state = pt.load_cpptraj_state('''
parm {0}
trajin {1} 1 {2} 2
trajin {1} 1 {2} 2
loadtraj name traj
crdaction traj rms crdframes 1,10
crdaction traj rms crdframes 1,30,2
'''.format(traj.top.filename, traj.filename, max_frames))
state.run()
rmsd_0 = pt.rmsd(traj, ref=0, frame_indices=range(10))
rmsd_crdframes = state.data[2].values
aa_eq(rmsd_0, rmsd_crdframes)
traj2 = traj.copy()
assert traj2.n_frames == traj.n_frames, 'must have the same n_frames'
traj2._load(traj.filename, frame_slice=(0, max_frames, 2))
assert traj2.n_frames == 2 * traj.n_frames, 'n_frames must be doubled after reload'
rmsd_1 = pt.rmsd(traj2, ref=0, frame_indices=range(0, 30, 2))
rmsd_crdframes = state.data[3].values
aa_eq(rmsd_1, rmsd_crdframes)
def test_iterframe(self):
'''test iterframe for both Trajectory and TrajectoryIterator
'''
orig_traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
# iterframe (already in doctest), just throwing raise to increase coverage score
for traj in [orig_traj, orig_traj[:]]:
self.assertRaises(ValueError,
lambda: traj.iterframe(rmsfit='crazy'))
# rmsfit is an int
t0 = orig_traj[:].rmsfit(ref=3)
aa_eq(
pt.rmsd_nofit(
traj(rmsfit=3),
ref=orig_traj[-1]),
pt.rmsd_nofit(t0, ref=orig_traj[-1]))
# test TypeError if not has n_frames info
t0 = orig_traj[:]
def int_gen(k):
for i in range(k):
yield i
fi = pt.iterframe(t0, frame_indices=int_gen(3))
aa_eq(pt.radgyr(fi, top=traj.top), pt.radgyr(orig_traj[:3]))
def test_0(self):
_ds1 = pt.calc_dssp(traj[:5], dtype='dataset')
ds1 = _ds1.grep('LYS')
_ds2 = pt.calc_dssp(traj[5:], dtype='dataset')
ds2 = _ds2.grep('LYS')
dstack = stack((ds1, ds2))
_d12 = pt.calc_dssp(traj, dtype='dataset')
d12 = _d12.grep("LYS")
dstack_dict = dstack.to_dict()
d12_dict = d12.to_dict()
assert sorted(dstack_dict.keys()) == sorted(d12_dict)
for key in dstack_dict.keys():
arr0 = dstack_dict[key]
arr1 = d12_dict[key]
if np.any(arr0 == arr1) == False:
pass
arr1 = ds1.to_ndarray()
ds2.to_ndarray()
arrstack = dstack.to_ndarray()
arr12 = d12.to_ndarray()
aa_eq(arrstack.flatten(), arr12.flatten())
def test_mass(self):
traj = self.traj
# cpptraj output
cm = '''
reference avg.pdb
rms R0 reference @CA,C,N,O savematrices mass
'''
state = pt.load_batch(traj, cm)
state.run()
saved_mat = state.data[-1].values
# pytraj output
avg = pt.mean_structure(traj)
mat = pt.calc_rotation_matrix(
traj, ref=avg, mask='@CA,C,N,O', mass=True)
assert mat.shape == (traj.n_frames, 3, 3), 'mat shape'
aa_eq(mat.flatten(), saved_mat.flatten())
# with rmsd
avg = pt.mean_structure(traj)
mat2, rmsd_ = pt.calc_rotation_matrix(
traj, ref=avg, mask='@CA,C,N,O', mass=True, with_rmsd=True)
aa_eq(mat2.flatten(), saved_mat.flatten())
assert pt.tools.rmsd(rmsd_, state.data['R0']) < 1E-3
def test_volmap(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
# raise if does not provide size
self.assertRaises(
AssertionError, lambda: pt.pmap(pt.volmap,
traj,
mask=':WAT@O',
grid_spacing=(0.5, 0.5, 0.5),
n_cores=2))
mask = ':WAT@O'
grid_spacing = (0.5, 0.5, 0.5)
for n_cores in [1, 2, 3]:
for size in [(20, 20, 20), (20, 40, 60)]:
serial_out = pt.volmap(traj,
mask=mask,
grid_spacing=grid_spacing,
size=size)
parallel_out = pt.pmap(pt.volmap,
traj,
mask=mask,
grid_spacing=grid_spacing,
size=size,
n_cores=n_cores)
self.assertEqual(serial_out.shape, tuple(2 * x for x in size))
aa_eq(serial_out, parallel_out)
def test_nomass(self):
traj = self.traj
# cpptraj output
cm = '''
reference avg.pdb
rms R0 reference @CA,C,N,O savematrices
'''
state = pt.load_batch(traj, cm)
state.run()
saved_mat = state.data[-1].values
# pytraj output
avg = pt.mean_structure(traj)
mat = pt.calc_rotation_matrix(traj, ref=avg, mask='@CA,C,N,O')
assert mat.shape == (traj.n_frames, 3, 3), 'mat shape'
aa_eq(mat.flatten(), saved_mat.flatten())
# not specify reference
cm = '''
rms R0 @CA,C,N,O savematrices
'''
state = pt.load_batch(traj, cm)
state.run()
saved_mat = state.data[-1].values
mat = pt.calc_rotation_matrix(traj, mask='@CA,C,N,O')
aa_eq(mat.flatten(), saved_mat.flatten())
def test_projection_with_None_average_coord(self):
cm = '''
matrix name correlmat {scalar_type} @CA
diagmatrix correlmat name evecs vecs 5
run
projection P1 modes evecs beg 1 end 2 @CA
'''
scalar_type = 'covar'
cm2 = cm.format(scalar_type=scalar_type)
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
state = pt.load_cpptraj_state(cm2, traj)
state.run()
mat = pt.matrix.covar(traj, '@CA')
eigenvalues, eigenvectors = pt.matrix.diagonalize(mat,
2,
dtype='tuple')
data = pt.all_actions.projection(traj,
'@CA',
eigenvalues,
eigenvectors,
scalar_type=scalar_type)
aa_eq(np.abs(state.data['Mode1'].values), np.abs(data[0]))
aa_eq(np.abs(state.data['Mode2'].values), np.abs(data[1]))
assert data.shape == (2, traj.n_frames)
def test_pca_with_ref_with_different_mask_from_matrix(self):
'''test_pca_with_ref_with_different_mask_from_matrix
has reference. Use !@H= for ref_mask and use * for covariance matrix and projection
from drroe: "You should be able to supply separate masks for fitting and creating the covariance matrix
It is common enough for example to only perform rms-fitting on heavy atoms while still wanting all atoms in eigenvectors."
pytraj: pt.pca(traj, mask=mask_matrix, n_vecs=2, ref=ref, ref_mask=mask_ref)
'''
command_ref_provided = '''
parm {}
trajin {} 1 3
reference {}
# only perform fitting on heavy atoms
rms reference !@H=
# all atoms
matrix covar name MyMatrix *
createcrd CRD1
run
# Step three. Diagonalize matrix.
runanalysis diagmatrix MyMatrix vecs 2 name MyEvecs
# Step four. Project saved fit coordinates along eigenvectors 1 and 2
# all atoms
crdaction CRD1 projection evecs MyEvecs * out project.dat beg 1 end 2
'''.format(tz2_top, tz2_trajin, fn('tz2.rst7'))
traj = pt.load(fn('tz2.nc'),
fn('tz2.parm7'),
frame_indices=range(0, 3))
traj_on_disk = pt.iterload(fn('tz2.nc'),
fn('tz2.parm7'),
frame_slice=[(0, 3)])
ref = pt.load(fn('tz2.rst7'), traj.top)
state = pt.load_cpptraj_state(command_ref_provided)
state.run()
mask_ref = '!@H='
mask_matrix = '*'
data = pt.pca(traj,
mask=mask_matrix,
n_vecs=2,
ref=ref,
ref_mask=mask_ref)
data2 = pt.pca(traj_on_disk,
mask=mask_matrix,
n_vecs=2,
ref=ref,
ref_mask=mask_ref)
cpp_data = state.data[-2:].values
# use absolute values
aa_eq(np.abs(data[0]), np.abs(cpp_data), decimal=3)
aa_eq(np.abs(data2[0]), np.abs(cpp_data), decimal=3)
def test_closest_compared_to_cpptraj(self):
trajin = fn('tz2.ortho.nc')
parm = fn('tz2.ortho.parm7')
traj = pt.iterload(trajin, parm, frame_slice=[
(0, 2),
])
state = pt.load_cpptraj_state('''
parm {}
trajin {} 1 2
autoimage
closest 100 :1-13
createcrd mycrd'''.format(parm, trajin))
state.run()
fi, top = pt.closest(traj(autoimage=True),
mask=':1-13',
n_solvents=100)
xyz = pt.get_coordinates(fi)
t0 = pt.Trajectory(xyz=xyz, top=top)
aa_eq(state.data['mycrd'].xyz, t0.xyz)
# dtype = 'trajectory'
t1 = pt.closest(traj(autoimage=True),
mask=':1-13',
n_solvents=100,
dtype='trajectory')
aa_eq(state.data['mycrd'].xyz, t1.xyz)
def test_get_velocity(self):
traj = pt.iterload(fn('issue807/trunc.nc'),
fn("issue807/system.prmtop"))
traj[0]
# no mask, no frame_indices
vels = pt.get_velocity(traj)
assert vels.shape == (traj.n_frames, traj.n_atoms, 3), 'vels.shape'
# string mask
vels = pt.get_velocity(traj, '@O', frame_indices=[0, 2])
fi = traj(frame_indices=[0, 2], mask='@O')
assert vels.shape == (fi.n_frames, fi.top.n_atoms, 3), 'vels.shape'
# atom indices
atm_indices = pt.select_atoms('@O', traj.top)
vels_ = pt.get_velocity(traj, atm_indices, frame_indices=[0, 2])
fi = traj(frame_indices=[0, 2], mask='@O')
assert vels_.shape == (fi.n_frames, fi.top.n_atoms, 3), 'vels.shape'
aa_eq(vels, vels_)
# raise if not having velocity
traj2 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
with pytest.raises(ValueError):
pt.get_velocity(traj2)
def test_drmsd(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
txt = '''
parm data/tz2.parm7
trajin data/tz2.nc
drmsd drms_nofit out drmsd.dat
rms rms_nofit out drmsd.dat nofit
rms rms_fit out drmsd.dat
drmsd drms_fit out drmsd.dat
'''
state = pt.load_cpptraj_state(txt)
state.run()
cpp_data = state.data[1:]
# distance_rmsd
data_drmsd = pt.distance_rmsd(traj)
aa_eq(data_drmsd, cpp_data[0])
aa_eq(pt.drmsd(traj), cpp_data[0])
# rms_nofit
aa_eq(cpp_data[1], pt.rmsd(traj, nofit=True))
# rms_fit
aa_eq(cpp_data[2], pt.rmsd(traj, nofit=False))
# drmsd with rmsfit
aa_eq(cpp_data[3], pt.distance_rmsd(traj(rmsfit=0), ref=traj[0]))
def test_crdframes(self):
'''test crdframes in cpptraj
'''
max_frames = 50
traj = pt.iterload(fn('tz2.nc'),
fn('tz2.parm7'),
frame_slice=(0, max_frames, 2))
state = pt.load_cpptraj_state('''
parm {0}
trajin {1} 1 {2} 2
trajin {1} 1 {2} 2
loadtraj name traj
crdaction traj rms crdframes 1,10
crdaction traj rms crdframes 1,30,2
'''.format(traj.top.filename, traj.filename, max_frames))
state.run()
rmsd_0 = pt.rmsd(traj, ref=0, frame_indices=range(10))
rmsd_crdframes = state.data[2].values
aa_eq(rmsd_0, rmsd_crdframes)
traj2 = traj.copy()
assert traj2.n_frames == traj.n_frames, 'must have the same n_frames'
traj2._load(traj.filename, frame_slice=(0, max_frames, 2))
assert traj2.n_frames == 2 * traj.n_frames, 'n_frames must be doubled after reload'
rmsd_1 = pt.rmsd(traj2, ref=0, frame_indices=range(0, 30, 2))
rmsd_crdframes = state.data[3].values
aa_eq(rmsd_1, rmsd_crdframes)
def test_0(self):
import numpy as np
from pytraj.datasetlist import stack as stack
_ds1 = pt.calc_dssp(traj[:5], dtype='dataset')
ds1 = _ds1.grep('LYS')
_ds2 = pt.calc_dssp(traj[5:], dtype='dataset')
ds2 = _ds2.grep('LYS')
dstack = stack((ds1, ds2))
_d12 = pt.calc_dssp(traj, dtype='dataset')
d12 = _d12.grep("LYS")
dstack_dict = dstack.to_dict()
d12_dict = d12.to_dict()
assert sorted(dstack_dict.keys()) == sorted(d12_dict)
for key in dstack_dict.keys():
arr0 = dstack_dict[key]
arr1 = d12_dict[key]
if np.any(arr0 == arr1) == False:
pass
arr1 = ds1.to_ndarray()
arr2 = ds2.to_ndarray()
arrstack = dstack.to_ndarray()
arr12 = d12.to_ndarray()
aa_eq(arrstack.flatten(), arr12.flatten())
def test_multiple_cores(self):
from multiprocessing import Pool
traj = pt.iterload('data/tz2.nc', 'data/tz2.parm7')
for _ in range(10):
traj._load(traj.filelist)
saved_angle = pt.angle(traj, ':3 :10 :11')
saved_dist = pt.distance(traj, ':3 :10')
for n_cores in [2, 3]:
lines = ['angle :3 :10 :11', 'distance :3 :10']
pfuncs = partial(worker_state,
n_cores=n_cores,
traj=traj,
dtype='dict',
lines=lines)
p = Pool(n_cores)
data_list = p.map(pfuncs, [rank for rank in range(n_cores)])
p.close()
p.join()
data_list_sorted_rank = (data[1]
for data in sorted(data_list,
key=lambda x: x[0]))
final_data = concat_dict(data_list_sorted_rank)
aa_eq(final_data['Ang_00002'], saved_angle)
aa_eq(final_data['Dis_00003'], saved_dist)
def test_progress_log():
"""test_progress_log: simple test, just to make sure it is runnable
"""
traj = pt.datafiles.load_tz2()
p = ProgressBarTrajectory(traj, style='basic', every=20)
pt.molsurf(p)
p = ProgressBarTrajectory(traj, style='bar', every=20)
pt.molsurf(p)
p = ProgressBarTrajectory(traj, style='circle', every=20)
pt.molsurf(p)
# need to run on Jupyter notebook
if tqdm is not None:
p = ProgressBarTrajectory(traj, style='tqdm')
pt.molsurf(p)
try:
p = ProgressBarTrajectory(traj, style=tqdm_notebook)
pt.molsurf(p)
except traitlets.traitlets.TraitError:
pass
# make sure not loading all coordinates from TrajectoryIterator
traj2 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
traj2._size_limit_in_GB = traj2._estimated_GB - 0.001
with pytest.raises(MemoryError):
traj2.xyz
p2 = ProgressBarTrajectory(traj2)
aa_eq(pt.rmsd(p2), pt.rmsd(traj2))
def test_center(self):
traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
pt.load(fn("tz2.center_mass.nc"), traj.top)
fa = traj[:]
fa2 = traj[:].copy()
pt.center(fa, mask=':1', mass=True)
aa_eq(fa.xyz, fa.xyz, decimal=5)
# raise if center not in 'origin', 'box'
self.assertRaises(ValueError, lambda: pt.center(fa, center='oh'))
# center to point
pt.center(fa, ':1', center=[0, 1, 2], mass=True)
aa_eq(pt.center_of_mass(fa, ':1')[0], [0, 1, 2])
pt.center(fa, ':1', center=[0, 1, 2], mass=False)
aa_eq(pt.center_of_geometry(fa, ':1')[0], [0, 1, 2])
fa2.center(':1', center=[0, 1, 2], mass=False)
aa_eq(pt.center_of_geometry(fa2, ':1')[0], [0, 1, 2])
# on_disk
traj_on_disk = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
pt.center(traj_on_disk, ':1', center=[0, 1, 2], mass=True)
aa_eq(pt.center_of_mass(traj_on_disk, ':1')[0], [0, 1, 2])
def test_velocity(self):
traj = pt.iterload("./data/issue807/trunc.nc",
"data/issue807/system.prmtop")
f = traj[0]
# no mask, no frame_indices
vels = pt.get_velocity(traj)
assert vels.shape == (traj.n_frames, traj.n_atoms, 3), 'vels.shape'
# string mask
vels = pt.get_velocity(traj, '@O', frame_indices=[0, 2])
fi = traj(frame_indices=[0, 2], mask='@O')
assert vels.shape == (fi.n_frames, fi.top.n_atoms, 3), 'vels.shape'
# atom indices
atm_indices = pt.select_atoms('@O', traj.top)
vels_ = pt.get_velocity(traj, atm_indices, frame_indices=[0, 2])
fi = traj(frame_indices=[0, 2], mask='@O')
assert vels_.shape == (fi.n_frames, fi.top.n_atoms, 3), 'vels.shape'
aa_eq(vels, vels_)
# raise if not having velocity
traj2 = pt.iterload('data/tz2.nc', 'data/tz2.parm7')
self.assertRaises(ValueError, lambda: pt.get_velocity(traj2))
def test_comprehensive(self):
traj = pt.iterload("data/Test_RemdTraj/rem.nc.000",
"data/Test_RemdTraj/ala2.99sb.mbondi2.parm7")
# temperature
aa_eq(traj.temperatures, [300., 630.5, 630.5, 630.5, 630.5, 630.5,
630.5, 630.5, 492.2, 384.3])
# iterframe (already in doctest), just throwing raise to increase coverage score
self.assertRaises(ValueError, lambda: traj.iterframe(rmsfit='crazy'))
# raise
# memory error if load larger than 1GB for xyz
traj = pt.datafiles.load_tz2_ortho()
for _ in range(11):
traj._load(traj.filelist)
self.assertRaises(MemoryError, lambda: traj.xyz)
# can not find filename
self.assertRaises(ValueError, lambda: traj._load(None))
# has filename but does not have Topology
self.assertRaises(
ValueError,
lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=None))
self.assertRaises(
ValueError,
lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000"))
# empty Topology
self.assertRaises(
ValueError,
lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=pt.Topology()))
# weird Topology
self.assertRaises(
ValueError,
lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=pt.Frame))
def test_rmsfit_with_autoimage_compared_to_cpptraj(self):
# assert to cpptraj: need to set mass
traj = self.traj.copy()
txt = '''
parm {0}
trajin {1}
autoimage
rms first {2} mass
trajout tmp.nc
'''.format(traj.top.filename, traj.filename, self.mask)
with tempfolder():
state = pt.datafiles.load_cpptraj_output(txt, dtype='state')
state.run()
# need to load to memory (not iterload)
saved_traj = pt.load('tmp.nc', traj.top)
fa1 = traj[:]
fa1.autoimage()
pt.superpose(fa1, ref=0, mask=self.mask, mass=True)
aa_eq(fa1.xyz, saved_traj.xyz)
fa_saved_nowat = saved_traj['!:WAT']
fa1_nowat = fa1['!:WAT']
aa_eq(fa_saved_nowat.xyz, fa1_nowat.xyz)
def test_general(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
# with mask
saved_data = pt.radgyr(traj, '@CA')
data = pt.pmap(pt.radgyr, traj, '@CA')
data = pt.tools.dict_to_ndarray(data)
aa_eq(saved_data, data)
# with a series of functions
func_list = [pt.radgyr, pt.molsurf, pt.rmsd]
ref = traj[-3]
for n_cores in [2, 3]:
for func in func_list:
if func in [pt.rmsd, ]:
pout = pt.tools.dict_to_ndarray(pt.pmap(func=func,
traj=traj,
ref=ref,
n_cores=n_cores))
serial_out = flatten(func(traj, ref=ref))
else:
pout = pt.tools.dict_to_ndarray(pt.pmap(n_cores=n_cores,
func=func,
traj=traj))
serial_out = flatten(func(traj))
aa_eq(pout[0], serial_out)
# test worker
# need to test this since coverages seems not recognize partial func
from pytraj.parallel.multiprocessing_ import worker_byfunc
data = worker_byfunc(rank=2, n_cores=8, func=pt.radgyr, traj=traj, args=(), kwd={'mask': '@CA'}, iter_options={})
assert data[0] == 2, 'rank must be 2'
assert data[2] == 1, 'n_frames for rank=2 should be 1 (only 10 frames in total)'
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_rmsfit_with_autoimage_compared_to_cpptraj(self):
# assert to cpptraj: need to set mass
traj = self.traj.copy()
txt = '''
parm {0}
trajin {1}
autoimage
rms first {2} mass
trajout tmp.nc
'''.format(traj.top.filename, traj.filename, self.mask)
with tempfolder():
state = pt.datafiles.load_cpptraj_output(txt, dtype='state')
state.run()
# need to load to memory (not iterload)
saved_traj = pt.load('tmp.nc', traj.top)
fa1 = traj[:]
fa1.autoimage()
pt.superpose(fa1, ref=0, mask=self.mask, mass=True)
aa_eq(fa1.xyz, saved_traj.xyz)
fa_saved_nowat = saved_traj['!:WAT']
fa1_nowat = fa1['!:WAT']
aa_eq(fa_saved_nowat.xyz, fa1_nowat.xyz)
def test_matrix_module(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
for n_cores in [2, 3]:
for func in [matrix.dist, matrix.idea]:
x = pt.pmap(func, traj, '@CA', n_cores=n_cores)
aa_eq(x, func(traj, '@CA'))
def test_combine_with_frame_iterator(self):
traj = pt.iterload(fn("tz2.ortho.nc"), fn("tz2.ortho.parm7"))
dslist = CpptrajDatasetList()
commands = [
'autoimage',
'rms',
]
actlist = ActionList(commands, top=traj.top, dslist=dslist)
def get_frameiter(actlist, traj):
for frame in traj:
actlist.compute(frame)
yield frame
def do_extra(fi):
a = []
for frame in fi:
frame.xyz = frame.xyz + 2.
a.append(frame.copy())
return a
new_list = do_extra(get_frameiter(actlist, traj))
t0 = traj[:].autoimage().superpose()
t0.xyz += 2.
aa_eq(np.array([frame.xyz for frame in new_list]), t0.xyz)
def test_different_cores(self):
# use different frame_slice with different n_cores to test
REF_0 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))[0]
for frame_slice in [
(0, 100),
(10, 100, 3),
(50, 80, 2),
(51, 80, 3),
]:
traj = pt.iterload(fn('tz2.nc'),
fn('tz2.parm7'),
frame_slice=frame_slice)
saved_angle = pt.angle(traj, ':3 :7 :9')
saved_dist = pt.distance(traj, ':2 :10')
saved_rmsd = pt.rmsd(traj, ref=REF_0, mask='@CA')
lines = ['angle :3 :7 :9', 'distance :2 :10']
for n_cores in [2, 3]:
data_list = [
worker_by_state(rank, n_cores, traj, lines)
for rank in range(n_cores)
]
final_data = concat_dict([x[1] for x in data_list])
aa_eq(final_data['Ang_00002'], saved_angle)
aa_eq(final_data['Dis_00003'], saved_dist)
def test_tools(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
aa_eq(pt.tools.as_2darray(traj), pt.tools.as_2darray(traj.xyz))
# as_2darray
assert pt.tools.as_2darray(traj).ndim == 2, 'ndim must be 2'
#
self.assertRaises(
ValueError,
lambda: pt.tools.rmsd_1darray([3, 2], [[2, 3], [4, 6]]))
self.assertRaises(
ValueError,
lambda: pt.tools.rmsd_1darray([[2, 3], [4, 6]], [2, 3]))
# rmsd
self.assertRaises(ValueError,
lambda: pt.tools.rmsd([[2, 3], [4, 6]], [2, 3]))
self.assertRaises(
ValueError, lambda: pt.tools.rmsd([[2, 3]], [2, 3], flatten=False))
#
for frame in pt.tools.split_traj_by_residues(traj, 0, 12):
pass
assert not [
name for name in pt.tools.dir_(frame) if name.startswith('_')
]
def test_iterframe(self):
'''test iterframe for both Trajectory and TrajectoryIterator
'''
orig_traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
# iterframe (already in doctest), just throwing raise to increase coverage score
for traj in [orig_traj, orig_traj[:]]:
self.assertRaises(ValueError,
lambda: traj.iterframe(rmsfit='crazy'))
# rmsfit is an int
t0 = orig_traj[:].rmsfit(ref=3)
aa_eq(pt.rmsd_nofit(traj(rmsfit=3), ref=orig_traj[-1]),
pt.rmsd_nofit(t0, ref=orig_traj[-1]))
# test TypeError if not has n_frames info
t0 = orig_traj[:]
def int_gen(k):
for i in range(k):
yield i
fi = pt.iterframe(t0, frame_indices=int_gen(3))
aa_eq(pt.radgyr(fi, top=traj.top), pt.radgyr(orig_traj[:3]))
def test_xyz(self):
traj = pt.datafiles.load_tz2_ortho()
t0 = traj[:]
def set_xyz_not_c_contiguous():
t0.xyz = np.asfortranarray(traj.xyz)
def append_2d():
traj1 = pt.load_sample_data('ala3')
def set_xyz_not_same_n_atoms():
traj1 = pt.load_sample_data('ala3')
t0.xyz = traj1.xyz
def append_2d():
traj1 = pt.load_sample_data('ala3')
t0.append_xyz(pt.tools.as_2darray(traj))
# fortran order, autoconvert
# make sure there is no TypeError
set_xyz_not_c_contiguous()
self.assertRaises(ValueError, lambda: set_xyz_not_same_n_atoms())
self.assertRaises(ValueError, lambda: append_2d())
# make sure to autoconvert from f4 to f8
xyz_f4 = np.array(traj.xyz, dtype='f4')
assert xyz_f4.itemsize == 4, 'must be f4'
t0.xyz = xyz_f4
aa_eq(t0.xyz, xyz_f4)
assert t0.xyz.itemsize == 8, 'must be converted from f4 to f8'
def test_1(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
trajectory_t0 = Trajectory(traj)
# __iter__
for f in trajectory_t0:
pass
f.xyz[0, 0] = 10.
assert f.xyz[0, 0] == 10.
assert trajectory_t0.xyz[-1, 0, 0] == 10.
# __getitem__
# make a new copy
trajectory_t0 = Trajectory(traj)
f0 = trajectory_t0[0]
f0.xyz[0, 0] = 200.
assert trajectory_t0.xyz[0, 0, 0] == 200.
# translate
# make a new copy
trajectory_t0 = Trajectory(traj)
t0 = traj[:]
pt.translate(trajectory_t0, 'x 1.2')
pt.translate(t0, 'x 1.2')
aa_eq(trajectory_t0.xyz, t0.xyz)
try:
pt.rmsd(t0, ref=0)
pt.rmsd(trajectory_t0, ref=0)
except ImportError:
pass
# test rmsfit
trajectory_t0.rmsfit(ref=0)
def test_vs_list_tuple(self):
traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
traj0 = pt.replicate_cell(traj, direction='dir 001 dir 0-10')
traj1 = pt.replicate_cell(traj, direction=('001', '0-10'))
aa_eq(traj0.xyz, traj1.xyz)
self.assertRaises(ValueError, lambda: pt.replicate_cell(traj, direction=traj[0]))
def test_pmap_average_structure(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
saved_frame = pt.mean_structure(traj, '@CA')
saved_xyz = saved_frame.xyz
for n_cores in [2, 3, 4]:
frame = pt.pmap(pt.mean_structure, traj, '@CA', n_cores=n_cores)
aa_eq(frame.xyz, saved_xyz)
def test_pieline(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
fi = pt.pipe(traj, ['autoimage', ])
aa_eq(pt.get_coordinates(fi), traj[:].autoimage().xyz)
fi = pt.pipe(traj, ['autoimage', ], frame_indices=[3, 5])
aa_eq(pt.get_coordinates(fi), traj[[3, 5]].autoimage().xyz)
def test_DatasetMatrix3x3(self):
# test _append_from_array
mat0 = pt.calc_rotation_matrix(self.traj, ref=0)
shape2d = (mat0.shape[0], mat0.shape[1] * mat0.shape[2])
dmat3x3 = c_datasets.DatasetMatrix3x3()
dmat3x3._append_from_array(mat0.reshape(shape2d))
aa_eq(mat0, dmat3x3.values)
def test_0(self):
from pytraj.core import mass_atomic_number_dict, mass_element_dict
top = mdio.load_topology("./data/tz2.parm7")
mass_list = []
for atom in top:
mass_list.append(mass_atomic_number_dict[atom.atomic_number])
aa_eq(mass_list, top.mass, 2)
def test_0(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
from pytraj.utils import convert as cv
arange = cv.array_to_cpptraj_range(range(7))
a0 = pt.multidihedral(traj, resrange='1-7').values
a1 = pt.multidihedral(traj, resrange=range(7)).values
aa_eq(a0.flatten(), a1.flatten())
def test_read_gaussian_output(self):
import cclib
filename = "./data/gaussian/GF2.log"
gau = cclib.parser.Gaussian(filename)
go = gau.parse()
traj = pt.tools.read_gaussian_output(filename, "./data/gaussian/GF2.pdb")
aa_eq(traj.xyz, go.atomcoords)
def test_issue_172_cpptraj(self):
'''Internal Error: Attempting to assign to Frame with external memory
'''
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
arr_out_of_memory = pt.atomiccorr(traj(0, 8, 2), '@CA')
arr_in_memory = pt.atomiccorr(traj[0: 8: 2], '@CA')
aa_eq(arr_out_of_memory, arr_in_memory)
def test_0(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
from pytraj.utils import convert as cv
cv.array_to_cpptraj_range(range(7))
a0 = pt.multidihedral(traj, resrange='1-7').values
a1 = pt.multidihedral(traj, resrange=range(7)).values
aa_eq(a0.flatten(), a1.flatten())
def test_trajiter_with_actionlist(self):
traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
traj.autoimage().center('origin').superpose('@CA')
pk_fn = 'output/test.pk'
pt.to_pickle(traj, pk_fn)
traj2 = pt.read_pickle(pk_fn)
print(traj2._transform_commands)
aa_eq(traj.xyz, traj2.xyz)
def test_volume(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
state = pt.load_cpptraj_state('''
volume''', traj)
state.run()
vol = pt.volume(traj)
aa_eq(state.data[-1], vol)
def test_0(self):
traj = pt.iterload("./data/Test_NAstruct/adh026.3.pdb")
d = pt.calc_delta(traj, resrange='1').values
d1 = pt.dihedral(traj, ":1@C5' :1@C4' :1@C3' :1@O3'")
d2 = pt._dihedral_res(traj, ("C5'", "C4'", "C3'", "O3'"))
aa_eq(d, d1)
aa_eq(d, d2)
def test_fancy_save(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
traj[1:8].save("./output/test_fancy_save_frame1_7.x", overwrite=True)
fanew = pt.iterload("./output/test_fancy_save_frame1_7.x", traj.top)
for idx, f0 in enumerate(traj[1:8]):
f0new = fanew[idx]
aa_eq(f0.xyz, f0new.xyz)
def test_atomiccorr(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
state = pt.load_batch(traj, '''
atomiccorr out test.dat :1-13 byres
''')
state.run()
data = pt.atomiccorr(traj, ':1-13', byres=True)
aa_eq(data, state.data[1].values)
def test_6(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
mask_list = ('@CB @CA', '@CA @H')
dslist = pt.calc_distance(traj, mask_list)
dslist3_0 = pt.calc_distance(traj, mask_list[0])
dslist3_1 = pt.calc_distance(traj, mask_list[1])
aa_eq(dslist3_0, dslist[0])
aa_eq(dslist3_1, dslist[1])
def test_sander_pmap_simple(self):
traj = pt.iterload('./data/Tc5b.x', './data/Tc5b.top')
fname = traj.top.filename
serial = pt.energy_decomposition(traj, prmtop=fname)['dihedral']
parallel = pt.pmap(n_cores=4,
func=pt.energy_decomposition,
traj=traj,
prmtop=fname)['dihedral']
aa_eq(serial, parallel)
def test_allocate_frames(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
traj2 = pt.Trajectory()
traj2._allocate(traj.n_frames, traj.n_atoms)
assert (traj2.shape == traj.shape)
traj2.top = traj.top.copy()
traj2.xyz = traj.xyz[:]
aa_eq(traj2.xyz, traj.xyz)
def test_center(self):
traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
saved_traj = pt.load("./data/tz2.center_mass.nc", traj.top)
fa = traj[:]
pt.center(fa, mask=':1', mass=True)
aa_eq(fa.xyz, fa.xyz, decimal=5)
# raise if center not in 'origin', 'box'
self.assertRaises(ValueError, lambda: pt.center(fa, center='oh'))
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_dataset_coords_ref(self):
traj = pt.iterload('data/tz2.nc', 'data/tz2.parm7')
avg_frame = pt.mean_structure(traj(rmsfit=(0, '!@H=')))
state = self.state
# need to loop several times to make sure this does not fail
# due to memory free
for _ in range(20):
cpp_ref = state.data['AVG'].get_frame()
aa_eq(avg_frame.xyz, cpp_ref.xyz)
def test_DatasetDouble(self):
dslist = CpptrajDatasetList()
d = dslist.add_new(dtype='double')
a = range(8)
# append
for i in a:
d.append(i)
aa_eq(a, d)
assert int(d[2]) == a[2] == 2, 'must be equal'
def test_insert_new_function(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
# create mutable Trajectory
t0 = traj[:]
for frame in t0:
frame.xyz[:10] += 1.
data_parallel = pt.pmap(pt.radgyr, traj, n_cores=2, apply=change_10_atoms)
data_serial = pt.radgyr(t0)
aa_eq(data_parallel['RoG_00000'], data_serial)
