def test_ParmFile():
top = pt.load_topology(tc5b_top)
with tempfolder():
pt.write_parm("test_io.top", top, overwrite=True)
newtop = pt.load_topology("test_io.top")
assert top.n_atoms == newtop.n_atoms
# test raise if file exists
with pytest.raises(RuntimeError):
pt.write_parm("test_io.top", top, overwrite=False)
def test_ParmFile(self):
top = pt.read_parm("./data/Tc5b.top")
pt.write_parm("./output/test_io.top", top, overwrite=True)
newtop = pt.read_parm("./output/test_io.top")
assert top.n_atoms == newtop.n_atoms
# test raise if file exists
self.assertRaises(
RuntimeError, lambda: pt.write_parm(
"./output/test_io.top", top, overwrite=False))
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 strip_single_traj(self, ligcom, run_path, lamda):
'''Strip waters and ions, use amber mask to select which atoms to keep'''
assert ligcom in ['ligands', 'complex']
parm = str(list(Path(run_path, lamda).glob('*.parm7'))[0])
trajin = str(list(Path(run_path, lamda).glob('*.nc'))[0])
traj_name = trajin.split('/')[-1]
parm_name = parm.split('/')[-1]
traj = pt.load(trajin, parm)
n_frames = traj.n_frames
# strip
if ligcom == 'ligands':
# ligand only
mask = f'{self.ligand_mask}'
paths = self.ligand_paths
elif ligcom == 'complex':
# ligand and protein
# residues from number of ca
n_residues = len(traj.top.select('@CA'))
mask = f':1-{n_residues}|{self.ligand_mask}'
paths = self.complex_paths
# get ion index
if self.ion_decharge:
atom_idx = self.get_decharged_ion(run_path)
mask = mask + f'|@{atom_idx}'
# select region to keep
traj = traj[mask]
logging.info(f'{run_path}, lambda: {lamda}, mask: {mask}')
# last half
if self.last_half_frames:
traj = traj[n_frames // 2:]
# subsample
if self.stride != 1:
traj = traj[::self.stride]
# save nc and restart
path_idx = paths.index(run_path)
if ligcom == 'ligands':
fold = self.lig_folds[path_idx]
out_traj = str(
Path(self.dest_dir, 'ligands', fold, lamda, traj_name))
out_parm = str(
Path(self.dest_dir, 'ligands', fold, lamda, parm_name))
elif ligcom == 'complex':
fold = self.com_folds[path_idx]
out_traj = str(
Path(self.dest_dir, 'complex', fold, lamda, traj_name))
out_parm = str(
Path(self.dest_dir, 'complex', fold, lamda, parm_name))
pt.write_traj(out_traj, traj)
pt.write_parm(out_parm, traj.top)
import pytraj as pt
top = pt.load_topology("../tests/data/Tc5b.top")
# save only CA atoms
pt.write_parm("./output/_Tcb5.onlyCA.top", top["@CA"])
def main():
print('\nWelcome to OCD.py!\n')
### Load arguments from parser
args = argumentParser()
### If no topology was provided, assume that the input
### can be used as its own topology (eg pdb files)
if args.top == '':
print('No topology was provided.'
'The input file will be used as its own topology.')
args.top = args.input
### Check if files exist
for file in [args.input, args.top]:
if not os.path.exists(file):
print('File {!s} doesn\'t exist.'
'Please check the filepath or add the file.\n'.format(file))
calc.sysexit(1)
### Format input arguments and generate some variables based on the input
# args.use needs to be a tuple for use in pytraj.
# If no argument was given, use every frame of the whole trajectory.
# Throws an error if input is not three arguments
if args.use is None:
args.use = (0, -1, 1)
elif len(args.use) == 3:
args.use = (args.use[0], args.use[1], args.use[2])
else:
print('Wrong number of arguments in -use.'
'Usage: -use <FirstFrame> <LastFrame> <Stride>')
calc.sysexit(1)
trajname = os.path.basename(args.input)
### Load trajectory
print('Loading data.')
traj_all = pt.iterload(args.input, args.top, frame_slice=args.use)
#Align trajectory on first frame before running the calculations
if args.align or args.vmd:
traj_all = traj_all.superpose()
### Apply atom mask to trajectory, split the domains
stripped_domains = [
traj_all.strip('!(' + args.A + ')'),
traj_all.strip('!(' + args.B + ')')
]
### Get reference coordinates from reference structure or first frame,
# if no reference structure was provided-
# Assuming the reference COM to be [0,0,0] is only valid if
# the standard_orientation function is used,
# so we need to reassign when args.abangleref == True
ref_com_A = np.array([0, 0, 0])
ref_com_B = np.array([0, 0, 0])
if args.refstruc != None:
ref_traj = pt.iterload(args.refstruc)
stripped_refs = [
ref_traj.strip('!(' + args.A + ')'),
ref_traj.strip('!(' + args.B + ')')
]
### If residues are missing, exclude them from the
# alignment with new_mask
# - necessary because alignment would break if
# the number of atoms supplied is different
new_mask_A, new_mask_B = calc.new_masks(stripped_domains,
stripped_refs)
new_masks = [new_mask_A, new_mask_B]
## Strip domains down for alignment, if
# reference structure has fewer residues.
# Same thing happens for the residue structure
# in the function standard_orientation.
for dix, domain in enumerate(stripped_domains):
try:
if new_masks[dix] != None:
stripped_domains[dix] = domain[:].strip('!(' +
new_masks[dix] +
')')
except ValueError as e:
print('WARNING: Trajectory domain {} could not be '
'stripped further. This usually happens when '
'the trajectory contains fewer residues than '
'the reference structure!'.format('A' if dix ==
0 else 'B'))
ref_coords_A, ref_coords_B = calc.standard_orientation(
stripped_refs, new_mask_A, new_mask_B, args.output)
else:
ref_coords_A, ref_coords_B = calc.standard_orientation(
stripped_domains, None, None, args.output)
### Extract data for the two domains
# TrajectoryIterator.strip as used here returns the trajectory stripped
# down to the masks given in args.A and args.B.
# This is somehow faster then applying masks to a normal trajectory object.
print('Preparing structural data from input.')
traj_A = stripped_domains[0]
traj_B = stripped_domains[1]
coords_A = traj_A.xyz
coords_B = traj_B.xyz
results = []
vectors_all, distances_all = [], [] #These are for the vmd visualization
### Calculate angles for each frame
print('Calculating angles.')
# Enumerate over all frames: Not very efficient,
# but the time limiting step is anyway the initiation of the trajectory.
for i, frame in enumerate(coords_A):
B_points, A_points = calc.apply_coordinatesystem(
ref_coords_A, ref_coords_B, coords_A[i], coords_B[i], ref_com_A,
ref_com_B)
# Calculate Angles, add time and RMSD to output
angles = list(calc.angle_calculation(B_points, A_points))
rmsd = list(
calc.orientational_rmsd(ref_coords_A, ref_coords_B, coords_A[i],
coords_B[i]))
simulated_time = [float(args.simtime) / len(coords_A) * i]
results.append(angles + rmsd + simulated_time)
### Calculate some data needed for the vmd visualization from each frame
if args.vmd or args.pymol:
vectors_frame = []
distance_frame = []
for domain in (B_points, A_points):
start_point = domain[0]
distance_frame.append(list(start_point))
for c in range(1, 4):
try:
# start point / end point / vec number
vector_data = [list(start_point), list(domain[c]), c]
vectors_frame.append(vector_data)
except:
pass
#append 0 as identifier of the distance axis
distance_frame.append(0)
vectors_all.append(vectors_frame)
distances_all.append([distance_frame])
### Output data
print('Outputing and visualizing data.')
cols = ['AB', 'AC1', 'BC1', 'AC2', 'BC2', 'dc', 'RMSD_A', 'RMSD_B', 'Time']
df = pd.DataFrame(results, columns=cols)
df = df.apply(pd.to_numeric, errors='ignore')
with open('OCD_{}.dat'.format(args.output), 'w+') as f:
df.to_csv(f, sep='\t', index=False, float_format='%.3f')
# Write a TCL script for the visualization in VMD -
# needs an aligned trajectory to work, so it outputs one as well
if args.vmd:
vis.vmd_script('./OCD_{}_vmd_ocd.nc'.format(trajname),
'./OCD_{}_vmd_ocd.parm7'.format(trajname), (0, -1, 1),
vectors_all, distances_all, args.output)
if not os.path.isfile('./OCD_{}_vmd.nc'.format(trajname)):
print('Writing trajectory for VMD visualization...')
pt.write_traj('./OCD_{}_vmd_ocd.nc'.format(trajname),
traj_all,
overwrite=True)
pt.write_parm('./OCD_{}_vmd_ocd.parm7'.format(trajname),
top=traj_all.top,
overwrite=True)
# Write out first frame as pdb for PyMol visualization
if args.pdb or args.pymol:
frame1 = pt.iterframe(traj_all, frame_indices=[0])
frame1.save('./OCD_{}.pdb'.format(args.output), overwrite=True)
# Write PyMOl input script for vizaulisation of
# the first frame and coordinate system
if args.pymol:
str_o = ''
str_o += vis.pymol_init()
str_o += vis.pymol_load('./OCD_{}.pdb'.format(args.output),
'Frame1')
str_o += vis.pymol_draw_vectors(vectors_all[0], distances_all[0])
str_o += vis.pymol_settings()
with open('./OCD_{}.pym'.format(args.output), 'w+') as f:
f.write(str_o)
if args.plot:
labeldict = {
'AB': 'AB Angle /$^\circ$',
'AC1': 'AC1 Angle /$^\circ$',
'BC1': 'BC1 Angle /$^\circ$',
'AC2': 'AC2 Angle /$^\circ$',
'BC2': 'BC2 Angle /$^\circ$',
'dc': 'dc Distance /$\AA$'
}
titledict = {
'AB': 'AB Torsion Angle',
'AC1': 'AC1 Tilt Angle',
'BC1': 'BC1 Tilt Angle',
'AC2': 'AC2 Tilt Angle',
'BC2': 'BC2 Tilt Angle',
'dc': 'dc Distance'
}
lim_dict = {
'AB': args.lim_AB,
'AC1': args.lim_AC1,
'BC1': args.lim_BC1,
'AC2': args.lim_AC2,
'BC2': args.lim_BC2,
'dc': args.lim_DC
}
for s in ['AB', 'AC1', 'BC1', 'AC2', 'BC2', 'dc']:
if s == 'dc':
binw = 0.1
else:
binw = 0.5
vis.angle_plots(x_data=df[s],
y_time=df['Time'],
xlim=lim_dict[s],
xlabel=labeldict[s],
title=titledict[s],
bin_dims=(np.floor(df[s].min()),
np.ceil(df[s].max()), binw),
norm=True,
hist2c='#08519c',
linec='darkblue',
plot_type=['hist', 'time'])
if not os.path.exists('./OCD_{}_Plots'.format(args.output)):
os.makedirs('./OCD_{}_Plots'.format(args.output))
plt.savefig('./OCD_{}_Plots/{}.png'.format(args.output, s),
bbox_inches='tight')
calc.sysexit(0)
