Diagramme en mosaïque avec résidus de Pearson

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from matplotlib.colors import TwoSlopeNorm
from scipy import stats

# Data: marriage rating (1-5, 1-2 pooled) × time spent in affairs (none / ≤ 1 / > 1)
df = sm.datasets.fair.load_pandas().data
rating = pd.cut(df.rate_marriage, [0, 2, 3, 4, 5], labels=["Poor", "Fair", "Good", "Very good"])
affairs = pd.cut(df.affairs, [-1, 0, 1, np.inf], labels=["None", "Occasional", "Frequent"])
table = pd.crosstab(rating, affairs)
obs = table.to_numpy()

# Chi-squared test and Pearson residuals
chi2, p, dof, expected = stats.chi2_contingency(obs, correction=False)
resid = (obs - expected) / np.sqrt(expected)

fig, ax = plt.subplots(figsize=(8.5, 5.6))
cmap = plt.get_cmap("RdBu")                         # red = fewer than expected, blue = more
norm = TwoSlopeNorm(0, vmin=-12, vmax=12)
gap = 0.012
widths = obs.sum(axis=1) / obs.sum()
x0 = 0
for i, w in enumerate(widths):
    heights = obs[i] / obs[i].sum()
    y0 = 0
    for j, h in enumerate(heights):
        ax.bar(x0, h - gap, width=w - gap, bottom=y0, align="edge",
               color=cmap(norm(resid[i, j])), edgecolor="none")
        sep = "  " if h < 0.12 else "\n"              # thin tiles get a one-line label
        ax.text(x0 + (w - gap) / 2, y0 + (h - gap) / 2, f"{obs[i, j]}{sep}{resid[i, j].round(1) + 0:+.1f}", ha="center", va="center",
                fontsize=9, linespacing=1.3, color="white" if abs(resid[i, j]) > 7 else "#1D1D1F")
        y0 += h
    ax.text(x0 + (w - gap) / 2, -0.03, f"{table.index[i]}\n(n = {obs[i].sum()})", ha="center", va="top",
            fontsize=10, color="#3A3A3C")
    x0 += w

# Row labels on the left, placed at the heights of the first column
y0 = 0
for j, h in enumerate(obs[0] / obs[0].sum()):
    ax.text(-0.015, y0 + (h - gap) / 2, table.columns[j], ha="right", va="center", fontsize=10, color="#3A3A3C")
    y0 += h

sm_ = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
cb = fig.colorbar(sm_, ax=ax, shrink=0.8, pad=0.02, ticks=[-12, -8, -4, 0, 4, 8, 12])
cb.set_label("Pearson residual", color="#3A3A3C")
cb.outline.set_visible(False)
ax.set_xlim(0, 1 - gap); ax.set_ylim(0, 1 - gap)
ax.axis("off")
ax.text(0.5, -0.15, "Self-rated marriage quality", transform=ax.transAxes, ha="center", fontsize=11)
ax.text(-0.19, 0.5, "Time spent in affairs", transform=ax.transAxes, rotation=90, va="center", fontsize=11)
ax.set_title("Unhappy marriages carry far more affairs than chance predicts", pad=26)
ax.text(0, 1.02, f"Pearson χ²({dof}) = {chi2:.0f}, p < .001 · n = {obs.sum():,} · "
        "tile labels: count and residual", transform=ax.transAxes, fontsize=9.5, color="#6E6E73")
plt.show()