Update src/streamlit_app.py

src/streamlit_app.py

@@ -626,7 +626,7 @@ def get_structure_viz2(atoms_obj, style='stick', show_unit_cell=True, width=400,
     view = py3Dmol.view(width=width, height=height)
     view.addModel(xyz_str, "xyz")
     
-    if style.lower() == 'ball_stick':
+    if style.lower() == 'ball-stick':
         view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.3}})
     elif style.lower() == 'stick':
         view.setStyle({'stick': {}})
@@ -865,7 +865,11 @@ if atoms is not None:
     
     with col1:
         st.markdown('### Structure Visualization', unsafe_allow_html=True)
-        view_3d = get_structure_viz2(atoms, style='stick', show_unit_cell=True, width=400, height=400)
+        viz_style = st.selectbox("Select Visualization Style:", 
+                           ["ball-stick", 
+                            "stick", 
+                            "ball"])
+        view_3d = get_structure_viz2(atoms, style=viz_style, show_unit_cell=True, width=400, height=400)
         st.components.v1.html(view_3d._make_html(), width=400, height=400)
         
         st.markdown("### Structure Information")
@@ -1070,7 +1074,7 @@ if atoms is not None:
                             view_viz.setBackgroundColor('white')
                             return view_viz
 
-                        opt_view = get_structure_viz2(calc_atoms, style='stick', show_unit_cell=True, width=400, height=400)
+                        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)
                         
                         with tempfile.NamedTemporaryFile(delete=False, suffix=".xyz", mode="w+") as tmp_file_opt: