add materials project ID support

src/streamlit_app.py

@@ -4,13 +4,16 @@ import io
 import tempfile
 import torch
 # FOR CPU only mode
-# torch._dynamo.config.suppress_errors = True
+#torch._dynamo.config.suppress_errors = True
 # Or disable compilation entirely
 # torch.backends.cudnn.enabled = False
 import numpy as np
 from ase import Atoms
 from ase.io import read, write
 from ase.optimize import BFGS, LBFGS, FIRE
+from ase.optimize.basin import BasinHopping
+from ase.optimize.minimahopping import MinimaHopping
+from ase.units import kB
 from ase.constraints import FixAtoms
 from ase.filters import FrechetCellFilter
 from ase.visualize import view
@@ -26,6 +29,9 @@ import subprocess
 import sys
 import pkg_resources
 from ase.vibrations import Vibrations
+from mp_api.client import MPRester
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+from pymatgen.io.ase import AseAtomsAdaptor
 import matplotlib.pyplot as plt
 mattersim_available = True
 if mattersim_available:
@@ -894,6 +900,7 @@ def check_atom_limit(atoms_obj, selected_model):
 MACE_MODELS = {
     "MACE MPA Medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model",
     "MACE OMAT Medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-medium.model",
+    "MACE OMAT Small": "https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-small.model",
     "MACE MATPES r2SCAN Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-r2scan-omat-ft.model",
     "MACE MATPES PBE Medium": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-pbe-omat-ft.model",
     "MACE MP 0a Small": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model",
@@ -952,59 +959,273 @@ def get_fairchem_model(selected_model_name, model_path_or_name, device, selected
         calc = FAIRChemCalculator(predictor, task_name="omol")
     return calc
 
+# st.sidebar.markdown("## Input Options")
+# input_method = st.sidebar.radio("Choose Input Method:", 
+#                                 ["Select Example", "Upload File", "Paste Content", "Materials Project ID"])
+# atoms = None
+
+# if input_method == "Upload File":
+#     uploaded_file = st.sidebar.file_uploader("Upload structure file", type=["xyz", "cif", "POSCAR", "mol", "tmol", "vasp", "sdf", "CONTCAR"])
+#     if uploaded_file:
+#         with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
+#             tmp_file.write(uploaded_file.getvalue())
+#             tmp_filepath = tmp_file.name
+#         try:
+#             atoms = read(tmp_filepath)
+#             st.sidebar.success(f"Successfully loaded structure with {len(atoms)} atoms!")
+#         except Exception as e:
+#             st.sidebar.error(f"Error loading file: {str(e)}")
+#         finally:
+#             if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath):
+#                  os.unlink(tmp_filepath)
+
+# elif input_method == "Select Example":
+#     example_name = st.sidebar.selectbox("Select Example Structure:", list(SAMPLE_STRUCTURES.keys()))
+#     if example_name:
+#         file_path = os.path.join(SAMPLE_DIR, SAMPLE_STRUCTURES[example_name])
+#         try:
+#             atoms = read(file_path)
+#             st.sidebar.success(f"Loaded {example_name} with {len(atoms)} atoms!")
+#         except Exception as e:
+#             st.sidebar.error(f"Error loading example: {str(e)}")
+
+# elif input_method == "Paste Content":
+#     file_format = st.sidebar.selectbox("File Format:", ["XYZ", "CIF", "extXYZ", "POSCAR (VASP)", "Turbomole", "MOL"])
+#     content = st.sidebar.text_area("Paste file content here:", height=200)
+#     if content:
+#         try:
+#             suffix_map = {"XYZ": ".xyz", "CIF": ".cif", "extXYZ": ".extxyz", "POSCAR (VASP)": ".vasp", "Turbomole": ".tmol", "MOL": ".mol"}
+#             suffix = suffix_map.get(file_format, ".xyz")
+#             with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
+#                 tmp_file.write(content.encode())
+#                 tmp_filepath = tmp_file.name
+#             atoms = read(tmp_filepath)
+#             st.sidebar.success(f"Successfully parsed structure with {len(atoms)} atoms!")
+#         except Exception as e:
+#             st.sidebar.error(f"Error parsing content: {str(e)}")
+#         finally:
+#             if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath):
+#                 os.unlink(tmp_filepath)
+# elif input_method == "Materials Project ID":
+#     # st.sidebar.info("Requires an API Key from [Materials Project](https://next-gen.materialsproject.org/api)")
+    
+#     # 1. Get API Key (You can also set this as an env variable)
+#     # mp_api_key = st.sidebar.text_input("MP API Key", type="password")
+#     mp_api_key = os.getenv("MP_API_KEY")
+#     # 2. Get Material ID
+#     material_id = st.sidebar.text_input("Enter Material ID:", value="mp-149")
+    
+#     # 3. Choose Cell Type
+#     cell_type = st.sidebar.radio("Unit Cell Type:", ['Primitive Cell', 'Conventional Unit Cell'])
+    
+#     # 1. Initialize the session state variable if it doesn't exist
+#     if "atoms" not in st.session_state:
+#         st.session_state.atoms = None
+
+#     if st.sidebar.button("Fetch Structure"):
+#         if not mp_api_key:
+#             st.sidebar.error("Please enter your Materials Project API Key.")
+#         elif not material_id:
+#             st.sidebar.error("Please enter a Material ID.")
+#         else:
+#             try:
+#                 with st.spinner(f"Fetching {material_id}..."):
+#                     # Connect to MP
+#                     with MPRester(mp_api_key) as mpr:
+#                         # Get Pymatgen Structure
+#                         pmg_structure = mpr.get_structure_by_material_id(material_id)
+                        
+#                         # Handle Symmetry (Primitive vs Conventional)
+#                         analyzer = SpacegroupAnalyzer(pmg_structure)
+                        
+#                         if cell_type == 'Conventional Unit Cell':
+#                             final_structure = analyzer.get_conventional_standard_structure()
+#                         else:
+#                             final_structure = analyzer.get_primitive_standard_structure()
+                            
+#                         # Convert Pymatgen Structure -> ASE Atoms
+#                         # 2. Store the result in session_state instead of a local variable
+#                         st.session_state.atoms = AseAtomsAdaptor.get_atoms(final_structure)
+                        
+#                         st.sidebar.success(f"Loaded {material_id} ({cell_type}) with {len(st.session_state.atoms)} atoms.")
+                        
+#             except Exception as e:
+#                 st.sidebar.error(f"Error fetching data: {str(e)}")
+#                 st.session_state.atoms = None
+
+#     # 3. Retrieve the data from session_state for use in the rest of your app
+#     atoms = st.session_state.atoms
+
+#     # --- Rest of your processing code ---
+#     # if atoms is not None:
+#     #     st.write(f"System is ready. Formula: {atoms.get_chemical_formula()}")
+#     # else:
+#     #     st.info("Please fetch a structure to begin.")
+# if atoms is not None:
+#     if not hasattr(atoms, 'info'):
+#         atoms.info = {}
+#     atoms.info["charge"] = atoms.info.get("charge", 0) # Default charge
+#     atoms.info["spin"] = atoms.info.get("spin", 1) # Default spin (usually 2S for ASE, model might want 2S+1)
+
+# --- INITIALIZATION (Must be run first) ---
+if "atoms" not in st.session_state:
+    st.session_state.atoms = None
+
+# Reset atoms state if input method changes, to prevent using old data
+# Use a key to track the currently active input method
+if 'current_input_method' not in st.session_state:
+    st.session_state.current_input_method = "Select Example"
+
 st.sidebar.markdown("## Input Options")
-input_method = st.sidebar.radio("Choose Input Method:", ["Select Example", "Upload File", "Paste Content"])
-atoms = None
+input_method = st.sidebar.radio("Choose Input Method:", 
+                                ["Select Example", "Upload File", "Paste Content", "Materials Project ID"])
 
+# If the input method changes, clear the loaded structure
+if input_method != st.session_state.current_input_method:
+    st.session_state.atoms = None
+    st.session_state.current_input_method = input_method
+
+# --- UPLOAD FILE ---
 if input_method == "Upload File":
     uploaded_file = st.sidebar.file_uploader("Upload structure file", type=["xyz", "cif", "POSCAR", "mol", "tmol", "vasp", "sdf", "CONTCAR"])
+    
+    # Load immediately upon file upload/change (no button needed)
     if uploaded_file:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
-            tmp_file.write(uploaded_file.getvalue())
-            tmp_filepath = tmp_file.name
         try:
-            atoms = read(tmp_filepath)
-            st.sidebar.success(f"Successfully loaded structure with {len(atoms)} atoms!")
+            # Check if this file content has already been loaded to prevent redundant temp file operations
+            if 'uploaded_file_hash' not in st.session_state or st.session_state.uploaded_file_hash != uploaded_file.name: 
+                
+                # Use tempfile to handle the uploaded file content
+                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
+                    tmp_file.write(uploaded_file.getvalue())
+                    tmp_filepath = tmp_file.name
+                    
+                atoms_to_store = read(tmp_filepath)
+                st.session_state.atoms = atoms_to_store
+                st.session_state.uploaded_file_hash = uploaded_file.name # Track the loaded file
+                st.sidebar.success(f"Successfully loaded structure with {len(atoms_to_store)} atoms!")
+            
         except Exception as e:
             st.sidebar.error(f"Error loading file: {str(e)}")
+            st.session_state.atoms = None
+            st.session_state.uploaded_file_hash = None # Clear hash on failure
         finally:
+            # Clean up the temporary file
             if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath):
                  os.unlink(tmp_filepath)
+    else:
+        # Clear structure if file uploader is empty
+        st.session_state.atoms = None
 
+# --- SELECT EXAMPLE ---
 elif input_method == "Select Example":
+    # Load immediately upon selection change (no button needed)
     example_name = st.sidebar.selectbox("Select Example Structure:", list(SAMPLE_STRUCTURES.keys()))
-    if example_name:
+    
+    # Only load if a valid example is selected and it's different from the current state
+    if example_name and (st.session_state.atoms is None or st.session_state.atoms.info.get('source_name') != example_name):
         file_path = os.path.join(SAMPLE_DIR, SAMPLE_STRUCTURES[example_name])
         try:
-            atoms = read(file_path)
-            st.sidebar.success(f"Loaded {example_name} with {len(atoms)} atoms!")
+            atoms_to_store = read(file_path)
+            atoms_to_store.info['source_name'] = example_name # Add a tag for tracking
+            st.session_state.atoms = atoms_to_store
+            st.sidebar.success(f"Loaded {example_name} with {len(atoms_to_store)} atoms!")
         except Exception as e:
             st.sidebar.error(f"Error loading example: {str(e)}")
+            st.session_state.atoms = None
 
+# --- PASTE CONTENT ---
 elif input_method == "Paste Content":
     file_format = st.sidebar.selectbox("File Format:", ["XYZ", "CIF", "extXYZ", "POSCAR (VASP)", "Turbomole", "MOL"])
-    content = st.sidebar.text_area("Paste file content here:", height=200)
+    content = st.sidebar.text_area("Paste file content here:", height=200, key="paste_content_input")
+    
+    # Load immediately upon content change (no button needed)
+    # Check if content is present and is different from the last successfully parsed content
     if content:
-        try:
-            suffix_map = {"XYZ": ".xyz", "CIF": ".cif", "extXYZ": ".extxyz", "POSCAR (VASP)": ".vasp", "Turbomole": ".tmol", "MOL": ".mol"}
-            suffix = suffix_map.get(file_format, ".xyz")
-            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
-                tmp_file.write(content.encode())
-                tmp_filepath = tmp_file.name
-            atoms = read(tmp_filepath)
-            st.sidebar.success(f"Successfully parsed structure with {len(atoms)} atoms!")
-        except Exception as e:
-            st.sidebar.error(f"Error parsing content: {str(e)}")
-        finally:
-            if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath):
-                os.unlink(tmp_filepath)
+        # Simple check to avoid parsing on every single character change
+        if 'last_parsed_content' not in st.session_state or st.session_state.last_parsed_content != content:
+            
+            try:
+                suffix_map = {"XYZ": ".xyz", "CIF": ".cif", "extXYZ": ".extxyz", "POSCAR (VASP)": ".vasp", "Turbomole": ".tmol", "MOL": ".mol"}
+                suffix = suffix_map.get(file_format, ".xyz")
+                
+                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
+                    tmp_file.write(content.encode())
+                    tmp_filepath = tmp_file.name
+                    
+                atoms_to_store = read(tmp_filepath)
+                st.session_state.atoms = atoms_to_store
+                st.session_state.last_parsed_content = content # Track the parsed content
+                st.sidebar.success(f"Successfully parsed structure with {len(atoms_to_store)} atoms!")
+            except Exception as e:
+                st.sidebar.error(f"Error parsing content: {str(e)}")
+                st.session_state.atoms = None
+                st.session_state.last_parsed_content = None
+            finally:
+                if 'tmp_filepath' in locals() and os.path.exists(tmp_filepath):
+                    os.unlink(tmp_filepath)
+    else:
+        # Clear structure if text area is empty
+        st.session_state.atoms = None
+
+
+# --- MATERIALS PROJECT ID ---
+elif input_method == "Materials Project ID":
+    mp_api_key = os.getenv("MP_API_KEY")
+    material_id = st.sidebar.text_input("Enter Material ID:", value="mp-149", key="mp_id_input")
+    cell_type = st.sidebar.radio("Unit Cell Type:", ['Primitive Cell', 'Conventional Unit Cell'], key="cell_type_radio")
+    
+    # Reactive Loading (No button needed)
+    # Check for valid inputs and if the current material_id/cell_type is different from the loaded one
+    if mp_api_key and material_id: 
+        
+        # Simple tracking to avoid API call if nothing has changed
+        current_mp_key = f"{material_id}_{cell_type}"
+        if 'last_fetched_mp_key' not in st.session_state or st.session_state.last_fetched_mp_key != current_mp_key:
+            
+            try:
+                with st.spinner(f"Fetching {material_id}..."):
+                    with MPRester(mp_api_key) as mpr:
+                        pmg_structure = mpr.get_structure_by_material_id(material_id)
+                        analyzer = SpacegroupAnalyzer(pmg_structure)
+                        
+                        if cell_type == 'Conventional Unit Cell':
+                            final_structure = analyzer.get_conventional_standard_structure()
+                        else:
+                            final_structure = analyzer.get_primitive_standard_structure()
+                            
+                        atoms_to_store = AseAtomsAdaptor.get_atoms(final_structure)
+                        st.session_state.atoms = atoms_to_store
+                        st.session_state.last_fetched_mp_key = current_mp_key # Update tracking key
+                        st.sidebar.success(f"Loaded {material_id} ({cell_type}) with {len(st.session_state.atoms)} atoms.")
+                        
+            except Exception as e:
+                st.sidebar.error(f"Error fetching data: {str(e)}")
+                st.session_state.atoms = None
+                st.session_state.last_fetched_mp_key = None # Clear key on failure
+    
+    # Handle error messages when inputs are missing
+    elif not mp_api_key:
+        st.sidebar.error("Please set your Materials Project API Key (MP_API_KEY environment variable).")
+    elif not material_id:
+        st.sidebar.error("Please enter a Material ID.")
+
+# ----------------------------------------------------
+# --- FINAL STRUCTURE RETRIEVAL (The persistent structure) ---
+# ----------------------------------------------------
+# This is the single source of truth for the rest of your app
+atoms = st.session_state.atoms
 
 if atoms is not None:
     if not hasattr(atoms, 'info'):
         atoms.info = {}
     atoms.info["charge"] = atoms.info.get("charge", 0) # Default charge
     atoms.info["spin"] = atoms.info.get("spin", 1) # Default spin (usually 2S for ASE, model might want 2S+1)
-
+    
+    # Display confirmation in the main area (optional, helps the user confirm what's loaded)
+    st.markdown(f"**Loaded Structure:** {atoms.get_chemical_formula()} ({len(atoms)} atoms)")
+    
+    # You can now add your processing buttons here, which will consistently use the 'atoms' variable.
 
 st.sidebar.markdown("## Model Selection")
 if mattersim_available:
@@ -1044,7 +1265,6 @@ if model_type == "FairChem":
         else:
             atoms.info["charge"] = 0
             atoms.info["spin"] = 1 # FairChem expects multiplicity
-            
 if model_type == "ORB":
     selected_model = st.sidebar.selectbox("Select ORB Model:", list(ORB_MODELS.keys()))
     model_path = ORB_MODELS[selected_model]
@@ -1076,19 +1296,43 @@ st.sidebar.markdown("## Task Selection")
 task = st.sidebar.selectbox("Select Calculation Task:", 
                            ["Energy Calculation", 
                             "Energy + Forces Calculation", 
-                            "Atomization/Cohesive Energy", # New Task Added
+                            "Atomization/Cohesive Energy", 
                             "Geometry Optimization", 
                             "Cell + Geometry Optimization",
+                            #"Global Optimization",
                             "Vibrational Mode Analysis",
                             #"Phonons"
-                           ])
+                            ])
 
 if "Optimization" in task:
+    # st.sidebar.markdown("### Optimization Parameters")
+    # max_steps = st.sidebar.slider("Maximum Steps:", min_value=10, max_value=200, value=50, step=1) # Increased max_steps
+    # fmax = st.sidebar.slider("Convergence Threshold (eV/Å):", min_value=0.001, max_value=0.1, value=0.01, step=0.001, format="%.3f") # Adjusted default fmax
+    # optimizer_type = st.sidebar.selectbox("Optimizer:", ["BFGS", "LBFGS", "FIRE"], index=1) # Renamed to optimizer_type
     st.sidebar.markdown("### Optimization Parameters")
-    max_steps = st.sidebar.slider("Maximum Steps:", min_value=10, max_value=200, value=50, step=1) # Increased max_steps
-    fmax = st.sidebar.slider("Convergence Threshold (eV/Å):", min_value=0.001, max_value=0.1, value=0.01, step=0.001, format="%.3f") # Adjusted default fmax
-    optimizer_type = st.sidebar.selectbox("Optimizer:", ["BFGS", "LBFGS", "FIRE"], index=1) # Renamed to optimizer_type
+    
+    # 1. Configuration for GLOBAL Optimization
+    if task == "Global Optimization":
+        global_method = st.sidebar.selectbox("Method:", ["Basin Hopping", "Minima Hopping"])
+        
+        # Common parameters
+        temperature_K = st.sidebar.number_input("Temperature (K):", min_value=10.0, max_value=2000.0, value=300.0, step=10.0)
+        global_steps = st.sidebar.number_input("Search Steps:", min_value=10, max_value=500, value=50, step=10)
+        # Basin Hopping specific
+        if global_method == "Basin Hopping":
+            dr_amp = st.sidebar.number_input("Displacement Amplitude (Å):", min_value=0.1, max_value=2.0, value=0.7, step=0.1)
+            fmax_local = st.sidebar.number_input("Local Relaxation Threshold (eV/Å):", value=0.05, format="%.3f")
+            
+        # Minima Hopping specific
+        elif global_method == "Minima Hopping":
+            st.sidebar.caption("Minima Hopping automates threshold adjustments to escape local minima.")
+            fmax_local = st.sidebar.number_input("Local Relaxation Threshold (eV/Å):", value=0.05, format="%.3f")
 
+    # 2. Configuration for LOCAL/CELL Optimization
+    else:
+        max_steps = st.sidebar.slider("Maximum Steps:", min_value=10, max_value=200, value=50, step=1)
+        fmax = st.sidebar.slider("Convergence Threshold (eV/Å):", min_value=0.001, max_value=0.1, value=0.01, step=0.001, format="%.3f")
+        optimizer_type = st.sidebar.selectbox("Optimizer:", ["BFGS", "LBFGS", "FIRE"], index=1)
 
 if "Vibration" in task:
     st.write("### Thermodynamic Quantities (Molecule Only)")
@@ -1131,7 +1375,7 @@ if atoms is not None:
         st.markdown("### Selected Task")
         st.write(f"**Task:** {task}")
         
-        if "Optimization" in task:
+        if "Geometry Optimization" in task:
             st.write(f"**Max Steps:** {max_steps}")
             st.write(f"**Convergence Threshold:** {fmax} eV/Å")
             st.write(f"**Optimizer:** {optimizer_type}")
@@ -1270,7 +1514,7 @@ if atoms is not None:
                                 results["System Energy ($E_{system}$)"] = f"{E_system:.6f} eV"
                                 results["Total Isolated Atom Energy ($\sum E_{atoms}$)"] = f"{E_isolated_atoms_total:.6f} eV"
 
-                    elif "Optimization" in task: # Handles both Geometry and Cell+Geometry Opt
+                    elif "Geometry Optimization" in task: # Handles both Geometry and Cell+Geometry Opt
                         is_periodic = any(calc_atoms.pbc)
                         opt_atoms_obj = FrechetCellFilter(calc_atoms) if task == "Cell + Geometry Optimization" else calc_atoms
                         # Create temporary trajectory file
@@ -1402,6 +1646,188 @@ if atoms is not None:
                             )
                         
                         show_trajectory_and_controls()
+                    elif task == "Global Optimization":
+                        st.info(f"Starting Global Optimization using {global_method}...")
+                        
+                        # Create temporary trajectory file to store the "hopping" steps
+                        traj_filename = tempfile.NamedTemporaryFile(delete=False, suffix=".traj").name
+                        
+                        # Container for live updates
+                        log_container = st.empty()
+                        global_min_energy = 0
+                        
+                        def global_log(opt_instance):
+                            """Helper to log global optimization steps."""
+                            global global_min_energy
+                            current_e = opt_instance.atoms.get_potential_energy()
+                            # For BasinHopping, nsteps is available. For others, we might need a counter.
+                            step = getattr(opt_instance, 'nsteps', 'N/A')
+                            log_container.write(f"Global Step: {step} | Energy: {current_e:.6f} eV")
+                            if current_e < global_min_energy:
+                                global_min_energy = current_e
+
+                        if global_method == "Basin Hopping":
+                            # Basin Hopping requires Temperature in eV (kB * T)
+                            kT = temperature_K * kB 
+                            
+                            # Create the wrapper for the hack needed to enforce the optimization to stop when it reaches a certain number of steps
+                            class LimitedLBFGS(LBFGS):
+                                def run(self, fmax=0.05, steps=None):
+                                    # 'steps' here overrides whatever BasinHopping tries to do.
+                                    # Set your desired max local steps (e.g., 200)
+                                    return super().run(fmax=fmax, steps=100)
+                            # Initialize Basin Hopping with the trajectory file
+                            bh = BasinHopping(calc_atoms, 
+                                            temperature=kT, 
+                                            dr=dr_amp, 
+                                            optimizer=LimitedLBFGS,
+                                            fmax=fmax_local,
+                                            trajectory=traj_filename) # Log steps to file automatically
+                            
+                            # Attach the live logger
+                            bh.attach(lambda: global_log(bh), interval=1)
+                            
+                            # Run the optimization
+                            bh.run(global_steps)
+                            
+                            results["Global Minimum Energy"] = f"{global_min_energy:.6f} eV"
+                            results["Steps Taken"] = global_steps
+                            results["Converged"] = "N/A (Global Search)"
+
+                        elif global_method == "Minima Hopping":
+                            # Minima Hopping manages its own internal optimizers and doesn't accept a 'trajectory' 
+                            # file argument in the same way BasinHopping does in __init__.
+                            opt = MinimaHopping(calc_atoms, 
+                                                T0=temperature_K,
+                                                fmax=fmax_local,
+                                                optimizer=LBFGS)
+                            
+                            # We run it. Live logging is harder here without subclassing, 
+                            # so we rely on the final output for the trajectory.
+                            opt(totalsteps=global_steps)
+                            
+                            results["Current Energy"] = f"{calc_atoms.get_potential_energy():.6f} eV"
+                            
+                            # Post-processing: MinimaHopping stores visited minima in an internal list usually.
+                            # We explicitly write the found minima to the trajectory file so the visualizer below works.
+                            # Note: opt.minima is a list of Atoms objects found during the hop.
+                            if hasattr(opt, 'minima'):
+                                from ase.io import write
+                                write(traj_filename, opt.minima)
+                            else:
+                                # Fallback if specific version doesn't store list, just save final
+                                write(traj_filename, calc_atoms)
+
+                        st.success("Global Optimization Complete!")
+                        
+                        st.markdown("### Results")
+                        for key, value in results.items():
+                            st.write(f"**{key}:** {value}")
+
+                        # --- Visualization and Downloading (Fragmented) ---
+                        
+                        # 1. Clean up the temp file path for reading
+                        # We define the visualizer function using @st.fragment to prevent full re-runs
+                        
+                        @st.fragment
+                        def show_global_trajectory_and_dl():
+                            from ase.io import read
+                            import py3Dmol
+                            
+                            # Helper to convert atoms list to XYZ string for the single download file
+                            def atoms_list_to_xyz_string(atoms_list):
+                                xyz_str = ""
+                                for i, atoms in enumerate(atoms_list):
+                                    xyz_str += f"{len(atoms)}\n"
+                                    xyz_str += f"Step {i}, 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
+
+                            if "global_traj_frames" not in st.session_state:
+                                if os.path.exists(traj_filename):
+                                    try:
+                                        # Read the trajectory we just created
+                                        trajectory = read(traj_filename, index=":")
+                                        st.session_state.global_traj_frames = trajectory
+                                        st.session_state.global_traj_index = 0
+                                    except Exception as e:
+                                        st.error(f"Error reading trajectory: {e}")
+                                        return
+                                else:
+                                    st.warning("Trajectory file not generated.")
+                                    return
+
+                            trajectory = st.session_state.global_traj_frames
+                            
+                            if not trajectory:
+                                st.warning("No steps recorded in trajectory.")
+                                return
+
+                            index = st.session_state.global_traj_index
+
+                            st.markdown("### Global Search Trajectory")
+                            st.write(f"Captured {len(trajectory)} hopping steps (Local Minima)")
+
+                            # Navigation Controls
+                            col1, col2, col3, col4 = st.columns(4)
+                            with col1:
+                                if st.button("⏮ First", key="g_first"):
+                                    st.session_state.global_traj_index = 0
+                            with col2:
+                                if st.button("◀ Previous", key="g_prev") and index > 0:
+                                    st.session_state.global_traj_index -= 1
+                            with col3:
+                                if st.button("Next ▶", key="g_next") and index < len(trajectory) - 1:
+                                    st.session_state.global_traj_index += 1
+                            with col4:
+                                if st.button("Last ⏭", key="g_last"):
+                                    st.session_state.global_traj_index = len(trajectory) - 1
+
+                            # Display Visualization
+                            current_atoms = trajectory[st.session_state.global_traj_index]
+                            st.write(f"Frame {st.session_state.global_traj_index + 1}/{len(trajectory)} | E = {current_atoms.get_potential_energy():.4f} eV")
+
+                            viz_view = get_structure_viz2(current_atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
+                            st.components.v1.html(viz_view._make_html(), width=400, height=400)
+
+                            # Download Logic
+                            full_xyz_content = atoms_list_to_xyz_string(trajectory)
+                            
+                            st.download_button(
+                                label="Download Trajectory (XYZ)",
+                                data=full_xyz_content,
+                                file_name="global_optimization_path.xyz",
+                                mime="chemical/x-xyz"
+                            )
+                            
+                            # Separate Download for just the Best Structure (Last frame usually in BH, or sorted)
+                            # Often in BH, the last frame is the accepted state, but not necessarily the global min seen *ever*.
+                            # But usually, we want the lowest energy one.
+                            energies = [a.get_potential_energy() for a in trajectory]
+                            best_idx = np.argmin(energies)
+                            best_atoms = trajectory[best_idx]
+                            
+                            # Create XYZ for single best
+                            with tempfile.NamedTemporaryFile(mode='w', suffix=".xyz", delete=False) as tmp_best:
+                                write(tmp_best.name, best_atoms)
+                                tmp_best_name = tmp_best.name
+                            
+                            with open(tmp_best_name, "r") as f:
+                                st.download_button(
+                                    label=f"Download Best Structure (E={energies[best_idx]:.4f} eV)",
+                                    data=f.read(),
+                                    file_name="best_global_structure.xyz",
+                                    mime="chemical/x-xyz"
+                                )
+                            os.unlink(tmp_best_name)
+
+                        # Call the fragment function
+                        show_global_trajectory_and_dl()
+                        
+                        # Cleanup main trajectory file after loading it into session state if desired, 
+                        # though keeping it until session end is safer for re-reads.
+                        # os.unlink(traj_filename)
                     elif task == "Vibrational Mode Analysis":
                         # Conversion factors
                         from ase.units import kB as kB_eVK, _Nav, J  # ASE's constants
@@ -1422,20 +1848,15 @@ if atoms is not None:
                                 vib.run()
                                 freqs = vib.get_frequencies()
                                 energies = vib.get_energies()
-
                                 
-
                                 print('\n\n\n\n\n\n\n\n')
                                 # vib.get_hessian_2d()
                                 # st.write(vib.summary())
                                 # print('\n')
                                 # vib.tabulate()
-
                             
-
                             freqs_cm = freqs
                             freqs_eV = energies
-                            
 
                             # Classify frequencies
                             mode_data = []
@@ -1601,6 +2022,7 @@ with st.expander('ℹ️ About This App & Foundational MLIPs'):
     - For **FairChem models**, isolated atom energies are based on pre-tabulated reference values (provided in a YAML-like structure within the app). Ensure the selected FairChem task type (`omol`, `omat`, etc. for UMA models) or model type (ESEN models use `omol` references) aligns with the system and has the necessary elemental references.
     """)
 
+
 with st.expander('🔧 Tech Stack & System Information'):
     import platform
     import psutil