aboutsummaryrefslogtreecommitdiff
path: root/scripts/quick-ligand-protein.py
blob: 6ea16c1ce7898f1d2a0bcc658f38aeef4a5c8585 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
#!/usr/bin/python3
import argparse

# import logzero
# import logging
# from logzero import logger as log
import pymol2
import time

import os

print(os.getcwd())

#################
# Configuration #
#################

startTime = time.time()
version = "1.0"
desc_text = "PyMol Quick Visualtion " + version
ligandColor = "red"
# logzero.loglevel(logging.INFO)
height = 1000
width = 800
dpi = 300
ray = 0


m1 = "target"
m2 = "ligand"

parser = argparse.ArgumentParser(description=desc_text)
parser.add_argument("-p", "--protein", help="Path to protein file")
parser.add_argument("-l", "--ligand", help="Path to ligand_out file")
parser.add_argument("-c", "--color", help="Color for ligand in visualisation")

args = parser.parse_args()

if args.protein == None:
    print("Error: Please specify protein file")
    exit(1)
if args.ligand == None:
    print("Error: Please specify ligand file")
    exit(1)
if args.color == None:
    print("No color was speciifed, using default settings.")

protein = args.protein
print("Protein: ", protein)
ligand = args.ligand


def loadMol(filename, name):
    print("Loading " + filename + " as " + name)
    cmd.load(filename, name)


def changeColor(name, colorName):
    print("Changed " + name + "'s color to " + colorName)
    cmd.color(colorName, name)


def orientEtZoom():
    cmd.orient()
    cmd.zoom()


def showSurface(name):
    cmd.show("surface", name)


def surfaceTransparency(amount):
    print("Changed surface transparency to " + str(amount * 100) + "%")
    cmd.set("transparency", amount)


def generatePNG(filename, height=height, width=width, dpi=dpi, ray=ray):
    print("Generating " + filename + ".png")
    cmd.png(filename, height, width, dpi=dpi, ray=ray)


def flipHorizontal():
    cmd.rotate("y", 180)


def zoomTo(name):
    cmd.zoom(name)


def generatePictures():
    generatePNG("output-front")
    flipHorizontal()
    generatePNG("output-back")
    zoomTo(m2)
    generatePNG("closeup-back")
    orientEtZoom()
    flipHorizontal()
    zoomTo(m2)
    generatePNG("closeup-front")


print("Initialising PyMol")
session = pymol2.PyMOL()
print("Starting PyMol Session")
session.start()
cmd = session.cmd

loadMol(protein, m1)
cmd.remove("resn hoh")  # remove water
loadMol(ligand, m2)
changeColor(m1, "grey60")
changeColor(m2, ligandColor)
cmd.color("blue", "hetatm")  # color heteroatoms
orientEtZoom()
showSurface(m1)
surfaceTransparency(0.6)

generatePictures()

endTime = time.time()
print("Finished Execution in " + str(round((endTime - startTime), 2)) + " seconds.")