Create app.py

app.py based on UMA

app.py

@@ -0,0 +1,219 @@
+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 orb_models.forcefield import pretrained
+from orb_models.forcefield.calculator import ORBCalculator
+
+# Global variable for the 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"
+            )
+            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
+
+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", ""
+        
+        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.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**:
+"""
+        
+        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 = [
+    ["""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"
+            )
+            
+            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]
+    )
+    
+    # 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
+    )