Update src/streamlit_app.py

src/streamlit_app.py

@@ -68,12 +68,6 @@ 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:
@@ -1251,53 +1245,37 @@ 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 st.session_state.optimization_results.items():
+                    for key, value in results.items():
                         st.write(f"**{key}:** {value}")
                     
-                    if hasattr(st.session_state, 'optimized_atoms'):
+                    if "Optimization" in task and "Final Energy" in results: # Check if opt was successful
                         st.markdown("### Optimized Structure")
-                        
-                        # Visualization
-                        opt_view = get_structure_viz2(st.session_state.optimized_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
+                        # 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)
                         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, st.session_state.optimized_atoms, format="xyz")
+                            write(tmp_file_opt.name, calc_atoms, format="xyz")
                             tmp_filepath_opt = tmp_file_opt.name
                         
                         with open(tmp_filepath_opt, 'r') as file_opt:
@@ -1307,23 +1285,48 @@ if atoms is not None:
                             label="Download Optimized Structure (XYZ)",
                             data=xyz_content_opt,
                             file_name="optimized_structure.xyz",
-                            mime="chemical/x-xyz",
-                            key="download_opt_structure"  # Unique key prevents conflicts
+                            mime="chemical/x-xyz"
                         )
                         os.unlink(tmp_filepath_opt)
+
+                        # Convert trajectory to XYZ for download
+                        from ase.io import read
                         
-                        # 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")
+                        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}")
                             
-                            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
-                            )
+                            finally:
+                                # Clean up trajectory file
+                                os.unlink(traj_filename)
                 
             except Exception as e:
                 st.error(f"🔴 Calculation error: {str(e)}")