Upload 5 files

scripts/visualize.py

@@ -0,0 +1,84 @@
+import matplotlib.pyplot as plt # type: ignore
+import pandas as pd
+import re
+
+def create_design_plot(fasta_file_path):
+    """
+    Parses the MPNN fasta file and returns a Matplotlib figure.
+    """
+    data = []
+    with open(fasta_file_path, 'r') as f:
+        content = f.read()
+    
+    # Extract score (local score, not global_score) and recovery from headers
+    # Example: >T=0.1, sample=1, score=0.7880, global_score=1.3897, seq_recovery=0.4295
+    # We want to capture the local "score=" value, not "global_score="
+    pattern = r"score=([\d\.]+)(?:,\s+global_score=[\d\.]+)?,\s+seq_recovery=([\d\.]+)"
+    
+    for line in content.split('\n'):
+        if line.startswith(">") and "score=" in line and "seq_recovery" in line:
+            # Make sure we're not matching global_score by checking the pattern order
+            match = re.search(pattern, line)
+            if match:
+                # Verify we got the local score (should be before global_score if present)
+                score_val = float(match.group(1))
+                recovery_val = float(match.group(2))
+                data.append({
+                    "Score": score_val, 
+                    "Recovery": recovery_val
+                })
+    
+    df = pd.DataFrame(data)
+    
+    # Create the Plot
+    fig, ax = plt.subplots(figsize=(8, 5))
+    ax.scatter(df['Score'], df['Recovery'], color='black', alpha=0.6, s=80, edgecolors='white')
+    
+    # Highlight the Lead Candidate (Lowest Score) with a gold star
+    best = df.loc[df['Score'].idxmin()]
+    ax.scatter(best['Score'], best['Recovery'], marker='*', color='gold', s=400, edgecolors='black', linewidths=1.5, label="Lead Candidate", zorder=10)
+    
+    ax.set_title("BroteinShake: Design Evolution (N=20)", fontsize=14)
+    ax.set_xlabel("ProteinMPNN Score (Lower = More Stable)", fontsize=10)
+    ax.set_ylabel("Sequence Recovery (%)", fontsize=10)
+    ax.legend()
+    ax.grid(True, linestyle='--', alpha=0.5)
+    
+    return fig
+
+def create_protein_viewer(pdb_file_path):
+    """
+    Generates HTML/JS for a 3D protein viewer using 3Dmol.js.
+    """
+    with open(pdb_file_path, 'r') as f:
+        pdb_content = f.read().replace('\n', '\\n')
+    
+    html_content = f"""
+    <div id="container-3d" style="height: 500px; width: 100%; position: relative;"></div>
+    <script>
+        (function() {{
+            // Wait for 3Dmol to be loaded (from head)
+            function initViewer() {{
+                if (typeof $3Dmol === 'undefined') {{
+                    setTimeout(initViewer, 100);
+                    return;
+                }}
+                let element = document.getElementById('container-3d');
+                if (!element) return;
+                let config = {{ backgroundColor: 'white' }};
+                let viewer = $3Dmol.createViewer(element, config);
+                let pdbData = `{pdb_content}`;
+                viewer.addModel(pdbData, "pdb");
+                viewer.setStyle({{}}, {{cartoon: {{color: 'spectrum'}}}});
+                viewer.zoomTo();
+                viewer.render();
+            }}
+            if (document.readyState === 'loading') {{
+                document.addEventListener('DOMContentLoaded', initViewer);
+            }} else {{
+                initViewer();
+            }}
+        }})();
+    </script>
+    """
+    return html_content