danielkovtun commited on
Commit
f17a933
·
1 Parent(s): 8aac36a

chore: lint

Browse files
Files changed (1) hide show
  1. inference_app.py +154 -134
inference_app.py CHANGED
@@ -18,58 +18,55 @@ from gradio_molecule3d import Molecule3D
18
 
19
  EVAL_METRICS = ["system", "L_rms", "I_rms", "F_nat", "DOCKQ", "CAPRI_class"]
20
 
 
21
  def predict(
22
- receptor_pdb: Path,
23
- ligand_pdb: Path,
24
- receptor_fasta: Path | None = None,
25
  ligand_fasta: Path | None = None,
26
  ) -> tuple[str, float]:
27
  start_time = time.time()
28
  # Do inference here
29
- # return an output pdb file with the protein and two chains R and L.
30
  receptor = atoms.atom_array_from_pdb_file(receptor_pdb, extra_fields=["b_factor"])
31
  ligand = atoms.atom_array_from_pdb_file(ligand_pdb, extra_fields=["b_factor"])
32
  receptor = atoms.normalize_orientation(receptor)
33
  ligand = atoms.normalize_orientation(ligand)
34
-
35
  # Number of random poses to generate
36
  M = 50
37
- # Inititalize an empty stack with shape (m x n x 3)
38
  stack = AtomArrayStack(M, ligand.shape[0])
39
-
40
  # copy annotations from ligand
41
  for annot in ligand.get_annotation_categories():
42
  stack.set_annotation(annot, np.copy(ligand.get_annotation(annot)))
43
-
44
  # Random translations sampled along 0-50 angstroms per axis
45
- translation_magnitudes = np.linspace(
46
- 0, 26,
47
- num=26,
48
- endpoint=False
49
- )
50
  # generate one pose at a time
51
  for i in range(M):
52
  q = R.random()
53
  translation_vec = [
54
- random.choice(translation_magnitudes), # x
55
- random.choice(translation_magnitudes), # y
56
- random.choice(translation_magnitudes), # z
57
  ]
58
  # transform the ligand chain
59
  stack.coord[i, ...] = q.apply(ligand.coord) + translation_vec
60
-
61
  # Find clashes (1.2 A contact radius)
62
  stack_conts = get_stack_contacts(receptor, stack, threshold=1.2)
63
-
64
  # Keep the "best" pose based on pose w/fewest clashes
65
  pose_clashes = []
66
  for i in range(stack_conts.shape[0]):
67
  pose_conts = stack_conts[i]
68
  pose_clashes.append((i, np.argwhere(pose_conts != -1).shape[0]))
69
-
70
  best_pose_idx = sorted(pose_clashes, key=lambda x: x[1])[0][0]
71
  best_pose = receptor + stack[best_pose_idx]
72
-
73
  output_dir = Path(receptor_pdb).parent
74
  # System ID
75
  pdb_name = Path(receptor_pdb).stem + "--" + Path(ligand_pdb).name
@@ -81,8 +78,8 @@ def predict(
81
 
82
 
83
  def evaluate(
84
- system_id: str,
85
- prediction_pdb: Path,
86
  ) -> tuple[pd.DataFrame, float]:
87
  start_time = time.time()
88
  system = PinderSystem(system_id)
@@ -90,7 +87,16 @@ def evaluate(
90
  bdq = BiotiteDockQ(native, Path(prediction_pdb), parallel_io=False)
91
  metrics = bdq.calculate()
92
  metrics = metrics[["system", "LRMS", "iRMS", "Fnat", "DockQ", "CAPRI"]].copy()
93
- metrics.rename(columns={"LRMS": "L_rms", "iRMS": "I_rms", "Fnat": "F_nat", "DockQ": "DOCKQ", "CAPRI": "CAPRI_class"}, inplace=True)
 
 
 
 
 
 
 
 
 
94
  end_time = time.time()
95
  run_time = end_time - start_time
96
  pred = Structure(Path(prediction_pdb))
@@ -102,117 +108,131 @@ def evaluate(
102
 
103
  with gr.Blocks() as app:
104
  with gr.Tab("🧬 PINDER inference template"):
105
- gr.Markdown("Title, description, and other information about the model")
106
- with gr.Row():
107
- with gr.Column():
108
- input_protein_1 = gr.File(label="Input Protein 1 monomer (PDB)")
109
- input_fasta_1 = gr.File(label="Input Protein 1 monomer sequence (FASTA)")
110
- with gr.Column():
111
- input_protein_2 = gr.File(label="Input Protein 2 monomer (PDB)")
112
- input_fasta_2 = gr.File(label="Input Protein 2 monomer sequence (FASTA)")
113
-
114
-
115
-
116
- # define any options here
117
-
118
- # for automated inference the default options are used
119
- # slider_option = gr.Slider(0,10, label="Slider Option")
120
- # checkbox_option = gr.Checkbox(label="Checkbox Option")
121
- # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")
122
-
123
- btn = gr.Button("Run Inference")
124
-
125
- gr.Examples(
126
- [
127
- [
128
- "8i5w_R.pdb",
129
- "8i5w_R.fasta",
130
- "8i5w_L.pdb",
131
- "8i5w_L.fasta",
132
- ],
133
- ],
134
- [input_protein_1, input_fasta_1, input_protein_2, input_fasta_2],
135
- )
136
- reps = [
137
- {
138
- "model": 0,
139
- "style": "cartoon",
140
- "chain": "R",
141
- "color": "whiteCarbon",
142
- },
143
- {
144
- "model": 0,
145
- "style": "cartoon",
146
- "chain": "L",
147
- "color": "greenCarbon",
148
- },
149
- {
150
- "model": 0,
151
- "chain": "R",
152
- "style": "stick",
153
- "sidechain": True,
154
- "color": "whiteCarbon",
155
- },
156
- {
157
- "model": 0,
158
- "chain": "L",
159
- "style": "stick",
160
- "sidechain": True,
161
- "color": "greenCarbon"
162
- }
163
- ]
164
-
165
- out = Molecule3D(reps=reps)
166
- run_time = gr.Textbox(label="Runtime")
167
-
168
- btn.click(predict, inputs=[input_protein_1, input_protein_2, input_fasta_1, input_fasta_2], outputs=[out, run_time])
 
 
 
 
 
 
169
  with gr.Tab("⚖️ PINDER evaluation template"):
170
- with gr.Row():
171
- with gr.Column():
172
- input_system_id = gr.Textbox(label="PINDER system ID")
173
- input_prediction_pdb = gr.File(label="Top ranked prediction (PDB with chains R and L)")
174
-
175
- eval_btn = gr.Button("Run Evaluation")
176
- gr.Examples(
177
- [
178
- [
179
- "3g9w__A1_Q71LX4--3g9w__D1_P05556",
180
- "3g9w_R--3g9w_L.pdb",
181
- ],
182
- ],
183
- [input_system_id, input_prediction_pdb],
184
- )
185
- reps = [
186
- {
187
- "model": 0,
188
- "style": "cartoon",
189
- "chain": "R",
190
- "color": "greenCarbon",
191
- },
192
- {
193
- "model": 0,
194
- "style": "cartoon",
195
- "chain": "L",
196
- "color": "cyanCarbon",
197
- },
198
- {
199
- "model": 1,
200
- "style": "cartoon",
201
- "chain": "R",
202
- "color": "grayCarbon",
203
- },
204
- {
205
- "model": 1,
206
- "style": "cartoon",
207
- "chain": "L",
208
- "color": "blueCarbon",
209
- },
210
- ]
211
-
212
- pred_native = Molecule3D(reps=reps, config={"backgroundColor": "black"})
213
- eval_run_time = gr.Textbox(label="Evaluation runtime")
214
- metric_table = gr.DataFrame(pd.DataFrame([], columns=EVAL_METRICS),label="Evaluation metrics")
215
-
216
- eval_btn.click(evaluate, inputs=[input_system_id, input_prediction_pdb], outputs=[metric_table, pred_native, eval_run_time])
 
 
 
 
 
 
 
 
217
 
218
  app.launch()
 
18
 
19
  EVAL_METRICS = ["system", "L_rms", "I_rms", "F_nat", "DOCKQ", "CAPRI_class"]
20
 
21
+
22
  def predict(
23
+ receptor_pdb: Path,
24
+ ligand_pdb: Path,
25
+ receptor_fasta: Path | None = None,
26
  ligand_fasta: Path | None = None,
27
  ) -> tuple[str, float]:
28
  start_time = time.time()
29
  # Do inference here
30
+ # return an output pdb file with the protein and two chains R and L.
31
  receptor = atoms.atom_array_from_pdb_file(receptor_pdb, extra_fields=["b_factor"])
32
  ligand = atoms.atom_array_from_pdb_file(ligand_pdb, extra_fields=["b_factor"])
33
  receptor = atoms.normalize_orientation(receptor)
34
  ligand = atoms.normalize_orientation(ligand)
35
+
36
  # Number of random poses to generate
37
  M = 50
38
+ # Inititalize an empty stack with shape (m x n x 3)
39
  stack = AtomArrayStack(M, ligand.shape[0])
40
+
41
  # copy annotations from ligand
42
  for annot in ligand.get_annotation_categories():
43
  stack.set_annotation(annot, np.copy(ligand.get_annotation(annot)))
44
+
45
  # Random translations sampled along 0-50 angstroms per axis
46
+ translation_magnitudes = np.linspace(0, 26, num=26, endpoint=False)
 
 
 
 
47
  # generate one pose at a time
48
  for i in range(M):
49
  q = R.random()
50
  translation_vec = [
51
+ random.choice(translation_magnitudes), # x
52
+ random.choice(translation_magnitudes), # y
53
+ random.choice(translation_magnitudes), # z
54
  ]
55
  # transform the ligand chain
56
  stack.coord[i, ...] = q.apply(ligand.coord) + translation_vec
57
+
58
  # Find clashes (1.2 A contact radius)
59
  stack_conts = get_stack_contacts(receptor, stack, threshold=1.2)
60
+
61
  # Keep the "best" pose based on pose w/fewest clashes
62
  pose_clashes = []
63
  for i in range(stack_conts.shape[0]):
64
  pose_conts = stack_conts[i]
65
  pose_clashes.append((i, np.argwhere(pose_conts != -1).shape[0]))
66
+
67
  best_pose_idx = sorted(pose_clashes, key=lambda x: x[1])[0][0]
68
  best_pose = receptor + stack[best_pose_idx]
69
+
70
  output_dir = Path(receptor_pdb).parent
71
  # System ID
72
  pdb_name = Path(receptor_pdb).stem + "--" + Path(ligand_pdb).name
 
78
 
79
 
80
  def evaluate(
81
+ system_id: str,
82
+ prediction_pdb: Path,
83
  ) -> tuple[pd.DataFrame, float]:
84
  start_time = time.time()
85
  system = PinderSystem(system_id)
 
87
  bdq = BiotiteDockQ(native, Path(prediction_pdb), parallel_io=False)
88
  metrics = bdq.calculate()
89
  metrics = metrics[["system", "LRMS", "iRMS", "Fnat", "DockQ", "CAPRI"]].copy()
90
+ metrics.rename(
91
+ columns={
92
+ "LRMS": "L_rms",
93
+ "iRMS": "I_rms",
94
+ "Fnat": "F_nat",
95
+ "DockQ": "DOCKQ",
96
+ "CAPRI": "CAPRI_class",
97
+ },
98
+ inplace=True,
99
+ )
100
  end_time = time.time()
101
  run_time = end_time - start_time
102
  pred = Structure(Path(prediction_pdb))
 
108
 
109
  with gr.Blocks() as app:
110
  with gr.Tab("🧬 PINDER inference template"):
111
+ gr.Markdown("Title, description, and other information about the model")
112
+ with gr.Row():
113
+ with gr.Column():
114
+ input_protein_1 = gr.File(label="Input Protein 1 monomer (PDB)")
115
+ input_fasta_1 = gr.File(
116
+ label="Input Protein 1 monomer sequence (FASTA)"
117
+ )
118
+ with gr.Column():
119
+ input_protein_2 = gr.File(label="Input Protein 2 monomer (PDB)")
120
+ input_fasta_2 = gr.File(
121
+ label="Input Protein 2 monomer sequence (FASTA)"
122
+ )
123
+
124
+ # define any options here
125
+
126
+ # for automated inference the default options are used
127
+ # slider_option = gr.Slider(0,10, label="Slider Option")
128
+ # checkbox_option = gr.Checkbox(label="Checkbox Option")
129
+ # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")
130
+
131
+ btn = gr.Button("Run Inference")
132
+
133
+ gr.Examples(
134
+ [
135
+ [
136
+ "8i5w_R.pdb",
137
+ "8i5w_R.fasta",
138
+ "8i5w_L.pdb",
139
+ "8i5w_L.fasta",
140
+ ],
141
+ ],
142
+ [input_protein_1, input_fasta_1, input_protein_2, input_fasta_2],
143
+ )
144
+ reps = [
145
+ {
146
+ "model": 0,
147
+ "style": "cartoon",
148
+ "chain": "R",
149
+ "color": "whiteCarbon",
150
+ },
151
+ {
152
+ "model": 0,
153
+ "style": "cartoon",
154
+ "chain": "L",
155
+ "color": "greenCarbon",
156
+ },
157
+ {
158
+ "model": 0,
159
+ "chain": "R",
160
+ "style": "stick",
161
+ "sidechain": True,
162
+ "color": "whiteCarbon",
163
+ },
164
+ {
165
+ "model": 0,
166
+ "chain": "L",
167
+ "style": "stick",
168
+ "sidechain": True,
169
+ "color": "greenCarbon",
170
+ },
171
+ ]
172
+
173
+ out = Molecule3D(reps=reps)
174
+ run_time = gr.Textbox(label="Runtime")
175
+
176
+ btn.click(
177
+ predict,
178
+ inputs=[input_protein_1, input_protein_2, input_fasta_1, input_fasta_2],
179
+ outputs=[out, run_time],
180
+ )
181
  with gr.Tab("⚖️ PINDER evaluation template"):
182
+ with gr.Row():
183
+ with gr.Column():
184
+ input_system_id = gr.Textbox(label="PINDER system ID")
185
+ input_prediction_pdb = gr.File(
186
+ label="Top ranked prediction (PDB with chains R and L)"
187
+ )
188
+
189
+ eval_btn = gr.Button("Run Evaluation")
190
+ gr.Examples(
191
+ [
192
+ [
193
+ "3g9w__A1_Q71LX4--3g9w__D1_P05556",
194
+ "3g9w_R--3g9w_L.pdb",
195
+ ],
196
+ ],
197
+ [input_system_id, input_prediction_pdb],
198
+ )
199
+ reps = [
200
+ {
201
+ "model": 0,
202
+ "style": "cartoon",
203
+ "chain": "R",
204
+ "color": "greenCarbon",
205
+ },
206
+ {
207
+ "model": 0,
208
+ "style": "cartoon",
209
+ "chain": "L",
210
+ "color": "cyanCarbon",
211
+ },
212
+ {
213
+ "model": 1,
214
+ "style": "cartoon",
215
+ "chain": "R",
216
+ "color": "grayCarbon",
217
+ },
218
+ {
219
+ "model": 1,
220
+ "style": "cartoon",
221
+ "chain": "L",
222
+ "color": "blueCarbon",
223
+ },
224
+ ]
225
+
226
+ pred_native = Molecule3D(reps=reps, config={"backgroundColor": "black"})
227
+ eval_run_time = gr.Textbox(label="Evaluation runtime")
228
+ metric_table = gr.DataFrame(
229
+ pd.DataFrame([], columns=EVAL_METRICS), label="Evaluation metrics"
230
+ )
231
+
232
+ eval_btn.click(
233
+ evaluate,
234
+ inputs=[input_system_id, input_prediction_pdb],
235
+ outputs=[metric_table, pred_native, eval_run_time],
236
+ )
237
 
238
  app.launch()