Update app.py

app.py

@@ -7,23 +7,43 @@ from ase.io.trajectory import Trajectory
 from gradio_molecule3d import Molecule3D
 import hashlib
 
-# ==== Molecule3D viewer preparation ====
+# ==== Molecule3D viewer preparation con debug ====
 def prepare_molecule_for_viewer(traj_path):
     """Convert trajectory to format compatible with Molecule3D viewer"""
     if not traj_path or not os.path.exists(traj_path):
+        print("No trajectory path provided or file doesn't exist")
         return None
     
     try:
         traj = Trajectory(traj_path)
         if len(traj) == 0:
+            print("Empty trajectory")
             return None
         
+        # Debug info
+        print(f"Preparing viewer: {len(traj)} frames, {len(traj[-1])} atoms")
+        
         # Save last frame as PDB for Molecule3D
         temp_pdb = tempfile.NamedTemporaryFile(suffix='.pdb', delete=False)
         write(temp_pdb.name, traj[-1], format='pdb')
-        return temp_pdb.name
+        
+        # Verify PDB was created
+        if os.path.exists(temp_pdb.name):
+            file_size = os.path.getsize(temp_pdb.name)
+            print(f"PDB created: {temp_pdb.name}, size: {file_size} bytes")
+            
+            # Read first few lines to verify content
+            with open(temp_pdb.name, 'r') as f:
+                first_lines = f.read(200)
+                print(f"PDB content preview: {first_lines[:100]}...")
+            
+            return temp_pdb.name
+        else:
+            print("PDB file was not created")
+            return None
+            
     except Exception as e:
-        print(f"Error preparing molecule: {e}")
+        print(f"Error preparing molecule for viewer: {e}")
         return None
 
 # ==== OrbMol SPE ====
@@ -49,10 +69,8 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         if not structure_file:
             return "Error: Please upload a structure file", "Error"
 
-        # structure_file es directamente el path del archivo en Gradio
         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"
         
@@ -63,8 +81,8 @@ def predict_molecule(structure_file, charge=0, spin_multiplicity=1):
         atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
         atoms.calc = calc
 
-        energy = atoms.get_potential_energy()  # eV
-        forces = atoms.get_forces()            # eV/Å
+        energy = atoms.get_potential_energy()
+        forces = atoms.get_forces()
 
         lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
         for i, fc in enumerate(forces):
@@ -82,14 +100,14 @@ from simulation_scripts_orbmol import (
     run_relaxation_simulation,
 )
 
-# ==== Wrappers: usan archivos subidos ====
+# ==== Wrappers con debug ====
 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)
 
-        # structure_file es directamente el path del archivo
         file_path = structure_file
+        print(f"MD Wrapper: Processing {file_path}")
 
         traj_path, log_text, script_text, explanation = run_md_simulation(
             file_path,
@@ -103,12 +121,16 @@ def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble
         )
         status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
 
+        print(f"MD completed, trajectory: {traj_path}")
+        
         # Preparar archivo PDB para el visor 3D
         pdb_file = prepare_molecule_for_viewer(traj_path)
+        print(f"PDB file for viewer: {pdb_file}")
         
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
+        print(f"MD Wrapper Error: {e}")
         return (f"Error: {e}", None, "", "", "", None)
 
 def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
@@ -116,8 +138,8 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
         if not structure_file:
             return ("Error: Please upload a structure file", None, "", "", "", None)
 
-        # structure_file es directamente el path del archivo
         file_path = structure_file
+        print(f"Relax Wrapper: Processing {file_path}")
 
         traj_path, log_text, script_text, explanation = run_relaxation_simulation(
             file_path,
@@ -129,12 +151,16 @@ def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
         )
         status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
 
+        print(f"Relaxation completed, trajectory: {traj_path}")
+        
         # Preparar archivo PDB para el visor 3D
         pdb_file = prepare_molecule_for_viewer(traj_path)
+        print(f"PDB file for viewer: {pdb_file}")
         
         return (status, traj_path, log_text, script_text, explanation, pdb_file)
 
     except Exception as e:
+        print(f"Relax Wrapper Error: {e}")
         return (f"Error: {e}", None, "", "", "", None)
 
 # ==== UI ====
@@ -191,7 +217,38 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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="3D Structure Viewer")
+                    
+                    # Molecule3D viewer con configuraciones adicionales
+                    md_viewer = Molecule3D(
+                        label="3D Structure Viewer",
+                        reps=[
+                            {
+                                "model": 0,
+                                "chain": "",
+                                "resname": "",
+                                "style": "stick",
+                                "color": "whiteCarbon",
+                                "residue_range": "",
+                                "around": 0,
+                                "byres": False,
+                                "visible": True,
+                                "opacity": 1.0
+                            },
+                            {
+                                "model": 0,
+                                "chain": "",
+                                "resname": "",
+                                "style": "sphere",
+                                "color": "whiteCarbon",
+                                "residue_range": "",
+                                "around": 0,
+                                "byres": False,
+                                "visible": True,
+                                "opacity": 0.7
+                            }
+                        ]
+                    )
+                    
                     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()
@@ -226,7 +283,38 @@ with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
                 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="3D Structure Viewer")
+                    
+                    # Molecule3D viewer con configuraciones adicionales
+                    rlx_viewer = Molecule3D(
+                        label="3D Structure Viewer",
+                        reps=[
+                            {
+                                "model": 0,
+                                "chain": "",
+                                "resname": "",
+                                "style": "stick",
+                                "color": "whiteCarbon",
+                                "residue_range": "",
+                                "around": 0,
+                                "byres": False,
+                                "visible": True,
+                                "opacity": 1.0
+                            },
+                            {
+                                "model": 0,
+                                "chain": "",
+                                "resname": "",
+                                "style": "sphere",
+                                "color": "whiteCarbon",
+                                "residue_range": "",
+                                "around": 0,
+                                "byres": False,
+                                "visible": True,
+                                "opacity": 0.7
+                            }
+                        ]
+                    )
+                    
                     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()