def test(self):
cmd.set('suspend_updates')
cmd.set('depth_cue', 0)
cmd.set('ambient', 1)
cmd.set('specular', 0)
cmd.set('stereo_angle', 10)
cmd.pseudoatom()
cmd.show_as('sphere')
cmd.zoom()
cmd.stereo('anaglyph')
cmd.set('suspend_updates', 0)
wh = (40, 40)
cmd.draw(*wh, antialias=0)
cmd.sync()
img = self.get_imagearray(prior=1)
self.assertEqual(wh, img.shape[:2])
# ray tracing and shaders in PyMOL 1.7.6+:
# 0xe50000
# 0x00e5e5
# -> delta=26 (0x1a)
self.assertImageHasColor('0xff0000', img, delta=0x1a)
self.assertImageHasColor('0x00ffff', img, delta=0x1a)
def test(self):
cmd.set('suspend_updates')
cmd.set('depth_cue', 0)
cmd.set('ambient', 1)
cmd.set('specular', 0)
cmd.set('stereo_angle', 10)
cmd.pseudoatom()
cmd.show_as('sphere')
cmd.zoom()
cmd.stereo('anaglyph')
cmd.set('suspend_updates', 0)
wh = (40, 40)
cmd.draw(*wh, antialias=0)
cmd.sync()
img = self.get_imagearray(prior=1)
self.assertEqual(wh, img.shape[:2])
# ray tracing and shaders in PyMOL 1.7.6+:
# 0xe50000
# 0x00e5e5
# -> delta=26 (0x1a)
self.assertImageHasColor('0xff0000', img, delta=0x1a)
self.assertImageHasColor('0x00ffff', img, delta=0x1a)
def nicecolor(method, selection='all'):
"""
DESCRIPTION
color a selection with unique colors based on atom properties exposed by iterate
USAGE
nicecolor method [, selection]
ARGUMENTS
method = property that will be used as a method for coloring (e.g.: "resi")
selection = selection to act upon
"""
stored.tmp_dict = {}
stored.r_choice = random.choice
cmd.iterate(
selection,
f'stored.tmp_dict[{method}] = stored.r_choice(stored.nice_colors)')
cmd.alter(selection, f'color = stored.tmp_dict[{method}]')
cmd.sync()
del stored.tmp_dict
cmd.recolor()
cmd.sync()
def nicesele(delete=None):
"""
DESCRIPTION
easily select relevant groups of atoms in a martini system
USAGE
nicesele [, delete]
ARGUMENTS
delete = if set, will delete the selections instead of creating them
"""
if delete:
if delete == 'delete':
for sel in stored.nice_selectors:
cmd.delete(sel)
else:
print(
'Unknown option. Type `nicesele delete` to delete all the pre-made selections'
)
else:
for sel, logic in stored.nice_selectors.items():
cmd.select(sel, logic)
cmd.sync()
# disable last selection to avoid accidental modifications to it
cmd.deselect()
cmd.sync()
def main():
parser=argparse.ArgumentParser()
parser.add_argument('-p', action='store', required=True, dest='pdbfile',
help='pathfile is a textfile of name of pdb')
parser.add_argument('-l', action='store', required=True, dest='loop',
help='loop is an int number')
inputs=parser.parse_args()
pdbfile=inputs.pdbfile
loop=int(inputs.loop)
print "loop is %s "% loop
dirname=os.getcwd()+'/interface_analyzer/pdbdata/'+pdbfile
pdbfileselect=pdbfile.split('.')
savefile=os.getcwd()+'/interface_analyzer/pdbdata/'+pdbfileselect[0]+'_antigen.pdb'
print "remove antibody of complex,so can caculate the value of antigen sasa "
pymol.finish_launching()
if loop!=0:
pymol.cmd.reinitialize()
cmd.load(dirname)
cmd.remove('hetatm')
cmd.select('target1',pdbfileselect[0])
cmd.save(dirname,(('target1')))
cmd.delete('target1')
cmd.remove('solvent')
cmd.remove('hydrogens')
cmd.remove('chain L+H')
cmd.select('target',pdbfileselect[0])
cmd.save(savefile,(('target')))
cmd.delete(all)
cmd.sync()
cmd.quit()
print "remove the antibody of complex have done"
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-p',
action='store',
required=True,
dest='pdbfile',
help='pathfile is a textfile of name of pdb')
parser.add_argument('-l',
action='store',
required=True,
dest='loop',
help='loop is an int number')
parser.add_argument('-c',
action='store',
required=True,
dest='chain',
help='the chain number is needed')
parser.add_argument('-c1',
action='store',
required=True,
dest='chain1',
help='the chain number is needed')
inputs = parser.parse_args()
pdbfile = inputs.pdbfile
removechains = 'not chain ' + inputs.chain + '+' + inputs.chain1
removechain = 'chain ' + inputs.chain
loop = int(inputs.loop)
print "loop is %s " % loop
dirname = os.getcwd() + '/interface_analyzer/pdbdata/' + pdbfile
pdbfileselect = pdbfile.split('.')
savefile = os.getcwd(
) + '/interface_analyzer/pdbdata/' + pdbfileselect[0] + '_antigen.pdb'
print "remove antibody of complex,so can caculate the value of antigen sasa "
pymol.finish_launching()
if loop != 0:
pymol.cmd.reinitialize()
cmd.fetch(pdbfileselect[0], async=0)
#cmd.select('targetfile',pdbfileselect[0])
#cmd.save(dirname,(('targetfile')))
#cmd.delete(all)
#cmd.load(dirname)
cmd.remove(removechains)
cmd.remove('solvent')
cmd.remove('hydrogens')
cmd.remove('hetatm')
cmd.select('target1', pdbfileselect[0])
cmd.save(dirname, (('target1')))
cmd.remove('solvent')
cmd.remove('hydrogens')
cmd.remove('not ' + removechain)
cmd.select('target', pdbfileselect[0])
cmd.save(savefile, (('target')))
cmd.delete(all)
cmd.sync()
cmd.quit()
print "remove the antibody of complex have done"
def ray_tracer(width, height, scenes="all", ray=1, directory=""):
"""
ray_tracer will call png for scenes in a pymol session. All these scenes will be
produced with the same dimensions. By default, all scenes are generated with ray
tracing turned on in the current directory. Files will be named the same as their
scene names.
width, height = The width and height in pixels of the resulting images.
scenes = A scene or pythonic list of scenes to raytrace (e.g. "scene1" or ["scene1", "scene2"]). By default, do all scenes.
raytrace = Should we raytrace? Passed on to ray in png command.
directory = By default, images will be saved in the current directory. A relative path can be given
here for a different directory to be used.
"""
#Check if 'all' happens to be a scene name. If so, print an error message and quit.
assert (
"all" not in cmd.get_scene_list()
), "You have 'all' as a scene name. This causes a conflict with ray_tracer. Change the scene name."
#The directory name must end in a forward / (unless it is an empty string). If it doesn't, add one for the user
if (directory != "" and directory[-1:] != "/"): directory = directory + "/"
#Figure out which scenes to png
#if scenes == all, get all scenes and use that as our list.
if (scenes == 'all'):
scene_list = cmd.get_scene_list()
#if scenes is a list object, use that as our list.
elif (isinstance(scenes, list)):
scene_list = scenes
#Otherwise, assume we have a string that corresponds to a scene name
else:
scene_list = [scenes]
#Make sure that the scene_animation setting is set to 0, or changing view will not work.
cmd.set("scene_animation", 0)
#scene_list should now have a list of all the scenes to render.
#Loop over scene_list and cmd.png each one.
for scene in scene_list:
#Check that the scene actually exists, or spit out an error.
if (scene not in cmd.get_scene_list()):
print("The scene " + scene +
" is not a scene in your pymol environment. Skipping.")
continue
#Change to the appropriate scene.
cmd.scene(scene)
cmd.sync()
#Figure out the full name of the output file.
imgname = directory + scene + ".png"
#Actually make the image.
cmd.png(filename=imgname, width=width, height=height, ray=ray)
def testPngExists(self, sync_in_main_thread, sync_in_thread):
'''
Save a PNG image with width/height specified and
check if the file exists. a valid workaround would
be to call sync(), but if sync() is called inside the
thread, then it still doesnt work. ideally we shouldnt
need to call sync() at all, the png call should just do
it
'''
cmd.pseudoatom('m1')
cmd.show('spheres')
with testing.mktemp('.png') as filename:
th = threading.Thread(target=self._pngExistsAsync, args=[filename, sync_in_thread])
th.start()
th.join()
if sync_in_main_thread:
cmd.sync()
self.assertTrue(os.path.exists(filename), 'png file not written')
def roving_density(self, cleanup=0):
if not cleanup:
try:
cmd.load("$PYMOL_DATA/demo/il2.pdb")
cmd.set("suspend_updates", 1, quiet=1)
cmd.remove("hydro")
cmd.disable()
cmd.enable("il2")
cmd.map_new("map", "gaussian", "0.75", "il2")
cmd.feedback("disable", "objectmesh", "actions")
cmd.set("ribbon_color", "purple", "il2")
cmd.set("roving_detail", 1)
cmd.set("roving_origin", 1)
cmd.set("stick_radius", 0.12, "il2")
cmd.set("roving_sticks", 0)
cmd.set("roving_polar_contacts", 0)
cmd.set("line_width", "3")
cmd.set("roving_map1_name", "map")
cmd.isomesh("rov_m1", "map", 9999.0, "il2")
cmd.color("density", "rov_m1")
cmd.set_view ((\
0.132852688, -0.729740858, 0.670686543,\
-0.228543565, 0.635894477, 0.737154961,\
-0.964425683, -0.251212329, -0.082298420,\
0.000062190, 0.000183226, -58.861488342,\
13.349151611, -1.565427899, 22.383148193,\
55.259441376, 63.259449005, 0.000000000 ))
finally:
cmd.set("suspend_updates", 0, quiet=1)
cmd.refresh()
else:
cmd.set("roving_detail", 0)
cmd.set("roving_map1_name", "")
cmd.set("roving_polar_contacts", 7)
cmd.set("roving_sticks", 6)
cmd.delete("il2")
cmd.delete("map")
cmd.set("line_width", 1.5)
cmd.refresh()
cmd.set("roving_detail", 0)
cmd.delete("rov_*")
cmd.sync()
def testPngExists(self, sync_in_main_thread, sync_in_thread):
'''
Save a PNG image with width/height specified and
check if the file exists. a valid workaround would
be to call sync(), but if sync() is called inside the
thread, then it still doesnt work. ideally we shouldnt
need to call sync() at all, the png call should just do
it
'''
cmd.pseudoatom('m1')
cmd.show('spheres')
with testing.mktemp('.png') as filename:
th = threading.Thread(target=self._pngExistsAsync,
args=[filename, sync_in_thread])
th.start()
th.join()
if sync_in_main_thread:
cmd.sync()
self.assertTrue(os.path.exists(filename), 'png file not written')
def set_chains(selection='all'):
"""
alters chain identifiers based on segi and type
"""
selection = f'{selection} and prot'
stored.tmp_set = set()
cmd.iterate(f'{selection}', f'stored.tmp_set.add(segi)')
cmd.sync()
# create dictionary to match each segi to a chain
let = string.ascii_uppercase
if len(let) > len(stored.tmp_set):
chains = dict(zip(stored.tmp_set, let))
else:
chains = dict.fromkeys(stored.tmp_set, 'A')
stored.chains = chains
cmd.alter(selection, f'chain=stored.chains[segi]')
cmd.sync()
def roving_density(self,cleanup=0):
if not cleanup:
try:
cmd.load("$PYMOL_DATA/demo/il2.pdb")
cmd.set("suspend_updates",1,quiet=1)
cmd.remove("hydro")
cmd.disable()
cmd.enable("il2")
cmd.map_new("map","gaussian","0.75","il2")
cmd.feedback("disable","objectmesh","actions")
cmd.set("ribbon_color","purple","il2")
cmd.set("roving_detail",1)
cmd.set("roving_origin",1)
cmd.set("stick_radius",0.12,"il2")
cmd.set("roving_sticks",0)
cmd.set("roving_polar_contacts",0)
cmd.set("line_width","3")
cmd.set("roving_map1_name","map")
cmd.isomesh("rov_m1","map",9999.0,"il2")
cmd.color("density","rov_m1")
cmd.set_view ((\
0.132852688, -0.729740858, 0.670686543,\
-0.228543565, 0.635894477, 0.737154961,\
-0.964425683, -0.251212329, -0.082298420,\
0.000062190, 0.000183226, -58.861488342,\
13.349151611, -1.565427899, 22.383148193,\
55.259441376, 63.259449005, 0.000000000 ))
finally:
cmd.set("suspend_updates",0,quiet=1)
cmd.refresh()
else:
cmd.set("roving_detail",0)
cmd.set("roving_map1_name","")
cmd.set("roving_polar_contacts",7)
cmd.set("roving_sticks",6)
cmd.delete("il2")
cmd.delete("map")
cmd.set("line_width",1.5)
cmd.refresh()
cmd.set("roving_detail",0)
cmd.delete("rov_*")
cmd.sync()
def set_vdw(selection='all'):
"""
alters vdw radii based on bead type provided by garnish
"""
def alter_vdw(elem, vdw):
vdw_patterns = {
re.compile('[QPNCX][\w\d]|W'): 2.35,
re.compile('S([QPNCX][\w\d]|W)'): 2.00,
re.compile('T([QPNCX][\w\d]|W)'): 1.65,
}
for pattern, radius in vdw_patterns.items():
if pattern.match(elem):
return radius
# TODO: if anything goes wrong, just return the original vdw for now
return vdw
stored.alter_vdw = alter_vdw
cmd.alter(selection, 'vdw=stored.alter_vdw(elem, vdw)')
cmd.sync()
def nice(style='clean', selection='all'):
"""
DESCRIPTION
apply a nice preset for the visualization of a martini molecular system
NOTE: ignores object names matching `*_elastics`
USAGE
nice [style [, selection]]
ARGUMENTS
style = one of clean|rainbow|balls (default='clean')
selection (default='all')
"""
# sanitize input
if style not in stored.nice_set.keys():
print(
f'Error: {style} is not a valid option. Choose from: {list(stored.nice_set.keys())}'
)
return
nicesele()
cmd.set('stick_radius', 0.7)
cmd.sync()
for sel_type, commands in stored.nice_set[style].items():
for command in commands.values():
if command is not None:
# run function with its arguments. All functions must have `selection` as
# valid argument for this to work!
command[0](*command[1:],
selection=f'{selection} and {sel_type}')
# show elastics as lines once again
cmd.show_as('lines', f'{selection} and *_elastics')
cmd.color('orange', f'{selection} and *_elastics')
cmd.set('line_width', 1)
def split(rule='molecules', selection='all'):
"""
DESCRIPTION
split the system into multiple objects according to a specific rule
USAGE
split [, rule [, selection]]
ARGUMENTS
rule = one of groups|molecules (default='groups')
groups: splits protein, lipids, solvent, nucleic
molecules: every molecule in a different object (also chains)
selection (default='all')
"""
# sanitize input
if rule not in ('groups', 'molecules'):
print(
f'Error: {rule} is not a valid option. Choose from: groups, molecules'
)
return
nicesele()
cmd.sync()
if rule == 'groups':
for obj in cmd.get_object_list(selection):
groups = [
sel for sel in ('prot', 'lip', 'solv', 'nucl')
if cmd.count_atoms(f'{obj} and {sel}') > 0
]
for sel in groups:
cmd.create(f'{obj}_{sel}', f'{obj} and {sel}')
cmd.sync()
elif rule == 'molecules':
for obj in cmd.get_object_list(selection):
stored.tmp_set = set()
cmd.iterate(f'{selection} and {obj}',
f'stored.tmp_set.add(int(segi))')
cmd.sync()
zeros = len(str(max(stored.tmp_set)))
for mol_id in stored.tmp_set:
cmd.create(f'{obj}_n{mol_id:0{zeros}}', f'segi {mol_id}')
def testSync(self):
cmd.sync()
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()
def _pngExistsAsync(self, filename, sync_in_thread):
cmd.png(filename, width=100, height=100, ray=0)
cmd.draw()
if sync_in_thread:
cmd.sync()
def cheese(render_type,
savefile=None,
duration=5,
mode='ray',
width=1920,
height=1080):
"""
DESCRIPTION
loads nice settings for rendering and take a picture or records a movie of the trajectory.
Before running, position the camera to have a good point of view.
USAGE
cheese render_type [, savefile [, duration [, mode [, width [, height]]]]]
ARGUMENTS
render_type = one of:
- set: only load render settings
- snap: take a snapshot of the current view
- traj: record the whole trajectory
- bullettime: make a whole revolution around the z axis
savefile = file name to use to save the movie, without extension. If not given,
the movie won't be save but only shown (default=None)
duration = duration in seconds of the final movie (default=5)
mode = set to 'draw' to disable ray tracing (default='ray')
width = width of the movie in pixel (default=1920)
height = height of the movie in pixel (default=1080)
"""
duration = int(duration)
# # store settings for later restoration
# store_settings.load()
# cmd.sync()
# settings_file = '/tmp/stir_settings.py'
# cmd.do(f'store_settings {settings_file}')
# load nice settings for rendering
config.rendering()
cmd.sync()
def init_movie():
cmd.mdelete()
cmd.mview('reset')
frames = duration * 30
cmd.mset(f'1x{frames}')
return frames
# only load settings
if render_type == 'set':
return
# take a simple picture
elif render_type == 'snap':
cmd.ray(width=width, height=height)
# make a still movie of the whole trajectory
elif render_type == 'traj':
frames = init_movie()
cmd.mview('store', 1, state=1)
cmd.sync()
states = cmd.count_states()
cmd.mview('store', frames, state=states)
cmd.sync()
cmd.mview('reinterpolate')
cmd.sync()
# make a movie rotating 360 around the z axis
elif render_type == 'bullettime':
frames = init_movie()
cmd.mview('store', 1)
cmd.mview('store', frames)
# TODO: moving objects seems a bit hacky. Any better solution?
for obj in cmd.get_object_list():
cmd.mview('store', object=obj)
cmd.sync()
for i in range(4):
frame = int(frames * (i + 1) / 4)
cmd.frame(frame)
cmd.sync()
for obj in cmd.get_object_list():
cmd.rotate('z', angle=-90, camera=0, object=obj)
cmd.sync()
cmd.mview('store', object=obj)
cmd.sync()
cmd.mview('reinterpolate', object=obj, power=1)
cmd.sync()
cmd.mview('reinterpolate')
cmd.sync()
else:
raise ValueError(
f'{render_type} is not a valid argument. See `help cheese`.')
# save if required, otherwise just play it
if savefile:
if render_type in (
'traj',
'bullettime',
):
cmd.viewport(width, height)
movie.produce(f'{savefile}.mp4',
mode=mode,
encoder='ffmpeg',
quality=100)
elif render_type in ('snap', ):
cmd.png(f'{savefile}.png', dpi=300)
cmd.sync()
else:
if render_type in (
'traj',
'bullettime',
):
cmd.mplay()
elif render_type in ('snap', ):
pass
def _pngExistsAsync(self, filename, sync_in_thread):
cmd.png(filename, width=100, height=100, ray=0)
cmd.draw()
if sync_in_thread:
cmd.sync()
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)
def visualize_path((path, freq, pdb, outfile)):
nodes = path.split('=>')
save_pse = True
#print outfile
#print freq
#sys.exit()
#print nodes
node_color = sns.color_palette("Reds", n_colors=len(nodes))
node_color.reverse()
#print node_color
cmd.reinitialize()
cmd.load(pdb)
cmd.hide("everything")
# cmd.show("ribbon")
cmd.show("cartoon")
#resa = re.sub(r':\w', ':', nodes[0])
(chain, resnum) = re.split(':\w', nodes[0])
a = "chain {} and resi {} and name CA".format(chain, resnum)
cmd.show("spheres", a)
colorname = "color" + str(0)
cmd.set_color(colorname, node_color[0])
cmd.color(colorname, a)
cmd.label(a, '" %s:%s" % (resi, resn)')
cmd.set("label_color", "red", a)
#cmd.label(a, '"%s" % (resi)')
for i in range(1, len(nodes)):
#resb = re.sub(r':\w', ':', nodes[i])
(chain, resnum) = re.split(':\w', nodes[i])
b = "chain {} and resi {} and name CA".format(chain, resnum)
#print a,b
cmd.bond(a, b)
cmd.set_bond("line_width", 5, a, b)
cmd.set_bond("line_color", "red", a, b)
cmd.show("lines", a)
cmd.show("spheres", a)
cmd.show("lines", b)
cmd.show("spheres", b)
colorname = "color" + str(i)
cmd.set_color(colorname, node_color[i])
cmd.color(colorname, b)
#cmd.label(b,'" %s:%s" % (resi, resn)')
cmd.label(b, '" %s%s" % (one_letter[resn],resi)')
#cmd.label(b, '"%s" % (resi)')
a = b
#cmd.color(node_color[len)],b)
#cmd.space("cmyk")
cmd.set("ray_shadow", "off")
#cmd.bg_color("white")
#cmd.label_position([3,2,1])
cmd.set("label_position", (2, 2, 2))
sele = "chain A and resi 131 and not name H*"
cmd.show("sticks", sele)
norm_factor = freq[nodes[0]]
node_freqs = {k: freq[k] for k in nodes}
#print node_freqs
#print freq[nodes]
#print sorted(node_freqs.values())
node_freq_min = min(node_freqs.values())
#print node_freq_min
#sys.exit()
scale_cutoff = min(0.25, min(node_freqs.values()) / norm_factor)
#print nodes
#print node_freqs
#print freq
#sys.exit()
for res in freq.keys():
scale = freq[res] / norm_factor
#print res,freq[res],scale,scale_cutoff
if scale > scale_cutoff:
(chain, resnum) = re.split(':\w', res)
sele = "chain {} and resi {} and name CA".format(chain, resnum)
cmd.show("spheres", sele)
cmd.set("sphere_scale", scale, selection=sele)
if not res in nodes: #give residue not on highest freq path a different color
last_color = len(nodes) - 1
#colorname="color"+str(last_color)
colorname = "grey60" #color" + str(last_color)
cmd.color(colorname, sele)
cmd.label(sele, '" %s%s" % (one_letter[resn],resi)')
cmd.util.cnc(sele)
#cmd.center("chain A")
cmd.deselect()
### cut below here and paste into script ###
cmd.set_view([\
0.354133159, -0.915425777, 0.191264138,\
-0.626513124, -0.384070098, -0.678214610,\
0.694314659, 0.120348796, -0.709537089,\
-0.000188543, -0.000061929, -119.472831726,\
33.501632690, 81.432159424, 143.041992188,\
87.118530273, 151.834289551, -20.000000000\
])
### cut above here and paste into script ###
#outfile="{}cluster{}-{}.png".format(outfolder,cut,freq)
# if ray:
# cmd.ray()
# cmd.sync(10)
cmd.save(outfile)
cmd.sync(10)
if save_pse:
outfile_pse = re.sub('.png', '.pse', outfile)
cmd.save(outfile_pse)
cmd.sync(10)
def load_traj(*args, **kwargs):
cmd.set('defer_builds_mode', 3)
cmd.sync()
cmd.load_traj(*args, **kwargs)
def testSync(self):
cmd.sync()
def refresh():
cmd.sync()
cmd.refresh()
time.sleep(0.1)
# set nice scene
obj_list = cmd.get_object_list()
cmd.set("pdb_retain_ids", 0)
offset = 0
for i, obj in enumerate(obj_list):
cmd.alter(obj, "chain=" + str(i)) #fix chain
cmd.alter(obj, "resi=" + str(i)) #fix res numbers
cmd.alter(obj, "ID=ID+" + str(offset)) #generate unique IDS
offset = len(cmd.get_model(obj).atom)
if (obj == "O71"):
cmd.alter(obj, "resn=\"O71\"") #a very special case
cmd.alter(obj,
"resn=resn.replace(\"_\", \"T\")") #replace underscores
cmd.sync()
cmd.save(out_pdb_path)
print(out_pdb_path)
## GET ATOMS
atom_list = up.pymol_selection_to_atom_list("all")
## Filtering for rings
try:
RingFilter = Filter.RingFilter(atom_list)
pdb_blocks = [
cmd.get_pdbstr(obj) for obj in cmd.get_object_list()
] #u.convert_atoms_to_pdb_molecules(atom_list)
RingFilter.get_args(lambda x: (pdb_blocks))
filtered_atoms = RingFilter.filter()