事前分布の学習

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

import jax
import jax.numpy as jnp
import jax.random as jr
from jax.scipy.special import betaln, gammaln

k = jnp.array([70, 28, 6, 3, 2, 0])
n = jnp.array([100, 40, 10, 5, 2, 1])

def log_binom_coeff(n, k):
    # gammaln = log of the Gamma function (a continuous factorial); this is log of "n choose k".
    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1)

def population_loglik(a, b):
    """log p(all students' counts | a, b), theta integrated out (Beta-Binomial)."""
    per_student = (log_binom_coeff(n, k)
                   + betaln(a + k, b + n - k)
                   - betaln(a, b))
    return per_student.sum()

# Importance sampling over (a, b): draw many candidate populations from a broad
# hyperprior, weight each by how well it explains the data, report the weighted mean.
key = jr.PRNGKey(0)
ka, kb = jr.split(key)
N = 20000
a_samples = jr.uniform(ka, (N,), minval=0.5, maxval=20.0)   # broad hyperprior on a
b_samples = jr.uniform(kb, (N,), minval=0.5, maxval=20.0)   # broad hyperprior on b

log_w = jax.vmap(population_loglik)(a_samples, b_samples)
w = jnp.exp(log_w - log_w.max())
w = w / w.sum()

a_post = jnp.sum(w * a_samples)
b_post = jnp.sum(w * b_samples)
print(f"inferred a ~= {float(a_post):.2f}")
print(f"inferred b ~= {float(b_post):.2f}")
print(f"implied population tonkatsu rate ~= {float(a_post / (a_post + b_post)):.3f}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

import jax
import jax.numpy as jnp
import jax.random as jr
from jax.scipy.special import betaln, gammaln

def log_beta_binom(k, n, a, b):
    # log p(k | n, a, b): the Beta-Binomial marginal (theta integrated out).
    return gammaln(n + 1) - gammaln(k + 1) - gammaln(n - k + 1) \
        + betaln(a + k, b + n - k) - betaln(a, b)

def infer_population(k, n, seed=0, n_samples=60000):
    """Infer (mu, lambda) by importance sampling, a = mu*lam, b = (1-mu)*lam.
    Hyperpriors: mu ~ Uniform(0,1); lambda log-uniform over [0.1, 100] so it can
    land BELOW 1 (students differ / U-shaped) or ABOVE 1 (students are alike)."""
    km, kl = jr.split(jr.key(seed))
    mu = jr.uniform(km, (n_samples,), minval=0.01, maxval=0.99)
    lam = jnp.exp(jr.uniform(kl, (n_samples,), minval=jnp.log(0.1), maxval=jnp.log(100.0)))
    a, b = mu * lam, (1 - mu) * lam
    log_w = jax.vmap(lambda a, b: log_beta_binom(k, n, a, b).sum())(a, b)
    w = jnp.exp(log_w - log_w.max()); w = w / w.sum()
    return float(jnp.sum(w * mu)), float(jnp.sum(w * lam))

# Both classrooms share the SAME two data-light students: Emi 2/2, Farid 0/1.
classrooms = {
    "A (mixed bringers)":    (jnp.array([70, 28, 6, 3, 2, 0]),  jnp.array([100, 40, 10, 5, 2, 1])),
    "B (creatures of habit)":(jnp.array([97, 2, 19, 0, 2, 0]),  jnp.array([100, 40, 20, 20, 2, 1])),
}

for label, (k, n) in classrooms.items():
    mu, lam = infer_population(k, n)
    a, b = mu * lam, (1 - mu) * lam
    print(f"Classroom {label}:  inferred mean mu={mu:.2f}, concentration lambda={lam:.1f}")
    for name, ki, ni in [("Emi", 2, 2), ("Farid", 0, 1)]:
        print(f"    {name} {ki}/{ni}: raw {ki/ni:.2f} -> shrunk {(a + ki) / (a + b + ni):.2f}")