ガウス混合モデル

 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
37

import jax
import jax.numpy as jnp
import jax.random as random
from genjax import gen, flip, normal

@gen
def bento_mixture_model():
    """Two-component Gaussian mixture for Chibany's mystery bentos."""

    # Choose component with flip(0.7): True = tonkatsu (70%), False = hamburger (30%).
    # We use flip() because the choice is binary; flip(p) takes a probability directly.
    is_tonkatsu = flip(0.7) @ "component"

    # Pick the chosen component's parameters. jnp.where() selects without an if/else,
    # which keeps the model JAX-traceable.
    mu = jnp.where(is_tonkatsu, 500.0, 350.0)
    sigma = jnp.where(is_tonkatsu, 2.0, 2.0)   # both stds are 2.0 here

    # Generate the weight from the chosen component's Gaussian.
    weight = normal(mu, sigma) @ "weight"

    return weight, is_tonkatsu

# Simulate 20 bentos. GenJAX runs a model with model.simulate(key, args);
# here args is the empty tuple () because bento_mixture_model takes no arguments.
key = random.PRNGKey(42)
keys = random.split(key, 20)

# jax.vmap runs simulate once per key, in parallel.
traces = jax.vmap(lambda k: bento_mixture_model.simulate(k, ()))(keys)
weights, is_tonkatsu = traces.get_retval()

n_tonkatsu = jnp.sum(is_tonkatsu)
n_hamburger = jnp.sum(~is_tonkatsu)

print(f"Generated {n_tonkatsu} tonkatsu and {n_hamburger} hamburger bentos")
print(f"Weights: {weights}")

 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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

import jax
import jax.numpy as jnp
import jax.random as random
from genjax import gen, flip, normal, ChoiceMap

# Mystery bento weights from earlier.
mystery_weights = jnp.array([
    498.0, 352.0, 501.0, 349.0, 497.0, 503.0, 351.0, 500.0, 348.0, 502.0,
    499.0, 350.0, 498.0, 353.0, 501.0, 347.0, 499.0, 502.0, 352.0, 500.0
])

# The number of observations is a *structural* constant of the model — the @gen
# function below loops `range(N_OBS)` to lay out one assignment + one observation
# per data point. It must be a plain Python int (not a traced model argument),
# because JAX cannot trace through `range()` of a traced value. We close over it.
N_OBS = len(mystery_weights)
SIGMA_KNOWN = 2.0  # known standard deviation, shared by both components

@gen
def bayesian_gmm():
    """Bayesian 2-component Gaussian Mixture Model.

    Latents: two component means (mu_0, mu_1) and one assignment per observation.
    The standard deviation is the fixed constant SIGMA_KNOWN. The number of
    observations is the fixed constant N_OBS — both are closed over, not passed
    as arguments, so the loop length is known at trace time.
    """
    # Priors on the two component means (vague Normal priors).
    mu_0 = normal(400.0, 50.0) @ "mu_0"
    mu_1 = normal(400.0, 50.0) @ "mu_1"

    # Generate each observation: pick a component, then sample from its Gaussian.
    for i in range(N_OBS):
        # Assignment for observation i: flip(0.5) — True = component 1, False = component 0.
        z_i = flip(0.5) @ f"z_{i}"

        # Pull the assigned component's mean with jnp.where (keeps the model traceable).
        mu_i = jnp.where(z_i, mu_1, mu_0)

        # The observation itself.
        x_i = normal(mu_i, SIGMA_KNOWN) @ f"x_{i}"

    return mu_0, mu_1

# Condition on the observed weights by building a ChoiceMap that fixes each "x_i".
# ChoiceMap.d({...}) builds a choice map from a plain Python dict.
observations = ChoiceMap.d({
    f"x_{i}": mystery_weights[i] for i in range(N_OBS)
})

# generate() runs the model with those choices forced, and returns a trace plus
# a log-importance-weight. Running it for many keys gives weighted posterior samples.
# The model takes no arguments, so the args tuple is empty: ().
key = random.PRNGKey(42)
num_particles = 1000
keys = random.split(key, num_particles)

def one_particle(k):
    trace, log_weight = bayesian_gmm.generate(k, observations, ())
    choices = trace.get_choices()
    return choices["mu_0"], choices["mu_1"], log_weight

mu_0_samples, mu_1_samples, log_weights = jax.vmap(one_particle)(keys)

# Convert log-weights to normalized weights (log-sum-exp trick for stability).
weights = jnp.exp(log_weights - jnp.max(log_weights))
weights = weights / jnp.sum(weights)

# Weighted posterior means.
post_mu_0 = jnp.sum(weights * mu_0_samples)
post_mu_1 = jnp.sum(weights * mu_1_samples)

print(f"Posterior mean for mu_0: {post_mu_0:.1f}")
print(f"Posterior mean for mu_1: {post_mu_1:.1f}")