|
楼主 |
发表于 2022-2-21 09:55:52
|
显示全部楼层
#!/usr/bin/env python
import sys
import pymol
from pymol.cgo import *
import argparse
def argParser():
# Parser implementation
parser = argparse.ArgumentParser( prog='RING-Viz.py',
description="RING-Viz: Visualize Residue Interaction Networks with Pymol",
epilog=" Example 1: python RING-Viz.py examples/3jqz.edges\n Example 2: python RING-Viz.py examples/3jqz.edges -pdbFile examples/3jqz.pdb_modified\n\n",
formatter_class=argparse.RawDescriptionHelpFormatter)
#formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('ringFile', help='The RIN file (edges) generated by RING-2.0')
parser.add_argument('-pdbFile','--pdbFile', help='By default it is automatically fetched')
parser.add_argument('-pdbId','--pdbIdentifier', help='By default it is inferred from the ringFile')
parser.add_argument('-dMin','--distCutOffMin', default=None, type=float, help='Minimum distance threshold')
parser.add_argument('-dMax','--distCutOffMax', default=None, type=float, help='Maximum distance threshold')
parser.add_argument('-deg','--degreeMin', default=None, type=int, help='Minimum node degree')
parser.add_argument('-rMC','--removeMC_MC',action='store_true',default=False,
help='Remove main-chain/main-chain interactions')
parser.add_argument('-ori','--orientation',nargs='*',default=False,
help='Define a list of orientations to visualize')
parser.add_argument('-lw','--lineWidth',default=2.0, type=float,
help='Define interaction line width in pixels')
parser.add_argument('-int','--interaction',nargs='*',default=False,
help='Define a list of interactions to visualize')
parser.add_argument('-n','--nodes',nargs='*',default=False,
help='Consider all contacts involving at least one of the nodes in the list, <chain><position>')
parser.add_argument('-ns','--nodesStrict',action='store_true',default=False,
help='Consider all contacts involving only nodes in the --nodes list (strictly)')
# Initialize the parser object
args=parser.parse_args()
return args
############################################################
# Color map
intTypeMap = {
"IONIC":{"color":"blue","rgb":(0.0,0.0,1.0)},
"SSBOND":{"color":"yellow","rgb":(1.0,1.0,0.0)},
"PIPISTACK":{"color":"orange","rgb":(1.0,0.5,0.0)},
"PICATION":{"color":"red","rgb":(1.0,0.0,0.0)},
"HBOND":{"color":"cyan","rgb":(0.0,1.0,1.0)},
"VDW":{"color":"gray50","rgb":(0.5,0.5,0.5)},
"IAC":{"color":"white","rgb":(1.0,1.0,1.0)}
}
# Get all coordinates
def getCoordinates(pdbName):
model = pymol.cmd.get_model(pdbName, 1)
coords = {}
for at in model.atom:
insetionCode = "_"
resi = at.resi
try:
int(at.resi)
except:
insetionCode = at.resi[-1]
resi = at.resi[0:-1]
atom_id = at.chain + ":" + resi + ":" + insetionCode + ":" + at.resn + ":" + at.name
coords[atom_id] = [at.coord[0], at.coord[1], at.coord[2]]
del model
return coords
def parseEdges(fileName, coords, distCutOffMin = None, distCutOffMax = None, rmMainChain = False, degreeMin = None, nodeList = [], ppOrientList = [], intList = [], nodeListStrict = False):
'''
{
interaction_type : [
(
(chain, position, atom, coordinates),
(chain, position, atom, coordinates)
),
]
}
'''
interactions = {}
degree = {}
missing = []
with open(fileName) as f:
for line in f:
if line:
if line[0]!="#" and not line.startswith("NodeId"):
# Parse one RING edges line
nodeId1,interaction,nodeId2,distance,angle,energy,atom1,atom2,donor,positive,cation,orientation = line[0:-1].split("\t")[0:12]
chain1,pos1,i1,res1 = nodeId1.split(":")
chain2,pos2,i2,res2 = nodeId2.split(":")
pos1 = int(pos1)
pos2 = int(pos2)
intType,intSubType = interaction.split(":")
sub1,sub2 = intSubType.split("_")
if orientation:
orientation = orientation.split()[0]
distance = float(distance)
start = True
# Filters
if distCutOffMin:
if (distance <= distCutOffMin):
start = False
if distCutOffMax:
if (distance > distCutOffMax):
start = False
if rmMainChain:
if sub1 == "MC" and sub2 == "MC":
start = False
if ppOrientList:
if orientation not in ppOrientList:
start = False
if intList:
if intType not in intList:
start = False
if start:
coord1 = None
if (atom1.find(","))!=-1:
coord1 = [float(i) for i in atom1.split(",")]
atom1 = None
else:
coord1 = coords.get(nodeId1 + ":" + atom1)
coord2 = None
if (atom2.find(","))!=-1:
coord2 = [float(i) for i in atom2.split(",")]
atom2 = None
else:
coord2 = coords.get(nodeId2 + ":" + atom2)
# Assign interaction elements
if coord1 and coord2:
# Pymol selection of the interacting atoms
selection1 = (chain1,pos1,atom1,coord1)
selection2 = (chain2,pos2,atom2,coord2)
# Collect atom pairs
if intType not in interactions:
interactions[intType] = []
interactions[intType].append((selection1,selection2))
else:
missingCoords = []
if not coord1:
missingCoords.append((nodeId1,atom1))
if not coord2:
missingCoords.append((nodeId2,atom2))
missing.append((intType,nodeId1,nodeId2,distance,missingCoords))
if degreeMin:
# Calculate node degree
if (chain1,pos1) not in degree:
degree[(chain1,pos1)] = 0
degree[(chain1,pos1)] += 1
if (chain2,pos2) not in degree:
degree[(chain2,pos2)] = 0
degree[(chain2,pos2)] += 1
validNodes = None
# Filter by node degree
if degreeMin:
validNodesDegree = set([NodeId for NodeId in degree if degree[NodeId] >= degreeMin])
validNodes = validNodesDegree
print("No. nodes after filtering by degree >= ",degreeMin,":",len(validNodesDegree))
# Filter by user provided nodes
if nodeList:
validNodesUser = set([(node.strip()[0],int(node.strip()[1:])) for node in nodeList])
if degreeMin:
validNodes.intersection_update(validNodesUser)
else:
validNodes = validNodesUser
print("No. nodes after filtering by the user provided list:",len(validNodes))
# Remove edges
if nodeList or degreeMin:
for intType in interactions:
for i in range(len(interactions[intType])-1,-1,-1):
selection1,selection2 = interactions[intType][i]
if ((selection1[0],selection1[1]) not in validNodes and (selection2[0],selection2[1]) not in validNodes) or (nodeListStrict and ((selection1[0],selection1[1]) not in validNodes or (selection2[0],selection2[1]) not in validNodes)):
del interactions[intType][i]
if missing:
print ("\nTotal missing connections:",len(missing),"\nFirst 10 missing:\n", "\n".join([repr(ele) for ele in missing][0:10]))
print ("\nInteractions parsed:",[(intType,len(interactions[intType])) for intType in interactions])
print
return interactions
def drawConnections(interactions,intTypeMap,linewidth):
for interaction in interactions:
if interactions[interaction]:
# Select involved residues
selName = interaction + "_nodes"
pymol.cmd.select(selName, " or ".join(frozenset(["{}/{}/".format(s[0],s[1]) for ele in interactions[interaction] for s in ele])))
pymol.cmd.show( "sticks", interaction + "_nodes")
#pymol.cmd.show( "lines", interaction + "_nodes")
#pymol.cmd.color( intTypeMap[interaction]["color"], interaction + "_nodes")
objName = interaction + "_edges"
cr,cg,cb = intTypeMap[interaction]["rgb"]
obj = [ BEGIN, LINES, COLOR, cr, cg, cb ]
for s,t in interactions[interaction]:
s_coords = s[3]
t_coords = t[3]
obj.append( VERTEX )
obj.append(s_coords[0] )
obj.append(s_coords[1] )
obj.append(s_coords[2])
obj.append( VERTEX )
obj.append(t_coords[0])
obj.append(t_coords[1])
obj.append(t_coords[2])
obj.append(END)
pymol.cmd.load_cgo( obj, objName )
pymol.cmd.set("cgo_line_width", float(linewidth), objName )
print interaction,"contacts calculated"
#######################################################################
# Main
args = argParser()
if not args.pdbFile:
if not args.pdbIdentifier:
pdbName = args.ringFile.replace("..","").split(".")[0].replace("pdb","").replace("_edges","").replace("\\","/").split("/")[-1]
else:
pdbName = args.pdbIdentifier
pymol.finish_launching()
try:
print("Trying to download:", pdbName)
pymol.cmd.fetch(pdbName, pdbName)
except:
print("Fail to download:",pdbName)
else:
pdbName = args.pdbFile.replace("..","").split(".")[0].replace("pdb","").replace("_edges","").replace("\\","/").split("/")[-1]
pymol.finish_launching()
pymol.cmd.load(args.pdbFile, pdbName)
print("pdbName:", pdbName)
# Remove ambiguous atoms by occupancy values
try:
pymol.cmd.remove(pdbName + " and q<0.5")
except:
print("Fail to remove alternative atoms (by occupancy)")
# Remove water molecules
pymol.cmd.remove("resn hoh")
pymol.cmd.hide( "lines", "all")
pymol.cmd.show( "cartoon", "all")
pymol.cmd.show( "sticks", "hetatm")
#pymol.cmd.spectrum(selection="all")
pymol.util.cbc(selection="all")
pymol.util.cnc(selection="all")
# Parse interaction file
interactions = parseEdges( args.ringFile,
getCoordinates(pdbName),
distCutOffMin = args.distCutOffMin,
distCutOffMax = args.distCutOffMax,
rmMainChain = args.removeMC_MC,
degreeMin = args.degreeMin,
nodeList = args.nodes,
nodeListStrict = args.nodesStrict,
ppOrientList = args.orientation,
intList = args.interaction )
# Draw connections
drawConnections(interactions,intTypeMap,args.lineWidth)
pymol.cmd.orient(pdbName) |
|