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 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_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_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 test_time_varying_model():
np.random.seed(1039)
data = gen_toy_data()
formula_str = "1 + C(weekday)"
X_df = patsy.dmatrix(formula_str, data, return_type="dataframe")
X_np = X_df.values
xi_shape = X_np.shape[1]
xi_0_true = np.array([2.0, -2.0, 2.0, -2.0, 2.0, -2.0,
2.0]).reshape(xi_shape, 1)
xi_1_true = np.array([2.0, -2.0, 2.0, -2.0, 2.0, -2.0,
2.0]).reshape(xi_shape, 1)
xis_rv_true = np.stack([xi_0_true, xi_1_true], axis=1)
with pm.Model(**TV_CONFIG) as sim_model:
_ = create_dirac_zero_hmm(X_np,
mu=1000,
xis=xis_rv_true,
observed=np.zeros(X_np.shape[0]))
sim_point = pm.sample_prior_predictive(samples=1, model=sim_model)
y_t = sim_point["Y_t"].squeeze().astype(int)
split = int(len(y_t) * 0.7)
train_y, test_V = y_t[:split], sim_point["V_t"].squeeze()[split:]
train_X, test_X = X_np[:split, :], X_np[split:, :]
X = shared(train_X, name="X", borrow=True)
Y = shared(train_y, name="y_t", borrow=True)
with pm.Model() as model:
xis_rv = pm.Normal("xis", 0, 10, shape=xis_rv_true.shape)
_ = create_dirac_zero_hmm(X, 1000, xis_rv, Y)
number_of_draws = 500
with model:
steps = [
FFBSStep([model.V_t]),
pm.NUTS(
vars=[
model.gamma_0,
model.Gamma,
],
target_accept=0.90,
),
]
with model:
posterior_trace = pm.sample(
draws=number_of_draws,
step=steps,
random_seed=100,
return_inferencedata=True,
chains=1,
cores=1,
progressbar=True,
idata_kwargs={"dims": {
"Y_t": ["date"],
"V_t": ["date"]
}},
)
# Update the shared variable values
Y.set_value(np.ones(test_X.shape[0], dtype=Y.dtype))
X.set_value(test_X)
model.V_t.distribution.shape = (test_X.shape[0], )
hdi_data = az.hdi(posterior_trace, hdi_prob=0.95,
var_names=["xis"]).to_dataframe()
hdi_data = hdi_data.unstack(level="hdi")
xis_true_flat = xis_rv_true.squeeze().flatten()
check_idx = ~np.in1d(np.arange(len(xis_true_flat)),
np.arange(3, len(xis_true_flat), step=4))
assert np.all(
xis_true_flat[check_idx] <= hdi_data["xis",
"higher"].values[check_idx])
assert np.all(
xis_true_flat[check_idx] >= hdi_data["xis", "lower"].values[check_idx])
trace = posterior_trace.posterior.drop_vars(["Gamma", "V_t"])
with aesara.config.change_flags(compute_test_value="off"):
adds_pois_ppc = pm.sample_posterior_predictive(
trace, var_names=["V_t", "Y_t", "Gamma"], model=model)
assert (np.abs(adds_pois_ppc["V_t"] - test_V) /
test_V.shape[0]).mean() < 1e-2