Additional sampling tests

tests/test_sampling.py

@@ -0,0 +1,62 @@
+from pymc3_hmm.distributions import HMMStateSeq, SwitchingProcess
+from tests.utils import simulate_poiszero_hmm, check_metrics
+from pymc3_hmm.step_methods import FFBSStep, TransMatConjugateStep
+import pymc3 as pm
+import theano.tensor as tt
+import numpy as np
+
+
+def test_sampling(N: int = 200, off_param=1):
+    np.random.seed(123)
+    kwargs = {
+        "N": N,
+        "mus": [5000, 7000],
+        "pi_0_a": np.r_[1, 1, 1],
+        "Gamma": np.r_["0,2,1", [5, 1, 1], [1, 3, 1], [1, 1, 5]],
+    }
+    simulation, _ = simulate_poiszero_hmm(**kwargs)
+
+    with pm.Model() as test_model:
+        p_0_rv = pm.Dirichlet("p_0", np.r_[1, 1, 1])
+        p_1_rv = pm.Dirichlet("p_1", np.r_[1, 1, 1])
+        p_2_rv = pm.Dirichlet("p_2", np.r_[1, 1, 1])
+
+        P_tt = tt.stack([p_0_rv, p_1_rv, p_2_rv])
+        P_rv = pm.Deterministic("P_tt", P_tt)
+
+        pi_0_tt = simulation["pi_0"]
+        y_test = simulation["Y_t"]
+
+        S_rv = HMMStateSeq("S_t", y_test.shape[0], P_rv, pi_0_tt)
+        S_rv.tag.test_value = (y_test > 0).astype(np.int)
+
+        mu_1, mu_2 = kwargs["mus"]
+
+        E_1_mu, Var_1_mu = mu_1 * off_param, mu_1 * off_param / 5
+        E_2_mu, Var_2_mu = (mu_2) * off_param, mu_2 * off_param / 5
+
+        mu_1_rv = pm.Gamma("mu_1", E_1_mu ** 2 / Var_1_mu, E_1_mu / Var_1_mu)
+        mu_2_rv = pm.Gamma("mu_2", E_2_mu ** 2 / Var_2_mu, E_2_mu / Var_2_mu)
+
+        Y_rv = SwitchingProcess(
+            "Y_t",
+            [pm.Constant.dist(0), pm.Poisson.dist(mu_1_rv), pm.Poisson.dist(mu_2_rv),],
+            S_rv,
+            observed=y_test,
+        )
+
+    with test_model:
+        mu_step = pm.NUTS([mu_1_rv, mu_2_rv])
+        ffbs = FFBSStep([S_rv])
+        transitions = TransMatConjugateStep([p_0_rv, p_1_rv, p_2_rv], S_rv)
+        steps = [ffbs, mu_step, transitions]
+        trace_ = pm.sample(N, step=steps, return_inferencedata=True, chains=1)
+        posterior = pm.sample_posterior_predictive(trace_.posterior)
+        check_metrics(trace_, posterior, simulation)
+
+
+# def test_PriorRobust():
+#     for j in np.linspace(0.8, 1.2, 3):
+#         print(f'off pram : {j}')
+#         test_sampling(200, j)
+#     assert True

tests/test_sampling_seasonality.py

@@ -0,0 +1,77 @@
+from pymc3_hmm.distributions import HMMStateSeq, SwitchingProcess
+from tests.utils import (
+    simulate_poiszero_hmm,
+    check_metrics,
+)
+from pymc3_hmm.step_methods import FFBSStep, TransMatConjugateStep
+import pymc3 as pm
+import theano.tensor as tt
+import numpy as np
+import pandas as pd
+import patsy
+
+# %%
+def gen_design_matrix(N):
+    t = pd.date_range(end=pd.to_datetime("today"), periods=N, freq="H").to_frame()
+    t["weekday"] = t[0].dt.dayofweek
+    t["hour"] = t[0].dt.hour
+    t.reset_index()
+    formula_str = " 1 +  C(hour) + C(weekday)"
+    X_df = patsy.dmatrix(formula_str, t, return_type="dataframe")
+    return X_df.values
+
+
+def test_seasonality_sampling(N: int = 200, off_param=1):
+    np.random.seed(2032)
+
+    X_t = gen_design_matrix(N)
+    betas_intercept = np.random.normal(np.log(3000), 1, size=1)
+    betas_hour = np.sort(np.random.normal(1, 0.5, size=23))
+    betas_week = np.sort(np.random.normal(1, 0.5, size=6))
+
+    betas = tt.concatenate([betas_intercept, betas_hour, betas_week])
+    eta_r = tt.exp(tt.dot(X_t, betas))
+
+    kwargs = {
+        "N": N,
+        "mus": np.r_[eta_r],
+        "pi_0_a": np.r_[1, 1],
+        "Gamma": np.r_["0,2,1", [10, 1], [5, 5]],
+    }
+    simulation, _ = simulate_poiszero_hmm(**kwargs)
+
+    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", P_tt)
+
+        pi_0_tt = simulation["pi_0"]
+        y_test = simulation["Y_t"]
+
+        S_rv = HMMStateSeq("S_t", y_test.shape[0], P_rv, pi_0_tt)
+        S_rv.tag.test_value = (y_test > 0).astype(np.int)
+
+        X = gen_design_matrix(N)
+        beta_s_intercept = pm.Normal("beta_s_intercept", 8, 1, shape=(1,))
+        beta_s_hour = pm.Normal("beta_s_hour", 1, 0.5, shape=(23,))
+        beta_s_week = pm.Normal("beta_s_week", 1, 0.5, shape=(6,))
+
+        beta_s = pm.Deterministic(
+            "beta_s", tt.concatenate([beta_s_intercept, beta_s_hour, beta_s_week])
+        )
+        mu = tt.exp(tt.dot(X, beta_s))
+
+        Y_rv = SwitchingProcess(
+            "Y_t", [pm.Constant.dist(0), pm.Poisson.dist(mu)], S_rv, observed=y_test,
+        )
+    with test_model:
+        mu_step = pm.NUTS([mu, beta_s])
+        ffbs = FFBSStep([S_rv])
+        transitions = TransMatConjugateStep([p_0_rv, p_1_rv], S_rv)
+        steps = [ffbs, mu_step, transitions]
+        trace_ = pm.sample(N, step=steps, return_inferencedata=True, chains=1)
+        posterior = pm.sample_posterior_predictive(trace_.posterior)
+
+    check_metrics(trace_, posterior, simulation)

tests/utils.py

@@ -1,29 +1,81 @@
 import numpy as np
-
 import theano.tensor as tt
-
 import pymc3 as pm
+import numbers
+import theano
+import arviz as az
+
+from pymc3_hmm.distributions import HMMStateSeq, SwitchingProcess
 
-from pymc3_hmm.distributions import PoissonZeroProcess, HMMStateSeq
+theano.config.compute_test_value = "warn"
 
 
 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]
+    N, mus=np.r_[10.0, 30.0], pi_0_a=np.r_[1, 1], Gamma=np.r_["0,2", [5, 1], [1, 3]]
 ):
+    if isinstance(mus, numbers.Number):
+        mus = [mus]
+    assert pi_0_a.size == len(mus) + 1 == Gamma.shape[0] == Gamma.shape[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])
+        trans_rows = [pm.Dirichlet(f"p_{i}", r) for i, r in enumerate(Gamma)]
+        P_tt = tt.stack(trans_rows)
         P_rv = pm.Deterministic("P_tt", P_tt)
 
         pi_0_tt = pm.Dirichlet("pi_0", pi_0_a)
 
         S_rv = HMMStateSeq("S_t", N, P_rv, pi_0_tt)
 
-        Y_rv = PoissonZeroProcess("Y_t", mu, S_rv, observed=np.zeros(N))
+        Y_rv = SwitchingProcess(
+            "Y_t",
+            [pm.Constant.dist(0)] + [pm.Poisson.dist(mu) for mu in mus],
+            S_rv,
+            observed=np.zeros(N),
+        )
 
         y_test_point = pm.sample_prior_predictive(samples=1)
 
     return y_test_point, test_model
+
+
+def check_metrics(trace_, posterior, simulation):
+
+    ## checking for state prediction
+    st_trace = trace_.posterior["S_t"].mean(axis=0).mean(axis=0)
+    mean_error_rate = (
+        1
+        - np.sum(np.equal(st_trace == 0, simulation["S_t"] == 0) * 1)
+        / len(simulation["S_t"])
+    ).values.tolist()
+
+    ## check for positive possion
+    positive_index = simulation["Y_t"] > 0
+    positive_sim = simulation["Y_t"][positive_index]
+    ## point metric
+    y_trace = posterior["Y_t"].mean(axis=0)
+    MAPE = np.mean(abs(y_trace[positive_index] - positive_sim) / positive_sim)
+
+    ## confidence_metrics_
+    az_post_trace = az.from_pymc3(posterior_predictive=posterior)
+    CI_CONF = 0.95
+    post_pred_imps_hpd_df = az.hdi(
+        az_post_trace, hdi_prob=CI_CONF, group="posterior_predictive", var_names=["Y_t"]
+    ).to_dataframe()
+
+    post_pred_imps_hpd_df = post_pred_imps_hpd_df.unstack(level="hdi")
+    post_pred_imps_hpd_df.columns = post_pred_imps_hpd_df.columns.set_levels(
+        ["upper", "lower"], level="hdi"
+    )
+    pred_range = post_pred_imps_hpd_df[positive_index]["Y_t"]
+    pred_range["T_Y"] = simulation["Y_t"][positive_index]
+
+    pred_CI = sum(
+        (pred_range["T_Y"] <= pred_range["upper"])
+        & (pred_range["T_Y"] >= pred_range["lower"]) * 1
+    ) / len(pred_range)
+
+    print(mean_error_rate, MAPE, pred_CI)
+    print(pred_range)
+    assert mean_error_rate <= 0.05
+    assert MAPE <= 0.05
+    assert pred_CI >= CI_CONF - 0.05