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)
