def test_FFBSStep():
with pm.Model(), pytest.raises(ValueError):
P_rv = np.eye(2)[None, ...]
S_rv = DiscreteMarkovChain("S_t", P_rv, np.r_[1.0, 0.0], shape=10)
S_2_rv = DiscreteMarkovChain("S_2_t", P_rv, np.r_[0.0, 1.0], shape=10)
PoissonZeroProcess("Y_t",
9.0,
S_rv + S_2_rv,
observed=np.random.poisson(9.0, size=10))
# Only one variable can be sampled by this step method
ffbs = FFBSStep([S_rv, S_2_rv])
with pm.Model(), pytest.raises(TypeError):
S_rv = pm.Categorical("S_t", np.r_[1.0, 0.0], shape=10)
PoissonZeroProcess("Y_t",
9.0,
S_rv,
observed=np.random.poisson(9.0, size=10))
# Only `DiscreteMarkovChains` can be sampled with this step method
ffbs = FFBSStep([S_rv])
with pm.Model(), pytest.raises(TypeError):
P_rv = np.eye(2)[None, ...]
S_rv = DiscreteMarkovChain("S_t", P_rv, np.r_[1.0, 0.0], shape=10)
pm.Poisson("Y_t", S_rv, observed=np.random.poisson(9.0, size=10))
# Only `SwitchingProcess`es can used as dependent variables
ffbs = FFBSStep([S_rv])
np.random.seed(2032)
poiszero_sim, _ = simulate_poiszero_hmm(30, 150)
y_test = poiszero_sim["Y_t"]
with pm.Model() as test_model:
p_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2)
p_1_rv = pm.Dirichlet("p_1", np.r_[1, 1], shape=2)
P_tt = at.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt))
pi_0_tt = compute_steady_state(P_rv)
S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_tt, shape=y_test.shape[0])
PoissonZeroProcess("Y_t", 9.0, S_rv, observed=y_test)
with test_model:
ffbs = FFBSStep([S_rv])
test_point = test_model.test_point.copy()
test_point["p_0_stickbreaking__"] = poiszero_sim["p_0_stickbreaking__"]
test_point["p_1_stickbreaking__"] = poiszero_sim["p_1_stickbreaking__"]
res = ffbs.step(test_point)
assert np.array_equal(res["S_t"], poiszero_sim["S_t"])
def test_FFBSStep():
np.random.seed(2032)
poiszero_sim, _ = simulate_poiszero_hmm(30, 150)
y_test = poiszero_sim["Y_t"]
with pm.Model() as test_model:
p_0_rv = pm.Dirichlet("p_0", np.r_[1, 1])
p_1_rv = pm.Dirichlet("p_1", np.r_[1, 1])
P_tt = tt.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
pi_0_tt = compute_steady_state(P_rv)
S_rv = HMMStateSeq("S_t", P_rv, pi_0_tt, shape=y_test.shape[0])
Y_rv = PoissonZeroProcess("Y_t", 9.0, S_rv, observed=y_test)
with test_model:
ffbs = FFBSStep([S_rv])
test_point = test_model.test_point.copy()
test_point["p_0_stickbreaking__"] = poiszero_sim["p_0_stickbreaking__"]
test_point["p_1_stickbreaking__"] = poiszero_sim["p_1_stickbreaking__"]
res = ffbs.step(test_point)
assert np.array_equal(res["S_t"], poiszero_sim["S_t"])
def simulate_poiszero_hmm(N,
mu=10.0,
pi_0_a=np.r_[1, 1],
p_0_a=np.r_[5, 1],
p_1_a=np.r_[1, 1]):
with pm.Model() as test_model:
p_0_rv = pm.Dirichlet("p_0", p_0_a)
p_1_rv = pm.Dirichlet("p_1", p_1_a)
P_tt = tt.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
pi_0_tt = pm.Dirichlet("pi_0", pi_0_a)
S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_tt, shape=N)
PoissonZeroProcess("Y_t", mu, S_rv, observed=np.zeros(N))
sample_point = pm.sample_prior_predictive(samples=1)
# Remove the extra "sampling" dimension from the sample results
sample_point = {k: v.squeeze(0) for k, v in sample_point.items()}
# Remove the extra dimension added due to `pm.sample_prior_predictive`
# forcing `size=1` in its call to `test_model.Y_t.random`.
sample_point["Y_t"] = sample_point["Y_t"].squeeze(0)
return sample_point, test_model
def simulate_poiszero_hmm(N,
mu=10.0,
pi_0_a=np.r_[1, 1],
p_0_a=np.r_[5, 1],
p_1_a=np.r_[1, 1]):
with pm.Model() as test_model:
p_0_rv = pm.Dirichlet("p_0", p_0_a)
p_1_rv = pm.Dirichlet("p_1", p_1_a)
P_tt = tt.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", tt.shape_padleft(P_tt))
pi_0_tt = pm.Dirichlet("pi_0", pi_0_a)
S_rv = HMMStateSeq("S_t", P_rv, pi_0_tt, shape=N)
Y_rv = PoissonZeroProcess("Y_t", mu, S_rv, observed=np.zeros(N))
sample_point = pm.sample_prior_predictive(samples=1)
# TODO FIXME: Why is `pm.sample_prior_predictive` adding an extra
# dimension to the `Y_rv` result?
sample_point[Y_rv.name] = sample_point[Y_rv.name].squeeze()
return sample_point, test_model
def test_FFBSStep_extreme():
"""Test a long series with extremely large mixture separation (and, thus, very small likelihoods).""" # noqa: E501
np.random.seed(2032)
mu_true = 5000
poiszero_sim, _ = simulate_poiszero_hmm(9000, mu_true)
y_test = poiszero_sim["Y_t"]
with pm.Model() as test_model:
p_0_rv = poiszero_sim["p_0"]
p_1_rv = poiszero_sim["p_1"]
P_tt = at.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt))
pi_0_tt = poiszero_sim["pi_0"]
S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_tt, shape=y_test.shape[0])
S_rv.tag.test_value = (y_test > 0).astype(int)
# This prior is very far from the true value...
E_mu, Var_mu = 100.0, 10000.0
mu_rv = pm.Gamma("mu", E_mu**2 / Var_mu, E_mu / Var_mu)
PoissonZeroProcess("Y_t", mu_rv, S_rv, observed=y_test)
with test_model:
ffbs = FFBSStep([S_rv])
test_point = test_model.test_point.copy()
test_point["p_0_stickbreaking__"] = poiszero_sim["p_0_stickbreaking__"]
test_point["p_1_stickbreaking__"] = poiszero_sim["p_1_stickbreaking__"]
with np.errstate(over="ignore", under="ignore"):
res = ffbs.step(test_point)
assert np.array_equal(res["S_t"], poiszero_sim["S_t"])
with test_model, np.errstate(over="ignore",
under="ignore"), warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
mu_step = pm.NUTS([mu_rv])
ffbs = FFBSStep([S_rv])
steps = [ffbs, mu_step]
trace = pm.sample(
20,
step=steps,
cores=1,
chains=1,
tune=100,
n_init=100,
progressbar=False,
)
assert not trace.get_sampler_stats("diverging").all()
assert trace["mu"].mean() > 1000.0
def test_PoissonZeroProcess_point():
test_states = np.r_[0, 0, 1, 1, 0, 1]
with pm.Model():
test_mean = pm.Constant("c", 1000.0)
test_point = {"c": 100.0}
test_sample = PoissonZeroProcess.dist(
test_mean, test_states).random(point=test_point)
assert np.all(0 < test_sample[..., test_states > 0])
assert np.all(test_sample[..., test_states > 0] < 200)
def test_PoissonZeroProcess_random():
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_dist = PoissonZeroProcess.dist(10.0, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random()
assert test_sample.shape == (test_states.shape[0], )
assert np.all(test_sample[test_states > 0] > 0)
test_sample = test_dist.random(size=5)
assert np.array_equal(test_sample.shape, (5, ) + test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.r_[0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0]
test_dist = PoissonZeroProcess.dist(100.0, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random(size=1)
assert np.array_equal(test_sample.shape, (1, ) + test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_mus = np.r_[10.0, 10.0, 10.0, 20.0, 20.0, 20.0]
test_dist = PoissonZeroProcess.dist(test_mus, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random()
assert np.array_equal(test_sample.shape, test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.c_[0, 0, 1, 1, 0, 1].T
test_dist = PoissonZeroProcess.dist(test_mus, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random()
# TODO: This seems bad, but also what PyMC3 would do
assert np.array_equal(test_sample.shape, test_states.squeeze().shape)
assert np.all(test_sample[..., test_states.squeeze() > 0] > 0)
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_sample = PoissonZeroProcess.dist(10.0, test_states).random(size=3)
assert np.array_equal(test_sample.shape, (3, ) + test_states.shape)
assert np.all(test_sample.sum(0)[..., test_states > 0] > 0)
def test_TransMatConjugateStep():
with pm.Model() as test_model, pytest.raises(ValueError):
p_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2)
transmat = TransMatConjugateStep(p_0_rv)
np.random.seed(2032)
poiszero_sim, _ = simulate_poiszero_hmm(30, 150)
y_test = poiszero_sim["Y_t"]
with pm.Model() as test_model:
p_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2)
p_1_rv = pm.Dirichlet("p_1", np.r_[1, 1], shape=2)
P_tt = at.stack([p_0_rv, p_1_rv])
P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt))
pi_0_tt = compute_steady_state(P_rv)
S_rv = DiscreteMarkovChain("S_t", P_rv, pi_0_tt, shape=y_test.shape[0])
PoissonZeroProcess("Y_t", 9.0, S_rv, observed=y_test)
with test_model:
transmat = TransMatConjugateStep(P_rv)
test_point = test_model.test_point.copy()
test_point["S_t"] = (y_test > 0).astype(int)
res = transmat.step(test_point)
p_0_smpl = get_test_value(
p_0_rv.distribution.transform.backward(res[p_0_rv.transformed.name]))
p_1_smpl = get_test_value(
p_1_rv.distribution.transform.backward(res[p_1_rv.transformed.name]))
sampled_trans_mat = np.stack([p_0_smpl, p_1_smpl])
true_trans_mat = (
compute_trans_freqs(poiszero_sim["S_t"], 2, counts_only=True) +
np.c_[[1, 1], [1, 1]])
true_trans_mat = true_trans_mat / true_trans_mat.sum(0)[..., None]
assert np.allclose(sampled_trans_mat, true_trans_mat, atol=0.3)
def test_PoissonZeroProcess_random():
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_dist = PoissonZeroProcess.dist(10.0, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random()
assert test_sample.shape == (test_states.shape[0], )
assert np.all(test_sample[test_states > 0] > 0)
test_sample = test_dist.random(size=5)
assert np.array_equal(test_sample.shape, (5, ) + test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.r_[0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0]
test_dist = PoissonZeroProcess.dist(100.0, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random(size=1)
assert np.array_equal(test_sample.shape, (1, ) + test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_mus = np.r_[10.0, 10.0, 10.0, 20.0, 20.0, 20.0]
test_dist = PoissonZeroProcess.dist(test_mus, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_sample = test_dist.random()
assert np.array_equal(test_sample.shape, test_states.shape)
assert np.all(test_sample[..., test_states > 0] > 0)
test_states = np.c_[0, 0, 1, 1, 0, 1].T
test_dist = PoissonZeroProcess.dist(test_mus, test_states)
# There are six Poisson means and six length one time/state sequence
# dimensions; the result should broadcast the *state sequence* along the
# six Poisson means.
assert np.array_equal(test_dist.shape, (6, 6))
test_sample = test_dist.random()
assert np.array_equal(test_sample.shape, test_states.squeeze().shape)
assert np.all(test_sample[..., test_states.squeeze() > 0] > 0)
test_states = np.c_[0, 0, 1, 1, 0, 1]
test_dist = PoissonZeroProcess.dist(test_mus, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
test_states = np.r_[0, 0, 1, 1, 0, 1]
test_sample = PoissonZeroProcess.dist(10.0, test_states).random(size=3)
assert np.array_equal(test_sample.shape, (3, ) + test_states.shape)
assert np.all(test_sample.sum(0)[..., test_states > 0] > 0)
