Upload 4 files

app.py

@@ -1,15 +1,81 @@
 import gradio as gr # type: ignore
 import os
 from gradio_molecule3d import Molecule3D
+from Bio.PDB import PDBParser
 
 # Import your custom modules from the /scripts folder
 from scripts.download import download_and_clean_pdb
 from scripts.generator import run_broteinshake_generator
-from scripts.refine import polish_design
+from scripts.refine import polish_design, process_results
 from scripts.visualize import create_design_plot
 
 # --- HELPER FUNCTIONS ---
 
+def get_pdb_chains(pdb_file):
+    """Extracts unique chain IDs from a PDB file."""
+    if not pdb_file or not os.path.exists(pdb_file):
+        return []
+    try:
+        parser = PDBParser(QUIET=True)
+        structure = parser.get_structure("temp", pdb_file)
+        chains = [chain.id for model in structure for chain in model]
+        return sorted(list(set(chains)))
+    except Exception as e:
+        print(f"Error extracting chains: {e}")
+        return []
+
+def load_pdb_and_extract_chains(pdb_id):
+    """Download PDB and extract chains for selection."""
+    if not pdb_id or not pdb_id.strip():
+        return gr.update(choices=[], value=[]), "⚠️ Please enter a PDB ID", gr.update(interactive=False), []
+    
+    try:
+        # Download the PDB
+        pdb_path = download_and_clean_pdb(pdb_id.strip(), data_dir="data")
+        
+        # Extract chains
+        chains = get_pdb_chains(pdb_path)
+        
+        if not chains:
+            return gr.update(choices=[], value=[]), f"⚠️ No chains found in {pdb_id.upper()}", gr.update(interactive=False), []
+        
+        # Single-chain proteins are supported (will use different ProteinMPNN command)
+        # For single-chain, the only chain will be automatically selected for redesign
+        if len(chains) == 1:
+            status_msg = f"✅ Loaded {pdb_id.upper()}: Single-chain protein - will redesign chain {chains[0]}"
+            # Auto-select the chain for single-chain proteins
+            return gr.update(choices=chains, value=chains), status_msg, gr.update(interactive=True), chains
+        
+        status_msg = f"✅ Loaded {pdb_id.upper()}: Found {len(chains)} chain(s) - {', '.join(chains)}"
+        # Initially disable button - user must select at least one chain
+        return gr.update(choices=chains, value=chains), status_msg, gr.update(interactive=False), chains
+    except Exception as e:
+        error_msg = f"❌ Error loading {pdb_id.upper()}: {str(e)}"
+        print(error_msg)
+        return gr.update(choices=[], value=[]), error_msg, gr.update(interactive=False), []
+
+def validate_chain_selection(selected_chains, available_chains_state):
+    """Validate that at least one chain is selected and at least one remains fixed (for multi-chain)."""
+    if not selected_chains or len(selected_chains) == 0:
+        warning = "⚠️ Please select at least one chain to redesign"
+        return gr.update(interactive=False), warning, available_chains_state
+    
+    # Get available chains from state
+    available_chains = available_chains_state if available_chains_state else []
+    
+    # For single-chain proteins, allow selecting the only chain
+    if len(available_chains) == 1:
+        warning = f"✅ Single-chain protein: Will redesign chain {available_chains[0]}"
+        return gr.update(interactive=True), warning, available_chains_state
+    
+    # For multi-chain: Check if all chains are selected (would leave no fixed chains)
+    if available_chains and len(selected_chains) >= len(available_chains):
+        warning = f"⚠️ Cannot select all chains - at least one chain must remain fixed. Selected: {', '.join(selected_chains)}"
+        return gr.update(interactive=False), warning, available_chains_state
+    
+    warning = f"✅ {len(selected_chains)} chain(s) selected for redesign: {', '.join(selected_chains)}"
+    return gr.update(interactive=True), warning, available_chains_state
+
 def get_all_sequences(fasta_file: str) -> str:
     """Get all designed sequences from FASTA file."""
     sequences = []
@@ -69,15 +135,52 @@ def extract_best_sequence(fasta_file: str) -> str:
     else:
         raise ValueError(f"No valid designs found in {fasta_file}")
 
-def run_part1(pdb_id, fixed_chains, variable_chains):
+def run_part1(pdb_id, fixed_chains, variable_chains, temperature=0.1, selected_chains=None):
     """Downloads the PDB and runs ProteinMPNN design."""
     try:
         # Step 1: Secure the template
         pdb_path = download_and_clean_pdb(pdb_id, data_dir="data")
         
+        # Handle chain selection logic
+        # If chains are selected via checkbox, use those as variable chains
+        # Otherwise, use the text input (backward compatibility)
+        all_chains = get_pdb_chains(pdb_path)
+        
+        # Check if single-chain protein
+        is_single_chain = len(all_chains) == 1
+        
+        if selected_chains and len(selected_chains) > 0:
+            # Selected chains = variable chains, rest = fixed
+            variable_chains = "".join(selected_chains)
+            fixed_chains = "".join([c for c in all_chains if c not in selected_chains])
+            
+            # For single-chain: no fixed chains (will use different ProteinMPNN command)
+            # For multi-chain: Validate must have at least one fixed chain
+            if not is_single_chain and (not fixed_chains or len(fixed_chains) == 0):
+                raise ValueError(f"Cannot redesign all chains - at least one chain must remain fixed. Selected: {', '.join(selected_chains)}, Available: {', '.join(all_chains)}")
+            
+            if is_single_chain:
+                print(f"📋 Single-chain mode: Redesigning chain {variable_chains}")
+            else:
+                print(f"📋 Using chain selector: Fixed={fixed_chains}, Variable={variable_chains}")
+        else:
+            # If no chains selected, use text inputs (default behavior)
+            # For single-chain, if variable_chains is empty, use the only chain
+            if is_single_chain and not variable_chains:
+                variable_chains = all_chains[0]
+                fixed_chains = ""
+            # For multi-chain: Validate text inputs don't select all chains
+            elif not is_single_chain and fixed_chains and variable_chains:
+                all_selected = set(fixed_chains + variable_chains)
+                if len(all_selected) >= len(all_chains):
+                    raise ValueError(f"Cannot redesign all chains - at least one chain must remain fixed.")
+            print(f"📋 Using text inputs: Fixed={fixed_chains}, Variable={variable_chains}")
+        
         # Step 2: Generate Optimized Sequences
         # This creates the .fa files you need for the ESM Atlas
-        run_broteinshake_generator(pdb_path, fixed_chains, variable_chains)
+        print(f"🌡️  Temperature: {temperature}")
+        print(f"🔧 Parameters: Fixed chains={fixed_chains}, Variable chains={variable_chains}, Temp={temperature}")
+        run_broteinshake_generator(pdb_path, fixed_chains, variable_chains, num_seqs=20, temp=temperature)
         
         # Get all sequences and the best one
         fa_file = os.path.join("generated", pdb_id.lower(), "seqs", f"{pdb_id.lower()}_clones.fa")
@@ -123,23 +226,20 @@ def run_part2(pdb_id, uploaded_esm_file):
         
         # Call the new lightweight Biopython-only script
         print(f"🔍 Starting alignment for PDB ID: {pdb_id}, File: {file_path}")
-        final_pdb_path, rmsd_val = polish_design(pdb_id, file_path)
+        final_pdb_path, global_rmsd, core_rmsd, high_conf_rmsd = polish_design(pdb_id, file_path)
         
         # Validate output
         if not final_pdb_path or not os.path.exists(final_pdb_path):
             return None, f"❌ Error: Alignment failed - output file not created: {final_pdb_path}"
         
-        if rmsd_val is None:
+        if high_conf_rmsd is None:
             return None, "❌ Error: Alignment failed - RMSD calculation returned None"
         
-        print(f"✅ Alignment successful: RMSD = {rmsd_val:.3f} Å, Output: {final_pdb_path}")
+        print(f"✅ Alignment successful: Global RMSD = {global_rmsd:.3f} Å, Core RMSD = {core_rmsd:.3f} Å, High-Conf RMSD = {high_conf_rmsd:.3f} Å")
+        
+        # Generate detailed validation report
+        report = process_results(pdb_id, final_pdb_path, global_rmsd, high_conf_rmsd)
         
-        # Update the report to focus on Backbone Alignment
-        report = (
-            f"✅ Validation Successful!\n"
-            f"🎯 RMSD: {rmsd_val:.3f} Å\n"
-            f"🧬 Status: High-Precision Backbone Match"
-        )
         return final_pdb_path, report
     except FileNotFoundError as e:
         error_msg = f"❌ File Error: {str(e)}"
@@ -205,12 +305,54 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
         # TAB 1: GENERATIVE DESIGN
         with gr.Tab("1. Sequence Generation"):
             gr.Markdown("Enter a PDB ID to 'repaint' its binder interface using ProteinMPNN.")
-            with gr.Row():
-                pdb_input = gr.Textbox(label="Target PDB ID", placeholder="e.g., 3kas", value="3kas")
-                f_chains = gr.Textbox(label="Fixed Chains (Lock)", value="A")
-                v_chains = gr.Textbox(label="Variable Chains (Key)", value="B")
             
-            gen_btn = gr.Button("🚀 Generate Optimized Sequences", variant="primary")
+            pdb_input = gr.Textbox(label="Target PDB ID", placeholder="e.g., 3kas", value="")
+            load_pdb_btn = gr.Button("📥 Load PDB", variant="secondary")
+            pdb_status = gr.Markdown("💡 Enter a PDB ID and click 'Load PDB' to begin")
+            
+            with gr.Column():
+                gr.Markdown("### ⚙️ Design Parameters")
+                
+                # Temperature (T) is the most critical knob for sequence recovery
+                sampling_temp = gr.Slider(
+                    minimum=0.05, maximum=1.0, value=0.1, step=0.05,
+                    label="Sampling Temperature (T)",
+                    info="T=0.1 for high-fidelity; T=0.3 for natural diversity"
+                )
+
+                # Dynamic Chain Handling
+                chain_options = gr.CheckboxGroup(
+                    choices=[], 
+                    label="Chains to Redesign",
+                    info="Identify which chains ProteinMPNN should modify (will populate after loading PDB)"
+                )
+                
+                chain_warning = gr.Markdown("💡 Select at least one chain to enable generation", visible=True)
+                
+                # Hidden state to track if we've successfully parsed the PDB
+                pdb_state = gr.State()
+                
+                # Legacy text inputs (hidden but kept for backward compatibility)
+                with gr.Row(visible=False):
+                    f_chains = gr.Textbox(label="Fixed Chains (Lock)", value="A")
+                    v_chains = gr.Textbox(label="Variable Chains (Key)", value="B")
+            
+            # Generate button (initially disabled)
+            gen_btn = gr.Button("🚀 Generate Optimized Sequences", variant="primary", interactive=False)
+            
+            # Load PDB and extract chains when button is clicked
+            load_pdb_btn.click(
+                fn=load_pdb_and_extract_chains,
+                inputs=[pdb_input],
+                outputs=[chain_options, pdb_status, gen_btn, pdb_state]
+            )
+            
+            # Validate chain selection and update button state
+            chain_options.change(
+                fn=validate_chain_selection,
+                inputs=[chain_options, pdb_state],
+                outputs=[gen_btn, chain_warning, pdb_state]
+            )
             
             # Stack components vertically
             fa_output = gr.Code(
@@ -226,7 +368,7 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
             gr.Markdown("### System Status")
             status1 = gr.Markdown()
             
-            gen_btn.click(run_part1, inputs=[pdb_input, f_chains, v_chains], 
+            gen_btn.click(run_part1, inputs=[pdb_input, f_chains, v_chains, sampling_temp, chain_options], 
                          outputs=[fa_output, plot_output, status1])
 
         # TAB 2: STRUCTURAL VALIDATION
@@ -284,6 +426,5 @@ with gr.Blocks(theme=dark_biohub, css=biohub_css) as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    # FORCE DARK MODE: Add the theme toggle directly to the URL
-    demo.launch(server_port=7860)
-    # Visit: http://127.0.0.1:7860/?__theme=dark
+    # Docker deployment for HuggingFace Spaces
+    demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -33,7 +33,7 @@ requests
 safetensors==0.7.0
 sympy==1.13.1
 tokenizers==0.22.1
-torch==2.5.1
+torch==2.8.0
 torchvision==0.20.1
 tqdm==4.67.1
 transformers==4.57.3