Update app.py

app.py

@@ -1,52 +1,13 @@
 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 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
 
@@ -57,8 +18,8 @@ def load_orbmol_model():
         try:
             print("Loading OrbMol model...")
             orbff = pretrained.orb_v3_conservative_inf_omat(
-                device="cpu",  
-                precision="float32-high"
+                device="cpu",
+                precision="float32"  # más seguro que "float32-high" según la versión
             )
             model_calc = ORBCalculator(orbff, device="cpu")
             print("✅ OrbMol model loaded successfully")
@@ -67,230 +28,76 @@ def load_orbmol_model():
             model_calc = None
     return model_calc
 
-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"""
+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)
-        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+
+        # Configure charge and spin (IMPORTANT for OrbMol!)
+        atoms.info = {
+            "charge": int(charge),
+            "spin": int(spin_multiplicity)
+        }
+
+        # Assign OrbMol calculator
         atoms.calc = calc
-        
-        energy = atoms.get_potential_energy()
-        forces = atoms.get_forces()
-        
-        result = f"""🔋 **Total Energy**: {energy:.6f} eV
+
+        # 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_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 = [
+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
+
+    ["""4
 Methane
 C 0.0000 0.0000 0.0000
 H 1.0890 0.0000 0.0000
@@ -299,223 +106,113 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1]
 ]
 
-# Main Gradio interface
+# 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")
-                
-                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!"
-                    )
-                
-                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"
-                    )
-                
-                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.Tab("4. Geometry Optimization"):
-                    gr.Markdown("Optimize molecular geometries with OrbMol:")
-                    
-                    xyz_input_opt = gr.Textbox(
+
+                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.0 0.0 0.0
-H 0.76 0.0 0.59
-H -0.76 0.0 0.59""",
-                        lines=8
+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"
                     )
-                    
-                    with gr.Row():
-                        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/Å]")
-                    
+
+                    gr.Markdown("OMol-specific settings for total charge and spin multiplicity")
                     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
+                        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 Simulation Results")
-            
-            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!")
-            
-            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",
-                )
-    
-    # Sidebar
+            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,
+            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
-            * Built on Orb-v3 architecture, trained on OMol25 dataset (100M+ DFT calculations)
-            * Supports charge and spin multiplicity for accurate molecular modeling
+            * 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("""
-            * OrbMol has limitations and may not work perfectly for all systems
+            * 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 training data and methodology
+            * 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: [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)
-            * Uses ASE, Gradio, and other open source packages
-            * Licensed under Apache 2.0
+            * 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
             """)
-    
-    # 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("🚀 Loading OrbMol model...")
+print("🚀 Starting OrbMol model loading...")
 load_orbmol_model()
 
 if __name__ == "__main__":
@@ -523,4 +220,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True
-    )
+    )