Update app.py

app.py

@@ -7,25 +7,64 @@ from ase.io import read
 from ase.io.trajectory import Trajectory
 import hashlib
 
-# ==== Intentar visor nativo como en UMA (opcional) ====
+# ==== Usar componente nativo Molecule3D ====
 try:
     from gradio_molecule3d import Molecule3D
-    HAVE_MOL3D = False  # Forzar HTML siempre
-except Exception:
+    HAVE_MOL3D = True  # Ahora SÍ usar el componente nativo
+    print("✅ gradio_molecule3d loaded successfully")
+except Exception as e:
     HAVE_MOL3D = False
+    print(f"❌ gradio_molecule3d not available: {e}")
 
-# ==== Fallback HTML con 3Dmol.js ====
+# ==== Función para convertir trayectoria a archivo temporal XYZ ====
+def traj_to_molecule3d_file(traj_path):
+    """
+    Convierte una trayectoria ASE (.traj) a un archivo XYZ temporal para Molecule3D.
+    Retorna el path del archivo temporal.
+    """
+    if not traj_path or not os.path.exists(traj_path):
+        return None
+    
+    try:
+        traj = Trajectory(traj_path)
+        if len(traj) == 0:
+            return None
+            
+        # Usar último frame para visualización estática
+        atoms = traj[-1]
+        
+        # Crear archivo XYZ temporal
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
+            symbols = atoms.get_chemical_symbols()
+            coords = atoms.get_positions()
+            
+            f.write(f"{len(symbols)}\n")
+            f.write("Generated by OrbMol\n")
+            for s, (x, y, z) in zip(symbols, coords):
+                f.write(f"{s} {x:.6f} {y:.6f} {z:.6f}\n")
+            
+            return f.name
+        
+    except Exception as e:
+        print(f"Error converting trajectory: {e}")
+        return None
+
+# ==== Fallback HTML con 3Dmol.js (por si acaso) ====
 def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
     """
-    Render de una trayectoria ASE (.traj) con 3Dmol.js (sin depender de Jupyter).
+    Fallback HTML viewer (solo si Molecule3D no está disponible).
     """
-    # ID único basado en el path
-    viewer_id = f"viewer_{abs(hash(traj_path))}"
+    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)) % 10000}"
     
     try:
         traj = Trajectory(traj_path)
+        if len(traj) == 0:
+            return "<div style='color:#555; padding:20px;'>Empty trajectory</div>"
     except Exception as e:
-        return f"<div style='color:#b00'>Error leyendo trayectoria: {e}</div>"
+        return f"<div style='color:#b00; padding:20px;'>Error: {e}</div>"
 
     xyz_frames = []
     for atoms in traj:
@@ -36,37 +75,27 @@ def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
             parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
         xyz_frames.append("\n".join(parts))
 
-    if not xyz_frames:
-        return "<div style='color:#555'>Empty trajectory</div>"
-
     frames_json = str(xyz_frames).replace("'", '"')
 
     html = f"""
-<div style="margin-bottom:10px;">
-    <strong>3D Molecular Viewer</strong> - {len(xyz_frames)} frames
+<div style="margin-bottom:10px; padding:10px; background:#f5f5f5; border-radius:5px;">
+    <strong>🧬 3D Molecular Viewer</strong> - {len(xyz_frames)} frames
 </div>
-<div id="{viewer_id}" style="width:{width}px; height:{height}px; position:relative; border:1px solid #ccc; background:white;"></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>
-// Load 3Dmol if not already loaded
 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_{abs(hash(traj_path))}(); }}, 100);
+        setTimeout(function() {{ initViewer_{viewer_id}(); }}, 100);
     }};
     document.head.appendChild(script);
 }} else {{
-    initViewer_{abs(hash(traj_path))}();
+    initViewer_{viewer_id}();
 }}
-
-function initViewer_{abs(hash(traj_path))}() {{
+function initViewer_{viewer_id}() {{
     var el = document.getElementById("{viewer_id}");
-    if (!el) return;
-    
-    if (typeof $3Dmol === "undefined") {{
-        el.innerHTML = '<div style="padding:20px;text-align:center;color:#666;">3Dmol.js no disponible</div>';
-        return;
-    }}
+    if (!el || typeof $3Dmol === "undefined") return;
     
     var viewer = $3Dmol.createViewer(el, {{backgroundColor: 'white'}});
     var frames = {frames_json};
@@ -93,7 +122,7 @@ function initViewer_{abs(hash(traj_path))}() {{
 """
     return html
 
-# ==== OrbMol SPE directo (tu calculadora NO se toca) ====
+# ==== OrbMol SPE directo (sin cambios) ====
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
@@ -142,13 +171,12 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
     except Exception as e:
         return f"Error during calculation: {e}", "Error"
 
-# ==== Simulaciones (helpers locales, ya los tienes en simulation_scripts_orbmol.py) ====
+# ==== Simulaciones (sin cambios) ====
 from simulation_scripts_orbmol import (
     run_md_simulation,
     run_relaxation_simulation,
 )
 
-# --- Utilidad: si el usuario pega XYZ en el textbox, guardamos a .xyz temporal ---
 def _string_looks_like_xyz(text: str) -> bool:
     try:
         first = (text or "").strip().splitlines()[0]
@@ -165,7 +193,7 @@ def _to_file_if_xyz(input_or_path: str):
         return tf.name, True
     return input_or_path, False
 
-# Wrappers: devuelven SIEMPRE (status, traj_path, log, script, explain, html_fallback)
+# Wrappers actualizados para devolver archivos para Molecule3D
 def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
     tmp_created = False
     path_or_str = xyz_content
@@ -184,12 +212,16 @@ def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
-        # SIEMPRE generar HTML ya que HAVE_MOL3D = False
-        html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, html_value)
+        # Usar Molecule3D si está disponible, sino HTML
+        if HAVE_MOL3D:
+            xyz_file = traj_to_molecule3d_file(traj_path)
+            return (status, traj_path, log_text, script_text, explanation, xyz_file, "")
+        else:
+            html_value = traj_to_html(traj_path)
+            return (status, traj_path, log_text, script_text, explanation, None, html_value)
 
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", "")
+        return (f"Error: {e}", None, "", "", "", None, "")
     finally:
         if tmp_created and isinstance(path_or_str, str) and os.path.exists(path_or_str):
             try: os.remove(path_or_str)
@@ -211,18 +243,22 @@ def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
         )
         status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
-        # SIEMPRE generar HTML ya que HAVE_MOL3D = False
-        html_value = traj_to_html(traj_path)
-        return (status, traj_path, log_text, script_text, explanation, html_value)
+        # Usar Molecule3D si está disponible, sino HTML
+        if HAVE_MOL3D:
+            xyz_file = traj_to_molecule3d_file(traj_path)
+            return (status, traj_path, log_text, script_text, explanation, xyz_file, "")
+        else:
+            html_value = traj_to_html(traj_path)
+            return (status, traj_path, log_text, script_text, explanation, None, html_value)
 
     except Exception as e:
-        return (f"Error: {e}", None, "", "", "", "")
+        return (f"Error: {e}", None, "", "", "", None, "")
     finally:
         if tmp_created and isinstance(path_or_str, str) and os.path.exists(path_or_str):
             try: os.remove(path_or_str)
             except Exception: pass
 
-# ==== Ejemplos ====
+# ==== Ejemplos (sin cambios) ====
 examples = [
     ["""2
 Hydrogen molecule
@@ -242,10 +278,10 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1],
 ]
 
-# ==== UI ====
+# ==== UI actualizada ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
-        # -------- SPE --------
+        # -------- SPE (sin cambios) --------
         with gr.Tab("Single Point Energy"):
             with gr.Row():
                 with gr.Column(scale=2):
@@ -263,7 +299,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
             run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status])
 
-        # -------- MD --------
+        # -------- MD (actualizada con Molecule3D) --------
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
@@ -283,8 +319,13 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     md_status = gr.Textbox(label="MD Status", interactive=False)
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     
-                    # Siempre usar HTML ya que HAVE_MOL3D = False
-                    md_html = gr.HTML(label="Trajectory Viewer")
+                    # Usar Molecule3D si está disponible
+                    if HAVE_MOL3D:
+                        md_viewer = Molecule3D(label="3D Molecular Viewer")
+                        md_html = gr.HTML(visible=False)  # Oculto cuando usamos Molecule3D
+                    else:
+                        md_viewer = gr.HTML(visible=False)  # Placeholder
+                        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)
@@ -293,10 +334,10 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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],
+                outputs=[md_status, md_traj, md_log, md_script, md_explain, md_viewer, md_html],
             )
 
-        # -------- Relax --------
+        # -------- Relax (actualizada con Molecule3D) --------
         with gr.Tab("Relaxation / Optimization"):
             with gr.Row():
                 with gr.Column(scale=2):
@@ -313,8 +354,13 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                     rlx_status = gr.Textbox(label="Status", interactive=False)
                     rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                     
-                    # Siempre usar HTML ya que HAVE_MOL3D = False
-                    rlx_html = gr.HTML(label="Final Structure")
+                    # Usar Molecule3D si está disponible
+                    if HAVE_MOL3D:
+                        rlx_viewer = Molecule3D(label="Final Structure")
+                        rlx_html = gr.HTML(visible=False)  # Oculto cuando usamos Molecule3D
+                    else:
+                        rlx_viewer = gr.HTML(visible=False)  # Placeholder
+                        rlx_html = gr.HTML(label="Final Structure")
                     
                     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)
@@ -323,7 +369,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             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],
+                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer, rlx_html],
             )
 
 print("Starting OrbMol model loading…")