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)