Update app.py

app.py

@@ -9,6 +9,7 @@ from gradio_molecule3d import Molecule3D
 from simulation_scripts_orbmol import load_orbmol_model, run_md_simulation, run_relaxation_simulation
 import hashlib
 
+# ==== Configuración Molecule3D ====
 DEFAULT_MOLECULAR_REPRESENTATIONS = [
     {
         "model": 0,
@@ -38,36 +39,50 @@ DEFAULT_MOLECULAR_SETTINGS = {
     "disableFog": False,
 }
 
+# ==== Conversión a PDB para Molecule3D ====
 def convert_to_pdb_for_viewer(file_path):
+    """Convierte cualquier archivo a PDB para Molecule3D"""
     if not file_path or not os.path.exists(file_path):
         return None
+    
     try:
         atoms = read(file_path)
+        
         cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
         os.makedirs(cache_dir, exist_ok=True)
+        
         pdb_path = os.path.join(cache_dir, f"mol_{hashlib.md5(file_path.encode()).hexdigest()[:12]}.pdb")
+        
         write(pdb_path, atoms, format="proteindatabank")
+        
         return pdb_path
     except Exception as e:
         print(f"Error converting to PDB: {e}")
         return None
 
+# ==== OrbMol SPE ====
 def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
+    """Single Point Energy + fuerzas (OrbMol)"""
     try:
         calc = load_orbmol_model(task_name)
         if not structure_file:
             return "Error: Please upload a structure file", "Error", None
+
         file_path = structure_file
         if not os.path.exists(file_path):
             return f"Error: File not found: {file_path}", "Error", None
         if os.path.getsize(file_path) == 0:
             return f"Error: Empty file: {file_path}", "Error", None
+
         atoms = read(file_path)
+        
         if task_name in ["OMol", "OMol-Direct"]:
             atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
+        
         atoms.calc = calc
         energy = atoms.get_potential_energy()
         forces = atoms.get_forces()
+
         lines = [
             f"Model: {task_name}",
             f"Total Energy: {energy:.6f} eV",
@@ -78,17 +93,22 @@ def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
             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/Å"]
+
         pdb_file = convert_to_pdb_for_viewer(file_path)
+        
         return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file
+        
     except Exception as e:
         import traceback
         traceback.print_exc()
         return f"Error during calculation: {e}", "Error", None
 
+# ==== Wrappers MD y Relax ====
 def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_fs, ensemble):
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
+
         traj_path, log_text, script_text, explanation = run_md_simulation(
             structure_file,
             int(steps),
@@ -101,8 +121,11 @@ def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_f
             int(spin),
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
+
         pdb_file = convert_to_pdb_for_viewer(traj_path)
+        
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
+
     except Exception as e:
         import traceback
         traceback.print_exc()
@@ -112,6 +135,7 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
     try:
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
+
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
             structure_file,
             int(steps),
@@ -122,59 +146,178 @@ def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_ce
             bool(relax_cell),
         )
         status = f"Relaxation finished (<={int(steps)} steps, fmax={float(fmax)} eV/Å)"
+
         pdb_file = convert_to_pdb_for_viewer(traj_path)
+        
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
+
     except Exception as e:
         import traceback
         traceback.print_exc()
         return (f"Error: {e}", None, "", "", "", None)
 
+# ==== UI ====
 with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
     with gr.Tabs():
+        # -------- HOME TAB --------
         with gr.Tab("Home"):
-            gr.Image("logo_color_text.png", show_share_button=False, show_download_button=False, show_label=False, show_fullscreen_button=False)
-            gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
-            gr.Markdown("Welcome to the OrbMol demo! Use the tabs above to access different functionalities:")
-            gr.Markdown("1. **Single Point Energy**: Calculate energies and forces for a given molecular structure.")
-            gr.Markdown("2. **Molecular Dynamics**: Run MD simulations using OrbMol-trained potentials.")
-            gr.Markdown("3. **Relaxation / Optimization**: Optimize molecular structures to their minimum-energy configurations.")
-            gr.Markdown("Supported file formats: `.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf`.")
-            gr.Markdown("## Pretrained Models")
-            gr.Markdown("**OMol** and **OMol-Direct**")
-            gr.Markdown("- **Training dataset**: OMol25 (>100M calculations on small molecules, biomolecules, metal complexes, and electrolytes).")
-            gr.Markdown("- **Level of theory**: wB97M-V/def2-TZVPD with non-local dispersion; solvation treated explicitly.")
-            gr.Markdown("- **Inputs**: total charge & spin multiplicity.")
-            gr.Markdown("- **Applications**: biology, organic chemistry, protein folding, small-molecule drugs, organic liquids, homogeneous catalysis.")
-            gr.Markdown("- **Caveats**: trained only on aperiodic systems → periodic/inorganic cases may not work well.")
-            gr.Markdown("- **Difference**: OMol enforces energy–force consistency; OMol-Direct relaxes this for efficiency.")
-            gr.Markdown("**OMat**")
-            gr.Markdown("- **Training dataset**: OMat24 (>100M inorganic calculations, from Materials Project, Alexandria, and far-from-equilibrium samples).")
-            gr.Markdown("- **Level of theory**: PBE/PBE+U with Materials Project settings; VASP 54 pseudopotentials; no dispersion.")
-            gr.Markdown("- **Inputs**: No support for spin and charge. Spin polarization included but magnetic state cannot be selected.")
-            gr.Markdown("- **Applications**: inorganic discovery, photovoltaics, alloys, superconductors, electronic/optical materials.")
-            gr.Markdown("- **Caveats**: magnetic effects may be incompletely captured.")
-            gr.Markdown("## Technical Foundation")
-            gr.Markdown("All models are based on the **Orb-v3 architecture**, the latest generation of Orb universal interatomic potentials, combining transferability with quantum-level accuracy.")
-            gr.Markdown("## Resources & Support")
-            gr.Markdown("-  [Orb-v3 paper](https://arxiv.org/abs/2504.06231)")
-            gr.Markdown("-  [Orb-Models GitHub repository](https://github.com/orbital-materials/orb-models)")
-            gr.Markdown("-  For issues/questions, please open a GitHub issue or contact the developers.")
+            with gr.Row():
+                # Columna izquierda con acordeones
+                with gr.Column(scale=1):
+                    gr.Markdown("## Learn more about OrbMol")
+                    
+                    with gr.Accordion("What is OrbMol?", open=False):
+                        gr.Markdown("""
+OrbMol is a suite of quantum-accurate machine learning models for molecular predictions. Built on the **Orb-v3 architecture**, OrbMol provides fast and accurate calculations of energies, forces, and molecular properties at the level of advanced quantum chemistry methods.
+
+The models combine the transferability of universal potentials with quantum-level accuracy, making them suitable for a wide range of applications in chemistry, materials science, and drug discovery.
+                        """)
+                    
+                    with gr.Accordion("Available Models", open=False):
+                        gr.Markdown("""
+**OMol** and **OMol-Direct**
+- **Training dataset**: OMol25 (>100M calculations on small molecules, biomolecules, metal complexes, and electrolytes)
+- **Level of theory**: ωB97M-V/def2-TZVPD with non-local dispersion; solvation treated explicitly
+- **Inputs**: total charge & spin multiplicity
+- **Applications**: biology, organic chemistry, protein folding, small-molecule drugs, organic liquids, homogeneous catalysis
+- **Caveats**: trained only on aperiodic systems → periodic/inorganic cases may not work well
+- **Difference**: OMol enforces energy–force consistency; OMol-Direct relaxes this for efficiency
+
+**OMat**
+- **Training dataset**: OMat24 (>100M inorganic calculations, from Materials Project, Alexandria, and far-from-equilibrium samples)
+- **Level of theory**: PBE/PBE+U with Materials Project settings; VASP 54 pseudopotentials; no dispersion
+- **Inputs**: No support for spin and charge. Spin polarization included but magnetic state cannot be selected
+- **Applications**: inorganic discovery, photovoltaics, alloys, superconductors, electronic/optical materials
+- **Caveats**: magnetic effects may be incompletely captured
+                        """)
+                    
+                    with gr.Accordion("Supported File Formats", open=False):
+                        gr.Markdown("""
+OrbMol supports the following molecular structure formats:
+- `.xyz` - XYZ coordinate files
+- `.pdb` - Protein Data Bank format
+- `.cif` - Crystallographic Information File
+- `.traj` - ASE trajectory format
+- `.mol` - MDL Molfile
+- `.sdf` - Structure Data File
+
+All formats are automatically converted internally for processing.
+                        """)
+                    
+                    with gr.Accordion("How to Use", open=False):
+                        gr.Markdown("""
+**Single Point Energy**: Upload a molecular structure and select a model to calculate energies and forces.
+
+**Molecular Dynamics**: Run time-dependent simulations to observe molecular behavior at different temperatures and conditions.
+
+**Relaxation/Optimization**: Find the minimum-energy configuration of your molecular structure.
+
+Each tab provides specific parameters you can adjust to customize your calculations.
+                        """)
+                    
+                    with gr.Accordion("Technical Foundation", open=False):
+                        gr.Markdown("""
+All models are based on the **Orb-v3 architecture**, the latest generation of Orb universal interatomic potentials. 
+
+Key features:
+- Graph neural network architecture
+- Equivariant message passing
+- Multi-task learning across different quantum chemistry methods
+- Billions of training examples across diverse chemical spaces
+- Sub-kcal/mol accuracy on test sets
+                        """)
+                    
+                    with gr.Accordion("Resources & Support", open=False):
+                        gr.Markdown("""
+- [Orb-v3 paper](https://arxiv.org/abs/2504.06231)
+- [Orb-Models GitHub repository](https://github.com/orbital-materials/orb-models)
+- For issues/questions, please open a GitHub issue or contact the developers
+
+**Citation**: If you use OrbMol in your research, please cite the Orb-v3 paper and the relevant dataset papers (OMol25/OMat24).
+                        """)
+                
+                # Columna derecha con contenido principal
+                with gr.Column(scale=2):
+                    gr.Image("logo_color_text.png",
+                             show_share_button=False,
+                             show_download_button=False,
+                             show_label=False,
+                             show_fullscreen_button=False)
+                    
+                    gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
+                    
+                    gr.Markdown("""
+Welcome to the OrbMol demo! This interactive platform allows you to explore the capabilities of our quantum-accurate machine learning models for molecular simulations.
+
+## Quick Start
+
+Use the tabs above to access different functionalities:
+
+1. **Single Point Energy**: Calculate energies and forces for a given molecular structure
+2. **Molecular Dynamics**: Run MD simulations using OrbMol-trained potentials
+3. **Relaxation / Optimization**: Optimize molecular structures to their minimum-energy configurations
+
+Simply upload a molecular structure file in any supported format (`.xyz`, `.pdb`, `.cif`, `.traj`, `.mol`, `.sdf`) and select the appropriate model for your system.
 
+## Model Selection Guide
+
+**Choose OMol/OMol-Direct for:**
+- Organic molecules and biomolecules
+- Drug-like compounds
+- Metal-organic complexes
+- Molecules in solution
+- Systems where you need to specify charge and spin
+
+**Choose OMat for:**
+- Inorganic crystals and materials
+- Periodic systems
+- Bulk materials and alloys
+- Solid-state compounds
+
+Explore the accordions on the left to learn more about each model's capabilities, training data, and limitations.
+                    """)
+                    
+                    gr.Markdown("## Try an Example")
+                    gr.Markdown("""
+To get started quickly, navigate to any of the calculation tabs above and try one of these examples:
+- **Single Point Energy**: Upload a small molecule to see energy and force predictions
+- **Molecular Dynamics**: Run a short simulation at 300K to observe thermal motion
+- **Relaxation**: Optimize a distorted structure to find its equilibrium geometry
+                    """)
+        
+        # -------- SPE -------- (SIN CAMBIOS)
         with gr.Tab("Single Point Energy"):
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
                     gr.Markdown("**Supported formats:** .xyz, .pdb, .cif, .traj, .mol, .sdf")
-                    xyz_input = gr.File(label="Upload Structure File", file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], file_count="single")
-                    task_name_spe = gr.Radio(["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type")
+                    
+                    xyz_input = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_spe = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
                         charge_input = gr.Slider(-10, 10, 0, step=1, label="Charge")
                         spin_input = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
+                    
                     run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
+                    
                 with gr.Column(variant="panel", min_width=500):
                     spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                     spe_status = gr.Textbox(label="Status", interactive=False)
-                    spe_viewer = Molecule3D(label="Input Structure Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS)
+                    
+                    spe_viewer = Molecule3D(
+                        label="Input Structure Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
                     task_name_spe.change(
                         lambda x: (
                             gr.update(visible=x in ["OMol", "OMol-Direct"]),
@@ -183,14 +326,29 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [task_name_spe],
                         [charge_input, spin_input]
                     )
-            run_spe.click(predict_molecule, [xyz_input, task_name_spe, charge_input, spin_input], [spe_out, spe_status, spe_viewer])
+                    
+            run_spe.click(
+                predict_molecule,
+                [xyz_input, task_name_spe, charge_input, spin_input],
+                [spe_out, spe_status, spe_viewer]
+            )
 
+        # -------- MD -------- (SIN CAMBIOS)
         with gr.Tab("Molecular Dynamics"):
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("## Molecular Dynamics Simulation")
-                    xyz_md = gr.File(label="Upload Structure File", file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"], file_count="single")
-                    task_name_md = gr.Radio(["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type")
+                    
+                    xyz_md = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_md = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
                         charge_md = gr.Slider(-10, 10, 0, step=1, label="Charge")
                         spin_md = gr.Slider(1, 11, 1, step=1, label="Spin Multiplicity")
@@ -201,13 +359,21 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         timestep_md = gr.Slider(0.1, 5.0, 1.0, step=0.1, label="Timestep (fs)")
                         ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
                     run_md_btn = gr.Button("Run MD Simulation", variant="primary")
+
                 with gr.Column(variant="panel", min_width=520):
                     md_status = gr.Textbox(label="MD Status", interactive=False)
                     md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
-                    md_viewer = Molecule3D(label="MD Result Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS)
+                    
+                    md_viewer = Molecule3D(
+                        label="MD Result Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
                     md_log = gr.Textbox(label="Log", interactive=False, lines=15)
                     md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
                     md_explain = gr.Markdown()
+                    
                     task_name_md.change(
                         lambda x: (
                             gr.update(visible=x in ["OMol", "OMol-Direct"]),
@@ -216,14 +382,29 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [task_name_md],
                         [charge_md, spin_md]
                     )
-            run_md_btn.click(md_wrapper, [xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md], [md_status, md_traj, md_log, md_script, md_explain, md_viewer])
+                    
+            run_md_btn.click(
+                md_wrapper,
+                [xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
+                [md_status, md_traj, md_log, md_script, md_explain, md_viewer]
+            )
 
+        # -------- Relax -------- (SIN CAMBIOS)
         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")
-                    task_name_rlx = gr.Radio(["OMol", "OMat", "OMol-Direct"], value="OMol", label="Model Type")
+                    
+                    xyz_rlx = gr.File(
+                        label="Upload Structure File",
+                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
+                        file_count="single"
+                    )
+                    task_name_rlx = gr.Radio(
+                        ["OMol", "OMat", "OMol-Direct"],
+                        value="OMol",
+                        label="Model Type"
+                    )
                     with gr.Row():
                         steps_rlx = gr.Slider(1, 2000, 300, step=1, label="Max Steps")
                         fmax_rlx = gr.Slider(0.001, 0.5, 0.05, step=0.001, label="Fmax (eV/Å)")
@@ -232,13 +413,21 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         spin_rlx = gr.Slider(1, 11, 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_viewer = Molecule3D(label="Optimized Structure Viewer", reps=DEFAULT_MOLECULAR_REPRESENTATIONS, config=DEFAULT_MOLECULAR_SETTINGS)
+                    
+                    rlx_viewer = Molecule3D(
+                        label="Optimized Structure Viewer",
+                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
+                        config=DEFAULT_MOLECULAR_SETTINGS
+                    )
+                    
                     rlx_log = gr.Textbox(label="Log", interactive=False, lines=15)
                     rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20)
                     rlx_explain = gr.Markdown()
+
                     task_name_rlx.change(
                         lambda x: (
                             gr.update(visible=x in ["OMol", "OMol-Direct"]),
@@ -247,7 +436,12 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                         [task_name_rlx],
                         [charge_rlx, spin_rlx]
                     )
-            run_rlx_btn.click(relax_wrapper, [xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell], [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer])
+                    
+            run_rlx_btn.click(
+                relax_wrapper,
+                [xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
+                [rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer]
+            )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)