Update my_pages/rashomon_effect.py

my_pages/rashomon_effect.py

@@ -66,16 +66,11 @@ def render():
         return fig
 
 
-    # Store highlight point based on selection
-    highlight_point = None
-    if vertical_selected:
-        highlight_point = (32, 902)
-    elif horizontal_selected:
-        highlight_point = (97, 370)  # Low credit score, high income
-    elif slant_selected:
-        highlight_point = (32, 902)  # Example point for slant model
+    higlight_point = None
+    if "highlight_point" in st.session_state:
+        highlight_point = st.session_state.highlight_point
 
-# Top scatter plot (centered to match smaller width)
+    # Top scatter plot (centered to match smaller width)
     col1, col2, col3 = st.columns([1.5, 1, 1.5])
     with col2:
         st.pyplot(plot_scatter(income, credit, colors, title="Original Data"))
@@ -83,13 +78,19 @@ def render():
     col1, col2, col3, col4, col5 = st.columns([0.5, 1, 1, 1, 0.5])
     with col2:
         st.pyplot(plot_scatter(income, credit, colors, boundary_type="vertical", highlight_point=highlight_point))
-        vertical_selected = st.button("Choose Model 1")
+        if st.button("Choose Model 1"):
+            st.session_state.highlight_point = (32, 902)
+            st.rerun()
     with col3:
         st.pyplot(plot_scatter(income, credit, colors, boundary_type="slant", highlight_point=highlight_point))
-        slant_selected = st.button("Choose Model 2")
+        if st.button("Choose Model 2"):
+            st.session_state.highlight_point = (97, 370)
+            st.rerun()
     with col4:
         st.pyplot(plot_scatter(income, credit, colors, boundary_type="horizontal", highlight_point=highlight_point))
-        horizontal_selected = st.button("Choose Model 3")
+        if st.button("Choose Model 3"):
+            st.session_state.highlight_point = (32, 902)
+            st.rerun()
 
     st.markdown("---")
     col1, col2, col3, col4 = st.columns([2, 1, 1, 1])