Update src/streamlit_app.py

src/streamlit_app.py

@@ -1099,6 +1099,9 @@ if atoms is not None:
         run_calculation = st.button("Run Calculation", type="primary")
         
         if run_calculation:
+            # Delete all the items in Session state
+            for key in st.session_state.keys():
+                del st.session_state[key]
             results = {}
             #global table_placeholder # Ensure they are accessible
             opt_log = [] # Reset log for each run
@@ -1382,8 +1385,8 @@ if atoms is not None:
                                     except Exception as e:
                                         st.error(f"Error reading trajectory: {e}")
                                         return
-                                    finally:
-                                        os.unlink(traj_filename)
+                                    # finally:
+                                    #     os.unlink(traj_filename)
                                 else:
                                     st.warning("Trajectory file not found.")
                                     return
@@ -1423,14 +1426,16 @@ if atoms is not None:
                                     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)
+                            # 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)
+                            # 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)
+                            traj_view = get_structure_viz2(current_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
+                            st.components.v1.html(traj_view._make_html(), width=400, height=400)
                         
                             # Download button for entire trajectory
                             trajectory_xyz = ""