Update src/streamlit_app.py

src/streamlit_app.py

@@ -1298,29 +1298,61 @@ if atoms is not None:
                         
                         if os.path.exists(traj_filename):
                             try:
-                                # Read trajectory and convert to XYZ
-                                trajectory = read(traj_filename, index=':')  # Read all frames
-                                
-                                # Create XYZ content
-                                trajectory_xyz = ""
-                                for i, atoms in enumerate(trajectory):
-                                    trajectory_xyz += f"{len(atoms)}\n"
-                                    # Try to get energy if available
-                                    try:
-                                        energy = atoms.get_potential_energy()
-                                        trajectory_xyz += f"Step {i}: Energy = {energy:.6f} eV\n"
-                                    except:
-                                        trajectory_xyz += f"Step {i}: Optimization trajectory\n"
-                                    
-                                    for atom in atoms:
-                                        trajectory_xyz += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
-                                
+                                from ase.io import read
+                                from ase.visualize import view
+                                import py3Dmol
+                        
+                                trajectory = read(traj_filename, index=':')
                                 st.markdown("### Optimization Trajectory")
                                 st.write(f"Captured {len(trajectory)} optimization steps")
+                        
+                                # Store the trajectory in session state
+                                if "traj_frames" not in st.session_state:
+                                    st.session_state.traj_frames = trajectory
+                                    st.session_state.traj_index = 0
+                        
+                                # Navigation Buttons
+                                col1, col2, col3, col4 = st.columns(4)
+                                with col1:
+                                    if st.button("⏮ First"):
+                                        st.session_state.traj_index = 0
+                                with col2:
+                                    if st.button("◀ Previous") and st.session_state.traj_index > 0:
+                                        st.session_state.traj_index -= 1
+                                with col3:
+                                    if st.button("Next ▶") and st.session_state.traj_index < len(st.session_state.traj_frames) - 1:
+                                        st.session_state.traj_index += 1
+                                with col4:
+                                    if st.button("Last ⏭"):
+                                        st.session_state.traj_index = len(st.session_state.traj_frames) - 1
+                        
+                                # Show current frame
+                                current_atoms = st.session_state.traj_frames[st.session_state.traj_index]
+                                st.write(f"Frame {st.session_state.traj_index + 1}/{len(st.session_state.traj_frames)}")
+                        
+                                # Convert to xyz string for py3Dmol
+                                def atoms_to_xyz_string(atoms):
+                                    xyz_str = f"{len(atoms)}\nStep {st.session_state.traj_index}, Energy = {atoms.get_potential_energy():.6f} eV\n"
+                                    for atom in atoms:
+                                        xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
+                                    return xyz_str
+                        
+                                xyz_str = atoms_to_xyz_string(current_atoms)
+                                view = py3Dmol.view(width=400, height=400)
+                                view.addModel(xyz_str, "xyz")
+                                view.setStyle({'stick': {}})
+                                view.zoomTo()
+                                view.setBackgroundColor("white")
+                                st.components.v1.html(view._make_html(), height=400, width=400)
+                        
+                                # Download entire trajectory
                                 @st.fragment
                                 def show_trajectory_download_button():
-                                    # st.button("Release the balloons", help="Fragment rerun")
-                                    # st.balloons()
+                                    trajectory_xyz = ""
+                                    for i, atoms in enumerate(st.session_state.traj_frames):
+                                        trajectory_xyz += f"{len(atoms)}\nStep {i}, Energy = {atoms.get_potential_energy():.6f} eV\n"
+                                        for atom in atoms:
+                                            trajectory_xyz += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
                                     st.download_button(
                                         label="Download Optimization Trajectory (XYZ)",
                                         data=trajectory_xyz,
@@ -1328,14 +1360,13 @@ if atoms is not None:
                                         mime="chemical/x-xyz"
                                     )
                                 show_trajectory_download_button()
-                                
-                                
+                        
                             except Exception as e:
                                 st.warning(f"Could not process trajectory: {e}")
-                            
+                        
                             finally:
-                                # Clean up trajectory file
                                 os.unlink(traj_filename)
+
                 
             except Exception as e:
                 st.error(f"🔴 Calculation error: {str(e)}")