Update app.py

app.py

@@ -1,90 +1,23 @@
-import gradio as gr
-import numpy as np
-import tempfile
 import os
+import tempfile
+import numpy as np
+import gradio as gr
 
 from ase.io import read
-from ase import units
 from ase.optimize import LBFGS
-from ase.md.verlet import VelocityVerlet
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase.md.verlet import VelocityVerlet
 from ase.io.trajectory import Trajectory
+from ase import units
 
-# Intentar importar Molecule3D para vista 3D nativa
+# Visualizador 3D “pro” si está disponible
 try:
     from gradio_molecule3d import Molecule3D
     HAVE_MOL3D = True
 except Exception:
     HAVE_MOL3D = False
 
-from orb_models.forcefield import pretrained
-from orb_models.forcefield.calculator import ORBCalculator
-
-
-# =========================
-#   OrbMol global model
-# =========================
-model_calc = None
-
-def load_orbmol_model():
-    global model_calc
-    if model_calc is None:
-        try:
-            print("Loading OrbMol model...")
-            orbff = pretrained.orb_v3_conservative_inf_omat(
-                device="cpu",
-                precision="float32-high"
-            )
-            model_calc = ORBCalculator(orbff, device="cpu")
-            print("OrbMol model loaded successfully")
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            model_calc = None
-    return model_calc
-
-
-# =========================
-#   Single-point (SPE)
-# =========================
-def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
-    try:
-        calc = load_orbmol_model()
-        if calc is None:
-            return "Error: Could not load OrbMol model", ""
-
-        if not xyz_content.strip():
-            return "Error: Please enter XYZ coordinates", ""
-
-        with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
-            f.write(xyz_content)
-            xyz_file = f.name
-
-        atoms = read(xyz_file)
-        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
-        atoms.calc = calc
-
-        energy = atoms.get_potential_energy()  # eV
-        forces = atoms.get_forces()            # eV/Å
-
-        lines = []
-        lines.append(f"Total Energy: {energy:.6f} eV\n")
-        lines.append("Atomic Forces:")
-        for i, f in enumerate(forces):
-            lines.append(f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å")
-
-        max_force = float(np.max(np.linalg.norm(forces, axis=1)))
-        lines.append(f"\nMax Force: {max_force:.4f} eV/Å")
-
-        os.unlink(xyz_file)
-        return "\n".join(lines), "Calculation completed with OrbMol"
-
-    except Exception as e:
-        return f"Error during calculation: {str(e)}", "Error"
-
-
-# =========================
-#   Trajectory → HTML 3D fallback
-# =========================
+# --- Helpers de visualización: fallback 3Dmol.js si no hay Molecule3D ---
 def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
     traj = Trajectory(traj_path)
     xyz_frames = []
@@ -120,21 +53,24 @@ def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
 """
     return html
 
+# --- OrbMol directo para SPE ---
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
 
-# =========================
-#   MD with OrbMol
-# =========================
-def run_md(xyz_content, charge=0, spin_multiplicity=1,
-           steps=100, temperature=300, timestep=1.0):
-    try:
-        calc = load_orbmol_model()
-        if calc is None:
-            return "Error: Could not load OrbMol model", ""
+_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(xyz_content, charge=0, spin_multiplicity=1):
+    try:
+        calc = _load_orbmol_calc()
         if not xyz_content.strip():
-            return "Error: Please enter XYZ coordinates", ""
+            return "Error: Please enter XYZ coordinates", "Error"
 
-        # Leer estructura
         with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
             f.write(xyz_content)
             xyz_file = f.name
@@ -143,56 +79,109 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1,
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        # Pre-relajación ligera
-        opt = LBFGS(atoms, logfile=None)
-        opt.run(fmax=0.05, steps=20)
-
-        MaxwellBoltzmannDistribution(atoms, temperature_K=2*float(temperature))
-
-        dyn = VelocityVerlet(atoms, timestep=float(timestep) * units.fs)
-
-        tf = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
-        tf.close()
-        traj = Trajectory(tf.name, "w", atoms)
-        dyn.attach(traj.write, interval=1)
-        dyn.run(int(steps))
+        energy = atoms.get_potential_energy()  # eV
+        forces = atoms.get_forces()            # eV/Å
 
-        if HAVE_MOL3D:
-            # Mostrar último frame en Molecule3D
-            last = traj[-1]
-            mol_xyz = f"{len(last)}\nFinal frame\n"
-            for s, (x, y, z) in zip(last.get_chemical_symbols(), last.get_positions()):
-                mol_xyz += f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
-            view = Molecule3D(value=mol_xyz, label="Final Frame (XYZ)")
-        else:
-            view = traj_to_html(tf.name)
+        lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
+        for i, f in enumerate(forces):
+            lines.append(f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å")
+        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
 
-        return f"MD completed: {int(steps)} steps at {int(temperature)} K", view
+        return "\n".join(lines), "Calculation completed with OrbMol"
+    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) ---
+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):
+    try:
+        traj_path, log_text, script_text, explanation = run_md_simulation(
+            xyz_content,            # acepta string XYZ
+            int(steps),
+            20,                     # pre-relax steps como en la UI de UMA
+            float(timestep_fs),
+            float(tempK),
+            "NVT" if ensemble == "NVT" else "NVE",
+            int(charge),
+            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
+            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
+        )
     except Exception as e:
-        return f"Error during MD simulation: {str(e)}", ""
+        return (f"Error: {e}", None, None, None, "", "", "")
 
+def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
+    try:
+        traj_path, log_text, script_text, explanation = run_relaxation_simulation(
+            xyz_content,
+            int(steps),
+            float(fmax),
+            int(charge),
+            int(spin),
+            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
+        else:
+            viewer_value = None
+            html_value = traj_to_html(traj_path)
+
+        return (
+            status,
+            viewer_value,
+            html_value,
+            traj_path,
+            log_text,
+            script_text,
+            explanation,
+        )
+    except Exception as e:
+        return (f"Error: {e}", None, None, None, "", "", "")
 
-# =========================
-#   Ejemplos
-# =========================
+# ------------------------
+#   Ejemplos rápidos
+# ------------------------
 examples = [
     ["""2
 Hydrogen molecule
 H 0.0 0.0 0.0
 H 0.0 0.0 0.74""", 0, 1],
-
     ["""3
-Water molecule  
+Water molecule
 O 0.0000 0.0000 0.0000
 H 0.7571 0.0000 0.5864
 H -0.7571 0.0000 0.5864""", 0, 1],
-
     ["""5
 Methane
 C 0.0000 0.0000 0.0000
@@ -202,78 +191,132 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1],
 ]
 
-
-# =========================
-#   Gradio UI
-# =========================
+# ------------------------
+#   UI Gradio
+# ------------------------
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
-        # -------- Tab 1: Single Point --------
+        # ========== SPE ==========
         with gr.Tab("Single Point Energy"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    with gr.Column(variant="panel"):
-                        gr.Markdown("# OrbMol Demo - Quantum-Accurate Predictions")
-                        gr.Markdown("OrbMol is a neural network potential trained on the OMol25 dataset.")
-
-                        xyz_input = gr.Textbox(
-                            label="XYZ Coordinates",
-                            placeholder="Paste XYZ here...",
-                            lines=12,
-                        )
-
-                        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")
-
-                        predict_btn = gr.Button("Run OrbMol Prediction", variant="primary")
+                    gr.Markdown("# OrbMol Demo — Quantum-Accurate Molecular Predictions")
+                    gr.Markdown(
+                        "Predict **energies** and **forces** with OrbMol (OMol25). "
+                        "Supports **charge** and **spin multiplicity**."
+                    )
+
+                    xyz_input = gr.Textbox(
+                        label="XYZ Coordinates",
+                        placeholder="Paste XYZ here...",
+                        lines=12,
+                    )
+                    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")
+                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
 
                 with gr.Column(variant="panel", min_width=500):
-                    results_output = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
-                    status_output = gr.Textbox(label="Status", interactive=False, max_lines=1)
-
-            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
+                    spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
+                    spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
 
-            predict_btn.click(
+            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=[results_output, status_output],
+                outputs=[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\n* Supports charge and spin")
+                    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("* <1 kcal/mol error on Wiggle150\n* Accurate protein–ligand binding energies\n* Stable MD on biomolecules >20k atoms")
+                    gr.Markdown(
+                        "* Near-DFT accuracy on **GMTKN55** and **Wiggle150**\n"
+                        "* Accurate **protein–ligand** interaction energies (PLA15)\n"
+                        "* Stable long MD on biomolecules grandes"
+                    )
                 with gr.Accordion("Disclaimers", open=False):
-                    gr.Markdown("* Validate results for your use case\n* Training level of theory may limit accuracy")
+                    gr.Markdown(
+                        "* Validate for your use case\n"
+                        "* Consider training **level of theory** and intended domain"
+                    )
 
-        # -------- Tab 2: MD --------
+        # ========== MD ==========
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    xyz_input_md = gr.Textbox(label="XYZ Coordinates", lines=12)
-                    charge_input_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-                    spin_input_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
-                    steps_input = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
-                    temp_input = gr.Slider(value=300, minimum=10, maximum=1000, step=10, label="Temperature (K)")
-                    timestep_input = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
+                    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")
+                    with gr.Row():
+                        steps_md = gr.Slider(value=100, minimum=10, maximum=1000, 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)")
+                        ensemble_md = gr.Radio(choices=["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", lines=2, interactive=False)
-                    md_view = gr.HTML() if not HAVE_MOL3D else Molecule3D(label="Trajectory Viewer")
+                    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
+                    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_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
+                    md_explain = gr.Markdown()
 
             run_md_btn.click(
-                run_md,
-                inputs=[xyz_input_md, charge_input_md, spin_input_md, steps_input, temp_input, timestep_input],
-                outputs=[md_status, md_view],
+                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],
             )
 
+        # ========== 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")
+                    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")
+                        spin_rlx = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin")
+                    relax_cell = gr.Checkbox(label="Relax Unit Cell", value=False)
+                    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)
+                    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_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],
+            )
 
-print("Starting OrbMol model loading...")
-load_orbmol_model()
+print("Starting OrbMol model loading…")
+_load_orbmol_calc()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)