Update app.py

app.py

@@ -1,14 +1,25 @@
 import gradio as gr
 import torch
 import numpy as np
-from ase import Atoms
-from ase.io import read
 import tempfile
 import os
+
+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 import MDLogger
+from ase.io.trajectory import Trajectory
+
+import py3Dmol
+
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-# Global variable for the model
+# -----------------------------
+# Global model
+# -----------------------------
 model_calc = None
 
 def load_orbmol_model():
@@ -19,7 +30,7 @@ def load_orbmol_model():
             print("Loading OrbMol model...")
             orbff = pretrained.orb_v3_conservative_inf_omat(
                 device="cpu",
-                precision="float32"  # más seguro que "float32-high" según la versión
+                precision="float32-high"
             )
             model_calc = ORBCalculator(orbff, device="cpu")
             print("✅ OrbMol model loaded successfully")
@@ -28,12 +39,11 @@ def load_orbmol_model():
             model_calc = None
     return model_calc
 
+# -----------------------------
+# Single-point calculation
+# -----------------------------
 def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
-    """
-    Main function: XYZ → OrbMol → Results
-    """
     try:
-        # Load model
         calc = load_orbmol_model()
         if calc is None:
             return "❌ Error: Could not load OrbMol model", ""
@@ -41,183 +51,127 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         if not xyz_content.strip():
             return "❌ Error: Please enter XYZ coordinates", ""
 
-        # Create temporary file with XYZ
         with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
             f.write(xyz_content)
             xyz_file = f.name
 
-        # Read molecular structure
         atoms = read(xyz_file)
-
-        # Configure charge and spin (IMPORTANT for OrbMol!)
-        atoms.info = {
-            "charge": int(charge),
-            "spin": int(spin_multiplicity)
-        }
-
-        # Assign OrbMol calculator
+        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        # Make the prediction!
-        energy = atoms.get_potential_energy()  # In eV
-        forces = atoms.get_forces()  # In eV/Å
-
-        # Format results nicely
-        result = f"""
-🔋 **Total Energy**: {energy:.6f} eV
-
-⚡ **Atomic Forces**:
-"""
+        energy = atoms.get_potential_energy()
+        forces = atoms.get_forces()
 
-        for i, force in enumerate(forces):
-            result += f"Atom {i+1}: [{force[0]:.4f}, {force[1]:.4f}, {force[2]:.4f}] eV/Å\n"
+        result = f"🔋 **Total Energy**: {energy:.6f} eV\n\n⚡ **Atomic Forces**:\n"
+        for i, f in enumerate(forces):
+            result += f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å\n"
 
-        # Additional statistics
         max_force = np.max(np.linalg.norm(forces, axis=1))
         result += f"\n📊 **Max Force**: {max_force:.4f} eV/Å"
 
-        # Clean up temporary file
         os.unlink(xyz_file)
-
         return result, "✅ Calculation completed with OrbMol"
-
     except Exception as e:
         return f"❌ Error during calculation: {str(e)}", "Error"
 
-# Predefined examples
-examples = [
-    ["""2
-Hydrogen molecule
-H 0.0 0.0 0.0
-H 0.0 0.0 0.74""", 0, 1],
-
-    ["""3
-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],
-
-    ["""4
-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]
-]
-
-# Gradio interface - using FAIR Chem UMA style
-with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
-
-    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 (100M+ high-accuracy DFT calculations).
-
-                Predicts **energies** and **forces** with quantum accuracy, optimized for:
-                * 🧬 Biomolecules  
-                * ⚗️ Metal complexes
-                * 🔋 Electrolytes
-                """)
-
-                gr.Markdown("## Simulation inputs")
-
-                with gr.Column(variant="panel"):
-                    gr.Markdown("### Input molecular structure")
-
-                    xyz_input = gr.Textbox(
-                        label="XYZ Coordinates",
-                        placeholder="""3
-Water molecule
-O 0.0000 0.0000 0.0000  
-H 0.7571 0.0000 0.5864
-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"
-            )
+# -----------------------------
+# Helper: convert trajectory → HTML animation
+# -----------------------------
+def traj_to_html(traj_file):
+    traj = Trajectory(traj_file)
+    view = py3Dmol.view(width=400, height=400)
+    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._make_html()
+
+# -----------------------------
+# Molecular dynamics simulation
+# -----------------------------
+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", ""
 
-            status_output = gr.Textbox(
-                label="Status",
-                interactive=False,
-                max_lines=1
-            )
+        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
+
+        # Pre-relaxation
+        opt = LBFGS(atoms)
+        opt.run(fmax=0.05, steps=20)
+
+        # Velocities
+        MaxwellBoltzmannDistribution(atoms, temperature_K=2 * temperature)
+
+        # MD setup
+        dyn = VelocityVerlet(atoms, timestep=timestep * units.fs)
+
+        traj_file = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
+        traj = Trajectory(traj_file.name, "w", atoms)
+        dyn.attach(traj.write, interval=1)
+
+        dyn.run(steps)
+
+        html = traj_to_html(traj_file.name)
+
+        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)}", ""
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol + MD Demo") as demo:
+    gr.Markdown("# OrbMol Demo with Molecular Dynamics")
+
+    with gr.Tab("Single Point Energy"):
+        xyz_input = gr.Textbox(label="XYZ Coordinates", lines=12)
+        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_btn = gr.Button("Run OrbMol Calculation")
+        results_output = gr.Textbox(label="Results", lines=15)
+        status_output = gr.Textbox(label="Status")
+        run_btn.click(
+            predict_molecule,
+            inputs=[xyz_input, charge_input, spin_input],
+            outputs=[results_output, status_output],
+        )
+
+    with gr.Tab("Molecular Dynamics"):
+        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")
+        md_status = gr.Textbox(label="MD Status", lines=2)
+        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],
+        )
 
-    # Examples section
-    gr.Markdown("### 🧪 Try These Examples")
-    gr.Examples(
-        examples=examples,
-        inputs=[
-            xyz_input,
-            gr.Slider(visible=False, minimum=-10, maximum=10),   # charge
-            gr.Slider(visible=False, minimum=1, maximum=11)      # spin
-        ],
-        label="Click any example to load it"
-    )
-
-    # Connect button to function
-    predict_btn.click(
-        predict_molecule,
-        inputs=[xyz_input, charge_input, spin_input],
-        outputs=[results_output, status_output]
-    )
-
-    # Footer info - matching FAIR Chem UMA style
-    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 with quantum-level accuracy
-            * Built on the Orb-v3 architecture and trained on OMol25 dataset (100M+ DFT calculations)  
-            * Optimized for biomolecules, metal complexes, and electrolytes
-            * Supports configurable charge and spin multiplicity
-
-            [Read more about OrbMol](https://orbitalmaterials.com/posts/orbmol-extending-orb-to-molecular-systems)
-            """)
-
-        with gr.Accordion("Model Disclaimers", open=False):
-            gr.Markdown("""
-            * While OrbMol represents significant progress in molecular ML potentials, the model has limitations
-            * Always validate results for your specific use case
-            * Consider the limitations of the ωB97M-V/def2-TZVPD level of theory used in training
-            """)
-
-        with gr.Accordion("Open source packages", open=False):
-            gr.Markdown("""
-            * Model code available at [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)
-            * This demo uses ASE, Gradio, and other open source packages
-            """)
-
-# Load model on startup
 print("🚀 Starting OrbMol model loading...")
 load_orbmol_model()
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)