Update app.py

app.py

@@ -10,12 +10,12 @@ from ase.optimize import LBFGS
 from ase.md.verlet import VelocityVerlet
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.io.trajectory import Trajectory
-
 import py3Dmol
 
 from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
+
 # -----------------------------
 # Global OrbMol model
 # -----------------------------
@@ -26,18 +26,18 @@ def load_orbmol_model():
     global model_calc
     if model_calc is None:
         try:
-            print("Loading OrbMol model...")
             orbff = pretrained.orb_v3_conservative_inf_omat(
                 device="cpu",
                 precision="float32-high"
             )
             model_calc = ORBCalculator(orbff, device="cpu")
-            print("✅ OrbMol model loaded successfully")
+            print("OrbMol model loaded successfully")
         except Exception as e:
-            print(f"❌ Error loading model: {e}")
+            print(f"Error loading OrbMol model: {e}")
             model_calc = None
     return model_calc
 
+
 # -----------------------------
 # Single-point calculation
 # -----------------------------
@@ -45,10 +45,10 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
     try:
         calc = load_orbmol_model()
         if calc is None:
-            return "❌ Error: Could not load OrbMol model", ""
+            return "Error: Could not load OrbMol model", ""
 
         if not xyz_content.strip():
-            return "❌ Error: Please enter XYZ coordinates", ""
+            return "Error: Please enter XYZ coordinates", ""
 
         with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
             f.write(xyz_content)
@@ -61,20 +61,21 @@ def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
         energy = atoms.get_potential_energy()
         forces = atoms.get_forces()
 
-        result = f"🔋 **Total Energy**: {energy:.6f} eV\n\n⚡ **Atomic Forces**:\n"
+        result = f"Total Energy: {energy:.6f} eV\n\nAtomic Forces:\n"
         for i, f in enumerate(forces):
             result += f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å\n"
 
         max_force = np.max(np.linalg.norm(forces, axis=1))
-        result += f"\n📊 **Max Force**: {max_force:.4f} eV/Å"
+        result += f"\nMax Force: {max_force:.4f} eV/Å"
 
         os.unlink(xyz_file)
-        return result, "✅ Calculation completed with OrbMol"
+        return result, "Calculation completed with OrbMol"
     except Exception as e:
-        return f"❌ Error during calculation: {str(e)}", "Error"
+        return f"Error during calculation: {str(e)}", "Error"
+
 
 # -----------------------------
-# Trajectory → HTML for Gradio
+# Trajectory → HTML
 # -----------------------------
 def traj_to_html(traj_file):
     traj = Trajectory(traj_file)
@@ -89,7 +90,8 @@ def traj_to_html(traj_file):
     view.setStyle({"stick": {}})
     view.zoomTo()
     view.animate({"loop": "forward"})
-    return view.render()  # ✅ en Gradio funciona mejor que _make_html()
+    return view.render()
+
 
 # -----------------------------
 # MD simulation
@@ -98,10 +100,10 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
     try:
         calc = load_orbmol_model()
         if calc is None:
-            return "❌ Error: Could not load OrbMol model", ""
+            return "Error: Could not load OrbMol model", ""
 
         if not xyz_content.strip():
-            return "❌ Error: Please enter XYZ coordinates", ""
+            return "Error: Please enter XYZ coordinates", ""
 
         with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
             f.write(xyz_content)
@@ -111,28 +113,27 @@ def run_md(xyz_content, charge=0, spin_multiplicity=1, steps=100, temperature=30
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        # Pre-relax
+        # Pre-relaxation
         opt = LBFGS(atoms)
         opt.run(fmax=0.05, steps=20)
 
         # Initialize velocities
         MaxwellBoltzmannDistribution(atoms, temperature_K=2 * temperature)
 
-        # MD setup
+        # Run MD
         dyn = VelocityVerlet(atoms, timestep=timestep * units.fs)
-
         traj_file = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
         traj = Trajectory(traj_file.name, "w", atoms)
         dyn.attach(traj.write, interval=1)
-
         dyn.run(steps)
 
         html = traj_to_html(traj_file.name)
 
         os.unlink(xyz_file)
-        return f"✅ MD completed: {steps} steps at {temperature} K", html
+        return f"MD completed: {steps} steps at {temperature} K", html
     except Exception as e:
-        return f"❌ Error during MD simulation: {str(e)}", ""
+        return f"Error during MD simulation: {str(e)}", ""
+
 
 # -----------------------------
 # Examples
@@ -142,13 +143,11 @@ examples = [
 Hydrogen molecule
 H 0.0 0.0 0.0
 H 0.0 0.0 0.74""", 0, 1],
-
     ["""3
 Water molecule
 O 0.0000 0.0000 0.0000
 H 0.7571 0.0000 0.5864
 H -0.7571 0.0000 0.5864""", 0, 1],
-
     ["""5
 Methane
 C 0.0000 0.0000 0.0000
@@ -158,23 +157,22 @@ H -0.3630 -0.5133 0.8887
 H -0.3630 -0.5133 -0.8887""", 0, 1]
 ]
 
+
 # -----------------------------
 # Gradio UI
 # -----------------------------
-with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol + MD Demo") as demo:
+with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
+
     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")
-                    xyz_input = gr.Textbox(
-                        label="XYZ Coordinates", lines=12,
-                        placeholder="Paste XYZ here"
-                    )
+                    xyz_input = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here")
                     with gr.Row():
                         charge_input = gr.Slider(value=0, label="Charge", minimum=-10, maximum=10, step=1)
                         spin_input = gr.Slider(value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity")
-                    run_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
+                    predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
             with gr.Column(variant="panel", min_width=500):
                 gr.Markdown("## Results")
                 results_output = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
@@ -182,29 +180,35 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol + MD Demo") as demo:
 
         gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])
 
-        run_btn.click(
+        predict_btn.click(
             predict_molecule,
             inputs=[xyz_input, charge_input, spin_input],
             outputs=[results_output, status_output]
         )
 
     with gr.Tab("Molecular Dynamics"):
-        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")
-        md_status = gr.Textbox(label="MD Status", lines=2)
-        md_view = gr.HTML()
+        with gr.Row():
+            with gr.Column(scale=2):
+                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")
+                md_status = gr.Textbox(label="MD Status", lines=2)
+            with gr.Column(scale=1, variant="panel"):
+                gr.Markdown("## MD Visualization")
+                md_view = gr.HTML()
+
         run_md_btn.click(
             run_md,
             inputs=[xyz_input_md, charge_input_md, spin_input_md, steps_input, temp_input, timestep_input],
             outputs=[md_status, md_view],
         )
 
-print("🚀 Starting OrbMol model loading...")
+
+print("Starting OrbMol model loading...")
 load_orbmol_model()
 
 if __name__ == "__main__":