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 astep(self, point, inputs):
states = getattr(self, "state_seq_obs", None)
if states is None:
states = inputs[self.state_seq_name]
N_mat = compute_trans_freqs(states, self.n_rows, counts_only=True)
trans_res = [
d.distribution.dist.transform.forward_val(
np.random.dirichlet(
test_value(d.distribution.dist.a) +
N_mat[self.row_remaps[i]][self.row_slices[i]]))
for i, d in enumerate(self.dists)
]
sample = np.stack(trans_res, 1)
return sample.reshape(point.shape)
def test_compute_trans_freqs():
res = compute_trans_freqs(np.r_[0, 1, 1, 1, 1, 0, 1], 2, counts_only=True)
assert np.array_equal(res, np.c_[[0, 1], [2, 3]])
