# -*- coding: utf-8 -*-
# Aqua-Duct, a tool facilitating analysis of the flow of solvent molecules in molecular dynamic simulations
# Copyright (C) 2016-2018 Tomasz Magdziarz, Alicja Płuciennik, Michał Stolarczyk <info@aquaduct.pl>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import numpy as np
import cPickle as pickle
import os
import tarfile
import aquaduct.visual.pymol_cgo as cgo
from aquaduct.geom import traces
from aquaduct.traj.paths import PathTypesCodes
from aquaduct.utils.helpers import list_blocks_to_slices
from aquaduct.visual.helpers import color_codes, cc
from aquaduct.utils.helpers import create_tmpfile, listify
[docs]class BasicPymolCGO(object):
cgo_entity_begin = []
cgo_entity_end = []
[docs] def __init__(self):
self.cgo_entity = None
self.previous = None
self.clean()
[docs] def clean(self,empty=False):
self.previous = None
self.cgo_entity = []
if not empty:
self.cgo_entity.extend(self.cgo_entity_begin)
[docs] def new(self):
self.previous = None
[docs] def get(self):
return self.cgo_entity + self.cgo_entity_end
[docs] @staticmethod
def make_color_triple(color_definition):
return tuple(color_definition)
[docs]class BasicPymolCGOLines(BasicPymolCGO):
cgo_entity_begin = [cgo.BEGIN, cgo.LINES]
cgo_entity_end = [cgo.END]
[docs] def add(self, coords=None, color=None):
if color is not None:
self.cgo_entity.append(cgo.COLOR)
# self.cgo_entity.extend(map(float, color))
self.cgo_entity.append(self.make_color_triple(map(float, color)))
if coords is not None:
for nr, coord in enumerate(coords):
if self.previous is not None:
self.cgo_entity.append(cgo.VERTEX)
self.cgo_entity.extend(map(float, self.previous))
self.cgo_entity.append(cgo.VERTEX)
self.cgo_entity.extend(map(float, coord))
self.previous = coord
[docs]class BasicPymolCGOSpheres(BasicPymolCGO):
cgo_entity_begin = []
cgo_entity_end = []
[docs] def add(self, coords=None, radius=None, color=None):
# color to colors...
if color is not None:
color = np.matrix(color).A
if radius is not None:
radius = np.matrix(radius).A1
if coords is not None:
for nr, coord in enumerate(coords):
if color is not None:
if len(color) > 1:
c = color[nr]
else:
c = color[0]
self.cgo_entity.append(cgo.COLOR)
self.cgo_entity.append(self.make_color_triple(map(float, c)))
# self.cgo_entity.extend(map(float, c))
self.cgo_entity.append(cgo.SPHERE)
self.cgo_entity.extend(map(float, coord))
if radius is not None:
if len(radius) > 1:
r = radius[nr]
else:
r = radius
else:
r = 1.
self.cgo_entity.append(float(r))
# [cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2
[docs]class BasicPymolCGOPointers(BasicPymolCGO):
cgo_entity_begin = []
cgo_entity_end = []
[docs] def add_cone(self, coords1=None, coords2=None, radius1=None, radius2=None, color1=None, color2=None):
# color to colors... ???
if coords1 is not None and coords2 is not None:
self.cgo_entity.append(cgo.CONE)
self.cgo_entity.extend(map(float, coords1))
self.cgo_entity.extend(map(float, coords2))
self.cgo_entity.append(float(radius1))
self.cgo_entity.append(float(radius2))
self.cgo_entity.append(self.make_color_triple(map(float, color1)))
self.cgo_entity.append(self.make_color_triple(map(float, color2)))
# self.cgo_entity.extend(map(float, color1))
# self.cgo_entity.extend(map(float, color2))
self.cgo_entity.extend([cgo.NULL, cgo.POINTS])
[docs] def add_pointer(self, point=None, direction=None, length=None, color=None, reverse=False):
vec = point - direction
vec_len = np.sqrt(np.sum(vec ** 2))
vec = (vec / vec_len) * length
vec = point + vec
if reverse:
self.add_cone(coords1=point, coords2=vec, radius1=length / 3., radius2=0, color1=color, color2=color)
else:
self.add_cone(coords1=vec, coords2=point, radius1=length / 3., radius2=0, color1=color, color2=color)
[docs]class SimpleTarWriteHelper(object):
[docs] def __init__(self):
self.tar_fh = None
self.tmp_file = create_tmpfile()
[docs] def open(self, filename):
self.tar_fh = tarfile.open(filename, 'w:gz')
[docs] def save_object2tar(self, obj, name):
with open(self.tmp_file, 'w') as f:
pickle.dump(obj, f)
self.save_file2tar(self.tmp_file, name)
[docs] def save_file2tar(self, filename, name):
self.tar_fh.add(filename, arcname=name)
[docs] def __del__(self):
if self.tar_fh is not None:
self.tar_fh.close()
os.unlink(self.tmp_file)
[docs]class ConnectToPymol(object):
cgo_line_width = 2.
ct_pymol = 'pymol'
ct_file = 'file'
[docs] def __init__(self):
self.connection_type = None # possible types are pymol and file
self.script_fh = None
self.data_fh = SimpleTarWriteHelper()
self.cmd = None
[docs] @listify
def decode_color(self, cgo_object):
for element in cgo_object:
if isinstance(element, tuple):
for e in element:
yield e
else:
yield element
[docs] def init_pymol(self):
import pymol
self.cmd = pymol.cmd
pymol.finish_launching()
self.cmd.set('cgo_line_width', ConnectToPymol.cgo_line_width)
self.connection_type = self.ct_pymol
[docs] def init_script(self, filename):
self.script_fh = open(filename, 'w')
data_filename = os.path.splitext(os.path.basename(filename))[0] + '.tar.gz'
self.data_fh.open(data_filename)
self.connection_type = self.ct_file
# TODO: Script generated by this object does not work properly if run via PyMOL.
# init lines, imports etc.
self.script_fh.write('''import argparse
parser=argparse.ArgumentParser(description="Aqua-Duct visualization script")
parser.add_argument("--save-session",action="store",dest="session",required=False,default=None,help="Pymol session file name.")
parser.add_argument("--discard",action="store",dest="discard",required=False,default='',help="Objects to discard.")
parser.add_argument("--keep",action="store",dest="keep",required=False,default='',help="Objects to keep.")
parser.add_argument("--force-color",action="store",dest="fc",required=False,default='',help="Force specific color.")
parser.add_argument("--fast",action="store_true",dest="fast",required=False,help="Disable all objects while loading.")
args,unknown=parser.parse_known_args()
import sys
if unknown: print >> sys.stderr, "WARNING: Unknown options were used: "+" ".join(unknown)
def _kd_order():
if args.keep=='' and args.discard!='': return 'd'
if args.keep!='' and args.discard=='': return 'k'
if args.keep=='' and args.discard=='': return None
if sys.argv.index('--keep')<sys.argv.index('--discard'): return 'k'
return 'd'
kd_order = _kd_order()
def discard(name):
if len([d for d in args.discard.split() if d in name])>0: return True
return False
def keep(name):
if len([k for k in args.keep.split() if k in name])>0: return True
return False
def proceed(name):
if kd_order == 'k':
if not keep(name): return False
elif discard(name): return False
elif kd_order == 'd':
if discard(name):
if not keep(name): return False
return True
from pymol import cmd,finish_launching
finish_launching()
print "Loading Aqua-Duct visualization..."
cmd.set("cgo_line_width",%d)
cmd.set("line_smooth","off")
from os import close,unlink
from os.path import splitext,isfile
import tarfile
import cPickle as pickle
from tempfile import mkstemp
fd, pdb_filename = mkstemp(suffix=".pdb")
close(fd)
max_state=0
arch_file="%s"
if not isfile(arch_file):
import pymol
if pymol.IS_WINDOWS:
print "Please open visualization script using 'Open with' context menu and choose PyMol executable."
print "Alternatively, if you have PyMol installed as Python module, open visulaization script with Python executable."
while (pymol._ext_gui is None): pymol = reload(pymol)
while (not hasattr(pymol._ext_gui,'root')): pymol = reload(pymol)
import tkFileDialog
arch_file=tkFileDialog.askopenfilename(filetypes=[("AQ Vis Arch","*.tar.gz")],title="Select AQ visualization archive",parent=pymol._ext_gui.root)
data_fh=tarfile.open(arch_file,"r:gz")
def decode_color(cgo_object,fc=None):
for element in cgo_object:
if isinstance(element,tuple):
if fc is None:
for e in element: yield e
else:
for e in fc: yield e
else:
yield element
def load_object(filename,name,state):
if not proceed(name): return
global max_state
print "Loading %s" % splitext(filename)[0]
obj=pickle.load(data_fh.extractfile(filename))
if name in args.fc.split():
forced_color=args.fc.split()[args.fc.split().index(name)+1]
forced_color=cmd.get_color_tuple(forced_color)
obj=decode_color(obj,fc=forced_color)
else:
obj=decode_color(obj)
cmd.load_cgo(obj,name,state)
if state<2:
if args.fast: cmd.disable("all")
else: cmd.refresh()
if state>max_state:
max_state=state
def load_pdb(filename,name,state):
if not proceed(name): return
global max_state
with open(pdb_filename,'w') as fpdb:
fpdb.write(data_fh.extractfile(filename).read())
cmd.load(pdb_filename,state=state,object=name)
if state>max_state:
max_state=state
''' % (self.cgo_line_width, data_filename, "%s", "% s"))
[docs] def add_cgo_object(self, name, cgo_object, state=None):
if state is None:
state = 1
if self.connection_type == self.ct_pymol:
self.cmd.load_cgo(self.decode_color(cgo_object), str(name), state)
elif self.connection_type == self.ct_file:
obj_name = '%s_%d.dump' % (name, state)
self.data_fh.save_object2tar(cgo_object, obj_name)
self.script_fh.write('''load_object("%s","%s",%d)''' % (obj_name, str(name), state))
self.script_fh.write(os.linesep)
# self.script_fh.write('''self.cmd.refresh()''')
# self.script_fh.write(os.linesep)
[docs] def del_cgo_object(self, name, state=None):
raise NotImplementedError("Deletion of CGO objects is not implemented yet.")
[docs] def load_pdb(self, name, filename, state=None):
if state is None:
state = 1
if self.connection_type == self.ct_pymol:
self.cmd.load(filename, state=state, object=name)
self.cmd.show_as('cartoon',name)
self.cmd.color('silver',name)
elif self.connection_type == self.ct_file:
# save pdblile as string
filename_new = '%s_%d.pdb' % (name, state)
self.data_fh.save_file2tar(filename, filename_new)
self.script_fh.write('''load_pdb("%s","%s",%d)''' % (filename_new, name, state))
self.script_fh.write(os.linesep)
self.script_fh.write('''if proceed("%s"): cmd.show_as('cartoon','%s')''' % (name,name))
self.script_fh.write(os.linesep)
self.script_fh.write('''if proceed("%s"): cmd.color('silver','%s')''' % (name,name))
self.script_fh.write(os.linesep)
[docs] def orient_on(self, name):
if self.connection_type == self.ct_pymol:
self.cmd.orient(name)
elif self.connection_type == self.ct_file:
self.script_fh.write('''if proceed("%s"): cmd.orient("%s")''' % (name, name))
self.script_fh.write(os.linesep)
[docs] def __del__(self):
if self.connection_type == self.ct_file:
self.script_fh.write('''data_fh.close()
unlink(pdb_filename)
if args.fast: cmd.enable("all")
print "Aqua-Duct visualization loaded."
if args.session:
print "Preparing data to save session..."
for state in range(max_state):
cmd.set_frame(state+1)
cmd.refresh()
if (state+1)%100==0:
print "wait... %d of %d done..." % (state+1,max_state)
print "%d of %d done." % (state+1,max_state)
print "Saving session..."
cmd.set_frame(1)
cmd.save(args.session,state=0)
print "Let the Valve be always open!"
print "Goodby!"
cmd.quit()
''')
self.script_fh.write(os.linesep)
self.script_fh.close()
[docs]class SinglePathPlotter(object):
[docs] def __init__(self, pymol_connector, linearize=None):
self.cgo_lines = BasicPymolCGOLines()
self.cgo_spheres = BasicPymolCGOSpheres()
self.cgo_pointers = BasicPymolCGOPointers()
# linearize have to be callable! or False or None
self.linearize = linearize
self.pymol_connector = pymol_connector
# TODO: take care of proper colors handling for smoothed and not smoothed traces!
[docs] def add_single_path_continous_trace(self,
spath,
smooth=None,
plot_in=True,
plot_object=True,
plot_out=True,
plot_walk=True,
**kwargs):
self.cgo_lines.new()
# get coords
coords_cont = spath.get_coords_cont(smooth)
# create slices
if smooth:
sls = tuple(list_blocks_to_slices(spath.types_cont))
else:
sls = tuple(list_blocks_to_slices(spath.etypes_cont))
# create traces
traces_list = tuple([coords_cont[sl] for sl in sls])
new_line = False
for trace, sl in zip(traces.midpoints(traces_list), sls):
if len(trace) > 0:
# now trace has midpoints
# get type and etype
t = spath.types_cont[sl][0]
et = spath.etypes_cont[sl][0]
if smooth:
et = et[0]
# plot, if allowed
if (plot_in and t == PathTypesCodes.path_in_code) or (
plot_object and t == PathTypesCodes.path_object_code) or (
plot_out and t == PathTypesCodes.path_out_code) or (plot_walk and t == PathTypesCodes.path_walk_code):
# get color
c = color_codes(et)
# now, it is possible to linearize!
if smooth and self.linearize:
trace = self.linearize(trace)
self.cgo_lines.add(trace, cc(c))
# new_line = True
if new_line:
self.cgo_lines.new()
new_line = False
[docs] def paths_trace(self, spaths,
smooth=None,
name='paths',
state=None,
**kwargs):
# if state_function is None
if state is None:
state = 1
self.cgo_lines.clean()
for nr, spath in enumerate(spaths):
self.add_single_path_continous_trace(spath, smooth=smooth, **kwargs)
if self.cgo_lines.cgo_entity == self.cgo_lines.cgo_entity_begin:
self.cgo_lines.clean(empty=True)
else:
self.pymol_connector.add_cgo_object(name, self.cgo_lines.get(), state=state)
[docs] def paths_inlets(self, spaths,
smooth=None,
color=None,
plot_in=True,
plot_out=True,
name='in-out-let',
state=None,
**kwargs):
if state is None:
state = 1
self.cgo_pointers.clean()
for nr, spath in enumerate(spaths):
coords = spath.get_coords_cont(smooth=smooth)
etypes = spath.etypes_cont
if plot_in:
if color is None:
c = color_codes(etypes[0])
else:
c = color
self.cgo_pointers.add_pointer(point=coords[0],
direction=coords[1],
length=1.,
color=cc(c))
if plot_out:
if color is None:
c = color_codes(etypes[-1])
else:
c = color
self.cgo_pointers.add_pointer(point=coords[-1],
direction=coords[-2],
length=1.,
color=cc(c),
reverse=True)
self.pymol_connector.add_cgo_object(name, self.cgo_pointers.get(), state=state)
[docs] def scatter(self, coords,
radius=0.4,
color='r',
name='scatter',
state=None):
if isinstance(color, str):
color = cc(color)
if state is None:
state = 1
self.cgo_spheres.clean()
self.cgo_spheres.add(coords=coords, radius=radius, color=color)
self.pymol_connector.add_cgo_object(name, self.cgo_spheres.get(), state=state)
[docs] def convexhull(self, chull,
color='m',
name='convexhull',
state=None):
if isinstance(color, str):
color = cc(color)
if state is None:
state = 1
self.cgo_lines.clean()
for nr, facet in enumerate(chull.facets):
if nr > 0:
self.cgo_lines.new()
self.cgo_lines.add(facet, color=color)
self.pymol_connector.add_cgo_object(name, self.cgo_lines.get(), state=state)
