# Experiments for "Describing Data": every number in the post comes from here (numpy 2.1, scipy 1.18, seeded).
#   python describing_data.py  ->  prints the results
import numpy as np
from scipy import stats

rng = np.random.default_rng(0)

# 1. variance with 1/n and 1/(n-1): 200,000 samples of size n from N(0, 1), true variance 1
print("1. variance estimators (true variance 1)")
for n in (2, 5, 10, 30):
    x = rng.standard_normal((200_000, n))
    print(f"   n={n:2d}: mean of 1/n estimator {x.var(axis=1, ddof=0).mean():.3f} (expected {(n - 1) / n:.3f}), "
          f"1/(n-1) estimator {x.var(axis=1, ddof=1).mean():.3f}")

# 2. Anscombe's quartet (Anscombe 1973)
x123 = [10, 8, 13, 9, 11, 14, 6, 4, 12, 7, 5]
Q = {"I": (x123, [8.04, 6.95, 7.58, 8.81, 8.33, 9.96, 7.24, 4.26, 10.84, 4.82, 5.68]),
     "II": (x123, [9.14, 8.14, 8.74, 8.77, 9.26, 8.10, 6.13, 3.10, 9.13, 7.26, 4.74]),
     "III": (x123, [7.46, 6.77, 12.74, 7.11, 7.81, 8.84, 6.08, 5.39, 8.15, 6.42, 5.73]),
     "IV": ([8, 8, 8, 8, 8, 8, 8, 19, 8, 8, 8], [6.58, 5.76, 7.71, 8.84, 8.47, 7.04, 5.25, 12.50, 5.56, 7.91, 6.89])}
print("2. Anscombe's quartet")
for k, (x, y) in Q.items():
    x, y = np.array(x, float), np.array(y, float); b, a = np.polyfit(x, y, 1)
    print(f"   {k:3s} mean x {x.mean():.2f} var x {x.var(ddof=1):.2f} mean y {y.mean():.2f} var y {y.var(ddof=1):.3f} "
          f"r {np.corrcoef(x, y)[0, 1]:.3f} spearman {stats.spearmanr(x, y)[0]:.3f} line y = {a:.2f} + {b:.3f} x")

# 3. strong dependence, no correlation: y = x^2 + noise
x = rng.uniform(-3, 3, 200); y = x ** 2 + rng.normal(0, 0.5, 200)
q = np.polyfit(x, y, 2); r2q = 1 - ((y - np.polyval(q, x)) ** 2).sum() / ((y - y.mean()) ** 2).sum()
print(f"3. U shape: pearson r {np.corrcoef(x, y)[0, 1]:.3f}, spearman {stats.spearmanr(x, y)[0]:.3f}, R^2 of a quadratic fit {r2q:.3f}")

# 4. a causal effect that cancels: the advert raises liking for one group and lowers it for the other
n = 1000; g = rng.integers(0, 2, n); ad = rng.uniform(0, 10, n)
like = np.where(g == 1, 0.5, -0.5) * ad + rng.normal(0, 1.5, n)
print(f"4. cancelling effects: overall r {np.corrcoef(ad, like)[0, 1]:.3f}; group A r {np.corrcoef(ad[g == 1], like[g == 1])[0, 1]:.3f}; "
      f"group B r {np.corrcoef(ad[g == 0], like[g == 0])[0, 1]:.3f}")

# 5. a confounder: city population drives both fire hydrants and dogs
pop = np.exp(rng.normal(11, 1, 300))                          # 300 towns, ~60,000 people on average
hyd = pop / 120 * np.exp(rng.normal(0, 0.25, 300)); dog = pop / 8 * np.exp(rng.normal(0, 0.25, 300))
lh, ld, lp = np.log(hyd), np.log(dog), np.log(pop)
res = lambda v: v - np.polyval(np.polyfit(lp, v, 1), lp)
print(f"5. confounder: r(hydrants, dogs) {np.corrcoef(hyd, dog)[0, 1]:.3f}, on log scale {np.corrcoef(lh, ld)[0, 1]:.3f}; "
      f"partial r given population {np.corrcoef(res(lh), res(ld))[0, 1]:.3f}")

# 6. robustness: one wrong value in 15 commute distances
d = np.round(np.exp(rng.normal(2.2, 0.6, 15)), 1); d2 = d.copy(); d2[np.argmax(d)] = 500.0
for lab, v in (("original", d), ("max set to 500 km", d2)):
    q1, q3 = np.percentile(v, [25, 75])
    print(f"6. {lab}: mean {v.mean():.2f} median {np.median(v):.2f} sd {v.std(ddof=1):.2f} IQR {q3 - q1:.2f} MAD {stats.median_abs_deviation(v):.2f}")
print("   data:", d.tolist())

# 7. the coefficient of variation needs a true zero
t = np.array([18.0, 20.0, 22.0, 24.0, 26.0])
for unit, v in (("Celsius", t), ("Fahrenheit", t * 9 / 5 + 32), ("Kelvin", t + 273.15)):
    print(f"7. same temperatures in {unit}: CV {v.std(ddof=1) / v.mean():.3f}")

# 8. Pearson vs Spearman on a monotone curve
x = np.linspace(0.1, 5, 50); y = np.exp(x)
print(f"8. y = exp(x): pearson {np.corrcoef(x, y)[0, 1]:.3f}, spearman {stats.spearmanr(x, y)[0]:.3f}")