Update src/streamlit_app.py

src/streamlit_app.py

@@ -87,6 +87,86 @@ SAMPLE_STRUCTURES = {
     "Ibuprofen": "ibuprofen.xyz"
 }
 
+def get_structure_viz2(atoms_obj, style='stick', show_unit_cell=True, width=400, height=400):
+    """
+    Generate visualization of atomic structure with optional unit cell display
+    
+    Parameters:
+    -----------
+    atoms_obj : ase.Atoms
+        ASE Atoms object containing the structure
+    style : str
+        Visualization style: 'ball_stick', 'stick', or 'ball'
+    show_unit_cell : bool
+        Whether to display unit cell for periodic systems
+    width, height : int
+        Dimensions of the visualization window
+    
+    Returns:
+    --------
+    py3Dmol.view object
+    """
+    
+    # Convert atoms to XYZ format
+    xyz_str = ""
+    xyz_str += f"{len(atoms_obj)}\n"
+    xyz_str += "Structure\n"
+    for atom in atoms_obj:
+        xyz_str += f"{atom.symbol} {atom.position[0]:.6f} {atom.position[1]:.6f} {atom.position[2]:.6f}\n"
+    
+    # Create a py3Dmol visualization
+    view = py3Dmol.view(width=width, height=height)
+    view.addModel(xyz_str, "xyz")
+    
+    # Set molecular style based on input
+    if style.lower() == 'ball_stick':
+        view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.3}})
+    elif style.lower() == 'stick':
+        view.setStyle({'stick': {'radius': 0.15}})
+    elif style.lower() == 'ball':
+        view.setStyle({'sphere': {'scale': 0.4}})
+    else:
+        # Default to stick if unknown style
+        view.setStyle({'stick': {'radius': 0.15}})
+    
+    # Add unit cell visualization for periodic systems
+    if show_unit_cell and any(atoms_obj.pbc):
+        cell = atoms_obj.get_cell()
+        
+        # Define unit cell edges
+        origin = np.array([0.0, 0.0, 0.0])
+        edges = [
+            # Bottom face
+            (origin, cell[0]),  # a
+            (origin, cell[1]),  # b
+            (cell[0], cell[0] + cell[1]),  # a+b from a
+            (cell[1], cell[0] + cell[1]),  # a+b from b
+            # Top face
+            (cell[2], cell[2] + cell[0]),  # a from c
+            (cell[2], cell[2] + cell[1]),  # b from c
+            (cell[2] + cell[0], cell[2] + cell[0] + cell[1]),  # a+b from c+a
+            (cell[2] + cell[1], cell[2] + cell[0] + cell[1]),  # a+b from c+b
+            # Vertical edges
+            (origin, cell[2]),  # c
+            (cell[0], cell[0] + cell[2]),  # c from a
+            (cell[1], cell[1] + cell[2]),  # c from b
+            (cell[0] + cell[1], cell[0] + cell[1] + cell[2])  # c from a+b
+        ]
+        
+        # Add unit cell lines
+        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
 
 
 # Custom logger that updates the table
@@ -302,7 +382,8 @@ if atoms is not None:
             return view
 
         # Display the 3D structure
-        view = get_structure_viz(atoms)
+        view = get_structure_viz2(atoms, style='stick', show_unit_cell=True, width=400, height=400)
+        # view = get_structure_viz(atoms)
         html_str = view._make_html()
         st.components.v1.html(html_str, width=400, height=400)