Update my_pages/rashomon_effect.py

my_pages/rashomon_effect.py

@@ -1,93 +1,76 @@
-# pages/rashomon_effect.py
 import streamlit as st
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 
 def render():
     st.title("Rashomon Effect")
-    
+
     # Generate synthetic data
     np.random.seed(42)
     n_points = 100
-    
-    # Mostly correlated data
-    annual_income = np.random.normal(50, 15, n_points)  # x-axis
-    credit_score = annual_income + np.random.normal(0, 10, n_points)  # y-axis
-    
-    # Default labels: 1 = paid back (green), 0 = default (red)
-    labels = (annual_income + credit_score > 120).astype(int)
-    colors = np.where(labels == 1, 'green', 'red')
-    
-    # --- Function to draw scatter with optional decision boundary ---
-    def draw_scatter(x, y, colors, boundary=None, boundary_type=None, extra_point=None):
-        fig, ax = plt.subplots()
-    
-        # Background for decision boundary
-        if boundary is not None and boundary_type:
-            x_min, x_max = ax.get_xlim()
-            y_min, y_max = ax.get_ylim()
-    
-            if boundary_type == "vertical":
-                ax.axvline(boundary, color="black", linestyle="--")
-                ax.fill_betweenx([y_min, y_max], x_min, boundary, color="red", alpha=0.1)
-                ax.fill_betweenx([y_min, y_max], boundary, x_max, color="green", alpha=0.1)
-    
-            elif boundary_type == "horizontal":
-                ax.axhline(boundary, color="black", linestyle="--")
-                ax.fill_between([x_min, x_max], y_min, boundary, color="red", alpha=0.1)
-                ax.fill_between([x_min, x_max], boundary, y_max, color="green", alpha=0.1)
-    
-        ax.scatter(x, y, c=colors, edgecolors='black')
-    
-        if extra_point:
-            ax.scatter(*extra_point["coords"], c=extra_point["color"], s=150, edgecolors="black", marker="o")
-    
+    income = np.random.normal(50, 15, n_points)
+    credit = np.random.normal(50, 15, n_points)
+    labels = (income + credit > 100).astype(int)  # 1 = paid back, 0 = default
+
+    colors = ['green' if label == 1 else 'red' for label in labels]
+
+    # Function to plot scatter
+    def plot_scatter(x, y, colors, title="", decision_boundary=None, boundary_type=None, highlight_point=None):
+        fig, ax = plt.subplots(figsize=(5, 5))
+        ax.scatter(x, y, c=colors, alpha=0.6)
         ax.set_xlabel("Annual Income")
         ax.set_ylabel("Credit Score")
-    
+        ax.set_title(title)
+
+        # Decision boundary
+        if decision_boundary is not None:
+            if boundary_type == "vertical":
+                ax.axvline(decision_boundary, color='blue', linestyle='--')
+                ax.fill_betweenx(np.arange(min(y), max(y)), decision_boundary, max(x), alpha=0.1, color='green')
+                ax.fill_betweenx(np.arange(min(y), max(y)), min(x), decision_boundary, alpha=0.1, color='red')
+            elif boundary_type == "horizontal":
+                ax.axhline(decision_boundary, color='blue', linestyle='--')
+                ax.fill_between(np.arange(min(x), max(x)), decision_boundary, max(y), alpha=0.1, color='green')
+                ax.fill_between(np.arange(min(x), max(x)), min(y), decision_boundary, alpha=0.1, color='red')
+
+        # Highlight specific point
+        if highlight_point is not None:
+            ax.scatter(*highlight_point, c='yellow', edgecolors='black', s=200, zorder=5)
+
         return fig
-    
-    # --- Layout ---
-    col_top = st.container()
-    col_bottom = st.columns(2)
-    
-    # Draw top scatter (raw data)
-    with col_top:
-        fig_top = draw_scatter(annual_income, credit_score, colors)
-        st.pyplot(fig_top)
-        st.markdown("⬇️ Data split into two different models ⬇️")
-    
-    # Session state to store choice
-    if "chosen_boundary" not in st.session_state:
-        st.session_state.chosen_boundary = None
-    
-    # Bottom left: vertical decision boundary
-    with col_bottom[0]:
-        fig_left = draw_scatter(annual_income, credit_score, colors, boundary=55, boundary_type="vertical")
-        st.pyplot(fig_left)
-        if st.button("Choose Vertical Model"):
-            st.session_state.chosen_boundary = "vertical"
-    
-    # Bottom right: horizontal decision boundary
-    with col_bottom[1]:
-        fig_right = draw_scatter(annual_income, credit_score, colors, boundary=55, boundary_type="horizontal")
-        st.pyplot(fig_right)
-        if st.button("Choose Horizontal Model"):
-            st.session_state.chosen_boundary = "horizontal"
-    
-    # After choice: add extra individual
-    if st.session_state.chosen_boundary:
-        st.markdown("---")
-        st.subheader("Your chosen model in action")
-    
-        if st.session_state.chosen_boundary == "vertical":
-            extra = {"coords": (40, 80), "color": "blue"}
-            fig_choice = draw_scatter(annual_income, credit_score, colors, boundary=55, boundary_type="vertical", extra_point=extra)
-            st.pyplot(fig_choice)
-            st.warning("This person has a high credit score but low income. You rejected them. Why not choose a model that would approve them?")
-    
-        elif st.session_state.chosen_boundary == "horizontal":
-            extra = {"coords": (80, 40), "color": "blue"}
-            fig_choice = draw_scatter(annual_income, credit_score, colors, boundary=55, boundary_type="horizontal", extra_point=extra)
-            st.pyplot(fig_choice)
-            st.warning("This person has high income but low credit score. You rejected them. Why not choose a model that would approve them?")
+
+    # Top scatter plot (centered to match smaller width)
+    col1, col2, col3 = st.columns([1, 2, 1])
+    with col2:
+        st.pyplot(plot_scatter(income, credit, colors, title="Original Data"))
+
+    # Side-by-side decision boundary plots
+    col_left, col_right = st.columns(2)
+    with col_left:
+        st.pyplot(plot_scatter(income, credit, colors, title="Vertical Boundary",
+                               decision_boundary=55, boundary_type="vertical"))
+        left_selected = st.button("Choose Vertical Boundary")
+    with col_right:
+        st.pyplot(plot_scatter(income, credit, colors, title="Horizontal Boundary",
+                               decision_boundary=55, boundary_type="horizontal"))
+        right_selected = st.button("Choose Horizontal Boundary")
+
+    # Show new individual based on selection
+    if left_selected:
+        new_point = (40, 80)  # High credit score, low income
+        col1, col2, col3 = st.columns([1, 2, 1])
+        with col2:
+            st.pyplot(plot_scatter(income, credit, colors,
+                                   title="Vertical Boundary + New Individual",
+                                   decision_boundary=55, boundary_type="vertical",
+                                   highlight_point=new_point))
+        st.warning("This individual was rejected by your chosen model. Why not choose a model that helps them?")
+    elif right_selected:
+        new_point = (80, 40)  # Low credit score, high income
+        col1, col2, col3 = st.columns([1, 2, 1])
+        with col2:
+            st.pyplot(plot_scatter(income, credit, colors,
+                                   title="Horizontal Boundary + New Individual",
+                                   decision_boundary=55, boundary_type="horizontal",
+                                   highlight_point=new_point))
+        st.warning("This individual was rejected by your chosen model. Why not choose a model that helps them?")