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 import hashlib import shutil # ==== util: PDB writer robusto para Molecule3D ==== def _pdb_cache_path(prefix: str, key: str) -> str: gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio") os.makedirs(gradio_cache_dir, exist_ok=True) return os.path.join(gradio_cache_dir, f"{prefix}_{key}.pdb") def _write_pdb_with_fallback(atoms, pdb_path: str) -> str | None: """ Intenta escribir PDB con ASE usando el id correcto ('proteindatabank'). Si falla, crea un PDB minimal válido (líneas ATOM + END) suficiente para gradio_molecule3d. Devuelve la ruta si existe y tiene tamaño > 0, en caso contrario None. """ try: # writer id correcto en ASE para PDB write(pdb_path, atoms, format="proteindatabank") except Exception as e: print(f"ASE PDB writer failed: {e} — trying manual minimal PDB") try: xyz = atoms.get_positions() syms = atoms.get_chemical_symbols() with open(pdb_path, "w") as f: serial = 1 for (sym, (x, y, z)) in zip(syms, xyz): # Campos esenciales: 'ATOM', serial, nombre (símbolo), resname, chain, res seq, coords y elemento f.write( f"ATOM {serial:5d} {sym:<3s} MOL A 1 " f"{x:8.3f}{y:8.3f}{z:8.3f} 1.00 0.00 {sym:>2s}\n" ) serial += 1 f.write("END\n") except Exception as e2: print(f"Manual PDB fallback failed: {e2}") return None # Validación if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0: print(f"✅ PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)") try: with open(pdb_path, "r") as f: preview = f.read(160) print(f"PDB preview: {preview[:120]}...") except Exception: pass return pdb_path print("❌ PDB file not created or empty") return None # ==== Molecule3D viewer preparation CORREGIDO ==== def prepare_molecule_for_viewer(traj_path): """Convert trajectory to format compatible with Molecule3D viewer""" if not traj_path or not os.path.exists(traj_path): print("No trajectory path provided or file doesn't exist") return None try: traj = Trajectory(traj_path) if len(traj) == 0: print("Empty trajectory") return None print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms") atoms = traj[-1] # clave única por posiciones del último frame pos_hash = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12] pdb_path = _pdb_cache_path("molecule", pos_hash) if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0: print(f"PDB already exists: {pdb_path}") return pdb_path return _write_pdb_with_fallback(atoms, pdb_path) except Exception as e: print(f"Error preparing molecule for viewer: {e}") import traceback; traceback.print_exc() return None # ==== Función para convertir archivo inicial a PDB para SPE ==== def prepare_input_for_viewer(structure_file): """Convert input structure file to PDB for Molecule3D viewer""" if not structure_file or not os.path.exists(structure_file): return None try: atoms = read(structure_file) key = hashlib.md5(str(structure_file).encode()).hexdigest()[:12] pdb_path = _pdb_cache_path("input", key) if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0: print(f"Input PDB already exists: {pdb_path}") return pdb_path return _write_pdb_with_fallback(atoms, pdb_path) except Exception as e: print(f"Error preparing input for viewer: {e}") return None # ==== OrbMol 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(structure_file, charge=0, spin_multiplicity=1): """ Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos. """ try: calc = _load_orbmol_calc() 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), "Calculation completed with OrbMol", 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, 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", 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, 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), 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(): 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=[ { "model": 0, "chain": "", "resname": "", "style": "stick", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 1.0 } ] ) run_spe.click(predict_molecule, [xyz_input, 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(): 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=[ { "model": 0, "chain": "", "resname": "", "style": "stick", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 1.0 }, { "model": 0, "chain": "", "resname": "", "style": "sphere", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 0.7 } ] ) 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() 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_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" ) 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=[ { "model": 0, "chain": "", "resname": "", "style": "stick", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 1.0 }, { "model": 0, "chain": "", "resname": "", "style": "sphere", "color": "whiteCarbon", "residue_range": "", "around": 0, "byres": False, "visible": True, "opacity": 0.7 } ] ) 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() run_rlx_btn.click( relax_wrapper, inputs=[xyz_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…") _ = _load_orbmol_calc() if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)