foreversheikh commited on
Commit
430159c
Β·
verified Β·
1 Parent(s): ac574b7

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +649 -34
src/streamlit_app.py CHANGED
@@ -1,40 +1,655 @@
1
- import altair as alt
2
- import numpy as np
3
- import pandas as pd
 
 
 
 
 
 
 
 
 
4
  import streamlit as st
 
 
 
5
 
6
- """
7
- # Welcome to Streamlit!
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
 
9
- Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
10
- If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
11
- forums](https://discuss.streamlit.io).
 
 
 
 
 
 
 
 
 
 
12
 
13
- In the meantime, below is an example of what you can do with just a few lines of code:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
 
16
- num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
17
- num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
18
-
19
- indices = np.linspace(0, 1, num_points)
20
- theta = 2 * np.pi * num_turns * indices
21
- radius = indices
22
-
23
- x = radius * np.cos(theta)
24
- y = radius * np.sin(theta)
25
-
26
- df = pd.DataFrame({
27
- "x": x,
28
- "y": y,
29
- "idx": indices,
30
- "rand": np.random.randn(num_points),
31
- })
32
-
33
- st.altair_chart(alt.Chart(df, height=700, width=700)
34
- .mark_point(filled=True)
35
- .encode(
36
- x=alt.X("x", axis=None),
37
- y=alt.Y("y", axis=None),
38
- color=alt.Color("idx", legend=None, scale=alt.Scale()),
39
- size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
40
- ))
 
1
+ # -*- coding: utf-8 -*-
2
+ """
3
+ AfDesign Peptide Binder Design β€” Streamlit App
4
+ Converted from ColabDesign notebook (binder_design.py)
5
+ """
6
+
7
+ import os
8
+ import re
9
+ import warnings
10
+ import tempfile
11
+ import subprocess
12
+ import tarfile
13
  import streamlit as st
14
+ import numpy as np
15
+ import requests
16
+ import plotly.express as px
17
 
18
+ warnings.simplefilter(action='ignore', category=FutureWarning)
19
+
20
+ # ───────────────────────────────────────────────────────────────────────
21
+ # Page config & custom CSS
22
+ # ───────────────────────────────────────────────────────────────────────
23
+ st.set_page_config(
24
+ page_title="AfDesign Β· Peptide Binder Designer",
25
+ page_icon="🧬",
26
+ layout="wide",
27
+ initial_sidebar_state="expanded",
28
+ )
29
+
30
+ CUSTOM_CSS = """
31
+ <style>
32
+ @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
33
+
34
+ :root {
35
+ --bg-primary: #0d1117;
36
+ --bg-secondary: #161b22;
37
+ --bg-card: rgba(22, 27, 34, 0.85);
38
+ --border-color: rgba(48, 54, 61, 0.7);
39
+ --accent: #58a6ff;
40
+ --accent-glow: rgba(88, 166, 255, 0.25);
41
+ --accent-secondary: #7ee787;
42
+ --text-primary: #e6edf3;
43
+ --text-secondary: #8b949e;
44
+ --gradient-1: linear-gradient(135deg, #58a6ff 0%, #bc8cff 50%, #f778ba 100%);
45
+ --gradient-2: linear-gradient(135deg, #1a1f35 0%, #0d1117 100%);
46
+ }
47
+
48
+ html, body, [class*="css"] {
49
+ font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
50
+ }
51
+
52
+ .stApp {
53
+ background: var(--gradient-2);
54
+ }
55
+
56
+ /* Sidebar */
57
+ section[data-testid="stSidebar"] {
58
+ background: var(--bg-secondary) !important;
59
+ border-right: 1px solid var(--border-color);
60
+ }
61
+
62
+ section[data-testid="stSidebar"] .stMarkdown h3 {
63
+ background: var(--gradient-1);
64
+ -webkit-background-clip: text;
65
+ -webkit-text-fill-color: transparent;
66
+ font-weight: 700;
67
+ letter-spacing: -0.02em;
68
+ margin-top: 1.2rem;
69
+ }
70
+
71
+ /* Hero header */
72
+ .hero-title {
73
+ font-size: 2.4rem;
74
+ font-weight: 700;
75
+ background: var(--gradient-1);
76
+ -webkit-background-clip: text;
77
+ -webkit-text-fill-color: transparent;
78
+ margin-bottom: 0;
79
+ letter-spacing: -0.03em;
80
+ animation: fadeInUp 0.8s ease-out;
81
+ }
82
+
83
+ .hero-sub {
84
+ color: var(--text-secondary);
85
+ font-size: 1.05rem;
86
+ margin-top: 0.3rem;
87
+ animation: fadeInUp 1s ease-out;
88
+ }
89
+
90
+ /* Glass cards */
91
+ .glass-card {
92
+ background: var(--bg-card);
93
+ backdrop-filter: blur(16px);
94
+ -webkit-backdrop-filter: blur(16px);
95
+ border: 1px solid var(--border-color);
96
+ border-radius: 16px;
97
+ padding: 1.5rem 1.8rem;
98
+ margin-bottom: 1.2rem;
99
+ transition: transform 0.25s ease, box-shadow 0.25s ease;
100
+ }
101
+ .glass-card:hover {
102
+ transform: translateY(-2px);
103
+ box-shadow: 0 8px 32px rgba(0, 0, 0, 0.35);
104
+ }
105
+
106
+ .card-title {
107
+ font-size: 1.15rem;
108
+ font-weight: 600;
109
+ color: var(--accent);
110
+ margin-bottom: 0.8rem;
111
+ display: flex;
112
+ align-items: center;
113
+ gap: 0.5rem;
114
+ }
115
+
116
+ /* Metric pills */
117
+ .metric-row {
118
+ display: flex;
119
+ gap: 1rem;
120
+ flex-wrap: wrap;
121
+ margin-bottom: 1rem;
122
+ }
123
+ .metric-pill {
124
+ background: rgba(88, 166, 255, 0.08);
125
+ border: 1px solid rgba(88, 166, 255, 0.2);
126
+ border-radius: 12px;
127
+ padding: 0.9rem 1.4rem;
128
+ min-width: 150px;
129
+ flex: 1;
130
+ text-align: center;
131
+ transition: border-color 0.3s ease;
132
+ }
133
+ .metric-pill:hover {
134
+ border-color: var(--accent);
135
+ }
136
+ .metric-pill .label {
137
+ font-size: 0.75rem;
138
+ text-transform: uppercase;
139
+ letter-spacing: 0.08em;
140
+ color: var(--text-secondary);
141
+ margin-bottom: 0.3rem;
142
+ }
143
+ .metric-pill .value {
144
+ font-size: 1.5rem;
145
+ font-weight: 700;
146
+ color: var(--text-primary);
147
+ }
148
 
149
+ /* Sequence display */
150
+ .seq-box {
151
+ font-family: 'Courier New', monospace;
152
+ font-size: 1.15rem;
153
+ letter-spacing: 0.15em;
154
+ word-break: break-all;
155
+ background: rgba(88,166,255,0.06);
156
+ border: 1px solid rgba(88,166,255,0.15);
157
+ border-radius: 12px;
158
+ padding: 1rem 1.2rem;
159
+ color: var(--accent-secondary);
160
+ line-height: 1.8;
161
+ }
162
 
163
+ /* Buttons */
164
+ .stButton > button {
165
+ background: var(--gradient-1) !important;
166
+ color: #fff !important;
167
+ border: none !important;
168
+ border-radius: 12px !important;
169
+ padding: 0.75rem 2rem !important;
170
+ font-weight: 600 !important;
171
+ font-size: 1rem !important;
172
+ letter-spacing: 0.02em;
173
+ transition: opacity 0.3s ease, transform 0.2s ease !important;
174
+ }
175
+ .stButton > button:hover {
176
+ opacity: 0.9 !important;
177
+ transform: translateY(-1px) !important;
178
+ }
179
+
180
+ /* Download button */
181
+ .stDownloadButton > button {
182
+ background: rgba(126, 231, 135, 0.12) !important;
183
+ border: 1px solid rgba(126, 231, 135, 0.35) !important;
184
+ color: var(--accent-secondary) !important;
185
+ border-radius: 12px !important;
186
+ font-weight: 600 !important;
187
+ transition: background 0.3s ease !important;
188
+ }
189
+ .stDownloadButton > button:hover {
190
+ background: rgba(126, 231, 135, 0.22) !important;
191
+ }
192
+
193
+ /* Animations */
194
+ @keyframes fadeInUp {
195
+ from { opacity: 0; transform: translateY(20px); }
196
+ to { opacity: 1; transform: translateY(0); }
197
+ }
198
+
199
+ /* Status indicator */
200
+ .status-dot {
201
+ display: inline-block;
202
+ width: 8px; height: 8px;
203
+ border-radius: 50%;
204
+ margin-right: 6px;
205
+ animation: pulse 2s infinite;
206
+ }
207
+ .status-dot.ready { background: var(--accent-secondary); }
208
+ .status-dot.running { background: #f0883e; }
209
+
210
+ @keyframes pulse {
211
+ 0%, 100% { opacity: 1; }
212
+ 50% { opacity: 0.4; }
213
+ }
214
+ </style>
215
  """
216
+ st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
217
+
218
+ # ───────────────────────────────────────────────────────────────────────
219
+ # Helper: fetch PDB file
220
+ # ───────────────────────────────────────────────────────────────────────
221
+ PARAMS_DIR = "params"
222
+
223
+ @st.cache_data(show_spinner=False)
224
+ def fetch_pdb(pdb_code: str) -> str:
225
+ """Download a PDB from RCSB or AlphaFoldDB and return the local path."""
226
+ if os.path.isfile(pdb_code):
227
+ return pdb_code
228
+ if len(pdb_code) == 4:
229
+ url = f"https://files.rcsb.org/view/{pdb_code}.pdb"
230
+ out_path = f"{pdb_code}.pdb"
231
+ else:
232
+ url = f"https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb"
233
+ out_path = f"AF-{pdb_code}-F1-model_v3.pdb"
234
+ if not os.path.isfile(out_path):
235
+ resp = requests.get(url, timeout=60)
236
+ resp.raise_for_status()
237
+ with open(out_path, "w") as f:
238
+ f.write(resp.text)
239
+ return out_path
240
+
241
+
242
+ def download_params(status_container):
243
+ """Download and extract AlphaFold parameters if not present."""
244
+ if os.path.isdir(PARAMS_DIR) and len(os.listdir(PARAMS_DIR)) > 0:
245
+ return True
246
+
247
+ PARAMS_URL = "https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar"
248
+ TAR_FILE = "alphafold_params_2022-12-06.tar"
249
+
250
+ try:
251
+ os.makedirs(PARAMS_DIR, exist_ok=True)
252
+
253
+ # Download with progress
254
+ status_container.write("πŸ“¦ Downloading AlphaFold parameters (~3.5 GB)... This only happens once.")
255
+ progress_bar = status_container.progress(0, text="Downloading...")
256
+
257
+ resp = requests.get(PARAMS_URL, stream=True, timeout=600)
258
+ resp.raise_for_status()
259
+ total = int(resp.headers.get('content-length', 0))
260
+ downloaded = 0
261
+
262
+ with open(TAR_FILE, 'wb') as f:
263
+ for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024): # 8MB chunks
264
+ f.write(chunk)
265
+ downloaded += len(chunk)
266
+ if total > 0:
267
+ pct = min(downloaded / total, 1.0)
268
+ progress_bar.progress(pct, text=f"Downloaded {downloaded / 1e9:.1f} / {total / 1e9:.1f} GB")
269
+
270
+ progress_bar.progress(1.0, text="Download complete!")
271
+
272
+ # Extract
273
+ status_container.write("πŸ“‚ Extracting parameters...")
274
+ with tarfile.open(TAR_FILE, 'r') as tar:
275
+ tar.extractall(path=PARAMS_DIR)
276
+
277
+ # Cleanup tar
278
+ if os.path.isfile(TAR_FILE):
279
+ os.remove(TAR_FILE)
280
+
281
+ status_container.write("βœ… AlphaFold parameters ready!")
282
+ return True
283
+
284
+ except Exception as e:
285
+ status_container.error(f"Failed to download AlphaFold parameters: {e}")
286
+ return False
287
+
288
+
289
+ # ───────────────────────────────────────────────────────────────────────
290
+ # Hero header
291
+ # ───────────────────────────────────────────────────────────────────────
292
+ st.markdown('<p class="hero-title">🧬 AfDesign · Peptide Binder Designer</p>', unsafe_allow_html=True)
293
+ st.markdown(
294
+ '<p class="hero-sub">'
295
+ 'Generate / hallucinate a protein binder sequence that AlphaFold predicts will bind your target structure. '
296
+ 'Maximises interface contacts and binder pLDDT.'
297
+ '</p>',
298
+ unsafe_allow_html=True,
299
+ )
300
+
301
+ # ───────────────────────────────────────────────────���───────────────────
302
+ # Sidebar β€” Inputs
303
+ # ───────────────────────────────────────────────────────────────────────
304
+ with st.sidebar:
305
+ st.markdown("### 🎯 Target Info")
306
+ pdb_code = st.text_input(
307
+ "PDB / UniProt Code",
308
+ value="5F9R",
309
+ help="Enter a 4-letter PDB code (e.g. 5F9R), a UniProt code (to fetch from AlphaFoldDB), or leave blank and upload a file below.",
310
+ )
311
+ uploaded_pdb = st.file_uploader("Or upload a PDB file", type=["pdb"])
312
+ target_chain = st.text_input("Target Chain", value="B", help="Chain identifier for the target protein.")
313
+ target_hotspot = st.text_input(
314
+ "Target Hotspot",
315
+ value="",
316
+ help='Restrict loss to specific positions on the target (e.g. "1-10,12,15"). Leave blank for no restriction.',
317
+ )
318
+ target_flexible = st.toggle("Flexible Target Backbone", value=False, help="Allow backbone of target to be flexible during design.")
319
+
320
+ st.markdown("---")
321
+ st.markdown("### πŸ”— Binder Info")
322
+ binder_len = st.slider("Binder Length", min_value=10, max_value=200, value=25, step=1, help="Length of the binder peptide to hallucinate.")
323
+ binder_seq_input = st.text_input(
324
+ "Initial Binder Sequence",
325
+ value="",
326
+ help="Optional amino acid sequence to initialize the design. If provided, binder length is set to its length.",
327
+ )
328
+ binder_chain_input = st.text_input(
329
+ "Binder Chain (supervised)",
330
+ value="",
331
+ help="If defined, supervised loss is used and binder length is ignored.",
332
+ )
333
+
334
+ st.markdown("---")
335
+ st.markdown("### βš™οΈ Model Configuration")
336
+ use_multimer = st.toggle("Use AlphaFold-Multimer", value=False, help="Use the multimer model for design.")
337
+ num_recycles = st.select_slider("Number of Recycles", options=[0, 1, 3, 6], value=1)
338
+ num_models_sel = st.selectbox("Number of Models", options=["1", "2", "3", "4", "5", "all"], index=0, help="Number of trained models to use during optimization.")
339
+
340
+ st.markdown("---")
341
+ st.markdown("### πŸš€ Optimization")
342
+ optimizer = st.selectbox(
343
+ "Optimizer",
344
+ options=["pssm_semigreedy", "3stage", "semigreedy", "pssm", "logits", "soft", "hard"],
345
+ index=0,
346
+ help=(
347
+ "β€’ pssm_semigreedy β€” uses designed PSSM to bias semigreedy opt (recommended)\n"
348
+ "β€’ 3stage β€” gradient descent: logits β†’ soft β†’ hard\n"
349
+ "β€’ semigreedy β€” random mutations, accepts if loss decreases\n"
350
+ "‒ pssm — GD logits→soft for a sequence profile\n"
351
+ "β€’ logits / soft / hard β€” raw GD optimization"
352
+ ),
353
+ )
354
+
355
+ with st.expander("Advanced GD Settings", expanded=False):
356
+ gd_method = st.selectbox(
357
+ "GD Method",
358
+ options=[
359
+ "sgd", "adam", "adamw", "adabelief", "adafactor", "adagrad",
360
+ "fromage", "lamb", "lars", "noisy_sgd", "dpsgd", "radam",
361
+ "rmsprop", "sm3", "yogi",
362
+ ],
363
+ index=0,
364
+ )
365
+ learning_rate = st.number_input("Learning Rate", min_value=0.0001, max_value=10.0, value=0.1, step=0.01, format="%.4f")
366
+ norm_seq_grad = st.toggle("Normalize Sequence Gradient", value=True)
367
+ dropout = st.toggle("Dropout", value=True)
368
+
369
+ st.markdown("---")
370
+ st.markdown("### 🎨 Visualization")
371
+ color_mode = st.selectbox("Color Scheme", options=["pLDDT", "chain", "rainbow"], index=0)
372
+ show_sidechains = st.toggle("Show Sidechains", value=False)
373
+ show_mainchains = st.toggle("Show Mainchains", value=False)
374
+
375
+ # ───────────────────────────────────────────────────────────────────────
376
+ # Process inputs
377
+ # ───────────────────────────────────────────────────────────────────────
378
+ binder_seq = re.sub("[^A-Z]", "", binder_seq_input.upper()) if binder_seq_input else ""
379
+ if len(binder_seq) > 0:
380
+ binder_len_final = len(binder_seq)
381
+ else:
382
+ binder_seq = None
383
+ binder_len_final = binder_len
384
+
385
+ binder_chain = binder_chain_input if binder_chain_input.strip() else None
386
+ hotspot = target_hotspot if target_hotspot.strip() else None
387
+ num_models_int = 5 if num_models_sel == "all" else int(num_models_sel)
388
+
389
+ # ───────────────────────────────────────────────────────────────────────
390
+ # Summary cards
391
+ # ───────────────────────────────────────────────────────────────────────
392
+ col1, col2, col3 = st.columns(3)
393
+ with col1:
394
+ st.markdown(
395
+ f"""<div class="glass-card">
396
+ <div class="card-title">🎯 Target</div>
397
+ <div class="metric-row">
398
+ <div class="metric-pill"><div class="label">PDB</div><div class="value">{pdb_code or "Upload"}</div></div>
399
+ <div class="metric-pill"><div class="label">Chain</div><div class="value">{target_chain}</div></div>
400
+ </div>
401
+ </div>""",
402
+ unsafe_allow_html=True,
403
+ )
404
+ with col2:
405
+ st.markdown(
406
+ f"""<div class="glass-card">
407
+ <div class="card-title">πŸ”— Binder</div>
408
+ <div class="metric-row">
409
+ <div class="metric-pill"><div class="label">Length</div><div class="value">{binder_len_final}</div></div>
410
+ <div class="metric-pill"><div class="label">Mode</div><div class="value">{"Supervised" if binder_chain else "Hallucinate"}</div></div>
411
+ </div>
412
+ </div>""",
413
+ unsafe_allow_html=True,
414
+ )
415
+ with col3:
416
+ st.markdown(
417
+ f"""<div class="glass-card">
418
+ <div class="card-title">βš™οΈ Model</div>
419
+ <div class="metric-row">
420
+ <div class="metric-pill"><div class="label">Optimizer</div><div class="value" style="font-size:1rem;">{optimizer}</div></div>
421
+ <div class="metric-pill"><div class="label">Models</div><div class="value">{num_models_sel}</div></div>
422
+ </div>
423
+ </div>""",
424
+ unsafe_allow_html=True,
425
+ )
426
+
427
+ # ───────────────────────────────────────────────────────────────────────
428
+ # Run button & pipeline
429
+ # ───────────────────────────────────────────────────────────────────────
430
+ st.markdown("---")
431
+ run_clicked = st.button("β–Ά Run Binder Design", use_container_width=True, type="primary")
432
+
433
+ if run_clicked:
434
+ # ── 1. Resolve PDB file ──────────────────────────────────────────
435
+ with st.status("πŸ”¬ Running AfDesign Pipeline...", expanded=True) as status:
436
+ try:
437
+ st.write("πŸ“₯ Resolving PDB structure...")
438
+ if uploaded_pdb is not None:
439
+ tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdb")
440
+ tmp.write(uploaded_pdb.read())
441
+ tmp.close()
442
+ pdb_path = tmp.name
443
+ else:
444
+ pdb_path = fetch_pdb(pdb_code)
445
+
446
+ # ── 2. Download params if needed ─────────────────────────
447
+ if not download_params(st):
448
+ status.update(label="❌ Failed to download AlphaFold params", state="error")
449
+ st.stop()
450
+
451
+ # ── 3. Build model ───────────────────────────────────────
452
+ st.write("πŸ—οΈ Building AfDesign model...")
453
+ from colabdesign import mk_afdesign_model, clear_mem
454
+ from colabdesign.shared.utils import copy_dict
455
+ from colabdesign.af.alphafold.common import residue_constants
456
+ from scipy.special import softmax
457
+
458
+ clear_mem()
459
+ model = mk_afdesign_model(
460
+ protocol="binder",
461
+ use_multimer=use_multimer,
462
+ num_recycles=num_recycles,
463
+ recycle_mode="sample",
464
+ )
465
+
466
+ prep_kwargs = {
467
+ "pdb_filename": pdb_path,
468
+ "chain": target_chain,
469
+ "binder_len": binder_len_final,
470
+ "binder_chain": binder_chain,
471
+ "hotspot": hotspot,
472
+ "use_multimer": use_multimer,
473
+ "rm_target_seq": target_flexible,
474
+ }
475
+ model.prep_inputs(**prep_kwargs, ignore_missing=False)
476
+ st.write(f" ↳ Target length: **{model._target_len}** Β· Binder length: **{model._binder_len}**")
477
+
478
+ # ── 4. Optimize ──────────────────────────────────────────
479
+ st.write(f"πŸ§ͺ Running optimization ({optimizer})...")
480
+ model.restart(seq=binder_seq)
481
+ model.set_optimizer(
482
+ optimizer=gd_method,
483
+ learning_rate=learning_rate,
484
+ norm_seq_grad=norm_seq_grad,
485
+ )
486
+ models_list = model._model_names[:num_models_int]
487
+ flags = {"num_recycles": num_recycles, "models": models_list, "dropout": dropout}
488
+
489
+ pssm = None
490
+
491
+ if optimizer == "3stage":
492
+ model.design_3stage(120, 60, 10, **flags)
493
+ pssm = softmax(model._tmp["seq_logits"], -1)
494
+
495
+ elif optimizer == "pssm_semigreedy":
496
+ model.design_pssm_semigreedy(120, 32, **flags)
497
+ pssm = softmax(model._tmp["seq_logits"], 1)
498
+
499
+ elif optimizer == "semigreedy":
500
+ model.design_pssm_semigreedy(0, 32, **flags)
501
+ pssm = None
502
+
503
+ elif optimizer == "pssm":
504
+ model.design_logits(120, e_soft=1.0, num_models=1, ramp_recycles=True, **flags)
505
+ model.design_soft(32, num_models=1, **flags)
506
+ flags.update({"dropout": False, "save_best": True})
507
+ model.design_soft(10, num_models=num_models_int, **flags)
508
+ pssm = softmax(model.aux["seq"]["logits"], -1)
509
+
510
+ else:
511
+ opt_map = {
512
+ "logits": model.design_logits,
513
+ "soft": model.design_soft,
514
+ "hard": model.design_hard,
515
+ }
516
+ if optimizer in opt_map:
517
+ opt_map[optimizer](120, num_models=1, ramp_recycles=True, **flags)
518
+ flags.update({"dropout": False, "save_best": True})
519
+ opt_map[optimizer](10, num_models=num_models_int, **flags)
520
+ pssm = softmax(model.aux["seq"]["logits"], -1)
521
+
522
+ st.write("βœ… Optimization complete!")
523
+
524
+ # ── 5. Save PDB ──────────────────────────────────────────
525
+ out_pdb = f"{model.protocol}.pdb"
526
+ model.save_pdb(out_pdb)
527
+
528
+ status.update(label="βœ… Design Complete!", state="complete")
529
+
530
+ except Exception as e:
531
+ status.update(label="❌ Error", state="error")
532
+ st.error(f"**Error:** {e}")
533
+ st.stop()
534
+
535
+ # ── Results ──────────────────────────────────────────────────────
536
+ st.markdown("---")
537
+ st.markdown('<p class="hero-title" style="font-size:1.6rem;">πŸ“Š Results</p>', unsafe_allow_html=True)
538
+
539
+ # Metrics
540
+ log = model._tmp.get("best", {}).get("aux", {}).get("log", {})
541
+ metric_cols = st.columns(4)
542
+ metrics_to_show = [
543
+ ("pLDDT (binder)", log.get("plddt", None)),
544
+ ("pAE", log.get("pae", None)),
545
+ ("i_pAE", log.get("i_pae", None)),
546
+ ("i_con", log.get("i_con", None)),
547
+ ]
548
+ for col, (label, val) in zip(metric_cols, metrics_to_show):
549
+ with col:
550
+ display_val = f"{val:.3f}" if val is not None else "β€”"
551
+ st.markdown(
552
+ f"""<div class="metric-pill" style="text-align:center;">
553
+ <div class="label">{label}</div>
554
+ <div class="value">{display_val}</div>
555
+ </div>""",
556
+ unsafe_allow_html=True,
557
+ )
558
+
559
+ # Designed sequence
560
+ st.markdown('<div class="glass-card"><div class="card-title">🧬 Designed Sequence</div>', unsafe_allow_html=True)
561
+ seqs = model.get_seqs()
562
+ if seqs:
563
+ seq_str = seqs[0] if isinstance(seqs, list) else str(seqs)
564
+ st.markdown(f'<div class="seq-box">{seq_str}</div>', unsafe_allow_html=True)
565
+ st.markdown("</div>", unsafe_allow_html=True)
566
+
567
+ # Download PDB
568
+ res_col1, res_col2 = st.columns([1, 1])
569
+ with res_col1:
570
+ if os.path.isfile(out_pdb):
571
+ with open(out_pdb, "r") as f:
572
+ pdb_data = f.read()
573
+ st.download_button(
574
+ label="⬇️ Download Designed PDB",
575
+ data=pdb_data,
576
+ file_name=out_pdb,
577
+ mime="chemical/x-pdb",
578
+ use_container_width=True,
579
+ )
580
+
581
+ # PSSM heatmap
582
+ if pssm is not None:
583
+ st.markdown(
584
+ '<div class="glass-card"><div class="card-title">πŸ“ˆ Amino Acid Probability (PSSM)</div>',
585
+ unsafe_allow_html=True,
586
+ )
587
+ pssm_2d = pssm.mean(0) if pssm.ndim == 3 else pssm
588
+ fig = px.imshow(
589
+ pssm_2d.T,
590
+ labels=dict(x="Position", y="Amino Acid", color="Probability"),
591
+ y=list(residue_constants.restypes),
592
+ zmin=0,
593
+ zmax=1,
594
+ color_continuous_scale="Viridis",
595
+ template="plotly_dark",
596
+ aspect="auto",
597
+ )
598
+ fig.update_layout(
599
+ paper_bgcolor="rgba(0,0,0,0)",
600
+ plot_bgcolor="rgba(0,0,0,0)",
601
+ font=dict(family="Inter", color="#e6edf3"),
602
+ margin=dict(l=50, r=20, t=30, b=40),
603
+ height=400,
604
+ )
605
+ fig.update_xaxes(side="top")
606
+ st.plotly_chart(fig, use_container_width=True)
607
+ st.markdown("</div>", unsafe_allow_html=True)
608
+
609
+ # Log details
610
+ with st.expander("πŸ“‹ Full Design Log", expanded=False):
611
+ st.json(log)
612
+
613
+ else:
614
+ # ── Landing / idle state ─────────────────────────────────────────
615
+ st.markdown(
616
+ """
617
+ <div class="glass-card" style="text-align:center; padding:3rem 2rem;">
618
+ <div style="font-size:3rem; margin-bottom:0.8rem;">🧬</div>
619
+ <div style="font-size:1.2rem; font-weight:600; color:#e6edf3; margin-bottom:0.5rem;">
620
+ Ready to Design
621
+ </div>
622
+ <div style="color:#8b949e; max-width:500px; margin:0 auto;">
623
+ Configure your target protein and binder parameters in the sidebar,
624
+ then click <strong style="color:#58a6ff;">Run Binder Design</strong> to begin.
625
+ </div>
626
+ </div>
627
+ """,
628
+ unsafe_allow_html=True,
629
+ )
630
+
631
+ st.markdown(
632
+ """
633
+ <div class="glass-card">
634
+ <div class="card-title">πŸ“– How It Works</div>
635
+ <div style="color:#8b949e; line-height:1.75;">
636
+ <strong style="color:#e6edf3;">1. Provide a target</strong> β€” Enter a PDB code, UniProt ID, or upload a PDB file.<br/>
637
+ <strong style="color:#e6edf3;">2. Set binder parameters</strong> β€” Choose length, optional initial sequence, and chain info.<br/>
638
+ <strong style="color:#e6edf3;">3. Configure model</strong> β€” Select recycles, model count, and optimizer strategy.<br/>
639
+ <strong style="color:#e6edf3;">4. Run design</strong> β€” AfDesign hallucinate a binder sequence and optimizes for interface contacts & pLDDT.<br/>
640
+ <strong style="color:#e6edf3;">5. Analyze results</strong> β€” View metrics, designed sequence, PSSM heatmap, and download the PDB.
641
+ </div>
642
+ </div>
643
+ """,
644
+ unsafe_allow_html=True,
645
+ )
646
 
647
+ st.markdown(
648
+ """
649
+ <div style="text-align:center; color:#484f58; font-size:0.8rem; margin-top:2rem;">
650
+ Powered by <a href="https://github.com/sokrypton/ColabDesign" target="_blank" style="color:#58a6ff; text-decoration:none;">ColabDesign</a>
651
+ &nbsp;Β·&nbsp; AlphaFold Parameters Β© DeepMind
652
+ </div>
653
+ """,
654
+ unsafe_allow_html=True,
655
+ )