pymc-extras 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

tests/statespace/test_coord_assignment.py

@@ -1,181 +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.models.structural import LevelTrendComponent
-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
-
-
-def make_model(index):
-    n = len(index)
-    a = pd.DataFrame(index=index, columns=["A", "B", "C", "D"], data=np.arange(n * 4).reshape(n, 4))
-
-    mod = LevelTrendComponent(order=2, innovations_order=[0, 1])
-    ss_mod = mod.build(name="a", verbose=False)
-
-    initial_trend_dims, sigma_trend_dims, P0_dims = ss_mod.param_dims.values()
-    coords = ss_mod.coords
-
-    with pm.Model(coords=coords) as model:
-        P0_diag = pm.Gamma("P0_diag", alpha=5, beta=5)
-        P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states) * P0_diag, dims=P0_dims)
-
-        initial_trend = pm.Normal("initial_trend", dims=initial_trend_dims)
-        sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=50, dims=sigma_trend_dims)
-
-        with pytest.warns(UserWarning, match="No time index found on the supplied data"):
-            ss_mod.build_statespace_graph(
-                a["A"],
-                mode="JAX",
-            )
-    return model
-
-
-def test_integer_index():
-    index = np.arange(8).astype(int)
-    model = make_model(index)
-    assert TIME_DIM in model.coords
-    np.testing.assert_allclose(model.coords[TIME_DIM], index)
-
-
-def test_float_index_raises():
-    index = np.linspace(0, 1, 8)
-
-    with pytest.raises(IndexError, match="Provided index is not an integer index"):
-        make_model(index)
-
-
-def test_non_strictly_monotone_index_raises():
-    # Decreases
-    index = [0, 1, 2, 1, 2, 3]
-    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
-        make_model(index)
-
-    # Has gaps
-    index = [0, 1, 2, 3, 5, 6]
-    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
-        make_model(index)
-
-    # Has duplicates
-    index = [0, 1, 1, 2, 3, 4]
-    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
-        make_model(index)
-
-
-def test_multiindex_raises():
-    index = pd.MultiIndex.from_tuples([(0, 0), (1, 1), (2, 2), (3, 3)])
-    with pytest.raises(
-        NotImplementedError, match="MultiIndex panel data is not currently supported"
-    ):
-        make_model(index)

tests/statespace/test_distributions.py

@@ -1,270 +0,0 @@
-import numpy as np
-import pymc as pm
-import pytensor
-import pytensor.tensor as pt
-import pytest
-
-from numpy.testing import assert_allclose
-from scipy.stats import multivariate_normal
-
-from pymc_extras.statespace import structural
-from pymc_extras.statespace.filters.distributions import (
-    LinearGaussianStateSpace,
-    SequenceMvNormal,
-    _LinearGaussianStateSpace,
-)
-from pymc_extras.statespace.utils.constants import (
-    ALL_STATE_DIM,
-    OBS_STATE_DIM,
-    TIME_DIM,
-)
-from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
-    rng,
-)
-from tests.statespace.utilities.test_helpers import (
-    delete_rvs_from_model,
-    fast_eval,
-    load_nile_test_data,
-)
-
-floatX = pytensor.config.floatX
-
-# TODO: These are pretty loose because of all the stabilizing of covariance matrices that is done inside the kalman
-#  filters. When that is improved, this should be tightened.
-ATOL = 1e-5 if floatX.endswith("64") else 1e-4
-RTOL = 1e-5 if floatX.endswith("64") else 1e-4
-
-filter_names = [
-    "standard",
-    "cholesky",
-    "univariate",
-]
-
-
-@pytest.fixture(scope="session")
-def data():
-    return load_nile_test_data()
-
-
-@pytest.fixture(scope="session")
-def pymc_model(data):
-    with pm.Model() as mod:
-        data = pm.Data("data", data.values)
-        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,))
-
-    return mod
-
-
-@pytest.fixture(scope="session")
-def pymc_model_2(data):
-    coords = {
-        ALL_STATE_DIM: ["level", "trend"],
-        OBS_STATE_DIM: ["level"],
-        TIME_DIM: np.arange(101, dtype="int"),
-    }
-
-    with pm.Model(coords=coords) 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,))
-        sigma_me = pm.Exponential("sigma_error", 1)
-
-    return mod
-
-
-@pytest.fixture(scope="session")
-def ss_mod_me():
-    ss_mod = structural.LevelTrendComponent(order=2)
-    ss_mod += structural.MeasurementError(name="error")
-    ss_mod = ss_mod.build("data", verbose=False)
-
-    return ss_mod
-
-
-@pytest.fixture(scope="session")
-def ss_mod_no_me():
-    ss_mod = structural.LevelTrendComponent(order=2)
-    ss_mod = ss_mod.build("data", verbose=False)
-
-    return ss_mod
-
-
-@pytest.mark.parametrize("kfilter", filter_names, ids=filter_names)
-def test_loglike_vectors_agree(kfilter, pymc_model):
-    # TODO: This test might be flakey, I've gotten random failures
-    ss_mod = structural.LevelTrendComponent(order=2).build(
-        "data", verbose=False, filter_type=kfilter
-    )
-    with pymc_model:
-        ss_mod._insert_random_variables()
-        matrices = ss_mod.unpack_statespace()
-
-        filter_outputs = ss_mod.kalman_filter.build_graph(pymc_model["data"], *matrices)
-        filter_mus, pred_mus, obs_mu, filter_covs, pred_covs, obs_cov, ll = filter_outputs
-
-    test_ll = fast_eval(ll)
-
-    # TODO: BUG: Why does fast eval end up with a 2d output when filter is "single"?
-    obs_mu_np = obs_mu.eval()
-    obs_cov_np = fast_eval(obs_cov)
-    data_np = fast_eval(pymc_model["data"])
-
-    scipy_lls = []
-    for y, mu, cov in zip(data_np, obs_mu_np, obs_cov_np):
-        scipy_lls.append(multivariate_normal.logpdf(y, mean=mu, cov=cov))
-    assert_allclose(test_ll, np.array(scipy_lls).ravel(), atol=ATOL, rtol=RTOL)
-
-
-def test_sequence_mvn_distribution():
-    # Base Case
-    mu_sequence = pt.tensor("mu_sequence", shape=(100, 3))
-    cov_sequence = pt.tensor("cov_sequence", shape=(100, 3, 3))
-    logp = pt.tensor("logp", shape=(100,))
-
-    dist = SequenceMvNormal.dist(mu_sequence, cov_sequence, logp)
-    assert dist.type.shape == (100, 3)
-
-    # With batch dimension
-    mu_sequence = pt.tensor("mu_sequence", shape=(10, 100, 3))
-    cov_sequence = pt.tensor("cov_sequence", shape=(10, 100, 3, 3))
-    logp = pt.tensor(
-        "logp",
-        shape=(
-            10,
-            100,
-        ),
-    )
-
-    dist = SequenceMvNormal.dist(mu_sequence, cov_sequence, logp)
-    assert dist.type.shape == (10, 100, 3)
-
-
-@pytest.mark.parametrize("output_name", ["states_latent", "states_observed"])
-def test_lgss_distribution_from_steps(output_name, ss_mod_me, pymc_model_2):
-    with pymc_model_2:
-        ss_mod_me._insert_random_variables()
-        matrices = ss_mod_me.unpack_statespace()
-
-        # pylint: disable=unpacking-non-sequence
-        latent_states, obs_states = LinearGaussianStateSpace("states", *matrices, steps=100)
-        # pylint: enable=unpacking-non-sequence
-
-        idata = pm.sample_prior_predictive(draws=10)
-        delete_rvs_from_model(["states_latent", "states_observed", "states_combined"])
-
-    assert idata.prior.coords["states_latent_dim_0"].shape == (101,)
-    assert not np.any(np.isnan(idata.prior[output_name].values))
-
-
-@pytest.mark.parametrize("output_name", ["states_latent", "states_observed"])
-def test_lgss_distribution_with_dims(output_name, ss_mod_me, pymc_model_2):
-    with pymc_model_2:
-        ss_mod_me._insert_random_variables()
-        matrices = ss_mod_me.unpack_statespace()
-
-        # pylint: disable=unpacking-non-sequence
-        latent_states, obs_states = LinearGaussianStateSpace(
-            "states",
-            *matrices,
-            steps=100,
-            dims=[TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM],
-            sequence_names=[],
-            k_endog=ss_mod_me.k_endog,
-        )
-        # pylint: enable=unpacking-non-sequence
-        idata = pm.sample_prior_predictive(draws=10)
-        delete_rvs_from_model(["states_latent", "states_observed", "states_combined"])
-
-    assert idata.prior.coords["time"].shape == (101,)
-    assert all(
-        [dim in idata.prior.states_latent.coords.keys() for dim in [TIME_DIM, ALL_STATE_DIM]]
-    )
-    assert all(
-        [dim in idata.prior.states_observed.coords.keys() for dim in [TIME_DIM, OBS_STATE_DIM]]
-    )
-    assert not np.any(np.isnan(idata.prior[output_name].values))
-
-
-@pytest.mark.parametrize("output_name", ["states_latent", "states_observed"])
-def test_lgss_with_time_varying_inputs(output_name, rng):
-    X = rng.random(size=(10, 3), dtype=floatX)
-    ss_mod = structural.LevelTrendComponent() + structural.RegressionComponent(
-        name="exog", k_exog=3
-    )
-    mod = ss_mod.build("data", verbose=False)
-
-    coords = {
-        ALL_STATE_DIM: ["level", "trend", "beta_1", "beta_2", "beta_3"],
-        OBS_STATE_DIM: ["level"],
-        TIME_DIM: np.arange(10, dtype="int"),
-    }
-
-    with pm.Model(coords=coords):
-        exog_data = pm.Data("data_exog", X)
-        P0_diag = pm.Exponential("P0_diag", 1, shape=(mod.k_states,))
-        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,))
-        beta_exog = pm.Normal("beta_exog", shape=(3,))
-
-        mod._insert_random_variables()
-        mod._insert_data_variables()
-        matrices = mod.unpack_statespace()
-
-        # pylint: disable=unpacking-non-sequence
-        latent_states, obs_states = LinearGaussianStateSpace(
-            "states",
-            *matrices,
-            steps=9,
-            sequence_names=["d", "Z"],
-            dims=[TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM],
-        )
-        # pylint: enable=unpacking-non-sequence
-        idata = pm.sample_prior_predictive(draws=10)
-
-    assert idata.prior.coords["time"].shape == (10,)
-    assert all(
-        [dim in idata.prior.states_latent.coords.keys() for dim in [TIME_DIM, ALL_STATE_DIM]]
-    )
-    assert all(
-        [dim in idata.prior.states_observed.coords.keys() for dim in [TIME_DIM, OBS_STATE_DIM]]
-    )
-    assert not np.any(np.isnan(idata.prior[output_name].values))
-
-
-def test_lgss_signature():
-    # Base case
-    x0 = pt.tensor("x0", shape=(None,))
-    P0 = pt.tensor("P0", shape=(None, None))
-    c = pt.tensor("c", shape=(None,))
-    d = pt.tensor("d", shape=(None,))
-    T = pt.tensor("T", shape=(None, None))
-    Z = pt.tensor("Z", shape=(None, None))
-    R = pt.tensor("R", shape=(None, None))
-    H = pt.tensor("H", shape=(None, None))
-    Q = pt.tensor("Q", shape=(None, None))
-
-    lgss = _LinearGaussianStateSpace.dist(x0, P0, c, d, T, Z, R, H, Q, steps=100)
-    assert (
-        lgss.owner.op.extended_signature
-        == "(s),(s,s),(s),(p),(s,s),(p,s),(s,r),(p,p),(r,r),[rng]->[rng],(t,n)"
-    )
-    assert lgss.owner.op.ndim_supp == 2
-    assert lgss.owner.op.ndims_params == [1, 2, 1, 1, 2, 2, 2, 2, 2]
-
-    # Case with time-varying matrices
-    T = pt.tensor("T", shape=(None, None, None))
-    lgss = _LinearGaussianStateSpace.dist(
-        x0, P0, c, d, T, Z, R, H, Q, steps=100, sequence_names=["T"]
-    )
-
-    assert (
-        lgss.owner.op.extended_signature
-        == "(s),(s,s),(s),(p),(t,s,s),(p,s),(s,r),(p,p),(r,r),[rng]->[rng],(t,n)"
-    )
-    assert lgss.owner.op.ndim_supp == 2
-    assert lgss.owner.op.ndims_params == [1, 2, 1, 1, 3, 2, 2, 2, 2]