Update src/streamlit_app.py

src/streamlit_app.py

@@ -1293,81 +1293,158 @@ if atoms is not None:
                         show_optimized_structure_download_button()
                         os.unlink(tmp_filepath_opt)
 
-                        # Convert trajectory to XYZ for download
-                        @st.fragment
-                        def show_trajectory():
-                            if os.path.exists(traj_filename):
-                                try:
-                                    from ase.io import read
-                                    from ase.visualize import view
-                                    import py3Dmol
+                        # # Convert trajectory to XYZ for download
+                        # @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")
+                        #             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
+                        #             # 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
+                        #             # 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)}")
+                        #             # 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
+                        #             # 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)
+                        #             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()
+                        #             # 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}")
+                        #         except Exception as e:
+                        #             st.warning(f"Could not process trajectory: {e}")
                             
-                                finally:
-                                    os.unlink(traj_filename)
+                        #         finally:
+                        #             os.unlink(traj_filename)
+
+                        # show_trajectory()
+                        @st.fragment
+                        def show_trajectory_and_controls():
+                            from ase.io import read
+                            import py3Dmol
+                        
+                            if "traj_frames" not in st.session_state:
+                                if os.path.exists(traj_filename):
+                                    try:
+                                        trajectory = read(traj_filename, index=":")
+                                        st.session_state.traj_frames = trajectory
+                                        st.session_state.traj_index = 0
+                                    except Exception as e:
+                                        st.error(f"Error reading trajectory: {e}")
+                                        return
+                                    finally:
+                                        os.unlink(traj_filename)
+                                else:
+                                    st.warning("Trajectory file not found.")
+                                    return
+                        
+                            trajectory = st.session_state.traj_frames
+                            index = st.session_state.traj_index
+                        
+                            st.markdown("### Optimization Trajectory")
+                            st.write(f"Captured {len(trajectory)} optimization steps")
+                        
+                            # 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 index > 0:
+                                    st.session_state.traj_index -= 1
+                            with col3:
+                                if st.button("Next ▶") and index < len(trajectory) - 1:
+                                    st.session_state.traj_index += 1
+                            with col4:
+                                if st.button("Last ⏭"):
+                                    st.session_state.traj_index = len(trajectory) - 1
+                        
+                            # Show current frame
+                            current_atoms = trajectory[st.session_state.traj_index]
+                            st.write(f"Frame {st.session_state.traj_index + 1}/{len(trajectory)}")
+                        
+                            def atoms_to_xyz_string(atoms, step_idx=None):
+                                xyz_str = f"{len(atoms)}\n"
+                                if step_idx is not None:
+                                    xyz_str += f"Step {step_idx}, Energy = {atoms.get_potential_energy():.6f} eV\n"
+                                else:
+                                    xyz_str += f"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, st.session_state.traj_index)
+                        
+                            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 button for entire trajectory
+                            trajectory_xyz = ""
+                            for i, atoms in enumerate(trajectory):
+                                trajectory_xyz += atoms_to_xyz_string(atoms, i)
+                            st.download_button(
+                                label="Download Optimization Trajectory (XYZ)",
+                                data=trajectory_xyz,
+                                file_name="optimization_trajectory.xyz",
+                                mime="chemical/x-xyz"
+                            )
+                        
+                        show_trajectory_and_controls()
 
-                        show_trajectory()
             except Exception as e:
                 st.error(f"🔴 Calculation error: {str(e)}")
                 st.error("Please check the structure, model compatibility, and parameters. For FairChem UMA, ensure the task type (omol, omat etc.) is appropriate for your system (e.g. omol for molecules, omat for materials).")