Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import torch
 import numpy as np
 import tempfile
 import os
@@ -10,22 +9,28 @@ from ase.optimize import LBFGS
 from ase.md.verlet import VelocityVerlet
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.io.trajectory import Trajectory
-import py3Dmol
+
+# py3Dmol es opcional; si no está instalado haremos fallback a 3Dmol.js
+try:
+    import py3Dmol  # noqa: F401
+    HAVE_PY3DMOL = True
+except Exception:
+    HAVE_PY3DMOL = False
 
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
 
-# -----------------------------
-# Global OrbMol model
-# -----------------------------
+# =========================
+#   OrbMol global model
+# =========================
 model_calc = None
 
 def load_orbmol_model():
-    """Load OrbMol model once"""
     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"
@@ -33,14 +38,14 @@ def load_orbmol_model():
             model_calc = ORBCalculator(orbff, device="cpu")
             print("OrbMol model loaded successfully")
         except Exception as e:
-            print(f"Error loading OrbMol model: {e}")
+            print(f"Error loading model: {e}")
             model_calc = None
     return model_calc
 
 
-# -----------------------------
-# Single-point calculation
-# -----------------------------
+# =========================
+#   Single-point (SPE)
+# =========================
 def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
     try:
         calc = load_orbmol_model()
@@ -50,7 +55,7 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         if not xyz_content.strip():
             return "Error: Please enter XYZ coordinates", ""
 
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
             f.write(xyz_content)
             xyz_file = f.name
 
@@ -58,45 +63,90 @@ 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()
-        forces = atoms.get_forces()
+        energy = atoms.get_potential_energy()  # eV
+        forces = atoms.get_forces()            # eV/Å
 
-        result = f"Total Energy: {energy:.6f} eV\n\nAtomic Forces:\n"
+        lines = []
+        lines.append(f"Total Energy: {energy:.6f} eV")
+        lines.append("")
+        lines.append("Atomic Forces:")
         for i, f in enumerate(forces):
-            result += f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å\n"
-
-        max_force = np.max(np.linalg.norm(forces, axis=1))
-        result += f"\nMax Force: {max_force:.4f} eV/Å"
+            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("")
+        lines.append(f"Max Force: {max_force:.4f} eV/Å")
 
         os.unlink(xyz_file)
-        return result, "Calculation completed with OrbMol"
+        return "\n".join(lines), "Calculation completed with OrbMol"
+
     except Exception as e:
         return f"Error during calculation: {str(e)}", "Error"
 
 
-# -----------------------------
-# Trajectory → HTML
-# -----------------------------
-def traj_to_html(traj_file):
-    traj = Trajectory(traj_file)
-    view = py3Dmol.view(width=400, height=400)
+# =========================
+#   Trajectory → HTML 3D
+#   (funciona con o sin py3Dmol)
+# =========================
+def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
+    """
+    Genera un bloque HTML con 3Dmol para animar una lista de frames en XYZ.
+    Sin dependencias de JupyterLab.
+    """
+    traj = Trajectory(traj_path)
+    xyz_frames = []
+
     for atoms in traj:
         symbols = atoms.get_chemical_symbols()
-        xyz = atoms.get_positions()
-        mol = ""
-        for s, (x, y, z) in zip(symbols, xyz):
-            mol += f"{s} {x} {y} {z}\n"
-        view.addModel(mol, "xyz")
-    view.setStyle({"stick": {}})
-    view.zoomTo()
-    view.animate({"loop": "forward"})
-    return view.render()
-
-
-# -----------------------------
-# MD simulation
-# -----------------------------
-def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=300, timestep=1.0):
+        coords = atoms.get_positions()
+        parts = [str(len(symbols)), "frame"]
+        for s, (x, y, z) in zip(symbols, coords):
+            parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
+        xyz_frames.append("\n".join(parts))
+
+    # HTML+JS autopackage
+    html = f"""
+<div id="viewer_md" style="width:{width}px; height:{height}px; position:relative;"></div>
+<script src="https://3dmol.org/build/3Dmol-min.js"></script>
+<script>
+(function() {{
+  var element = document.getElementById('viewer_md');
+  if (!element) return;
+  var viewer = $3Dmol.createViewer(element, {{backgroundColor: 'white'}});
+  var frames = {xyz_frames!r};
+  var i = 0;
+
+  function show(i) {{
+    viewer.clear();
+    viewer.addModel(frames[i], 'xyz');
+    viewer.setStyle({{}}, {{stick: {{}} }});
+    viewer.zoomTo();
+    viewer.render();
+  }}
+
+  if (frames.length === 0) {{
+    element.innerHTML = '<div style="padding:8px;color:#555">Empty trajectory</div>';
+  }} else if (frames.length === 1) {{
+    show(0);
+  }} else {{
+    show(0);
+    setInterval(function() {{
+      i = (i + 1) % frames.length;
+      show(i);
+    }}, {int(interval_ms)});
+  }}
+}})();
+</script>
+"""
+    return html
+
+
+# =========================
+#   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:
@@ -105,7 +155,8 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
         if not xyz_content.strip():
             return "Error: Please enter XYZ coordinates", ""
 
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
+        # Leer estructura
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
             f.write(xyz_content)
             xyz_file = f.name
 
@@ -113,99 +164,184 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        # Pre-relaxation
-        opt = LBFGS(atoms)
+        # Pre-relajación ligera
+        opt = LBFGS(atoms, logfile=None)
         opt.run(fmax=0.05, steps=20)
 
-        # Initialize velocities
-        MaxwellBoltzmannDistribution(atoms, temperature_K=2 * temperature)
+        # Velocidades (2*T por partición de energía)
+        MaxwellBoltzmannDistribution(atoms, temperature_K=2*float(temperature))
 
-        # Run MD
-        dyn = VelocityVerlet(atoms, timestep=timestep * units.fs)
-        traj_file = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
-        traj = Trajectory(traj_file.name, "w", atoms)
+        # Integrador NVE
+        dyn = VelocityVerlet(atoms, timestep=float(timestep) * units.fs)
+
+        # Trayectoria temporal
+        tf = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
+        tf.close()
+        traj = Trajectory(tf.name, "w", atoms)
         dyn.attach(traj.write, interval=1)
-        dyn.run(steps)
 
-        html = traj_to_html(traj_file.name)
+        # Ejecutar MD
+        dyn.run(int(steps))
+
+        # Visualización 3D
+        html = traj_to_html(tf.name)
+
+        # Limpieza archivo xyz temporal
+        try:
+            os.unlink(xyz_file)
+        except Exception:
+            pass
+
+        return f"MD completed: {int(steps)} steps at {int(temperature)} K", html
 
-        os.unlink(xyz_file)
-        return f"MD completed: {steps} steps at {temperature} K", html
     except Exception as e:
         return f"Error during MD simulation: {str(e)}", ""
 
 
-# -----------------------------
-# Examples
-# -----------------------------
+# =========================
+#   Ejemplos (tu set)
+# =========================
 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
 H 1.0890 0.0000 0.0000
 H -0.3630 1.0267 0.0000
 H -0.3630 -0.5133 0.8887
-H -0.3630 -0.5133 -0.8887""", 0, 1]
+H -0.3630 -0.5133 -0.8887""", 0, 1],
 ]
 
 
-# -----------------------------
-# Gradio UI
-# -----------------------------
+# =========================
+#   Gradio UI
+# =========================
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
+    with gr.Tabs():
+        # -------- Tab 1: tu layout original --------
+        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 Molecular Predictions")
+
+                        gr.Markdown(
+                            "OrbMol is a neural network potential trained on the OMol25 dataset. "
+                            "Predict energies and forces with configurable charge and spin."
+                        )
+
+                        gr.Markdown("## Simulation inputs")
+                        with gr.Column(variant="panel"):
+                            gr.Markdown("### Input molecular structure")
+
+                            xyz_input = gr.Textbox(
+                                label="XYZ Coordinates",
+                                placeholder=(
+                                    "3\nWater molecule\n"
+                                    "O 0.0000 0.0000 0.0000\n"
+                                    "H 0.7571 0.0000 0.5864\n"
+                                    "H -0.7571 0.0000 0.5864"
+                                ),
+                                lines=12,
+                                info="Paste XYZ coordinates of your molecule here",
+                            )
+
+                            gr.Markdown("OMol-specific settings for total charge and spin multiplicity")
+                            with gr.Row():
+                                charge_input = gr.Slider(
+                                    value=0, label="Total Charge", minimum=-10, maximum=10, step=1
+                                )
+                                spin_input = gr.Slider(
+                                    value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity"
+                                )
+
+                            predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
+
+                with gr.Column(variant="panel", elem_id="results", min_width=500):
+                    gr.Markdown("## OrbMol Prediction Results")
+
+                    results_output = gr.Textbox(
+                        label="Energy & Forces",
+                        lines=15,
+                        interactive=False,
+                        info="OrbMol energy and force predictions",
+                    )
+
+                    status_output = gr.Textbox(
+                        label="Status",
+                        interactive=False,
+                        max_lines=1,
+                    )
+
+            gr.Markdown("### Examples")
+            gr.Examples(
+                examples=examples,
+                inputs=[xyz_input, charge_input, spin_input],
+                label="Click any example to load it",
+            )
+
+            predict_btn.click(
+                predict_molecule,
+                inputs=[xyz_input, charge_input, spin_input],
+                outputs=[results_output, status_output],
+            )
 
-    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 Molecular Predictions")
-                    xyz_input = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here")
-                    with gr.Row():
-                        charge_input = gr.Slider(value=0, label="Charge", minimum=-10, maximum=10, step=1)
-                        spin_input = gr.Slider(value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity")
-                    predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
-            with gr.Column(variant="panel", min_width=500):
-                gr.Markdown("## Results")
-                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])
-
-        predict_btn.click(
-            predict_molecule,
-            inputs=[xyz_input, charge_input, spin_input],
-            outputs=[results_output, status_output]
-        )
-
-    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)")
-                run_md_btn = gr.Button("Run MD Simulation", variant="primary")
-                md_status = gr.Textbox(label="MD Status", lines=2)
-            with gr.Column(scale=1, variant="panel"):
-                gr.Markdown("## MD Visualization")
-                md_view = gr.HTML()
-
-        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],
-        )
+            with gr.Sidebar(open=True):
+                gr.Markdown("## Learn more about OrbMol")
+                with gr.Accordion("What is OrbMol?", open=False):
+                    gr.Markdown(
+                        "* OrbMol is a neural network potential for molecular property prediction\n"
+                        "* Built on the Orb-v3 architecture and trained on OMol25 dataset\n"
+                        "* Supports configurable charge and spin multiplicity"
+                    )
+                with gr.Accordion("Model Disclaimers", open=False):
+                    gr.Markdown(
+                        "* Validate results for your use case\n"
+                        "* Consider limitations of the training level of theory"
+                    )
+                with gr.Accordion("Open source packages", open=False):
+                    gr.Markdown(
+                        "* Model code: orbital-materials/orb-models\n"
+                        "* This demo uses ASE and Gradio"
+                    )
+
+        # -------- Tab 2: MD + Visualización 3D --------
+        with gr.Tab("Molecular Dynamics"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    with gr.Column(variant="panel"):
+                        xyz_input_md = gr.Textbox(
+                            label="XYZ Coordinates",
+                            lines=12,
+                            placeholder="Paste XYZ here",
+                        )
+                        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)")
+                        run_md_btn = gr.Button("Run MD Simulation", variant="primary")
+
+                with gr.Column(variant="panel", min_width=520):
+                    gr.Markdown("## MD Visualization")
+                    md_status = gr.Textbox(label="MD Status", lines=2, interactive=False)
+                    md_view = gr.HTML()
+
+            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],
+            )
 
 
 print("Starting OrbMol model loading...")