Update simulation_scripts_orbmol.py

simulation_scripts_orbmol.py

@@ -1,24 +1,11 @@
-# simulation_scripts_orbmol.py
 """
-Minimal FAIRChem-like simulation helpers for OrbMol (local inference).
-
-Usage from app.py:
-from simulation_scripts_orbmol import (
-    load_orbmol_model,
-    validate_ase_atoms,
-    run_md_simulation,
-    run_relaxation_simulation,
-    atoms_to_xyz,
-    last_frame_xyz_from_traj,
-)
+Simulaciones OrbMol con interfaz estilo Facebook FAIRChem
 """
 
 from __future__ import annotations
 import os
 import tempfile
 from pathlib import Path
-from typing import Tuple
-
 import numpy as np
 import ase
 import ase.io
@@ -35,16 +22,13 @@ from ase.md.nose_hoover_chain import NoseHooverChainNVT
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-
 # -----------------------------
-# Global model (lazy singleton)
+# Global model
 # -----------------------------
 _model_calc: ORBCalculator | None = None
 
 def load_orbmol_model(device: str = "cpu", precision: str = "float32-high") -> ORBCalculator:
-    """
-    Load OrbMol once and reuse the same calculator.
-    """
+    """Load OrbMol once and reuse the same calculator."""
     global _model_calc
     if _model_calc is None:
         orbff = pretrained.orb_v3_conservative_inf_omat(
@@ -54,182 +38,136 @@ def load_orbmol_model(device: str = "cpu", precision: str = "float32-high") -> O
         _model_calc = ORBCalculator(orbff, device=device)
     return _model_calc
 
-
 # -----------------------------
-# Helpers
+# FUNCIONES ESTILO FACEBOOK - COPIADAS EXACTAS
 # -----------------------------
-def atoms_to_xyz(atoms: ase.Atoms) -> str:
-    """
-    Convert ASE Atoms to an XYZ string for quick visualization.
-    """
-    lines = [str(len(atoms)), "generated by simulation_scripts_orbmol"]
-    for s, (x, y, z) in zip(atoms.get_chemical_symbols(), atoms.get_positions()):
-        lines.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
-    return "\n".join(lines)
 
-def last_frame_xyz_from_traj(traj_path: str | Path) -> str:
-    """
-    Read the last frame of an ASE .traj and return it as XYZ string.
-    """
-    tr = Trajectory(str(traj_path))
-    last = tr[-1]
-    return atoms_to_xyz(last)
-
-def _center_atoms(atoms: ase.Atoms) -> None:
-    """
-    Center coordinates for nicer visualization (no effect on energies).
-    """
-    atoms.positions -= atoms.get_center_of_mass()
-    if atoms.cell is not None and np.array(atoms.cell).any():
-        cell_center = atoms.get_cell().sum(axis=0) / 2
-        atoms.positions += cell_center
+def load_check_ase_atoms(structure_file):
+    """COPIA EXACTA de Facebook - valida y carga estructura"""
+    if not structure_file:
+        raise Exception("You need an input structure file to run a simulation!")
 
-def _string_looks_like_xyz(text: str) -> bool:
-    """
-    Heurística simple para detectar si un input es un XYZ en texto.
-    """
-    if not isinstance(text, str):
-        return False
-    lines = [l for l in text.strip().splitlines() if l.strip()]
-    if len(lines) < 2:
-        return False
-    # primera línea: número de átomos
     try:
-        _ = int(lines[0].split()[0])
-        return True
-    except Exception:
-        return False
+        atoms = ase.io.read(structure_file)
 
-def _materialize_input_to_file(input_or_path: str | Path) -> Tuple[str, bool]:
-    """
-    Devuelve (file_path, is_temp). Si input es un string XYZ, lo guarda a un .xyz temporal.
-    Si es una ruta existente, la devuelve tal cual.
-    """
-    # Caso: dict de Gradio File {'path': ...}
-    if isinstance(input_or_path, dict) and "path" in input_or_path:
-        p = input_or_path["path"]
-        return p, False
-
-    # Caso: Path o ruta existente
-    if isinstance(input_or_path, (str, Path)) and os.path.exists(str(input_or_path)):
-        return str(input_or_path), False
-
-    # Caso: probablemente es un string XYZ
-    if isinstance(input_or_path, str) and _string_looks_like_xyz(input_or_path):
-        tf = tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False)
-        tf.write(input_or_path)
-        tf.flush()
-        tf.close()
-        return tf.name, True
-
-    raise ValueError("Input must be an existing file path or a valid XYZ string.")
-
-def validate_ase_atoms(structure_file: str | Path, max_atoms: int = 5000) -> ase.Atoms:
-    """
-    Read & validate an ASE-compatible file; ensures uniform PBC and non-empty.
-    Returns a centered Atoms object.
-    """
-    if not structure_file:
-        raise ValueError("Missing input structure file path.")
-    atoms = ase.io.read(str(structure_file))
-
-    if len(atoms) == 0:
-        raise ValueError("No atoms found in the input structure.")
+        if not (all(atoms.pbc) or np.all(~np.array(atoms.pbc))):
+            raise Exception(
+                "Mixed PBC are not supported yet - please set PBC all True or False in your structure before uploading"
+            )
 
-    # Uniform PBC (all True or all False). Mixed PBC often breaks MD settings.
-    pbc = np.array(atoms.pbc, dtype=bool)
-    if not (pbc.all() or (~pbc).all()):
-        raise ValueError(f"Mixed PBC {atoms.pbc} not supported. Set all True or all False.")
+        if len(atoms) == 0:
+            raise Exception("Error: Structure file contains no atoms.")
 
-    if len(atoms) > max_atoms:
-        raise ValueError(
-            f"Structure has {len(atoms)} atoms, exceeding the limit of {max_atoms} for this demo."
-        )
+        if len(atoms) > 2000:
+            raise Exception(
+                f"Error: Structure file contains {len(atoms)}, which is more than 2000 atoms. Please use a smaller structure for this demo, or run this on a local machine!"
+            )
 
-    _center_atoms(atoms)
-    return atoms
+        # Centrar para visualización
+        atoms.positions -= atoms.get_center_of_mass()
+        cell_center = atoms.get_cell().sum(axis=0) / 2
+        atoms.positions += cell_center
 
+        return atoms
+    except Exception as e:
+        raise Exception(f"Error loading structure with ASE: {str(e)}")
 
-# -----------------------------
-# Molecular Dynamics (MD)
-# -----------------------------
 def run_md_simulation(
-    structure_file_or_xyz: str | Path,
-    num_steps: int,
-    num_prerelax_steps: int,
-    md_timestep: float,       # fs
-    temperature_k: float,     # K
-    md_ensemble: str,         # "NVE" or "NVT"
-    total_charge: int,
-    spin_multiplicity: int,
+    structure_file,
+    num_steps,
+    num_prerelax_steps,
+    md_timestep,
+    temperature_k,
+    md_ensemble,
+    task_name="OMol",  # Siempre OMol para OrbMol
+    total_charge=0,
+    spin_multiplicity=1,
     explanation: str | None = None,
-) -> tuple[str, str, str, str]:
+    oauth_token=None,  # Ignorado
+    progress=None,  # Para compatibilidad Gradio
+):
     """
-    Run short MD using OrbMol.
-    Accepts a path or an XYZ string.
-    Returns: (traj_path, md_log_text, reproduction_script, explanation)
+    MD simulation estilo Facebook pero con OrbMol
     """
+    temp_path = None
     traj_path = None
     md_log_path = None
     atoms = None
-    realized_path = None
-    is_temp = False
 
     try:
-        # Permitir tanto ruta como string XYZ
-        realized_path, is_temp = _materialize_input_to_file(structure_file_or_xyz)
-        atoms = validate_ase_atoms(realized_path)
+        # Cargar átomos (igual que Facebook)
+        atoms = load_check_ase_atoms(structure_file)
+
+        # Configurar charge y spin
+        atoms.info["charge"] = total_charge
+        atoms.info["spin"] = spin_multiplicity
 
-        # Attach the calculator
+        # AQUÍ EL CAMBIO: OrbMol en lugar de HFEndpointCalculator
         calc = load_orbmol_model()
-        atoms.info["charge"] = int(total_charge)
-        atoms.info["spin"]   = int(spin_multiplicity)
         atoms.calc = calc
 
-        # Output files
-        with tempfile.NamedTemporaryFile(suffix=".traj", delete=False) as tf:
-            traj_path = tf.name
-        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as lf:
-            md_log_path = lf.name
+        # Progress callback si existe
+        interval = 1
+        steps = [0]
+        expected_steps = num_steps + num_prerelax_steps
+
+        def update_progress():
+            steps[-1] += interval
+            if progress:
+                progress(steps[-1] / expected_steps)
 
-        # Quick pre-relaxation to remove bad contacts
+        # Archivos temporales (igual que Facebook)
+        with tempfile.NamedTemporaryFile(suffix=".traj", delete=False) as traj_f:
+            traj_path = traj_f.name
+        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as log_f:
+            md_log_path = log_f.name
+
+        # Pre-relaxación (igual que Facebook)
         opt = LBFGS(atoms, logfile=md_log_path, trajectory=traj_path)
-        if int(num_prerelax_steps) > 0:
-            opt.run(fmax=0.05, steps=int(num_prerelax_steps))
+        if progress:
+            opt.attach(update_progress, interval=interval)
+        opt.run(fmax=0.05, steps=num_prerelax_steps)
 
-        # Initialize velocities (double T after relaxation as in UMA demo)
-        MaxwellBoltzmannDistribution(atoms, temperature_K=2 * float(temperature_k))
+        # Velocidades (igual que Facebook - x2 después de relajación)
+        MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_k * 2)
 
-        # Choose integrator/ensemble
-        if md_ensemble.upper() == "NVT":
+        # Integrador (igual que Facebook)
+        if md_ensemble == "NVE":
+            dyn = VelocityVerlet(atoms, timestep=md_timestep * units.fs)
+        elif md_ensemble == "NVT":
             dyn = NoseHooverChainNVT(
                 atoms,
-                timestep=float(md_timestep) * units.fs,
-                temperature_K=float(temperature_k),
-                tdamp=10 * float(md_timestep) * units.fs,
+                timestep=md_timestep * units.fs,
+                temperature_K=temperature_k,
+                tdamp=10 * md_timestep * units.fs,
             )
-        else:
-            dyn = VelocityVerlet(atoms, timestep=float(md_timestep) * units.fs)
 
-        # Attach trajectory writer and MD logger
+        # Trajectory y logging (igual que Facebook)
         traj = Trajectory(traj_path, "a", atoms)
         dyn.attach(traj.write, interval=1)
+        if progress:
+            dyn.attach(update_progress, interval=interval)
         dyn.attach(
             MDLogger(
-                dyn, atoms, md_log_path, header=True, stress=False, peratom=True, mode="a"
+                dyn,
+                atoms,
+                md_log_path,
+                header=True,
+                stress=False,
+                peratom=True,
+                mode="a",
             ),
             interval=10,
         )
 
-        # Run MD
-        dyn.run(int(num_steps))
+        # Ejecutar MD
+        dyn.run(num_steps)
 
-        # Prepare reproduction script (using OrbMol locally)
-        reproduction_script = f"""\
+        # Script de reproducción (estilo Facebook pero con OrbMol)
+        reproduction_script = f"""
 import ase.io
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.md.verlet import VelocityVerlet
-from ase.md.nose_hoover_chain import NoseHooverChainNVT
 from ase.optimize import LBFGS
 from ase.io.trajectory import Trajectory
 from ase.md import MDLogger
@@ -237,134 +175,193 @@ from ase import units
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
-atoms = ase.io.read('input_file.traj')  # any ASE-readable file
-atoms.info['charge'] = {int(total_charge)}
-atoms.info['spin']   = {int(spin_multiplicity)}
-
+# Read the atoms object from ASE read-able file
+atoms = ase.io.read('input_file.traj')
+# Set the total charge and spin multiplicity
+atoms.info["charge"] = {total_charge}
+atoms.info["spin"] = {spin_multiplicity}
+# Set up the OrbMol calculator
 orbff = pretrained.orb_v3_conservative_inf_omat(device='cpu', precision='float32-high')
 atoms.calc = ORBCalculator(orbff, device='cpu')
-
-opt = LBFGS(atoms, trajectory='relaxation_output.traj')
-opt.run(fmax=0.05, steps={int(num_prerelax_steps)})
-
-MaxwellBoltzmannDistribution(atoms, temperature_K={float(temperature_k)}*2)
-
-ensemble = '{md_ensemble.upper()}'
-if ensemble == 'NVT':
-    dyn = NoseHooverChainNVT(atoms, timestep={float(md_timestep)}*units.fs,
-                             temperature_K={float(temperature_k)}, tdamp=10*{float(md_timestep)}*units.fs)
-else:
-    dyn = VelocityVerlet(atoms, timestep={float(md_timestep)}*units.fs)
-
-dyn.attach(MDLogger(dyn, atoms, 'md.log', header=True, stress=False, peratom=True, mode='w'), interval=10)
-traj = Trajectory('md_output.traj', 'w', atoms)
+# Do a quick pre-relaxation to make sure the system is stable
+opt = LBFGS(atoms, trajectory="relaxation_output.traj")
+opt.run(fmax=0.05, steps={num_prerelax_steps})
+# Initialize the velocity distribution; we set twice the temperature since we did a relaxation and
+# much of the kinetic energy will partition to the potential energy right away
+MaxwellBoltzmannDistribution(atoms, temperature_K={temperature_k}*2)
+# Initialize the integrator; NVE is shown here as an example
+dyn = VelocityVerlet(atoms, timestep={md_timestep} * units.fs)
+# Set up trajectory and MD logger
+dyn.attach(MDLogger(dyn, atoms, 'md.log', header=True, stress=False, peratom=True, mode="w"), interval=10)
+traj = Trajectory("md_output.traj", "w", atoms)
 dyn.attach(traj.write, interval=1)
-dyn.run({int(num_steps)})
-"""
+# Run the simulation!
+dyn.run({num_steps})
+        """
 
-        md_log_text = Path(md_log_path).read_text(encoding="utf-8", errors="ignore")
+        # Leer log
+        with open(md_log_path, "r") as md_log_file:
+            md_log = md_log_file.read()
 
         if explanation is None:
-            explanation = (
-                f"MD of {len(atoms)} atoms for {int(num_steps)} steps at "
-                f"{float(temperature_k)} K, timestep {float(md_timestep)} fs, "
-                f"ensemble {md_ensemble.upper()} (prerelax {int(num_prerelax_steps)} steps)."
-            )
+            explanation = f"MD simulation of {len(atoms)} atoms for {num_steps} steps with a timestep of {md_timestep} fs at {temperature_k} K in the {md_ensemble} ensemble using OrbMol. You submitted this simulation, so I hope you know what you're looking for or what it means!"
 
-        return traj_path, md_log_text, reproduction_script, explanation
+        return traj_path, md_log, reproduction_script, explanation
 
     except Exception as e:
-        raise RuntimeError(f"Error running MD: {e}") from e
+        raise Exception(
+            f"Error running MD simulation: {str(e)}. Please try again or report this error."
+        )
     finally:
-        # Detach calculator to free memory
+        # Limpieza (igual que Facebook)
+        if temp_path and os.path.exists(temp_path):
+            os.remove(temp_path)
+
+        if md_log_path and os.path.exists(md_log_path):
+            os.remove(md_log_path)
+
         if atoms is not None and getattr(atoms, "calc", None) is not None:
             atoms.calc = None
-        # Limpieza del .xyz temporal si lo generamos nosotros
-        if is_temp and realized_path and os.path.exists(realized_path):
-            try:
-                os.remove(realized_path)
-            except Exception:
-                pass
-
 
-# -----------------------------
-# Geometry optimization
-# -----------------------------
 def run_relaxation_simulation(
-    structure_file_or_xyz: str | Path,
-    num_steps: int,
-    fmax: float,               # eV/Å
-    total_charge: int,
-    spin_multiplicity: int,
-    relax_unit_cell: bool,
+    structure_file,
+    num_steps,
+    fmax,
+    task_name="OMol",  # Siempre OMol para OrbMol
+    total_charge: float = 0,
+    spin_multiplicity: float = 1,
+    relax_unit_cell=False,
     explanation: str | None = None,
-) -> tuple[str, str, str, str]:
+    oauth_token=None,  # Ignorado
+    progress=None,
+):
     """
-    Run LBFGS relaxation (with optional cell relaxation).
-    Accepts a path or an XYZ string.
-    Returns: (traj_path, log_text, reproduction_script, explanation)
+    Relaxation simulation estilo Facebook pero con OrbMol
     """
+    temp_path = None
     traj_path = None
     opt_log_path = None
     atoms = None
-    realized_path = None
-    is_temp = False
 
     try:
-        realized_path, is_temp = _materialize_input_to_file(structure_file_or_xyz)
-        atoms = validate_ase_atoms(realized_path)
+        # Cargar átomos (igual que Facebook)
+        atoms = load_check_ase_atoms(structure_file)
+
+        # Configurar charge y spin
+        atoms.info["charge"] = total_charge
+        atoms.info["spin"] = spin_multiplicity
 
+        # AQUÍ EL CAMBIO: OrbMol en lugar de HFEndpointCalculator
         calc = load_orbmol_model()
-        atoms.info["charge"] = int(total_charge)
-        atoms.info["spin"]   = int(spin_multiplicity)
         atoms.calc = calc
 
-        with tempfile.NamedTemporaryFile(suffix=".traj", delete=False) as tf:
-            traj_path = tf.name
-        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as lf:
-            opt_log_path = lf.name
+        # Archivos temporales
+        with tempfile.NamedTemporaryFile(suffix=".traj", delete=False) as traj_f:
+            traj_path = traj_f.name
+        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as log_f:
+            opt_log_path = log_f.name
+
+        # Optimizador (igual que Facebook)
+        optimizer = LBFGS(
+            FrechetCellFilter(atoms) if relax_unit_cell else atoms,
+            trajectory=traj_path,
+            logfile=opt_log_path,
+        )
+
+        # Progress callback si existe
+        if progress:
+            interval = 1
+            steps = [0]
+            def update_progress(steps):
+                steps[-1] += interval
+                progress(steps[-1] / num_steps)
+            optimizer.attach(update_progress, interval=interval, steps=steps)
 
-        subject = FrechetCellFilter(atoms) if relax_unit_cell else atoms
-        optimizer = LBFGS(subject, trajectory=traj_path, logfile=opt_log_path)
-        optimizer.run(fmax=float(fmax), steps=int(num_steps))
+        # Ejecutar optimización
+        optimizer.run(fmax=fmax, steps=num_steps)
 
-        reproduction_script = f"""\
+        # Script de reproducción
+        reproduction_script = f"""
 import ase.io
 from ase.optimize import LBFGS
 from ase.filters import FrechetCellFilter
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
+# Read the atoms object from ASE read-able file
 atoms = ase.io.read('input_file.traj')
-atoms.info['charge'] = {int(total_charge)}
-atoms.info['spin']   = {int(spin_multiplicity)}
-
+# Set the total charge and spin multiplicity
+atoms.info["charge"] = {total_charge}
+atoms.info["spin"] = {spin_multiplicity}
+# Set up the OrbMol calculator
 orbff = pretrained.orb_v3_conservative_inf_omat(device='cpu', precision='float32-high')
 atoms.calc = ORBCalculator(orbff, device='cpu')
+# Initialize the optimizer
+relax_unit_cell = {relax_unit_cell}
+optimizer = LBFGS(FrechetCellFilter(atoms) if relax_unit_cell else atoms, trajectory="relaxation_output.traj")
+# Run the optimization!
+optimizer.run(fmax={fmax}, steps={num_steps})
+        """
 
-relax_unit_cell = {bool(relax_unit_cell)}
-subject = FrechetCellFilter(atoms) if relax_unit_cell else atoms
-optimizer = LBFGS(subject, trajectory='relaxation_output.traj')
-optimizer.run(fmax={float(fmax)}, steps={int(num_steps)})
-"""
-
-        log_text = Path(opt_log_path).read_text(encoding="utf-8", errors="ignore")
+        # Leer log
+        with open(opt_log_path, "r") as opt_log_file:
+            opt_log = opt_log_file.read()
 
         if explanation is None:
-            explanation = (
-                f"Relaxation of {len(atoms)} atoms for up to {int(num_steps)} steps "
-                f"with fmax {float(fmax)} eV/Å (relax_cell={bool(relax_unit_cell)})."
-            )
+            explanation = f"Relaxation of {len(atoms)} atoms for {num_steps} steps with a force tolerance of {fmax} eV/Å using OrbMol. You submitted this simulation, so I hope you know what you're looking for or what it means!"
 
-        return traj_path, log_text, reproduction_script, explanation
+        return traj_path, opt_log, reproduction_script, explanation
 
     except Exception as e:
-        raise RuntimeError(f"Error running relaxation: {e}") from e
+        raise Exception(
+            f"Error running relaxation: {str(e)}. Please try again or report this error."
+        )
     finally:
+        # Limpieza (igual que Facebook)
+        if temp_path and os.path.exists(temp_path):
+            os.remove(temp_path)
+        if opt_log_path and os.path.exists(opt_log_path):
+            os.remove(opt_log_path)
         if atoms is not None and getattr(atoms, "calc", None) is not None:
             atoms.calc = None
-        if is_temp and realized_path and os.path.exists(realized_path):
-            try:
-                os.remove(realized_path)
-            except Exception:
-                pass
+
+# -----------------------------
+# Helper functions para compatibilidad
+# -----------------------------
+
+def atoms_to_xyz(atoms: ase.Atoms) -> str:
+    """Convert ASE Atoms to an XYZ string for quick visualization."""
+    lines = [str(len(atoms)), "generated by simulation_scripts_orbmol"]
+    for s, (x, y, z) in zip(atoms.get_chemical_symbols(), atoms.get_positions()):
+        lines.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
+    return "\n".join(lines)
+
+def last_frame_xyz_from_traj(traj_path: str | Path) -> str:
+    """Read the last frame of an ASE .traj and return it as XYZ string."""
+    tr = Trajectory(str(traj_path))
+    last = tr[-1]
+    return atoms_to_xyz(last)
+
+# Función de validación simplificada (sin autenticación)
+def validate_ase_atoms_and_login(structure_file, login_button_value="", oauth_token=None):
+    """Validación simplificada - sin login UMA"""
+    if not structure_file:
+        return (False, False, "Missing input structure!")
+    
+    if isinstance(structure_file, dict):
+        structure_file = structure_file["path"]
+    
+    try:
+        atoms = ase.io.read(structure_file)
+        
+        if len(atoms) == 0:
+            return (False, False, "No atoms in the structure file!")
+        elif not (all(atoms.pbc) or np.all(~np.array(atoms.pbc))):
+            return (False, False, f"Mixed PBC {atoms.pbc} not supported!")
+        elif len(atoms) > 2000:
+            return (False, False, f"Too many atoms ({len(atoms)}), max 2000!")
+        else:
+            return (True, True, "Structure loaded successfully - ready for OrbMol simulation!")
+            
+    except Exception as e:
+        return (False, False, f"Failed to load structure: {str(e)}")