def testLoadTOP(self):
'''
Data files from:
http://ambermd.org/tutorials/basic/tutorial2/section6.htm
'''
cmd.load(self.datafile("TRPcage.top"))
cmd.load_traj(self.datafile("heat1.crd"), "TRPcage")
self.assertEqual(304, cmd.count_atoms())
self.assertEqual(10, cmd.count_states())
def testLoadTOP(self):
'''
Data files from:
http://ambermd.org/tutorials/basic/tutorial2/section6.htm
'''
cmd.load(self.datafile("TRPcage.top"))
cmd.load_traj(self.datafile("heat1.crd"), "TRPcage")
self.assertEqual(304, cmd.count_atoms())
self.assertEqual(10, cmd.count_states())
def load_file(opts):
cmd.load(opts.topology)
cmd.load_traj(
opts.trajectory,
start=opts.frames[0] if opts.frames else 1,
stop=opts.frames[1] if opts.frames else -1,
interval=opts.frames[2] if opts.frames else 1,
selection="{} or {}".format(opts.ligand_sel, opts.protein_sel),
)
def testSpheresForSpheroidTrajectory(self, use_shader):
pdbfile = self.datafile("sampletrajectory.pdb")
dcdfile = self.datafile("sampletrajectory.dcd")
cmd.load(pdbfile)
cmd.load_traj(dcdfile)
cmd.set('use_shaders', use_shader)
cmd.spheroid("sampletrajectory", 3)
self.ambientOnly()
cmd.color('blue')
cmd.show_as('spheres')
self.assertImageHasColor('blue')
def testSpheresForSpheroidTrajectory(self, use_shader):
pdbfile = self.datafile("sampletrajectory.pdb")
dcdfile = self.datafile("sampletrajectory.dcd")
cmd.load(pdbfile)
cmd.load_traj(dcdfile)
cmd.set('use_shaders', use_shader)
cmd.spheroid("sampletrajectory", 3)
self.ambientOnly()
cmd.color('blue')
cmd.show_as('spheres')
self.assertImageHasColor('blue')
def _view(mol, viewname):
topo = tempname(suffix=".cif")
xtc = tempname(suffix=".xtc")
mol.write(topo)
mol.write(xtc)
cmd.delete(viewname)
cmd.load(topo, viewname)
cmd.load_traj(xtc, viewname, state=1)
cmd.dss() # Guess SS
os.remove(topo)
os.remove(xtc)
showing[viewname] = True
def test_save_traj(tmp_path):
cmd.reinitialize()
cmd.fragment('gly', 'm1')
for state in range(1, 4):
cmd.create('m2', 'm1', 1, state)
for ext in ['dcd', 'crd']:
filename = str(tmp_path / "out.") + ext
# export
psico.exporting.save_traj(filename, 'm2')
# import (verify export)
cmd.create('m3', 'm1')
cmd.load_traj(filename, 'm3', 1)
assert cmd.count_states('m3', 3)
cmd.delete('m3')
def testLoadTraj(self, topext, trjext):
pdbfile = self.datafile("sampletrajectory" + topext)
dcdfile = self.datafile("sampletrajectory" + trjext)
cmd.load(pdbfile)
cmd.load_traj(dcdfile)
self.assertEqual(11, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, interval=2, max=3)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, start=3, stop=5)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, stop=9, average=3)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
def testLoadTraj(self):
pdbfile = self.datafile("sampletrajectory.pdb")
dcdfile = self.datafile("sampletrajectory.dcd")
cmd.load(pdbfile)
cmd.load_traj(dcdfile)
self.assertEqual(501, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, interval=5, max=20)
self.assertEqual(20, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, start=31, stop=40)
self.assertEqual(10, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, stop=30, average=3)
self.assertEqual(10, cmd.count_states())
cmd.delete('*')
def testLoadTraj(self, topext, trjext):
pdbfile = self.datafile("sampletrajectory" + topext)
dcdfile = self.datafile("sampletrajectory" + trjext)
cmd.load(pdbfile)
cmd.load_traj(dcdfile)
self.assertEqual(11, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, interval=2, max=3)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, start=3, stop=5)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
cmd.load(pdbfile)
cmd.load_traj(dcdfile, state=1, stop=9, average=3)
self.assertEqual(3, cmd.count_states())
cmd.delete('*')
cmd.set('async_builds', 1)
cmd.set('ray_transparency_contrast', 3.0)
cmd.set('ray_transparency_shadows', 0)
model = cmd.get_model('system')
#for atom in model.atom:
# print "%8d %4s %3s %5d %8.3f %8.3f %8.3f" % (atom.index, atom.name, atom.resn, int(atom.resi), atom.coord[0], atom.coord[1], atom.coord[2])
#pymol.finish_launching()
cmd.viewport(width,height)
#niterations = 10 # DEBUG
# Load trajectory
cmd.load_traj(trajectory_temporary_filename, object='system')
# Align all states
if solute == 'dna':
cmd.intra_fit('name P') # DNA
else:
cmd.intra_fit('solute') # protein
# Zoom viewport
cmd.zoom('solute')
cmd.orient('solute')
# Create one-to-one mapping between states and frames.
cmd.mset("1 -%d" % nframes)
else:
selection = f'inp and ({args.select})'
if args.frames is not None:
stop = max(args.frames)
chunks = [c[c <= stop] for c in chunks if min(c) <= stop]
else:
stop = -1
args.frames = range(1, nframes + 1)
for chunkid, chunk in enumerate(chunks):
start, stop = min(chunk), max(chunk)
cmd.reinitialize()
cmd.load(args.top, 'inp')
cmd.load_traj(args.traj,
'inp',
state=1,
start=start,
stop=stop,
selection=selection)
if args.align is not None:
if len(chunks) == 1:
rmsds = cmd.intra_fit(selection, args.align)
rmsds[rmsds == -1.] = 0.
outrmsdfile = f"{os.path.splitext(args.out)[0]}_rmsd.txt"
numpy.savetxt(outrmsdfile, rmsds, fmt="%.4f")
states = numpy.where(numpy.isin(chunk, args.frames))[0] + 1
nstates = cmd.count_states('inp')
if len(states) == nstates:
cmd.create('out', selection=selection, source_state=0, target_state=-1)
else:
for s in states:
sys.stdout.write(f'Getting state {s}/{max(states)}\r')
def load_traj(*args, **kwargs):
cmd.set('defer_builds_mode', 3)
cmd.sync()
cmd.load_traj(*args, **kwargs)
def testLoadTraj_selection(self):
base = self.datafile("sampletrajectory")
cmd.load(base + ".gro")
cmd.load_traj(base + ".xtc", selection="backbone")
self.assertEqual(55, cmd.count_atoms('state 10'))
def testLoadTraj_selection_trj(self):
base = self.datafile("sampletrajectory")
cmd.load(base + ".pdb")
cmd.load_traj(base + ".crd", selection="backbone", format="trj")
self.assertEqual(55, cmd.count_atoms('state 10'))
def extract_psa3d(self,
traj_file,
gro_file,
obj_list=None,
include_SandP=None,
cmpd_name=None):
"""
Calculates the 3d polar surface area (3D-PSA) of molecules in Interface_Pymol for all the snapshots in a MD trajectory.
(Contribution by Benjamin Schroeder)
Parameters:
-------------
traj_file: str
trajectory filename
gro_file: str
coordinates filename (.gro or .pdb)
cmpd_name: str, optional
Name of the compound used as object name (default is "cmpd1")
include_SandP: bool, optional
Set to False to exclude the S and P atoms from the calculation of the 3D-PSA. (Default = True)
atom_to_remove: str, optional
Single atom name of the atom to remove from the selection (Default = None).
Useful if you want to include only S or only P in the calculation of the 3D-PSA.
Returns
----------
dict_psa3d: dict
Keys are mean (3d_psa_av), standard deviation (3d_psa_sd), and median (3d_psa_med) of the 3D-PSA calculated over the simulation time.
If cmpd_name is specified, it is returned in the dictionary.
"""
if cmpd_name == None:
cmpd_name = "cmpd1"
# Load trajectory and remove solvent and salts
obj1 = cmpd_name
cmd.reinitialize()
cmd.load(gro_file, object=obj1)
cmd.load_traj(traj_file, object=obj1)
cmd.remove("solvent")
cmd.remove("resn Cl-")
cmd.remove("resn Na+")
atom_names = []
cmd.iterate_state(-1,
selection=obj1 + " and not elem H",
expression="atom_names.append(name)",
space=locals())
total_psa = self.calc_psa3d(include_SandP=include_SandP)
dict_psa3d = {
'3d_psa_av': total_psa[0],
'3d_psa_sd': total_psa[1],
'3d_psa_med': total_psa[2]
}
if cmpd_name == "cmpd1":
return (dict_psa3d)
else:
dict_psa3d.update({"cmpd_name": cmpd_name})
return (dict_psa3d)
def main():
"""
parses arguments user input and initializes stir based on user input
provides help for usage
"""
parser = HelpfulParser(
prog='stir',
formatter_class=argparse.RawDescriptionHelpFormatter,
add_help=False,
description=
'A python wrapper for pymol and several helpful tools and scripts\n'
'mainly focused on martini coarse-grained trajectories.\n\n'
'The accepted file formats are:\n'
'- structure: gro, pdb\n'
'- scene: pse\n'
'- topology: top, itp, tpr\n'
'- trajectory: xtc',
epilog='Examples:\n'
'\tstir system.gro topol.top md.xtc\n'
'\tstir system.gro --keep-water -r supercell 3,3,1\n'
'\tstir system.gro topol.tpr --pymol -qi myscript.pml',
)
files_group = parser.add_argument_group('file arguments')
files_group.add_argument(
dest='files',
action=FilesAction,
nargs='+',
help=
'a structure or scene file is required. Topology files allow for topology '
'reconstruction. Trajectory files are also accepted. Order does not matter'
)
opt_group = parser.add_argument_group('optional stir arguments')
opt_group.add_argument(
'--keep-water',
dest='keepwater',
action='store_true',
help='do not delete waters from the system. Decreases performance')
opt_group.add_argument('-g',
'--gmx',
dest='gmx',
type=str,
default=None,
help='path to the gromacs executable')
opt_group.add_argument(
'-r',
'--run-tool',
dest='runtool',
metavar='token',
type=str,
default=[],
nargs='*',
action='append',
help='a command to be run after loading. (e.g.: supercell 3,3,1). '
'Can be specified multiple times')
gar_group = parser.add_argument_group('optional garnish arguments')
gar_group.add_argument(
'--no-fix',
dest='nofix',
action='store_false', # store false!
help='disable the atom-id-based fix for the elastic network in garnish '
'(use if your system has messy, non-sequential numbering.')
gar_group.add_argument(
'--no-prot',
dest='noprot',
action='store_false', # store false!
help=
'do not guess protein backbone beads (use if normal guessing makes mistakes'
)
gar_group.add_argument(
'--no-garnish',
dest='nogarnish',
action='store_true',
help='do not run garnish on the system (use with atomistic systems)')
traj_group = parser.add_argument_group('optional trajectory arguments')
traj_group.add_argument(
'-s',
'--skip',
dest='skip',
metavar='n',
type=int,
default=1,
help=
'load trajectory frames skipping this interval. Useful to reduce memory load'
)
traj_group.add_argument(
'-b',
'--begin',
dest='begin',
metavar='frame',
type=int,
default=1,
help=
'first frame to load from trajectory. Only acts on first trajectory file'
)
traj_group.add_argument(
'-e',
'--end',
dest='end',
metavar='frame',
type=int,
default=-1,
help=
'last frame to load from trajectory. Only acts on last trajectory file'
)
traj_group.add_argument('-m',
'--max',
dest='max',
metavar='n',
type=int,
default=0,
help='maximum number of frames to load')
more_group = parser.add_argument_group('pymol arguments')
more_group.add_argument(
'-p',
'--pymol',
dest='pymol',
default=[],
nargs=argparse.REMAINDER,
help='all following arguments will be passed directly to pymol. '
'Accepts options and .pml scripts')
help_group = parser.add_argument_group('info')
help_group.add_argument('-h',
'--help',
action='help',
help='show this help message and exit')
help_group.add_argument('-V',
'--version',
action='version',
version=f'%(prog)s {__version__}')
args = parser.parse_args()
struct = args.files['struct']
scene = args.files['scene']
topol = args.files['topol']
traj = args.files['traj']
# make sure we only have 1 structure OR scene file and at most one topol
if not bool(struct) ^ bool(scene): # not xor
parser.error('you must provide either a structure or scene file')
elif len(struct) > 1 or len(scene) > 1 or len(topol) > 1:
parser.error(
'only one system can be opened at once... for now!') # TODO?
# sanitize trajectory args
if args.skip < 1:
args.skip = 1
if args.begin < 1:
args.begin = 1
if args.end < 1 and args.end != -1:
args.end = 1
if args.max < 0:
args.max = 0
pymol_args = []
scripts = []
for arg in args.pymol:
p = clean_path(arg)
if p.suffix in ('.pml', '.py'):
if not p.is_file():
raise FileNotFoundError(f'{p} does not exist')
scripts.append(str(p))
else:
pymol_args.append(str(p))
# initialize pymol
__main__.pymol_argv = ['pymol'] + pymol_args
pymol.finish_launching()
# run pymolrc and load all the stir tools
config.pymolrc()
view.load()
supercell.load()
render.load()
edit.load()
# load garnish
garnish.extend_garnish()
cmd.sync()
# open the structure
if scene:
cmd.load(scene[0])
elif struct:
cmd.load(struct[0])
cmd.sync()
# get the loaded object's name, so we can load the traj into it as new states
sys_obj = cmd.get_object_list()[0]
# load trajectories, leaving out waters if not asked for
if traj:
skip = args.skip
max_states = args.max
selection = 'all'
if not args.keepwater:
selection = 'not resname W+WN'
config.trajectory()
for i, t in enumerate(traj):
cmd.sync()
start = 1
if i == 0:
start = args.begin
stop = -1
if i == len(traj) - 1:
stop = args.end
if args.max != 0:
max_states = args.max - cmd.count_states()
if max_states < 1:
break
cmd.load_traj(t,
sys_obj,
interval=skip,
start=start,
stop=stop,
max=max_states,
selection=selection)
cmd.sync()
# also, delete waters from first frame
if not args.keepwater:
cmd.remove('resname W+WN')
cmd.sync()
if not args.nogarnish:
# sanitize topol
if not topol:
topol = [None]
garnish.garnish(file=topol[0],
gmx=args.gmx,
fix_elastics=args.nofix,
guess_prot=args.noprot,
show=False)
cmd.sync()
# load garnish data into pymol
view.nicesele()
cmd.sync()
view.set_vdw()
cmd.sync()
view.set_chains()
cmd.sync()
# run nice with the default settings, or with balls if no topol was given
view.nice()
cmd.sync()
# finally run user-requested tools
for tool in args.runtool:
command = ' '.join(tool)
cmd.do(command)
cmd.sync()
# and user-provided scripts
for scr in scripts:
cmd.run(scr)
cmd.sync()
# add command for simple help function
cmd.extend('stir', stir_help)
# print some help after everything is loaded
stir_help()
cmd.sync()
pdbfilename = caption['pdb']
else:
pdbfilename = args.pdb
if 'traj' in caption:
trajfilename = caption['traj']
else:
trajfilename = None
try:
print(colored(f' • {pdbfilename} and {sel}', color))
except KeyError:
print(f' • {pdbfilename} and {sel} -> {color}')
pdbfilenames = glob.glob(pdbfilename)
for pdbfilename in pdbfilenames:
cmd.load(pdbfilename, '_inp_')
if trajfilename is not None:
cmd.load_traj(trajfilename, '_inp_', state=1)
nstates = cmd.count_states('_inp_')
if 'project' in caption: # Projection on scattered dots
project = np.genfromtxt(caption['project'])
args.atomic = True
args.geom = True
assert len(project) == nstates
else:
project = None
if 'label' in caption: # label of the colorbar for the data projected
clb_proj_label = caption['label']
else:
clb_proj_label = None
if 'sort' in caption: # Sort the data based on the 'project' field
dosort = caption['sort'] # 1: sort | -1: reverse sort
else:
import numpy as np
if __name__ == '__main__':
import argparse
# argparse.ArgumentParser(prog=None, usage=None, description=None, epilog=None, parents=[], formatter_class=argparse.HelpFormatter, prefix_chars='-', fromfile_prefix_chars=None, argument_default=None, conflict_handler='error', add_help=True, allow_abbrev=True, exit_on_error=True)
parser = argparse.ArgumentParser(description='')
# parser.add_argument(name or flags...[, action][, nargs][, const][, default][, type][, choices][, required][, help][, metavar][, dest])
parser.add_argument('-p', '--pdb')
parser.add_argument('-t', '--traj')
parser.add_argument('-s', '--sel', default='all')
parser.add_argument('-r', '--refs', nargs='+')
parser.add_argument('-o', '--out', default='rmsds.txt')
args = parser.parse_args()
cmd.load(args.pdb, 'trajin')
cmd.load_traj(args.traj, 'trajin', state=1)
nframes = cmd.count_states('trajin')
print(f'Number of frames: {nframes}')
for ref in args.refs:
print(f'Loading {ref}')
cmd.load(ref, 'trajin')
cmd.remove('hydro and trajin')
nref = len(args.refs)
rmsds_all = []
for ind, frameref in enumerate([nframes + i for i in range(1, nref + 1)]):
print(f'Computing RMSD for {args.refs[ind]}')
rmsds = cmd.intra_fit(f'trajin and {args.sel}', state=frameref)[:nframes]
rmsds_all.append(rmsds)
rmsds_all = np.asarray(rmsds_all, dtype=object).T
min_rmsd = rmsds_all.min(axis=1)
args_min = rmsds_all.argmin(axis=1)
import pickle
import os
import sys
import subprocess
import time
with open("frames.file", "rb") as f:
frame_list = pickle.load(f)
f.close()
os.remove('frames.file')
pymol_argv = ['pymol', '-q -x ']
finish_launching(pymol_argv)
cmd.load(sys.argv[1], 'loaded_protein')
cmd.load_traj(sys.argv[2], 'loaded_protein')
cmd.util.performance(100)
cmd.set('defer_builds_mode', 10)
cmd.remove('solvent')
for frame in frame_list:
try:
CapsLock_status = subprocess.getoutput('xset q | grep Caps')[21:24]
if CapsLock_status == 'on ':
print('Frame ', frame.number)
cmd.set('state', frame.number)
for command in frame.commands_list:
cmd.color(command.color_name, selection=command.atoms_to_color)
elif CapsLock_status == 'off':
print('Wait for CapsLock...')
parser.add_argument('-r',
'--residues',
default=None,
nargs='+',
type=int,
help='Specific residues to show.')
args = parser.parse_args()
# Basic formatting
cmd.set('valence', 0)
cmd.set('stick_h_scale', 1)
# Load trajectory
cmd.load(args.topology, 'system')
if args.trajectory is not None:
cmd.load_traj(args.trajectory, 'system')
cmd.show('cartoon', 'system')
# Format ligands, if any
if args.ligands is not None:
for resname in args.ligands:
cmd.show('sticks', 'resn {}'.format(resname))
# Show any residues the user might be interested in
if args.residues is not None:
for resid in args.residues:
cmd.show('sticks', 'polymer.protein and resi {}'.format(resid))
# Hide nonpolar hydrogen atoms - this may have issues sometimes (may need to remove brackets)
cmd.hide('(h. and (e. c extend 1))')
def main(*args, **kwargs):
parser = argparse.ArgumentParser(prog='mt_pymol')
parser.add_argument(
dest='struct',
type=valid_str,
help='gro or similar file containing a martini structure')
parser.add_argument(dest='topol',
type=valid_top,
default=None,
nargs='?',
help='top or tpr file with the topology of the system')
parser.add_argument(
dest='traj',
type=valid_traj,
default=None,
nargs='*',
help='corresponding trajectory file. If multiple files are given, '
'they are concatenated')
parser.add_argument(
'-s',
'--skip',
dest='skip',
type=int,
default=1,
help='when loading a trajectory, load frames with this rate')
parser.add_argument('-g',
'--gmx',
dest='gmx',
type=str,
default=None,
help='path to the gromacs executable')
parser.add_argument(
'--keepwater',
dest='keepwater',
action='store_true',
help='do not delete waters from the system. Decreases performance')
# TODO: add more options (load_traj start/end...)
# TODO: passing arguments to pymol
args = parser.parse_args()
# check if there's enough memory to load the requested traj and warn the user if needed
if args.traj:
freemem = psutil.virtual_memory().available
traj_size = 0
for traj in args.traj:
traj_size += os.path.getsize(clean_path(traj))
water_ratio = 1
if not args.keepwater:
# TODO: VERY arbitrary number. When garnish's parsing is a module, use that!
# EDIT: I will probably leave it like this. Unnecessary and complex to use parse_tpr
water_ratio = 1 / 2
# check if there's enough free memory: 5 is based on some testing
if freemem < 5 * (traj_size / args.skip):
ok = False
inp = input(
'WARNING: You may not have enough free memory to open this big trajectory.\n'
'Consider using the trajectory options (-s, ...).\n'
'Otherwise, continue at your own risk ;) [y/N] ')
while not ok:
if inp.lower() in ['yes', 'y']:
ok = True
elif inp.lower() in ['no', 'n']:
parser.print_help()
exit(0)
else:
print(f'"{inp}" is not a valid choice. [y/N]')
# initialize pymol
__main__.pymol_argv = ['pymol']
pymol.finish_launching()
# run pymolrc and load all the mtools
config.pymolrc()
mt_nice.load()
mt_supercell.load()
mt_movie.load()
# load garnish
garnish.extend_garnish()
cmd.sync()
# open the structure
cmd.load(clean_path(args.struct))
cmd.sync()
# get the loaded object's name, so we can load the traj into it as new states
sys_obj = cmd.get_object_list()[0]
# load trajectories
if args.traj:
config.trajectory()
for traj in args.traj:
cmd.sync()
cmd.load_traj(clean_path(traj), sys_obj, interval=args.skip)
cmd.sync()
# TODO: "selection" in load_traj seems not to work as planned. Can we get it to work?
# Other option: call trjconv to get rid of the waters before loading
# delete waters, unless they are needed
if not args.keepwater:
cmd.remove('resname W or resname WN')
cmd.sync()
# run garnish with as many arguments as we got
garnish_args = []
if args.topol:
garnish_args.append(clean_path(args.topol))
if args.gmx:
garnish_args.append(f'gmx={args.gmx}')
garnish_args = ', '.join(garnish_args)
cmd.do(f'garnish {garnish_args}')
cmd.sync()
# run mt_nice with the `clean` setting
cmd.do(f'mt_nice not *_elastics')
cmd.sync()
# print some help after everything is loaded
mt_help = '''
Martini Tools functions:
- garnish
- mt_nice, mt_sele, mt_color
- mt_supercell
- mt_movie
'''
cmd.sync()
print(mt_help)