-
Notifications
You must be signed in to change notification settings - Fork 2
/
generate_pml.py
executable file
·102 lines (79 loc) · 2.48 KB
/
generate_pml.py
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
from glob import glob
import os
import numpy as np
from pathlib import Path
import sys
print("Basic pyton script to generate pymol file for alphafold models")
print("Improvement will come late...")
print("USAGE : ")
print(" 1. Go to your working directory (where you have the folder 'predictions' but do not go inside'")
print(" 2. pythyon generate_pml.py")
print(" 3. pymol visualisation.pml")
def keep_relax(models):
keep = []
for model in models:
if "_unrelaxed" in model:
relaxed = model.replace('_unrelaxed', "_relaxed")
if relaxed in models :
if relaxed not in keep:
keep.append(relaxed)
else:
keep.append(model)
return keep
models=glob("predictions/**/*.pdb")
if len(models) == 0:
models = glob("predictions/*.pdb")
if len(models) == 0:
print("no models found")
sys.exit(0)
models = keep_relax(models)
numfolder = len(models[0].split('/'))
relaxed_or_notrelaxed = {}
cmdload = ''
rename = ''
firstname = ''
groupcmd = ""
print(models)
if numfolder > 2:
group=True
folders=list(np.unique([x.split('/')[1] for x in models]))
for f in folders:
tmp = [x for x in models if f in x.split('/')[1]]
if "_relaxed" in tmp[0]:
relaxed_or_notrelaxed[f] = "relaxed"
labelRelaxed = "mini"
cmdload = cmdload+f'loadall predictions/{f}/*_relaxed*.pdb\n'
else:
relaxed_or_notrelaxed[f] = "unrelaxed"
cmdload = cmdload+f'loadall predictions/{f}/*_unrelaxed*.pdb\n'
#rename = "alter chain B, chain='A'\nalter chain C, chain='B'\n"
labelRelaxed = "noMini"
#cmdload = cmdload+f'loadall predictions/{f}/*_relaxed*.pdb, {f}\n'
groupcmd = groupcmd + f"group {f}, {f}_*\n"
else:
group = False
cmdload = ''
if "_relaxed" in models[0]:
cmdload = cmdload+f'loadall predictions/*_relaxed*.pdb'
labelRelaxed = "mini"
else:
cmdload = cmdload+f'loadall predictions/*_unrelaxed*.pdb'
rename = "alter chain B, chain='A'\nalter chain C, chain='B'\n"
labelRelaxed = "noMini"
pml = f"""
from pymol import cmd
{cmdload}
{rename}
#Align on first model
extra_fit b>70, {Path(models[0]).stem}
{groupcmd}
as cartoon
set grid_mode, 1
set precomputed_lighting, 1
util.cnc
set opaque_background; 0
remove symbol H
spectrum b, rainbow_rev, minimum=10, maximum=90
"""
with open('visualisation.pml','w') as out:
out.write(pml)