Update app.py

separate logic (orbmol and visualization )

app.py

@@ -10,12 +10,12 @@ from ase.md.verlet import VelocityVerlet
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.io.trajectory import Trajectory
 
-# py3Dmol es opcional; si no está instalado haremos fallback a 3Dmol.js
+# Intentar importar Molecule3D para vista 3D nativa
 try:
-    import py3Dmol  # noqa: F401
-    HAVE_PY3DMOL = True
+    from gradio_molecule3d import Molecule3D
+    HAVE_MOL3D = True
 except Exception:
-    HAVE_PY3DMOL = False
+    HAVE_MOL3D = False
 
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
@@ -67,16 +67,13 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         forces = atoms.get_forces()            # eV/Å
 
         lines = []
-        lines.append(f"Total Energy: {energy:.6f} eV")
-        lines.append("")
+        lines.append(f"Total Energy: {energy:.6f} eV\n")
         lines.append("Atomic Forces:")
         for i, f in enumerate(forces):
-            lines.append(
-                f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å"
-            )
+            lines.append(f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å")
+
         max_force = float(np.max(np.linalg.norm(forces, axis=1)))
-        lines.append("")
-        lines.append(f"Max Force: {max_force:.4f} eV/Å")
+        lines.append(f"\nMax Force: {max_force:.4f} eV/Å")
 
         os.unlink(xyz_file)
         return "\n".join(lines), "Calculation completed with OrbMol"
@@ -86,17 +83,11 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
 
 
 # =========================
-#   Trajectory → HTML 3D
-#   (funciona con o sin py3Dmol)
+#   Trajectory → HTML 3D fallback
 # =========================
 def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
-    """
-    Genera un bloque HTML con 3Dmol para animar una lista de frames en XYZ.
-    Sin dependencias de JupyterLab.
-    """
     traj = Trajectory(traj_path)
     xyz_frames = []
-
     for atoms in traj:
         symbols = atoms.get_chemical_symbols()
         coords = atoms.get_positions()
@@ -105,37 +96,25 @@ 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))
 
-    # HTML+JS autopackage
     html = f"""
-<div id="viewer_md" style="width:{width}px; height:{height}px; position:relative;"></div>
+<div id="viewer_md" style="width:{width}px; height:{height}px;"></div>
 <script src="https://3dmol.org/build/3Dmol-min.js"></script>
 <script>
 (function() {{
-  var element = document.getElementById('viewer_md');
-  if (!element) return;
-  var viewer = $3Dmol.createViewer(element, {{backgroundColor: 'white'}});
+  var viewer = $3Dmol.createViewer("viewer_md", {{backgroundColor: 'white'}});
   var frames = {xyz_frames!r};
   var i = 0;
-
   function show(i) {{
     viewer.clear();
-    viewer.addModel(frames[i], 'xyz');
-    viewer.setStyle({{}}, {{stick: {{}} }});
+    viewer.addModel(frames[i], "xyz");
+    viewer.setStyle({{}}, {{stick: {{}}}});
     viewer.zoomTo();
     viewer.render();
   }}
-
-  if (frames.length === 0) {{
-    element.innerHTML = '<div style="padding:8px;color:#555">Empty trajectory</div>';
-  }} else if (frames.length === 1) {{
-    show(0);
-  }} else {{
-    show(0);
-    setInterval(function() {{
-      i = (i + 1) % frames.length;
-      show(i);
-    }}, {int(interval_ms)});
-  }}
+  if(frames.length>0) show(0);
+  if(frames.length>1) setInterval(function(){{
+    i=(i+1)%frames.length; show(i);
+  }}, {int(interval_ms)});
 }})();
 </script>
 """
@@ -168,38 +147,39 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1,
         opt = LBFGS(atoms, logfile=None)
         opt.run(fmax=0.05, steps=20)
 
-        # Velocidades (2*T por partición de energía)
         MaxwellBoltzmannDistribution(atoms, temperature_K=2*float(temperature))
 
-        # Integrador NVE
         dyn = VelocityVerlet(atoms, timestep=float(timestep) * units.fs)
 
-        # Trayectoria temporal
         tf = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
         tf.close()
         traj = Trajectory(tf.name, "w", atoms)
         dyn.attach(traj.write, interval=1)
-
-        # Ejecutar MD
         dyn.run(int(steps))
 
-        # Visualización 3D
-        html = traj_to_html(tf.name)
+        if HAVE_MOL3D:
+            # Mostrar último frame en Molecule3D
+            last = traj[-1]
+            mol_xyz = f"{len(last)}\nFinal frame\n"
+            for s, (x, y, z) in zip(last.get_chemical_symbols(), last.get_positions()):
+                mol_xyz += f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
+            view = Molecule3D(value=mol_xyz, label="Final Frame (XYZ)")
+        else:
+            view = traj_to_html(tf.name)
 
-        # Limpieza archivo xyz temporal
         try:
             os.unlink(xyz_file)
         except Exception:
             pass
 
-        return f"MD completed: {int(steps)} steps at {int(temperature)} K", html
+        return f"MD completed: {int(steps)} steps at {int(temperature)} K", view
 
     except Exception as e:
         return f"Error during MD simulation: {str(e)}", ""
 
 
 # =========================
-#   Ejemplos (tu set)
+#   Ejemplos
 # =========================
 examples = [
     ["""2
@@ -228,67 +208,31 @@ H -0.3630 -0.5133 -0.8887""", 0, 1],
 # =========================
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
-        # -------- Tab 1: tu layout original --------
+        # -------- Tab 1: Single Point --------
         with gr.Tab("Single Point Energy"):
             with gr.Row():
                 with gr.Column(scale=2):
                     with gr.Column(variant="panel"):
-                        gr.Markdown("# OrbMol Demo - Quantum-Accurate Molecular Predictions")
+                        gr.Markdown("# OrbMol Demo - Quantum-Accurate Predictions")
+                        gr.Markdown("OrbMol is a neural network potential trained on the OMol25 dataset.")
 
-                        gr.Markdown(
-                            "OrbMol is a neural network potential trained on the OMol25 dataset. "
-                            "Predict energies and forces with configurable charge and spin."
+                        xyz_input = gr.Textbox(
+                            label="XYZ Coordinates",
+                            placeholder="Paste XYZ here...",
+                            lines=12,
                         )
 
-                        gr.Markdown("## Simulation inputs")
-                        with gr.Column(variant="panel"):
-                            gr.Markdown("### Input molecular structure")
-
-                            xyz_input = gr.Textbox(
-                                label="XYZ Coordinates",
-                                placeholder=(
-                                    "3\nWater molecule\n"
-                                    "O 0.0000 0.0000 0.0000\n"
-                                    "H 0.7571 0.0000 0.5864\n"
-                                    "H -0.7571 0.0000 0.5864"
-                                ),
-                                lines=12,
-                                info="Paste XYZ coordinates of your molecule here",
-                            )
-
-                            gr.Markdown("OMol-specific settings for total charge and spin multiplicity")
-                            with gr.Row():
-                                charge_input = gr.Slider(
-                                    value=0, label="Total Charge", minimum=-10, maximum=10, step=1
-                                )
-                                spin_input = gr.Slider(
-                                    value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity"
-                                )
-
-                            predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
-
-                with gr.Column(variant="panel", elem_id="results", min_width=500):
-                    gr.Markdown("## OrbMol Prediction Results")
-
-                    results_output = gr.Textbox(
-                        label="Energy & Forces",
-                        lines=15,
-                        interactive=False,
-                        info="OrbMol energy and force predictions",
-                    )
-
-                    status_output = gr.Textbox(
-                        label="Status",
-                        interactive=False,
-                        max_lines=1,
-                    )
-
-            gr.Markdown("### Examples")
-            gr.Examples(
-                examples=examples,
-                inputs=[xyz_input, charge_input, spin_input],
-                label="Click any example to load it",
-            )
+                        with gr.Row():
+                            charge_input = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
+                            spin_input = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
+
+                        predict_btn = gr.Button("Run OrbMol Prediction", variant="primary")
+
+                with gr.Column(variant="panel", min_width=500):
+                    results_output = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
+                    status_output = gr.Textbox(label="Status", interactive=False, max_lines=1)
+
+            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
 
             predict_btn.click(
                 predict_molecule,
@@ -299,43 +243,27 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
             with gr.Sidebar(open=True):
                 gr.Markdown("## Learn more about OrbMol")
                 with gr.Accordion("What is OrbMol?", open=False):
-                    gr.Markdown(
-                        "* OrbMol is a neural network potential for molecular property prediction\n"
-                        "* Built on the Orb-v3 architecture and trained on OMol25 dataset\n"
-                        "* Supports configurable charge and spin multiplicity"
-                    )
-                with gr.Accordion("Model Disclaimers", open=False):
-                    gr.Markdown(
-                        "* Validate results for your use case\n"
-                        "* Consider limitations of the training level of theory"
-                    )
-                with gr.Accordion("Open source packages", open=False):
-                    gr.Markdown(
-                        "* Model code: orbital-materials/orb-models\n"
-                        "* This demo uses ASE and Gradio"
-                    )
-
-        # -------- Tab 2: MD + Visualización 3D --------
+                    gr.Markdown("* Neural network potential for molecules\n* Built on Orb-v3, trained on OMol25\n* Supports charge and spin")
+                with gr.Accordion("Benchmarks", open=False):
+                    gr.Markdown("* <1 kcal/mol error on Wiggle150\n* Accurate protein–ligand binding energies\n* Stable MD on biomolecules >20k atoms")
+                with gr.Accordion("Disclaimers", open=False):
+                    gr.Markdown("* Validate results for your use case\n* Training level of theory may limit accuracy")
+
+        # -------- Tab 2: MD --------
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    with gr.Column(variant="panel"):
-                        xyz_input_md = gr.Textbox(
-                            label="XYZ Coordinates",
-                            lines=12,
-                            placeholder="Paste XYZ here",
-                        )
-                        charge_input_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
-                        spin_input_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
-                        steps_input = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
-                        temp_input = gr.Slider(value=300, minimum=10, maximum=1000, step=10, label="Temperature (K)")
-                        timestep_input = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
-                        run_md_btn = gr.Button("Run MD Simulation", variant="primary")
+                    xyz_input_md = gr.Textbox(label="XYZ Coordinates", lines=12)
+                    charge_input_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
+                    spin_input_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
+                    steps_input = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
+                    temp_input = gr.Slider(value=300, minimum=10, maximum=1000, step=10, label="Temperature (K)")
+                    timestep_input = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
+                    run_md_btn = gr.Button("Run MD Simulation", variant="primary")
 
                 with gr.Column(variant="panel", min_width=520):
-                    gr.Markdown("## MD Visualization")
                     md_status = gr.Textbox(label="MD Status", lines=2, interactive=False)
-                    md_view = gr.HTML()
+                    md_view = gr.HTML() if not HAVE_MOL3D else Molecule3D(label="Trajectory Viewer")
 
             run_md_btn.click(
                 run_md,