caffeine_analysis.py

import pymc3 as pm
import numpy as np
import arviz as az
%config InlineBackend.figure_format = 'retina'

# Data from https://twitter.com/tomstafford/status/1456914037195907079?s=20
N = np.array([1258, 280])  # total number of caffiene and non-caffiene drinkers
k = np.array([966, 168])  # total number of those who have favourite mugs

def cohens_h(p):
    return (2 * np.arcsin(np.sqrt(p[0])) 
            - 2 * np.arcsin(np.sqrt(p[1])))

COORDS = {"group": ["caffiene", "no caffiene"]}

with pm.Model(coords = COORDS) as model:
    θ = pm.Beta("θ", 1, 1, dims="group")    
    h = pm.Deterministic("h", cohens_h(θ))
    pm.Binomial("k", n=N, p=θ, observed=k)
    trace = pm.sample(return_inferencedata=True)

# Plot posterior of proportion having favourite mug
az.plot_posterior(trace, var_names="θ");

# Plot posterior distribution of effect size
ax = az.plot_posterior(trace, var_names="h")
ax.set(title="Cohen's h")