Update app.py

app.py

@@ -4,21 +4,17 @@ import numpy as np
 import gradio as gr
 
 from ase.io import read
-from ase.optimize import LBFGS
-from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
-from ase.md.verlet import VelocityVerlet
 from ase.io.trajectory import Trajectory
-from ase import units
 
-# Visualizador 3D “pro” si está disponible
+# ==== Visualizador 3D ====
 try:
     from gradio_molecule3d import Molecule3D
     HAVE_MOL3D = True
 except Exception:
     HAVE_MOL3D = False
 
-# --- Helpers de visualización: fallback 3Dmol.js si no hay Molecule3D ---
 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:
@@ -53,22 +49,27 @@ def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
 """
     return html
 
-# --- OrbMol directo para SPE ---
+
+# ==== OrbMol SPE directo (rápido) ====
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-_model_calc = None
+_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
+    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(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.strip():
+        if not xyz_content or not xyz_content.strip():
             return "Error: Please enter XYZ coordinates", "Error"
 
         with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
@@ -97,20 +98,21 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
     except Exception as e:
         return f"Error during calculation: {e}", "Error"
 
-# --- Importa las rutinas FAIRChem-like para MD/Relax (ya soportan string XYZ o ruta) ---
+
+# ==== Simulación (estilo UMA) vía helpers locales ====
 from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
     last_frame_xyz_from_traj,
 )
 
-# Wrappers para conectar outputs de Gradio correctamente (string XYZ / HTML, file, logs...)
 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)."""
     try:
         traj_path, log_text, script_text, explanation = run_md_simulation(
-            xyz_content,            # acepta string XYZ
+            xyz_content,            # acepta string XYZ o ruta
             int(steps),
-            20,                     # pre-relax steps como en la UI de UMA
+            20,                     # prerelax por defecto, como UMA
             float(timestep_fs),
             float(tempK),
             "NVT" if ensemble == "NVT" else "NVE",
@@ -118,28 +120,28 @@ def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
             int(spin),
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
-        # Viewer
+
         if HAVE_MOL3D:
-            xyz_final = last_frame_xyz_from_traj(traj_path)  # value para Molecule3D
-            viewer_value = xyz_final
+            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,          # Molecule3D value (o None)
-            html_value,            # HTML fallback (o None)
-            traj_path,             # File download
-            log_text,              # Log
-            script_text,           # Script
-            explanation,           # Explanation
+            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
         )
     except Exception as e:
         return (f"Error: {e}", None, None, None, "", "", "")
 
 def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
+    """Conecta Gradio con run_relaxation_simulation y adapta outputs."""
     try:
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
             xyz_content,
@@ -150,6 +152,7 @@ def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
             bool(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
@@ -158,20 +161,19 @@ def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
             html_value = traj_to_html(traj_path)
 
         return (
-            status,
-            viewer_value,
-            html_value,
-            traj_path,
-            log_text,
-            script_text,
-            explanation,
+            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
         )
     except Exception as e:
         return (f"Error: {e}", None, None, None, "", "", "")
 
-# ------------------------
-#   Ejemplos rápidos
-# ------------------------
+
+# ==== Ejemplos ====
 examples = [
     ["""2
 Hydrogen molecule
@@ -191,21 +193,19 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1],
 ]
 
-# ------------------------
-#   UI Gradio
-# ------------------------
+
+# ==== UI Gradio ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
-        # ========== SPE ==========
+        # --- Tab 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 (OMol25). "
-                        "Supports **charge** and **spin multiplicity**."
+                        "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...",
@@ -237,17 +237,17 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     )
                 with gr.Accordion("Benchmarks", open=False):
                     gr.Markdown(
-                        "* Near-DFT accuracy on **GMTKN55** and **Wiggle150**\n"
-                        "* Accurate **protein–ligand** interaction energies (PLA15)\n"
-                        "* Stable long MD on biomolecules grandes"
+                        "* 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(
-                        "* Validate for your use case\n"
-                        "* Consider training **level of theory** and intended domain"
+                        "* Verifica resultados para tu caso\n"
+                        "* Considera el **nivel de teoría** de entrenamiento"
                     )
 
-        # ========== MD ==========
+        # --- Tab MD ---
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
@@ -256,7 +256,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         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")
                     with gr.Row():
-                        steps_md = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
+                        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)")
                     with gr.Row():
                         timestep_md = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
@@ -265,15 +265,16 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
 
                 with gr.Column(variant="panel", min_width=520):
                     md_status = gr.Textbox(label="MD Status", interactive=False)
+
                     if HAVE_MOL3D:
                         md_viewer = Molecule3D(label="Trajectory Viewer")
-                        md_html = gr.HTML(visible=False)  # placeholder para consistencia de outputs
+                        md_html = gr.HTML(visible=False)  # placeholder para layout consistente
                     else:
                         md_viewer = gr.Textbox(visible=False)  # placeholder
                         md_html = gr.HTML()
 
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                    md_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25)
+                    md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
                     md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     md_explain = gr.Markdown()
 
@@ -283,12 +284,12 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 outputs=[md_status, md_viewer, md_html, md_traj, md_log, md_script, md_explain],
             )
 
-        # ========== Relaxation ==========
+        # --- Tab Relaxation ---
         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=1000, step=1, label="Max Steps")
+                    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/Å)")
                     with gr.Row():
                         charge_rlx = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
@@ -298,14 +299,16 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
 
                 with gr.Column(variant="panel", min_width=520):
                     rlx_status = gr.Textbox(label="Status", interactive=False)
+
                     if HAVE_MOL3D:
                         rlx_viewer = Molecule3D(label="Final Structure")
                         rlx_html = gr.HTML(visible=False)
                     else:
                         rlx_viewer = gr.Textbox(visible=False)
                         rlx_html = gr.HTML()
+
                     rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                    rlx_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25)
+                    rlx_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
                     rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                     rlx_explain = gr.Markdown()
 
@@ -316,7 +319,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             )
 
 print("Starting OrbMol model loading…")
-_load_orbmol_calc()
+_LOAD = _load_orbmol_calc()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)