Retention plots prototypes

tabs/retention_plots.py

@@ -0,0 +1,98 @@
+import plotly.express as px
+import plotly.graph_objects as go
+import seaborn as sns
+import matplotlib.pyplot as plt
+from matplotlib.ticker import PercentFormatter
+
+
+def plot_wow_retention_by_type(wow_retention):
+    wow_retention["week"] = pd.to_datetime(wow_retention["week"])
+    wow_retention = wow_retention.sort_values(["trader_type", "week"])
+    fig = px.line(
+        wow_retention,
+        x="week",
+        y="retention_rate",
+        color="trader_type",
+        markers=True,
+        title="Weekly Retention Rate by Trader Type",
+        labels={
+            "week": "Week",
+            "retention_rate": "Retention Rate (%)",
+            "trader_type": "Trader Type",
+        },
+    )
+
+    fig.update_layout(
+        hovermode="x unified",
+        legend=dict(
+            yanchor="middle",
+            y=0.5,
+            xanchor="left",
+            x=1.02,  # Move legend outside
+            orientation="v",
+        ),
+        yaxis=dict(
+            ticksuffix="%",
+            range=[
+                0,
+                max(wow_retention["retention_rate"]) * 1.1,
+            ],  # Add 10% padding to y-axis
+        ),
+        xaxis=dict(tickformat="%Y-%m-%d"),
+        margin=dict(r=150),  # Add right margin to make space for legend
+    )
+
+    # Add hover template
+    fig.update_traces(
+        hovertemplate="<b>%{y:.1f}%</b><br>Week: %{x|%Y-%m-%d}<extra></extra>"
+    )
+
+    return fig
+
+
+def plot_cohort_retention_heatmap(retention_matrix):
+    # Create a copy of the matrix to avoid modifying the original
+    retention_matrix = retention_matrix.copy()
+
+    # Convert index to datetime and format to date string
+    retention_matrix.index = pd.to_datetime(retention_matrix.index).strftime("%Y-%m-%d")
+
+    # Create figure and axes with specified size
+    plt.figure(figsize=(12, 8))
+
+    # Create mask for NaN values
+    mask = retention_matrix.isna()
+
+    # Create heatmap
+    ax = sns.heatmap(
+        data=retention_matrix,
+        annot=True,  # Show numbers in cells
+        fmt=".1f",  # Format numbers to 1 decimal place
+        cmap="YlOrRd",  # Yellow to Orange to Red color scheme
+        vmin=0,
+        vmax=100,
+        center=50,
+        cbar_kws={"label": "Retention Rate (%)", "format": PercentFormatter()},
+        mask=mask,
+        annot_kws={"size": 8},
+    )
+
+    # Customize the plot
+    plt.title("Cohort Retention Analysis", pad=20, size=14)
+    plt.xlabel("Weeks Since First Trade", size=12)
+    plt.ylabel("Cohort Starting Week", size=12)
+
+    # Format week numbers on x-axis
+    x_labels = [f"Week {i}" for i in retention_matrix.columns]
+    ax.set_xticklabels(x_labels, rotation=45, ha="right")
+
+    # Set y-axis labels rotation
+    plt.yticks(rotation=0)
+
+    # Add gridlines
+    ax.set_axisbelow(True)
+
+    # Adjust layout to prevent label cutoff
+    plt.tight_layout()
+
+    return plt