自分でモデルを作る

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

@gen
def my_model(parameters):
    # Hidden variables (causes)
    hidden = distribution(...) @ "hidden"

    # Observed variables (effects)
    # Usually depend on hidden variables
    if hidden:
        observed = distribution_A(...) @ "observed"
    else:
        observed = distribution_B(...) @ "observed"

    return hidden  # Or whatever you want to predict

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

import jax
import jax.numpy as jnp
from genjax import gen, flip, uniform, categorical, ChoiceMap

# The number of flips is a Python constant captured by the @gen, NOT a model
# argument. JAX traces model arguments into abstract values, so a Python
# `for i in range(n)` loop can't use a count that arrives as a traced argument
# (it raises TracerIntegerConversionError). Fixing the count as a module-level
# constant keeps the teaching loop readable and runnable.
N_FLIPS = 10

@gen
def coin_flips(bias=0.5):
    """Generate N_FLIPS independent coin flips."""

    results = []
    for i in range(N_FLIPS):
        # Each flip is independent
        result = flip(bias) @ f"flip_{i}"
        results.append(result)

    return jnp.array(results).astype(int)

1
2
3
4

key = jax.random.key(42)
trace = coin_flips.simulate(key, (0.7,))
flips = trace.get_retval()
print(f"Flips: {flips}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

@gen
def coin_with_unknown_bias():
    """Coin with unknown bias — infer it from N_FLIPS flips."""

    # Hidden: the coin's true bias (uniform between 0 and 1)
    bias = uniform(0.0, 1.0) @ "bias"

    # Observations: flip outcomes (N_FLIPS is the module constant from above)
    flips = []
    for i in range(N_FLIPS):
        result = flip(bias) @ f"flip_{i}"
        flips.append(result)

    return bias  # Want to infer this!

 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

# Observe 7 heads out of 10 flips
observations = ChoiceMap.d({
    "flip_0": 1, "flip_1": 1, "flip_2": 0,
    "flip_3": 1, "flip_4": 1, "flip_5": 0,
    "flip_6": 1, "flip_7": 1, "flip_8": 0,
    "flip_9": 1
})

# Infer bias
key = jax.random.key(42)
keys = jax.random.split(key, 1000)

def infer_bias(k):
    # generate(key, CONSTRAINTS, ARGS) — the model takes no args, so ()
    trace, weight = coin_with_unknown_bias.generate(k, observations, ())
    return trace.get_retval(), weight

results = jax.vmap(infer_bias)(keys)
posterior_bias = results[0]
weights = results[1]

# Use importance sampling
normalized_weights = jnp.exp(weights - jnp.max(weights))
normalized_weights = normalized_weights / jnp.sum(normalized_weights)
mean_bias = jnp.sum(posterior_bias * normalized_weights)

print(f"Estimated bias: {mean_bias:.2f}")
# Should be around 0.70 (7 heads / 10 flips)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

import jax.numpy as jnp

@gen
def mood_model():
    """Weather affects Chibany's mood."""

    # Hidden: today's weather
    is_sunny = flip(0.7) @ "is_sunny"  # 70% sunny days

    # Observable: Chibany's mood depends on weather
    # Sunny → happy 90% of the time, Rainy → happy only 30% of the time
    happy_prob = jnp.where(is_sunny, 0.9, 0.3)
    is_happy = flip(happy_prob) @ "is_happy"

    return is_sunny

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

observation = ChoiceMap.d({"is_happy": 1})

def infer_weather(k):
    # generate(key, CONSTRAINTS, ARGS)
    trace, weight = mood_model.generate(k, observation, ())
    return trace.get_retval(), weight

key = jax.random.key(42)
keys = jax.random.split(key, 10000)
results = jax.vmap(infer_weather)(keys)
posterior_sunny = results[0]
weights = results[1]

# Use importance sampling
normalized_weights = jnp.exp(weights - jnp.max(weights))
normalized_weights = normalized_weights / jnp.sum(normalized_weights)
prob_sunny = jnp.sum(posterior_sunny * normalized_weights)

print(f"P(Sunny | Happy) ≈ {prob_sunny:.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

# Like N_FLIPS above, the number of days is a Python constant, not a model arg.
DAYS = 7

@gen
def weekly_meals():
    """Model a week of meals with memory."""

    meals = []

    # First day is random
    prev_meal = flip(0.5) @ "day_0"
    meals.append(prev_meal)

    # Each subsequent day depends on the previous day. prev_meal is a traced
    # value, so we pick the probability with jnp.where (not a Python if):
    #   tonkatsu yesterday (1) → want variety → 0.3; hamburger (0) → craving → 0.8
    for day in range(1, DAYS):
        current_prob = jnp.where(prev_meal == 1, 0.3, 0.8)
        current_meal = flip(current_prob) @ f"day_{day}"
        meals.append(current_meal)
        prev_meal = current_meal

    return jnp.array(meals).astype(int)

# Simulate one week
meals = weekly_meals.simulate(jax.random.key(0), ()).get_retval()
print(f"Week of meals (1=tonkatsu, 0=hamburger): {meals}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import jax.numpy as jnp

@gen
def mixed_days():
    """Different behavior on weekends vs weekdays."""

    # Hidden: is it a weekend?
    is_weekend = flip(2/7) @ "is_weekend"  # 2 out of 7 days

    # Weekend: high chance of tonkatsu (relaxed), Weekday: lower chance (busy)
    tonkatsu_prob = jnp.where(is_weekend, 0.9, 0.3)
    lunch = flip(tonkatsu_prob) @ "lunch"

    return is_weekend

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

import jax.numpy as jnp

@gen
def disease_model(prevalence=0.01, fever_if_disease=0.9, cough_if_disease=0.8,
                  fever_if_healthy=0.1, cough_if_healthy=0.2):
    """Medical diagnosis model."""

    # Hidden: disease status
    has_disease = flip(prevalence) @ "has_disease"

    # Symptoms depend on disease status
    fever_prob = jnp.where(has_disease, fever_if_disease, fever_if_healthy)
    cough_prob = jnp.where(has_disease, cough_if_disease, cough_if_healthy)
    fever = flip(fever_prob) @ "fever"
    cough = flip(cough_prob) @ "cough"

    return has_disease

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# Patient has both symptoms
observation = ChoiceMap.d({"fever": 1, "cough": 1})

def infer_disease(k):
    # generate(key, CONSTRAINTS, ARGS)
    trace, weight = disease_model.generate(k, observation, ())
    return trace.get_retval(), weight

key = jax.random.key(42)
keys = jax.random.split(key, 10000)
results = jax.vmap(infer_disease)(keys)
posterior = results[0]
weights = results[1]

# Use importance sampling
normalized_weights = jnp.exp(weights - jnp.max(weights))
normalized_weights = normalized_weights / jnp.sum(normalized_weights)
prob_disease = jnp.sum(posterior * normalized_weights)

print(f"=== MEDICAL DIAGNOSIS ===")
print(f"Prevalence: 1%")
print(f"Symptoms: Fever + Cough")
print(f"P(Disease | Symptoms) ≈ {prob_disease:.3f}")

1
2
3
4
5

# Good
is_diseased = flip(0.01) @ "is_diseased"

# Bad
x = flip(0.01) @ "x"

1
2
3
4
5
6
7
8
9

# Good
@gen
def model(disease_prevalence=0.01, test_accuracy=0.95):
    ...

# Bad
@gen
def model(p1=0.01, p2=0.95):
    ...

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

@gen
def weather_mood(sunny_prior=0.7):
    """Model how weather affects mood.

    Args:
        sunny_prior: Base rate of sunny days (default 0.7)

    Returns:
        is_sunny: Whether it's sunny today
    """

1
2
3
4
5

# Bad — can't condition on this!
x = flip(0.5)

# Good
x = flip(0.5) @ "x"

1
2
3
4
5
6
7

# Bad — name collision!
flip1 = flip(0.5) @ "flip"
flip2 = flip(0.5) @ "flip"  # ERROR!

# Good — unique names
flip1 = flip(0.5) @ "flip_1"
flip2 = flip(0.5) @ "flip_2"