Update app.py

upload universal functions from gradio_molecule3d.

app.py

@@ -6,159 +6,42 @@ 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
+from simulation_scripts_orbmol import load_orbmol_model, run_md_simulation, run_relaxation_simulation
 import hashlib
-import shutil
 
-# ==== Configuración UNIVERSAL para Molecule3D ====
-UNIVERSAL_MOLECULE3D_REPS = [
+# ==== Configuración Molecule3D (igual que Facebook UMA) ====
+DEFAULT_MOLECULAR_REPRESENTATIONS = [
     {
         "model": 0,
         "chain": "",
         "resname": "",
         "style": "sphere",
-        "color": "element",  # Colores automáticos por elemento químico
-        "residue_range": "",
+        "color": "Jmol",
         "around": 0,
         "byres": False,
-        "visible": True,
-        "opacity": 0.8
+        "scale": 0.3,
     },
     {
         "model": 0,
         "chain": "",
         "resname": "",
         "style": "stick",
-        "color": "element",
-        "residue_range": "",
+        "color": "Jmol",
         "around": 0,
         "byres": False,
-        "visible": True,
-        "opacity": 1.0
-    }
+        "scale": 0.2,
+    },
 ]
 
-# ==== Función mejorada para crear PDB universal ====
-def create_universal_pdb(atoms, cache_key):
-    """
-    Crea un PDB optimizado que funciona bien con gradio_molecule3d para cualquier molécula
-    """
-    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:
-        return pdb_path
-    
-    try:
-        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"✅ 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
-
-# ==== 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 in prepare_universal_molecule_viewer: {e}")
-        import traceback
-        traceback.print_exc()
-        return None
-
-# ==== Funciones de preparación universales ====
-def prepare_input_for_viewer(structure_file):
-    """Universal input preparation"""
-    if not structure_file or not os.path.exists(structure_file):
-        return None
-    try:
-        atoms = read(structure_file)
-        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)
+DEFAULT_MOLECULAR_SETTINGS = {
+    "backgroundColor": "white",
+    "orthographic": False,
+    "disableFog": False,
+}
 
 # ==== OrbMol SPE ====
-from orb_models.forcefield import pretrained
-from orb_models.forcefield.calculator import ORBCalculator
-
 def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
-    """
-    Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
-    """
+    """Single Point Energy + fuerzas (OrbMol)"""
     try:
         calc = load_orbmol_model(task_name)
         if not structure_file:
@@ -171,46 +54,44 @@ def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
             return f"Error: Empty file: {file_path}", "Error", None
 
         atoms = read(file_path)
-        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+        
+        # Solo aplicar charge/spin para OMol
+        if task_name in ["OMol", "OMol-Direct"]:
+            atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+        
         atoms.calc = calc
-
         energy = atoms.get_potential_energy()
         forces = atoms.get_forces()
 
-        lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
+        lines = [
+            f"Model: {task_name}",
+            f"Total Energy: {energy:.6f} eV",
+            "",
+            "Atomic Forces:"
+        ]
         for i, fc in enumerate(forces):
             lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å")
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
         lines += ["", f"Max Force: {max_force:.4f} eV/Å"]
 
-        # Preparar PDB para visualización
-        pdb_file = prepare_input_for_viewer(file_path)
-
-        return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file
+        # CLAVE: Devolver archivo original, NO un PDB procesado
+        return "\n".join(lines), f"✅ Calculation completed with {task_name}", file_path
+        
     except Exception as e:
         import traceback
         traceback.print_exc()
         return f"Error during calculation: {e}", "Error", None
 
-# ==== Simulaciones (helpers) ====
-from simulation_scripts_orbmol import (
-    run_md_simulation,
-    run_relaxation_simulation,
-)
-
-# ==== Wrappers con debug y Molecule3D ====
+# ==== Wrappers MD y Relax ====
 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)
 
-        file_path = structure_file
-        print(f"MD Wrapper: Processing {file_path}")
-
         traj_path, log_text, script_text, explanation = run_md_simulation(
-            file_path,
+            structure_file,
             int(steps),
-            20,  # pre-relax steps
+            20,
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -218,16 +99,12 @@ def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_f
             int(charge),
             int(spin),
         )
-        status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
-        print(f"MD completed, trajectory: {traj_path}")
+        status = f"✅ MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
-        pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for Molecule3D: {pdb_file}")
-
-        return (status, traj_path, log_text, script_text, explanation, pdb_file)
+        # CLAVE: Devolver traj_path directamente
+        return (status, traj_path, log_text, script_text, explanation, traj_path)
 
     except Exception as e:
-        print(f"MD Wrapper Error: {e}")
         import traceback
         traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
@@ -237,11 +114,8 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
 
-        file_path = structure_file
-        print(f"Relax Wrapper: Processing {file_path}")
-
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
-            file_path,
+            structure_file,
             int(steps),
             float(fmax),
             str(task_name),
@@ -249,16 +123,12 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
             int(spin),
             bool(relax_cell),
         )
-        status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
-        print(f"Relaxation completed, trajectory: {traj_path}")
+        status = f"✅ Relaxation finished (≤{int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
-        pdb_file = prepare_molecule_for_viewer(traj_path)
-        print(f"PDB file for Molecule3D: {pdb_file}")
-
-        return (status, traj_path, log_text, script_text, explanation, pdb_file)
+        # CLAVE: Devolver traj_path directamente
+        return (status, traj_path, log_text, script_text, explanation, traj_path)
 
     except Exception as e:
-        print(f"Relax Wrapper Error: {e}")
         import traceback
         traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
@@ -271,103 +141,113 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
-                    gr.Markdown("Upload molecular structure files (.xyz, .pdb, .cif, .traj) for energy and force calculations.")
+                    gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf")
                     
                     xyz_input = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
+                        label="Upload Structure File",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
+                    task_name_spe = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
-                        task_name_spe = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type",
-                            info="Choose the OrbMol model variant for the calculation."
-                        )
-                    with gr.Row():
-                        charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
-                        spin_input = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
+                        charge_input = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_input = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
                     
                     run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
                     
                 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)
+                    spe_status = gr.Textbox(label="Status", interactive=False)
+                    
+                    # PATRÓN FACEBOOK UMA: inputs + value=lambda
                     spe_viewer = Molecule3D(
                         label="Input Structure Viewer",
-                        reps=UNIVERSAL_MOLECULE3D_REPS
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS,
+                        render=False,
+                        inputs=[xyz_input],
+                        value=lambda x: x,
+                        interactive=False
                     )
                     
-                    # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name_spe.input(
+                    task_name_spe.change(
                         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))
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
                         ),
                         [task_name_spe],
-                        [charge_input, spin_input],
+                        [charge_input, spin_input]
                     )
                     
-            run_spe.click(predict_molecule, [xyz_input, task_name_spe, 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"):
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("## Molecular Dynamics Simulation")
-                    gr.Markdown("Upload your molecular structure and configure MD parameters.")
                     
                     xyz_md = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
+                        label="Upload Structure File",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
+                    task_name_md = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
-                        task_name_md = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type",
-                            info="Choose the OrbMol model variant for the calculation."
-                        )
-                    with gr.Row():
-                        charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
-                        spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
+                        charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
                     with gr.Row():
-                        steps_md = gr.Slider(minimum=10, maximum=2000, value=100, step=10, label="Steps")
-                        temp_md = gr.Slider(minimum=10, maximum=1500, value=300, step=10, label="Temperature (K)")
+                        steps_md = gr.Slider(10, 2000, 100, step=10, label="Steps")
+                        temp_md = gr.Slider(10, 1500, 300, step=10, label="Temperature (K)")
                     with gr.Row():
-                        timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Timestep (fs)")
+                        timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)")
                         ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
                     run_md_btn = gr.Button("Run MD Simulation", variant="primary")
 
                 with gr.Column(variant="panel", min_width=520):
                     md_status = gr.Textbox(label="MD Status", interactive=False)
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+                    
+                    # PATRÓN FACEBOOK UMA
                     md_viewer = Molecule3D(
-                        label="Final Structure Viewer (Last MD Frame)",
-                        reps=UNIVERSAL_MOLECULE3D_REPS
+                        label="MD Result Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS,
+                        render=False,
+                        inputs=[md_traj],
+                        value=lambda x: x,
+                        interactive=False
                     )
-                    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_log = gr.Textbox(label="Log", interactive=False, lines=15)
+                    md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
                     md_explain = gr.Markdown()
                     
-                    # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name_md.input(
+                    task_name_md.change(
                         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))
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
                         ),
                         [task_name_md],
-                        [charge_md, spin_md],
+                        [charge_md, spin_md]
                     )
                     
             run_md_btn.click(
                 md_wrapper,
-                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],
+                [xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
+                [md_status, md_traj, md_log, md_script, md_explain, md_viewer]
             )
 
         # -------- Relax --------
@@ -375,58 +255,59 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("## Structure Relaxation/Optimization")
-                    gr.Markdown("Upload your molecular structure for geometry optimization.")
                     
                     xyz_rlx = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
+                        label="Upload Structure File",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single"
                     )
+                    task_name_rlx = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
-                        task_name_rlx = gr.Radio(
-                            ["OMol", "OMat", "OMol-Direct"],
-                            value="OMol",
-                            label="Model Type",
-                            info="Choose the OrbMol model variant for the calculation."
-                        )
-                    with gr.Row():   
-                        steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps")
-                        fmax_rlx = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)")
+                        steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps")
+                        fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)")
                     with gr.Row():
-                        charge_rlx = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
-                        spin_rlx = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
+                        charge_rlx = gr.Slider(-10, 10, 0, step=1, label="Charge")
+                        spin_rlx = gr.Slider(1, 11, 1, step=1, label="Spin")
                     relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
                     run_rlx_btn = gr.Button("Run Optimization", variant="primary")
 
                 with gr.Column(variant="panel", min_width=520):
                     rlx_status = gr.Textbox(label="Status", interactive=False)
                     rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+                    
+                    # PATRÓN FACEBOOK UMA
                     rlx_viewer = Molecule3D(
                         label="Optimized Structure Viewer",
-                        reps=UNIVERSAL_MOLECULE3D_REPS
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS,
+                        render=False,
+                        inputs=[rlx_traj],
+                        value=lambda x: x,
+                        interactive=False
                     )
-                    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_log = gr.Textbox(label="Log", interactive=False, lines=15)
+                    rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
                     rlx_explain = gr.Markdown()
 
-                    # Charge and Spin are only applicable to OMol and OMol-Direct
-                    task_name_rlx.input(
+                    task_name_rlx.change(
                         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))
+                            gr.update(visible=x in ["OMol", "OMol-Direct"]),
+                            gr.update(visible=x in ["OMol", "OMol-Direct"])
                         ),
                         [task_name_rlx],
-                        [charge_rlx, spin_rlx],
+                        [charge_rlx, spin_rlx]
                     )
                     
             run_rlx_btn.click(
                 relax_wrapper,
-                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],
+                [xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
+                [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer]
             )
 
-print("Starting OrbMol model loading…")
-
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)