Update app.py

app.py

@@ -2,19 +2,19 @@ import os
 import tempfile
 import numpy as np
 import gradio as gr
-
 from ase.io import read
 from ase.io.trajectory import Trajectory
 
-# ==== Visualizador 3D ====
+# ==== Intentar visor nativo como en UMA ====
 try:
     from gradio_molecule3d import Molecule3D
     HAVE_MOL3D = True
 except Exception:
     HAVE_MOL3D = False
 
+
+# ==== Fallback HTML con 3Dmol.js ====
 def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
-    """Fallback de visualización con 3Dmol.js si no hay Molecule3D."""
     traj = Trajectory(traj_path)
     xyz_frames = []
     for atoms in traj:
@@ -50,7 +50,7 @@ def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
     return html
 
 
-# ==== OrbMol SPE directo (rápido) ====
+# ==== OrbMol SPE directo ====
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
@@ -59,14 +59,13 @@ def _load_orbmol_calc():
     global _MODEL_CALC
     if _MODEL_CALC is None:
         orbff = pretrained.orb_v3_conservative_inf_omat(
-            device="cpu",
-            precision="float32-high"
+            device="cpu", precision="float32-high"
         )
         _MODEL_CALC = ORBCalculator(orbff, device="cpu")
     return _MODEL_CALC
 
+
 def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
-    """Single Point Energy/Forces con OrbMol (input en Textbox XYZ)."""
     try:
         calc = _load_orbmol_calc()
         if not xyz_content or not xyz_content.strip():
@@ -80,8 +79,8 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        energy = atoms.get_potential_energy()  # eV
-        forces = atoms.get_forces()            # eV/Å
+        energy = atoms.get_potential_energy()
+        forces = atoms.get_forces()
 
         lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
         for i, f in enumerate(forces):
@@ -89,30 +88,49 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
         lines += ["", f"Max Force: {max_force:.4f} eV/Å"]
 
-        try:
-            os.unlink(xyz_file)
-        except Exception:
-            pass
+        try: os.unlink(xyz_file)
+        except Exception: pass
 
         return "\n".join(lines), "Calculation completed with OrbMol"
     except Exception as e:
         return f"Error during calculation: {e}", "Error"
 
 
-# ==== Simulación (estilo UMA) vía helpers locales ====
+# ==== Simulaciones (helpers locales) ====
 from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
-    last_frame_xyz_from_traj,
 )
 
+# Convierte textbox XYZ a fichero temporal si es necesario (para ASE)
+def _string_looks_like_xyz(text: str) -> bool:
+    try:
+        first = (text or "").strip().splitlines()[0]
+        int(first.split()[0])
+        return True
+    except Exception:
+        return False
+
+def _to_file_if_xyz(input_or_path: str):
+    if isinstance(input_or_path, str) and _string_looks_like_xyz(input_or_path):
+        tf = tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False)
+        tf.write(input_or_path)
+        tf.flush(); tf.close()
+        return tf.name, True
+    return input_or_path, False
+
+
+# Wrappers: devuelven SIEMPRE (status, traj_path, log, script, explain, html_fallback)
 def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
-    """Conecta Gradio con run_md_simulation y adapta outputs (viewer/HTML/file/log/script)."""
+    tmp_created = False
+    path_or_str = xyz_content
     try:
+        path_or_str, tmp_created = _to_file_if_xyz(xyz_content)
+
         traj_path, log_text, script_text, explanation = run_md_simulation(
-            xyz_content,            # acepta string XYZ o ruta
+            path_or_str,
             int(steps),
-            20,                     # prerelax por defecto, como UMA
+            20,                          # pre-relax
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -121,30 +139,25 @@ def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
-        if HAVE_MOL3D:
-            viewer_value = last_frame_xyz_from_traj(traj_path)  # string XYZ como value
-            html_value = None
-        else:
-            viewer_value = None
-            html_value = traj_to_html(traj_path)
-
-        return (
-            status,         # md_status
-            viewer_value,   # md_viewer (Molecule3D value) o None
-            html_value,     # md_html (fallback) o None
-            traj_path,      # md_traj file
-            log_text,       # md_log (Textbox)
-            script_text,    # md_script (Code py)
-            explanation,    # md_explain
-        )
+        html_value = "" if HAVE_MOL3D else traj_to_html(traj_path)
+        return (status, traj_path, log_text, script_text, explanation, html_value)
+
     except Exception as e:
-        return (f"Error: {e}", None, None, None, "", "", "")
+        return (f"Error: {e}", None, "", "", "", "")
+    finally:
+        if tmp_created and isinstance(path_or_str, str) and os.path.exists(path_or_str):
+            try: os.remove(path_or_str)
+            except Exception: pass
+
 
 def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
-    """Conecta Gradio con run_relaxation_simulation y adapta outputs."""
+    tmp_created = False
+    path_or_str = xyz_content
     try:
+        path_or_str, tmp_created = _to_file_if_xyz(xyz_content)
+
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
-            xyz_content,
+            path_or_str,
             int(steps),
             float(fmax),
             int(charge),
@@ -153,24 +166,15 @@ def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
         )
         status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
-        if HAVE_MOL3D:
-            viewer_value = last_frame_xyz_from_traj(traj_path)
-            html_value = None
-        else:
-            viewer_value = None
-            html_value = traj_to_html(traj_path)
-
-        return (
-            status,         # rlx_status
-            viewer_value,   # rlx_viewer (Molecule3D value) o None
-            html_value,     # rlx_html (fallback) o None
-            traj_path,      # rlx_traj file
-            log_text,       # rlx_log (Textbox)
-            script_text,    # rlx_script (Code py)
-            explanation,    # rlx_explain
-        )
+        html_value = "" if HAVE_MOL3D else traj_to_html(traj_path)
+        return (status, traj_path, log_text, script_text, explanation, html_value)
+
     except Exception as e:
-        return (f"Error: {e}", None, None, None, "", "", "")
+        return (f"Error: {e}", None, "", "", "", "")
+    finally:
+        if tmp_created and isinstance(path_or_str, str) and os.path.exists(path_or_str):
+            try: os.remove(path_or_str)
+            except Exception: pass
 
 
 # ==== Ejemplos ====
@@ -194,132 +198,114 @@ H -0.3630 -0.5133 -0.8887""", 0, 1],
 ]
 
 
-# ==== UI Gradio ====
+# ==== UI ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
-        # --- Tab SPE ---
+        # -------- SPE --------
         with gr.Tab("Single Point Energy"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    gr.Markdown("# OrbMol Demo — Quantum-Accurate Molecular Predictions")
-                    gr.Markdown(
-                        "Predict **energies** and **forces** with OrbMol (trained on **OMol25**, "
-                        "ωB97M-V/def2-TZVPD). Supports **charge** and **spin multiplicity**."
-                    )
-                    xyz_input = gr.Textbox(
-                        label="XYZ Coordinates",
-                        placeholder="Paste XYZ here...",
-                        lines=12,
-                    )
+                    gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
+                    gr.Markdown("Energías y fuerzas con **charge** y **spin multiplicity**.")
+                    xyz_input = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here…")
                     with gr.Row():
-                        charge_input = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-                        spin_input = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
+                        charge_input = gr.Slider(0, -10, 10, 1, label="Charge")
+                        spin_input = gr.Slider(1, 1, 11, 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)
 
-            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input], label="Examples")
-            run_spe.click(
-                predict_molecule,
-                inputs=[xyz_input, charge_input, spin_input],
-                outputs=[spe_out, spe_status],
-            )
+            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
+            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status])
 
-            with gr.Sidebar(open=True):
-                gr.Markdown("## Learn more about OrbMol")
-                with gr.Accordion("What is OrbMol?", open=False):
-                    gr.Markdown(
-                        "* Neural network potential for molecules\n"
-                        "* Built on Orb-v3, trained on OMol25 (ωB97M-V/def2-TZVPD)\n"
-                        "* Supports charge and spin multiplicity"
-                    )
-                with gr.Accordion("Benchmarks", open=False):
-                    gr.Markdown(
-                        "* Strong results on **GMTKN55** y **Wiggle150**\n"
-                        "* Accurate **protein–ligand** energies (PLA15)\n"
-                        "* Stable MD en biomoléculas grandes"
-                    )
-                with gr.Accordion("Disclaimers", open=False):
-                    gr.Markdown(
-                        "* Verifica resultados para tu caso\n"
-                        "* Considera el **nivel de teoría** de entrenamiento"
-                    )
-
-        # --- Tab MD ---
+        # -------- MD --------
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    xyz_md = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...")
+                    xyz_md = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here…")
                     with gr.Row():
-                        charge_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-                        spin_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
+                        charge_md = gr.Slider(0, -10, 10, 1, label="Charge")
+                        spin_md = gr.Slider(1, 1, 11, 1, label="Spin Multiplicity")
                     with gr.Row():
-                        steps_md = gr.Slider(value=100, minimum=10, maximum=2000, step=10, label="Steps")
-                        temp_md = gr.Slider(value=300, minimum=10, maximum=1500, step=10, label="Temperature (K)")
+                        steps_md = gr.Slider(100, 10, 2000, 10, label="Steps")
+                        temp_md = gr.Slider(300, 10, 1500, 10, label="Temperature (K)")
                     with gr.Row():
-                        timestep_md = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
-                        ensemble_md = gr.Radio(choices=["NVE", "NVT"], value="NVE", label="Ensemble")
+                        timestep_md = gr.Slider(1.0, 0.1, 5.0, 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)
 
+                    # Archivo de trayectoria SIEMPRE (Molecule3D lo consume)
+                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+
+                    # Visor nativo si está disponible
                     if HAVE_MOL3D:
-                        md_viewer = Molecule3D(label="Trajectory Viewer")
-                        md_html = gr.HTML(visible=False)  # placeholder para layout consistente
+                        md_viewer = Molecule3D(
+                            label="Trajectory Viewer",
+                            inputs=[md_traj],
+                            value=lambda x: x,
+                            interactive=False,
+                        )
+                        md_html = gr.HTML(visible=False)
                     else:
-                        md_viewer = gr.Textbox(visible=False)  # placeholder
-                        md_html = gr.HTML()
+                        md_html = gr.HTML(label="Trajectory Viewer")
+                        md_viewer = gr.Textbox(visible=False)
 
-                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                    md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
+                    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()
 
+            # NOTA: no actualizamos md_viewer directamente; se refresca al cambiar md_traj
             run_md_btn.click(
                 md_wrapper,
                 inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
-                outputs=[md_status, md_viewer, md_html, md_traj, md_log, md_script, md_explain],
+                outputs=[md_status, md_traj, md_log, md_script, md_explain, md_html],
             )
 
-        # --- Tab Relaxation ---
+        # -------- Relax --------
         with gr.Tab("Relaxation / Optimization"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    xyz_rlx = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...")
-                    steps_rlx = gr.Slider(value=300, minimum=1, maximum=2000, step=1, label="Max Steps")
-                    fmax_rlx = gr.Slider(value=0.05, minimum=0.001, maximum=0.5, step=0.001, label="Fmax (eV/Å)")
+                    xyz_rlx = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here…")
+                    steps_rlx = gr.Slider(300, 1, 2000, 1, label="Max Steps")
+                    fmax_rlx = gr.Slider(0.05, 0.001, 0.5, 0.001, label="Fmax (eV/Å)")
                     with gr.Row():
-                        charge_rlx = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-                        spin_rlx = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin")
-                    relax_cell = gr.Checkbox(label="Relax Unit Cell", value=False)
+                        charge_rlx = gr.Slider(0, -10, 10, 1, label="Charge")
+                        spin_rlx = gr.Slider(1, 1, 11, 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)
                     if HAVE_MOL3D:
-                        rlx_viewer = Molecule3D(label="Final Structure")
+                        rlx_viewer = Molecule3D(
+                            label="Final Structure",
+                            inputs=[rlx_traj],
+                            value=lambda x: x,
+                            interactive=False,
+                        )
                         rlx_html = gr.HTML(visible=False)
                     else:
+                        rlx_html = gr.HTML(label="Final Structure")
                         rlx_viewer = gr.Textbox(visible=False)
-                        rlx_html = gr.HTML()
 
-                    rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                    rlx_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
+                    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()
 
             run_rlx_btn.click(
                 relax_wrapper,
                 inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
-                outputs=[rlx_status, rlx_viewer, rlx_html, rlx_traj, rlx_log, rlx_script, rlx_explain],
+                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_html],
             )
 
 print("Starting OrbMol model loading…")
-_LOAD = _load_orbmol_calc()
+_ = _load_orbmol_calc()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)