Update app.py

app.py

@@ -7,8 +7,9 @@ from ase.io import read, write
 from ase.io.trajectory import Trajectory
 from gradio_molecule3d import Molecule3D
 import hashlib
+import shutil
 
-# ==== Molecule3D viewer preparation con debug ====
+# ==== Molecule3D viewer preparation CORREGIDO ====
 def prepare_molecule_for_viewer(traj_path):
     """Convert trajectory to format compatible with Molecule3D viewer"""
     if not traj_path or not os.path.exists(traj_path):
@@ -24,25 +25,99 @@ def prepare_molecule_for_viewer(traj_path):
         # Debug info
         print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms")
         
-        # Crear PDB temporal persistente (no se elimina automáticamente)
-        temp_pdb = tempfile.NamedTemporaryFile(suffix='.pdb', delete=False, mode='w')
+        # Crear directorio cache de Gradio si no existe
+        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+        os.makedirs(gradio_cache_dir, exist_ok=True)
         
-        # Escribir contenido PDB
+        # Crear PDB temporal PERSISTENTE en el directorio cache de Gradio
+        unique_id = hashlib.md5(str(traj_path).encode() + str(os.path.getmtime(traj_path)).encode()).hexdigest()[:12]
+        pdb_filename = f"molecule_{unique_id}.pdb"
+        pdb_path = os.path.join(gradio_cache_dir, pdb_filename)
+        
+        # Si ya existe, no lo recreamos
+        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+            print(f"PDB already exists: {pdb_path}")
+            return pdb_path
+        
+        # Tomar el último frame de la trayectoria
         atoms = traj[-1]
-        write(temp_pdb.name, atoms, format='pdb')
-        temp_pdb.close()  # Cerrar el archivo para que se pueda leer
         
-        # Verificar que el archivo se creó
-        if os.path.exists(temp_pdb.name):
-            file_size = os.path.getsize(temp_pdb.name)
-            print(f"PDB created: {temp_pdb.name}, size: {file_size} bytes")
-            return temp_pdb.name
+        # Escribir PDB con elementos en minúsculas (como requiere gradio_molecule3d)
+        try:
+            # Usar ASE para escribir PDB
+            write(pdb_path, atoms, format='pdb')
+            
+            # Post-procesar para asegurar formato correcto
+            with open(pdb_path, 'r') as f:
+                content = f.read()
+            
+            # Asegurar que los símbolos de elementos estén en formato correcto
+            lines = content.split('\n')
+            processed_lines = []
+            
+            for line in lines:
+                if line.startswith(('ATOM', 'HETATM')):
+                    # Formato PDB: columnas específicas para elemento
+                    processed_lines.append(line)
+                else:
+                    processed_lines.append(line)
+            
+            with open(pdb_path, 'w') as f:
+                f.write('\n'.join(processed_lines))
+                
+        except Exception as write_error:
+            print(f"Error writing PDB: {write_error}")
+            return None
+        
+        # Verificar que el archivo se creó correctamente
+        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+            file_size = os.path.getsize(pdb_path)
+            print(f"PDB created successfully: {pdb_path}, size: {file_size} bytes")
+            
+            # Verificar contenido básico del PDB
+            with open(pdb_path, 'r') as f:
+                first_lines = f.read(200)
+                print(f"PDB content preview: {first_lines[:100]}...")
+                
+            return pdb_path
         else:
-            print("PDB file was not created")
+            print("PDB file was not created properly")
             return None
             
     except Exception as e:
         print(f"Error preparing molecule for viewer: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+# ==== Función para convertir archivo inicial a PDB para SPE ====
+def prepare_input_for_viewer(structure_file):
+    """Convert input structure file to PDB for Molecule3D viewer"""
+    if not structure_file or not os.path.exists(structure_file):
+        return None
+    
+    try:
+        atoms = read(structure_file)
+        
+        # Crear directorio cache de Gradio si no existe
+        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+        os.makedirs(gradio_cache_dir, exist_ok=True)
+        
+        # Crear PDB para visualización
+        unique_id = hashlib.md5(str(structure_file).encode()).hexdigest()[:12]
+        pdb_filename = f"input_{unique_id}.pdb"
+        pdb_path = os.path.join(gradio_cache_dir, pdb_filename)
+        
+        write(pdb_path, atoms, format='pdb')
+        
+        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
+            print(f"Input PDB created: {pdb_path}")
+            return pdb_path
+        else:
+            return None
+            
+    except Exception as e:
+        print(f"Error preparing input for viewer: {e}")
         return None
 
 # ==== OrbMol SPE ====
@@ -66,15 +141,15 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
     try:
         calc = _load_orbmol_calc()
         if not structure_file:
-            return "Error: Please upload a structure file", "Error"
+            return "Error: Please upload a structure file", "Error", None
 
         file_path = structure_file
 
         if not os.path.exists(file_path):
-            return f"Error: File not found: {file_path}", "Error"
+            return f"Error: File not found: {file_path}", "Error", None
         
         if os.path.getsize(file_path) == 0:
-            return f"Error: Empty file: {file_path}", "Error"
+            return f"Error: Empty file: {file_path}", "Error", None
 
         atoms = read(file_path)
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
@@ -89,9 +164,12 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
         lines += ["", f"Max Force: {max_force:.4f} eV/Å"]
 
-        return "\n".join(lines), "Calculation completed with OrbMol"
+        # Preparar PDB para visualización
+        pdb_file = prepare_input_for_viewer(file_path)
+
+        return "\n".join(lines), "Calculation completed with OrbMol", pdb_file
     except Exception as e:
-        return f"Error during calculation: {e}", "Error"
+        return f"Error during calculation: {e}", "Error", None
 
 # ==== Simulaciones (helpers) ====
 from simulation_scripts_orbmol import (
@@ -99,7 +177,7 @@ from simulation_scripts_orbmol import (
     run_relaxation_simulation,
 )
 
-# ==== Wrappers con debug ====
+# ==== Wrappers con debug y Molecule3D ====
 def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble):
     try:
         if not structure_file:
@@ -111,7 +189,7 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
         traj_path, log_text, script_text, explanation = run_md_simulation(
             file_path,
             int(steps),
-            20,
+            20,  # pre-relax steps
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -122,14 +200,16 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
 
         print(f"MD completed, trajectory: {traj_path}")
         
-        # Preparar archivo PDB para el visor 3D
+        # Preparar archivo PDB para el visor Molecule3D
         pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for viewer: {pdb_file}")
+        print(f"PDB file for Molecule3D: {pdb_file}")
         
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
         print(f"MD Wrapper Error: {e}")
+        import traceback
+        traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
 def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
@@ -152,14 +232,16 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
 
         print(f"Relaxation completed, trajectory: {traj_path}")
         
-        # Preparar archivo PDB para el visor 3D
+        # Preparar archivo PDB para el visor Molecule3D
         pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for viewer: {pdb_file}")
+        print(f"PDB file for Molecule3D: {pdb_file}")
         
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
         print(f"Relax Wrapper Error: {e}")
+        import traceback
+        traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
 # ==== UI ====
@@ -186,8 +268,27 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 with gr.Column(variant="panel", min_width=500):
                     spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                     spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
+                    
+                    # Viewer para SPE
+                    spe_viewer = Molecule3D(
+                        label="Input Structure Viewer",
+                        reps=[
+                            {
+                                "model": 0,
+                                "chain": "",
+                                "resname": "",
+                                "style": "stick",
+                                "color": "whiteCarbon",
+                                "residue_range": "",
+                                "around": 0,
+                                "byres": False,
+                                "visible": True,
+                                "opacity": 1.0
+                            }
+                        ]
+                    )
 
-            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status])
+            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status, spe_viewer])
 
         # -------- MD --------
         with gr.Tab("Molecular Dynamics"):
@@ -217,9 +318,9 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     md_status = gr.Textbox(label="MD Status", interactive=False)
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     
-                    # Molecule3D viewer con configuraciones adicionales
+                    # Molecule3D viewer para MD
                     md_viewer = Molecule3D(
-                        label="3D Structure Viewer",
+                        label="Final Structure Viewer (Last MD Frame)",
                         reps=[
                             {
                                 "model": 0,
@@ -283,9 +384,9 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     rlx_status = gr.Textbox(label="Status", interactive=False)
                     rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     
-                    # Molecule3D viewer con configuraciones adicionales
+                    # Molecule3D viewer para Relaxation
                     rlx_viewer = Molecule3D(
-                        label="3D Structure Viewer",
+                        label="Optimized Structure Viewer",
                         reps=[
                             {
                                 "model": 0,