pymc-extras 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

tests/statespace/test_SARIMAX.py

@@ -1,405 +0,0 @@
-from itertools import combinations
-
-import numpy as np
-import pymc as pm
-import pytensor
-import pytensor.tensor as pt
-import pytest
-import statsmodels.api as sm
-
-from numpy.testing import assert_allclose, assert_array_less
-
-from pymc_extras.statespace import BayesianSARIMA
-from pymc_extras.statespace.models.utilities import (
-    make_harvey_state_names,
-    make_SARIMA_transition_matrix,
-)
-from pymc_extras.statespace.utils.constants import (
-    SARIMAX_STATE_STRUCTURES,
-    SHORT_NAME_TO_LONG,
-)
-from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
-    rng,
-)
-from tests.statespace.utilities.test_helpers import (
-    load_nile_test_data,
-    make_stationary_params,
-    simulate_from_numpy_model,
-)
-
-floatX = pytensor.config.floatX
-ATOL = 1e-8 if floatX.endswith("64") else 1e-6
-RTOL = 0 if floatX.endswith("64") else 1e-6
-
-test_state_names = [
-    ["data", "state_1", "state_2"],
-    ["data", "data_star", "state_star_1", "state_star_2"],
-    ["data", "D1.data", "data_star", "state_star_1", "state_star_2"],
-    ["data", "D1.data", "D1^2.data", "data_star", "state_star_1", "state_star_2"],
-    ["data", "state_1", "state_2", "state_3"],
-    ["data", "state_1", "state_2", "state_3", "state_4", "state_5", "state_6", "state_7"],
-    [
-        "data",
-        "state_1",
-        "state_2",
-        "state_3",
-        "state_4",
-        "state_5",
-        "state_6",
-        "state_7",
-        "state_8",
-        "state_9",
-        "state_10",
-    ],
-    [
-        "data",
-        "L1.data",
-        "L2.data",
-        "L3.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-    ],
-    [
-        "data",
-        "L1.data",
-        "L2.data",
-        "L3.data",
-        "D4.data",
-        "L1D4.data",
-        "L2D4.data",
-        "L3D4.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-    ],
-    [
-        "data",
-        "D1.data",
-        "L1D1.data",
-        "L2D1.data",
-        "L3D1.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-        "state_star_4",
-        "state_star_5",
-    ],
-    [
-        "data",
-        "D1.data",
-        "D1^2.data",
-        "L1D1^2.data",
-        "L2D1^2.data",
-        "L3D1^2.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-        "state_star_4",
-        "state_star_5",
-    ],
-    [
-        "data",
-        "D1.data",
-        "D1^2.data",
-        "L1D1^2.data",
-        "L2D1^2.data",
-        "L3D1^2.data",
-        "D1^2D4.data",
-        "L1D1^2D4.data",
-        "L2D1^2D4.data",
-        "L3D1^2D4.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-        "state_star_4",
-        "state_star_5",
-    ],
-    [
-        "data",
-        "D1.data",
-        "L1D1.data",
-        "L2D1.data",
-        "D1D3.data",
-        "L1D1D3.data",
-        "L2D1D3.data",
-        "D1D3^2.data",
-        "L1D1D3^2.data",
-        "L2D1D3^2.data",
-        "data_star",
-        "state_star_1",
-        "state_star_2",
-        "state_star_3",
-        "state_star_4",
-    ],
-    ["data", "data_star"] + [f"state_star_{i+1}" for i in range(26)],
-]
-
-test_orders = [
-    (2, 0, 2, 0, 0, 0, 0),
-    (2, 1, 2, 0, 0, 0, 0),
-    (2, 2, 2, 0, 0, 0, 0),
-    (2, 3, 2, 0, 0, 0, 0),
-    (0, 0, 0, 1, 0, 0, 4),
-    (0, 0, 0, 2, 0, 1, 4),
-    (2, 0, 2, 2, 0, 2, 4),
-    (0, 0, 0, 1, 1, 0, 4),
-    (0, 0, 0, 1, 2, 0, 4),
-    (1, 1, 1, 1, 1, 1, 4),
-    (1, 2, 1, 1, 1, 1, 4),
-    (1, 2, 1, 1, 2, 1, 4),
-    (1, 1, 1, 1, 3, 1, 3),
-    (2, 1, 2, 2, 0, 2, 12),
-]
-
-ids = [f"p={p},d={d},q={q},P={P},D={D},Q={Q},S={S}" for (p, d, q, P, D, Q, S) in test_orders]
-
-
-@pytest.fixture
-def data():
-    return load_nile_test_data()
-
-
-@pytest.fixture(scope="session")
-def arima_mod():
-    return BayesianSARIMA(order=(2, 0, 1), stationary_initialization=True, verbose=False)
-
-
-@pytest.fixture(scope="session")
-def pymc_mod(arima_mod):
-    data = load_nile_test_data()
-
-    with pm.Model(coords=arima_mod.coords) as pymc_mod:
-        # x0  = pm.Normal('x0', dims=['state'])
-        # P0_diag = pm.Gamma('P0_diag', alpha=2, beta=1, dims=['state'])
-        # P0 = pm.Deterministic('P0', pt.diag(P0_diag), dims=['state', 'state_aux'])
-        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["ar_lag"])
-        ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
-        sigma_state = pm.Exponential("sigma_state", 0.5)
-        arima_mod.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
-
-    return pymc_mod
-
-
-@pytest.fixture(scope="session")
-def arima_mod_interp():
-    return BayesianSARIMA(
-        order=(3, 0, 3),
-        stationary_initialization=False,
-        verbose=False,
-        state_structure="interpretable",
-        measurement_error=True,
-    )
-
-
-@pytest.fixture(scope="session")
-def pymc_mod_interp(arima_mod_interp):
-    data = load_nile_test_data()
-
-    with pm.Model(coords=arima_mod_interp.coords) as pymc_mod:
-        x0 = pm.Normal("x0", dims=["state"])
-        P0_sigma = pm.Exponential("P0_sigma", 1)
-        P0 = pm.Deterministic(
-            "P0", pt.eye(arima_mod_interp.k_states) * P0_sigma, dims=["state", "state_aux"]
-        )
-        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["ar_lag"])
-        ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
-        sigma_state = pm.Exponential("sigma_state", 0.5)
-        sigma_obs = pm.Exponential("sigma_obs", 0.1)
-
-        arima_mod_interp.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
-
-    return pymc_mod
-
-
-@pytest.mark.parametrize(
-    "p,d,q,P,D,Q,S,expected_names",
-    [(*order, name) for order, name in zip(test_orders, test_state_names)],
-    ids=ids,
-)
-def test_harvey_state_names(p, d, q, P, D, Q, S, expected_names):
-    if all([x == 0 for x in [p, d, q, P, D, Q, S]]):
-        pytest.skip("Skip all zero case")
-
-    k_states = max(p + P * S, q + Q * S + 1) + (S * D + d)
-    states = make_harvey_state_names(p, d, q, P, D, Q, S)
-
-    assert len(states) == k_states
-    missing_from_expected = set(expected_names) - set(states)
-    assert len(missing_from_expected) == 0
-
-
-@pytest.mark.parametrize("p,d,q,P,D,Q,S", test_orders)
-def test_make_SARIMA_transition_matrix(p, d, q, P, D, Q, S):
-    T = make_SARIMA_transition_matrix(p, d, q, P, D, Q, S)
-    mod = sm.tsa.SARIMAX(np.random.normal(size=100), order=(p, d, q), seasonal_order=(P, D, Q, S))
-    T2 = mod.ssm["transition"]
-
-    if D > 2:
-        pytest.skip("Statsmodels has a bug when D > 2, skip this test.")
-    else:
-        assert_allclose(T, T2, err_msg="Transition matrix does not match statsmodels")
-
-
-@pytest.mark.parametrize("p, d, q, P, D, Q, S", test_orders, ids=ids)
-@pytest.mark.filterwarnings(
-    "ignore:Non-invertible starting MA parameters found.",
-    "ignore:Non-stationary starting autoregressive parameters found",
-    "ignore:Non-invertible starting seasonal moving average",
-    "ignore:Non-stationary starting seasonal autoregressive",
-)
-def test_SARIMAX_update_matches_statsmodels(p, d, q, P, D, Q, S, data, rng):
-    sm_sarimax = sm.tsa.SARIMAX(data, order=(p, d, q), seasonal_order=(P, D, Q, S))
-
-    param_names = sm_sarimax.param_names
-    param_d = {name: getattr(np, floatX)(rng.normal(scale=0.1) ** 2) for name in param_names}
-
-    res = sm_sarimax.fit_constrained(param_d)
-    mod = BayesianSARIMA(
-        order=(p, d, q), seasonal_order=(P, D, Q, S), verbose=False, stationary_initialization=False
-    )
-
-    with pm.Model() as pm_mod:
-        x0 = pm.Normal("x0", shape=(mod.k_states,))
-        P0 = pm.Deterministic("P0", pt.eye(mod.k_states, dtype=floatX))
-
-        if q > 0:
-            pm.Deterministic(
-                "ma_params",
-                pt.as_tensor_variable(
-                    np.array([param_d[k] for k in param_d if k.startswith("ma.") and "S." not in k])
-                ),
-            )
-        if p > 0:
-            pm.Deterministic(
-                "ar_params",
-                pt.as_tensor_variable(
-                    np.array([param_d[k] for k in param_d if k.startswith("ar.") and "S." not in k])
-                ),
-            )
-        if P > 0:
-            pm.Deterministic(
-                "seasonal_ar_params",
-                pt.as_tensor_variable(
-                    np.array([param_d[k] for k in param_d if k.startswith("ar.S.")])
-                ),
-            )
-
-        if Q > 0:
-            pm.Deterministic(
-                "seasonal_ma_params",
-                pt.as_tensor_variable(
-                    np.array([param_d[k] for k in param_d if k.startswith("ma.S.")])
-                ),
-            )
-
-        pm.Deterministic("sigma_state", pt.as_tensor_variable(np.sqrt(param_d["sigma2"])))
-
-        mod._insert_random_variables()
-        matrices = pm.draw(mod.subbed_ssm)
-        matrix_dict = dict(zip(SHORT_NAME_TO_LONG.values(), matrices))
-
-    for matrix in ["transition", "selection", "state_cov", "obs_cov", "design"]:
-        if matrix == "transition" and D > 2:
-            pytest.skip("Statsmodels has a bug when D > 2, skip this test.)")
-        assert_allclose(matrix_dict[matrix], sm_sarimax.ssm[matrix], err_msg=f"{matrix} not equal")
-
-
-@pytest.mark.parametrize("filter_output", ["filtered", "predicted", "smoothed"])
-def test_all_prior_covariances_are_PSD(filter_output, pymc_mod, rng):
-    rv = pymc_mod[f"{filter_output}_covariance"]
-    cov_mats = pm.draw(rv, 100, random_seed=rng)
-    w, v = np.linalg.eig(cov_mats)
-    assert_array_less(0, w, err_msg=f"Smallest eigenvalue: {min(w.ravel())}")
-
-
-def test_interpretable_raises_if_d_nonzero():
-    with pytest.raises(
-        ValueError, match="Cannot use interpretable state structure with statespace differencing"
-    ):
-        BayesianSARIMA(
-            order=(2, 1, 1),
-            stationary_initialization=True,
-            verbose=False,
-            state_structure="interpretable",
-        )
-
-
-def test_interpretable_states_are_interpretable(arima_mod_interp, pymc_mod_interp):
-    with pymc_mod_interp:
-        prior = pm.sample_prior_predictive(draws=10)
-
-    prior_outputs = arima_mod_interp.sample_unconditional_prior(prior)
-    ar_lags = prior.prior.coords["ar_lag"].values - 1
-    ma_lags = prior.prior.coords["ma_lag"].values - 1
-
-    # Check the first p states are lags of the previous state
-    for t, tm1 in zip(ar_lags[1:], ar_lags[:-1]):
-        assert_allclose(
-            prior_outputs.prior_latent.isel(state=t).values[1:],
-            prior_outputs.prior_latent.isel(state=tm1).values[:-1],
-            err_msg=f"State {tm1} is not a lagged version of state {t} (AR lags)",
-        )
-
-    # Check the next p+q states are lags of the innovations
-    n = len(ar_lags)
-    for t, tm1 in zip(ma_lags[1:], ma_lags[:-1]):
-        assert_allclose(
-            prior_outputs.prior_latent.isel(state=n + t).values[1:],
-            prior_outputs.prior_latent.isel(state=n + tm1).values[:-1],
-            err_msg=f"State {n + tm1} is not a lagged version of state {n + t} (MA lags)",
-        )
-
-
-@pytest.mark.parametrize("p, d, q, P, D, Q, S", test_orders, ids=ids)
-@pytest.mark.filterwarnings(
-    "ignore:Non-invertible starting MA parameters found.",
-    "ignore:Non-stationary starting autoregressive parameters found",
-    "ignore:Maximum Likelihood optimization failed to converge.",
-)
-def test_representations_are_equivalent(p, d, q, P, D, Q, S, data, rng):
-    if (d + D) > 0:
-        pytest.skip('state_structure = "interpretable" cannot include statespace differences')
-
-    shared_params = make_stationary_params(data, p, d, q, P, D, Q, S)
-    test_values = {}
-
-    for representation in SARIMAX_STATE_STRUCTURES:
-        rng = np.random.default_rng(sum(map(ord, "representation test")))
-        mod = BayesianSARIMA(
-            order=(p, d, q),
-            seasonal_order=(P, D, Q, S),
-            stationary_initialization=False,
-            verbose=False,
-            state_structure=representation,
-        )
-        shared_params.update(
-            {
-                "x0": np.zeros(mod.k_states, dtype=floatX),
-                "initial_state_cov": np.eye(mod.k_states, dtype=floatX) * 100,
-            }
-        )
-        x, y = simulate_from_numpy_model(mod, rng, shared_params)
-        test_values[representation] = y
-
-    all_pairs = combinations(SARIMAX_STATE_STRUCTURES, r=2)
-    for rep_1, rep_2 in all_pairs:
-        assert_allclose(
-            test_values[rep_1],
-            test_values[rep_2],
-            err_msg=f"{rep_1} and {rep_2} are not the same",
-            atol=ATOL,
-            rtol=RTOL,
-        )
-
-
-@pytest.mark.parametrize("order, name", [((4, 1, 0, 0), "AR"), ((0, 0, 4, 1), "MA")])
-def test_invalid_order_raises(order, name):
-    p, P, q, Q = order
-    with pytest.raises(ValueError, match=f"The following {name} and seasonal {name} terms overlap"):
-        BayesianSARIMA(order=(p, 0, q), seasonal_order=(P, 0, Q, 4))

tests/statespace/test_VARMAX.py

@@ -1,184 +0,0 @@
-from itertools import pairwise, product
-
-import numpy as np
-import pandas as pd
-import pymc as pm
-import pytensor
-import pytensor.tensor as pt
-import pytest
-import statsmodels.api as sm
-
-from numpy.testing import assert_allclose, assert_array_less
-
-from pymc_extras.statespace import BayesianVARMAX
-from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG
-from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
-    rng,
-)
-
-floatX = pytensor.config.floatX
-ps = [0, 1, 2, 3]
-qs = [0, 1, 2, 3]
-orders = list(product(ps, qs))[1:]
-ids = [f"p={x[0]}, q={x[1]}" for x in orders]
-
-
-@pytest.fixture(scope="session")
-def data():
-    df = pd.read_csv(
-        "tests/statespace/test_data/statsmodels_macrodata_processed.csv",
-        index_col=0,
-        parse_dates=True,
-    ).astype(floatX)
-    df.index.freq = df.index.inferred_freq
-    return df
-
-
-@pytest.fixture(scope="session")
-def varma_mod(data):
-    return BayesianVARMAX(
-        endog_names=data.columns,
-        order=(2, 0),
-        stationary_initialization=True,
-        verbose=False,
-        measurement_error=True,
-    )
-
-
-@pytest.fixture(scope="session")
-def pymc_mod(varma_mod, data):
-    with pm.Model(coords=varma_mod.coords) as pymc_mod:
-        # x0 = pm.Normal("x0", dims=["state"])
-        # P0_diag = pm.Exponential("P0_diag", 1, size=varma_mod.k_states)
-        # P0 = pm.Deterministic(
-        #     "P0", pt.diag(P0_diag), dims=["state", "state_aux"]
-        # )
-        state_chol, *_ = pm.LKJCholeskyCov(
-            "state_chol", n=varma_mod.k_posdef, eta=1, sd_dist=pm.Exponential.dist(1)
-        )
-        ar_params = pm.Normal(
-            "ar_params", mu=0, sigma=0.1, dims=["observed_state", "ar_lag", "observed_state_aux"]
-        )
-        state_cov = pm.Deterministic(
-            "state_cov", state_chol @ state_chol.T, dims=["shock", "shock_aux"]
-        )
-        sigma_obs = pm.Exponential("sigma_obs", 1, dims=["observed_state"])
-
-        varma_mod.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
-
-    return pymc_mod
-
-
-@pytest.fixture(scope="session")
-def idata(pymc_mod, rng):
-    with pymc_mod:
-        idata = pm.sample_prior_predictive(draws=10, random_seed=rng)
-
-    return idata
-
-
-@pytest.mark.parametrize("order", orders, ids=ids)
-@pytest.mark.parametrize("var", ["AR", "MA", "state_cov"])
-@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.EstimationWarning")
-def test_VARMAX_param_counts_match_statsmodels(data, order, var):
-    p, q = order
-
-    mod = BayesianVARMAX(k_endog=data.shape[1], order=(p, q), verbose=False)
-    sm_var = sm.tsa.VARMAX(data, order=(p, q))
-
-    count = mod.param_counts[var]
-    if var == "state_cov":
-        # Statsmodels only counts the lower triangle
-        count = mod.k_posdef * (mod.k_posdef - 1)
-    assert count == sm_var.parameters[var.lower()]
-
-
-@pytest.mark.parametrize("order", orders, ids=ids)
-@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.EstimationWarning")
-@pytest.mark.filterwarnings("ignore::FutureWarning")
-def test_VARMAX_update_matches_statsmodels(data, order, rng):
-    p, q = order
-
-    sm_var = sm.tsa.VARMAX(data, order=(p, q))
-
-    param_counts = [None, *np.cumsum(list(sm_var.parameters.values())).tolist()]
-    param_slices = [slice(a, b) for a, b in pairwise(param_counts)]
-    param_lists = [trend, ar, ma, reg, state_cov, obs_cov] = [
-        sm_var.param_names[idx] for idx in param_slices
-    ]
-    param_d = {
-        k: getattr(np, floatX)(rng.normal(scale=0.1) ** 2)
-        for param_list in param_lists
-        for k in param_list
-    }
-
-    res = sm_var.fit_constrained(param_d)
-
-    mod = BayesianVARMAX(
-        k_endog=data.shape[1],
-        order=(p, q),
-        verbose=False,
-        measurement_error=False,
-        stationary_initialization=False,
-    )
-
-    ar_shape = (mod.k_endog, mod.p, mod.k_endog)
-    ma_shape = (mod.k_endog, mod.q, mod.k_endog)
-
-    with pm.Model() as pm_mod:
-        x0 = pm.Deterministic("x0", pt.zeros(mod.k_states, dtype=floatX))
-        P0 = pm.Deterministic("P0", pt.eye(mod.k_states, dtype=floatX))
-        ma_params = pm.Deterministic(
-            "ma_params",
-            pt.as_tensor_variable(np.array([param_d[var] for var in ma])).reshape(ma_shape),
-        )
-        ar_params = pm.Deterministic(
-            "ar_params",
-            pt.as_tensor_variable(np.array([param_d[var] for var in ar])).reshape(ar_shape),
-        )
-        state_chol = np.zeros((mod.k_posdef, mod.k_posdef), dtype=floatX)
-        state_chol[np.tril_indices(mod.k_posdef)] = np.array([param_d[var] for var in state_cov])
-        state_cov = pm.Deterministic("state_cov", pt.as_tensor_variable(state_chol @ state_chol.T))
-        mod._insert_random_variables()
-
-        matrices = pm.draw(mod.subbed_ssm)
-        matrix_dict = dict(zip(SHORT_NAME_TO_LONG.values(), matrices))
-
-    for matrix in ["transition", "selection", "state_cov", "obs_cov", "design"]:
-        assert_allclose(matrix_dict[matrix], sm_var.ssm[matrix])
-
-
-@pytest.mark.parametrize("filter_output", ["filtered", "predicted", "smoothed"])
-def test_all_prior_covariances_are_PSD(filter_output, pymc_mod, rng):
-    rv = pymc_mod[f"{filter_output}_covariance"]
-    cov_mats = pm.draw(rv, 100, random_seed=rng)
-    w, v = np.linalg.eig(cov_mats)
-    assert_array_less(0, w, err_msg=f"Smallest eigenvalue: {min(w.ravel())}")
-
-
-parameters = [
-    {"n_steps": 10, "shock_size": None},
-    {"n_steps": 10, "shock_size": 1.0},
-    {"n_steps": 10, "shock_size": np.array([1.0, 0.0, 0.0])},
-    {
-        "n_steps": 10,
-        "shock_cov": np.array([[1.38, 0.58, -1.84], [0.58, 0.99, -0.82], [-1.84, -0.82, 2.51]]),
-    },
-    {
-        "shock_trajectory": np.r_[
-            np.zeros((3, 3), dtype=floatX),
-            np.array([[1.0, 0.0, 0.0]]).astype(floatX),
-            np.zeros((6, 3), dtype=floatX),
-        ]
-    },
-]
-
-ids = ["from-posterior-cov", "scalar_shock_size", "array_shock_size", "user-cov", "trajectory"]
-
-
-@pytest.mark.parametrize("parameters", parameters, ids=ids)
-@pytest.mark.skipif(floatX == "float32", reason="Impulse covariance not PSD if float32")
-def test_impulse_response(parameters, varma_mod, idata, rng):
-    irf = varma_mod.impulse_response_function(idata.prior, random_seed=rng, **parameters)
-
-    assert not np.any(np.isnan(irf.irf.values))

tests/statespace/test_coord_assignment.py

@@ -1,116 +0,0 @@
-from contextlib import nullcontext as does_not_raise
-
-import numpy as np
-import pandas as pd
-import pymc as pm
-import pytensor
-import pytensor.tensor as pt
-import pytest
-
-from pymc_extras.statespace.models import structural
-from pymc_extras.statespace.utils.constants import (
-    FILTER_OUTPUT_DIMS,
-    FILTER_OUTPUT_NAMES,
-    SMOOTHER_OUTPUT_NAMES,
-    TIME_DIM,
-)
-from pymc_extras.statespace.utils.data_tools import (
-    NO_FREQ_INFO_WARNING,
-    NO_TIME_INDEX_WARNING,
-)
-from tests.statespace.utilities.test_helpers import load_nile_test_data
-
-function_names = ["pandas_date_freq", "pandas_date_nofreq", "pandas_nodate", "numpy", "pytensor"]
-expected_warning = [
-    does_not_raise(),
-    pytest.warns(UserWarning, match=NO_FREQ_INFO_WARNING),
-    pytest.warns(UserWarning, match=NO_TIME_INDEX_WARNING),
-    pytest.warns(UserWarning, match=NO_TIME_INDEX_WARNING),
-    pytest.warns(UserWarning, match=NO_TIME_INDEX_WARNING),
-]
-func_inputs = list(zip(function_names, expected_warning))
-floatX = pytensor.config.floatX
-
-
-@pytest.fixture
-def load_dataset():
-    data = load_nile_test_data()
-
-    def _load_dataset(f):
-        if f == "pandas_date_freq":
-            data.index.freq = data.index.inferred_freq
-            return data
-        if f == "pandas_date_nofreq":
-            data.index.freq = None
-            return data
-        elif f == "pandas_nodate":
-            return data.reset_index(drop=True)
-        elif f == "numpy":
-            return data.values
-        elif f == "pytensor":
-            return pt.as_tensor_variable(data.values)
-        else:
-            raise ValueError
-
-    return _load_dataset
-
-
-@pytest.fixture()
-def generate_timeseries():
-    def _generate_timeseries(freq):
-        index = pd.date_range(start="2000-01-01", freq=freq, periods=100)
-        data = np.random.normal(size=100).astype(floatX)
-        df = pd.DataFrame(data, index=index, columns=["level"])
-        return df
-
-    return _generate_timeseries
-
-
-@pytest.fixture()
-def create_model(load_dataset):
-    ss_mod = structural.LevelTrendComponent(order=2).build("data", verbose=False)
-
-    def _create_model(f):
-        data = load_dataset(f)
-        with pm.Model(coords=ss_mod.coords) as mod:
-            P0_diag = pm.Exponential(
-                "P0_diag",
-                1,
-                dims="state",
-            )
-            P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=("state", "state_aux"))
-            initial_trend = pm.Normal("initial_trend", dims="trend_state")
-            sigma_trend = pm.Exponential("sigma_trend", 1, dims="trend_shock")
-            ss_mod.build_statespace_graph(data, save_kalman_filter_outputs_in_idata=True)
-        return mod
-
-    return _create_model
-
-
-@pytest.mark.parametrize("f, warning", func_inputs, ids=function_names)
-def test_filter_output_coord_assignment(f, warning, create_model):
-    with warning:
-        pymc_model = create_model(f)
-
-    for output in FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES + ["predicted_observed_state"]:
-        assert pymc_model.named_vars_to_dims[output] == FILTER_OUTPUT_DIMS[output]
-
-
-def test_model_build_without_coords(load_dataset):
-    ss_mod = structural.LevelTrendComponent().build(verbose=False)
-    data = load_dataset("numpy")
-    with pm.Model() as mod:
-        P0_diag = pm.Exponential("P0_diag", 1, shape=(2,))
-        P0 = pm.Deterministic("P0", pt.diag(P0_diag))
-        initial_trend = pm.Normal("initial_trend", shape=(2,))
-        sigma_trend = pm.Exponential("sigma_trend", 1, shape=(2,))
-        ss_mod.build_statespace_graph(data, register_data=False)
-
-    assert mod.coords == {}
-
-
-@pytest.mark.parametrize("f, warning", func_inputs, ids=function_names)
-def test_data_index_is_coord(f, warning, create_model):
-    with warning:
-        pymc_model = create_model(f)
-    assert TIME_DIM in pymc_model.coords