def test_formats():
'''make sure to save correct format
'''
def assert_has_exptected_line_textfile(expected_line,
line_index,
fn='test'):
with open(fn, 'r') as fh:
lines = [fh.readline() for _ in range(3)]
# print(lines)
assert expected_line in lines[line_index]
def assert_has_exptected_line_binaryfile(expected_line, fn='test'):
with open(fn, 'rb') as fh:
buffer_ = fh.read()
# print(buffer_)
assert expected_line.encode() in buffer_
traj = pt.datafiles.load_tz2()
fn = 'test'
with tempfolder():
# pdb
traj.save(fn, format='pdb', overwrite=True)
expected_line = 'ATOM 1 N SER 1 -1.889 9.159 7.569'
assert_has_exptected_line_textfile(expected_line, 0, fn)
# mol2
traj.save(fn, format='mol2', overwrite=True)
expected_line = 'Cpptraj generated mol2 file'
assert_has_exptected_line_textfile(expected_line, 1, fn)
# mdcrd
traj.save(fn, format='mdcrd', overwrite=True)
expected_line = 'Cpptraj Generated trajectory'
assert_has_exptected_line_textfile(expected_line, 0, fn)
# netcdf
traj.save(fn, format='netcdf', overwrite=True)
expected_line = 'CDF'
assert_has_exptected_line_binaryfile(expected_line, fn)
# dcd
traj.save(fn, format='dcd', overwrite=True)
expected_line = 'Cpptraj generated dcd file'
assert_has_exptected_line_binaryfile(expected_line, fn)
# cif, not supported
# traj.save(fn, format='cif', overwrite=True)
# expected_line = 'Cpptraj generated cif file'
# assert_has_exptected_line_textfile(expected_line, 1, fn)
# trr
traj.save(fn, format='trr', overwrite=True)
expected_line = 'GMX_trn_file'
assert_has_exptected_line_binaryfile(expected_line, fn)
def test_trajectory_writer__with_statement():
frame0 = farray[0]
with tempfolder():
with TrajectoryWriter(
filename="test_trajout_withstatement.x",
top=farray.top,
overwrite=True) as trajout:
trajout.write(frame0)
# reload
farray2 = Trajectory("test_trajout_withstatement.x", fn('Tc5b.top'))
farray2[0]
def test_write_data():
traj = pt.iterload(fn("tz2.ortho.nc"), fn("tz2.ortho.parm7"))
rmsd_data = pt.rmsd(traj)
with tempfolder():
pt.io.write_data("test.agr", rmsd_data)
pt.io.write_data("test.gnu", rmsd_data)
assert os.path.exists("test.agr")
assert os.path.exists("test.gnu")
saved_dx_file = cpptraj_test_dir + '/Test_CCP4/fav8.dx.save'
output_file = 'test.ccp4'
pt.datafiles.convert(saved_dx_file, output_file)
assert os.path.exists("test.ccp4")
def test_pickle_progress_bar():
traj = pt.datafiles.load_tz2()
with tempfolder():
t2 = ProgressBarTrajectory(traj)
pt.to_pickle(t2, 'test.pk')
t3 = pt.read_pickle('test.pk')
data0 = pt.rmsd(traj)
data1 = pt.rmsd(t2)
data2 = pt.rmsd(t3)
aa_eq(data0, data1)
aa_eq(data0, data2)
def test_io_load_and_save_0():
traj = pt.iterload(filename=tc5b_trajin, top=tc5b_top)[:10]
indices = list(range(2, 3, 5)) + [3, 8, 9, 8]
with tempfolder():
pt.write_traj(filename="test_io_saved_.x",
traj=traj,
frame_indices=indices,
overwrite=True)
# check frames
traj2 = pt.iterload(filename="test_io_saved_.x", top=tc5b_top)
assert traj2.n_frames == len(indices)
def test_run_2(self):
# load traj
traj = pt.iterload(tc5b_trajin, tc5b_top)
dslist = CpptrajDatasetList()
dflist = DataFileList()
# creat ActionList to hold actions
alist = ActionList()
alist.add(adict['distance'], ":2@CA :10@CA out _dist.out", traj.top,
dslist, dflist)
with tempfolder():
alist.compute([traj.iterchunk()])
assert len(dslist) == 1
assert dslist[0].size == traj.n_frames
def test_load_and_save_0():
# need to load to Trajectory to save
traj = pt.iterload(filename=tc5b_trajin, top=tc5b_top)[:]
indices = list(range(2, 3, 5)) + [3, 7, 9, 8]
with tempfolder():
pt.write_traj(filename="test_io_saved_.x",
traj=traj[:],
frame_indices=indices,
overwrite=True)
# check frames
traj2 = pt.iterload(filename="test_io_saved_.x", top=tc5b_top)
# about 50% failures
assert traj2.n_frames == len(indices)
def test_write_force_from_scratch():
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
assert traj[0].force is None
def add_force(traj):
for frame in traj:
frame.force = np.zeros(traj.n_atoms * 3) + 1.
yield frame
with tempfolder():
out_fn = 'out.nc'
pt.write_traj(out_fn, traj=add_force(traj), top=traj.top, force=True)
traj2 = pt.iterload(out_fn, top=traj.top)
assert traj2.metadata['has_force']
assert not traj2.metadata['has_velocity']
def test_load_and_save_topology():
top = traj_tz2_ortho.top
with tempfolder():
top.save('test.prmtop')
top2 = pt.load_topology('test.prmtop')
assert top.n_atoms == top2.n_atoms, 'must have the same n_atoms'
# shortcut
pt.load_topology(fn('tz2.parm7')).save('tz2_0.parm7')
top3 = pt.load_topology('tz2_0.parm7')
assert top3.n_atoms == 223, 'must have 223 atoms'
assert top3.n_residues == 13, 'must have 13 residues'
# raise
with pytest.raises(ValueError):
pt.load_topology(100)
def test_save_traj_from_file():
traj = pt.iterload(tc5b_trajin, tc5b_top)[:5]
with tempfolder():
pt.write_traj(filename="test_0.binpos",
traj=traj,
top=tc5b_top,
overwrite=True)
savedtraj = pt.iterload("test_0.binpos", traj.top)
assert savedtraj.n_frames == traj.n_frames
# write_xyz
pt.write_traj("test_0.nc", traj.xyz, top=tc5b_top, overwrite=True)
aa_eq(pt.iterload('test_0.nc', traj.top).xyz, traj.xyz)
# write single Frame
pt.write_traj("test_0.nc", traj[0], top=traj.top, overwrite=True)
aa_eq(pt.iterload('test_0.nc', traj.top).xyz, traj[0].xyz)
# raise if traj is None
with pytest.raises(ValueError):
pt.write_traj("test_0.nc", None, overwrite=True)
# raise if _top is None
fi = pt.pipe(traj, [
'autoimage',
])
with pytest.raises(ValueError):
pt.write_traj("test_0.nc", traj=fi, overwrite=True)
# raise if Frame with frame_indices:
with pytest.raises(ValueError):
pt.write_traj("test_0.nc",
traj[0],
top=tc5b_top,
frame_indices=[3, 2],
overwrite=True)
# raise if Frame with no Topology
with pytest.raises(ValueError):
pt.write_traj("test_0.nc", traj[0], overwrite=True)
# test if xyz is not c-contiguous
# pytraj will autoconvert to c-contiguous
xyz = np.asfortranarray(traj.xyz)
# make sure no ValueError or TypeError is raised
pt.write_traj('xyz.nc', xyz, top=traj.top, overwrite=True)
def test_trajectory_writer_write_Trajectory():
"""test write Trajectory"""
farray = pt.load(
fn("Tc5b.x"), fn("Tc5b.top"), frame_indices=list(range(10)))
with tempfolder():
pt.write_traj(
"test_write_output.x", farray, top=farray.top, overwrite=True)
pt.write_traj(
"test_pdb_1.dummyext", farray[0], top=farray.top, overwrite=True)
# test 'save'
farray.save("test_write_output_save_method.x", overwrite=True)
# reproduce result?
f0 = pt.iterload("test_write_output.x", fn('Tc5b.top'))
f1 = pt.iterload("test_write_output_save_method.x", fn('Tc5b.top'))
aa_eq(f0[:, :, :].xyz, f1[:, :, :].xyz)
def test_write_force_from_scratch():
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
assert traj[0].force is None
def add_force(traj):
for frame in traj:
frame.force = np.zeros(traj.n_atoms * 3) + 1.
yield frame
with tempfolder():
out_fn = 'out.nc'
with TrajectoryWriter(out_fn, top=traj.top, crdinfo={'has_force': True}) as writer:
for frame in add_force(traj):
writer.write(frame)
traj2 = pt.iterload(out_fn, top=traj.top)
assert traj2.metadata['has_force']
assert not traj2.metadata['has_velocity']
def test_cluster_hieragglo():
command = """
parm {}
trajin {}
createcrd crd1
cluster crdset crd1 C0 !@H hieragglo epsilon 0.8 averagelinkage
""".format(tz2_top, tz2_trajin)
with tempfolder():
state = pt.load_cpptraj_state(command)
with capture_stdout() as (cpp_out, _):
state.run()
traj = pt.iterload(tz2_trajin, tz2_top)
data = pt.cluster.hieragglo(traj,
mask='!@H=',
options='epsilon 0.8 averagelinkage')
aa_eq(state.data[-2], data.cluster_index)
def test_cluster_dbscan():
command = """
parm {}
trajin {}
createcrd crd1
cluster crdset crd1 C0 @CA dbscan epsilon 1.7 minpoints 5
""".format(tz2_top, tz2_trajin)
with tempfolder():
state = pt.load_cpptraj_state(command)
with capture_stdout() as (out, _):
state.run()
traj = pt.iterload(tz2_trajin, tz2_top)
data = pt.cluster.dbscan(traj,
mask='@CA',
options='epsilon 1.7 minpoints 5')
aa_eq(state.data[-2], data.cluster_index)
def test_reference_with_different_topology_basic(self):
traj1 = pt.iterload(filename=tc5b_trajin, top=tc5b_top)
traj2 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
# re-establish ActionList
dslist = CpptrajDatasetList()
dslist.add('reference', name='myref')
dslist[0].top = traj2.top
dslist[0].add_frame(traj2[0])
actlist = pt.ActionList(['rmsd @1-11 @CB ref myref'],
top=traj1.top,
dslist=dslist)
for frame in traj1:
actlist.compute(frame)
# raise if ref_mask is given but not mask
self.assertRaises(ValueError,
lambda: pt.rmsd(traj1, ref=3, ref_mask='@CB'))
self.assertRaises(
ValueError, lambda: pt.rmsd(traj1, ref=traj2[:1], ref_mask='@CB'))
# assert to cpptraj
tc5b_traj = traj1[:]
tz2_traj = traj2[:1]
cm = '''
parm {} [tc5b]
trajin {}
parm {} [tz2]
reference {} parm [tz2] 1 [myref]
rms myrmsd ref [myref] @1-10 @11-20
'''.format(tc5b_top, tc5b_trajin, tz2_top, tz2_trajin)
print(cm)
state = pt.load_cpptraj_state(cm)
with tempfolder():
state.run()
expected_rmsd = state.data[-1].values
rmsd_data = pt.rmsd(tc5b_traj,
mask='@1-10',
ref=tz2_traj,
ref_mask='@11-20')
aa_eq(expected_rmsd, rmsd_data)
def test_trajectory_writer_open_close():
farray = pt.load(
fn("Tc5b.x"), fn("Tc5b.top"), frame_indices=list(range(10)))
frame0 = farray[0]
with tempfolder():
trajout = TrajectoryWriter()
trajout.open(filename="test.x", top=farray.top, overwrite=True)
trajout.write(frame0)
# add more frames
for i in range(5, 8):
trajout.write(farray[i])
trajout.close()
farray = Trajectory()
farray.top = pt.load_topology(fn('Tc5b.top'))
farray.load("test.x")
def test_load_and_save_1():
# do not support frame_indices for TrajectoryIterator
with pytest.raises(ValueError):
pt.iterload(filename=tc5b_trajin, top=tc5b_top, frame_indices=[0, 5])
traj = pt.iterload(filename=tc5b_trajin, top=tc5b_top)
indices = list(range(2, 4)) + [3, 7, 9, 8]
with tempfolder():
pt.write_traj(filename="test_io_saved.pdb",
traj=traj,
frame_indices=indices,
overwrite=True)
# check frames
traj = pt.iterload(filename="test_io_saved.pdb", top=tc5b_top)
assert traj.n_frames == len(indices)
assert traj.top.n_atoms == 304
def test_write_force_and_velocity():
trajin = cpptraj_test_dir + '/Test_systemVF/systemVF.nc'
tn = cpptraj_test_dir + '/Test_systemVF/systemVF.parm7'
traj = pt.iterload(trajin, tn)
assert traj.metadata['has_force']
assert traj.metadata['has_velocity']
with tempfolder():
fn2 = 'test.nc'
traj.save(fn2, overwrite=True, crdinfo=traj.metadata)
traj2 = pt.iterload(fn2, traj.top)
assert traj2.metadata['has_force']
assert traj2.metadata['has_velocity']
forces_traj = np.array([frame.force.copy() for frame in traj])
forces_traj2 = np.array([frame.force.copy() for frame in traj2])
aa_eq(forces_traj, forces_traj2)
def test_reference(self):
traj = pt.iterload(fn("tz2.truncoct.nc"), fn("tz2.truncoct.parm7"))
txt = '''
reference {} 2 2
rmsd :2-11 refindex 0 perres perresout center.agr range 1 perrescenter
'''.format(fn('tz2.truncoct.nc'))
with tempfolder():
state = pt.load_batch(traj, txt).run()
# state.data has 3 datasets: ref, rmsd, rmsd perres
# cpptraj use 2nd reference
rmsd0 = pt.rmsd(traj, ref=1, mask=':2-11')
rmsdperres = pt.rmsd_perres(traj,
ref=1,
mask=':2-11',
perres_mask='*',
resrange='1',
perres_center=True)
aa_eq(rmsd0, state.data[2])
aa_eq(rmsdperres[1], state.data[3].values)
def test_write_velocity_from_scratch():
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
assert traj[0].velocity is None
def add_velocity(traj):
for frame in traj:
frame.velocity = np.zeros(traj.n_atoms * 3) + 1.
yield frame
with tempfolder():
out_fn = 'out.nc'
with TrajectoryWriter(out_fn,
top=traj.top,
crdinfo={'has_velocity': True}) as writer:
for frame in add_velocity(traj):
writer.write(frame)
traj2 = pt.iterload(out_fn, top=traj.top)
assert traj2.metadata['has_velocity']
assert not traj2.metadata['has_force']
# make sure no segmentation fault
# issues/1486
assert traj2[:].n_frames == traj2.n_frames
def test_write_time():
top_fname = os.path.join(cpptraj_test_dir, 'FtuFabI.NAD.TCL.parm7')
traj_fname = os.path.join(cpptraj_test_dir, 'FtuFabI.NAD.TCL.nc')
traj = pt.iterload(traj_fname, top_fname)
time_arr = np.array([
37031.0, 37032.0, 37033.0, 37034.0, 37035.0, 37036.0, 37037.0, 37038.0,
37039.0, 37040.0
])
aa_eq([f.time for f in traj], time_arr)
traj2 = traj[:] # in memory
aa_eq(np.array([f.time for f in traj2]), time_arr)
traj3 = traj2[:] # slice in memory trajectory again
aa_eq(np.array([f.time for f in traj3]), time_arr)
with tempfolder():
pt.write_traj('test.nc', traj, time=True)
out_traj = pt.iterload('test.nc', traj.top)
aa_eq(np.array([f.time for f in out_traj]), time_arr)
# load method
out_traj2 = pt.load('test.nc', traj.top)
aa_eq(np.array([f.time for f in out_traj2]), time_arr)
def test_trajlist(self):
trajlist = []
N = 4
for i in range(N):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
trajlist.append(traj)
traj0 = trajlist[0]
with tempfolder():
pt.write_traj("test_savedtrajlist.x",
traj=trajlist,
top=traj.top,
overwrite=True)
# test loading
traj2 = pt.iterload("test_savedtrajlist.x", traj.top)
Nsize = int(traj2.n_frames / 4)
traj0_new = traj2[:Nsize]
for frame0, frame0_new in zip(traj0, traj0_new):
assert (frame0.rmsd(frame0_new) < 1E-3)
def test_short_save():
with tempfolder():
(pt.iterload(fn('tz2.nc'), fn('tz2.parm7')).save('mini.nc',
overwrite=True))
assert pt.iterload(
'mini.nc', fn('tz2.parm7')).n_frames == 101, 'must be 101 frames'
def test_rdf(self):
traj = pt.iterload(fn("tz2.truncoct.nc"),
fn("tz2.truncoct.parm7"),
frame_slice=(0, 10))
command = '''
radial Radial.agr 0.5 10.0 :5@CD :WAT@O
radial cRadial.agr 0.5 10.0 :5 :WAT@O center1
radial cRadial.agr 0.5 10.0 :5 :WAT@O center2
radial cRadial.agr 0.5 20.0 :3 :WAT@O
radial cRadial.agr 0.5 20.0 :3 :WAT@O noimage
radial radial.dat 0.5 10.0 :5@CD :WAT@O
radial radial2.dat 0.25 10.0 :5@CD :WAT@O
radial radial2.dat 0.25 10.0 :5@CD :WAT@O volume
'''
state = pt.load_batch(traj, command)
with tempfolder():
# get data directly from cpptraj
state.run()
# get data from pytraj
data0 = pt.rdf(traj,
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=10.0,
solvent_mask=':5@CD')
data01 = pt.rdf(traj,
solvent_mask=':5@CD',
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=10.0)
data1 = pt.rdf(traj,
solvent_mask=':5',
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=10.0,
center_solvent=True)
data2 = pt.rdf(traj,
solvent_mask=':5',
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=10.0,
center_solute=True)
data3 = pt.rdf(traj,
solvent_mask=':3',
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=20.0,
center_solute=False)
data4 = pt.rdf(traj,
solvent_mask=':3',
solute_mask=':WAT@O',
bin_spacing=0.5,
maximum=20.0,
center_solute=False,
image=False)
data5 = pt.rdf(traj,
solute_mask=':WAT@O',
bin_spacing=0.25,
maximum=10.0,
solvent_mask=':5@CD')
# solvent_mask is array
solvent_indices = pt.select(':WAT@O', traj.top)
data6 = pt.rdf(traj,
solvent_mask=':5@CD',
solute_mask=solvent_indices,
bin_spacing=0.25,
maximum=10.0)
# volume
data7 = pt.rdf(traj,
solvent_mask=':5@CD',
solute_mask=':WAT@O',
bin_spacing=0.25,
maximum=10.0,
volume=True)
# do assertion
aa_eq(data0[1], state.data[1], decimal=7)
aa_eq(data1[1], state.data[2], decimal=7)
aa_eq(data2[1], state.data[3], decimal=7)
aa_eq(data3[1], state.data[4], decimal=7)
aa_eq(data4[1], state.data[5], decimal=7)
aa_eq(data7[1], state.data[8], decimal=7)
# default solvent mask :WAT@O
aa_eq(data01[1], state.data[1], decimal=7)
steps = np.loadtxt('radial.dat').T[0]
aa_eq(data0[0], steps)
steps2 = np.loadtxt('radial2.dat').T[0]
aa_eq(data5[0], steps2)
aa_eq(data6[0], steps2)
