Update src/streamlit_app.py

src/streamlit_app.py

@@ -616,6 +616,202 @@ SAMPLE_STRUCTURES = {
     "hBN Monolayer (4x4)": "hBN_monolayer_4x4_supercell.extxyz",
 }
 
+def get_trajectory_viz(trajectory, style='stick', show_unit_cell=True, width=400, height=400, 
+                      show_path=True, path_color='red', path_radius=0.02):
+    """
+    Visualize optimization trajectory with multiple frames
+    
+    Args:
+        trajectory: List of ASE atoms objects representing the optimization steps
+        style: Visualization style ('stick', 'ball', 'ball-stick')
+        show_unit_cell: Whether to show unit cell
+        show_path: Whether to show trajectory paths for each atom
+        path_color: Color of trajectory paths
+        path_radius: Radius of trajectory path cylinders
+    """
+    if not trajectory:
+        return None
+    
+    view = py3Dmol.view(width=width, height=height)
+    
+    # Add all frames to the viewer
+    for frame_idx, atoms_obj in enumerate(trajectory):
+        xyz_str = ""
+        xyz_str += f"{len(atoms_obj)}\n"
+        xyz_str += f"Frame {frame_idx}\n"
+        for atom in atoms_obj:
+            xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
+        
+        view.addModel(xyz_str, "xyz")
+    
+    # Set style for all models
+    if style.lower() == 'ball-stick':
+        view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.3}})
+    elif style.lower() == 'stick':
+        view.setStyle({'stick': {}})
+    elif style.lower() == 'ball':
+        view.setStyle({'sphere': {'scale': 0.4}})
+    else:
+        view.setStyle({'stick': {'radius': 0.15}})
+    
+    # Add trajectory paths
+    if show_path and len(trajectory) > 1:
+        for atom_idx in range(len(trajectory[0])):
+            for frame_idx in range(len(trajectory) - 1):
+                start_pos = trajectory[frame_idx][atom_idx].position
+                end_pos = trajectory[frame_idx + 1][atom_idx].position
+                
+                view.addCylinder({
+                    'start': {'x': start_pos[0], 'y': start_pos[1], 'z': start_pos[2]},
+                    'end': {'x': end_pos[0], 'y': end_pos[1], 'z': end_pos[2]},
+                    'radius': path_radius,
+                    'color': path_color,
+                    'alpha': 0.5
+                })
+    
+    # Add unit cell for the last frame
+    if show_unit_cell and trajectory[-1].pbc.any():
+        cell = trajectory[-1].get_cell()
+        origin = np.array([0.0, 0.0, 0.0])
+        if cell is not None and cell.any():
+            edges = [
+                (origin, cell[0]), (origin, cell[1]), (cell[0], cell[0] + cell[1]), (cell[1], cell[0] + cell[1]),
+                (cell[2], cell[2] + cell[0]), (cell[2], cell[2] + cell[1]),
+                (cell[2] + cell[0], cell[2] + cell[0] + cell[1]), (cell[2] + cell[1], cell[2] + cell[0] + cell[1]),
+                (origin, cell[2]), (cell[0], cell[0] + cell[2]), (cell[1], cell[1] + cell[2]),
+                (cell[0] + cell[1], cell[0] + cell[1] + cell[2])
+            ]
+            for start, end in edges:
+                view.addCylinder({
+                    'start': {'x': start[0], 'y': start[1], 'z': start[2]},
+                    'end': {'x': end[0], 'y': end[1], 'z': end[2]},
+                    'radius': 0.05, 'color': 'black', 'alpha': 0.7
+                })
+    
+    view.zoomTo()
+    view.setBackgroundColor('white')
+    return view
+
+
+def get_animated_trajectory_viz(trajectory, style='stick', show_unit_cell=True, width=400, height=400):
+    """
+    Create an animated trajectory visualization
+    """
+    if not trajectory:
+        return None
+    
+    view = py3Dmol.view(width=width, height=height)
+    
+    # Add all frames
+    for frame_idx, atoms_obj in enumerate(trajectory):
+        xyz_str = ""
+        xyz_str += f"{len(atoms_obj)}\n"
+        xyz_str += f"Frame {frame_idx}\n"
+        for atom in atoms_obj:
+            xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
+        
+        view.addModel(xyz_str, "xyz")
+    
+    # Set style
+    if style.lower() == 'ball-stick':
+        view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.3}})
+    elif style.lower() == 'stick':
+        view.setStyle({'stick': {}})
+    elif style.lower() == 'ball':
+        view.setStyle({'sphere': {'scale': 0.4}})
+    else:
+        view.setStyle({'stick': {'radius': 0.15}})
+    
+    # Add unit cell for last frame
+    if show_unit_cell and trajectory[-1].pbc.any():
+        cell = trajectory[-1].get_cell()
+        origin = np.array([0.0, 0.0, 0.0])
+        if cell is not None and cell.any():
+            edges = [
+                (origin, cell[0]), (origin, cell[1]), (cell[0], cell[0] + cell[1]), (cell[1], cell[0] + cell[1]),
+                (origin, cell[2]), (cell[0], cell[0] + cell[2]), (cell[1], cell[1] + cell[2]),
+                (cell[0] + cell[1], cell[0] + cell[1] + cell[2]),
+                (cell[2], cell[2] + cell[0]), (cell[2], cell[2] + cell[1]),
+                (cell[2] + cell[0], cell[2] + cell[0] + cell[1]), (cell[2] + cell[1], cell[2] + cell[0] + cell[1])
+            ]
+            for start, end in edges:
+                view.addCylinder({
+                    'start': {'x': start[0], 'y': start[1], 'z': start[2]},
+                    'end': {'x': end[0], 'y': end[1], 'z': end[2]},
+                    'radius': 0.05, 'color': 'black', 'alpha': 0.7
+                })
+    
+    view.zoomTo()
+    view.setBackgroundColor('white')
+    
+    # Enable animation
+    view.animate({'loop': 'forward', 'reps': 0, 'interval': 500})
+    
+    return view
+
+
+# Streamlit implementation example
+def display_optimization_trajectory(trajectory, viz_style='stick'):
+    """
+    Display optimization trajectory in Streamlit with controls
+    """
+    if not trajectory:
+        st.error("No trajectory data available")
+        return
+    
+    st.subheader(f"Optimization Trajectory ({len(trajectory)} steps)")
+    
+    # Trajectory options
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        viz_mode = st.selectbox(
+            "Visualization Mode",
+            ["Static with paths", "Animation", "Step-by-step"],
+            key="viz_mode"
+        )
+    
+    with col2:
+        if viz_mode == "Static with paths":
+            show_paths = st.checkbox("Show trajectory paths", value=True)
+            path_color = st.selectbox("Path color", ["red", "blue", "green", "orange"], index=0)
+        elif viz_mode == "Step-by-step":
+            frame_idx = st.slider("Frame", 0, len(trajectory)-1, 0, key="frame_slider")
+    
+    # Display visualization based on mode
+    if viz_mode == "Static with paths":
+        opt_view = get_trajectory_viz(
+            trajectory, 
+            style=viz_style, 
+            show_unit_cell=True, 
+            width=400, 
+            height=400,
+            show_path=show_paths,
+            path_color=path_color
+        )
+        st.components.v1.html(opt_view._make_html(), width=400, height=400)
+        
+    elif viz_mode == "Animation":
+        opt_view = get_animated_trajectory_viz(
+            trajectory, 
+            style=viz_style, 
+            show_unit_cell=True, 
+            width=400, 
+            height=400
+        )
+        st.components.v1.html(opt_view._make_html(), width=400, height=400)
+        
+    elif viz_mode == "Step-by-step":
+        opt_view = get_structure_viz2(
+            trajectory[frame_idx], 
+            style=viz_style, 
+            show_unit_cell=True, 
+            width=400, 
+            height=400
+        )
+        st.components.v1.html(opt_view._make_html(), width=400, height=400)
+        st.write(f"Step {frame_idx + 1} of {len(trajectory)}")
+
 def get_structure_viz2(atoms_obj, style='stick', show_unit_cell=True, width=400, height=400):
     xyz_str = ""
     xyz_str += f"{len(atoms_obj)}\n"