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 | |
| import hashlib | |
| import shutil | |
| # ==== Configuración UNIVERSAL para Molecule3D ==== | |
| UNIVERSAL_MOLECULE3D_REPS = [ | |
| { | |
| "model": 0, | |
| "chain": "", | |
| "resname": "", | |
| "style": "sphere", | |
| "color": "element", # Colores automáticos por elemento químico | |
| "residue_range": "", | |
| "around": 0, | |
| "byres": False, | |
| "visible": True, | |
| "opacity": 0.8 | |
| }, | |
| { | |
| "model": 0, | |
| "chain": "", | |
| "resname": "", | |
| "style": "stick", | |
| "color": "element", | |
| "residue_range": "", | |
| "around": 0, | |
| "byres": False, | |
| "visible": True, | |
| "opacity": 1.0 | |
| } | |
| ] | |
| # ==== Función mejorada para crear PDB universal ==== | |
| def create_universal_pdb(atoms, cache_key): | |
| """ | |
| Crea un PDB optimizado que funciona bien con gradio_molecule3d para cualquier molécula | |
| """ | |
| cache_dir = os.path.join(tempfile.gettempdir(), "gradio") | |
| os.makedirs(cache_dir, exist_ok=True) | |
| pdb_path = os.path.join(cache_dir, f"mol_{cache_key}.pdb") | |
| if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0: | |
| return pdb_path | |
| try: | |
| positions = atoms.get_positions() | |
| symbols = atoms.get_chemical_symbols() | |
| with open(pdb_path, "w") as f: | |
| # Header estándar | |
| f.write("HEADER MOLECULAR STRUCTURE 01-JAN-25 MOL \n") | |
| f.write("TITLE CALCULATION RESULT\n") | |
| f.write("MODEL 1\n") | |
| # Escribir átomos con formato PDB correcto | |
| for i, (symbol, pos) in enumerate(zip(symbols, positions)): | |
| # Asegurar formato correcto del elemento | |
| element = symbol.strip().upper() | |
| # Nombre del átomo (4 caracteres, alineado a la izquierda) | |
| atom_name = f"{element:<4s}" | |
| # Línea ATOM con formato PDB estricto | |
| f.write( | |
| f"ATOM {i+1:5d} {atom_name} UNL A 1 " | |
| f"{pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}" | |
| f" 1.00 30.00 {element:>2s} \n" | |
| ) | |
| f.write("ENDMDL\n") | |
| f.write("END\n") | |
| # Verificar creación | |
| if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0: | |
| print(f"✅ Universal PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)") | |
| # Debug: mostrar contenido | |
| with open(pdb_path, 'r') as f: | |
| content = f.read() | |
| print(f"PDB content preview:\n{content[:300]}...") | |
| return pdb_path | |
| else: | |
| print(f"❌ Universal PDB creation failed") | |
| return None | |
| except Exception as e: | |
| print(f"❌ Error in create_universal_pdb: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return None | |
| # ==== Función actualizada para preparar cualquier molécula ==== | |
| def prepare_universal_molecule_viewer(traj_path_or_atoms): | |
| """ | |
| Función universal que funciona con cualquier molécula | |
| """ | |
| try: | |
| # Determinar si es trayectoria o átomos | |
| if isinstance(traj_path_or_atoms, str): | |
| if not os.path.exists(traj_path_or_atoms): | |
| print("Trajectory file doesn't exist") | |
| return None | |
| traj = Trajectory(traj_path_or_atoms) | |
| if len(traj) == 0: | |
| print("Empty trajectory") | |
| return None | |
| atoms = traj[-1] | |
| cache_key = hashlib.md5(str(traj_path_or_atoms).encode()).hexdigest()[:12] | |
| else: | |
| atoms = traj_path_or_atoms | |
| cache_key = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12] | |
| print(f"Preparing universal viewer: {len(atoms)} atoms") | |
| print(f"Chemical symbols: {atoms.get_chemical_symbols()}") | |
| # Crear PDB universal | |
| pdb_path = create_universal_pdb(atoms, cache_key) | |
| return pdb_path | |
| except Exception as e: | |
| print(f"Error in prepare_universal_molecule_viewer: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return None | |
| # ==== Funciones de preparación universales ==== | |
| def prepare_input_for_viewer(structure_file): | |
| """Universal input preparation""" | |
| if not structure_file or not os.path.exists(structure_file): | |
| return None | |
| try: | |
| atoms = read(structure_file) | |
| return prepare_universal_molecule_viewer(atoms) | |
| except Exception as e: | |
| print(f"Error preparing input for viewer: {e}") | |
| return None | |
| def prepare_molecule_for_viewer(traj_path): | |
| """Universal trajectory preparation""" | |
| return prepare_universal_molecule_viewer(traj_path) | |
| # ==== OrbMol SPE ==== | |
| from orb_models.forcefield import pretrained | |
| from orb_models.forcefield.calculator import ORBCalculator | |
| def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1): | |
| """ | |
| Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos. | |
| """ | |
| 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) | |
| atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)} | |
| atoms.calc = calc | |
| energy = atoms.get_potential_energy() | |
| forces = atoms.get_forces() | |
| lines = [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/Å"] | |
| # Preparar PDB para visualización | |
| pdb_file = prepare_input_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 | |
| # ==== Simulaciones (helpers) ==== | |
| from simulation_scripts_orbmol import ( | |
| run_md_simulation, | |
| run_relaxation_simulation, | |
| ) | |
| # ==== Wrappers con debug y Molecule3D ==== | |
| 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) | |
| file_path = structure_file | |
| print(f"MD Wrapper: Processing {file_path}") | |
| traj_path, log_text, script_text, explanation = run_md_simulation( | |
| file_path, | |
| int(steps), | |
| 20, # pre-relax steps | |
| 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})" | |
| print(f"MD completed, trajectory: {traj_path}") | |
| pdb_file = prepare_molecule_for_viewer(traj_path) | |
| print(f"PDB file for Molecule3D: {pdb_file}") | |
| return (status, traj_path, log_text, script_text, explanation, pdb_file) | |
| except Exception as e: | |
| print(f"MD Wrapper Error: {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) | |
| file_path = structure_file | |
| print(f"Relax Wrapper: Processing {file_path}") | |
| traj_path, log_text, script_text, explanation = run_relaxation_simulation( | |
| file_path, | |
| 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/Å)" | |
| print(f"Relaxation completed, trajectory: {traj_path}") | |
| pdb_file = prepare_molecule_for_viewer(traj_path) | |
| print(f"PDB file for Molecule3D: {pdb_file}") | |
| return (status, traj_path, log_text, script_text, explanation, pdb_file) | |
| except Exception as e: | |
| print(f"Relax Wrapper Error: {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(): | |
| # -------- SPE -------- | |
| with gr.Tab("Single Point Energy"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions") | |
| gr.Markdown("Upload molecular structure files (.xyz, .pdb, .cif, .traj) for energy and force calculations.") | |
| xyz_input = gr.File( | |
| label="Upload Structure File (.xyz/.pdb/.cif/.traj)", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single" | |
| ) | |
| with gr.Row(): | |
| task_name_spe = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| info="Choose the OrbMol model variant for the calculation." | |
| ) | |
| with gr.Row(): | |
| charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge") | |
| spin_input = gr.Slider(minimum=1, maximum=11, value=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, max_lines=1) | |
| spe_viewer = Molecule3D( | |
| label="Input Structure Viewer", | |
| reps=UNIVERSAL_MOLECULE3D_REPS | |
| ) | |
| # Charge and Spin are only applicable to OMol and OMol-Direct | |
| task_name_spe.input( | |
| lambda x: ( | |
| (gr.Number(visible=True), gr.Number(visible=True)) | |
| if x == "OMol" or x == "OMol-Direct" | |
| else (gr.Number(visible=False), gr.Number(visible=False)) | |
| ), | |
| [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") | |
| gr.Markdown("Upload your molecular structure and configure MD parameters.") | |
| xyz_md = gr.File( | |
| label="Upload Structure File (.xyz/.pdb/.cif/.traj)", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single" | |
| ) | |
| with gr.Row(): | |
| task_name_md = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| info="Choose the OrbMol model variant for the calculation." | |
| ) | |
| with gr.Row(): | |
| charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge") | |
| spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity") | |
| with gr.Row(): | |
| steps_md = gr.Slider(minimum=10, maximum=2000, value=100, step=10, label="Steps") | |
| temp_md = gr.Slider(minimum=10, maximum=1500, value=300, step=10, label="Temperature (K)") | |
| with gr.Row(): | |
| timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=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="Final Structure Viewer (Last MD Frame)", | |
| reps=UNIVERSAL_MOLECULE3D_REPS | |
| ) | |
| md_log = gr.Textbox(label="Log", 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() | |
| # Charge and Spin are only applicable to OMol and OMol-Direct | |
| task_name_md.input( | |
| lambda x: ( | |
| (gr.Number(visible=True), gr.Number(visible=True)) | |
| if x == "OMol" or x == "OMol-Direct" | |
| else (gr.Number(visible=False), gr.Number(visible=False)) | |
| ), | |
| [task_name_md], | |
| [charge_md, spin_md], | |
| ) | |
| run_md_btn.click( | |
| md_wrapper, | |
| inputs=[xyz_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], | |
| ) | |
| # -------- Relax -------- | |
| with gr.Tab("Relaxation / Optimization"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("## Structure Relaxation/Optimization") | |
| gr.Markdown("Upload your molecular structure for geometry optimization.") | |
| xyz_rlx = gr.File( | |
| label="Upload Structure File (.xyz/.pdb/.cif/.traj)", | |
| file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], | |
| file_count="single" | |
| ) | |
| with gr.Row(): | |
| task_name_rlx = gr.Radio( | |
| ["OMol", "OMat", "OMol-Direct"], | |
| value="OMol", | |
| label="Model Type", | |
| info="Choose the OrbMol model variant for the calculation." | |
| ) | |
| with gr.Row(): | |
| steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps") | |
| fmax_rlx = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)") | |
| with gr.Row(): | |
| charge_rlx = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge") | |
| spin_rlx = gr.Slider(minimum=1, maximum=11, value=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=UNIVERSAL_MOLECULE3D_REPS | |
| ) | |
| rlx_log = gr.Textbox(label="Log", 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() | |
| # Charge and Spin are only applicable to OMol and OMol-Direct | |
| task_name_rlx.input( | |
| lambda x: ( | |
| (gr.Number(visible=True), gr.Number(visible=True)) | |
| if x == "OMol" or x == "OMol-Direct" | |
| else (gr.Number(visible=False), gr.Number(visible=False)) | |
| ), | |
| [task_name_rlx], | |
| [charge_rlx, spin_rlx], | |
| ) | |
| run_rlx_btn.click( | |
| relax_wrapper, | |
| inputs=[xyz_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], | |
| ) | |
| print("Starting OrbMol model loading…") | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True) |