-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathmain.py
147 lines (128 loc) · 6.92 KB
/
main.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
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
132
133
134
135
136
137
138
import os
import mrcfile
import numpy as np
from ops.argparser import argparser
from ops.os_operation import mkdir
import time
def init_save_path(origin_map_path):
save_path = os.path.join(os.getcwd(), 'Predict_Result')
mkdir(save_path)
map_name = os.path.split(origin_map_path)[1].replace(".mrc", "")
map_name = map_name.replace(".map", "")
map_name = map_name.replace("(","").replace(")","")
save_path = os.path.join(save_path, map_name)
mkdir(save_path)
return save_path,map_name
if __name__ == "__main__":
params = argparser()
if params['mode']==0:
#configure the running path
running_dir = os.path.dirname(os.path.abspath(__file__))
gpu_id = params['gpu']
if gpu_id is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
cur_map_path = os.path.abspath(params['F'])
#process the map path if it's ending with .gz
if ".gz"==cur_map_path[-3:]:
from ops.os_operation import unzip_gz
cur_map_path = unzip_gz(cur_map_path)
model_dir = os.path.abspath(params['M'])
if params['resolution']<3 or params['resolution']>5:
if params['resolution']<3:
model_dir_candidate = os.path.join(running_dir,"best_model_0_3A")
else:
model_dir_candidate = os.path.join(running_dir,"best_model_5_10A")
model_1st_stage_path = os.path.join(model_dir_candidate,"stage1_network.pth")
model_2nd_stage_path = os.path.join(model_dir_candidate,"stage2_network.pth")
if os.path.exists(model_1st_stage_path) and os.path.exists(model_2nd_stage_path):
model_dir = model_dir_candidate
print("with resolution %f"%params['resolution'])
print("use resolution model in %s"%model_dir)
if params['resolution']>10:
print("!!!Warning!!!Please provide map at resolution lower than 10A to run!")
exit()
if params['prediction_only'] or params['no_seqinfo']:
fasta_path = None
else:
fasta_path = os.path.abspath(params['P'])
if params['output'] is None:
save_path,map_name = init_save_path(cur_map_path)
else:
save_path=params['output']
map_name="input_cryoread"
mkdir(save_path)
os.chdir(running_dir)
from data_processing.Unify_Map import Unify_Map
cur_map_path = Unify_Map(cur_map_path,os.path.join(save_path,map_name+"_unified.mrc"))
from data_processing.Resize_Map import Resize_Map
cur_map_path = Resize_Map(cur_map_path,os.path.join(save_path,map_name+".mrc"))
if params['contour']<0:
#change contour level to 0 and increase all the density
from data_processing.map_utils import increase_map_density
cur_map_path = increase_map_density(cur_map_path,os.path.join(save_path,map_name+"_increase.mrc"),params['contour'])
params['contour']=0
#segment map to remove those useless regions
from data_processing.map_utils import segment_map
cur_new_map_path = os.path.join(save_path,map_name+"_segment.mrc")
cur_map_path = segment_map(cur_map_path,cur_new_map_path,contour=1e-6) #save the final prediction prob array space
#do a pre check to notify user errors for contour
with mrcfile.open(cur_map_path,permissive=True) as mrc:
data=mrc.data
if np.max(data)<=params['contour']:
print("!!!Warning!!!Please provide contour level lower than maximum density value to run!")
#exit()
from data_processing.map_utils import automate_contour
params['contour']=automate_contour(data)
print("!!!Warning!!!reset contour level: %f"%params['contour'])
from predict.predict_1st_stage import Predict_1st_Stage
#1st stage cascad prediciton
model_1st_stage_path = os.path.join(model_dir,"stage1_network.pth")
save_path_1st_stage = os.path.join(save_path,"1st_stage_detection")
mkdir(save_path_1st_stage)
Predict_1st_Stage(cur_map_path,model_1st_stage_path,save_path_1st_stage,
params['box_size'],params['stride'],params['batch_size'],params['contour'],params)
#2nd stage refine prediction
from predict.predict_2nd_stage import Predict_2nd_Stage
model_2nd_stage_path = os.path.join(model_dir,"stage2_network.pth")
save_path_2nd_stage = os.path.join(save_path,"2nd_stage_detection")
mkdir(save_path_2nd_stage)
Predict_2nd_Stage(cur_map_path,save_path_1st_stage,model_2nd_stage_path,save_path_2nd_stage,
params['box_size'],params['stride'],params['batch_size'],params['contour'],params)
#0 gen protein map for other programs to run
cur_predict_path = os.path.join(save_path_2nd_stage,"Input")
chain_predict_path = os.path.join(cur_predict_path,"chain_predictprob.npy")
chain_prob = np.load(chain_predict_path)#[sugar,phosphate,A,UT,C,G,protein,base]
mask_map_path = os.path.join(save_path,"mask_protein.mrc")
from data_processing.Gen_MaskDRNA_map import Gen_MaskProtein_map
Gen_MaskProtein_map(chain_prob,cur_map_path,mask_map_path,params['contour'],threshold=0.3)
if params['prediction_only']:
print("Our prediction results are saved in %s with mrc format for visualization check."%save_path_2nd_stage)
exit()
#graph based atomic structure modeling
gaussian_bandwidth = params['g'] #use 3
dcut = params['m']#after meanshifting merge points distance<[float]
rdcut = params['f']#remove ldp threshold 0.01
from graph.Build_Unet_Graph import Build_Unet_Graph
graph_save_path = os.path.join(save_path,"graph_atomic_modeling")
mkdir(graph_save_path)
Build_Unet_Graph(cur_map_path,chain_predict_path,fasta_path,graph_save_path,
gaussian_bandwidth,dcut, rdcut,params)
elif params['mode']==1:
#structure refinement pipeline
input_pdb = os.path.abspath(params['F'])
input_map = os.path.abspath(params['M'])
output_dir = os.path.abspath(params['P'])
running_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(running_dir)
#resolution param also needed
from graph.refine_structure import refine_structure
refine_structure(input_pdb,input_map,output_dir,params)
elif params['mode']==2:
#add evaluation script for predicted structure and native structure
# python3 main.py --mode=2 -F=predicted.cif[.pdb] -M=target.cif[.pdb]
query_pdb = os.path.abspath(params['F']) # predicted pdb/cif file
target_pdb = os.path.abspath(params['M']) # native pdb/cif file
cutoff = 5.0
#output all the metrics reported in CryoREAD's paper
from evaluation.evaluate_structure import evaluate_structure
evaluate_structure(query_pdb,target_pdb,cutoff)