Update app.py

app.py

@@ -6,122 +6,156 @@ import gradio as gr
 from ase.io import read, write
 from ase.io.trajectory import Trajectory
 from gradio_molecule3d import Molecule3D
-from simulation_scripts_orbmol import load_orbmol_model  # CORREGIDO
+from simulation_scripts_orbmol import load_orbmol_model
 import hashlib
 import shutil
 
-# ==== util: PDB writer robusto para Molecule3D ====
-def _pdb_cache_path(prefix: str, key: str) -> str:
-    gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
-    os.makedirs(gradio_cache_dir, exist_ok=True)
-    return os.path.join(gradio_cache_dir, f"{prefix}_{key}.pdb")
-
-def _write_pdb_with_fallback(atoms, pdb_path: str) -> str | None:
+# ==== Configuración UNIVERSAL para Molecule3D ====
+UNIVERSAL_MOLECULE3D_REPS = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "sphere",
+        "color": "element",  # Colores automáticos por elemento químico
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True,
+        "opacity": 0.8
+    },
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "stick",
+        "color": "element",
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True,
+        "opacity": 1.0
+    }
+]
+
+# ==== Función mejorada para crear PDB universal ====
+def create_universal_pdb(atoms, cache_key):
     """
-    Intenta escribir PDB con ASE usando el id correcto ('proteindatabank').
-    Si falla, crea un PDB minimal válido (líneas ATOM + END) suficiente para gradio_molecule3d.
-    Devuelve la ruta si existe y tiene tamaño > 0, en caso contrario None.
+    Crea un PDB optimizado que funciona bien con gradio_molecule3d para cualquier molécula
     """
-    try:
-        # writer id correcto en ASE para PDB
-        write(pdb_path, atoms, format="proteindatabank")
-    except Exception as e:
-        print(f"ASE PDB writer failed: {e} — trying manual minimal PDB")
-        try:
-            xyz = atoms.get_positions()
-            syms = atoms.get_chemical_symbols()
-            with open(pdb_path, "w") as f:
-                serial = 1
-                for (sym, (x, y, z)) in zip(syms, xyz):
-                    # Campos esenciales: 'ATOM', serial, nombre (símbolo), resname, chain, res seq, coords y elemento
-                    f.write(
-                        f"ATOM  {serial:5d} {sym:<3s} MOL A   1    "
-                        f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00          {sym:>2s}\n"
-                    )
-                    serial += 1
-                f.write("END\n")
-        except Exception as e2:
-            print(f"Manual PDB fallback failed: {e2}")
-            return None
-
-    # Validación
+    cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
+    os.makedirs(cache_dir, exist_ok=True)
+    pdb_path = os.path.join(cache_dir, f"mol_{cache_key}.pdb")
+    
     if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-        print(f"✅ PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)")
-        try:
-            with open(pdb_path, "r") as f:
-                preview = f.read(160)
-            print(f"PDB preview: {preview[:120]}...")
-        except Exception:
-            pass
         return pdb_path
-    print("❌ PDB file not created or empty")
-    return None
-
-# ==== 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):
-        print("No trajectory path provided or file doesn't exist")
-        return None
+    
     try:
-        traj = Trajectory(traj_path)
-        if len(traj) == 0:
-            print("Empty trajectory")
-            return None
-
-        print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms")
-        atoms = traj[-1]
-
-        # clave única por posiciones del último frame
-        pos_hash = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12]
-        pdb_path = _pdb_cache_path("molecule", pos_hash)
-
+        positions = atoms.get_positions()
+        symbols = atoms.get_chemical_symbols()
+        
+        with open(pdb_path, "w") as f:
+            # Header estándar
+            f.write("HEADER    MOLECULAR STRUCTURE                     01-JAN-25   MOL \n")
+            f.write("TITLE     CALCULATION RESULT\n")
+            f.write("MODEL        1\n")
+            
+            # Escribir átomos con formato PDB correcto
+            for i, (symbol, pos) in enumerate(zip(symbols, positions)):
+                # Asegurar formato correcto del elemento
+                element = symbol.strip().upper()
+                
+                # Nombre del átomo (4 caracteres, alineado a la izquierda)
+                atom_name = f"{element:<4s}"
+                
+                # Línea ATOM con formato PDB estricto
+                f.write(
+                    f"ATOM  {i+1:5d} {atom_name} UNL A   1    "
+                    f"{pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}"
+                    f"  1.00 30.00           {element:>2s}  \n"
+                )
+            
+            f.write("ENDMDL\n")
+            f.write("END\n")
+        
+        # Verificar creación
         if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-            print(f"PDB already exists: {pdb_path}")
+            print(f"✅ Universal PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)")
+            
+            # Debug: mostrar contenido
+            with open(pdb_path, 'r') as f:
+                content = f.read()
+                print(f"PDB content preview:\n{content[:300]}...")
+            
             return pdb_path
+        else:
+            print(f"❌ Universal PDB creation failed")
+            return None
+            
+    except Exception as e:
+        print(f"❌ Error in create_universal_pdb: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
 
-        return _write_pdb_with_fallback(atoms, pdb_path)
-
+# ==== Función actualizada para preparar cualquier molécula ====
+def prepare_universal_molecule_viewer(traj_path_or_atoms):
+    """
+    Función universal que funciona con cualquier molécula
+    """
+    try:
+        # Determinar si es trayectoria o átomos
+        if isinstance(traj_path_or_atoms, str):
+            if not os.path.exists(traj_path_or_atoms):
+                print("Trajectory file doesn't exist")
+                return None
+            
+            traj = Trajectory(traj_path_or_atoms)
+            if len(traj) == 0:
+                print("Empty trajectory")
+                return None
+            
+            atoms = traj[-1]
+            cache_key = hashlib.md5(str(traj_path_or_atoms).encode()).hexdigest()[:12]
+        else:
+            atoms = traj_path_or_atoms
+            cache_key = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12]
+        
+        print(f"Preparing universal viewer: {len(atoms)} atoms")
+        print(f"Chemical symbols: {atoms.get_chemical_symbols()}")
+        
+        # Crear PDB universal
+        pdb_path = create_universal_pdb(atoms, cache_key)
+        
+        return pdb_path
+        
     except Exception as e:
-        print(f"Error preparing molecule for viewer: {e}")
-        import traceback; traceback.print_exc()
+        print(f"Error in prepare_universal_molecule_viewer: {e}")
+        import traceback
+        traceback.print_exc()
         return None
 
-# ==== Función para convertir archivo inicial a PDB para SPE ====
+# ==== Funciones de preparación universales ====
 def prepare_input_for_viewer(structure_file):
-    """Convert input structure file to PDB for Molecule3D viewer"""
+    """Universal input preparation"""
     if not structure_file or not os.path.exists(structure_file):
         return None
     try:
         atoms = read(structure_file)
-        key = hashlib.md5(str(structure_file).encode()).hexdigest()[:12]
-        pdb_path = _pdb_cache_path("input", key)
-
-        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
-            print(f"Input PDB already exists: {pdb_path}")
-            return pdb_path
-
-        return _write_pdb_with_fallback(atoms, pdb_path)
-
+        return prepare_universal_molecule_viewer(atoms)
     except Exception as e:
         print(f"Error preparing input for viewer: {e}")
         return None
 
+def prepare_molecule_for_viewer(traj_path):
+    """Universal trajectory preparation"""
+    return prepare_universal_molecule_viewer(traj_path)
+
 # ==== OrbMol SPE ====
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-#_MODEL_CALC = None
-#def _load_orbmol_calc():
-#    global _MODEL_CALC
-#    if _MODEL_CALC is None:
-#        orbff = pretrained.orb_v3_conservative_inf_omat(
-#            device="cpu", precision="float32-high"
-#        )
-#        _MODEL_CALC = ORBCalculator(orbff, device="cpu")
-#    return _MODEL_CALC
-
-def predict_molecule(structure_file,task_name,charge=0, spin_multiplicity=1):
+def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
     """
     Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
     """
@@ -154,7 +188,8 @@ def predict_molecule(structure_file,task_name,charge=0, spin_multiplicity=1):
 
         return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file
     except Exception as e:
-        import traceback; traceback.print_exc()
+        import traceback
+        traceback.print_exc()
         return f"Error during calculation: {e}", "Error", None
 
 # ==== Simulaciones (helpers) ====
@@ -164,7 +199,7 @@ from simulation_scripts_orbmol import (
 )
 
 # ==== Wrappers con debug y Molecule3D ====
-def md_wrapper(structure_file,task_name, charge, spin, steps, tempK, timestep_fs, ensemble):
+def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_fs, ensemble):
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
@@ -193,10 +228,11 @@ def md_wrapper(structure_file,task_name, charge, spin, steps, tempK, timestep_fs
 
     except Exception as e:
         print(f"MD Wrapper Error: {e}")
-        import traceback; traceback.print_exc()
+        import traceback
+        traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
-def relax_wrapper(structure_file,task_name, steps, fmax, charge, spin, relax_cell):
+def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_cell):
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
@@ -223,7 +259,8 @@ def relax_wrapper(structure_file,task_name, steps, fmax, charge, spin, relax_cel
 
     except Exception as e:
         print(f"Relax Wrapper Error: {e}")
-        import traceback; traceback.print_exc()
+        import traceback
+        traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
 # ==== UI ====
@@ -242,7 +279,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_count="single"
                     )
                     with gr.Row():
-                        task_name = gr.Radio(
+                        task_name_spe = gr.Radio(
                             ["OMol", "OMat", "OMol-Direct"],
                             value="OMol",
                             label="Model Type",
@@ -259,26 +296,21 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
                     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
-                            }
-                        ]
+                        reps=UNIVERSAL_MOLECULE3D_REPS
                     )
+                    
                     # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name.input(
+                    task_name_spe.input(
                         lambda x: (
                             (gr.Number(visible=True), gr.Number(visible=True))
                             if x == "OMol" or x == "OMol-Direct"
                             else (gr.Number(visible=False), gr.Number(visible=False))
-                            ),
-                            [task_name],
-                            [charge_input, spin_input],
-                            )
+                        ),
+                        [task_name_spe],
+                        [charge_input, spin_input],
+                    )
                     
-            run_spe.click(predict_molecule, [xyz_input, task_name,charge_input, spin_input], [spe_out, spe_status, spe_viewer])
+            run_spe.click(predict_molecule, [xyz_input, task_name_spe, charge_input, spin_input], [spe_out, spe_status, spe_viewer])
 
         # -------- MD --------
         with gr.Tab("Molecular Dynamics"):
@@ -293,7 +325,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_count="single"
                     )
                     with gr.Row():
-                        task_name = gr.Radio(
+                        task_name_md = gr.Radio(
                             ["OMol", "OMat", "OMol-Direct"],
                             value="OMol",
                             label="Model Type",
@@ -315,37 +347,26 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     md_viewer = Molecule3D(
                         label="Final Structure Viewer (Last MD Frame)",
-                        reps=[
-                            {
-                                "model": 0, "chain": "", "resname": "",
-                                "style": "stick", "color": "whiteCarbon",
-                                "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 1.0
-                            },
-                            {
-                                "model": 0, "chain": "", "resname": "",
-                                "style": "sphere", "color": "whiteCarbon",
-                                "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 0.7
-                            }
-                        ]
+                        reps=UNIVERSAL_MOLECULE3D_REPS
                     )
                     md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                     md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     md_explain = gr.Markdown()
                     
                     # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name.input(
+                    task_name_md.input(
                         lambda x: (
                             (gr.Number(visible=True), gr.Number(visible=True))
                             if x == "OMol" or x == "OMol-Direct"
                             else (gr.Number(visible=False), gr.Number(visible=False))
-                            ),
-                            [task_name],
-                            [charge_md, spin_md],
-                            )
+                        ),
+                        [task_name_md],
+                        [charge_md, spin_md],
+                    )
                     
             run_md_btn.click(
                 md_wrapper,
-                inputs=[xyz_md, task_name,charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
+                inputs=[xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
                 outputs=[md_status, md_traj, md_log, md_script, md_explain, md_viewer],
             )
 
@@ -362,7 +383,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         file_count="single"
                     )
                     with gr.Row():
-                        task_name = gr.Radio(
+                        task_name_rlx = gr.Radio(
                             ["OMol", "OMat", "OMol-Direct"],
                             value="OMol",
                             label="Model Type",
@@ -382,42 +403,30 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     rlx_viewer = Molecule3D(
                         label="Optimized Structure Viewer",
-                        reps=[
-                            {
-                                "model": 0, "chain": "", "resname": "",
-                                "style": "stick", "color": "whiteCarbon",
-                                "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 1.0
-                            },
-                            {
-                                "model": 0, "chain": "", "resname": "",
-                                "style": "sphere", "color": "whiteCarbon",
-                                "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 0.7
-                            }
-                        ]
+                        reps=UNIVERSAL_MOLECULE3D_REPS
                     )
                     rlx_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                     rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     rlx_explain = gr.Markdown()
 
                     # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name.input(
+                    task_name_rlx.input(
                         lambda x: (
                             (gr.Number(visible=True), gr.Number(visible=True))
                             if x == "OMol" or x == "OMol-Direct"
                             else (gr.Number(visible=False), gr.Number(visible=False))
-                            ),
-                            [task_name],
-                            [charge_rlx, spin_rlx],
-                            )
+                        ),
+                        [task_name_rlx],
+                        [charge_rlx, spin_rlx],
+                    )
                     
             run_rlx_btn.click(
                 relax_wrapper,
-                inputs=[xyz_rlx,task_name, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
+                inputs=[xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
                 outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer],
             )
 
 print("Starting OrbMol model loading…")
-#_ = load_orbmol_model(task_name)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)