structure-scripts/dlg_to_pdb_V2.py at master · akshaips/structure-scripts · GitHub

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
122
123
124
125
126
127
128
129
130
131
import os
import sys

runs = []

if len(sys.argv) < 2:
	print("USAGE:\n")
	print("   python dlg_to_pdb.py filename.dlg")
	print("\n   eg: python dlg_to_pdb.py dock.dlg")
	print("\nGive inputs : \n\n   'All' for all to be extracted separated\
		\n	'0' to extract everything in one file \n\
	'2' to extract only second structure")
	exit()
else:
	dlg_file = sys.argv[1]
	cluster_extraction = input("Cluster extract (y/n):")
	if cluster_extraction == "n" or cluster_extraction == "N":
		model_number = input('Model number ("All"/0(to write in one pdb)/\'model number to be extracted\'(eg:1,2,3)) :   ')
	elif cluster_extraction == "y" or cluster_extraction == "Y":
		print ("Enter cluster number")
	else:
		print ("Enter valid input")
		exit()
	os.system("mkdir " + dlg_file[:-4])

def cluster_extract(cluster_number):
	file_open = open(dlg_file, "r").readlines()

	conformations = []
	for entry in file_open:
		if "RANKING" in entry:
			entry_split = entry.split()
			if int(entry_split[0]) == cluster_number:
				conformations.append(str(entry_split[2]))
			else:
				None

	return (conformations)


def Read_file(runs,model_number):
	file_open = open(dlg_file, "r").readlines()
	count = 0
	n = 0
	conformation_number = []
	while n < len(file_open):
		file_split = file_open[n].split()
		if len(file_split) >= 1 \
				and (file_split[0] == "DOCKED:" and file_split[1] == "ATOM"):
			runs.append(file_open[n])
		if len(file_split) >= 1 and file_split[0] == "Run:":
			count += 1
			conformation_number.append(str(count))
		n += 1


	if cluster_extraction == "n" or cluster_extraction == "N":
		Split_to_pdb(runs, count,dlg_file[:-4],model_number)
		print(count)
	elif cluster_extraction == "y" or cluster_extraction == "Y":
		conformations = cluster_extract(cluster_number)
		for model_number in conformations:
			Split_to_pdb(runs, count,dlg_file[:-4] + "/" + "Cluster_" + str(cluster_number),model_number)
		print(str(len(conformations)) + "/" + str(count) + " conformations")


def Split_to_pdb(runs, count,folder_name,model_number):
	total_atoms = 0
	n = 0
	total_atoms = len(runs)
	if model_number == "0":
		writing = open(folder_name + "/all_models.pdb", "w")
		models = ""
		mdl_number = 1
		m = 0
		while m < len(runs):
			models += ("MODEL	" + str(mdl_number) + "\n")
			n = 0
			while n < int(total_atoms/count):
				models += (runs[m+n][8:])
				n += 1
			mdl_number += 1
			models += ("ENDMDL\n")
			m += int(total_atoms/count)
		writing.write(models)
	elif model_number == "All" or model_number == "ALL"\
			or model_number == "all":
		mdl_number = 1
		m = 0
		while m < len(runs):
			writing = open(folder_name + "/" + str(mdl_number) +
						   "_" + str(count) + ".pdb", "w")
			models = ""
			models += ("MODEL	" + str(mdl_number) + "\n")
			n = 0
			while n < int(total_atoms/count):
				models += (runs[m+n][8:])
				n += 1
			mdl_number += 1
			models += ("ENDMDL\n")
			writing.write(models)
			writing.close()
			m += int(total_atoms/count)
	else:
		mdl_number = 1
		m = int((int(model_number) * (total_atoms/count)) -
				(total_atoms/count))
		while m < int(model_number) * (total_atoms/count):
			writing = open(folder_name + "/" + str(model_number) +
						   "_" + str(count) + ".pdb", "w")
			models = ""
			models += ("MODEL	" + str(model_number) + "\n")
			n = 0
			while n < int(total_atoms/count):
				models += (runs[m+n][8:])
				n += 1
			mdl_number += 1
			models += ("ENDMDL\n")
			writing.write(models)
			writing.close()
			m += int(total_atoms/count)


if cluster_extraction == "n" or cluster_extraction == "N":
	Read_file(runs,model_number)
elif cluster_extraction == "y" or cluster_extraction == "Y":
	cluster_number = int(input("Cluster number :"))
	os.system("mkdir " + dlg_file[:-4] + "/Cluster_" + str(cluster_number))
	Read_file(runs,"0")