 codex/positive-covariance-from-grouping.py

"""Exact four-point example: positive pooled, negative within-group covariance."""
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def main():
    g = np.array([-1, -1, 1, 1])
    e = np.array([-1, 1, -1, 1])
    x, y = 2*g + e, 2*g - e
    assert np.allclose(np.cov(x, y, bias=True), [[5, 3], [3, 5]])
    fig, ax = plt.subplots(figsize=(7.6, 4.2), dpi=100, layout="constrained")
    fig.set_facecolor("white")
    for group, color in [(-1, "#1565c0"), (1, "#c45114")]:
        keep = g == group
        assert np.cov(x[keep], y[keep], bias=True)[0, 1] == -1
        ax.plot(x[keep], y[keep], "o-", color=color, linewidth=2,
                markersize=8, label=f"G = {group:+d}: covariance −1")
    t = np.linspace(-3.5, 3.5, 100)
    ax.plot(t, .6*t, "--", color="#444444", label="Pooled fit: slope 3/5")
    ax.axhline(0, color="#bbbbbb", lw=.7)
    ax.axvline(0, color="#bbbbbb", lw=.7)
    ax.set(xlabel="X = 2G + E", ylabel="Y = 2G − E",
           title="Pooled covariance +3; within each group −1",
           xlim=(-3.6, 3.6), ylim=(-3.6, 3.6))
    ax.legend(loc="upper left", fontsize=9)
    ax.grid(alpha=.15)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()