def test_write_xyz(self):
xyz = self.traj.xyz
fname = 'test_xyz.nc'
with tempfolder():
pt.write_traj(fname, xyz, top=self.traj.top, overwrite=True)
t0 = pt.iterload(fname, top=self.traj.top)
aa_eq(self.traj.xyz, t0.xyz)
def test_overwrite():
trajin, tn = fn("tz2.nc"), fn("tz2.parm7")
with tempfolder():
traj = pt.iterload(trajin, tn)
pt.write_traj('what.nc', traj)
# ensure no IOError is raised.
pt.write_traj('what.nc', traj, overwrite=True)
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 pickframe(thread, direction):
# Picks a random frame from the last valid trajectory of thread in the given direction (either 'fwd' or 'bwd'). The
# selection is chosen from within the first n_adjust steps. Once the desired frame is chosen, pytraj is used to
# extract it from the given trajectory and the name of the new shooting point .rst file is returned sans extension.
if thread.last_valid == '0': # to catch case where 'last_valid' is still the initial shooting point because a new valid transition path has not been found
return thread.start_name # will cause the new start_name to be unchanged from the previous run
frame_number = random.randint(0, n_adjust)
this_suffix = thread.last_valid
traj = pytraj.iterload(thread.basename + '_' + this_suffix + '_' +
direction + '.nc',
thread.prmtop,
format='.nc',
frame_slice=[(frame_number, frame_number)])
new_suffix = str(int(thread.suffix) + 1)
pytraj.write_traj(thread.basename + '_' + new_suffix + '.rst',
traj,
format='.rst7')
open(
'as.log',
'a').write('\nInitializing next shooting point from shooting run ' +
thread.basename + '_' + str(thread.last_valid) + ' in ' +
direction + ' direction, frame number ' + str(frame_number))
return thread.basename + '_' + new_suffix
def test_buffer_not_c_contiguous(self):
# source code was lightly adapted from jmborr
# https://github.com/Amber-MD/pytraj/issues/991
traj = pt.load('data/issue991/short.dcd',
'data/issue991/pdb.gz',
mask='(!:1-256)&(@H1,@H2,@H3,@H4,@H5)')
# Trajectory of the center of mass of the first two residues
minitraj = np.empty((2, traj.n_frames, 3))
minitraj[0] = pt.center_of_mass(traj, mask=':1')
minitraj[1] = pt.center_of_mass(traj, mask=':2')
minitraj = minitraj.transpose((1, 0, 2))
# "topology" using the first two atoms
minitop = traj.top['@1,2']
# Save trajectory
# make sure there is no ValueError
# something is wrong with pdb, crd extension when loading with
# minitop (poor topology?)
# make issue in cpptraj too?
for ext in ['nc', 'dcd', 'mdcrd', 'crd', 'pdb', 'trr']:
fn = 'output/junk.' + ext
pt.write_traj(filename=fn,
traj=minitraj,
top=minitop,
overwrite=True)
# load coord back to make sure we correctly write it
new_traj = pt.iterload(fn, minitop)
# mdcrd, crd, pdb has only 3 digits after decimal
decimal = 5 if ext in ['nc', 'dcd', 'trr'] else 3
aa_eq(minitraj, new_traj.xyz, decimal=decimal)
def create_traj(top, mol_name, nconf, Flag, filename=None):
''' top is topology
mol_name is molecule name, must be set of structures number mol_name001 to mol_namexn1
nconf is conformation number
Flag:
0 if using index moleculename.000 to moleculename.nconf (moleculename.012)
1 filename where each line has a file path to the rst7 file
'''
crd = [] # holds the filename for list of structures
nconf = int(nconf)
if Flag == 0:
for n in np.arange(1, nconf + 1, 1):
n = "%03d" % (n, ) # make n 000 001 002 etc
string = "%s%s.energy.rst7" % (mol_name, n)
crd.append(string)
# read from a file name that have list of names with no extension
elif Flag == 1:
file_list = np.genfromtxt(filename, dtype=str, usecols=0)
for string in file_list:
string = "%s.rst7" % (string)
crd.append(string)
# read from a list of structures
elif Flag == 2:
for string in filename:
crd.append(string)
# write trajectory
traj = pt.iterload(crd, top)
pt.write_traj('%s.nc' % (mol_name), traj, overwrite=True)
print("trajectory name %s.nc is completed" % (mol_name))
def test_buffer_not_c_contiguous(self):
# source code was lightly adapted from jmborr
# https://github.com/Amber-MD/pytraj/issues/991
traj = pt.load('data/issue991/short.dcd',
'data/issue991/pdb.gz',
mask='(!:1-256)&(@H1,@H2,@H3,@H4,@H5)')
# Trajectory of the center of mass of the first two residues
minitraj = np.empty((2, traj.n_frames, 3))
minitraj[0] = pt.center_of_mass(traj, mask=':1')
minitraj[1] = pt.center_of_mass(traj, mask=':2')
minitraj = minitraj.transpose((1, 0, 2))
# "topology" using the first two atoms
minitop = traj.top['@1,2']
# Save trajectory
# make sure there is no ValueError
# something is wrong with pdb, crd extension when loading with
# minitop (poor topology?)
# make issue in cpptraj too?
for ext in ['nc', 'dcd', 'mdcrd', 'crd', 'pdb', 'trr']:
fn = 'output/junk.' + ext
pt.write_traj(filename=fn,
traj=minitraj,
top=minitop,
overwrite=True)
# load coord back to make sure we correctly write it
new_traj = pt.iterload(fn, minitop)
# mdcrd, crd, pdb has only 3 digits after decimal
decimal = 5 if ext in ['nc', 'dcd', 'trr'] else 3
aa_eq(minitraj, new_traj.xyz, decimal=decimal)
def test_trajectory_writer():
traj = pt.iterload(tc5b_trajin, tc5b_top)
with tempfolder():
pt.write_traj("test_1.pdb", traj[0], top=traj.top, overwrite=True)
pt.write_traj("test_1.dcd", traj[0], top=traj.top, overwrite=True)
with TrajectoryWriter("test_1", overwrite=True) as trajout:
trajout.write(traj[0])
def test_save_frameiter():
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
traj(0, 8, 2, mask='@CA').save('dummy_test0.nc', overwrite=True)
pt.write_traj('dummy_test1.nc', traj(0, 8, 2, mask='@CA'), overwrite=True)
new_top = traj.top._get_new_from_mask('@CA')
t0 = pt.iterload('dummy_test0.nc', new_top)
t1 = pt.iterload('dummy_test1.nc', new_top)
aa_eq(t0.xyz, t1.xyz)
def align(self, ligcom, lamda):
'''Concat all replicate trajectories, RMSD fit to first frame and autoimage'''
# glob all trajs
# keep corresponding parm at each lambda to maintain partial charges
if ligcom == 'ligands':
paths = sorted(
list(Path(self.dest_dir, 'ligands').glob(f'*/{lamda}/*.nc')))
parm = sorted(
list(
Path(self.dest_dir, 'ligands').glob(f'*/{lamda}/*.parm7')))
elif ligcom == 'complex':
paths = sorted(
list(Path(self.dest_dir, 'complex').glob(f'*/{lamda}/*.nc')))
parm = sorted(
list(
Path(self.dest_dir, 'complex').glob(f'*/{lamda}/*.parm7')))
paths = list(map(str, paths))
# read first parm
parm = str(parm[0])
# iterload all
traj = pt.iterload(paths, parm)
# rmsd align to first
if ligcom == 'ligands':
mask = self.ligand_mask
elif ligcom == 'complex':
n_residues = len(traj.top.select('@CA'))
mask = f':1-{n_residues}@CA'
rmsd = pt.rmsd(traj, mask=mask, ref=0)
# autoimage
traj = traj.autoimage()
# save concat traj, restart, parm
# bug? only works with ncrst restarts, not rst7
traj_save = str(
Path(self.dest_dir, f'concat_{ligcom}', lamda, 'ti001.nc'))
restart_save = Path(self.dest_dir, f'concat_{ligcom}', lamda,
'ti.ncrst')
parm_save = str(
Path(self.dest_dir, f'concat_{ligcom}', lamda, 'ti.parm'))
pt.write_traj(traj_save, traj)
pt.write_traj(str(restart_save), traj, frame_indices=[0])
pt.save(parm_save, traj.top)
# fix restart name, remove suffix
Path(f'{restart_save}.1').rename(restart_save)
logging.info(f'align {self.dest_path} {ligcom} {lamda}')
def calc_alignment(trajfiles, topfile, rdc_types, masks, step=1, seg=None):
Aa = dict()
Ar = dict()
eigx = dict()
eigz = dict()
for n, tr in enumerate(trajfiles):
traj = pt.iterload(tr, topfile, frame_slice=(0, -1, step))
if seg==None:
residues = [r.index+1 for r in traj.topology.residues if r.name!='PRO' and r.index>0]
else:
residues = [r.index+1 for r in traj.topology.residues
if r.name!='PRO' and r.index>=1 and r.index+1>=seg[0] and r.index+1<=seg[1]]
Aa_temp = []
Ar_temp = []
eigx_temp = []
eigz_temp = []
for f, frame in enumerate(traj):
pt.write_traj("temp.pdb", frame, top=traj.top, overwrite=True)
if seg==None:
os.system("pales -pdb temp.pdb > pales.out")
else:
os.system("pales -pdb temp.pdb -s1 "+str(seg[0])+" -sN "+str(seg[1])+" > pales.out")
os.system("grep \"DATA EIGENVALUES (Sxx_d,Syy_d,Szz_d)\" pales.out > eigenvalues.out")
with open('eigenvalues.out', 'r') as tf:
data = tf.readlines()
data = data[0].split()
Sxx_d=float(data[-3])
Syy_d=float(data[-2])
Szz_d=float(data[-1])
Aa_temp.append(Szz_d)
Ar_temp.append((2./3.)*(Sxx_d-Syy_d))
os.system("grep \"DATA EIGENVECTORS X_AXIS\" pales.out > eigenvalues1.out")
with open('eigenvalues1.out', 'r') as tf:
data = tf.readlines()
data = data[0].split()
eigx_temp.append([float(data[-3]), float(data[-2]), float(data[-1])])
os.system("grep \"DATA EIGENVECTORS Z_AXIS\" pales.out > eigenvalues3.out")
with open('eigenvalues3.out', 'r') as tf:
data = tf.readlines()
data = data[0].split()
eigz_temp.append([float(data[-3]), float(data[-2]), float(data[-1])])
eigz[n] = eigz_temp
eigx[n] = eigx_temp
Aa[n] = Aa_temp
Ar[n] = Ar_temp
return Aa, Ar, eigx, eigz
def pack(traj, mol, n_copies, unitcells=None, ig=8888, grid_spacing=0.2):
'''
Parameters
----------
traj : pytraj.Trajectory
mol : pytraj.Trajectory
unitcells : None or np.ndarray, dim=2
if None, use box info from traj else use it
n_copies : number of `mol`
ig : int
randome seed
grid_spacing : float
'''
add_to_box_exe = amberbin('AddToBox') or 'AddToBox'
input_pdb = 'input.pdb'
mol_pdb = 'mol.pdb'
out_pdb = 'out.pdb'
assert mol.n_frames == 1
total_n_atoms = traj.n_atoms + mol.n_atoms * n_copies
new_traj_xyz = np.empty((traj.n_frames, total_n_atoms, 3), dtype='f8')
with tempfolder():
mol.save(mol_pdb, overwrite=True)
for index, frame in enumerate(traj):
pytraj.write_traj(input_pdb, traj=frame, top=traj.top, overwrite=True)
parm = parmed.load_file(input_pdb)
if unitcells is not None:
parm.box = unitcells[index]
else:
assert frame.has_box(), 'must have box info'
parm.box = frame.box
# add remark 290
parm.save(input_pdb, overwrite=True)
command = [
add_to_box_exe,
'-c', input_pdb,
'-a', mol_pdb,
'-na', str(n_copies),
'-IG', str(ig),
'-G', str(grid_spacing),
'-o', out_pdb
]
try:
subprocess.check_output(command, stderr=subprocess.STDOUT)
new_traj_xyz[index] = pytraj.load(out_pdb).xyz
except subprocess.CalledProcessError as e:
return e.output.decode()
top = pytraj.load_topology(out_pdb)
return pytraj.Trajectory(xyz=new_traj_xyz, top=top)
def test_0(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
traj(0, 8, 2, mask='@CA').save('output/test0.nc', overwrite=True)
pt.write_traj('./output/test1.nc',
traj(0, 8, 2, mask='@CA'),
overwrite=True)
new_top = traj.top._get_new_from_mask('@CA')
t0 = pt.iterload('./output/test0.nc', new_top)
t1 = pt.iterload('./output/test1.nc', new_top)
aa_eq(t0.xyz, t1.xyz)
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_makestructure(self):
# FIXME: What does this test?
# https://github.com/Amber-MD/cpptraj/issues/27
# load only 1st frame
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
# polyproline II dihedral to residues 1-13
t0 = traj[:1].copy()
pt.make_structure(t0, 'pp2:1-13')
pt.write_traj('dummy_test0.pdb', t0, options='model', overwrite=True)
t1 = traj[:1].copy()
pt.make_structure(t1, "chi1:8:N:CA:CB:CG:35")
pt.write_traj('dummy_test1.pdb', t1, options='model', overwrite=True)
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_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_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_write_pdb():
# TODO: get absolute path so we can use `tempfolder`
# if not: wrong dir if using TrajectoryIterator
traj = pt.iterload(tc5b_trajin, tc5b_top)[:]
TrajectoryWriter()
# multiple pdb in multiple files, using `save` method in traj
with tempfolder():
basename = "test_pdb_files.pdb"
traj.save(basename, overwrite=True, options="multi")
for i in range(10):
fname = basename + "." + str(i + 1) # cpptraj use `1`
frame = pt.iterload(fname, traj.top)[0]
aa_eq(frame.xyz, traj[i].xyz)
# multiple pdb in multiple files, using `pt.write_traj`
with tempfolder():
basename = "test_pdb_files_pt_write_traj.pdb"
pt.write_traj(basename, traj, overwrite=True, options="multi")
for i in range(10):
fname = basename + "." + str(i + 1) # cpptraj use `1`
frame = pt.iterload(fname, traj.top)[0]
aa_eq(frame.xyz, traj[i].xyz)
with pytest.raises(IOError):
# write files again, raise if file exists
pt.write_traj(basename, traj, overwrite=False, options="multi")
# keepext
with tempfolder():
basename = 'test_pdb_files_pt_write_traj.pdb'
basename2 = 'test_pdb_files_pt_write_traj'
pt.write_traj(basename, traj, overwrite=True, options="multi keepext")
for i in range(10):
fname = "{}.{}.pdb".format(basename2, i + 1)
frame = pt.iterload(fname, traj.top)[0]
aa_eq(frame.xyz, traj[i].xyz)
with pytest.raises(IOError):
pt.write_traj(basename,
traj,
overwrite=False,
options="multi keepext")
# multiple pdb in SINGLE file
with tempfolder():
basename = "test_pdb_files.pdb"
traj.save(basename, overwrite=True)
traj2 = pt.load(basename, traj.top)
aa_eq(traj.xyz, traj2.xyz)
# multiple pdb in SINGLE file with `model` keyword
# write to output so we can manually check
basename = "test_pdb_files_model.pdb"
with tempfolder():
traj.save(basename, overwrite=True, options='model')
traj3 = pt.load(basename, traj.top)
aa_eq(traj.xyz, traj3.xyz)
def main(traj):
'''
最低能量（势能）构象帧提取
Parameters
----------
traj: MD轨迹 Pytraj对象
'''
print('\nStart Lowest Energy Structure Analysis...')
# 调用perl脚本分析MD out数据
# process_mdout.perl脚本为AMBER官方脚本
os.system(
'cd ./pynalysis/lowestenergy && chmod 777 ./process_mdout.perl && ./process_mdout.perl ../../npt/npt.out'
)
print('\nProcessing Complete.', end='\n')
print('\nSearching Lowest Energy Time...', end='\n')
temp = os.popen('''
cat ./pynalysis/lowestenergy/summary.EPTOT | awk '{if($2<min) {min=$2;print $1" "min}}'
''') #EPTOT总势能
lines = temp.readlines()
match_time = re.match('\d{1,5}.000', lines[-1].strip()) # 匹配最后一行（最低能量）的时间
lowest_energy_time = match_time.group() # 获取时间具体值
for info in lines:
print(info.strip())
print('\nLowest Energy Time:', lowest_energy_time, 'ps')
os.environ['lowest_energy_time'] = str(
lowest_energy_time) # 向shell传递变量 搜索最低能量帧
temp2 = os.popen('grep $lowest_energy_time ./npt/npt.out')
frame_temp = temp2.read().strip() # 最低能量帧信息
print('\nLowest Energy Frame info: \n', frame_temp)
# 提取最低能量帧数: 帧数 = 步数 / 1000
match_frame = re.search('(?<=NSTEP\s=\s)\d*(?=000)', str(frame_temp))
lowest_energy_frame = int(match_frame.group())
print('\nLowest Energy Frame:', lowest_energy_frame)
pt.write_traj('./pynalysis/lowestenergy/lowest_energy_structure.pdb',
traj,
frame_indices=[lowest_energy_frame],
overwrite=True) # 输出最低能量构象PDB文件
print('\nLowest Energy Structure PDB file saved.')
def test_0(self):
# TrajectoryIterrator
# status: failed
from pytraj.compat import zip
traj = pt.iterload("./data/tz2.ortho.nc", "./data/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('data/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_trajectory_writer_write_Trajectory(tmpdir):
"""test write Trajectory"""
farray = pt.load(fn("Tc5b.x"),
fn("Tc5b.top"),
frame_indices=list(range(10)))
with tmpdir.as_cwd():
pt.write_traj("test_write_output.x", farray, overwrite=True)
pt.write_traj("test_pdb_1.dummyext", farray[:1], 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_on_disk_trajectory(self, tmpdir):
with tmpdir.as_cwd():
# 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_0(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
traj(0, 8, 2, mask='@CA').save('output/test0.nc', overwrite=True)
pt.write_traj('./output/test1.nc',
traj(0,
8,
2,
mask='@CA'),
overwrite=True)
new_top = traj.top._get_new_from_mask('@CA')
t0 = pt.iterload('./output/test0.nc', new_top)
t1 = pt.iterload('./output/test1.nc', new_top)
aa_eq(t0.xyz, t1.xyz)
def test_load_and_save_0(self):
# need to load to Trajectory to save
traj = pt.iterload(filename="./data/Tc5b.x", top="./data/Tc5b.top")[:]
indices = list(range(2, 3, 5)) + [3, 7, 9, 8]
pt.write_traj(filename="./output/test_io_saved_.x",
traj=traj[:],
top="./data/Tc5b.top",
frame_indices=indices,
overwrite=True)
# check frames
traj2 = pt.iterload(filename="./output/test_io_saved_.x",
top="./data/Tc5b.top")
# about 50% failures
assert traj2.n_frames == len(indices)
def pdb_extract(index_tuple,traj_list, top_path):
"""Extracts frames from trajectories and saves them in pdb format
Parameters:
index_tuple: Gives list of tuples of length==n;
first and second element of tuple correspond to trajectory and frame number respectively.
traj_list: List of string; each string is address of each trajectory file
top_path: String; path of topology file
Returns:
Creates pdb files
"""
import pytraj as pt
L=len(index_tuple)
for i in range(0,L):
traj=traj_list[index_tuple[i][0]]
frame=index_tuple[i][1]
traj=pt.iterload(traj,top_path,frame_slice=(frame, frame+1))
pt.write_traj(str(i)+'.pdb', traj, overwrite=True)
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_regular(self):
traj = self.traj.copy()
assert traj[0].has_box() == True
with tempfolder():
# write traj with nobox info
fname = "traj_nobox.nc"
pt.write_traj(fname, traj, options='nobox')
t = pt.load(fname, traj.top)
assert t[0].has_box() == False
aa_eq(t.xyz, traj.xyz)
# write from frame_iter, need to provide top
fname = "traj_frame_iter.nc"
# raise ValueError if there is no Topology
pt.write_traj(fname, traj(), top=traj.top, overwrite=True)
t = pt.iterload(fname, traj.top)
aa_eq(t.xyz, traj.xyz)
def test_regular(self):
traj = self.traj.copy()
assert traj[0].has_box() == True
with tempfolder():
# write traj with nobox info
fname = "traj_nobox.nc"
pt.write_traj(fname, traj, options='nobox')
t = pt.load(fname, traj.top)
assert t[0].has_box() == False
aa_eq(t.xyz, traj.xyz)
# write from frame_iter, need to provide top
fname = "traj_frame_iter.nc"
# raise ValueError if there is no Topology
pt.write_traj(fname, traj(), top=traj.top, overwrite=True)
t = pt.iterload(fname, traj.top)
aa_eq(t.xyz, traj.xyz)
def do_clustering():
# set random seed for reproducibility
np.random.seed(1)
traj = pt.iterload('GAAC3.5000frames.nc', 'GAAC.topo')
n_frames = 50
# as an example, we extract only first 50 frames and do clustering for @P atoms
t0 = traj[:n_frames, '@P']
# need to convert to 2D array, shape=(n_frames, n_features)
X = pt.tools.as_2darray(t0)
# specify the total number of clusters
n_clusters = 10
# do it
data = k_means(X,
n_clusters,
init='k-means++',
precompute_distances='auto',
n_init=10,
max_iter=300,
verbose=False,
tol=0.0001,
random_state=None,
copy_x=False,
n_jobs=-1,
return_n_iter=False)
# convert 2D to 3D array
x0 = pt.tools.as_3darray(data[0])
print('cluster index: ', data[1].tolist())
# write output, dump all frames to a single pdb file, each snapshot is seperated by
# 'MODEL' keyword. good for visualizing in VMD.
# first, need to create a frame iterator for ``pytraj.write_traj``
frame_iter = pt.iterframe_from_array(x0, t0.n_atoms, range(n_clusters))
pt.write_traj('output.pdb',
frame_iter,
top=t0.top,
options='model',
overwrite=True)
def test_closest(self):
# raise if not has solvent
traj0 = pt.iterload(fn('tz2.nc'),
fn('tz2.parm7'),
frame_slice=[(0, 2)])
self.assertRaises(RuntimeError, lambda: pt.closest(traj0))
traj = pt.iterload(fn('tz2.ortho.nc'),
fn('tz2.ortho.parm7'),
frame_slice=[(0, 2)])
fi, top = pt.closest(traj)
coords = []
for frame in fi:
coords.append(frame.xyz.copy())
assert isinstance(frame, pt.Frame), 'must be Frame'
# make a Trajectory
fi, top = pt.closest(traj)
xyz = pt.get_coordinates(fi)
t0 = pt.Trajectory(xyz=xyz, top=top)
aa_eq(np.array(coords), t0.xyz)
# test write to disk
fi, top = pt.closest(traj)
pt.write_traj(outputname('fi.nc'), fi, top=top, overwrite=True)
# load back
t1 = pt.load(outputname('fi.nc'), top=top)
aa_eq(t0.xyz, t1.xyz)
# make sure n_sovent=10 (default)
n_solvents = 0
for mol in top.mols:
if mol.is_solvent():
n_solvents += 1
assert n_solvents == 10, 'must be 10 solvents'
fi, top = pt.closest(traj)
pt.write_traj(outputname('test.pdb'),
next(fi),
top=top,
overwrite=True)
def test_makestructure(self):
# https://github.com/Amber-MD/cpptraj/issues/27
# load only 1st frame
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
# pply polyproline II dihedral to residues 1-13
t0 = traj[:1].copy()
pt.make_structure(t0, 'pp2:1-13')
pt.write_traj('./output/test0.pdb',
t0,
options='model',
overwrite=True)
t1 = traj[:1].copy()
pt.make_structure(t1, "chi1:8:N:CA:CB:CG:35")
pt.write_traj('./output/test1.pdb',
t1,
options='model',
overwrite=True)
def download_pdb(pdb_list):
'''
function takes list of pdb code, you can download from pdb website
pdb_list is like 1HHP, 3TKG, 3TKW, 3TL9
'''
data=[]
with open(pdb_list) as prlist:
pr = csv.reader(prlist, delimiter=' ')
for row in pr:
data.append(row)
num_of_pdb = len(data[0])
# make a directory PDB
os.system("mkdir PDB")
#get each pdb
for i in np.arange(0, num_of_pdb, 1):
t = data[0][i].strip(',') #get the fourlettercode
lower = t.lower() #change to lower case
pdb = pt.io.loadpdb_rcsb("%s" %lower)
#write pdb
pt.write_traj("PDB/%s.pdb" %lower, pdb, overwrite=True)
def do_clustering():
# set random seed for reproducibility
np.random.seed(1)
traj = pt.iterload('GAAC3.5000frames.nc', 'GAAC.topo')
n_frames = 50
# as an example, we extract only first 50 frames and do clustering for @P atoms
t0 = traj[:n_frames, '@P']
# need to convert to 2D array, shape=(n_frames, n_features)
X = pt.tools.as_2darray(t0)
# specify the total number of clusters
n_clusters = 10
# do it
data = k_means(X, n_clusters,
init='k-means++',
precompute_distances='auto',
n_init=10,
max_iter=300,
verbose=False,
tol=0.0001,
random_state=None,
copy_x=False,
n_jobs=-1,
return_n_iter=False)
# convert 2D to 3D array
x0 = pt.tools.as_3darray(data[0])
print('cluster index: ', data[1].tolist())
# write output, dump all frames to a single pdb file, each snapshot is seperated by
# 'MODEL' keyword. good for visualizing in VMD.
# first, need to create a frame iterator for ``pytraj.write_traj``
frame_iter = pt.iterframe_from_array(x0, t0.n_atoms, range(n_clusters))
pt.write_traj('output.pdb', frame_iter,
top=t0.top,
options='model',
overwrite=True)
def filter_seperated(trajectory, topology, targetdir, separated_frame, masks,
outprefix):
cwd = os.getcwd()
trajectory = lib.paths_for_there(trajectory, targetdir)
topology = lib.paths_for_there(topology, targetdir)
separated_frame = lib.paths_for_there(separated_frame, targetdir)
os.system('mkdir -p ' + targetdir)
os.chdir(targetdir)
cpptrajtext = ('parm ' + topology + ' \n'
'trajin ' + trajectory + ' \n'
'reference ' + separated_frame + ' [ref] \n'
'nativecontacts ' + masks + ' distance 3.3 out '
'contacts series ref [ref] \n'
'go \n')
with open('cpptraj.in', 'w') as f:
f.write(cpptrajtext)
os.system('mpirun -np 4 cpptraj.MPI -i cpptraj.in')
top = pt.load_topology(topology)
traj = pt.iterload(trajectory, top)
d = np.loadtxt('contacts', usecols=(0, 2), dtype=int)
nocontact = []
contact = []
for i, k in d:
if k == 0:
nocontact.append(i - 1)
else:
contact.append(i - 1)
pt.write_traj('nocont_' + outprefix + '.nc',
traj[nocontact],
top=top,
overwrite=True)
conttraj = 'cont_' + outprefix + '.nc'
pt.write_traj(conttraj, traj[contact], top=top, overwrite=True)
os.chdir(cwd)
return lib.paths_for_here(targetdir, conttraj)
def test_load_and_save_1(self):
# do not support frame_indices for TrajectoryIterator
self.assertRaises(
ValueError, lambda: pt.iterload(filename="./data/Tc5b.x",
top="./data/Tc5b.top",
frame_indices=[0, 5]))
traj = pt.iterload(filename="./data/Tc5b.x", top="./data/Tc5b.top")
indices = list(range(2, 4)) + [3, 7, 9, 8]
pt.write_traj(filename="./output/test_io_saved.pdb",
traj=traj,
top="./data/Tc5b.top",
frame_indices=indices,
overwrite=True)
# check frames
traj = pt.iterload(filename="./output/test_io_saved.pdb",
top="./data/Tc5b.top")
assert traj.n_frames == len(indices)
assert traj.top.n_atoms == 304
def _concat_trajectory(self, indices: List[int]):
# get filenames
traj_files = [
os.path.join(self.path, f"node_{index}",
SETTINGS["md"]["trj_file"]) for index in indices
if index != 0
]
top_file = os.path.join(self.path, SETTINGS["data_files"]["top_file"])
ref_file = os.path.join(self.path, SETTINGS["data_files"]["ref_file"])
ref_top_file = os.path.join(self.path,
SETTINGS["data_files"]["ref_top_file"])
full_traj_file = os.path.join(self.report_dir,
SETTINGS["results"]["trj_file"])
# load, align write
load_mask = SETTINGS["system"]["load_mask"]
traj = pytraj.iterload(traj_files, top=top_file, mask=load_mask)
ref = pytraj.load(ref_file, top=ref_top_file, mask=load_mask)
mask = SETTINGS["system"]["protein_mask"]
pytraj.rmsd(traj, mask=mask, ref=ref)
pytraj.write_traj(full_traj_file, traj)
def test_4(self):
"""test write Trajectory"""
farray = pt.load("data/Tc5b.x",
"./data/Tc5b.top",
frame_indices=list(range(10)))
pt.write_traj("./output/test_write_output.x",
farray,
top=farray.top,
overwrite=True)
pt.write_traj("./output/test_pdb_1.dummyext",
farray[0],
top=farray.top,
overwrite=True)
# test 'save'
farray.save("./output/test_write_output_save_method.x", overwrite=True)
# reproduce result?
f0 = mdio.iterload("./output/test_write_output.x", "./data/Tc5b.top")
f1 = mdio.iterload("./output/test_write_output_save_method.x",
"./data/Tc5b.top")
aa_eq(f0[:, :, :].xyz, f1[:, :, :].xyz)
def minimize(traj, engine='sander', input=None, top=None):
"""
>>> from pytraj.amber_wrap import minimize
>>> minimize(traj)
>>> minimize(traj, engine='pmemd')
"""
from pytraj import Trajectory
if not isinstance(traj, Trajectory):
raise ValueError("support only mutable Trajectory")
_top = get_topology(traj, top)
if input is not None:
min_in = input
else:
min_in = MIN_IN
if engine in ['sander', 'pmemd']:
_engine = "$AMBERHOME/bin/" + engine
else:
_engine = engine
with tempfolder():
with open("min.in", 'w') as fh:
fh.write(min_in)
pt.write_parm("tmp.prmtop", _top)
for frame in traj:
pt.write_traj("tmp_frame.rst7", frame, top=_top, overwrite=True)
os.system(
"%s -O -p tmp.prmtop -c tmp_frame.rst7.1 -r min.r -i min.in" %
(_engine))
f0 = pt.load("min.r", traj.top)[0]
# update coords
frame.xyz[:] = f0.xyz
def test_raise(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
# list
self.assertRaises(
NotImplementedError,
lambda: pt.write_traj('output/test.pdb', [traj[0], traj[1]], top=traj.top, frame_indices=range(3), overwrite=True))
# single Frame
self.assertRaises(
ValueError,
lambda: pt.write_traj('output/test.pdb', traj[0], top=traj.top, frame_indices=range(3), overwrite=True))
# single Frame, no Topology
self.assertRaises(
ValueError,
lambda: pt.write_traj('output/test.pdb', traj[0], top=None, overwrite=True))
# traj is None
self.assertRaises(
ValueError,
lambda: pt.write_traj('output/test.pdb', None, top=traj.top, overwrite=True))
def test_closest(self):
# raise if not has solvent
traj0 = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
self.assertRaises(RuntimeError, lambda: pt.closest(traj0))
traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
fi, top = pt.closest(traj)
coords = []
for frame in fi:
coords.append(frame.xyz.copy())
assert isinstance(frame, pt.Frame), 'must be Frame'
# make a Trajectory
fi, top = pt.closest(traj)
xyz = pt.get_coordinates(fi)
t0 = pt.Trajectory(xyz=xyz, top=top)
aa_eq(np.array(coords), t0.xyz)
# test write to disk
fi, top = pt.closest(traj)
pt.write_traj('output/fi.nc', fi, top=top, overwrite=True)
# load back
t1 = pt.load('output/fi.nc', top=top)
aa_eq(t0.xyz, t1.xyz)
# make sure n_sovent=10 (default)
n_solvents = 0
for mol in top.mols:
if mol.is_solvent():
n_solvents += 1
assert n_solvents == 10, 'must be 10 solvents'
fi, top = pt.closest(traj)
pt.write_traj('output/test.pdb', next(fi), top=top, overwrite=True)
# for example you're only interested in a single Frame
# let's load it to memory
frame = traj[1]
print(frame)
# now you want to rotate omega dihedral angle for residue 3?
# make a list to store new frames
framelist = []
for deg in range(-180, 180, 5):
# since Frame does not hold Topology information, you need to
# pass it
pt.rotate_dihedral(frame, '3:omega:' + str(deg), top=traj.top)
# always make a copy since
# the `rotate_dihedral` inplace-change coords of frame
framelist.append(frame.copy())
# write multiple pdbs to a single file
# just like multiple pdbs when you load from rcsb
# you can use this pdb file to view in VMD without loading
# prmtop/gro/psf/...
pt.write_traj("./output/test0.pdb", framelist,
top=traj.top,
options='model',
overwrite=True)
# you can use VMD to open the new file
# vmd ./output/test0.pdb
import pytraj as pt
import os
# use iterload for memory saving
traj = pt.iterload("../tests/data/Tc5b.x", "../tests/data/Tc5b.top")
pt.write_traj("./output/trj.dcd", traj, overwrite=True)
os.system("ls ./output/trj.dcd")
from itertools import chain
import pytraj as pt
import os
traj = pt.iterload("../tests/data/Tc5b.x", "../tests/data/Tc5b.top")
# write new traj with frames from 0 to 2, and 5 to 8
# skip last frame to follow python's convention.
# you can use traj[0:3], traj[5:9] but this does not save memory
# traj(0,3) like `range` in python3 or `xrange` in python2
pt.write_traj('traj_append.nc',
traj=traj,
frame_indices=chain(range(0, 3), range(5, 9)),
top=traj.top,
overwrite=True, )
# you can load many files to a single traj and write specific frames too
# traj = pt.iterload(['remd.x.000', 'remd.x.001', 'remd.x.009'], your_topology_file)
import pytraj as pt
# turn on this to see cpptraj's stdout
# pt.set_cpptraj_verbose(True)
# use iterload for memory saving
traj = pt.iterload("../tests/data/Tc5b.x", "../tests/data/Tc5b.top")
# write multiple pdbs to a single file
# just like multiple pdbs when you load from rcsb
# you can use this pdb file to view in VMD without loading
# prmtop/gro/psf/...
pt.write_traj("./output/test.pdb", traj, options='model', overwrite=True)
def test_write():
pt.write_traj('test.nc',
traj,
frame_indices=range(10000),
overwrite=True)
print(pt.iterload('./test.nc', traj.top))
en = pt.energy_decomposition(t0,
igb=8,
verbose=False,
parm=traj.top.filename)['dihedral'][0]
print(deg, en)
deg_ene.append((deg, en))
arr = np.array(deg_ene).T
print(len(flist))
# write multiple pdb to a single file to visualize
# rotation with vmd.
pt.write_traj("output/test.pdb", flist,
top=traj.top,
overwrite=True,
options='model')
# plot the PMF
try:
import matplotlib
from pytraj import plot
arr[1] = arr[1] - np.min(arr[1])
pt.plot(0, 1, data=arr)
pt.savefig("output/test.png")
except ImportError:
pass
except ImportError:
pass
import pytraj as pt
traj = pt.iterload('../tests//data/tz2.ortho.nc',
'../tests/data/tz2.ortho.parm7')
pt.write_traj("output/test.nc", traj(autoimage=True,
rmsfit=(0, '@CA,N,O')),
overwrite=True)
import pytraj as pt
traj = pt.load("../tests/data/Tc5b.x", "../tests/data/Tc5b.top")
# etract first 10 frames and write to CHARMM format
subtraj = traj[:10]
pt.write_traj(filename="./output/subtraj_0_CHARMM.dcd",
traj=subtraj,
top=traj.top,
overwrite=True)
# make sure we can load the traj,
# use AMBER top is fine
charmm_traj = pt.load("./output/subtraj_0_CHARMM.dcd",
"../tests/data/Tc5b.top")
# calculate rmsd between old and saved traj for 1st frame
print(charmm_traj[0].rmsd(subtraj[0]))
assert charmm_traj[0].rmsd(subtraj[0]) < 1E-6
# another way
subtraj.save("./output/subtraj_1_CHARMM.dcd", overwrite=True)
charmm_traj_1 = pt.load("./output/subtraj_1_CHARMM.dcd",
"../tests/data/Tc5b.top")
print(charmm_traj_1[0].rmsd(subtraj[0]))
assert charmm_traj_1[0].rmsd(subtraj[0]) < 1E-6
def test_write_xyz(self):
xyz = self.traj.xyz
fname = './output/test_xyz.nc'
pt.write_traj(fname, xyz, top=self.traj.top, overwrite=True)
t0 = pt.iterload(fname, top=self.traj.top)
aa_eq(self.traj.xyz, t0.xyz)
#!/usr/bin/env python
import pytraj as pt
traj = pt.iterload('./RAN.rst7', './RAN.parm7')
t0 = traj[:]
flist = []
for deg in range(-180, 175, 5):
pt._rotate_dih(t0, resid='1', dihtype='chin', deg=deg)
flist.append(t0[0].copy())
print(pt.calc_chin(flist, top=t0.top))
pt.write_traj('combined_traj.nc', flist, top=t0.top)
# Aim: read Gaussian output and write to pdb file as a Trajectory
# to view in VMD
# require: cclib
# install: "pip install cclib"
import pytraj as pt
traj = pt.tools.read_gaussian_output(filename="../tests/data/gaussian/GF2.log",
top="../tests/data/gaussian/GF2.pdb")
# use mode='model' to write multiple pdbs to a single file
pt.write_traj("output/traj_vmd.pdb", traj, mode='model')
# to view by VMD, just
# vmd ./output/traj_vmd.pdb
# you will see a movie with 56 frames.
# datafiles and this test.py were provided by jmborr
# issue #991: # Error "ValueError: Buffer not C contiguous"
# https://github.com/Amber-MD/pytraj/issues/991
import os
import argparse
import numpy
import pytraj
from pdb import set_trace as tr
traj = pytraj.load('short.dcd', 'pdb', mask='(!:1-256)&(@H1,@H2,@H3,@H4,@H5)')
# Trajectory of the center of mass of the first two residues
minitraj = numpy.empty((2,traj.n_frames,3))
minitraj[0] = pytraj.center_of_mass(traj,mask=':1')
minitraj[1] = pytraj.center_of_mass(traj,mask=':2')
minitraj = minitraj.transpose((1,0,2))
# "topology" using the first two atoms
minitop = traj.top['@1,2']
# Save trajectory
pytraj.write_traj(filename='/tmp/junk.crd', traj=minitraj, top=minitop, overwrite=True)
# Aim: read Gaussian output and write to pdb file as a Trajectory
# to view in VMD
# require: cclib
# install: "pip install cclib"
import pytraj as pt
traj = pt.tools.read_gaussian_output(filename="../tests/data/gaussian/GF2.log",
top="../tests/data/gaussian/GF2.pdb")
# use options='model' to write multiple pdbs to a single file
pt.write_traj("output/traj_vmd.pdb", traj, options='model', overwrite=True)
# to view by VMD, just
# vmd ./output/traj_vmd.pdb
# you will see a movie with 56 frames.
nres = masktop.n_residues
CoM = numpy.empty((nres,traj.n_frames,3))
ires = 0
for residue in masktop.residues:
try:
resnum = residue.resnum
except:
resnum = int(residue.name)
presidue = ptop.residues[resnum-1]
first_atom = presidue[0] # selected atom to represent the whole residue in the output PDB file
CMselection.append('@{0}'.format(first_atom.number))
#Next, calculate CM for given residue in all frames
v = pytraj.center_of_mass(traj,mask='({0})&(:{1})'.format(p_args.mask,resnum))
CoM[ires] = v
ires += 1
CoM = CoM.transpose((1,0,2)) #shape=(n_frames, n_residues, 3)
CMselection = ','.join(CMselection) # mask containing the first atoms of each considered residue
else:
CoM = pytraj.center_of_mass(traj)
CoM = CoM.reshape((traj.n_frames, 1, 3))
first_residue = ptop.residues[0]
first_atom = first_residue[0]
CMselection = '@{0}'.format(first_atom.number)
outtop = traj.top[CMselection]
for i in range(outtop.n_atoms):
atom = outtop[i]
atom.set_mol(0)
stripped_traj = traj[CMselection]
pytraj.write_traj(filename=p_args.outtop, traj=stripped_traj, frame_indices=[0,], top=outtop, overwrite=True)
pytraj.write_traj(p_args.outtraj, CoM, top=outtop, overwrite=True)
