Spaces:
Sleeping
Sleeping
| import os | |
| import tempfile | |
| import numpy as np | |
| import gradio as gr | |
| from ase.io import read | |
| 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 import units | |
| # Visualizador 3D “pro” si está disponible | |
| try: | |
| from gradio_molecule3d import Molecule3D | |
| HAVE_MOL3D = True | |
| except Exception: | |
| HAVE_MOL3D = False | |
| # --- Helpers de visualización: fallback 3Dmol.js si no hay Molecule3D --- | |
| def traj_to_html(traj_path, width=520, height=520, interval_ms=200): | |
| traj = Trajectory(traj_path) | |
| xyz_frames = [] | |
| for atoms in traj: | |
| symbols = atoms.get_chemical_symbols() | |
| coords = atoms.get_positions() | |
| parts = [str(len(symbols)), "frame"] | |
| for s, (x, y, z) in zip(symbols, coords): | |
| parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}") | |
| xyz_frames.append("\n".join(parts)) | |
| html = f""" | |
| <div id="viewer_md" style="width:{width}px; height:{height}px;"></div> | |
| <script src="https://3dmol.org/build/3Dmol-min.js"></script> | |
| <script> | |
| (function() {{ | |
| var viewer = $3Dmol.createViewer("viewer_md", {{backgroundColor: 'white'}}); | |
| var frames = {xyz_frames!r}; | |
| var i = 0; | |
| function show(i) {{ | |
| viewer.clear(); | |
| viewer.addModel(frames[i], "xyz"); | |
| viewer.setStyle({{}}, {{stick: {{}}}}); | |
| viewer.zoomTo(); | |
| viewer.render(); | |
| }} | |
| if(frames.length>0) show(0); | |
| if(frames.length>1) setInterval(function(){{ | |
| i=(i+1)%frames.length; show(i); | |
| }}, {int(interval_ms)}); | |
| }})(); | |
| </script> | |
| """ | |
| return html | |
| # --- OrbMol directo para SPE --- | |
| from orb_models.forcefield import pretrained | |
| from orb_models.forcefield.calculator import ORBCalculator | |
| _model_calc = None | |
| def _load_orbmol_calc(): | |
| global _model_calc | |
| if _model_calc is None: | |
| orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu", precision="float32-high") | |
| _model_calc = ORBCalculator(orbff, device="cpu") | |
| return _model_calc | |
| def predict_molecule(xyz_content, charge=0, spin_multiplicity=1): | |
| try: | |
| calc = _load_orbmol_calc() | |
| if not xyz_content.strip(): | |
| return "Error: Please enter XYZ coordinates", "Error" | |
| 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() # eV | |
| forces = atoms.get_forces() # eV/Å | |
| lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"] | |
| for i, f in enumerate(forces): | |
| lines.append(f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å") | |
| max_force = float(np.max(np.linalg.norm(forces, axis=1))) | |
| lines += ["", f"Max Force: {max_force:.4f} eV/Å"] | |
| try: | |
| os.unlink(xyz_file) | |
| except Exception: | |
| pass | |
| return "\n".join(lines), "Calculation completed with OrbMol" | |
| except Exception as e: | |
| return f"Error during calculation: {e}", "Error" | |
| # --- Importa las rutinas FAIRChem-like para MD/Relax (ya soportan string XYZ o ruta) --- | |
| from simulation_scripts_orbmol import ( | |
| run_md_simulation, | |
| run_relaxation_simulation, | |
| last_frame_xyz_from_traj, | |
| ) | |
| # Wrappers para conectar outputs de Gradio correctamente (string XYZ / HTML, file, logs...) | |
| def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble): | |
| try: | |
| traj_path, log_text, script_text, explanation = run_md_simulation( | |
| xyz_content, # acepta string XYZ | |
| int(steps), | |
| 20, # pre-relax steps como en la UI de UMA | |
| float(timestep_fs), | |
| float(tempK), | |
| "NVT" if ensemble == "NVT" else "NVE", | |
| int(charge), | |
| int(spin), | |
| ) | |
| status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})" | |
| # Viewer | |
| if HAVE_MOL3D: | |
| xyz_final = last_frame_xyz_from_traj(traj_path) # value para Molecule3D | |
| viewer_value = xyz_final | |
| html_value = None | |
| else: | |
| viewer_value = None | |
| html_value = traj_to_html(traj_path) | |
| return ( | |
| status, # MD Status | |
| viewer_value, # Molecule3D value (o None) | |
| html_value, # HTML fallback (o None) | |
| traj_path, # File download | |
| log_text, # Log | |
| script_text, # Script | |
| explanation, # Explanation | |
| ) | |
| except Exception as e: | |
| return (f"Error: {e}", None, None, None, "", "", "") | |
| def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell): | |
| try: | |
| traj_path, log_text, script_text, explanation = run_relaxation_simulation( | |
| xyz_content, | |
| int(steps), | |
| float(fmax), | |
| int(charge), | |
| int(spin), | |
| bool(relax_cell), | |
| ) | |
| status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)" | |
| if HAVE_MOL3D: | |
| viewer_value = last_frame_xyz_from_traj(traj_path) | |
| html_value = None | |
| else: | |
| viewer_value = None | |
| html_value = traj_to_html(traj_path) | |
| return ( | |
| status, | |
| viewer_value, | |
| html_value, | |
| traj_path, | |
| log_text, | |
| script_text, | |
| explanation, | |
| ) | |
| except Exception as e: | |
| return (f"Error: {e}", None, None, None, "", "", "") | |
| # ------------------------ | |
| # Ejemplos rápidos | |
| # ------------------------ | |
| 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], | |
| ] | |
| # ------------------------ | |
| # UI Gradio | |
| # ------------------------ | |
| with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo: | |
| with gr.Tabs(): | |
| # ========== SPE ========== | |
| with gr.Tab("Single Point Energy"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("# OrbMol Demo — Quantum-Accurate Molecular Predictions") | |
| gr.Markdown( | |
| "Predict **energies** and **forces** with OrbMol (OMol25). " | |
| "Supports **charge** and **spin multiplicity**." | |
| ) | |
| xyz_input = gr.Textbox( | |
| label="XYZ Coordinates", | |
| placeholder="Paste XYZ here...", | |
| lines=12, | |
| ) | |
| with gr.Row(): | |
| 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_spe = gr.Button("Run OrbMol Prediction", variant="primary") | |
| with gr.Column(variant="panel", min_width=500): | |
| spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False) | |
| spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1) | |
| gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input], label="Examples") | |
| run_spe.click( | |
| predict_molecule, | |
| inputs=[xyz_input, charge_input, spin_input], | |
| outputs=[spe_out, spe_status], | |
| ) | |
| with gr.Sidebar(open=True): | |
| gr.Markdown("## Learn more about OrbMol") | |
| with gr.Accordion("What is OrbMol?", open=False): | |
| gr.Markdown( | |
| "* Neural network potential for molecules\n" | |
| "* Built on Orb-v3, trained on OMol25 (ωB97M-V/def2-TZVPD)\n" | |
| "* Supports charge and spin multiplicity" | |
| ) | |
| with gr.Accordion("Benchmarks", open=False): | |
| gr.Markdown( | |
| "* Near-DFT accuracy on **GMTKN55** and **Wiggle150**\n" | |
| "* Accurate **protein–ligand** interaction energies (PLA15)\n" | |
| "* Stable long MD on biomolecules grandes" | |
| ) | |
| with gr.Accordion("Disclaimers", open=False): | |
| gr.Markdown( | |
| "* Validate for your use case\n" | |
| "* Consider training **level of theory** and intended domain" | |
| ) | |
| # ========== MD ========== | |
| with gr.Tab("Molecular Dynamics"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| xyz_md = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...") | |
| with gr.Row(): | |
| charge_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge") | |
| spin_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity") | |
| with gr.Row(): | |
| steps_md = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps") | |
| temp_md = gr.Slider(value=300, minimum=10, maximum=1500, step=10, label="Temperature (K)") | |
| with gr.Row(): | |
| timestep_md = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)") | |
| ensemble_md = gr.Radio(choices=["NVE", "NVT"], value="NVE", label="Ensemble") | |
| run_md_btn = gr.Button("Run MD Simulation", variant="primary") | |
| with gr.Column(variant="panel", min_width=520): | |
| md_status = gr.Textbox(label="MD Status", interactive=False) | |
| if HAVE_MOL3D: | |
| md_viewer = Molecule3D(label="Trajectory Viewer") | |
| md_html = gr.HTML(visible=False) # placeholder para consistencia de outputs | |
| else: | |
| md_viewer = gr.Textbox(visible=False) # placeholder | |
| md_html = gr.HTML() | |
| md_traj = gr.File(label="Trajectory (.traj)", interactive=False) | |
| md_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25) | |
| md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30) | |
| md_explain = gr.Markdown() | |
| run_md_btn.click( | |
| md_wrapper, | |
| inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md], | |
| outputs=[md_status, md_viewer, md_html, md_traj, md_log, md_script, md_explain], | |
| ) | |
| # ========== Relaxation ========== | |
| with gr.Tab("Relaxation / Optimization"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| xyz_rlx = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...") | |
| steps_rlx = gr.Slider(value=300, minimum=1, maximum=1000, step=1, label="Max Steps") | |
| fmax_rlx = gr.Slider(value=0.05, minimum=0.001, maximum=0.5, step=0.001, label="Fmax (eV/Å)") | |
| with gr.Row(): | |
| charge_rlx = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge") | |
| spin_rlx = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin") | |
| relax_cell = gr.Checkbox(label="Relax Unit Cell", value=False) | |
| run_rlx_btn = gr.Button("Run Optimization", variant="primary") | |
| with gr.Column(variant="panel", min_width=520): | |
| rlx_status = gr.Textbox(label="Status", interactive=False) | |
| if HAVE_MOL3D: | |
| rlx_viewer = Molecule3D(label="Final Structure") | |
| rlx_html = gr.HTML(visible=False) | |
| else: | |
| rlx_viewer = gr.Textbox(visible=False) | |
| rlx_html = gr.HTML() | |
| rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False) | |
| rlx_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25) | |
| rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30) | |
| rlx_explain = gr.Markdown() | |
| run_rlx_btn.click( | |
| relax_wrapper, | |
| inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell], | |
| outputs=[rlx_status, rlx_viewer, rlx_html, rlx_traj, rlx_log, rlx_script, rlx_explain], | |
| ) | |
| print("Starting OrbMol model loading…") | |
| _load_orbmol_calc() | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) | |