def plot(self, coords, freqs, min, max, min_r, max_r, col_min, col_max, out_obj):
widths = []
cgo = []
nsteps = 200
gradient = []
cols = []
#gradient = self.make_gradient("sbcol", array(cmd.get_color_tuple(self.interpret_color(col_min))), array(cmd.get_color_tuple(self.interpret_color(col_max))), nsteps )
gradient = self.make_gradient("sbcol", array(col_min), array(col_max), nsteps )
for i,freq in enumerate(freqs):
if freq < min:
freqs[i] = min
elif freq > max:
freqs[i] = max
widths.append( (freqs[i] - min)*((max_r - min_r) / (max - min)) + min_r )
cols.append(gradient[int(round(freqs[i]/max*(nsteps-1)))])
for i in range(0, len(coords), 2):
cgo.append(CYLINDER)
cgo.extend(coords[i])
cgo.extend(coords[i+1])
cgo.append(widths[i/2])
cgo.extend(cmd.get_color_tuple(self.interpret_color(cols[i/2])))
cgo.extend(cmd.get_color_tuple(self.interpret_color(cols[i/2])))
cmd.load_cgo(cgo,out_obj)
def visualize_orientation(direction, center=[0, 0, 0], scale=1.0, symmetric=False, color="green", color2="red"):
"""
Draw an arrow. Helper function for "helix_orientation" etc.
"""
from pymol import cgo
color_list = cmd.get_color_tuple(color)
color2_list = cmd.get_color_tuple(color2)
if symmetric:
scale *= 0.5
end = cpv.add(center, cpv.scale(direction, scale))
radius = 0.3
obj = [cgo.SAUSAGE]
obj.extend(center)
obj.extend(end)
obj.extend([radius, 0.8, 0.8, 0.8])
obj.extend(color_list)
if symmetric:
start = cpv.sub(center, cpv.scale(direction, scale))
obj.append(cgo.SAUSAGE)
obj.extend(center)
obj.extend(start)
obj.extend([radius, 0.8, 0.8, 0.8])
obj.extend(color2_list)
coneend = cpv.add(end, cpv.scale(direction, 4.0 * radius))
if cmd.get_version()[1] >= 1.2:
obj.append(cgo.CONE)
obj.extend(end)
obj.extend(coneend)
obj.extend([radius * 1.75, 0.0])
obj.extend(color_list * 2)
obj.extend([1.0, 1.0]) # Caps
cmd.load_cgo(obj, get_unused_name("oriVec"), zoom=0)
def cylinder(obj, coords, colors, radius):
"""draw the bonds between rings"""
for idx, coord in enumerate(coords):
r1, g1, b1 = cmd.get_color_tuple(colors[idx][0])
r2, g2, b2 = cmd.get_color_tuple(colors[idx][1])
x1, y1, z1, x2, y2, z2 = coord
obj.extend([CYLINDER, x1, y1, z1, x2, y2, z2, radius, r1, g1, b1, r2, g2, b2])
return obj
def setselection(self):
if (self.selectionentry.get() != ""):
self.selection=self.selectionentry.get()
#Color of each residue is stored in stored.colorlist to check if the molecule has a colorgradient
stored.colorlist = []
cmd.iterate (self.selection+" & name CA", "stored.colorlist.append(int(color))")
if (len(stored.colorlist)==0):
#for other objects (e.g. density...)
stored.colorlist.append(cmd.get_object_color_index(self.selection))
stored.colorlist.append(cmd.get_object_color_index(self.selection))
initialcolornterm=cmd.get_color_tuple(stored.colorlist[0])
initialcolorcterm=cmd.get_color_tuple(stored.colorlist[len(stored.colorlist)-1])
self.farbe1=initialcolornterm[0]*255, initialcolornterm[1]*255, initialcolornterm[2]*255
self.farbe2=initialcolorcterm[0]*255, initialcolorcterm[1]*255, initialcolorcterm[2]*255
#Set active object to label
self.label.config(text=self.selection)
#check if there is a gradient and adjust Mono/Gradbutton
if (initialcolornterm==initialcolorcterm):
self.Monobutton.select()
self.Mono()
elif (initialcolornterm!=initialcolorcterm):
self.Gradbutton.select()
self.Grad()
#adjust colorfields
self.colorfield1.config(bg=self.RGBToHTMLColor(self.farbe1))
self.colorfield2.config(bg=self.RGBToHTMLColor(self.farbe2))
self.Farbe1button.select()
self.Farbe1()
#Set scales to initialcolor of the new object
if (self.colorsystem=='rgb'):
startred=255*initialcolornterm[0]
startgreen=255*initialcolornterm[1]
startblue=255*initialcolornterm[2]
self.ScaleRed.set(startred)
self.ScaleGreen.set(startgreen)
self.ScaleBlue.set(startblue)
elif (self.colorsystem=='hsv'):
hsvcolor=colorsys.rgb_to_hsv(initialcolornterm[0], initialcolornterm[1], initialcolornterm[2])
h=hsvcolor[0]
s=hsvcolor[1]
v=hsvcolor[2]
self.ScaleRed.set(h)
self.ScaleGreen.set(s)
self.ScaleBlue.set(v)
def _getColorTuple(self, color):
if isinstance(color, tuple):
return color
if isinstance(color, list):
return tuple(color)
elif isinstance(color, int):
if color in dict((k,v) for (v,k) in cmd.get_color_indices()):
return cmd.get_color_tuple(color)
else:
hexval = color & 0xffffff
return tuple(map(lambda x: x/255., [ hexval >> 16, (hexval >> 8) & 0xff, hexval & 0xff ]))
elif isinstance(color, str):
return cmd.get_color_tuple(color)
def visualize_orientation(direction, center=[0.0]*3, scale=1.0, symmetric=False, color='green', color2='red'):
'''
DESCRIPTION
Draw an arrow. Helper function for "helix_orientation" etc.
'''
from pymol import cgo
color_list = cmd.get_color_tuple(color)
color2_list = cmd.get_color_tuple(color2)
if symmetric:
scale *= 0.5
end = cpv.add(center, cpv.scale(direction, scale))
radius = 0.3
obj = [cgo.SAUSAGE]
obj.extend(center)
obj.extend(end)
obj.extend([
radius,
0.8, 0.8, 0.8,
])
obj.extend(color_list)
if symmetric:
start = cpv.sub(center, cpv.scale(direction, scale))
obj.append(cgo.SAUSAGE)
obj.extend(center)
obj.extend(start)
obj.extend([
radius,
0.8, 0.8, 0.8,
])
obj.extend(color2_list)
coneend = cpv.add(end, cpv.scale(direction, 4.0*radius/cpv.length(direction)))
obj.append(cgo.CONE)
obj.extend(end)
obj.extend(coneend)
obj.extend([
radius * 1.75,
0.0,
])
obj.extend(color_list * 2)
obj.extend([
1.0, 1.0, # Caps
])
cmd.load_cgo(obj, get_unused_name('oriVec'), zoom=0)
def parseColor(color):
"""Parse various ways of specifying colors. Returns an array of 3-tuples
Examples:
"red" -> [(1,0,0)]
"red green" -> [(1,0,0),(0,1,0)]
(1,0,0) -> (1,0,0)
[(1,0,0),(0,1,0)] -> [(1,0,0),(0,1,0)]
"""
if type(color) == str:
colors = [cmd.get_color_tuple(c) for c in color.split()]
return colors
colors = None
if len(color) == 3:
# Try treating as a 3-tuple
try:
colors = [map(float,color)]
except TypeError:
pass # not a 3-tuple
if colors is None:
# Assume list of 3-tuples
if any([len(c) != 3 for c in color]):
raise ValueError("Unrecognized color format")
try:
colors = [map(float,c) for c in color]
except TypeError:
raise ValueError("Unrecognized color format")
# validate color
for col in colors:
if any([c < 0 or 1 < c for c in col]):
raise ValueError("Color value out of bounds: %s"%col)
return colors
def sele2Color(self, sele):
globalShader = self.optionMenu_shader.getvalue()
newShader = self.seleShaderDict[sele]
if newShader != globalShader: # should be different from global shader, otherwise do nothing
newColorInc = ShaderFactory.seleSlot[newShader]
stored.idcolor_list = set()
cmd.iterate(sele, 'stored.idcolor_list.add(int(color))')
if(len(stored.idcolor_list)>1):
self.selectionConsole.set('Warning: Selection [%s] contains more than one color.' % sele)
color = stored.idcolor_list.pop()
rgb_color = cmd.get_color_tuple(color)
if rgb_color[2] >= 0.990:
color_id = '%s %s %s' % (str(rgb_color[0])[0:5].ljust(5, '0'), str(rgb_color[1])[0:5].ljust(5, '0'), str(rgb_color[2]-newColorInc)[0:5])
else:
color_id = '%s %s %s' % (str(rgb_color[0])[0:5].ljust(5, '0'), str(rgb_color[1])[0:5].ljust(5, '0'), str(rgb_color[2]+newColorInc)[0:5])
# color_id from all the sele:shader dictionary
# determine color_id:shader pair
if color_id not in self.spColorShaderDict:
self.spColorShaderDict[color_id] = [newShader, sele]
# apply new color to atom set
newRGB = color_id.split(' ')
newColor = [float(newRGB[0]), float(newRGB[1]), float(newRGB[2])]
cmd.set_color(color_id, newColor)
cmd.color(color_id, sele)
print 'apply shader [%s] to selection [%s].' % (newShader, sele)
def load_msms_surface(filename, name="", _colors=None):
"""
DESCRIPTION
Load MSMS .vert and .face files as a CGO
"""
from pymol import cgo
from pymol.cgo import NORMAL, VERTEX, COLOR
if _colors:
_colors = [cmd.get_color_tuple(c) for c in _colors]
if filename.endswith(".vert") or filename.endswith(".face"):
filename = filename[:-5]
# vertex file
line_iter = iter(open(filename + ".vert"))
# skip header
for line in line_iter:
if not line.startswith("#"):
break
# read vertices
vertices = [None] # make 1-indexable
for line in line_iter:
data = line.split()
vertex = [float(x) for x in data[0:3]]
normal = [float(x) for x in data[3:6]]
sphere = int(data[7]) - 1
vertices.append((vertex, normal, sphere))
# faces file
line_iter = iter(open(filename + ".face"))
# skip header
for line in line_iter:
if not line.startswith("#"):
break
cgobuf = [cgo.BEGIN, cgo.TRIANGLES]
# read triangles
for line in line_iter:
for index in line.split()[:3]:
data = vertices[int(index)]
if _colors:
cgobuf.append(COLOR)
cgobuf.extend(_colors[data[2]])
cgobuf.append(NORMAL)
cgobuf.extend(data[1])
cgobuf.append(VERTEX)
cgobuf.extend(data[0])
cgobuf.append(cgo.END)
if not name:
name = cmd.get_unused_name("msmssurf")
cmd.load_cgo(cgobuf, name)
def fade_color( sele = "all", fade_extent = 0.7, by_chain = True ):
"""
Fade out RGB color values associated with a selection.
The by_chain variable assumes that chains are uniform color, to
allow speedy color setting based on chain, rather than based on
residue (which can take a long time!)
"""
pymol.stored.colors = []
cmd.iterate( sele, "stored.colors.append( (chain, resi, name, color))")
res_colors = {}
stored_colors = pymol.stored.colors;
cols = []
colorname = ''
chain = ''
prev_chain = ''
for chain, resi, name, color in stored_colors:
if name == 'CA' or name == 'P': # c-alpha atom
if by_chain and chain != prev_chain and len( colorname ) > 0:
cmd.color( colorname, 'chain %s and %s' % (prev_chain,sele) )
color_tuple = cmd.get_color_tuple(color)
cols = [];
for i in range(3):
cols.append( fade_extent + ( 1.0 - fade_extent ) * float(color_tuple[ i ]) )
colorname = 'temp_chain%s' % chain
cmd.set_color( colorname, cols )
if not by_chain: cmd.color( colorname, 'resi %s and chain %s and %s' % (resi,chain,sele) )
prev_chain = chain
if by_chain: cmd.color( colorname, 'chain %s and %s' % (chain,sele) )
def beads(obj, coords, colors, radius):
"""draw beads"""
for idx, coord in enumerate(coords):
r, g, b = cmd.get_color_tuple(colors[idx])
x1, y1, z1 = np.mean(coord, axis=0)
obj.extend([COLOR, r, g, b, SPHERE, x1, y1, z1, radius])
return obj
def cgo_arrow(atom1='pk1',
atom2='pk2',
radius=0.07,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + [
cgo.CONE
] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + [1.0, 0.0]
return obj
def spectrumproperty(propname, color_list, selection='(all)', minimum=None, maximum=None, quiet=1):
'''
DESCRIPTION
Define a color spectrum with as many color-stops as you like (at least 2).
ARGUMENTS
propname = string: property name which has a float value.
color_list = string: Space separated list of colors
... all other arguments like with `spectrum` command
'''
quiet = int(quiet)
colors = color_list.split()
if len(colors) < 2:
print 'failed! please provide at least 2 colors'
return
colvec = [cmd.get_color_tuple(i) for i in colors]
parts = len(colvec) - 1
value_list = []
cmd.iterate(selection, 'value_list.append(properties[propname])', space=locals())
if len(value_list) == 0:
print 'empty selection'
return
if minimum is None:
minimum = min(value_list)
if maximum is None:
maximum = max(value_list)
minimum, maximum = float(minimum), float(maximum)
val_range = (maximum - minimum) * (1 + 1e-6)
if not quiet:
print ' Spectrum: range (%.5f to %.5f)' % (minimum, maximum)
if maximum == minimum:
print 'no spectrum possible, only equal values'
return
rgb = lambda i, p, j: int(255 * (colvec[i+1][j] * p + colvec[i][j] * (1.0 - p)))
col_list = []
for value in value_list:
p = (value - minimum) / val_range * parts
i = int(p)
p -= i
col_list.append(0x40000000 +
rgb(i, p, 0) * 0x10000 +
rgb(i, p, 1) * 0x100 +
rgb(i, p, 2))
cmd.alter(selection, 'color = col_iter.next()', space={'col_iter': iter(col_list)})
cmd.recolor()
def get_cgo_vector_list(r, color_name, eps=1e-5):
rgb = cmd.get_color_tuple(color_name)
r = r / np.linalg.norm(r)
r = np.where(np.isclose(r, 0.0, atol=eps), 0., r) #see https://github.com/schrodinger/pymol-open-source/issues/220
cgo_list = [
cgo.CYLINDER, 0.0, 0.0, 0.0, l*r[0], l*r[1], l*r[2], w, *rgb, *rgb,
cgo.CONE, l*r[0], l*r[1], l*r[2], (h+l)*r[0], (h+l)*r[1], (h+l)*r[2], d, 0.0, *rgb, *rgb, 1.0, 1.0
]
return cgo_list
def draw_arrow(xyz1,
xyz2,
radius=0.5,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
'''
Draw an arrow; borrows heavily from cgi arrows.
'''
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
xyz1 = list(xyz1)
xyz2 = list(xyz2)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def visualize_orientation(direction, center=[0, 0, 0], scale=1.0, symmetric=False, color='green', color2='red'):
'''
Draw an arrow. Helper function for "helix_orientation" etc.
'''
from pymol import cgo
color_list = cmd.get_color_tuple(color)
color2_list = cmd.get_color_tuple(color2)
if symmetric:
scale *= 0.5
end = cpv.add(center, cpv.scale(direction, scale))
radius = 0.3
obj = [cgo.SAUSAGE]
obj.extend(center)
obj.extend(end)
obj.extend([
radius,
0.8, 0.8, 0.8,
])
obj.extend(color_list)
if symmetric:
start = cpv.sub(center, cpv.scale(direction, scale))
obj.append(cgo.SAUSAGE)
obj.extend(center)
obj.extend(start)
obj.extend([
radius,
0.8, 0.8, 0.8,
])
obj.extend(color2_list)
coneend = cpv.add(end, cpv.scale(direction, 4.0 * radius))
if cmd.get_version()[1] >= 1.2:
obj.append(cgo.CONE)
obj.extend(end)
obj.extend(coneend)
obj.extend([
radius * 1.75,
0.0,
])
obj.extend(color_list * 2)
obj.extend([
1.0, 1.0, # Caps
])
cmd.load_cgo(obj, get_unused_name('oriVec'), zoom=0)
def parseDistObj(obj):
if (obj[5][0][3][10] != 1): # 'show dashed' flag
return ""
N = obj[5][2][0][0]
points = obj[5][2][0][1]
ret = []
for p in points:
ret.append("%.3f" % p)
color = cmd.get_color_tuple(obj[5][0][2])
return "\ndists:%.3f,%.3f,%.3f:" % color + ','.join(ret)
def parseDistObj(obj):
if (obj[5][0][3][10] != 1): # 'show dashed' flag
return ""
N = obj[5][2][0][0]
points = obj[5][2][0][1]
ret = []
for p in points:
ret.append("%.3f" % p)
color = cmd.get_color_tuple(obj[5][0][2]);
return "\ndists:%.3f,%.3f,%.3f:" % color + ','.join(ret)
def _imageHasColor(self, color, img, delta=0):
if isinstance(color, str):
color = [int(v*255) for v in cmd.get_color_tuple(color)]
else:
color = list(color)
dim = img.shape[-1]
if dim == len(color) + 1:
dim -= 1
img = img[...,:dim]
diff = abs(img.reshape((-1, dim)) - color)
return (diff - delta <= 0).prod(1).sum()
def _imageHasColor(self, color, img, delta=0):
if isinstance(color, str):
color = [int(v * 255) for v in cmd.get_color_tuple(color)]
else:
color = list(color)
dim = img.shape[-1]
if dim == len(color) + 1:
dim -= 1
img = img[..., :dim]
diff = abs(img.reshape((-1, dim)) - color)
return (diff - delta <= 0).prod(1).sum()
def sele2Color_deprecated(self):
globalShader = self.optionMenu_shader.getvalue()
# (color:shader) dictionary
# be used in shaderFactory.SCString()
#self.spColorShaderDict = {}
# get all the selection name
sess = cmd.get_session()['names']
for i in sess:
if type(i) is ListType:
# for each selection in the scene
# also should exist in seleShaderDict
if i[0] in self.seleShaderDict:
newShader = self.seleShaderDict[i[0]]
if newShader != globalShader: # should be different from global shader, otherwise do nothing
newColorInc = ShaderFactory.seleSlot[newShader]
# pymol api routine
# get the color for all the atoms in current selection
stored.idcolor_list = []
cmd.iterate(i[0], 'stored.idcolor_list.append((ID, int(color)))')
# for a set of atoms have the same new color
# apply color all at once for better performance
# cmd.color is slow
# example result: (color_index: [atom ids])
# {25: [2, 3, 6], 6: [1], 9: [4, 5]}
colorSetDict = {}
for (atom_id, color) in stored.idcolor_list:
colorSetDict.setdefault(color, []).append(str(atom_id))
#print repr(colorSetDict)
# for each color
#for (atom_id, color) in stored.idcolor_list:
for color in colorSetDict:
atom_set = colorSetDict[color]
rgb_color = cmd.get_color_tuple(color)
#newColorInc = 0.001
if rgb_color[2] >= 0.990:
color_id = '%s %s %s' % (str(rgb_color[0])[0:5].ljust(5, '0'), str(rgb_color[1])[0:5].ljust(5, '0'), str(rgb_color[2]-newColorInc)[0:5])
else:
color_id = '%s %s %s' % (str(rgb_color[0])[0:5].ljust(5, '0'), str(rgb_color[1])[0:5].ljust(5, '0'), str(rgb_color[2]+newColorInc)[0:5])
#print atom_id, color, color_id, repr(rgb_color)
# color_id from all the sele:shader dictionary
# determine color_id:shader pair
if color_id not in self.spColorShaderDict:
self.spColorShaderDict[color_id] = [newShader, i[0]]
# apply new color to atom set
newRGB = color_id.split(' ')
newColor = [float(newRGB[0]), float(newRGB[1]), float(newRGB[2])]
cmd.set_color(color_id, newColor)
cmd.color(color_id, ('ID %s' % '+'.join(atom_set)))
print 'apply shader [%s] to selection [%s].' % (newShader, i[0])
def color_shift2(start="cterm",end="cterm2",divide=100):
stored.start_colors=[]
cmd.iterate("%s and name ca and chain A" % start,"stored.start_colors.append(color)")
stored.end_colors=[]
cmd.iterate("%s and name ca and chain A" % end,"stored.end_colors.append(color)")
start_color_indice=cmd.get_color_tuple(stored.start_colors[0])
end_color_indice=cmd.get_color_tuple(stored.end_colors[0])
graduate_color_indice=(np.array(end_color_indice)-np.array(start_color_indice))/float(divide)
i=0
while i<=divide:
cmd.refresh()
cmd.recolor()
cmd.set_color("temp%s" %i,list(np.array(start_color_indice)+graduate_color_indice*i))
#cmd.color("temp%s" % i,"%s and name ca and chain A and not (e;N,O)" % start)
cmd.color("temp%s" % i,"start%s2_end%s_%s and chain A and not (e;N,O)" % (start,end,i))
#cmd.set("sphere_color","temp%s" % i,"start%s2_end%s_%s and elem C" % (start,end,i))
cmd.set("stick_color","red","start%s2_end%s_%s and elem O" % (start,end,i))
cmd.set("stick_color","blue","start%s2_end%s_%s and elem N" % (start,end,i))
cmd.set("cartoon_color","red","start%s2_end%s_%s and chain C" % (start,end,i))
cmd.set("cartoon_color","tv_blue","start%s2_end%s_%s and chain B" % (start,end,i))
i+=1
def get_residue_colors(sele):
"""
Get RGB color values associated with a selection.
Useful if you want to exactly match coloring of 3D models
with coloring in, say, a MATLAB script.
"""
pymol.stored.colors = []
cmd.iterate(sele, "stored.colors.append( (chain, resi, name, color))")
res_colors = {}
for chain, resi, name, color in pymol.stored.colors:
if name == 'CA': # c-alpha atom
res_colors[(chain, resi)] = cmd.get_color_tuple(color)
print res_colors
return res_colors
def assertImageHasColor(self, color, img=None, delta=0):
if isinstance(color, str):
color = [int(v*255) for v in cmd.get_color_tuple(color)]
else:
color = list(color)
img = self.get_imagearray(img)
dim = img.shape[-1]
if dim == len(color) + 1:
color.append(255)
if isinstance(delta, list) and dim == len(delta) + 1:
delta.append(0)
diff = abs(img.reshape((-1, dim)) - color)
self.assertTrue((diff - delta <= 0).prod(1).sum(),
'no such color: ' + str(color))
def get_residue_colors( sele ):
"""
Get RGB color values associated with a selection.
Useful if you want to exactly match coloring of 3D models
with coloring in, say, a MATLAB script.
"""
pymol.stored.colors = []
cmd.iterate( sele, "stored.colors.append( (chain, resi, name, color))")
res_colors = {}
for chain, resi, name, color in pymol.stored.colors:
if name == 'CA': # c-alpha atom
res_colors[(chain, resi)] = cmd.get_color_tuple(color)
print res_colors
return res_colors
def color_shift(start="Apo2",end="cterm2",divide=100):
stored.start_colors=[]
cmd.iterate("%s and name ca and chain A" % start,"stored.start_colors.append(color)")
stored.end_colors=[]
cmd.iterate("%s and name ca and chain A" % end,"stored.end_colors.append(color)")
print stored.start_colors[0],stored.end_colors[0]
print cmd.get_color_tuple(stored.start_colors[0]),cmd.get_color_tuple(stored.end_colors[0])
start_color_indice=cmd.get_color_tuple(stored.start_colors[0])
end_color_indice=cmd.get_color_tuple(stored.end_colors[0])
print "start ,end"
print start_color_indice,end_color_indice
print np.array(end_color_indice)-np.array(start_color_indice)
#print (np.array(end_color_indice)-np.array(start_color_indice))/float(divide)
graduate_color_indice=(np.array(end_color_indice)-np.array(start_color_indice))/float(divide)
i=1
while i<divide:
i+=1
#print list(np.array(start_color_indice)+graduate_color_indice*i)
time.sleep(0.1)
cmd.refresh()
cmd.recolor()
cmd.set_color("temp%s" %i,list(np.array(start_color_indice)+graduate_color_indice*i))
cmd.color("temp%s" % i,"%s and name ca and chain A" % start)
def sphere(name, model_and_center_atom, radius, color, tr):
'''
DESCRIPTION
"sphere" creates a sphere with the given center-coordinates and radius
USAGE
sphere name, x_center, y_center, z_center, radius, color, tr
name = name of sphere
center_atom = center of sphere
radius = radius of sphere
color = color name
tr = transparency value
'''
color_rgb = cmd.get_color_tuple(cmd.get_color_index(color))
r = color_rgb[0]
g = color_rgb[1]
b = color_rgb[2]
str_list = []
#str_list.append(str(center_atom))
#res_str = str(center_atom)
#str_list.append(str(model))
#str_list.append(str("and resi"))
#str_list.append(str(res_str))
#str_list.append(str("and name Ca"))
sel_str = model_and_center_atom #string.join(str_list, ' ')
print sel_str
stored.xyz = []
#stored.xyz.append([x_center,y_center,z_center])
cmd.create("sphere", sel_str)
cmd.iterate_state(1,"sphere","stored.xyz.append([x,y,z])")
cmd.delete("sphere")
print stored.xyz
obj = []
obj.extend([cgo.ALPHA, tr])
obj.extend([
BEGIN, SPHERE,
COLOR, r, g, b,
SPHERE, stored.xyz[0][0], stored.xyz[0][1], stored.xyz[0][2], radius,
END
])
cmd.load_cgo(obj, name)
def construct_palette(color_list=None, max_value=7, min_value=0):
""" Construct a palette
Args:
color_list ([str]): list of PyMOL color strings (Default value = None)
max_value (int): max palette index (Default value = 7)
min_value (int): min palette index (Default value = 1)
Returns:
palette ([str]): list of color definitions
"""
if color_list is None:
color_list = ['tv_red', 'tv_orange', 'tv_yellow', 'tv_green', 'tv_blue', 'aquamarine', 'violet']
colors = []
for color in color_list:
if isinstance(color, str):
colors.append(cmd.get_color_tuple(color))
else:
colors.append(tuple(color))
output_range = max_value - min_value
palette = []
if output_range <= len(colors):
for color in colors[:max_value]:
palette.append('0x%02x%02x%02x' % tuple([int(255 * x) for x in color]))
elif (output_range > len(colors)) and (len(colors) > 1):
step = float(len(colors)) / float(output_range)
for i in range(output_range):
ix = float(i) * step
# get the indices of the surrounding colors correcting for floating point imprecision
lower_ind = max(int(math.floor(ix)), 0)
upper_ind = min(int(math.ceil(ix)), len(colors) - 1)
fx = ix - lower_ind
if lower_ind == upper_ind:
# special case where interpolation is exactly at an input color
palette.append('0x%02x%02x%02x' % tuple([int(255 * x) for x in colors[lower_ind]]))
else:
color = [fx * colors[lower_ind][i] + (1 - fx) * colors[upper_ind][i] for i in range(3)]
palette.append('0x%02x%02x%02x' % tuple([int(255 * x) for x in color]))
else:
logger.error("Palette overriden from default but only provided with one value for a multi-sphere object output. Either provide multiple colors to permit interpolation or leave at default.")
raise ValueError
logger.debug("Palette constructed with {0} colors".format(len(palette)))
return palette
def callback(resv, resn, color):
if stored.resv is None:
stored.resv = resv - (resv % 70)
if gapped:
while stored.resv+1 < resv:
callback(stored.resv+1, '-', 25)
stored.resv += 1
if stored.resv % 70 == 1:
html.append(('</font>\n<br>%4d <font>' % (resv)).replace(' ', ' '))
stored.color = None
c = cmd.get_color_tuple(color)
color = '#%02x%02x%02x' % tuple(int(0xFF * v) for v in c)
aa = one_letter.get(resn, '-')
if color != stored.color:
html.append('</font><font color="' + color + '">')
stored.color = color
html.append(aa)
def callback(resv, resn, color):
if stored.resv is None:
stored.resv = resv - (resv % 70)
if gapped:
while stored.resv+1 < resv:
callback(stored.resv+1, '-', 25)
stored.resv += 1
if stored.resv % 70 == 1:
html.append(('</font>\n<br>%4d <font>' % (resv)).replace(' ', ' '))
stored.color = None
c = cmd.get_color_tuple(color)
color = '#%02x%02x%02x' % (c[0]*255, c[1]*255, c[2]*255)
aa = one_letter.get(resn, '-')
if color != stored.color:
html.append('</font><font color="' + color + '">')
stored.color = color
html.append(aa)
def sphere(name, model_and_center_atom, radius, color, tr):
'''
DESCRIPTION
"sphere" creates a sphere with the given center-coordinates and radius
USAGE
sphere name, x_center, y_center, z_center, radius, color, tr
name = name of sphere
center_atom = center of sphere
radius = radius of sphere
color = color name
tr = transparency value
'''
color_rgb = cmd.get_color_tuple(cmd.get_color_index(color))
r = color_rgb[0]
g = color_rgb[1]
b = color_rgb[2]
str_list = []
#str_list.append(str(center_atom))
#res_str = str(center_atom)
#str_list.append(str(model))
#str_list.append(str("and resi"))
#str_list.append(str(res_str))
#str_list.append(str("and name Ca"))
sel_str = model_and_center_atom #string.join(str_list, ' ')
print sel_str
stored.xyz = []
#stored.xyz.append([x_center,y_center,z_center])
cmd.create("sphere", sel_str)
cmd.iterate_state(1, "sphere", "stored.xyz.append([x,y,z])")
cmd.delete("sphere")
print stored.xyz
obj = []
obj.extend([cgo.ALPHA, tr])
obj.extend([
BEGIN, SPHERE, COLOR, r, g, b, SPHERE, stored.xyz[0][0],
stored.xyz[0][1], stored.xyz[0][2], radius, END
])
cmd.load_cgo(obj, name)
def eval_color(v):
try:
if not v:
v=eval(cmd.get('bg_rgb'))
v=list(map(sum, list(zip(v,[-1,-1,-1]))))
v=list(map(abs, v))
if v[0]==v[1]==v[2]==0.5: # grey
v=[0,0,0]
return v
if isinstance(v, list):
return v[0:3]
if not isinstance(v, str):
return v[0:3]
if v.startswith('['):
return cmd.safe_list_eval(v)[0:3]
return list(cmd.get_color_tuple(v))
except:
return [random.random(),random.random(),random.random()]
def get_residue_colors( sele = "all", outfile = "colors.txt" ):
"""
Get RGB color values associated with a selection.
Useful if you want to exactly match coloring of 3D models
with coloring in, say, a MATLAB script.
"""
pymol.stored.colors = []
cmd.iterate( sele, "stored.colors.append( (chain, resi, name, color))")
res_colors = {}
stored_colors = pymol.stored.colors;
fid = open( outfile, 'w' )
for chain, resi, name, color in stored_colors:
if name == 'CA' or name == 'P': # c-alpha atom
color_tuple = cmd.get_color_tuple(color)
res_colors[chain+resi] = color_tuple
fid.write( '%s%s %f %f %f\n' % (chain,resi,color_tuple[0],color_tuple[1],color_tuple[2]) )
print "Outputted RGB colors to: ", outfile
return res_colors
def eval_color(v):
try:
if not v:
v = eval(cmd.get('bg_rgb'))
v = map(sum, zip(v, [-1, -1, -1]))
v = map(abs, v)
if v[0] == v[1] == v[2] == 0.5: # grey
v = [0, 0, 0]
return v
if isinstance(v, list):
return v[0:3]
if not isinstance(v, str):
return v[0:3]
if v.startswith('['):
return cmd.safe_list_eval(v)[0:3]
return list(cmd.get_color_tuple(v))
except:
return [random.random(), random.random(), random.random()]
def get_color(selection, which=0, mode=0):
'''
DESCRIPTION
API only. Returns the color of the first/middle/... guide atom in
selection.
ARGUMENTS
which = 0: color of first atom
which = 1: color of middle atom
which = 2: most frequent color
mode = 0: color index or color string
mode = 1: color tuple
mode = 2: color string in hash-hex format (for HTML, matplotlib, ...)
'''
s_first = 'first' if which == 0 else ''
try:
colors = []
for s_guide in ('guide', 'elem C', 'all'):
cmd.iterate('{} (({}) & {})'.format(s_first, selection, s_guide),
'colors.append(color)',
space=locals())
if colors:
break
if which == 2:
color = max((colors.count(color), color) for color in colors)[1]
else:
color = colors[len(colors) // 2]
if color >= 0x40000000:
color = '0x%06x' % (color & 0xFFFFFF)
except:
print(' Warning: could not get color for ' + str(selection))
color = 'gray'
if mode > 0:
color = cmd.get_color_tuple(color)
if mode == 2:
return '#%02x%02x%02x' % tuple(int(0xFF * v) for v in color)
return color
def get_color(selection, which=0, mode=0):
'''
DESCRIPTION
API only. Returns the color of the first/middle/... guide atom in
selection.
ARGUMENTS
which = 0: color of first atom
which = 1: color of middle atom
which = 2: most frequent color
mode = 0: color index or color string
mode = 1: color tuple
mode = 2: color string in hash-hex format (for HTML, matplotlib, ...)
'''
s_first = 'first' if which == 0 else ''
try:
colors = []
for s_guide in ('guide', 'elem C', 'all'):
cmd.iterate('{} (({}) & {})'.format(s_first, selection, s_guide),
'colors.append(color)', space=locals())
if colors:
break
if which == 2:
color = max((colors.count(color), color) for color in colors)[1]
else:
color = colors[len(colors) // 2]
if color >= 0x40000000:
color = '0x%06x' % (color & 0xFFFFFF)
except:
print(' Warning: could not get color for ' + str(selection))
color = 'gray'
if mode > 0:
color = cmd.get_color_tuple(color)
if mode == 2:
return '#%02x%02x%02x' % tuple(int(0xFF * v) for v in color)
return color
def plot(self, xp, yp, meta):
# Convert from 'label' space to 'pixel' space
x = self.convertToPixel("X", xp)
y = self.convertToPixel("Y", yp)
resn, color, ss = self.idx2resn.get(meta)
if self.symbols == 0:
# symbols by amino acid (G/P/other)
mark = {'GLY': 'Tri', 'PRO': 'Rect'}.get(resn, 'Oval')
else:
# symbols by secondary structure
mark = {'H': 'Oval', 'S': 'Rect'}.get(ss, 'Tri')
if mark == 'Oval':
create_shape = self.create_oval
coords = [
x - self.mark_size, y - self.mark_size, x + self.mark_size,
y + self.mark_size
]
elif mark == 'Tri':
create_shape = self.create_polygon
coords = [
x, y - self.mark_size, x + self.mark_size, y + self.mark_size,
x - self.mark_size, y + self.mark_size
]
else:
create_shape = self.create_rectangle
coords = [
x - self.mark_size, y - self.mark_size, x + self.mark_size,
y + self.mark_size
]
if color >= 0x40000000:
color = '#%06x' % (color & 0xffffff)
else:
color = '#%02x%02x%02x' % tuple(
[255 * i for i in cmd.get_color_tuple(color)])
oval = create_shape(width=1, outline="black", fill=color, *coords)
self.shapes[oval] = [x, y, 0, xp, yp, meta]
def get_color(selection, which=0, mode=0):
'''
DESCRIPTION
API only. Returns the color of the first/middle/... guide atom in
selection.
ARGUMENTS
which = 0: color of first atom
which = 1: color of middle atom
which = 2: most frequent color
mode = 0: color index or color string
mode = 1: color tuple
mode = 2: color string in hash-hex format (for HTML, matplotlib, ...)
'''
try:
colors = []
if which == 0:
cmd.iterate('first ((%s) and guide)' % (selection),
'colors.append(color)',
space=locals())
color = colors[0]
else:
cmd.iterate('(%s) and guide' % (selection),
'colors.append(color)',
space=locals())
if which == 1:
color = colors[len(colors) / 2]
else:
color = max((colors.count(color), color) for color in colors)
if color >= 0x40000000:
color = '0x%06x' % (color & 0xFFFFFF)
except:
print(' Warning: could not get color for ' + str(selection))
color = 'gray'
if mode > 0:
color = cmd.get_color_tuple(color)
if mode == 2:
return '#%02x%02x%02x' % tuple(255 * i for i in color)
return color
def hexagon(obj, coords, colors, rep, radius):
"""draw the rings"""
for idx, coord in enumerate(coords):
r, g, b = cmd.get_color_tuple(colors[idx])
coord.append(coord[0])
coord = np.array(coord)
mean = np.mean(coord[:-1], axis=0)
x3, y3, z3 = mean
# average the normals of two neighbouring triangles
normals = [0.5 * (
np.cross((coord[i] - coord[i - 1]), (mean - coord[i])) +
np.cross((coord[i + 1] - coord[i]), (mean - coord[i])))
for i in range(len(coord) - 1)]
centernormal = np.mean(normals, axis=0).tolist()
# add first value to be able to cycle trought the triangles
normals.append(normals[0])
for i in range(len(coord) - 1):
x1, y1, z1 = coord[i]
x2, y2, z2 = coord[i + 1]
# Triangles
if rep == 'cartoon':
tri = [
BEGIN, TRIANGLES,
COLOR, r, g, b,
NORMAL, normals[i][0], normals[i][1], normals[i][2],
NORMAL, normals[i + 1][0], normals[i + 1][1], normals[i + 1][2],
NORMAL, centernormal[0], centernormal[1], centernormal[2],
VERTEX, x1, y1, z1,
VERTEX, x2, y2, z2,
VERTEX, x3, y3, z3,
END
]
obj.extend(tri)
obj.extend([CYLINDER, x1, y1, z1, x2, y2,
z2, radius, r, g, b, r, g, b])
obj.extend([COLOR, r, g, b, SPHERE, x1, y1, z1, radius])
return obj
def get_color(selection, which=0, mode=0):
'''
DESCRIPTION
API only. Returns the color of the first/middle/... guide atom in
selection.
ARGUMENTS
which = 0: color of first atom
which = 1: color of middle atom
which = 2: most frequent color
mode = 0: color index or color string
mode = 1: color tuple
mode = 2: color string in hash-hex format (for HTML, matplotlib, ...)
'''
try:
colors = []
if which == 0:
cmd.iterate('first ((%s) and guide)' % (selection),
'colors.append(color)', space=locals())
color = colors[0]
else:
cmd.iterate('(%s) and guide' % (selection),
'colors.append(color)', space=locals())
if which == 1:
color = colors[len(colors)/2]
else:
color = max((colors.count(color), color) for color in colors)
if color >= 0x40000000:
color = '0x%06x' % (color & 0xFFFFFF)
except:
print(' Warning: could not get color for ' + str(selection))
color = 'gray'
if mode > 0:
color = cmd.get_color_tuple(color)
if mode == 2:
return '#%02x%02x%02x' % tuple(255 * i for i in color)
return color
def hexagon(obj, coords, colors, rep, radius):
"""draw the rings"""
for idx, coord in enumerate(coords):
r, g, b = cmd.get_color_tuple(colors[idx])
coord.append(coord[0])
coord = np.array(coord)
mean = np.mean(coord[:-1], axis=0)
x3, y3, z3 = mean
# average the normals of two neighbouring triangles
normals = [0.5 * (
np.cross((coord[i] - coord[i - 1]), (mean - coord[i])) +
np.cross((coord[i + 1] - coord[i]), (mean - coord[i])))
for i in range(len(coord) - 1)]
centernormal = np.mean(normals, axis=0).tolist()
# add first value to be able to cycle trought the triangles
normals.append(normals[0])
for i in range(len(coord)-1):
x1, y1, z1 = coord[i]
x2, y2, z2 = coord[i+1]
# Triangles
if rep == 'cartoon':
tri = [
BEGIN, TRIANGLES,
COLOR, r, g, b,
NORMAL, normals[i][0], normals[i][1], normals[i][2],
NORMAL, normals[i+1][0], normals[i+1][1], normals[i+1][2],
NORMAL, centernormal[0], centernormal[1], centernormal[2],
VERTEX, x1, y1, z1,
VERTEX, x2, y2, z2,
VERTEX, x3, y3, z3,
END
]
obj.extend(tri)
obj.extend([CYLINDER, x1, y1, z1, x2, y2, z2, radius, r, g, b, r, g, b])
obj.extend([COLOR, r, g, b, SPHERE, x1, y1, z1, radius])
return obj
def plot(self,xp,yp,meta):
# Convert from 'label' space to 'pixel' space
x = self.convertToPixel("X",xp)
y = self.convertToPixel("Y",yp)
resn, color, ss = self.idx2resn.get(meta)
if self.symbols == 0:
# symbols by amino acid (G/P/other)
mark = {'GLY': 'Tri', 'PRO': 'Rect'}.get(resn, 'Oval')
else:
# symbols by secondary structure
mark = {'H': 'Oval', 'S': 'Rect'}.get(ss, 'Tri')
if mark == 'Oval':
create_shape = self.create_oval
coords = [x-self.mark_size, y-self.mark_size,
x+self.mark_size, y+self.mark_size]
elif mark == 'Tri':
create_shape = self.create_polygon
coords = [x, y-self.mark_size,
x+self.mark_size, y+self.mark_size,
x-self.mark_size, y+self.mark_size]
else:
create_shape = self.create_rectangle
coords = [x-self.mark_size, y-self.mark_size,
x+self.mark_size, y+self.mark_size]
if color >= 0x40000000:
color = '#%06x' % (color & 0xffffff)
else:
color = '#%02x%02x%02x' % tuple([255*i
for i in cmd.get_color_tuple(color)])
oval = create_shape(*coords,
width=1, outline="black", fill=color)
self.shapes[oval] = [x,y,0,xp,yp,meta]
def construct_palette(color_list=None, max_value=7, min_value=1):
""" Construct a palette
Args:
color_list ([str]): list of PyMOL color strings (Default value = None)
max_value (int): max palette index (Default value = 7)
min_value (int): min palette index (Default value = 1)
Returns:
palette ([str]): list of color definitions
"""
if color_list is None:
color_list = [
'tv_red', 'tv_orange', 'tv_yellow', 'tv_green', 'tv_blue',
'aquamarine', 'violet'
]
if max_value <= len(color_list):
return color_list
colors = [cmd.get_color_tuple(x) for x in color_list]
output_range = max_value - min_value + 1
palette = []
color_vectors = len(colors) - 1
steps = output_range / color_vectors
for cv in range(color_vectors):
for x in range(steps):
fx = float(x) / steps
cl = [
fx * colors[cv + 1][i] + (1 - fx) * colors[cv][i]
for i in range(3)
]
cn = '0x%02x%02x%02x' % tuple([255 * x for x in cl])
palette.append(cn)
logger.debug("Palette constructed with {0} colors".format(len(palette)))
return palette
def SetPartitionColor(self, Sel):
try:
mycolors = {'colors': []}
cmd.iterate(Sel, 'colors.append(color)', space=mycolors)
for color in mycolors['colors']:
one = list(cmd.get_color_tuple(color))
break
for i in range(0, 3):
if one[i] <= 0.80:
one[i] += 0.20
else:
one[i] -= 0.20
cmd.set_color(self.PartitionColor, one)
partition_rgb = General.one_to_rgb(one)
except:
return tuple()
return tuple(partition_rgb)
def SetPartitionColor(self, Sel):
try:
mycolors = {'colors': []}
cmd.iterate( Sel, 'colors.append(color)', space=mycolors)
for color in mycolors['colors']:
one = list( cmd.get_color_tuple(color) )
break
for i in range(0,3):
if one[i] <= 0.80:
one[i] += 0.20
else:
one[i] -= 0.20
cmd.set_color(self.PartitionColor, one)
partition_rgb = General.one_to_rgb(one)
except:
return tuple()
return tuple(partition_rgb)
def cgo_arrow(atom1='pk1',
atom2='pk2',
radius=0.5,
gap=0.0,
hlength=-1,
hradius=-1,
color='blue red',
name=''):
'''
DESCRIPTION
Create a CGO arrow between two picked atoms.
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: single atom selection or list of 3 floats {default: pk2}
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
#obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
#[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
#[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def testGetColorTuple(self):
self.assertEqual((0.0, 0.0, 1.0), cmd.get_color_tuple("blue"))
def spectrumany(expression, color_list, selection='(all)', minimum=None, maximum=None, quiet=1):
'''
DESCRIPTION
Define a color spectrum with as many color-stops as you like (at least 2).
USAGE
spectrumany expression, color_list [, selection [, minimum [, maximum ]]]
ARGUMENTS
expression = count, resi, b, q, or pc: respectively, atom count, residue
index, temperature factor, occupancy, or partial charge {default: count}
color_list = string: Space separated list of colors
... all other arguments like with `spectrum` command
EXAMPLE
spectrumany count, forest green yellow white
spectrumany b, red yellow white, (polymer), maximum=100.0
SEE ALSO
spectrum
'''
quiet = int(quiet)
colors = color_list.split()
if len(colors) < 2:
print 'failed! please provide at least 2 colors'
return
colvec = [cmd.get_color_tuple(i) for i in colors]
parts = len(colvec) - 1
expression = {'pc': 'partial_charge', 'fc': 'formal_charge',
'count': 'index'}.get(expression, expression)
minmax_expr = {'resi': 'resv'}.get(expression, expression)
discrete_expr = ['index', 'resi']
if cmd.count_atoms(selection) == 0:
print 'empty selection'
return
if None in [minimum, maximum]:
stored.e = list()
cmd.iterate(selection, 'stored.e.append(%s)' % (minmax_expr))
if minimum is None:
minimum = min(stored.e)
if maximum is None:
maximum = max(stored.e)
minimum, maximum = float(minimum), float(maximum)
if not quiet:
print ' Spectrum: range (%.5f to %.5f)' % (minimum, maximum)
if maximum == minimum:
print 'no spectrum possible, only equal %s values' % (expression)
return
if expression in discrete_expr:
val_range = int(maximum - minimum + 1)
else:
val_range = maximum - minimum
cmd.color(colors[0], selection)
steps = 60 / parts
steps_total = steps * parts
val_start = minimum
for p in range(parts):
for i in range(steps):
ii = float(i)/steps
col_list = [colvec[p+1][j] * ii + colvec[p][j] * (1.0 - ii) for j in range(3)]
col_name = '0x%02x%02x%02x' % (col_list[0] * 255, col_list[1] * 255, col_list[2] * 255)
val_end = val_range * (i + 1 + p * steps) / steps_total + minimum
if expression in discrete_expr:
cmd.color(col_name, '(%s) and %s %d-%d' % (selection, expression, val_start, val_end))
else:
cmd.color(col_name, '(%s) and %s > %f' % (selection, expression, val_start))
val_start = val_end
def spectrum_states(selection='all', representations='cartoon ribbon',
color_list='blue cyan green yellow orange red',
first=1, last=0, quiet=1):
'''
DESCRIPTION
Color each state in a multi-state object different.
(c) 2011 Takanori Nakane and Thomas Holder
USAGE
spectrum_states [ selection [, representations [, color_list [, first [, last ]]]]]
ARGUMENTS
selection = string: object names (works with complete objects only)
{default: all}
representations = string: space separated list of representations
{default: cartoon ribbon}
color_list = string: space separated list of colors {default: blue cyan
green yellow orange red}
SEE ALSO
spectrum, spectrumany
'''
from math import floor, ceil
first, last, quiet = int(first), int(last), int(quiet)
colors = color_list.split()
if len(colors) < 2:
print(' Error: please provide at least 2 colors')
raise CmdException
colvec = [cmd.get_color_tuple(i) for i in colors]
# filter for valid <repr>_color settings
settings = []
for r in representations.split():
if r[-1] == 's':
r = r[:-1]
s = r + '_color'
if s in cmd.setting.name_list:
settings.append(s)
elif not quiet:
print(' Warning: no such setting: ' + repr(s))
# object names only
selection = ' '.join(cmd.get_object_list('(' + selection + ')'))
if cmd.count_atoms(selection) == 0:
print(' Error: empty selection')
raise CmdException
if last < 1:
last = cmd.count_states(selection)
val_range = int(last - first + 1)
if val_range < 2:
print(' Error: no spectrum possible, need more than 1 state')
raise CmdException
for i in range(val_range):
p = float(i) / (val_range - 1) * (len(colvec) - 1)
p0, p1 = int(floor(p)), int(ceil(p))
ii = (p - p0)
col_list = [colvec[p1][j] * ii + colvec[p0][j] * (1.0 - ii) for j in range(3)]
col_name = '0x%02x%02x%02x' % tuple(int(0xFF * v) for v in col_list)
for s in settings:
cmd.set(s, col_name, selection, state=i+first)
def spectrumany(expression, color_list, selection='(all)', minimum=None, maximum=None, quiet=1):
'''
DESCRIPTION
Define a color spectrum with as many color-stops as you like (at least 2).
USAGE
spectrumany expression, color_list [, selection [, minimum [, maximum ]]]
ARGUMENTS
expression = count, resi, b, q, or pc: respectively, atom count, residue
index, temperature factor, occupancy, or partial charge {default: count}
color_list = string: Space separated list of colors
... all other arguments like with `spectrum` command
EXAMPLE
spectrumany count, forest green yellow white
spectrumany b, red yellow white, (polymer), maximum=100.0
SEE ALSO
spectrum
'''
quiet = int(quiet)
colors = color_list.split()
if len(colors) < 2:
print('failed! please provide at least 2 colors')
return
colvec = [cmd.get_color_tuple(i) for i in colors]
parts = len(colvec) - 1
expression = {'pc': 'partial_charge', 'fc': 'formal_charge',
'count': 'index'}.get(expression, expression)
minmax_expr = {'resi': 'resv'}.get(expression, expression)
discrete_expr = ['index', 'resi']
if cmd.count_atoms(selection) == 0:
print('empty selection')
return
if None in [minimum, maximum]:
e_list = list()
cmd.iterate(selection, 'e_list.append(%s)' % (minmax_expr), space=locals())
if minimum is None:
minimum = min(e_list)
if maximum is None:
maximum = max(e_list)
minimum, maximum = float(minimum), float(maximum)
if not quiet:
print(' Spectrum: range (%.5f to %.5f)' % (minimum, maximum))
if maximum == minimum:
print('no spectrum possible, only equal %s values' % (expression))
return
if expression in discrete_expr:
val_range = int(maximum - minimum + 1)
else:
val_range = maximum - minimum
cmd.color(colors[0], selection)
steps = 60 / parts
steps_total = steps * parts
val_start = minimum
for p in range(parts):
for i in range(steps):
ii = float(i)/steps
col_list = [colvec[p+1][j] * ii + colvec[p][j] * (1.0 - ii) for j in range(3)]
col_name = '0x%02x%02x%02x' % tuple(int(0xFF * v) for v in col_list)
val_end = val_range * (i + 1 + p * steps) / steps_total + minimum
if expression in discrete_expr:
cmd.color(col_name, '(%s) and %s %d-%d' % (selection, expression, val_start, val_end))
else:
cmd.color(col_name, '(%s) and %s > %f' % (selection, expression, val_start))
val_start = val_end
colors = [
'red',
'blue',
'green',
'yellow',
'magenta',
'cyan',
'orange',
'marine',
'chartreuse',
'purpleblue',
'violet',
'limon',
]
colors_value = [
tuple(int(i * 255) for i in cmd.get_color_tuple(color)) for color in colors
]
def __init__(self):
self.menuBar.addmenuitem('Plugin',
'command',
'Contact Map Visualizer',
label='Contact Map Visualizer',
command=lambda s=self: CMVDialog(s))
def CMVDialog(self):
import tkFileDialog, tkMessageBox
try:
def parseObjMol(obj):
name = obj[0]
ids = []
sphere = []
trace = []
ribbon = []
stick = []
surface = []
line = []
cross = []
smallSphere = []
helix = []
sheet = []
colors = {}
for atom in obj[5][7]:
rep = atom[20] + [0] * 12
serial = atom[22]
ss = atom[10]
bonded = (atom[25] == 1)
if (rep[5] == 1):
ribbon.append(serial)
if (rep[1] == 1):
sphere.append(serial)
if (rep[2] == 1):
surface.append(serial)
if (rep[7] == 1):
line.append(serial)
if (rep[6] == 1):
trace.append(serial)
if (rep[4] == 1 and not bonded):
smallSphere.append(serial)
if (rep[11] == 1 and not bonded):
cross.append(serial)
if (rep[0] == 1 and bonded):
stick.append(serial)
if (ss == 'S'):
sheet.append(serial)
if (ss == 'H'):
helix.append(serial)
c = cmd.get_color_tuple(atom[21])
if (not c in colors):
colors[c] = []
colors[c].append(serial)
ids.append("ID %d is %s in resi %s %s at chain %s"\
% (atom[22], atom[6], atom[3], atom[5], atom[1]))
for c in colors.keys(): # TODO: better compression
colors[c] = compactSeq(colors[c])
ret = ''
ret += "\nsheet:" + compactSeq(sheet)
ret += "\nhelix:" + compactSeq(helix)
ret += "\nsurface:" + compactSeq(surface)
ret += "\nsphere:" + compactSeq(sphere)
ret += "\ntrace:" + compactSeq(trace)
ret += "\nribbon:" + compactSeq(ribbon)
ret += "\nstick:" + compactSeq(stick)
ret += "\nline:" + compactSeq(line)
ret += "\nsmallSphere:" + compactSeq(smallSphere)
ret += "\ncross:" + compactSeq(cross)
for c in colors.keys():
ret += "\ncolor:%.3f,%.3f,%.3f:%s" % (c[0], c[1], c[2], colors[c])
return ret
def get_color_rgb(color):
try:
return cmd.get_color_tuple(color)
except:
print('Could not find a reference for that color')
return None
def edge(name,
i_node,
j_node,
color=None,
r=1.0,
g=0.0,
b=0.0,
dg=0.3,
dl=0.5,
dr=0.2,
dir=1,
dir_color=None,
dir_r=0.0,
dir_g=1.0,
dir_b=0.0):
'''
DESCRIPTION
"edge" creates a cylinder (actually sausage) between the two
selections that correspond to the 2 nodes. If the edge is
directed, only half of the user-formatted cylinder will be
drawn towards the target node n2 and the rest will be drawn as
a thin cylinder.
USAGE
edge name, i_node, j_node [, color, r, g, b, dg, dl, dr, dir,
dir_color, dir_r, dir_g, dir_b]
name = name of edge
i_node, j_node = atom selections for node 1 and node 2
color = color name (overwrites rgb)
r, g, b = rgb color (default red)
dg = dash gap (default 0 - alternative 0.3)
dl = dash length (default 0.5)
dr = dash radius (default 0.2)
dir = directed edge (default 1-yes)
dir_color = color name for the other half (overwrites dir_rgb)
dir_[r, g, b] = rgb color for the other half (default green)
'''
if color is not None:
color_rgb = cmd.get_color_tuple(cmd.get_color_index(color))
r = color_rgb[0]
g = color_rgb[1]
b = color_rgb[2]
else:
# Convert arguments into floating point values
r = float(r)
g = float(g)
b = float(b)
if dir_color is not None:
dir_color_rgb = cmd.get_color_tuple(cmd.get_color_index(dir_color))
dir_r = dir_color_rgb[0]
dir_g = dir_color_rgb[1]
dir_b = dir_color_rgb[2]
else:
dir_r = float(dir_r)
dir_g = float(dir_g)
dir_b = float(dir_b)
dg = float(dg)
dl = float(dl)
dr = float(dr)
directed = int(dir)
frag = directed + 1
# Get tuple containing object and index of atoms in these
# selections
x1 = cmd.index(i_node, 1)
x2 = cmd.index(j_node, 1)
# Get number of atoms in each selection
n1 = len(x1)
n2 = len(x2)
if (n1 < 1):
print "Error: node " + n1 + " has no atoms"
return
if (n2 < 1):
print "Error: node " + n2 + " has no atoms"
return
# Get objects and atom indices
o1 = x1[0][0]
i1 = x1[0][1]
o2 = x2[0][0]
i2 = x2[0][1]
# Get ChemPy models
m1 = cmd.get_model(o1)
m2 = cmd.get_model(o2)
# Get atoms
a1 = m1.atom[i1 - 1]
a2 = m2.atom[i2 - 1]
# Get coords
x1 = a1.coord[0]
y1 = a1.coord[1]
z1 = a1.coord[2]
x2 = a2.coord[0]
y2 = a2.coord[1]
z2 = a2.coord[2]
# Make some nice strings for user feedback
#n1 = o1 + "/" + a1.segi + "/" + a1.chain + "/" + a1.resn + "." + a1.resi + "/" + a1.name
#print n1 + "(" + str(x1) + "," + str(y1) + "," + str(z1) + ")"
#n2 = o2 + "/" + a2.segi + "/" + "/" + a2.chain + "/" + a2.resn + "." + a2.resi + "/" + a2.name
#print n2 + "(" + str(x2) + "," + str(y2) + "," + str(z2) + ")"
# Calculate distances
dx = (x2 - x1) / frag
dy = (y2 - y1) / frag
dz = (z2 - z1) / frag
d = math.sqrt((dx * dx) + (dy * dy) + (dz * dz))
#print "distance = " + str(d) + "A"
# Work out how many times (dash_len + gap_len) fits into d
dash_tot = dl + dg
n_dash = math.floor(d / dash_tot)
# Work out step lengths
dx1 = (dl / dash_tot) * (dx / n_dash)
dy1 = (dl / dash_tot) * (dy / n_dash)
dz1 = (dl / dash_tot) * (dz / n_dash)
dx2 = (dx / n_dash)
dy2 = (dy / n_dash)
dz2 = (dz / n_dash)
# Generate dashes
x = x1
y = y1
z = z1
# Empty CGO object
obj = []
for i in range(n_dash):
# Generate a sausage
obj.extend([
SAUSAGE, x, y, z, x + dx1, y + dy1, z + dz1, dr, r, g, b, r, g, b
])
# Move to start of next dash
x = x + dx2
y = y + dy2
z = z + dz2
if directed == 1:
obj.extend([
SAUSAGE, x, y, z, x2, y2, z2, 0.05, dir_r, dir_g, dir_b, dir_r,
dir_g, dir_b
])
cmd.set("stick_quality", 24)
# Load the object into PyMOL
cmd.load_cgo(obj, name)
def triangle(name, i_node, j_node, k_node, contact_type, color, tr):
'''
DESCRIPTION
"triangle" creates a triangle with the given nodes
USAGE
triangle name, i_node, j_node, k_node, color, tr
name = name of triangle
i_node, j_node, k_node = the residue numbers
color = color name
tr = transparency value
'''
color_rgb = cmd.get_color_tuple(cmd.get_color_index(color))
r = color_rgb[0]
g = color_rgb[1]
b = color_rgb[2]
str_list = []
str_list.append(str(i_node))
str_list.append(str(j_node))
str_list.append(str(k_node))
res_str = string.join(str_list, '+')
str_list[0] = "resi"
str_list[1] = res_str
str_list[2] = "and name"
str_list.append(str(contact_type))
sel_str = string.join(str_list, ' ')
#print sel_str
stored.xyz = []
cmd.create("triangle", sel_str)
cmd.iterate_state(1, "triangle", "stored.xyz.append([x,y,z])")
cmd.delete("triangle")
#print stored.xyz
#1st way doesn't work
normalx = ((stored.xyz[1][1] - stored.xyz[0][1]) *
(stored.xyz[2][2] - stored.xyz[0][2])) - (
(stored.xyz[2][1] - stored.xyz[0][1]) *
(stored.xyz[1][2] - stored.xyz[0][2]))
normaly = ((stored.xyz[1][2] - stored.xyz[0][2]) *
(stored.xyz[2][0] - stored.xyz[0][0])) - (
(stored.xyz[2][2] - stored.xyz[0][2]) *
(stored.xyz[1][0] - stored.xyz[0][0]))
normalz = ((stored.xyz[1][0] - stored.xyz[0][0]) *
(stored.xyz[2][1] - stored.xyz[0][1])) - (
(stored.xyz[2][0] - stored.xyz[0][0]) *
(stored.xyz[1][1] - stored.xyz[0][1]))
obj = [
BEGIN, TRIANGLES, stored.xyz[0][0], stored.xyz[0][1], stored.xyz[0][2],
stored.xyz[1][0], stored.xyz[1][1], stored.xyz[1][2], stored.xyz[2][0],
stored.xyz[2][1], stored.xyz[2][2], normalx - stored.xyz[0][0],
normaly - stored.xyz[0][1], normalz - stored.xyz[0][2],
normalx - stored.xyz[1][0], normaly - stored.xyz[1][1],
normalz - stored.xyz[1][2], normalx - stored.xyz[2][0],
normaly - stored.xyz[2][1], normalz - stored.xyz[2][2], 1.0, 1.0, 1.0,
1.0, 1.0, 1.0, 1.0, 1.0, 1.0, END
]
#2nd way
obj = []
obj.extend([cgo.ALPHA, tr])
obj.extend([
BEGIN, TRIANGLES, COLOR, r, g, b, VERTEX, stored.xyz[0][0],
stored.xyz[0][1], stored.xyz[0][2], VERTEX, stored.xyz[1][0],
stored.xyz[1][1], stored.xyz[1][2], VERTEX, stored.xyz[2][0],
stored.xyz[2][1], stored.xyz[2][2], END
])
cmd.load_cgo(obj, name)
def torus(center=(0., 0., 0.), normal=(0., 0., 1.), radius=1., color='',
cradius=.25, samples=20, csamples=20):
'''
Generate and return a torus CGO with given center, normal
and ring radius.
'''
from math import cos, sin, pi
if color and isinstance(color, str):
color = list(cmd.get_color_tuple(color))
obj = []
axis = cpv.cross_product(normal, (0., 0., 1.))
angle = -cpv.get_angle(normal, (0., 0., 1.))
matrix = cpv.rotation_matrix(angle, cpv.normalize(axis))
obj_vertex = lambda x, y, z: obj.extend([VERTEX] + cpv.add(center,
cpv.transform(matrix, [x, y, z])))
obj_normal = lambda x, y, z: obj.extend([NORMAL] +
cpv.transform(matrix, [x, y, z]))
r = radius
cr = cradius
rr = 1.5 * cr
dv = 2 * pi / csamples
dw = 2 * pi / samples
v = 0.0
w = 0.0
while w < 2 * pi:
v = 0.0
c_w = cos(w)
s_w = sin(w)
c_wdw = cos(w + dw)
s_wdw = sin(w + dw)
obj.append(BEGIN)
obj.append(TRIANGLE_STRIP)
if color:
obj.append(COLOR)
obj.extend(color)
while v < 2 * pi + dv:
c_v = cos(v)
s_v = sin(v)
c_vdv = cos(v + dv)
s_vdv = sin(v + dv)
obj_normal(
(r + rr * c_v) * c_w - (r + cr * c_v) * c_w,
(r + rr * c_v) * s_w - (r + cr * c_v) * s_w,
(rr * s_v - cr * s_v))
obj_vertex(
(r + cr * c_v) * c_w,
(r + cr * c_v) * s_w,
cr * s_v)
obj_normal(
(r + rr * c_vdv) * c_wdw - (r + cr * c_vdv) * c_wdw,
(r + rr * c_vdv) * s_wdw - (r + cr * c_vdv) * s_wdw,
rr * s_vdv - cr * s_vdv)
obj_vertex(
(r + cr * c_vdv) * c_wdw,
(r + cr * c_vdv) * s_wdw,
cr * s_vdv)
v += dv
obj.append(END)
w += dw
return obj
def dump_rep(name):
if 'PYMOL_GIT_MOD' in os.environ:
import shutil
try:
shutil.copytree(
os.path.join(os.environ['PYMOL_GIT_MOD'], 'pymol2glmol', 'js'),
os.path.join(os.getcwd(), 'js'))
except OSError:
pass
try:
cmd.set('pse_export_version', 1.74)
except:
pass
names = cmd.get_session()['names']
cmd.set('pdb_retain_ids', 1)
ret = ''
for obj in names:
if (obj == None):
continue
if (obj[2] == 0): # not visible
continue
if (obj[1] == 0 and obj[4] == 1 and obj[0] == name):
ret += parseObjMol(obj)
if (obj[1] == 0
and obj[4] == 4): # currently all dist objects are exported
ret += parseDistObj(obj)
cmd.turn('z', 180)
view = cmd.get_view()
cmd.turn('z', 180)
cx = -view[12]
cy = -view[13]
cz = -view[14]
cameraZ = -view[11] - 150
fov = float(cmd.get("field_of_view"))
fogStart = float(cmd.get("fog_start"))
slabNear = view[15] + view[11]
slabFar = view[16] + view[11]
ret += "\nview:%.3f,%.3f,%.3f,%.3f,%.3f,%.3f,%.3f,%.3f" % \
(cx, cy, cz, cameraZ, slabNear, slabFar, fogStart, fov)
for i in range(9):
ret += ",%.3f" % view[i]
bgcolor = cmd.get_setting_tuple('bg_rgb')[1]
if len(bgcolor) == 1:
bgcolor = cmd.get_color_tuple(bgcolor[0])
ret += "\nbgcolor:%02x%02x%02x" % (int(255 * float(bgcolor[0])), \
int(255 * float(bgcolor[1])), int(255 * float(bgcolor[2])))
if 'PYMOL_GIT_MOD' in os.environ:
template = open(os.path.join(os.environ['PYMOL_GIT_MOD'], 'pymol2glmol', 'imported.html')).read().\
replace("###INCLUDE_PDB_FILE_HERE###", cmd.get_pdbstr(name)).\
replace('###INCLUDE_REPRESENTATION_HERE###', ret)
else:
template = open('imported.html').read().\
replace("###INCLUDE_PDB_FILE_HERE###", cmd.get_pdbstr(name)).\
replace('###INCLUDE_REPRESENTATION_HERE###', ret)
f = open(name + '.html', 'w')
f.write(template)
f.close()
def cgo_arrow(atom1='pk1', atom2='pk2', radius=0.5, gap=0.0, hlength=-1, hradius=-1,
color='black black', name=''):
'''
#I modify the color the line just before
DESCRIPTION
Create a CGO arrow between two picked atoms.
ARGUMENTS
atom1 = string: single atom selection or list of 3 floats {default: pk1}
atom2 = string: single atom selection or list of 3 floats {default: pk2}
radius = float: arrow radius {default: 0.5}
gap = float: gap between arrow tips and the two atoms {default: 0.0}
hlength = float: length of head
hradius = float: radius of head
color = string: one or two color names {default: blue red}
name = string: name of CGO object
'''
from chempy import cpv
radius, gap = float(radius), float(gap)
hlength, hradius = float(hlength), float(hradius)
try:
color1, color2 = color.split()
except:
color1 = color2 = color
color1 = list(cmd.get_color_tuple(color1))
color2 = list(cmd.get_color_tuple(color2))
def get_coord(v):
if not isinstance(v, str):
return v
if v.startswith('['):
return cmd.safe_list_eval(v)
return cmd.get_atom_coords(v)
xyz1 = get_coord(atom1)
xyz2 = get_coord(atom2)
normal = cpv.normalize(cpv.sub(xyz1, xyz2))
if hlength < 0:
hlength = radius * 3.0
if hradius < 0:
hradius = hlength * 0.6
if gap:
diff = cpv.scale(normal, gap)
xyz1 = cpv.sub(xyz1, diff)
xyz2 = cpv.add(xyz2, diff)
xyz3 = cpv.add(cpv.scale(normal, hlength), xyz2)
obj = [cgo.CYLINDER] + xyz1 + xyz3 + [radius] + color1 + color2 + \
[cgo.CONE] + xyz3 + xyz2 + [hradius, 0.0] + color2 + color2 + \
[1.0, 0.0]
if not name:
name = cmd.get_unused_name('arrow')
cmd.load_cgo(obj, name)
def construct_palette(color_list=None, max_value=7, min_value=0):
""" Construct a palette
Args:
color_list ([str]): list of PyMOL color strings (Default value = None)
max_value (int): max palette index (Default value = 7)
min_value (int): min palette index (Default value = 1)
Returns:
palette ([str]): list of color definitions
"""
output_range = max_value - min_value
default_color_list = [
'tv_red', 'tv_orange', 'tv_yellow', 'tv_green', 'tv_blue',
'aquamarine', 'violet'
]
if color_list is None:
color_list = default_color_list
elif (output_range > 1) and len(colors) == 1:
logger.warning(
"Only a single color has been provided for multi-output visualization--supplementing the input palette with default values"
)
color_list.extend(default_color_list)
colors = []
for color in color_list:
if isinstance(color, str):
colors.append(cmd.get_color_tuple(color))
else:
colors.append(tuple(color))
palette = []
if output_range <= len(colors):
for color in colors[:max_value]:
palette.append('0x%02x%02x%02x' %
tuple([int(255 * x) for x in color]))
elif (output_range > len(colors)) and (len(colors) > 1):
step = float(len(colors)) / float(output_range)
for i in range(output_range):
ix = float(i) * step
# get the indices of the surrounding colors correcting for floating point imprecision
lower_ind = max(int(math.floor(ix)), 0)
upper_ind = min(int(math.ceil(ix)), len(colors) - 1)
fx = ix - lower_ind
if lower_ind == upper_ind:
# special case where interpolation is exactly at an input color
palette.append('0x%02x%02x%02x' %
tuple([int(255 * x)
for x in colors[lower_ind]]))
else:
color = [
fx * colors[lower_ind][i] + (1 - fx) * colors[upper_ind][i]
for i in range(3)
]
palette.append('0x%02x%02x%02x' %
tuple([int(255 * x) for x in color]))
logger.debug("Palette constructed with {0} colors".format(len(palette)))
return palette
