Update app.py

app.py

@@ -1,13 +1,52 @@
 import gradio as gr
 import torch
 import numpy as np
-from ase import Atoms
-from ase.io import read, write
 import tempfile
 import os
+from ase import Atoms
+from ase.io import read, write
+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.md import MDLogger
+from ase import units
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
+# Install and import gradio_molecule3d
+try:
+    from gradio_molecule3d import Molecule3D
+    HAS_3D = True
+except ImportError:
+    HAS_3D = False
+
+# Default molecular representations for 3D visualization
+DEFAULT_MOLECULAR_REPRESENTATIONS = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "sphere",
+        "color": "Jmol",
+        "scale": 0.3,
+    },
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "stick",
+        "color": "Jmol",
+        "scale": 0.2,
+    },
+]
+
+DEFAULT_MOLECULAR_SETTINGS = {
+    "backgroundColor": "white",
+    "orthographic": False,
+    "disableFog": False,
+}
+
 # Global variable for the model
 model_calc = None
 
@@ -28,12 +67,177 @@ def load_orbmol_model():
             model_calc = None
     return model_calc
 
-def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
-    """
-    Main function: XYZ → OrbMol → Results
-    """
+def run_md_simulation(input_file, md_steps, prerelax_steps, md_timestep, temperature_k, md_ensemble, charge, spin_multiplicity, explanation_buffer=""):
+    """Run molecular dynamics simulation similar to FAIR Chem UMA"""
+    try:
+        calc = load_orbmol_model()
+        if calc is None:
+            return None, "❌ Error: Could not load OrbMol model", "", explanation_buffer
+        
+        # Parse input file content
+        if isinstance(input_file, str):
+            # Handle XYZ string input
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
+                f.write(input_file)
+                xyz_file = f.name
+            atoms = read(xyz_file)
+            os.unlink(xyz_file)
+        else:
+            # Handle uploaded file
+            atoms = read(input_file)
+        
+        # Set charge and spin
+        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+        atoms.calc = calc
+        
+        # Pre-relaxation
+        if prerelax_steps > 0:
+            opt = LBFGS(atoms)
+            opt.run(fmax=0.05, steps=prerelax_steps)
+        
+        # Create trajectory file
+        traj_path = tempfile.NamedTemporaryFile(suffix='.traj', delete=False).name
+        
+        # Initialize velocity distribution
+        MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_k)
+        
+        # Set up MD
+        if md_ensemble == "NVE":
+            dyn = VelocityVerlet(atoms, timestep=md_timestep * units.fs)
+        else:  # NVT
+            from ase.md.langevin import Langevin
+            dyn = Langevin(atoms, md_timestep * units.fs, temperature_K=temperature_k, friction=0.001/units.fs)
+        
+        # Attach trajectory
+        traj = Trajectory(traj_path, "w", atoms)
+        dyn.attach(traj.write, interval=1)
+        
+        # Run simulation
+        dyn.run(md_steps)
+        traj.close()
+        
+        # Generate log
+        log_content = f"""OrbMol Molecular Dynamics Simulation
+=====================================
+System: {len(atoms)} atoms
+MD Steps: {md_steps}
+Temperature: {temperature_k} K
+Ensemble: {md_ensemble}
+Timestep: {md_timestep} fs
+Charge: {charge}
+Spin Multiplicity: {spin_multiplicity}
+
+Final Energy: {atoms.get_potential_energy():.6f} eV
+Simulation completed successfully!
+"""
+        
+        # Generate reproduction script
+        script_content = f"""# OrbMol MD Simulation Script
+import ase.io
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase.md.verlet import VelocityVerlet
+from ase.optimize import LBFGS
+from ase.io.trajectory import Trajectory
+from ase import units
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
+
+# Load structure
+atoms = ase.io.read('input_structure.xyz')  # Your input file
+atoms.info = {{"charge": {charge}, "spin": {spin_multiplicity}}}
+
+# Set up OrbMol calculator
+orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu")
+atoms.calc = ORBCalculator(orbff, device="cpu")
+
+# Pre-relaxation
+opt = LBFGS(atoms)
+opt.run(fmax=0.05, steps={prerelax_steps})
+
+# Initialize velocities
+MaxwellBoltzmannDistribution(atoms, temperature_K={temperature_k})
+
+# Set up MD
+dyn = VelocityVerlet(atoms, timestep={md_timestep} * units.fs)
+traj = Trajectory("md_output.traj", "w", atoms)
+dyn.attach(traj.write, interval=1)
+
+# Run simulation
+dyn.run({md_steps})
+"""
+        
+        explanation = explanation_buffer if explanation_buffer else f"Molecular dynamics simulation completed! This shows {len(atoms)} atoms moving over {md_steps} steps at {temperature_k} K. The atoms are vibrating due to thermal motion, and you can see the molecular structure evolving over time."
+        
+        return traj_path, log_content, script_content, explanation
+        
+    except Exception as e:
+        return None, f"❌ Error during MD simulation: {str(e)}", "", "Error occurred"
+
+def run_optimization(input_file, optimization_steps, fmax, charge, spin_multiplicity):
+    """Run geometry optimization"""
+    try:
+        calc = load_orbmol_model()
+        if calc is None:
+            return None, "❌ Error: Could not load OrbMol model", ""
+        
+        # Parse input
+        if isinstance(input_file, str):
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
+                f.write(input_file)
+                xyz_file = f.name
+            atoms = read(xyz_file)
+            os.unlink(xyz_file)
+        else:
+            atoms = read(input_file)
+        
+        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+        atoms.calc = calc
+        
+        # Create trajectory file
+        traj_path = tempfile.NamedTemporaryFile(suffix='.traj', delete=False).name
+        
+        # Optimize
+        opt = LBFGS(atoms, trajectory=traj_path)
+        opt.run(fmax=fmax, steps=optimization_steps)
+        
+        # Generate log
+        log_content = f"""OrbMol Geometry Optimization
+===========================
+System: {len(atoms)} atoms
+Max Steps: {optimization_steps}
+Force Convergence: {fmax} eV/Å
+Charge: {charge}
+Spin Multiplicity: {spin_multiplicity}
+
+Final Energy: {atoms.get_potential_energy():.6f} eV
+Max Force: {np.max(np.linalg.norm(atoms.get_forces(), axis=1)):.6f} eV/Å
+Optimization completed!
+"""
+        
+        script_content = f"""# OrbMol Geometry Optimization Script
+import ase.io
+from ase.optimize import LBFGS
+from orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
+
+atoms = ase.io.read('input_structure.xyz')
+atoms.info = {{"charge": {charge}, "spin": {spin_multiplicity}}}
+
+orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu")
+atoms.calc = ORBCalculator(orbff, device="cpu")
+
+opt = LBFGS(atoms, trajectory="optimization.traj")
+opt.run(fmax={fmax}, steps={optimization_steps})
+"""
+        
+        return traj_path, log_content, script_content
+        
+    except Exception as e:
+        return None, f"❌ Error during optimization: {str(e)}", ""
+
+def predict_single_point(xyz_content, charge=0, spin_multiplicity=1):
+    """Single point energy and forces calculation"""
     try:
-        # Load model
         calc = load_orbmol_model()
         if calc is None:
             return "❌ Error: Could not load OrbMol model", ""
@@ -41,51 +245,40 @@ 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/Å
+        energy = atoms.get_potential_energy()
+        forces = atoms.get_forces()
         
-        # Format results nicely
-        result = f"""
-🔋 **Total Energy**: {energy:.6f} eV
+        result = f"""🔋 **Total Energy**: {energy:.6f} eV
 
 ⚡ **Atomic Forces**:
 """
-        
         for i, force in enumerate(forces):
             result += f"Atom {i+1}: [{force[0]:.4f}, {force[1]:.4f}, {force[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 = [
+examples_md = [
+    ["2\nHydrogen molecule\nH 0.0 0.0 0.0\nH 0.0 0.0 0.74", 100, 20, 1.0, 300.0, "NVE", 0, 1, "Simple H2 molecule dynamics showing thermal vibrations"],
+    ["3\nWater molecule\nO 0.0 0.0 0.0\nH 0.7571 0.0 0.5864\nH -0.7571 0.0 0.5864", 200, 20, 1.0, 300.0, "NVE", 0, 1, "Water molecule thermal motion and vibrations"],
+]
+
+examples_sp = [
     ["""2
 Hydrogen molecule
 H 0.0 0.0 0.0
@@ -97,7 +290,7 @@ O 0.0000 0.0000 0.0000
 H 0.7571 0.0000 0.5864
 H -0.7571 0.0000 0.5864""", 0, 1],
     
-    ["""4
+    ["""5
 Methane
 C 0.0000 0.0000 0.0000
 H 1.0890 0.0000 0.0000
@@ -106,7 +299,7 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1]
 ]
 
-# Gradio interface - using FAIR Chem UMA style
+# Main Gradio interface
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     
     with gr.Row():
@@ -114,101 +307,215 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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).
+                with gr.Tab("1. OrbMol Intro"):
+                    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
+                    
+                    Try the examples below to see OrbMol in action!
+                    """)
+                    
+                    # Quick examples for MD
+                    gr.Examples(
+                        examples=examples_md,
+                        inputs=[
+                            gr.Textbox(visible=False, value=""),  # Will be set by interface
+                            gr.Slider(visible=False),
+                            gr.Slider(visible=False),
+                            gr.Slider(visible=False),
+                            gr.Slider(visible=False),
+                            gr.Radio(visible=False),
+                            gr.Slider(visible=False),
+                            gr.Slider(visible=False),
+                            gr.Textbox(visible=False)
+                        ],
+                        outputs=[
+                            gr.File(visible=False),
+                            gr.Code(visible=False),
+                            gr.Code(visible=False),
+                            gr.Markdown(visible=False)
+                        ],
+                        fn=run_md_simulation,
+                        cache_examples=True,
+                        label="Try molecular dynamics examples!"
+                    )
                 
-                Predicts **energies** and **forces** with quantum accuracy, optimized for:
-                * 🧬 Biomolecules  
-                * ⚗️ Metal complexes
-                * 🔋 Electrolytes
-                """)
+                with gr.Tab("2. Single Point Calculations"):
+                    gr.Markdown("Calculate energy and forces for a single molecular geometry:")
+                    
+                    # Single point interface
+                    with gr.Row():
+                        with gr.Column():
+                            xyz_input_sp = gr.Textbox(
+                                label="XYZ Coordinates",
+                                placeholder="""3
+Water molecule
+O 0.0 0.0 0.0
+H 0.76 0.0 0.59
+H -0.76 0.0 0.59""",
+                                lines=8
+                            )
+                            
+                            with gr.Row():
+                                charge_sp = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
+                                spin_sp = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
+                            
+                            predict_btn_sp = gr.Button("Calculate Energy & Forces", variant="primary")
+                        
+                        with gr.Column():
+                            results_sp = gr.Textbox(label="Results", lines=12, interactive=False)
+                            status_sp = gr.Textbox(label="Status", max_lines=1, interactive=False)
+                    
+                    gr.Examples(
+                        examples=examples_sp,
+                        inputs=[xyz_input_sp, charge_sp, spin_sp],
+                        label="Single point examples"
+                    )
                 
-                gr.Markdown("## Simulation inputs")
+                with gr.Tab("3. Molecular Dynamics"):
+                    gr.Markdown("Run molecular dynamics simulations with OrbMol:")
+                    
+                    xyz_input_md = gr.Textbox(
+                        label="XYZ Coordinates",
+                        placeholder="""3
+Water molecule
+O 0.0 0.0 0.0
+H 0.76 0.0 0.59
+H -0.76 0.0 0.59""",
+                        lines=8
+                    )
+                    
+                    with gr.Row():
+                        md_steps = gr.Slider(minimum=10, maximum=500, value=100, label="MD Steps")
+                        prerelax_steps = gr.Slider(minimum=0, maximum=100, value=20, label="Pre-Relaxation Steps")
+                    
+                    with gr.Row():
+                        temperature_k = gr.Slider(minimum=0, maximum=1500, value=300, label="Temperature [K]")
+                        md_timestep = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, label="Timestep [fs]")
+                    
+                    md_ensemble = gr.Radio(choices=["NVE", "NVT"], value="NVE", label="Ensemble")
+                    
+                    with gr.Row():
+                        charge_md = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
+                        spin_md = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
+                    
+                    md_button = gr.Button("Run MD Simulation", variant="primary")
                 
-                with gr.Column(variant="panel"):
-                    gr.Markdown("### Input molecular structure")
+                with gr.Tab("4. Geometry Optimization"):
+                    gr.Markdown("Optimize molecular geometries with OrbMol:")
                     
-                    xyz_input = gr.Textbox(
+                    xyz_input_opt = 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"
+O 0.0 0.0 0.0
+H 0.76 0.0 0.59
+H -0.76 0.0 0.59""",
+                        lines=8
                     )
                     
-                    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"
-                        )
+                        opt_steps = gr.Slider(minimum=1, maximum=500, value=300, label="Max Steps")
+                        fmax = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, label="Force Tolerance [eV/Å]")
                     
-                    predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
+                    with gr.Row():
+                        charge_opt = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
+                        spin_opt = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
                     
+                    opt_button = gr.Button("Run Optimization", variant="primary")
+        
+        # Results panel
         with gr.Column(variant="panel", elem_id="results", min_width=500):
-            gr.Markdown("## OrbMol Prediction Results")
+            gr.Markdown("## OrbMol Simulation Results")
             
-            results_output = gr.Textbox(
-                label="Energy & Forces",
-                lines=15,
-                interactive=False,
-                info="OrbMol energy and force predictions"
-            )
+            with gr.Tab("Visualization"):
+                if HAS_3D:
+                    output_structure = Molecule3D(
+                        label="Simulation Visualization",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS,
+                        height=500,
+                        interactive=False,
+                    )
+                else:
+                    gr.Markdown("3D visualization not available. Install gradio-molecule3d for 3D viewing.")
+                
+                output_traj = gr.File(label="Trajectory File", interactive=False)
+                
+                explanation = gr.Markdown("Run a simulation to see results here!")
             
-            status_output = gr.Textbox(
-                label="Status",
-                interactive=False,
-                max_lines=1
-            )
-    
-    # Examples section
-    gr.Markdown("### 🧪 Try These Examples")
-    gr.Examples(
-        examples=examples,
-        inputs=[xyz_input, charge_input, spin_input],
-        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]
-    )
+            with gr.Tab("Log"):
+                output_text = gr.Code(lines=20, max_lines=30, label="Simulation Log", interactive=False)
+            
+            with gr.Tab("Script"):
+                reproduction_script = gr.Code(
+                    interactive=False,
+                    max_lines=30,
+                    language="python",
+                    label="Reproduction Script",
+                )
     
-    # Footer info - matching FAIR Chem UMA style
+    # Sidebar
     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)  
+            * OrbMol is a neural network potential for molecular property prediction
+            * Built on Orb-v3 architecture, trained on OMol25 dataset (100M+ DFT calculations)
+            * Supports charge and spin multiplicity for accurate molecular modeling
             * 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
+            * OrbMol has limitations and may not work perfectly for all systems
             * Always validate results for your specific use case
-            * Consider the limitations of the ωB97M-V/def2-TZVPD level of theory used in training
+            * Consider the limitations of the training data and methodology
             """)
-            
+        
         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
+            * Model: [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)
+            * Uses ASE, Gradio, and other open source packages
+            * Licensed under Apache 2.0
             """)
+    
+    # Connect buttons to functions
+    predict_btn_sp.click(
+        predict_single_point,
+        inputs=[xyz_input_sp, charge_sp, spin_sp],
+        outputs=[results_sp, status_sp]
+    )
+    
+    md_button.click(
+        run_md_simulation,
+        inputs=[xyz_input_md, md_steps, prerelax_steps, md_timestep, temperature_k, md_ensemble, charge_md, spin_md],
+        outputs=[output_traj, output_text, reproduction_script, explanation]
+    )
+    
+    opt_button.click(
+        run_optimization,
+        inputs=[xyz_input_opt, opt_steps, fmax, charge_opt, spin_opt],
+        outputs=[output_traj, output_text, reproduction_script]
+    )
+    
+    # Update 3D visualization when trajectory changes
+    if HAS_3D:
+        output_traj.change(
+            lambda x: x,
+            inputs=[output_traj],
+            outputs=[output_structure]
+        )
 
 # Load model on startup
-print("🚀 Starting OrbMol model loading...")
+print("🚀 Loading OrbMol model...")
 load_orbmol_model()
 
 if __name__ == "__main__":