Update src/streamlit_app.py

src/streamlit_app.py

@@ -68,6 +68,12 @@ except Exception as e:
 
 os.environ["STREAMLIT_WATCHER_TYPE"] = "none"
 
+# Initialize session state
+if 'optimization_results' not in st.session_state:
+    st.session_state.optimization_results = None
+if 'trajectory_data' not in st.session_state:
+    st.session_state.trajectory_data = None
+
 # YAML data for FairChem reference energies
 ELEMENT_REF_ENERGIES_YAML = """
 oc20_elem_refs:
@@ -1245,37 +1251,53 @@ if atoms is not None:
                         if task == "Cell + Geometry Optimization":
                             results["Final Cell Parameters"] = np.round(calc_atoms.cell.cellpar(), 4).tolist()
                 
+                    st.success("Calculation completed successfully!")
+                    # Store results and trajectory in session state
+                    st.session_state.optimization_results = results
+                    st.session_state.optimized_atoms = calc_atoms.copy()
+                    
+                    # Process trajectory
+                    if os.path.exists(traj_filename):
+                        try:
+                            trajectory = read(traj_filename, index=':')
+                            trajectory_xyz = ""
+                            for i, atoms in enumerate(trajectory):
+                                trajectory_xyz += f"{len(atoms)}\n"
+                                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"
+                            
+                            st.session_state.trajectory_data = {
+                                'xyz_content': trajectory_xyz,
+                                'num_steps': len(trajectory)
+                            }
+                        except Exception as e:
+                            st.warning(f"Could not process trajectory: {e}")
+                        finally:
+                            os.unlink(traj_filename)
+                
+                # Display results if available
+                if st.session_state.optimization_results:
                     st.success("Calculation completed successfully!")
                     st.markdown("### Results")
-                    for key, value in results.items():
+                    for key, value in st.session_state.optimization_results.items():
                         st.write(f"**{key}:** {value}")
                     
-                    if "Optimization" in task and "Final Energy" in results: # Check if opt was successful
+                    if hasattr(st.session_state, 'optimized_atoms'):
                         st.markdown("### Optimized Structure")
-                        # Need get_structure_viz function that takes atoms obj
-                        def get_structure_viz_simple(atoms_obj_viz):
-                            xyz_str_viz = f"{len(atoms_obj_viz)}\nStructure\n"
-                            for atom_viz in atoms_obj_viz:
-                                xyz_str_viz += f"{atom_viz.symbol} {atom_viz.position[0]:.6f} {atom_viz.position[1]:.6f} {atom_viz.position[2]:.6f}\n"
-                            view_viz = py3Dmol.view(width=400, height=400)
-                            view_viz.addModel(xyz_str_viz, "xyz")
-                            view_viz.setStyle({'stick': {}})
-                            if any(atoms_obj_viz.pbc): # Show cell for optimized periodic structures
-                                cell_viz = atoms_obj_viz.get_cell()
-                                if cell_viz is not None and cell_viz.any():
-                                     # Simplified cell drawing for brevity, use get_structure_viz2 if full cell needed
-                                     view_viz.addUnitCell({'box': {'lx':cell_viz.lengths()[0],'ly':cell_viz.lengths()[1],'lz':cell_viz.lengths()[2], 
-                                                                   'hx':cell_viz.cellpar()[3],'hy':cell_viz.cellpar()[4],'hz':cell_viz.cellpar()[5]}})
-
-                            view_viz.zoomTo()
-                            view_viz.setBackgroundColor('white')
-                            return view_viz
-
-                        opt_view = get_structure_viz2(calc_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
+                        
+                        # Visualization
+                        opt_view = get_structure_viz2(st.session_state.optimized_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
                         st.components.v1.html(opt_view._make_html(), width=400, height=400)
                         
+                        # Download optimized structure
                         with tempfile.NamedTemporaryFile(delete=False, suffix=".xyz", mode="w+") as tmp_file_opt:
-                            write(tmp_file_opt.name, calc_atoms, format="xyz")
+                            write(tmp_file_opt.name, st.session_state.optimized_atoms, format="xyz")
                             tmp_filepath_opt = tmp_file_opt.name
                         
                         with open(tmp_filepath_opt, 'r') as file_opt:
@@ -1285,48 +1307,23 @@ if atoms is not None:
                             label="Download Optimized Structure (XYZ)",
                             data=xyz_content_opt,
                             file_name="optimized_structure.xyz",
-                            mime="chemical/x-xyz"
+                            mime="chemical/x-xyz",
+                            key="download_opt_structure"  # Unique key prevents conflicts
                         )
                         os.unlink(tmp_filepath_opt)
-
-                        # Convert trajectory to XYZ for download
-                        from ase.io import read
                         
-                        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"
-                                
-                                st.markdown("### Optimization Trajectory")
-                                st.write(f"Captured {len(trajectory)} optimization steps")
-                                
-                                st.download_button(
-                                    label="Download Optimization Trajectory (XYZ)",
-                                    data=trajectory_xyz,
-                                    file_name="optimization_trajectory.xyz",
-                                    mime="chemical/x-xyz"
-                                )
-                                
-                            except Exception as e:
-                                st.warning(f"Could not process trajectory: {e}")
+                        # Download trajectory
+                        if st.session_state.trajectory_data:
+                            st.markdown("### Optimization Trajectory")
+                            st.write(f"Captured {st.session_state.trajectory_data['num_steps']} optimization steps")
                             
-                            finally:
-                                # Clean up trajectory file
-                                os.unlink(traj_filename)
+                            st.download_button(
+                                label="Download Optimization Trajectory (XYZ)",
+                                data=st.session_state.trajectory_data['xyz_content'],
+                                file_name="optimization_trajectory.xyz",
+                                mime="chemical/x-xyz",
+                                key="download_trajectory"  # Unique key prevents conflicts
+                            )
                 
             except Exception as e:
                 st.error(f"🔴 Calculation error: {str(e)}")