Update src/streamlit_app.py

src/streamlit_app.py

@@ -1219,13 +1219,14 @@ if atoms is not None:
 
                     elif "Optimization" in task: # Handles both Geometry and Cell+Geometry Opt
                         opt_atoms_obj = FrechetCellFilter(calc_atoms) if task == "Cell + Geometry Optimization" else calc_atoms
-                        
+                        # Create temporary trajectory file
+                        traj_filename = tempfile.NamedTemporaryFile(delete=False, suffix=".traj").name
                         if optimizer_type == "BFGS":
-                            opt = BFGS(opt_atoms_obj)
+                            opt = BFGS(opt_atoms_obj, trajectory=traj_filename)
                         elif optimizer_type == "LBFGS":
-                            opt = LBFGS(opt_atoms_obj)
+                            opt = LBFGS(opt_atoms_obj, trajectory=traj_filename)
                         else: # FIRE
-                            opt = FIRE(opt_atoms_obj)
+                            opt = FIRE(opt_atoms_obj, trajectory=traj_filename)
                         
                         # opt.attach(streamlit_log, interval=1) # Removed lambda for simplicity if streamlit_log is defined correctly
                         opt.attach(lambda: streamlit_log(opt), interval=1)
@@ -1287,6 +1288,45 @@ if atoms is not None:
                             mime="chemical/x-xyz"
                         )
                         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}")
+                            
+                            finally:
+                                # Clean up trajectory file
+                                os.unlink(traj_filename)
                 
             except Exception as e:
                 st.error(f"🔴 Calculation error: {str(e)}")