Update app.py

app.py

@@ -3,9 +3,8 @@ import numpy as np
 import gradio as gr
 import matplotlib.pyplot as plt
 from transformers import AutoTokenizer, EsmModel
-from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.decomposition import PCA
 from Bio.PDB import PDBParser, PDBIO
-import freesasa
 import tempfile
 import os
 
@@ -13,48 +12,25 @@ import os
 model = EsmModel.from_pretrained("facebook/esm1b_t33_650M_UR50S", output_hidden_states=True)
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
 
-# Compute per-residue cosine similarity scores (ASA-aware)
-def compute_asa_filtered_scores(seq, pdb_path):
+# Compute scaled PCA values for a selected component
+def compute_scaled_pca_scores(seq, component=0):
     inputs = tokenizer(seq, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        embedding = outputs.last_hidden_state[0]
+        embedding = outputs.last_hidden_state[0]  # shape (L+2, d)
 
     L = len(seq)
-    embedding = embedding[1:L+1]  # Remove CLS/EOS
+    embedding = embedding[1:L+1]  # remove CLS and EOS
 
-    # Compute ASA using freesasa
-    structure = freesasa.Structure(pdb_path)
-    result = freesasa.calc(structure)
-    rASA = []
-    for i in range(L):
-        try:
-            res_id = structure.residueNumber(i)
-            chain = structure.chainLabel(i)
-            area = result.residueAreas()[chain][res_id]['total']
-            # Estimate max ASA for normalization (simplified)
-            max_acc = 200.0  # Conservative estimate for normalization
-            rASA.append(area / max_acc)
-        except:
-            rASA.append(0.0)
-    rASA = np.array(rASA)
+    # Run PCA
+    pca = PCA(n_components=component + 1)
+    pca_result = pca.fit_transform(embedding.detach().cpu().numpy())
+    selected_component = pca_result[:, component]
 
-    # Bin into buried (<= 0.25) and exposed (> 0.25)
-    buried_idx = np.where(rASA <= 0.25)[0]
-    exposed_idx = np.where(rASA > 0.25)[0]
-
-    # Compute cosine similarity matrix
-    sim_matrix = cosine_similarity(embedding.detach().cpu().numpy())
-
-    # Sum similarities only within ASA bins
-    filtered_scores = np.zeros(L)
-    for i in range(L):
-        group = buried_idx if i in buried_idx else exposed_idx
-        filtered_scores[i] = np.sum(sim_matrix[i, group])
-
-    # Normalize
-    norm_scores = 100 * (filtered_scores - np.min(filtered_scores)) / (np.max(filtered_scores) - np.min(filtered_scores))
-    return norm_scores
+    # Scale between 0 and 100 for B-factor compatibility
+    scaled = (selected_component - selected_component.min()) / (selected_component.max() - selected_component.min())
+    scaled *= 100
+    return scaled
 
 # Inject scores into B-factor column
 def inject_bfactors_into_pdb(pdb_file, scores):
@@ -75,21 +51,23 @@ def inject_bfactors_into_pdb(pdb_file, scores):
     io.save(out_path)
     return out_path
 
-# Combined Gradio interface
-def process(seq, pdb_file):
-    scores = compute_asa_filtered_scores(seq, pdb_file.name)
+# Gradio interface logic
+def process(seq, pdb_file, component):
+    scores = compute_scaled_pca_scores(seq, component)
     pdb_with_scores = inject_bfactors_into_pdb(pdb_file, scores)
     return pdb_with_scores
 
-# Gradio Interface
+# Gradio UI
 demo = gr.Interface(
     fn=process,
     inputs=[
         gr.Textbox(label="Input Protein Sequence (1-letter code)"),
-        gr.File(label="Upload PDB File", file_types=[".pdb"])
+        gr.File(label="Upload PDB File", file_types=[".pdb"]),
+        gr.Number(label="PCA Component (0 = first PC)", value=0, precision=0)
     ],
-    outputs=gr.File(label="Modified PDB with ASA-filtered Embedding Scores in B-factor Column"),
-    title="ESM-1b ASA-Aware Residue Scoring for Structural Visualization"
+    outputs=gr.File(label="Modified PDB with PCA Component in B-factor Column"),
+    title="ESM-1b PCA Component Projection for Structural Mapping"
 )
 
 demo.launch()
+