fix pdb utils

app.py

@@ -126,6 +126,39 @@ def add_noise_to_coordinates(protein: ESMProtein, noise_level: float) -> ESMProt
     noisy_coordinates = coordinates + noise
     return ESMProtein(sequence=protein.sequence, coordinates=noisy_coordinates)
 
+def run_structure_prediction(protein: ESMProtein) -> ESMProtein:
+    structure_prediction_config = GenerationConfig(
+        track="structure",
+        num_steps=40,
+        temperature=0.7,
+    )
+    try:
+        structure_prediction = model.generate(protein, structure_prediction_config)
+        if isinstance(structure_prediction, ESMProtein):
+            return structure_prediction
+        else:
+            raise ValueError(f"Unexpected result type: {type(structure_prediction)}")
+    except Exception as e:
+        print(f"Error during structure prediction: {str(e)}")
+        return None
+
+def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein) -> tuple[ESMProtein, float]:
+    if structure_prediction is None:
+        return None, float('inf')
+    
+    try:
+        structure_prediction_chain = structure_prediction.to_protein_chain()
+        protein_chain = protein.to_protein_chain()
+        structure_indices = np.arange(0, len(structure_prediction_chain.sequence))
+        aligned_chain = structure_prediction_chain.align(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)
+        crmsd = structure_prediction_chain.rmsd(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)    
+        return ESMProtein.from_protein_chain(aligned_chain), crmsd
+    except AttributeError as e:
+        print(f"Error during alignment: {str(e)}")
+        print(f"Structure prediction type: {type(structure_prediction)}")
+        print(f"Structure prediction attributes: {dir(structure_prediction)}")
+        return None, float('inf')
+
 def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_frames: int):
     protein = get_protein(pdb_file)
     runs = []
@@ -136,30 +169,20 @@ def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_fr
         for i in range(num_runs):
             structure_prediction = run_structure_prediction(noisy_protein)
             aligned, crmsd = align_after_prediction(protein, structure_prediction)
-            runs.append((crmsd, aligned))
+            if aligned is not None:
+                runs.append((crmsd, aligned))
+    
+    if not runs:
+        return None, "No successful predictions"
     
     best_aligned = sorted(runs)[0]
     view = visualize_after_pred(protein, best_aligned[1])
     return view, f"Best cRMSD: {best_aligned[0]:.4f}"
 
-def run_structure_prediction(protein: ESMProtein) -> ESMProtein:
-    structure_prediction_config = GenerationConfig(
-        track="structure",
-        num_steps=40,
-        temperature=0.7,
-    )
-    structure_prediction = model.generate(protein, structure_prediction_config)
-    return structure_prediction
-
-def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein) -> tuple[ESMProtein, float]:
-    structure_prediction_chain = structure_prediction.to_protein_chain()
-    protein_chain = protein.to_protein_chain()
-    structure_indices = np.arange(0, len(structure_prediction_chain.sequence))
-    aligned_chain = structure_prediction_chain.align(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)
-    crmsd = structure_prediction_chain.rmsd(protein_chain, mobile_inds=structure_indices, target_inds=structure_indices)    
-    return ESMProtein.from_protein_chain(aligned_chain), crmsd
-
 def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
+    if aligned is None:
+        return py3Dmol.view(width=800, height=600)
+    
     view = py3Dmol.view(width=800, height=600)
     view.addModel(protein.to_pdb_string(), "pdb")
     view.setStyle({"cartoon": {"color": "lightgrey"}})
@@ -175,6 +198,9 @@ def run_prediction(pdb_file, num_runs, noise_level, num_frames):
             return "Please upload a PDB file.", "No file uploaded"
         
         view, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames)
+        if view is None:
+            return "No successful predictions were made. Try adjusting the parameters.", crmsd_text
+        
         html = view._make_html()
         return f"""
         <div style="height: 600px;">