Conditional probability in plain words
P(A | B) reads "the probability of A given B": the share of A within the subgroup where B is true.
If older patients are more likely to have a comorbidity, then P(comorbidity | age > 60) is larger than
P(comorbidity | age ≤ 60).
When attributes depend on each other like this, comparing groups becomes tricky. That is the whole problem of causal inference in one sentence.
!pip install -q numpy polars matplotlib seaborn
import numpy as np, polars as pl
import matplotlib.pyplot as plt, seaborn as sns
rng = np.random.default_rng(101)
sns.set_theme(style="whitegrid")
A DGP where risk depends on age
We make comorbidity risk rise with age through a logistic (sigmoid) function. The sigmoid turns any number (log-odds) into a probability between 0 and 1; you will see it again for propensity scores.
N = 50_000
age = np.clip(rng.normal(45, 12, N), 18, 90)
log_odds = -4.0 + 0.07 * age
p_comorbidity = 1 / (1 + np.exp(-log_odds)) # sigmoid
comorbidity = rng.binomial(1, p_comorbidity)
df = pl.DataFrame({"age": age, "comorbidity": comorbidity})
risk = (
df.with_columns((pl.col("age") > 60).alias("over_60"))
.group_by("over_60")
.agg(pl.col("comorbidity").mean().round(3).alias("P(comorbidity | group)"),
pl.len().alias("patients"))
.sort("over_60")
)
print(risk)
xs = np.linspace(18, 90, 200)
plt.plot(xs, 1 / (1 + np.exp(-(-4 + 0.07 * xs))))
plt.xlabel("Age"); plt.ylabel("P(comorbidity)"); plt.title("Risk rises with age (logistic DGP)"); plt.show()
Independence and how to check it
Two variables are independent if knowing one tells you nothing about the other:
P(A | B) = P(A). A chi-square test on a cross-tabulation checks that formally.
from scipy.stats import chi2_contingency
age_group = (age > 60).astype(int)
table = np.array([[np.sum((age_group == g) & (comorbidity == c)) for c in (0, 1)] for g in (0, 1)])
chi2, pval, *_ = chi2_contingency(table)
print(table)
print(f"Chi-square = {chi2:.1f}, p-value = {pval:.2g} -> age and comorbidity are NOT independent")
Your first confounder
Now a treatment enters the story. Doctors prescribe a new drug more often to older patients, and older patients also have worse outcomes (more hospital days) regardless of the drug. The drug truly reduces hospital days by 2.
Age is a confounder: it affects both who gets treated and the outcome.
TRUE_EFFECT = -2.0
p_treat = 1 / (1 + np.exp(-(-5 + 0.11 * age))) # older -> more likely treated
drug = rng.binomial(1, p_treat)
hospital_days = 2 + 0.3 * age + TRUE_EFFECT * drug + rng.normal(0, 2, N)
d = pl.DataFrame({"age": age, "drug": drug, "days": hospital_days})
naive = d.filter(pl.col("drug") == 1)["days"].mean() - d.filter(pl.col("drug") == 0)["days"].mean()
print(f"True effect: {TRUE_EFFECT:+.2f} days | Naive difference in means: {naive:+.2f} days")
The naive comparison says the drug adds hospital days. Treated patients are older, and age alone adds days. Compare like with like instead: within narrow age bands, then average the band-specific differences.
strata = (
d.with_columns((pl.col("age") // 5 * 5).alias("age_band"))
.group_by("age_band")
.agg(
(pl.col("days").filter(pl.col("drug") == 1).mean()
- pl.col("days").filter(pl.col("drug") == 0).mean()).alias("diff"),
pl.len().alias("n"),
)
.drop_nans().drop_nulls()
)
stratified = (strata["diff"] * strata["n"]).sum() / strata["n"].sum()
print(f"Stratified (age-adjusted) estimate: {stratified:+.2f} days")
Stratifying is the simplest form of adjustment: the idea behind regression controls (Module 3), matching and weighting (Module 6) and Double Machine Learning (Module 7). It only works if you adjust for the right variables, the question DAGs answer in Module 5.
Exercise: Simpson's paradox
Build a dataset with two hospitals. Hospital A treats mostly severe cases and has a higher death rate overall, yet within each severity group it has a lower death rate than Hospital B. Show both the pooled and the within-severity comparisons.
Show a solution
n = 20_000
hospital_a = rng.binomial(1, 0.5, n)
severe = rng.binomial(1, np.where(hospital_a == 1, 0.8, 0.2)) # A gets the severe cases
p_death = 0.05 + 0.25 * severe - 0.03 * hospital_a # A is better within each group
death = rng.binomial(1, p_death)
h = pl.DataFrame({"hospital_a": hospital_a, "severe": severe, "death": death})
print(h.group_by("hospital_a").agg(pl.col("death").mean().round(3)).sort("hospital_a"))
print(h.group_by("severe", "hospital_a").agg(pl.col("death").mean().round(3)).sort("severe", "hospital_a"))
Checklist and recap
- I can compute
P(A | B)withgroup_byand test independence. - I can write a DGP with a confounder and show that a naive comparison is biased.
- I can explain why stratifying on the confounder recovers the true effect.