Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| 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.io.trajectory import Trajectory | |
| import py3Dmol | |
| from orb_models.forcefield import pretrained | |
| from orb_models.forcefield.calculator import ORBCalculator | |
| # ----------------------------- | |
| # Global OrbMol model | |
| # ----------------------------- | |
| model_calc = None | |
| def load_orbmol_model(): | |
| """Load OrbMol model once""" | |
| global model_calc | |
| if model_calc is None: | |
| try: | |
| 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 OrbMol model: {e}") | |
| model_calc = None | |
| return model_calc | |
| # ----------------------------- | |
| # Single-point calculation | |
| # ----------------------------- | |
| def predict_molecule(xyz_content, charge=0, spin_multiplicity=1): | |
| try: | |
| 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", "" | |
| 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 | |
| energy = atoms.get_potential_energy() | |
| forces = atoms.get_forces() | |
| result = f"Total Energy: {energy:.6f} eV\n\nAtomic Forces:\n" | |
| 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/Å" | |
| os.unlink(xyz_file) | |
| return result, "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) | |
| 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): | |
| try: | |
| 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", "" | |
| 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) | |
| # Initialize velocities | |
| MaxwellBoltzmannDistribution(atoms, temperature_K=2 * 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) | |
| 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)}", "" | |
| # ----------------------------- | |
| # 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], | |
| ["""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] | |
| ] | |
| # ----------------------------- | |
| # Gradio UI | |
| # ----------------------------- | |
| with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo: | |
| 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], | |
| ) | |
| 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) | |