Spaces:
Running
Running
| import os | |
| os.environ["TORCH_DYNAMO_DISABLE"] = "1" | |
| import tempfile | |
| import numpy as np | |
| import gradio as gr | |
| from ase.io import read, write | |
| from ase.io.trajectory import Trajectory | |
| from gradio_molecule3d import Molecule3D | |
| from simulation_scripts_orbmol import load_orbmol_model, run_md_simulation, run_relaxation_simulation | |
| import hashlib | |
| from pathlib import Path | |
| # ==== Configuración Molecule3D ==== | |
| DEFAULT_MOLECULAR_REPRESENTATIONS = [ | |
| { | |
| "model": 0, | |
| "chain": "", | |
| "resname": "", | |
| "style": "sphere", | |
| "color": "Jmol", | |
| "around": 0, | |
| "byres": False, | |
| "scale": 0.3, | |
| }, | |
| { | |
| "model": 0, | |
| "chain": "", | |
| "resname": "", | |
| "style": "stick", | |
| "color": "Jmol", | |
| "around": 0, | |
| "byres": False, | |
| "scale": 0.2, | |
| }, | |
| ] | |
| DEFAULT_MOLECULAR_SETTINGS = { | |
| "backgroundColor": "white", | |
| "orthographic": False, | |
| "disableFog": False, | |
| } | |
| # ==== UI definition ==== | |
| # ==== SPE Inputs and Outputs ==== | |
| input_sp = gr.File( | |
| label="Upload Structure File", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single", | |
| render=False | |
| ) | |
| task_name_sp = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| render=False | |
| ) | |
| total_charge_sp = gr.Slider(-10, 10, 0, step=1, label="Charge",render=False) | |
| spin_multiplicity_sp = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity",render=False) | |
| run_sp = gr.Button("Run OrbMol Prediction", variant="primary",render=False) | |
| spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False,render=False) | |
| spe_status = gr.Textbox(label="Status", interactive=False,render=False) | |
| spe_viewer = Molecule3D( | |
| label="Input Structure Viewer", | |
| reps=DEFAULT_MOLECULAR_REPRESENTATIONS, | |
| config=DEFAULT_MOLECULAR_SETTINGS, | |
| render=False | |
| ) | |
| #==== MD Inputs and Outputs ==== | |
| input_md = gr.File( | |
| label="Upload Structure File", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single", | |
| render=False | |
| ) | |
| task_name_md = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| render=False | |
| ) | |
| charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge",render=False) | |
| spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity",render=False) | |
| steps_md = gr.Slider(10, 2000, 100, step=10, label="Steps",render=False) | |
| temp_md = gr.Slider(10, 1500, 300, step=10, label="Temperature (K)",render=False) | |
| timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)",render=False) | |
| ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble",render=False) | |
| run_md_btn = gr.Button("Run MD Simulation", variant="primary",render=False) | |
| md_status = gr.Textbox(label="MD Status", interactive=False,render=False) | |
| md_traj = gr.File(label="Trajectory (.traj)", interactive=False,render=False) | |
| md_viewer = Molecule3D( | |
| label="MD Result Viewer", | |
| reps=DEFAULT_MOLECULAR_REPRESENTATIONS, | |
| config=DEFAULT_MOLECULAR_SETTINGS, | |
| render=False | |
| ) | |
| md_log = gr.Textbox(label="Log", interactive=False, lines=15,render=False) | |
| md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20,render=False) | |
| md_explain = gr.Markdown(render=False) | |
| #==== Relax Inputs and Outputs ==== | |
| input_rlx = gr.File( | |
| label="Upload Structure File", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single", | |
| render=False | |
| ) | |
| task_name_rlx = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| render=False | |
| ) | |
| steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps",render=False) | |
| fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)",render=False) | |
| charge_rlx = gr.Slider(-10, 10, 0, step=1, label="Charge",render=False) | |
| spin_rlx = gr.Slider(1, 11, 1, step=1, label="Spin",render=False) | |
| relax_cell = gr.Checkbox(False, label="Relax Unit Cell",render=False) | |
| run_rlx_btn = gr.Button("Run Optimization", variant="primary",render=False) | |
| rlx_status = gr.Textbox(label="Status", interactive=False,render=False) | |
| rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False,render=False) | |
| rlx_viewer = Molecule3D( | |
| label="Optimized Structure Viewer", | |
| reps=DEFAULT_MOLECULAR_REPRESENTATIONS, | |
| config=DEFAULT_MOLECULAR_SETTINGS, | |
| render=False | |
| ) | |
| rlx_log = gr.Textbox(label="Log", interactive=False, lines=15,render=False) | |
| rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20,render=False) | |
| rlx_explain = gr.Markdown(render=False) | |
| # ==== Conversión a PDB para Molecule3D ==== | |
| def convert_to_pdb_for_viewer(file_path): | |
| """Convierte cualquier archivo a PDB para Molecule3D""" | |
| if not file_path or not os.path.exists(file_path): | |
| return None | |
| try: | |
| atoms = read(file_path) | |
| cache_dir = os.path.join(tempfile.gettempdir(), "gradio") | |
| os.makedirs(cache_dir, exist_ok=True) | |
| pdb_path = os.path.join(cache_dir, f"mol_{hashlib.md5(file_path.encode()).hexdigest()[:12]}.pdb") | |
| write(pdb_path, atoms, format="proteindatabank") | |
| return pdb_path | |
| except Exception as e: | |
| print(f"Error converting to PDB: {e}") | |
| return None | |
| # ==== OrbMol SPE ==== | |
| def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1): | |
| """Single Point Energy + fuerzas (OrbMol)""" | |
| try: | |
| calc = load_orbmol_model(task_name) | |
| if not structure_file: | |
| return "Error: Please upload a structure file", "Error", None | |
| file_path = structure_file | |
| if not os.path.exists(file_path): | |
| return f"Error: File not found: {file_path}", "Error", None | |
| if os.path.getsize(file_path) == 0: | |
| return f"Error: Empty file: {file_path}", "Error", None | |
| atoms = read(file_path) | |
| if task_name in ["OMol", "OMol-Direct"]: | |
| atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)} | |
| atoms.calc = calc | |
| energy = atoms.get_potential_energy() | |
| forces = atoms.get_forces() | |
| lines = [ | |
| f"Model: {task_name}", | |
| f"Total Energy: {energy:.6f} eV", | |
| "", | |
| "Atomic Forces:" | |
| ] | |
| for i, fc in enumerate(forces): | |
| lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å") | |
| max_force = float(np.max(np.linalg.norm(forces, axis=1))) | |
| lines += ["", f"Max Force: {max_force:.4f} eV/Å"] | |
| pdb_file = convert_to_pdb_for_viewer(file_path) | |
| return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| return f"Error during calculation: {e}", "Error", None | |
| # ==== Wrappers MD y Relax ==== | |
| def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_fs, ensemble): | |
| try: | |
| if not structure_file: | |
| return ("Error: Please upload a structure file", None, "", "", "", None) | |
| traj_path, log_text, script_text, explanation = run_md_simulation( | |
| structure_file, | |
| int(steps), | |
| 20, | |
| float(timestep_fs), | |
| float(tempK), | |
| "NVT" if ensemble == "NVT" else "NVE", | |
| str(task_name), | |
| int(charge), | |
| int(spin), | |
| ) | |
| status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})" | |
| pdb_file = convert_to_pdb_for_viewer(traj_path) | |
| return (status, traj_path, log_text, script_text, explanation, pdb_file) | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| return (f"Error: {e}", None, "", "", "", None) | |
| def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_cell): | |
| try: | |
| if not structure_file: | |
| return ("Error: Please upload a structure file", None, "", "", "", None) | |
| traj_path, log_text, script_text, explanation = run_relaxation_simulation( | |
| structure_file, | |
| int(steps), | |
| float(fmax), | |
| str(task_name), | |
| int(charge), | |
| int(spin), | |
| bool(relax_cell), | |
| ) | |
| status = f"Relaxation finished (<={int(steps)} steps, fmax={float(fmax)} eV/Å)" | |
| pdb_file = convert_to_pdb_for_viewer(traj_path) | |
| return (status, traj_path, log_text, script_text, explanation, pdb_file) | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| return (f"Error: {e}", None, "", "", "", None) | |
| # ==== UI ==== | |
| with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo: | |
| with gr.Tabs(): | |
| # -------- HOME TAB (NUEVA) -------- | |
| with gr.Tab("Home"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## Learn more about OrbMol") | |
| with gr.Accordion("What is OrbMol?", open=False): | |
| gr.Markdown(""" | |
| OrbMol is a suite of quantum-accurate machine learning models for molecular predictions. Built on the **Orb-v3 architecture**, OrbMol provides fast and accurate calculations of energies, forces, and molecular properties at the level of advanced quantum chemistry methods. | |
| """) | |
| with gr.Accordion("Available Models", open=False): | |
| gr.Markdown(""" | |
| **OMol** and **OMol-Direct** | |
| - **Training dataset**: OMol25 (>100M calculations on small molecules, biomolecules, metal complexes, and electrolytes) | |
| - **Level of theory**: ωB97M-V/def2-TZVPD with non-local dispersion; solvation treated explicitly | |
| - **Inputs**: total charge & spin multiplicity | |
| - **Applications**: biology, organic chemistry, protein folding, small-molecule drugs, organic liquids, homogeneous catalysis | |
| - **Caveats**: trained only on aperiodic systems → periodic/inorganic cases may not work well | |
| - **Difference**: OMol enforces energy–force consistency; OMol-Direct relaxes this for efficiency | |
| **OMat** | |
| - **Training dataset**: OMat24 (>100M inorganic calculations, from Materials Project, Alexandria, and far-from-equilibrium samples) | |
| - **Level of theory**: PBE/PBE+U with Materials Project settings; VASP 54 pseudopotentials; no dispersion | |
| - **Inputs**: No support for spin and charge | |
| - **Applications**: inorganic discovery, photovoltaics, alloys, superconductors, electronic/optical materials | |
| """) | |
| with gr.Accordion("Supported File Formats", open=False): | |
| gr.Markdown(""" | |
| OrbMol supports: `.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf` | |
| """) | |
| with gr.Accordion("Resources & Support", open=False): | |
| gr.Markdown(""" | |
| - [Orb-v3 paper](https://arxiv.org/abs/2504.06231) | |
| - [Orb-Models GitHub repository](https://github.com/orbital-materials/orb-models) | |
| """) | |
| with gr.Column(scale=2): | |
| gr.Image("logo_color_text.png", | |
| show_share_button=False, | |
| show_download_button=False, | |
| show_label=False, | |
| show_fullscreen_button=False) | |
| gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions") | |
| gr.Markdown(""" | |
| Welcome to the OrbMol demo! Use the tabs above to access different functionalities: | |
| 1. **Single Point Energy**: Calculate energies and forces | |
| 2. **Molecular Dynamics**: Run MD simulations | |
| 3. **Relaxation / Optimization**: Optimize structures | |
| Supported formats: `.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf` | |
| """) | |
| # -------- SPE (IDÉNTICA A LA QUE FUNCIONABA) -------- | |
| with gr.Tab("Single Point Energy"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("# OrbMol: Single Point Energy") | |
| gr.Markdown("Run a quantum-accurate single point energy calculation with OrbMol models.") | |
| gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf") | |
| gr.Examples( | |
| examples=[[str(Path(__file__).parent/ "./examples/10mer_b-DNA.pdb"), | |
| "OMol", | |
| 0, | |
| 1, | |
| "", | |
| ], | |
| [str(Path(__file__).parent/ "./examples/PtO2.cif"), | |
| "OMat", | |
| 0, | |
| 1, | |
| "", | |
| ], | |
| ], | |
| example_labels=[ | |
| "A 10-mer DNA (A tiny DNA fragment)", | |
| "An inorganic crystal (PtO2)", | |
| ], | |
| inputs=[ | |
| input_sp, | |
| task_name_sp, | |
| total_charge_sp, | |
| spin_multiplicity_sp | |
| ], | |
| outputs=[ | |
| spe_out, | |
| spe_status, | |
| spe_viewer | |
| ], | |
| fn=predict_molecule, | |
| run_on_click=True, | |
| cache_examples=True, | |
| label="Try an example!", | |
| ) | |
| input_sp.render() | |
| task_name_sp.render() | |
| with gr.Row(): | |
| total_charge_sp.render() | |
| spin_multiplicity_sp.render() | |
| run_sp.render() | |
| with gr.Column(variant="panel", min_width=500): | |
| spe_out.render() | |
| spe_status.render() | |
| spe_viewer.render() | |
| task_name_sp.change( | |
| lambda x: ( | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]), | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]) | |
| ), | |
| [task_name_sp], | |
| [total_charge_sp, spin_multiplicity_sp] | |
| ) | |
| run_sp.click( | |
| predict_molecule, | |
| [input_sp, task_name_sp, total_charge_sp, spin_multiplicity_sp], | |
| [spe_out, spe_status, spe_viewer] | |
| ) | |
| # -------- MD (IDÉNTICA A LA QUE FUNCIONABA) -------- | |
| with gr.Tab("Molecular Dynamics"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("## OrbMol: Molecular Dynamics Simulation") | |
| gr.Markdown("Run a quantum-accurate molecular dynamics simulation with OrbMol models.") | |
| gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf") | |
| gr.Examples( | |
| examples=[[str(Path(__file__).parent/ "./examples/Cyclohexane.sdf"), | |
| "OMol", | |
| 0, | |
| 1, | |
| 1000, | |
| 300, | |
| 0.5, | |
| "NVT", | |
| "", | |
| ], | |
| [str(Path(__file__).parent/ "./examples/Cyclohexane.sdf"), | |
| "OMol-Direct", | |
| 0, | |
| 1, | |
| 1000, | |
| 300, | |
| 0.5, | |
| "NVT", | |
| "", | |
| ], | |
| ], | |
| example_labels=[ | |
| "Cyclohexane with OMol", | |
| "Cyclohexane with OMol-Direct", | |
| ], | |
| inputs=[ | |
| input_md, | |
| task_name_md, | |
| charge_md, | |
| spin_md, | |
| steps_md, | |
| temp_md, | |
| timestep_md, | |
| ensemble_md | |
| ], | |
| outputs=[ | |
| md_status, | |
| md_traj, | |
| md_log, | |
| md_script, | |
| md_explain, | |
| md_viewer | |
| ], | |
| fn=md_wrapper, | |
| run_on_click=True, | |
| cache_examples=True, | |
| label="Try an example!", | |
| ) | |
| input_md.render() | |
| task_name_md.render() | |
| with gr.Row(): | |
| charge_md.render() | |
| spin_md.render() | |
| with gr.Row(): | |
| steps_md.render() | |
| temp_md.render() | |
| with gr.Row(): | |
| timestep_md.render() | |
| ensemble_md.render() | |
| run_md_btn.render() | |
| with gr.Column(variant="panel", min_width=520): | |
| md_status.render() | |
| md_traj.render() | |
| md_viewer.render() | |
| md_log.render() | |
| md_script.render() | |
| md_explain.render() | |
| task_name_md.change( | |
| lambda x: ( | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]), | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]) | |
| ), | |
| [task_name_md], | |
| [charge_md, spin_md] | |
| ) | |
| run_md_btn.click( | |
| md_wrapper, | |
| [input_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md], | |
| [md_status, md_traj, md_log, md_script, md_explain, md_viewer] | |
| ) | |
| # -------- Relax (IDÉNTICA A LA QUE FUNCIONABA) -------- | |
| with gr.Tab("Relaxation / Optimization"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("## OrbMol: Structure Relaxation/Optimization") | |
| gr.Markdown("Run a quantum-accurate structure relaxation with OrbMol models.") | |
| gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf") | |
| gr.Examples( | |
| examples=[[str(Path(__file__).parent/ "./examples/dioctyl_sebacate.xyz"), | |
| "OMol", | |
| 1000, | |
| 0.05, | |
| 0, | |
| 1, | |
| False, | |
| "", | |
| ], | |
| [str(Path(__file__).parent/ "./examples/Fe2O3.cif"), | |
| "OMat", | |
| 1000, | |
| 0.05, | |
| 0, | |
| 1, | |
| False, | |
| "", | |
| ], | |
| ], | |
| example_labels=[ | |
| "A branched Alkane (Dioctyl Sebacate)", | |
| "An inorganic crystal (Fe2O3)", | |
| ], | |
| inputs=[ | |
| input_rlx, | |
| task_name_rlx, | |
| steps_rlx, | |
| fmax_rlx, | |
| charge_rlx, | |
| spin_rlx, | |
| relax_cell | |
| ], | |
| outputs=[ | |
| rlx_status, | |
| rlx_traj, | |
| rlx_log, | |
| rlx_script, | |
| rlx_explain, | |
| rlx_viewer | |
| ], | |
| fn=relax_wrapper, | |
| run_on_click=True, | |
| cache_examples=True, | |
| label="Try an example!", | |
| ) | |
| input_rlx.render() | |
| task_name_rlx.render() | |
| with gr.Row(): | |
| steps_rlx.render() | |
| fmax_rlx.render() | |
| with gr.Row(): | |
| charge_rlx.render() | |
| spin_rlx.render() | |
| relax_cell.render() | |
| run_rlx_btn.render() | |
| with gr.Column(variant="panel", min_width=520): | |
| rlx_status.render() | |
| rlx_traj.render() | |
| rlx_viewer.render() | |
| rlx_log.render() | |
| rlx_script.render() | |
| rlx_explain.render() | |
| task_name_rlx.change( | |
| lambda x: ( | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]), | |
| gr.update(visible=x in ["OMol", "OMol-Direct"]) | |
| ), | |
| [task_name_rlx], | |
| [charge_rlx, spin_rlx] | |
| ) | |
| run_rlx_btn.click( | |
| relax_wrapper, | |
| [input_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell], | |
| [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) |