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 import subprocess, sys from pathlib import Path # === BUILD AND INSTALL LOCAL gradio_molecule3d === try: print("🔧 Building and installing local gradio_molecule3d fork...") base_path = Path(__file__).parent local_pkg = base_path / "gradio_molecule3d" # Step 1 — gradio cc install subprocess.call(["gradio", "cc", "install"], cwd=local_pkg) # Step 2 — gradio cc build subprocess.call(["gradio", "cc", "build"], cwd=local_pkg) # Step 3 — pip install generated wheel wheel_path = local_pkg / "dist" / "gradio_molecule3d-0.0.7-py3-none-any.whl" if not wheel_path.exists(): print("Wheel not found, listing dist contents:") subprocess.call(["ls", "-R", str(local_pkg / "dist")]) subprocess.call( [ sys.executable, "-m", "pip", "install", str(wheel_path), ], cwd=base_path.parent, ) print("gradio_molecule3d built and installed successfully!") except Exception as e: print(f"Error building gradio_molecule3d: {e}") # === Import only after it's installed === from gradio_molecule3d import Molecule3D from gradio_molecule3d import Molecule3D from simulation_scripts_orbmol import load_orbmol_model, run_md_simulation, run_relaxation_simulation import hashlib # ==== 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, } # ==== 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 -------- with gr.Tab("Home"): with gr.Row(): # Columna izquierda con acordeones 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. The models combine the transferability of universal potentials with quantum-level accuracy, making them suitable for a wide range of applications in chemistry, materials science, and drug discovery. """) 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. Spin polarization included but magnetic state cannot be selected - **Applications**: inorganic discovery, photovoltaics, alloys, superconductors, electronic/optical materials - **Caveats**: magnetic effects may be incompletely captured """) with gr.Accordion("Supported File Formats", open=False): gr.Markdown(""" OrbMol supports the following molecular structure formats: - `.xyz` - XYZ coordinate files - `.pdb` - Protein Data Bank format - `.cif` - Crystallographic Information File - `.traj` - ASE trajectory format - `.mol` - MDL Molfile - `.sdf` - Structure Data File All formats are automatically converted internally for processing. """) with gr.Accordion("How to Use", open=False): gr.Markdown(""" **Single Point Energy**: Upload a molecular structure and select a model to calculate energies and forces. **Molecular Dynamics**: Run time-dependent simulations to observe molecular behavior at different temperatures and conditions. **Relaxation/Optimization**: Find the minimum-energy configuration of your molecular structure. Each tab provides specific parameters you can adjust to customize your calculations. """) with gr.Accordion("Technical Foundation", open=False): gr.Markdown(""" All models are based on the **Orb-v3 architecture**, the latest generation of Orb universal interatomic potentials. Key features: - Graph neural network architecture - Equivariant message passing - Multi-task learning across different quantum chemistry methods - Billions of training examples across diverse chemical spaces - Sub-kcal/mol accuracy on test sets """) 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) - For issues/questions, please open a GitHub issue or contact the developers **Citation**: If you use OrbMol in your research, please cite the Orb-v3 paper and the relevant dataset papers (OMol25/OMat24). """) # Columna derecha con contenido principal 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! This interactive platform allows you to explore the capabilities of our quantum-accurate machine learning models for molecular simulations. ## Quick Start Use the tabs above to access different functionalities: 1. **Single Point Energy**: Calculate energies and forces for a given molecular structure 2. **Molecular Dynamics**: Run MD simulations using OrbMol-trained potentials 3. **Relaxation / Optimization**: Optimize molecular structures to their minimum-energy configurations Simply upload a molecular structure file in any supported format (`.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf`) and select the appropriate model for your system. ## Model Selection Guide **Choose OMol/OMol-Direct for:** - Organic molecules and biomolecules - Drug-like compounds - Metal-organic complexes - Molecules in solution - Systems where you need to specify charge and spin **Choose OMat for:** - Inorganic crystals and materials - Periodic systems - Bulk materials and alloys - Solid-state compounds Explore the accordions on the left to learn more about each model's capabilities, training data, and limitations. """) gr.Markdown("## Try an Example") gr.Markdown(""" To get started quickly, navigate to any of the calculation tabs above and try one of these examples: - **Single Point Energy**: Upload a small molecule to see energy and force predictions - **Molecular Dynamics**: Run a short simulation at 300K to observe thermal motion - **Relaxation**: Optimize a distorted structure to find its equilibrium geometry """) # -------- SPE -------- with gr.Tab("Single Point Energy"): with gr.Row(): with gr.Column(scale=2): gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions") gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf") xyz_input = gr.File( label="Upload Structure File", file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], file_count="single" ) task_name_spe = gr.Radio( ["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type" ) with gr.Row(): charge_input = gr.Slider(-10, 10, 0, step=1, label="Charge") spin_input = gr.Slider(1, 11, 1, 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) spe_viewer = Molecule3D( label="Input Structure Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS ) task_name_spe.change( lambda x: ( gr.update(visible=x in ["OMol", "OMol-Direct"]), gr.update(visible=x in ["OMol", "OMol-Direct"]) ), [task_name_spe], [charge_input, spin_input] ) run_spe.click( predict_molecule, [xyz_input, task_name_spe, charge_input, spin_input], [spe_out, spe_status, spe_viewer] ) # -------- MD -------- with gr.Tab("Molecular Dynamics"): with gr.Row(): with gr.Column(scale=2): gr.Markdown("## Molecular Dynamics Simulation") xyz_md = gr.File( label="Upload Structure File", file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], file_count="single" ) task_name_md = gr.Radio( ["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type" ) with gr.Row(): charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge") spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity") with gr.Row(): steps_md = gr.Slider(10, 2000, 100, step=10, label="Steps") temp_md = gr.Slider(10, 1500, 300, step=10, label="Temperature (K)") with gr.Row(): timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)") ensemble_md = gr.Radio(["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) md_traj = gr.File(label="Trajectory (.traj)", interactive=False) md_viewer = Molecule3D( label="MD Result Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS ) md_log = gr.Textbox(label="Log", interactive=False, lines=15) md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20) md_explain = gr.Markdown() 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, [xyz_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 -------- with gr.Tab("Relaxation / Optimization"): with gr.Row(): with gr.Column(scale=2): gr.Markdown("## Structure Relaxation/Optimization") xyz_rlx = gr.File( label="Upload Structure File", file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], file_count="single" ) task_name_rlx = gr.Radio( ["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type" ) with gr.Row(): steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps") fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)") with gr.Row(): charge_rlx = gr.Slider(-10, 10, 0, step=1, label="Charge") spin_rlx = gr.Slider(1, 11, 1, step=1, label="Spin") relax_cell = gr.Checkbox(False, label="Relax Unit Cell") 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) rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False) rlx_viewer = Molecule3D( label="Optimized Structure Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS ) rlx_log = gr.Textbox(label="Log", interactive=False, lines=15) rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20) rlx_explain = gr.Markdown() 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, [xyz_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)