Update app.py

app.py

@@ -1,24 +1,60 @@
 import os
-import tempfile
+import re
 import numpy as np
 import gradio as gr
 from ase.io import read, write
 from ase.io.trajectory import Trajectory
-import hashlib
 
-# ==== Forzar visor HTML con 3Dmol.js (sin gradio_molecule3d) ====
-HAVE_MOL3D = False
+# =========================
+# 3Dmol.js helpers (preview y trayectorias)
+# =========================
+
+def _atoms_to_xyz_block(atoms):
+    symbols = atoms.get_chemical_symbols()
+    coords = atoms.get_positions()
+    parts = [str(len(symbols)), "frame"]
+    for s, (x, y, z) in zip(symbols, coords):
+        parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
+    return "\n".join(parts)
+
+def structure_to_html(file_path, width=520, height=520):
+    """Renderiza estructura única en 3Dmol.js sin correr cálculos."""
+    if not file_path or not os.path.exists(file_path):
+        return "<div style='color:#b00; padding:20px;'>No structure file found</div>"
+    try:
+        atoms = read(file_path)
+    except Exception as e:
+        return f"<div style='color:#b00; padding:20px;'>Error reading structure: {e}</div>"
+
+    xyz_block = _atoms_to_xyz_block(atoms).replace("`", "\\`")
+    viewer_id = f"viewer_{abs(hash(('single', file_path))) % 100000}"
+    return f"""
+<div style="margin-bottom:10px; padding:10px; background:#f5f5f5; border-radius:5px;">
+  <strong>🧬 3D Molecular Viewer</strong> — preview (no compute)
+</div>
+<div id="{viewer_id}" style="width:{width}px; height:{height}px; position:relative; border:2px solid #ddd; border-radius:8px; background:#fafafa;"></div>
+<script>
+const _load3Dmol = (ok) => {{
+  if (typeof window.$3Dmol !== 'undefined') return ok();
+  const s=document.createElement('script');
+  s.src='https://3dmol.org/build/3Dmol-min.js';
+  s.onload=ok; document.head.appendChild(s);
+}};
+_load3Dmol(() => {{
+  const el=document.getElementById("{viewer_id}");
+  if(!el||typeof $3Dmol==='undefined') return;
+  const v=$3Dmol.createViewer(el, {{backgroundColor:'white'}});
+  v.addModel(`{xyz_block}`, "xyz");
+  v.setStyle({{}}, {{stick:{{}}, sphere:{{}}}});
+  v.zoomTo(); v.render();
+}});
+</script>
+"""
 
-# ==== Fallback HTML con 3Dmol.js ====
 def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
-    """
-    Lee una .traj de ASE y genera un visor HTML (3Dmol.js) con animación.
-    """
+    """Animación 3Dmol.js desde .traj de ASE."""
     if not traj_path or not os.path.exists(traj_path):
         return "<div style='color:#b00; padding:20px;'>No trajectory file found</div>"
-
-    viewer_id = f"viewer_{abs(hash(traj_path)) % 100000}"
-
     try:
         traj = Trajectory(traj_path)
         if len(traj) == 0:
@@ -26,57 +62,43 @@ def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
     except Exception as e:
         return f"<div style='color:#b00; padding:20px;'>Error: {e}</div>"
 
-    xyz_frames = []
-    for atoms in traj:
-        symbols = atoms.get_chemical_symbols()
-        coords = atoms.get_positions()
-        parts = [str(len(symbols)), "frame"]
-        for s, (x, y, z) in zip(symbols, coords):
-            parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
-        xyz_frames.append("\n".join(parts))
-
-    frames_json = str(xyz_frames).replace("'", '"')
-
-    html = f"""
+    frames = [_atoms_to_xyz_block(at) for at in traj]
+    frames_json = str(frames).replace("'", '"')
+    viewer_id = f"viewer_{abs(hash(traj_path)) % 100000}"
+    return f"""
 <div style="margin-bottom:10px; padding:10px; background:#f5f5f5; border-radius:5px;">
-  <strong>🧬 3D Molecular Viewer</strong> — {len(xyz_frames)} frames
+  <strong>🧬 3D Molecular Viewer</strong> — {len(frames)} frames
 </div>
 <div id="{viewer_id}" style="width:{width}px; height:{height}px; position:relative; border:2px solid #ddd; border-radius:8px; background:#fafafa;"></div>
 <script>
-if (typeof window.$3Dmol === 'undefined') {{
-  var script = document.createElement('script');
-  script.src = 'https://3dmol.org/build/3Dmol-min.js';
-  script.onload = function() {{ setTimeout(function() {{ initViewer_{viewer_id}(); }}, 100); }};
-  document.head.appendChild(script);
-}} else {{
-  initViewer_{viewer_id}();
-}}
-function initViewer_{viewer_id}() {{
-  var el = document.getElementById("{viewer_id}");
-  if (!el || typeof $3Dmol === "undefined") return;
-  var viewer = $3Dmol.createViewer(el, {{backgroundColor: 'white'}});
-  var frames = {frames_json};
-  var currentFrame = 0;
-  function showFrame(i) {{
-    viewer.clear();
-    viewer.addModel(frames[i], "xyz");
-    viewer.setStyle({{}}, {{stick: {{}}, sphere: {{}}}});
-    viewer.zoomTo();
-    viewer.render();
+const _load3Dmol = (ok) => {{
+  if (typeof window.$3Dmol !== 'undefined') return ok();
+  const s=document.createElement('script');
+  s.src='https://3dmol.org/build/3Dmol-min.js';
+  s.onload=ok; document.head.appendChild(s);
+}};
+_load3Dmol(() => {{
+  const el=document.getElementById("{viewer_id}");
+  if(!el||typeof $3Dmol==='undefined') return;
+  const v=$3Dmol.createViewer(el, {{backgroundColor:'white'}});
+  const frames={frames_json};
+  let idx=0;
+  function draw(i){{
+    v.clear(); v.addModel(frames[i], "xyz");
+    v.setStyle({{}}, {{stick:{{}}, sphere:{{}}}});
+    v.zoomTo(); v.render();
   }}
-  showFrame(0);
-  if (frames.length > 1) {{
-    setInterval(function() {{
-      currentFrame = (currentFrame + 1) % frames.length;
-      showFrame(currentFrame);
-    }}, {interval_ms});
+  draw(0);
+  if(frames.length>1){{
+    setInterval(()=>{{ idx=(idx+1)%frames.length; draw(idx); }}, {interval_ms});
   }}
-}}
+}});
 </script>
 """
-    return html
 
-# ==== OrbMol SPE ====
+# =========================
+# OrbMol (SPE)
+# =========================
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
@@ -84,25 +106,18 @@ _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"
-        )
+        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). Maneja archivos de Gradio correctamente.
-    """
+    """Single Point Energy + fuerzas (OrbMol)."""
     try:
         calc = _load_orbmol_calc()
         if not structure_file:
             return "Error: Please upload a structure file", "Error"
+        file_path = structure_file  # gr.File(type='filepath') -> str
 
-        # En Gradio (con type='filepath'), structure_file es directamente un path (str)
-        file_path = structure_file
-
-        # Verificar que el archivo existe y no está vacío
         if not os.path.exists(file_path):
             return f"Error: File not found: {file_path}", "Error"
         if os.path.getsize(file_path) == 0:
@@ -116,34 +131,70 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         forces = atoms.get_forces()            # eV/Å
 
         lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
+        norms = np.linalg.norm(forces, axis=1)
         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/Å"]
+        lines += ["", f"Max Force: {float(np.max(norms)):.4f} eV/Å"]
 
         return "\n".join(lines), "Calculation completed with OrbMol"
     except Exception as e:
         return f"Error during calculation: {e}", "Error"
 
-# ==== Simulaciones (helpers) ====
+# =========================
+# Simulaciones (helpers) y parsing de log
+# =========================
 from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
 )
 
-# ==== Wrappers (MD / Relax): usan directamente el path ====
+ENERGY_PATTERNS = [
+    re.compile(r"(?:^|\\s)(?:step|iter|i)\\s*[:=]?\\s*(\\d+).*?(?:E|energy)\\s*[:=]?\\s*(-?\\d+(?:\\.\\d+)?(?:[eE][+-]?\\d+)?)"),
+    re.compile(r"(?:^|\\s)E(?:nergy)?\\s*[:=]?\\s*(-?\\d+(?:\\.\\d+)?(?:[eE][+-]?\\d+)?)"),
+]
+
+def extract_energy_series(log_text):
+    steps, energies = [], []
+    if not log_text:
+        return steps, energies
+    for line in log_text.splitlines():
+        m = ENERGY_PATTERNS[0].search(line)
+        if m:
+            steps.append(int(m.group(1)))
+            energies.append(float(m.group(2)))
+            continue
+        m2 = ENERGY_PATTERNS[1].search(line)
+        if m2:
+            steps.append(len(steps))
+            energies.append(float(m2.group(1)))
+    return steps, energies
+
+def plot_energy(steps, energies, title):
+    import matplotlib.pyplot as plt
+    fig = plt.figure(figsize=(5,3.2))
+    if steps:
+        plt.plot(steps, energies, linewidth=1.6)
+        plt.xlabel("Step")
+    else:
+        plt.plot(range(len(energies)), energies, linewidth=1.6)
+        plt.xlabel("Index")
+    plt.ylabel("Energy (eV)")
+    plt.title(title)
+    plt.tight_layout()
+    return fig
+
+# =========================
+# Wrappers MD / Relax (devuelven: status, traj, log, script, explain, html, plot)
+# =========================
 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  # ya es un path (str)
-
-        # Validaciones rápidas
+            return ("Error: Please upload a structure file", None, "", "", "", "", None)
+        file_path = structure_file
         if not os.path.exists(file_path):
-            return ("Error: File not found: " + str(file_path), None, "", "", "", None, "")
+            return ("Error: File not found: " + str(file_path), None, "", "", "", "", None)
         if os.path.getsize(file_path) == 0:
-            return ("Error: Empty file: " + str(file_path), None, "", "", "", None, "")
+            return ("Error: Empty file: " + str(file_path), None, "", "", "", "", None)
 
         traj_path, log_text, script_text, explanation = run_md_simulation(
             file_path,
@@ -156,25 +207,22 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
             int(spin),
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
-
         html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, None, html_value)
-
+        s, e = extract_energy_series(log_text)
+        fig = plot_energy(s, e, "MD: Energy vs Step")
+        return (status, traj_path, log_text, script_text, explanation, html_value, fig)
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", None, "")
+        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  # ya es un path (str)
-
-        # Validaciones rápidas
+            return ("Error: Please upload a structure file", None, "", "", "", "", None)
+        file_path = structure_file
         if not os.path.exists(file_path):
-            return ("Error: File not found: " + str(file_path), None, "", "", "", None, "")
+            return ("Error: File not found: " + str(file_path), None, "", "", "", "", None)
         if os.path.getsize(file_path) == 0:
-            return ("Error: Empty file: " + str(file_path), None, "", "", "", None, "")
+            return ("Error: Empty file: " + str(file_path), None, "", "", "", "", None)
 
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
             file_path,
@@ -185,14 +233,26 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
             bool(relax_cell),
         )
         status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
-
         html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, None, html_value)
-
+        s, e = extract_energy_series(log_text)
+        fig = plot_energy(s, e, "Relaxation: Energy vs Step")
+        return (status, traj_path, log_text, script_text, explanation, html_value, fig)
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", None, "")
-
-# ==== UI ====
+        return (f"Error: {e}", None, "", "", "", "", None)
+
+def preview_structure(structure_file):
+    """Preview inmediato al subir archivo en SPE."""
+    if not structure_file:
+        return "<div style='color:#b00; padding:20px;'>Upload a file to preview</div>"
+    if not os.path.exists(structure_file):
+        return "<div style='color:#b00; padding:20px;'>File not found</div>"
+    if os.path.getsize(structure_file) == 0:
+        return "<div style='color:#b00; padding:20px;'>Empty file</div>"
+    return structure_to_html(structure_file)
+
+# =========================
+# UI
+# =========================
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
         # -------- SPE --------
@@ -200,25 +260,25 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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.")
-                    
-                    # IMPORTANTE: type='filepath'
+                    gr.Markdown("Upload molecular structure files (.xyz, .pdb, .cif, .traj, .mol, .sdf) for preview and calculation.")
                     xyz_input = gr.File(
-                        label="Upload Structure File (.xyz/.pdb/.cif/.traj/.mol/.sdf)",
+                        label="Upload Structure File",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single",
-                        type="filepath"
+                        type="filepath",
                     )
-                    
                     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):
+                with gr.Column(variant="panel", min_width=520):
+                    spe_preview = gr.HTML(label="Structure Preview")
                     spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                     spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
 
+            # Preview inmediato (sin cálculo)
+            xyz_input.change(preview_structure, inputs=[xyz_input], outputs=[spe_preview])
+            # Cálculo SPE
             run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status])
 
         # -------- MD --------
@@ -226,15 +286,12 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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/.mol/.sdf)",
+                        label="Upload Structure File",
                         file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                         file_count="single",
-                        type="filepath"
+                        type="filepath",
                     )
-                    
                     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")
@@ -247,4 +304,56 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     run_md_btn = gr.Button("Run MD Simulation", variant="primary")
 
                 with gr.Column(variant="panel", min_width=520):
-                    md_status = gr.Text_
+                    md_status = gr.Textbox(label="MD Status", interactive=False)
+                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
+                    md_html = gr.HTML(label="Trajectory Viewer")
+                    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()
+                    md_plot = gr.Plot(label="Energy vs Step")
+
+            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_html, md_plot],
+            )
+
+        # -------- 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",
+                        type="filepath",
+                    )
+                    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_html = gr.HTML(label="Final Structure / Trajectory")
+                    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()
+                    rlx_plot = gr.Plot(label="Energy vs Step")
+
+            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_html, rlx_plot],
+            )
+
+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)