def mesh_to_wireframe_CGO(mesh, color_tuple, alpha=1.0): """Creates a wireframe CGO object for a mesh for display in PyMOL Args: mesh (Trimesh): Trimesh mesh object color (str): PyMOL color string (Default value = 'gray60') alpha (float): transparency value (Default value = 1.0) Returns: cgobuffer (str): CGO buffer that contains the instruction to load a wireframe object """ cgobuffer = [cgo.BEGIN, cgo.LINES, cgo.ALPHA, alpha] cgobuffer.append(cgo.COLOR) cgobuffer.extend(cmd.get_color_tuple(cmd.get_color_index(color))) for edge in mesh.edges: cgobuffer.append(cgo.VERTEX) cgobuffer.extend(mesh.vertices[edge[0]]) cgobuffer.append(cgo.VERTEX) cgobuffer.extend(mesh.vertices[edge[1]]) cgobuffer.append(cgo.END) logger.debug("CGO wireframe object created for mesh: {0}".format(mesh)) return cgobuffer
def mesh_to_solid_CGO(mesh, color, alpha=1.0): """Creates a solid CGO object for a mesh for display in PyMOL Args: mesh (Trimesh): Trimesh mesh object color (str): PyMOL color string (Default value = 'gray60') alpha (float): transparency value (Default value = 1.0) Returns: cgobuffer (str): CGO buffer that contains the instruction to load a solid object """ cgobuffer = [cgo.BEGIN, cgo.TRIANGLES, cgo.ALPHA, alpha] color_values = cmd.get_color_tuple(cmd.get_color_index(color)) for face in mesh.faces: for v_index in face: cgobuffer.append(cgo.COLOR) cgobuffer.extend(color_values) cgobuffer.append(cgo.NORMAL) cgobuffer.extend( [mesh.vertex_normals[v_index][i] for i in range(3)]) cgobuffer.append(cgo.VERTEX) cgobuffer.extend([mesh.vertices[v_index][i] for i in range(3)]) cgobuffer.append(cgo.END) logger.debug("CGO solid object created for mesh: {0}".format(mesh)) return cgobuffer
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 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 add_menu(): '''Add a color blind-friendly list of colors to the PyMOL OpenGL menu.''' # Make sure cb_colors are installed. print('Checking for colorblindfriendly colors...') try: if cmd.get_color_index('cb_red') == -1: # mimic pre-1.7.4 behavior raise pymol.CmdException except pymol.CmdException: print('Adding colorblindfriendly colors...') set_colors() # Abort if PyMOL is too old. try: from pymol.menu import all_colors_list except ImportError: print('PyMOL version too old for cb_colors menu. Requires 1.6.0 or later.') return # Add the menu print('Adding cb_colors menu...') # mimic pymol.menu.all_colors_list format # first color in list is used for menu item color cb_colors = ('cb_colors', [ ('830', 'cb_red'), ('064', 'cb_green'), ('046', 'cb_blue'), ('882', 'cb_yellow'), ('746', 'cb_magenta'), ('368', 'cb_skyblue'), ('860', 'cb_orange'), ]) # First `pymol` is the program instance, second is the Python module all_colors_list = pymol.pymol.menu.all_colors_list if cb_colors in all_colors_list: print('Menu was already added!') else: all_colors_list.append(cb_colors) print(' done.')
def set(self, v): if isinstance(v, str) and v.strip(): v = cmd.get_color_index(v) super(ColorVar, self).set(v)
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 updateColor(self): selection = 'all' stored.atoms_charge = [] stored.atoms_colors = [] cmd.map_new('chargyc_map', selection="(all)") if not self.colorNeutral: tkMessageBox.showerror("Error", "Set Neutral Color, Please") return if not self.colorNegative: tkMessageBox.showerror("Error", "Set Negative Color, Please") return if not self.colorPositive: tkMessageBox.showerror("Error", "Set Positive Color, Please") return cmd.iterate_state(1, '(' + selection + ')', 'stored.atoms_charge.append(partial_charge)') _i = 0 minValue = None maxValue = None while _i < len(stored.atoms_charge): color = [] if _i == 0: maxValue = stored.atoms_charge[_i] minValue = stored.atoms_charge[_i] if(stored.atoms_charge[_i] > maxValue): maxValue = stored.atoms_charge[_i] if stored.atoms_charge[_i] < minValue: minValue = stored.atoms_charge[_i] _i += 1 self.minNegativeLbl["text"] = 'Min Found: ' + str(round(minValue, 3)) self.maxPositiveLbl["text"] = 'Max Found: ' + str(round(maxValue, 3)) if(self.scaleNegative == 0.0 and self.scalePositive == 0.0): self.scaleNegative = round(minValue, 3) self.scalePositive = round(maxValue, 3) self.sclSelectNegative.delete(0, "end") self.sclSelectPositive.delete(0, "end") self.sclSelectNegative.insert(0, round(minValue, 3)) self.sclSelectPositive.insert(0, round(maxValue, 3)) else: self.scaleNegative = float(self.sclSelectNegative.get()) self.scalePositive = float(self.sclSelectPositive.get()) minValue = float(self.sclSelectNegative.get()) maxValue = float(self.sclSelectPositive.get()) self.scaleLbl["text"] = 'Scale: ' + str( self.scaleNegative) + ' to ' + str(self.scalePositive) middleValue = 0 if(maxValue < 0): maxValue = 0 if(minValue > 0): minValue = 0 _i = 0 while _i < len(stored.atoms_charge): color = [] cmd.set_color("neutral_color", self.colorNeutral[0]) cmd.set_color("positive_color", self.colorPositive[0]) cmd.set_color("negative_color", self.colorNegative[0]) if(stored.atoms_charge[_i] >= middleValue): if(stored.atoms_charge[_i] == middleValue): cmd.set_color(str(_i) + "_color", self.colorNeutral[0]) else: cmd.set_color(str(_i) + "_color", self.getColor( self.colorNeutral[0], self.colorPositive[0], maxValue, stored.atoms_charge[_i] if stored.atoms_charge[_i] < maxValue else maxValue)) else: cmd.set_color(str(_i) + "_color", self.getColor( self.colorNeutral[0], self.colorNegative[0], abs(minValue), abs(stored.atoms_charge[_i]) if abs(stored.atoms_charge[_i]) < abs(minValue) else abs(minValue))) index = cmd.get_color_index(str(_i) + "_color") stored.atoms_colors.append(index) _i += 1 cmd.alter_state(1, '(' + selection + ')', "color=stored.atoms_colors.pop(0)") cmd.ramp_new('chargy_ramp', 'chargyc_map', range=[self.scaleNegative, ((self.scaleNegative+self.scalePositive)/2.0), self.scalePositive], color=['negative_color', 'neutral_color', 'positive_color'])
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 read_moestr(contents, object, state=0, finish=1, discrete=1, quiet=1, zoom=-1, _self=cmd): moestr = contents name = object import sys if sys.version_info[0] > 2 and isinstance(moestr, bytes): moestr = moestr.decode() cmd = _self mr = MOEReader() mr.appendFromStr(moestr) split_chains = cmd.get_setting_int("moe_separate_chains") cmd.group(name) if hasattr(mr, 'system'): have_valences = 0 chain_count = 0 cmd.set_color("_aseg0", [1.0, 1.0, 1.0]) aseg_color = cmd.get_color_index("_aseg0") aseg_flag = 0 aseg_rep = {} model = Indexed() molecule = mr.system['molecule'] if 'atoms' in molecule: n_atom = molecule['atoms'] model.atom = [Atom() for x in range(n_atom)] residues = {} chains = {} for columns, data in molecule['attr']: for row in data: cur_atom = None for key, value in zip(columns, row): key = key[0] if key == 'ID': ID = value else: aProp = _atom_prop_map.get(key, None) if aProp != None: setattr(model.atom[ID - 1], aProp, value) else: xyz = _atom_coord_map.get(key, None) if xyz != None: coord = list(model.atom[ID - 1].coord) coord[xyz] = value model.atom[ID - 1].coord = coord elif key in _atom_vis_map: atom = model.atom[ID - 1] if hasattr(atom, 'visible'): visible = atom.visible else: visible = _default_visible if key == 'aBondLook': if value == 'cylinder': atom.visible = 0x00000001 | visible elif value == 'line': atom.visible = 0x00000080 | visible elif value == 'none': atom.visible = -129 & Visible # 0xFFFFFF7F elif key == 'aNucleusLook': if value == 'sphere': atom.visible = 0x00000002 | visible elif value == 'small-sphere': # need to set sphere_scale=0.2 for these atoms atom.visible = 0x00000002 | visible atom.sphere_scale = 0.2 elif value == 'point': # nonbonded atom.visible = 0x00000800 | visible elif value == 'none': atom.visible = -2067 & visible # 0xFFFFF7ED elif key == 'aHidden': atom.visible = 0 atom.hidden = 1 if hasattr( atom, 'hidden' ): # be sure that hidden atoms aren't shown atom.visible = 0 elif key in _atom_color_map: if key == 'aRGB': model.atom[ID - 1].trgb = value elif key == 'aColorBy': model.atom[ID - 1].aColorBy = value elif key in _atom_label_map: atom = model.atom[ID - 1] if hasattr(atom, 'label_dict'): atom.label_dict[key] = None else: atom.label_dict = {key: None} elif key in _residue_prop_map: resi_dict = residues.get(ID, {}) resi_dict[key] = value residues[ID] = resi_dict elif key in _chain_prop_map: chain_dict = chains.get(ID, {}) if ID not in chains: chain_count = chain_count + 1 chain_dict['count'] = chain_count chain_dict[key] = value chains[ID] = chain_dict chn_keys = list(chains.keys()) chn_keys.sort() res_keys = list(residues.keys()) res_keys.sort() # map chain properties onto residues chn_resi = 0 ch_colors = copy.deepcopy(_ch_colors) unique_chain_names = {} for chn_idx in chn_keys: chain_dict = chains[chn_idx] cName = make_valid_name(chain_dict.get('cName', '')) segi = cName[0:4] chain = cName[-1:] if not len(cName): if 'count' in chain_dict: cName = "chain_" + str(chain_dict['count']) else: cName = str(chn_idx) if cName not in unique_chain_names: unique_chain_names[cName] = cName else: cnt = 2 while (cName + "_" + str(cnt)) in unique_chain_names: cnt = cnt + 1 newCName = cName + "_" + str(cnt) unique_chain_names[newCName] = cName cName = newCName chain_dict['chain_color'] = ch_colors[0] ch_colors = ch_colors[1:] + [ch_colors[0]] cResCount = chain_dict.get('cResidueCount', 0) for res_idx in range(chn_resi, chn_resi + cResCount): resi_dict = residues[res_keys[res_idx]] resi_dict['chain'] = chain resi_dict['segi'] = segi resi_dict['cName'] = cName resi_dict['chain_dict'] = chain_dict chn_resi = chn_resi + cResCount # map residue properties onto atoms res_atom = 0 for res_idx in res_keys: resi_dict = residues[res_idx] rRibbonMode = resi_dict.get('rRibbonMode', 'none') rAtomCount = resi_dict['rAtomCount'] rType = resi_dict.get('rType', '') if rAtomCount > 0: for at_idx in range(res_atom, res_atom + rAtomCount): atom = model.atom[at_idx] setattr( atom, 'resi', string.strip( str(resi_dict.get('rUID', '')) + resi_dict.get('rINS', ''))) setattr(atom, 'resn', resi_dict.get('rName', '')) setattr(atom, 'chain', resi_dict.get('chain', '')) setattr(atom, 'segi', resi_dict.get('segi', '')) setattr(atom, 'custom', resi_dict.get('cName', '')) # add labels if hasattr(atom, 'label_dict'): label = '' label_dict = atom.label_dict if 'aLabelElement' in label_dict: label = atom.symbol if 'aLabelRes' in label_dict: if len(label): label = label + "," label = label + atom.resn + "_" + atom.resi if 'aLabelName' in label_dict: if len(label): label = label + "." label = label + atom.name atom.label = label if rType not in ['none', 'heme']: atom.hetatm = 0 # all normal atoms else: atom.flags = 0x02000000 # hetatom or ligand -- ignore when surfacing if rRibbonMode != 'none': if hasattr(atom, 'visible'): visible = atom.visible else: visible = _default_visible rRibbonColorBy = resi_dict['rRibbonColorBy'] repeat = 1 while repeat: repeat = 0 if rRibbonColorBy in ['rgb', 'r:rgb' ]: # handled automatically pass elif rRibbonColorBy == 'chain': chain_dict = resi_dict['chain_dict'] cColorBy = chain_dict['cColorBy'] if cColorBy == 'rgb': cColorBy = 'c:rgb' rRibbonColorBy = cColorBy repeat = 1 rRibbon_color = 0 rRibbon_trgb = 0xFFFFFF # default -- white # now color ribbon if rRibbonColorBy == 'r:rgb': rRibbon_trgb = resi_dict.get('rRGB', 0xFFFFFF) elif rRibbonColorBy == 'rgb': rRibbon_trgb = resi_dict.get( 'rRibbonRGB', 0xFFFFFF) elif rRibbonColorBy == 'c:rgb': chain_dict = resi_dict['chain_dict'] rRibbon_trgb = chain_dict.get('cRGB', 0xFFFFFF) elif rRibbonColorBy == 'r:aseg': rRibbon_trgb = None rRibbon_color = aseg_color aseg_flag = 1 elif rRibbonColorBy == 'tempfactor': pass # TO DO elif rRibbonColorBy == 'c:number': # per chain colors rRibbon_trgb = chain_dict['chain_color'] if rRibbonMode in ['line', 'trace']: atom.visible = 0x00000040 | visible # PyMOL ribbon if rRibbon_trgb != None: atom.ribbon_trgb = rRibbon_trgb else: atom.ribbon_color = rRibbon_color aseg_rep['ribbon'] = 1 else: atom.visible = 0x00000020 | visible # PyMOL cartoon if rRibbon_trgb != None: atom.cartoon_trgb = rRibbon_trgb else: atom.cartoon_color = rRibbon_color aseg_rep['cartoon'] = 1 if hasattr(atom, 'label'): if hasattr(atom, 'visible'): visible = atom.visible else: visible = _default_visible atom.visible = 0x00000028 | visible # labels if not hasattr(atom, 'aColorBy'): atom.aColorBy = 'element' if hasattr(atom, 'aColorBy'): aColorBy = atom.aColorBy repeat = 1 while repeat: repeat = 0 if aColorBy == 'ribbon': aColorBy = resi_dict.get('rRibbonColorBy') if aColorBy == 'rgb': aColorBy = 'rib:rgb' else: repeat = 1 # TO DO still need to handle "cartoon_color", "ribbon_color" elif aColorBy == 'element': if hasattr(atom, 'trgb'): del atom.trgb elif aColorBy in ['rgb', 'a:rgb' ]: # handled automatically pass elif aColorBy == 'residue': rColorBy = resi_dict.get('rColorBy') if rColorBy == 'rgb': rColorBy = 'r:rgb' aColorBy = rColorBy repeat = 1 elif aColorBy == 'chain': chain_dict = resi_dict['chain_dict'] cColorBy = chain_dict['cColorBy'] if cColorBy == 'rgb': cColorBy = 'c:rgb' aColorBy = cColorBy repeat = 1 # now color atom... if aColorBy == 'r:rgb': atom.trgb = resi_dict.get('rRGB', 0xFFFFFF) elif aColorBy == 'rib:rgb': atom.trgb = resi_dict.get('rRibbonRGB', 0xFFFFFF) elif aColorBy == 'c:rgb': chain_dict = resi_dict['chain_dict'] atom.trgb = chain_dict.get('cRGB', 0xFFFFFF) elif aColorBy == 'r:aseg': pass # TO DO elif aColorBy == 'tempfactor': pass # TO DO elif aColorBy == 'c:number': # per chain colors atom.trgb = chain_dict['chain_color'] res_atom = res_atom + rAtomCount bond_list = molecule.get('bond', []) for bond in bond_list: new_bond = Bond() new_bond.index = [bond[0] - 1, bond[1] - 1] if len(bond) > 2: new_bond.order = bond[2] if bond[2] == 2: # work around .MOE bug with triple bonds if model.atom[new_bond.index[0]].hybridization == 'sp': if model.atom[new_bond.index[1]].hybridization == 'sp': new_bond.order = 3 have_valences = 1 model.bond.append(new_bond) if 'ViewOrientationY' in mr.system: vy = mr.system['ViewOrientationY'] vz = mr.system['ViewOrientationZ'] pos = mr.system['ViewLookAt'] scale = mr.system['ViewScale'] vx = cpv.cross_product(vy, vz) m = [cpv.normalize(vx), cpv.normalize(vy), cpv.normalize(vz)] m = cpv.transpose(m) asp_rat = 0.8 # MOE default (versus 0.75 for PyMOL) cmd.set("field_of_view", 25.0) fov = float(cmd.get("field_of_view")) window_height = scale * asp_rat dist = (0.5 * window_height) / math.atan(3.1415 * (0.5 * fov) / 180.0) new_view = tuple(m[0] + m[1] + m[2] + [0.0, 0.0, -dist] + pos + [dist * 0.5, dist * 1.5, 0.0]) cmd.set_view(new_view) zoom = 0 cmd.set("auto_color", 0) cmd.set_color("carbon", [0.5, 0.5, 0.5]) # default MOE grey obj_name = name + ".system" if split_chains < 0: # by default, don't split chains if over 50 objects would be created if len(unique_chain_names) > 50: split_chains = 0 if not split_chains: cmd.load_model(model, obj_name, state=state, finish=finish, discrete=discrete, quiet=quiet, zoom=zoom) obj_name_list = [obj_name] else: cmd.load_model(model, obj_name, state=state, finish=finish, discrete=discrete, quiet=1, zoom=zoom) obj_name_list = [] system_name = obj_name for chain in unique_chain_names.keys(): obj_name = name + "." + chain obj_name_list.append(obj_name) cmd.select("_moe_ext_tmp", "custom %s" % chain, domain=system_name) cmd.extract(obj_name, "_moe_ext_tmp", quiet=quiet, zoom=0) # cmd.extract(obj_name,system_name+" and text_type %s"%chain,quiet=quiet) cmd.delete("_moe_ext_tmp") if not cmd.count_atoms(system_name): cmd.delete(system_name) else: obj_name_list.append(system_name) cmd.order(name + ".*", sort=1) for obj_name in obj_name_list: cmd.set("stick_radius", 0.1, obj_name) cmd.set("line_width", 2.0, obj_name) cmd.set("label_color", "white", obj_name) cmd.set("nonbonded_size", 0.05, obj_name) # temporary workaround... if have_valences: # if this MOE file has valences, then show em! cmd.set("valence", 1, obj_name) cmd.set("stick_valence_scale", 1.25, obj_name) if aseg_flag: cmd.dss(obj_name) if 'cartoon' in aseg_rep: cmd.set("cartoon_color", "red", obj_name + " and cartoon_color _aseg0 and ss h") cmd.set("cartoon_color", "yellow", obj_name + " and cartoon_color _aseg0 and ss s") cmd.set( "cartoon_color", "cyan", obj_name + " and cartoon_color _aseg0 and not ss h+s") if 'ribbon' in aseg_rep: cmd.set("ribbon_color", "red", obj_name + " and ribbon_color _aseg0 and ss h" ) # need selection op ribbon_color cmd.set("ribbon_color", "yellow", obj_name + " and ribbon_color _aseg0 and ss s") cmd.set( "ribbon_color", "cyan", obj_name + " and ribbon_color _aseg0 and not ss h+s") if 'ViewZFront' in mr.system: moe_front = mr.system['ViewZFront'] moe_width = mr.system['ViewZWidth'] extent = cmd.get_extent( name) # will this work with groups? TO DO: check! dx = (extent[0][0] - extent[1][0]) dy = (extent[0][1] - extent[1][1]) dz = (extent[0][2] - extent[1][2]) half_width = 0.5 * math.sqrt(dx * dx + dy * dy + dz * dz) cmd.clip("atoms", 0.0, name) cur_view = cmd.get_view() center = (cur_view[-3] + cur_view[-2]) * 0.5 # center of clipping slab front = center - half_width back = center + half_width width = half_width * 2 new_view = tuple( list(cur_view[0:15]) + [ front + width * moe_front, front + width * (moe_front + moe_width), 0.0 ]) cmd.set_view(new_view) if 'graphics' in mr.system: cgo_cnt = 1 lab_cnt = 1 unique_cgo_names = {} for graphics in mr.system['graphics']: cgo = [] for gvertex in graphics.get('gvertex', []): vrt = gvertex[0] seg_list = gvertex[1]['seg'] idx = gvertex[1]['idx'] len_idx = len(idx) if not cmd.is_list(seg_list): seg_list = [seg_list] * (len_idx / seg_list) last_seg = None ix_start = 0 for seg in seg_list: if seg != last_seg: if last_seg != None: cgo.append(END) if seg == 3: cgo.extend([BEGIN, TRIANGLES]) elif seg == 2: cgo.extend([BEGIN, LINES]) elif seg == 1: cgo.extend([BEGIN, POINTS]) ix_stop = seg + ix_start if seg == 3: for s in idx[ix_start:ix_stop]: v = vrt[s - 1] cgo.extend([ COLOR, (0xFF & (v[0] >> 16)) / 255.0, (0xFF & (v[0] >> 8)) / 255.0, (0xFF & (v[0])) / 255.0 ]) if len(v) > 4: cgo.extend([NORMAL, v[4], v[5], v[6]]) cgo.extend([VERTEX, v[1], v[2], v[3]]) elif seg == 2: for s in idx[ix_start:ix_stop]: v = vrt[s - 1] cgo.extend([ COLOR, (0xFF & (v[0] >> 16)) / 255.0, (0xFF & (v[0] >> 8)) / 255.0, (0xFF & (v[0])) / 255.0 ]) if len(v) > 4: cgo.extend([NORMAL, v[4], v[5], v[6]]) cgo.extend([VERTEX, v[1], v[2], v[3]]) elif seg == 1: for s in idx[ix_start:ix_stop]: v = vrt[s - 1] cgo.extend([ COLOR, (0xFF & (v[0] >> 16)) / 255.0, (0xFF & (v[0] >> 8)) / 255.0, (0xFF & (v[0])) / 255.0 ]) if len(v) > 4: cgo.extend([NORMAL, v[4], v[5], v[6]]) cgo.extend([VERTEX, v[1], v[2], v[3]]) ix_start = ix_stop last_seg = seg if last_seg != None: cgo.append(END) for gtext in graphics.get('gtext', []): lab_name = name + ".label_%02d" % lab_cnt exists = 0 for entry in gtext: exists = 1 cmd.pseudoatom(lab_name, pos=[ float(entry[1]), float(entry[2]), float(entry[3]) ], color="0x%06x" % entry[0], label=entry[4]) if exists: cmd.set('label_color', -1, lab_name) lab_cnt = lab_cnt + 1 # TO DO -- via CGO's? if len(cgo): cgo_name = name + "." + make_valid_name( graphics.get('GTitle', 'graphics')) if cgo_name not in unique_cgo_names: unique_cgo_names[cgo_name] = cgo_name else: cnt = 2 while cgo_name + "_" + str(cnt) in unique_cgo_names: cnt = cnt + 1 new_name = cgo_name + "_" + str(cnt) unique_cgo_names[new_name] = new_name cgo_name = new_name cmd.load_cgo(cgo, cgo_name, state=state, quiet=quiet, zoom=0) cgo_cnt = cgo_cnt + 1 cmd.set("two_sided_lighting", 1) # global setting... cmd.set("cgo_line_width", 2, cgo_name) if 'GTransparency' in graphics: g_trans = graphics['GTransparency'] if len(g_trans) >= 2: if g_trans[0] != 0: cmd.set('cgo_transparency', '%1.6f' % (g_trans[0] / 255.0), cgo_name) cmd.set('transparency_global_sort') if 'meter' in mr.system: meter_name = name + ".meter" exists = 0 for meter_block in mr.system['meter']: if meter_block[0][0:2] == ['type', 'atoms']: for meter in meter_block[1]: (type, atoms) = meter[0:2] arg = tuple([meter_name] + list( map(lambda x, o=name: o + " and id " + str(x - 1), atoms))) getattr(cmd, type)(*arg) exists = 1 if exists: cmd.color("green", meter_name) # print mr.system['meter'] elif hasattr(mr, 'feature'): model = Indexed() cols = mr.feature[0] rows = mr.feature[1] col = {} cnt = 0 for a in cols: col[a] = cnt cnt = cnt + 1 for row in rows: atom = Atom() atom.coord = [row[col['x']], row[col['y']], row[col['z']]] atom.vdw = row[col['r']] atom.custom = row[col['expr']] model.atom.append(atom) obj_name = name + ".feature" cmd.load_model(model, obj_name, state=state, finish=finish, discrete=discrete, quiet=quiet, zoom=zoom) rank = 1 for row in rows: cmd.color("0x%06x" % row[col['color']], obj_name + " & id %d" % (rank - 1)) rank = rank + 1 cmd.show("mesh", obj_name) cmd.set("min_mesh_spacing", 0.55, obj_name) cmd.label(obj_name, "' '+custom") cmd.set("label_color", "yellow", obj_name) else: print(dir(mr))
def get_color_rgb(color): """Given a color name, returns the RGB values.""" index = cmd.get_color_index(color) rgb = cmd.get_color_tuple(index) return rgb # e.g. (1.0, 1.0, 1.0)
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 testGetColorIndex(self): self.assertEqual(cmd.get_color_index("blue"), 2)
def nice_settings(): """ sets up a series of settings and stores in pymol.stored to be used at later points by stir parts: - a set of nicely distinguishable colors for nicecolor - a dictionary of selection expressions for nicesele - a dictionary of representation information for nice """ # colors are based on this SO answer: # https://graphicdesign.stackexchange.com/questions/3682/ # where-can-i-find-a-large-palette-set-of-contrasting-colors-for-coloring-many-d colors = [ (240, 163, 255), (0, 117, 220), (153, 63, 0), (76, 0, 92), (25, 25, 25), (0, 92, 49), (43, 206, 72), (255, 204, 153), (128, 128, 128), (148, 255, 181), (143, 124, 0), (157, 204, 0), (194, 0, 136), (0, 51, 128), (255, 164, 5), (255, 168, 187), (66, 102, 0), (255, 0, 16), (94, 241, 242), (0, 153, 143), (224, 255, 102), (116, 10, 255), (153, 0, 0), (255, 255, 128), (255, 255, 0), (255, 80, 5), ] nicecolors = {} for i, rgb in enumerate(colors): col_name = f'nicecolor_{i}' cmd.set_color(col_name, rgb) nicecolors[col_name] = rgb # store color ids in a list to be used by nicecolor stored.nice_colors = [] for color in nicecolors: idx = cmd.get_color_index(color) stored.nice_colors.append(idx) # selection expressions stored.nice_selectors = { 'prot': 'polymer.protein', 'BB': 'polymer.protein and name BB', 'SC': 'polymer.protein and not name BB', 'solv': 'resn W+WN+ION', 'ions': 'resn ION', 'lip': 'organic and not ions', 'nucl': 'polymer.nucleic' } # settings to be used by nice. Values are lists with a function as first element and # its arguments following stored.nice_set = { 'clean': { 'prot': { 'color': None, 'style': None, }, 'BB': { 'color': [nicecolor, 'chain'], 'style': [cmd.show_as, 'sticks'], }, 'SC': { 'color': None, 'style': [cmd.hide, 'everything'], }, 'solv': { 'color': None, 'style': [cmd.hide, 'everything'], }, 'ions': { 'color': None, 'style': [cmd.hide, 'everything'], }, 'lip': { 'color': [nicecolor, 'resn'], 'style': [cmd.show_as, 'sticks'], }, 'nucl': { 'color': [nicecolor, 'resi'], 'style': [cmd.show_as, 'sticks'], }, }, 'rainbow': { 'prot': { 'color': None, 'style': None, }, 'BB': { 'color': [nicecolor, 'chain'], 'style': [cmd.show_as, 'sticks'], }, 'SC': { 'color': [nicecolor, 'resn'], 'style': [cmd.show_as, 'sticks'], }, 'solv': { 'color': [cmd.color, 'blue'], 'style': [cmd.show_as, 'nb_spheres'], }, 'ions': { 'color': [nicecolor, 'name'], 'style': [cmd.show_as, 'nb_spheres'], }, 'lip': { 'color': [nicecolor, 'resi'], 'style': [cmd.show_as, 'sticks'], }, 'nucl': { 'color': [nicecolor, 'resn'], 'style': [cmd.show_as, 'sticks'], }, }, 'balls': { 'prot': { 'color': None, 'style': None, }, 'BB': { 'color': [cmd.color, 'purple'], 'style': [cmd.show_as, 'spheres'], }, 'SC': { 'color': [cmd.color, 'red'], 'style': [cmd.show_as, 'spheres'], }, 'solv': { 'color': [cmd.color, 'blue'], 'style': [cmd.show_as, 'nb_spheres'], }, 'ions': { 'color': [nicecolor, 'resn'], 'style': [cmd.show_as, 'nb_spheres'], }, 'lip': { 'color': [nicecolor, 'resn'], 'style': [cmd.show_as, 'spheres'], }, 'nucl': { 'color': [nicecolor, 'resi'], 'style': [cmd.show_as, 'spheres'], }, }, }