def test_SwitchingProcess_random():
test_states = np.r_[0, 0, 1, 1, 0, 1]
mu_zero_nonzero = [Constant.dist(0), Constant.dist(1)]
test_dist = SwitchingProcess.dist(mu_zero_nonzero, 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 = SwitchingProcess.dist(mu_zero_nonzero, 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 = [Constant.dist(i) for i in range(6)]
test_dist = SwitchingProcess.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_mus = np.arange(1, 6).astype(float)
# One of the states has emissions that are a sequence of five Dirac delta
# distributions on the values 1 to 5 (i.e. the one with values
# `test_mus`), and the other is just a single delta at 0. A single state
# sample from this emissions mixture is a length five array of zeros or the
# values 1 to 5.
# Instead of specifying a state sequence containing only one state, we use
# six state sequences--each of length one. This should give us six samples
# of either five zeros or the values 1 to 5.
test_dist = SwitchingProcess.dist(
[Constant.dist(0), Constant.dist(test_mus)], test_states
)
assert np.array_equal(test_dist.shape, (6, 5))
test_sample = test_dist.random()
assert np.array_equal(test_sample.shape, test_dist.shape)
assert np.all(test_sample[np.where(test_states > 0)[0]] == test_mus)
test_states = np.c_[0, 0, 1, 1, 0, 1]
test_mus = np.arange(1, 7).astype(float)
test_dist = SwitchingProcess.dist(
[Constant.dist(0), Constant.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 = SwitchingProcess.dist(
[Constant.dist(0), Constant.dist(test_mus)], test_states
).random(size=3)
assert np.array_equal(test_sample.shape, (3,) + test_mus.shape)
assert np.all(test_sample.sum(0)[..., test_states > 0] > 0)
def test_PoissonZeroProcess_point():
test_states = np.r_[0, 0, 1, 1, 0, 1]
with pm.Model():
test_mean = 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)
assert np.all(test_sample[..., test_states == 0] == 0)
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 = at.stack([at.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, )), shape=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",
[Constant.dist(np.array(0, dtype=np.int64)),
Constant.dist(mu)],
V_rv,
observed=observed,
)
return Y_rv
def test_subset_args():
test_dist = Constant.dist(c=np.r_[0.1, 1.2, 2.3])
test_idx = np.r_[0, 2]
res = distribution_subset_args(test_dist, shape=[3], idx=test_idx)
assert np.array_equal(res["c"].eval(), np.r_[0.1, 2.3])
test_dist = pm.Normal.dist(mu=np.r_[0.1, 1.2, 2.3], sigma=np.r_[10.0])
test_idx = np.r_[0, 2]
res = distribution_subset_args(test_dist, shape=[3], idx=test_idx)
assert np.array_equal(res["mu"].eval(), np.r_[0.1, 2.3])
assert np.array_equal(res["sigma"].eval(), np.r_[10.0, 10.0])
test_dist = pm.Poisson.dist(mu=np.r_[0.1, 1.2, 2.3])
test_idx = np.r_[0, 2]
res = distribution_subset_args(test_dist, shape=[3], idx=test_idx)
assert np.array_equal(res["mu"].eval(), np.r_[0.1, 2.3])
test_dist = pm.NegativeBinomial.dist(mu=np.r_[0.1, 1.2, 2.3], alpha=2)
test_idx = np.r_[0, 2]
res = distribution_subset_args(test_dist, shape=[3], idx=test_idx)
assert np.array_equal(res["mu"].eval(), np.r_[0.1, 2.3])
assert np.array_equal(res["alpha"].eval(), np.r_[2.0, 2.0])
def test_SwitchingProcess():
np.random.seed(2023532)
test_states = np.r_[2, 0, 1, 2, 0, 1]
test_dists = [
Constant.dist(0),
pm.Poisson.dist(100.0),
pm.Poisson.dist(1000.0)
]
test_dist = SwitchingProcess.dist(test_dists, 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)
assert np.all(0 < test_sample[test_states == 1])
assert np.all(test_sample[test_states == 1] < 1000)
assert np.all(100 < test_sample[test_states == 2])
test_mus = np.r_[100, 100, 500, 100, 100, 100]
test_dists = [
Constant.dist(0),
pm.Poisson.dist(test_mus),
pm.Poisson.dist(10000.0),
]
test_dist = SwitchingProcess.dist(test_dists, 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(200 < test_sample[2] < 600)
assert np.all(0 < test_sample[5] < 200)
assert np.all(5000 < test_sample[test_states == 2])
test_dists = [
Constant.dist(0),
pm.Poisson.dist(100.0),
pm.Poisson.dist(1000.0)
]
test_dist = SwitchingProcess.dist(test_dists, test_states)
for i in range(len(test_dists)):
test_logp = test_dist.logp(
np.tile(test_dists[i].mode.eval(), test_states.shape)).eval()
assert test_logp[test_states != i].max() < test_logp[test_states ==
i].min()
# Try a continuous mixture
test_states = np.r_[2, 0, 1, 2, 0, 1]
test_dists = [
pm.Normal.dist(0.0, 1.0),
pm.Normal.dist(100.0, 1.0),
pm.Normal.dist(1000.0, 1.0),
]
test_dist = SwitchingProcess.dist(test_dists, 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] < 10)
assert np.all(50 < test_sample[test_states == 1])
assert np.all(test_sample[test_states == 1] < 150)
assert np.all(900 < test_sample[test_states == 2])
# Make sure we can use a large number of distributions in the mixture
test_states = np.ones(50)
test_dists = [Constant.dist(i) for i in range(50)]
test_dist = SwitchingProcess.dist(test_dists, test_states)
assert np.array_equal(test_dist.shape, test_states.shape)
with pytest.raises(TypeError):
SwitchingProcess.dist([1], test_states)
with aesara.change_flags(compute_test_value="off"):
# Test for the case when a default can't be computed
test_dist = pm.Poisson.dist(at.scalar())
# Confirm that there's no default
with pytest.raises(AttributeError):
test_dist.default()
# Let it try to sample using `Distribution.random` and fail
with pytest.raises(ValueError):
SwitchingProcess.dist([test_dist], test_states)
# Evaluate multiple observed state sequences in an extreme case
test_states = at.imatrix("states")
test_states.tag.test_value = np.zeros((10, 4)).astype("int32")
test_dist = SwitchingProcess.dist(
[Constant.dist(0), Constant.dist(1)], test_states)
test_obs = np.tile(np.arange(4), (10, 1)).astype("int32")
test_logp = test_dist.logp(test_obs)
exp_logp = np.tile(
np.array([0.0] + [-np.inf] * 3, dtype=aesara.config.floatX), (10, 1))
assert np.array_equal(test_logp.tag.test_value, exp_logp)