Update simulation_scripts_orbmol.py

simulation_scripts_orbmol.py

@@ -1,3 +1,4 @@
+# simulation_scripts_orbmol.py
 """
 Minimal FAIRChem-like simulation helpers for OrbMol (local inference).
 
@@ -8,6 +9,7 @@ from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
     atoms_to_xyz,
+    last_frame_xyz_from_traj,
 )
 """
 
@@ -15,6 +17,7 @@ from __future__ import annotations
 import os
 import tempfile
 from pathlib import Path
+from typing import Tuple
 
 import numpy as np
 import ase
@@ -44,8 +47,6 @@ def load_orbmol_model(device: str = "cpu", precision: str = "float32-high") -> O
     """
     global _model_calc
     if _model_calc is None:
-        # NOTE: orb_v3_conservative_inf_omat is the conservative Orb family entry point
-        # used in OrbMol blog; works for molecules (aperiodic).
         orbff = pretrained.orb_v3_conservative_inf_omat(
             device=device,
             precision=precision,
@@ -66,15 +67,63 @@ def atoms_to_xyz(atoms: ase.Atoms) -> str:
         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 atoms.cell.any():
+    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 _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
+
+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.
@@ -105,7 +154,7 @@ def validate_ase_atoms(structure_file: str | Path, max_atoms: int = 5000) -> ase
 # Molecular Dynamics (MD)
 # -----------------------------
 def run_md_simulation(
-    structure_file: str | Path,
+    structure_file_or_xyz: str | Path,
     num_steps: int,
     num_prerelax_steps: int,
     md_timestep: float,       # fs
@@ -117,14 +166,19 @@ def run_md_simulation(
 ) -> tuple[str, str, str, str]:
     """
     Run short MD using OrbMol.
+    Accepts a path or an XYZ string.
     Returns: (traj_path, md_log_text, reproduction_script, explanation)
     """
     traj_path = None
     md_log_path = None
     atoms = None
+    realized_path = None
+    is_temp = False
 
     try:
-        atoms = validate_ase_atoms(structure_file)
+        # Permitir tanto ruta como string XYZ
+        realized_path, is_temp = _materialize_input_to_file(structure_file_or_xyz)
+        atoms = validate_ase_atoms(realized_path)
 
         # Attach the calculator
         calc = load_orbmol_model()
@@ -140,7 +194,7 @@ def run_md_simulation(
 
         # Quick pre-relaxation to remove bad contacts
         opt = LBFGS(atoms, logfile=md_log_path, trajectory=traj_path)
-        if num_prerelax_steps > 0:
+        if int(num_prerelax_steps) > 0:
             opt.run(fmax=0.05, steps=int(num_prerelax_steps))
 
         # Initialize velocities (double T after relaxation as in UMA demo)
@@ -220,22 +274,24 @@ dyn.run({int(num_steps)})
         return traj_path, md_log_text, reproduction_script, explanation
 
     except Exception as e:
-        # Bubble up a clean error
         raise RuntimeError(f"Error running MD: {e}") from e
     finally:
         # Detach calculator to free memory
         if atoms is not None and getattr(atoms, "calc", None) is not None:
             atoms.calc = None
-        if md_log_path and not os.path.exists(md_log_path):
-            md_log_path = None
-        # (No deletion of traj/log here; the UI needs the files.)
+        # 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: str | Path,
+    structure_file_or_xyz: str | Path,
     num_steps: int,
     fmax: float,               # eV/Å
     total_charge: int,
@@ -245,14 +301,18 @@ def run_relaxation_simulation(
 ) -> tuple[str, str, str, str]:
     """
     Run LBFGS relaxation (with optional cell relaxation).
+    Accepts a path or an XYZ string.
     Returns: (traj_path, log_text, reproduction_script, explanation)
     """
     traj_path = None
     opt_log_path = None
     atoms = None
+    realized_path = None
+    is_temp = False
 
     try:
-        atoms = validate_ase_atoms(structure_file)
+        realized_path, is_temp = _materialize_input_to_file(structure_file_or_xyz)
+        atoms = validate_ase_atoms(realized_path)
 
         calc = load_orbmol_model()
         atoms.info["charge"] = int(total_charge)
@@ -303,3 +363,8 @@ optimizer.run(fmax={float(fmax)}, steps={int(num_steps)})
     finally:
         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