import numpy as np
import matplotlib.pyplot as plt
from pymc import Normal, Beta, Bernoulli, deterministic, MCMC, Matplot
%matplotlib inline

n = 1000
y_sample = np.random.binomial(1, 0.5, size=(n,))
x_sample = np.empty(n)
x_sample[y_sample == 0] = np.random.normal(-1, 1, size=(n,))[y_sample == 0]
x_sample[y_sample == 1] = np.random.normal(1, 1, size=(n,))[y_sample == 1]

p = Beta('p', alpha=1.0, beta=1.0) 
y = Bernoulli('y', p=p, value=y_sample, observed=True)

mu0 = Normal('mu0', 0.0, 1.0)
mu1 = Normal('mu1', 0.0, 1.0)

@deterministic(plot=False)
def mu(y=y, mu0=mu0, mu1=mu1):
    out = np.empty_like(y, dtype=np.float)
    out[y == 0] = mu0
    out[y == 1] = mu1
    return out

x = Normal('x', mu, 1.0, value=x_sample, observed=True)

M = MCMC(input=[p, y, mu0, mu1, mu, x])
M.sample(iter=10000)

Matplot.plot(p)

Matplot.plot(mu0)

Matplot.plot(mu1)