There are two problems in your code:

pm.Exponential
uses a different parameterization compared to Distributions
or Turing
, so you should specify alpha = mean(count_data)
instead of alpha = 1 / mean(count_data)
.

tau
is in fact a discrete value, though you can use a continuous approximation in some sense, the HMC
or gradient will not work for tau
. You should specify a socalled CompoundStep
in terms of pymc
.
For example, if you replace the pm.Metropolis
with pm.HamiltonianMC
, you will find that parameters use different sampler:
### Mysterious code to be explained in Chapter 3.
with model:
step = pm.Metropolis()
trace = pm.sample(10000, tune=5000,step=step)
"""
Sequential sampling (2 chains in 1 job)
CompoundStep
>Metropolis: [tau]
>Metropolis: [lambda_2]
>Metropolis: [lambda_1]
100%██████████ 15000/15000 [00:06<00:00, 2194.08it/s]
100%██████████ 15000/15000 [00:06<00:00, 2469.70it/s]
The number of effective samples is smaller than 25% for some parameters.
"""
### Mysterious code to be explained in Chapter 3.
with model:
step = pm.HamiltonianMC()
trace = pm.sample(10000, tune=5000,step=step)
"""
Sequential sampling (2 chains in 1 job)
CompoundStep
>HamiltonianMC: [lambda_2, lambda_1]
>Metropolis: [tau]
100%██████████ 15000/15000 [00:07<00:00, 1899.67it/s]
100%██████████ 15000/15000 [00:07<00:00, 2078.96it/s]
The estimated number of effective samples is smaller than 200 for some parameters.
"""
My solution, while I don’t know much how Gibbs
work and MH
is not exactly equivalent to pymc.Metropolis
:
# alpha = 1 / mean(count_data)
alpha = mean(count_data)
@model function λ(y)
lambda1 ~ Exponential(alpha)
lambda2 ~ Exponential(alpha)
τ ~ Uniform(1, n_count_data)
for i in 1:length(y)
y[i] ~ Poisson(τ > i ? lambda1 : lambda2)
end
# return τ
end
sampler = Gibbs(MH(:τ), HMC(0.01, 5, :lambda1, :lambda2))
e = sample(λ(count_data), sampler, 10000)
plot(e)
EDIT: Ok, it seems that just using HMC does not matter much. Random momentum is valid proposal though:
sampler = HMC(0.01, 10)
e = sample(λ(count_data), sampler, 10000)