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_DiscreteMarkovChain_point(): test_Gammas = at.as_tensor_variable(np.array([[[1.0, 0.0], [0.0, 1.0]]])) with pm.Model(): # XXX: `draw_values` won't use the `Deterministic`s values in the `point` map! # Also, `Constant` is only for integer types (?!), so we can't use that. test_gamma_0 = pm.Dirichlet("gamma_0", np.r_[1.0, 1000.0], shape=2) test_point = {"gamma_0": np.r_[1.0, 0.0]} assert np.all( DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=10).random(point=test_point) == 0) assert np.all( DiscreteMarkovChain.dist(test_Gammas, 1.0 - test_gamma_0, shape=10).random( point=test_point) == 1)
def test_DiscreteMarkovChain_logp(Gammas, gamma_0, obs, exp_res): aesara.config.compute_test_value = "warn" test_dist = DiscreteMarkovChain.dist(Gammas, gamma_0, shape=obs.shape[-1]) test_logp_tt = test_dist.logp(obs) test_logp_val = test_logp_tt.eval() if exp_res is None: def logp_single_chain(Gammas, gamma_0, obs): state_transitions = np.stack([obs[:-1], obs[1:]]).T p_S_0_to_1 = gamma_0.dot(Gammas[0]) p_S_obs = np.empty_like(obs, dtype=np.float64) p_S_obs[0] = p_S_0_to_1[obs[0]] for t, (S_tm1, S_t) in enumerate(state_transitions): p_S_obs[t + 1] = Gammas[t, S_tm1, S_t] return np.log(p_S_obs) logp_fn = np.vectorize(logp_single_chain, signature="(n,m,m),(m),(n)->(n)") Gammas = np.broadcast_to(Gammas, (obs.shape[0], ) + Gammas.shape[-2:]) exp_res = logp_fn(Gammas, gamma_0, obs) exp_res = exp_res.sum(-1) assert np.allclose(test_logp_val, exp_res)
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 = DiscreteMarkovChain("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 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_DiscreteMarkovChain_str(): Gammas = at.as_tensor(np.eye(2)[None, ...], name="Gammas") gamma_0 = at.as_tensor(np.r_[0, 1], name="gamma_0") with pm.Model(): test_dist = DiscreteMarkovChain("P_rv", Gammas, gamma_0, shape=(2, )) assert str(test_dist) == "P_rv ~ DiscreteMarkovChain"
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_only_positive_state(): number_of_draws = 50 S = 2 mu = 10 y_t = np.repeat(0, 100) with pm.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]) Gammas_tt = pm.Deterministic("P_tt", at.shape_padleft(P_tt)) gamma_0_rv = pm.Dirichlet("gamma_0", np.ones((S, )), shape=S) V_rv = DiscreteMarkovChain("V_t", Gammas_tt, gamma_0_rv, shape=y_t.shape[0]) V_rv.tag.test_value = (y_t > 0) * 1 _ = SwitchingProcess( "Y_t", [Constant.dist(np.array(0, dtype=np.int64)), Constant.dist(mu)], V_rv, observed=y_t, ) posterior_trace = pm.sample( chains=1, draws=number_of_draws, return_inferencedata=True, step=FFBSStep([V_rv]), ) posterior_pred_trace = pm.sample_posterior_predictive( posterior_trace.posterior, var_names=["Y_t"]) assert np.all(posterior_pred_trace["Y_t"] == 0)
def create_dirac_zero_hmm(X, mu, xis, observed): S = 2 z_tt = tt.stack([tt.dot(X, xis[..., s, :]) for s in range(S)], axis=1) Gammas_tt = pm.Deterministic("Gamma", multilogit_inv(z_tt)) gamma_0_rv = pm.Dirichlet("gamma_0", np.ones((S, ))) if type(observed) == np.ndarray: T = X.shape[0] else: T = X.get_value().shape[0] V_rv = DiscreteMarkovChain("V_t", Gammas_tt, gamma_0_rv, shape=T) if type(observed) == np.ndarray: V_rv.tag.test_value = (observed > 0) * 1 else: V_rv.tag.test_value = (observed.get_value() > 0) * 1 Y_rv = SwitchingProcess( "Y_t", [pm.Constant.dist(0), pm.Constant.dist(mu)], V_rv, observed=observed, ) return Y_rv
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) 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_DiscreteMarkovChain_random(): # A single transition matrix and initial probabilities vector for each # element in the state sequence test_Gamma = np.array([[[1.0, 0.0], [0.0, 1.0]]]) test_gamma_0 = np.r_[0.0, 1.0] test_sample = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=10).random() assert np.all(test_sample == 1) test_sample = DiscreteMarkovChain.dist(test_Gamma, 1.0 - test_gamma_0, shape=10).random() assert np.all(test_sample == 0) test_sample = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=10).random(size=12) assert test_sample.shape == ( 12, 10, ) test_sample = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=10).random(size=2) assert np.array_equal(test_sample, np.stack([np.ones(10), np.ones(10)], 0).astype(int)) # Now, the same set-up, but--this time--generate two state sequences # samples test_Gamma = np.array([[[0.8, 0.2], [0.2, 0.8]]]) test_gamma_0 = np.r_[0.2, 0.8] test_sample = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=10).random(size=2) # TODO: Fix the seed, and make sure there's at least one 0 and 1? assert test_sample.shape == (2, 10) # Two transition matrices--for two distinct state sequences--and one vector # of initial probs. test_Gamma = np.stack([ np.array([[[1.0, 0.0], [0.0, 1.0]]]), np.array([[[1.0, 0.0], [0.0, 1.0]]]) ]) test_gamma_0 = np.r_[0.0, 1.0] test_dist = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=(2, 10)) test_sample = test_dist.random() assert np.array_equal(test_sample, np.stack([np.ones(10), np.ones(10)], 0).astype(int)) assert test_sample.shape == (2, 10) # Now, the same set-up, but--this time--generate three state sequence # samples test_sample = test_dist.random(size=3) assert np.array_equal( test_sample, np.tile( np.stack([np.ones(10), np.ones(10)], 0).astype(int), (3, 1, 1)), ) assert test_sample.shape == (3, 2, 10) # Two transition matrices and initial probs. for two distinct state # sequences test_Gamma = np.stack([ np.array([[[1.0, 0.0], [0.0, 1.0]]]), np.array([[[1.0, 0.0], [0.0, 1.0]]]) ]) test_gamma_0 = np.stack([np.r_[0.0, 1.0], np.r_[1.0, 0.0]]) test_dist = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=(2, 10)) test_sample = test_dist.random() assert np.array_equal(test_sample, np.stack([np.ones(10), np.zeros(10)], 0).astype(int)) assert test_sample.shape == (2, 10) # Now, the same set-up, but--this time--generate three state sequence # samples test_sample = test_dist.random(size=3) assert np.array_equal( test_sample, np.tile( np.stack([np.ones(10), np.zeros(10)], 0).astype(int), (3, 1, 1)), ) assert test_sample.shape == (3, 2, 10) # "Time"-varying transition matrices with a single vector of initial # probabilities test_Gamma = np.stack( [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), ], axis=0, ) test_gamma_0 = np.r_[1, 0] test_dist = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=3) test_sample = test_dist.random() assert np.array_equal(test_sample, np.r_[1, 0, 0]) # Now, the same set-up, but--this time--generate three state sequence # samples test_sample = test_dist.random(size=3) assert np.array_equal(test_sample, np.tile(np.r_[1, 0, 0].astype(int), (3, 1))) # "Time"-varying transition matrices with two initial # probabilities vectors test_Gamma = np.stack( [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), ], axis=0, ) test_gamma_0 = np.array([[1, 0], [0, 1]]) test_dist = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=(2, 3)) test_sample = test_dist.random() assert np.array_equal(test_sample, np.array([[1, 0, 0], [0, 1, 1]])) # Now, the same set-up, but--this time--generate three state sequence # samples test_sample = test_dist.random(size=3) assert np.array_equal( test_sample, np.tile(np.array([[1, 0, 0], [0, 1, 1]]).astype(int), (3, 1, 1))) # Two "Time"-varying transition matrices with two initial # probabilities vectors test_Gamma = np.stack( [ [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), ], [ np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), ], ], axis=0, ) test_gamma_0 = np.array([[1, 0], [0, 1]]) test_dist = DiscreteMarkovChain.dist(test_Gamma, test_gamma_0, shape=(2, 3)) test_sample = test_dist.random() assert np.array_equal(test_sample, np.array([[1, 0, 0], [1, 1, 0]])) # Now, the same set-up, but--this time--generate three state sequence # samples test_sample = test_dist.random(size=3) assert np.array_equal( test_sample, np.tile(np.array([[1, 0, 0], [1, 1, 0]]).astype(int), (3, 1, 1)))
def test_TransMatConjugateStep_subtensors(): # Confirm that Dirichlet/non-Dirichlet mixed rows can be # parsed with pm.Model(): d_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2) d_1_rv = pm.Dirichlet("p_1", np.r_[1, 1], shape=2) p_0_rv = at.as_tensor([0, 0, 1]) p_1_rv = at.zeros(3) p_1_rv = at.set_subtensor(p_0_rv[[0, 2]], d_0_rv) p_2_rv = at.zeros(3) p_2_rv = at.set_subtensor(p_1_rv[[1, 2]], d_1_rv) P_tt = at.stack([p_0_rv, p_1_rv, p_2_rv]) P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt)) DiscreteMarkovChain("S_t", P_rv, np.r_[1, 0, 0], shape=(10, )) transmat = TransMatConjugateStep(P_rv) assert transmat.row_remaps == {0: 1, 1: 2} exp_slices = {0: np.r_[0, 2], 1: np.r_[1, 2]} assert exp_slices.keys() == transmat.row_slices.keys() assert all( np.array_equal(transmat.row_slices[i], exp_slices[i]) for i in exp_slices.keys()) # Same thing, just with some manipulations of the transition matrix with pm.Model(): d_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2) d_1_rv = pm.Dirichlet("p_1", np.r_[1, 1], shape=2) p_0_rv = at.as_tensor([0, 0, 1]) p_1_rv = at.zeros(3) p_1_rv = at.set_subtensor(p_0_rv[[0, 2]], d_0_rv) p_2_rv = at.zeros(3) p_2_rv = at.set_subtensor(p_1_rv[[1, 2]], d_1_rv) P_tt = at.horizontal_stack(p_0_rv[..., None], p_1_rv[..., None], p_2_rv[..., None]) P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt.T)) DiscreteMarkovChain("S_t", P_rv, np.r_[1, 0, 0], shape=(10, )) transmat = TransMatConjugateStep(P_rv) assert transmat.row_remaps == {0: 1, 1: 2} exp_slices = {0: np.r_[0, 2], 1: np.r_[1, 2]} assert exp_slices.keys() == transmat.row_slices.keys() assert all( np.array_equal(transmat.row_slices[i], exp_slices[i]) for i in exp_slices.keys()) # Use an observed `DiscreteMarkovChain` and check the conjugate results with pm.Model(): d_0_rv = pm.Dirichlet("p_0", np.r_[1, 1], shape=2) d_1_rv = pm.Dirichlet("p_1", np.r_[1, 1], shape=2) p_0_rv = at.as_tensor([0, 0, 1]) p_1_rv = at.zeros(3) p_1_rv = at.set_subtensor(p_0_rv[[0, 2]], d_0_rv) p_2_rv = at.zeros(3) p_2_rv = at.set_subtensor(p_1_rv[[1, 2]], d_1_rv) P_tt = at.horizontal_stack(p_0_rv[..., None], p_1_rv[..., None], p_2_rv[..., None]) P_rv = pm.Deterministic("P_tt", at.shape_padleft(P_tt.T)) DiscreteMarkovChain("S_t", P_rv, np.r_[1, 0, 0], shape=(4, ), observed=np.r_[0, 1, 0, 2]) transmat = TransMatConjugateStep(P_rv)
def test_DiscreteMarkovChain_logp(): theano.config.compute_test_value = "warn" # A single transition matrix and initial probabilities vector for each # element in the state sequence test_Gammas = np.array([[[0.0, 1.0], [1.0, 0.0]]]) test_gamma_0 = np.r_[1.0, 0.0] test_obs = np.r_[1, 0, 1, 0] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) assert test_logp_tt.eval() == 0 # "Time"-varying transition matrices with a single vector of initial # probabilities test_Gammas = np.stack( [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), ], axis=0, ) test_gamma_0 = np.r_[1.0, 0.0] test_obs = np.r_[1, 0, 1, 0] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) assert test_logp_tt.eval() == 0 # Static transition matrix and two state sequences test_Gammas = np.array([[[0.0, 1.0], [1.0, 0.0]]]) test_obs = np.array([[1, 0, 1, 0], [0, 1, 0, 1]]) test_gamma_0 = np.r_[0.5, 0.5] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) test_logp = test_logp_tt.eval() assert test_logp[0] == test_logp[1] # Time-varying transition matrices and two state sequences test_Gammas = np.stack( [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), ], axis=0, ) test_obs = np.array([[1, 0, 1, 0], [0, 1, 0, 1]]) test_gamma_0 = np.r_[0.5, 0.5] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) test_logp = test_logp_tt.eval() assert test_logp[0] == test_logp[1] # Two sets of time-varying transition matrices and two state sequences test_Gammas = np.stack( [ [ np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]]), ], [ np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), ], ], axis=0, ) test_obs = np.array([[1, 0, 1, 0], [0, 0, 0, 0]]) test_gamma_0 = np.r_[0.5, 0.5] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) test_logp = test_logp_tt.eval() assert test_logp[0] == test_logp[1] # Two sets of time-varying transition matrices--via `gamma_0` # broadcasting--and two state sequences test_gamma_0 = np.array([[0.5, 0.5], [0.5, 0.5]]) test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) test_logp = test_logp_tt.eval() assert test_logp[0] == test_logp[1] # "Time"-varying transition matrices with a single vector of initial # probabilities, but--this time--with better test values test_Gammas = np.stack( [ np.array([[0.1, 0.9], [0.5, 0.5]]), np.array([[0.2, 0.8], [0.6, 0.4]]), np.array([[0.3, 0.7], [0.7, 0.3]]), np.array([[0.4, 0.6], [0.8, 0.2]]), ], axis=0, ) test_gamma_0 = np.r_[0.3, 0.7] test_obs = np.r_[1, 0, 1, 0] test_dist = DiscreteMarkovChain.dist(test_Gammas, test_gamma_0, shape=test_obs.shape[-1]) test_logp_tt = test_dist.logp(test_obs) logp_res = test_logp_tt.eval() logp_exp = np.concatenate( [ test_gamma_0.dot(test_Gammas[0])[None, ...], test_Gammas[(np.ogrid[1:4], test_obs[:-1])], ], axis=-2, ) logp_exp = logp_exp[(np.ogrid[:4], test_obs)] logp_exp = np.log(logp_exp).sum() assert np.allclose(logp_res, logp_exp)