Update src/streamlit_app.py

src/streamlit_app.py

@@ -1294,78 +1294,78 @@ if atoms is not None:
                         os.unlink(tmp_filepath_opt)
 
                         # Convert trajectory to XYZ for download
-                        from ase.io import read
-                        
-                        if os.path.exists(traj_filename):
-                            try:
-                                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.fragment
+                        def show_trajectory():
+                            if os.path.exists(traj_filename):
+                                try:
+                                    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
-                                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():
-                                    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"
+                            
+                                    # 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:
-                                            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,
-                                        file_name="optimization_trajectory.xyz",
-                                        mime="chemical/x-xyz"
-                                    )
-                                show_trajectory_download_button()
-                        
-                            except Exception as e:
-                                st.warning(f"Could not process trajectory: {e}")
-                        
-                            finally:
-                                os.unlink(traj_filename)
+                                            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():
+                                        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,
+                                            file_name="optimization_trajectory.xyz",
+                                            mime="chemical/x-xyz"
+                                        )
+                                    show_trajectory_download_button()
+                            
+                                except Exception as e:
+                                    st.warning(f"Could not process trajectory: {e}")
+                            
+                                finally:
+                                    os.unlink(traj_filename)
 
                 
             except Exception as e: