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

tests/statespace/test_kalman_filter.py

@@ -1,326 +0,0 @@
-import numpy as np
-import pytensor
-import pytensor.tensor as pt
-import pytest
-
-from numpy.testing import assert_allclose, assert_array_less
-
-from pymc_extras.statespace.filters import (
-    KalmanSmoother,
-    SquareRootFilter,
-    StandardFilter,
-    UnivariateFilter,
-)
-from pymc_extras.statespace.filters.kalman_filter import BaseFilter
-from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
-    rng,
-)
-from tests.statespace.utilities.test_helpers import (
-    get_expected_shape,
-    get_sm_state_from_output_name,
-    initialize_filter,
-    make_test_inputs,
-    nile_test_test_helper,
-)
-
-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-6 if floatX.endswith("64") else 1e-3
-RTOL = 1e-6 if floatX.endswith("64") else 1e-3
-
-standard_inout = initialize_filter(StandardFilter())
-cholesky_inout = initialize_filter(SquareRootFilter())
-univariate_inout = initialize_filter(UnivariateFilter())
-
-f_standard = pytensor.function(*standard_inout, on_unused_input="ignore")
-f_cholesky = pytensor.function(*cholesky_inout, on_unused_input="ignore")
-f_univariate = pytensor.function(*univariate_inout, on_unused_input="ignore")
-
-filter_funcs = [f_standard, f_cholesky, f_univariate]
-
-filter_names = [
-    "StandardFilter",
-    "CholeskyFilter",
-    "UnivariateFilter",
-]
-
-output_names = [
-    "filtered_states",
-    "predicted_states",
-    "smoothed_states",
-    "filtered_covs",
-    "predicted_covs",
-    "smoothed_covs",
-    "log_likelihood",
-    "ll_obs",
-]
-
-
-def test_base_class_update_raises():
-    filter = BaseFilter()
-    inputs = [None] * 7
-    with pytest.raises(NotImplementedError):
-        filter.update(*inputs)
-
-
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-def test_output_shapes_one_state_one_observed(filter_func, rng):
-    p, m, r, n = 1, 1, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-    outputs = filter_func(*inputs)
-
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-def test_output_shapes_when_all_states_are_stochastic(filter_func, rng):
-    p, m, r, n = 1, 2, 2, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-
-    outputs = filter_func(*inputs)
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-def test_output_shapes_when_some_states_are_deterministic(filter_func, rng):
-    p, m, r, n = 1, 5, 2, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-
-    outputs = filter_func(*inputs)
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.fixture
-def f_standard_nd():
-    ksmoother = KalmanSmoother()
-    data = pt.tensor(name="data", dtype=floatX, shape=(None, None))
-    a0 = pt.vector(name="a0", dtype=floatX)
-    P0 = pt.matrix(name="P0", dtype=floatX)
-    c = pt.vector(name="c", dtype=floatX)
-    d = pt.vector(name="d", dtype=floatX)
-    Q = pt.tensor(name="Q", dtype=floatX, shape=(None, None, None))
-    H = pt.tensor(name="H", dtype=floatX, shape=(None, None, None))
-    T = pt.tensor(name="T", dtype=floatX, shape=(None, None, None))
-    R = pt.tensor(name="R", dtype=floatX, shape=(None, None, None))
-    Z = pt.tensor(name="Z", dtype=floatX, shape=(None, None, None))
-
-    inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
-
-    (
-        filtered_states,
-        predicted_states,
-        observed_states,
-        filtered_covs,
-        predicted_covs,
-        observed_covs,
-        ll_obs,
-    ) = StandardFilter().build_graph(*inputs)
-
-    smoothed_states, smoothed_covs = ksmoother.build_graph(T, R, Q, filtered_states, filtered_covs)
-
-    outputs = [
-        filtered_states,
-        predicted_states,
-        smoothed_states,
-        filtered_covs,
-        predicted_covs,
-        smoothed_covs,
-        ll_obs.sum(),
-        ll_obs,
-    ]
-
-    f_standard = pytensor.function(inputs, outputs)
-
-    return f_standard
-
-
-def test_output_shapes_with_time_varying_matrices(f_standard_nd, rng):
-    p, m, r, n = 1, 5, 2, 10
-    data, a0, P0, c, d, T, Z, R, H, Q = make_test_inputs(p, m, r, n, rng)
-    T = np.concatenate([np.expand_dims(T, 0)] * n, axis=0)
-    Z = np.concatenate([np.expand_dims(Z, 0)] * n, axis=0)
-    R = np.concatenate([np.expand_dims(R, 0)] * n, axis=0)
-    H = np.concatenate([np.expand_dims(H, 0)] * n, axis=0)
-    Q = np.concatenate([np.expand_dims(Q, 0)] * n, axis=0)
-
-    outputs = f_standard_nd(data, a0, P0, c, d, T, Z, R, H, Q)
-
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-def test_output_with_deterministic_observation_equation(filter_func, rng):
-    p, m, r, n = 1, 5, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-
-    outputs = filter_func(*inputs)
-
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize(
-    ("filter_func", "filter_name"), zip(filter_funcs, filter_names), ids=filter_names
-)
-def test_output_with_multiple_observed(filter_func, filter_name, rng):
-    p, m, r, n = 5, 5, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-
-    outputs = filter_func(*inputs)
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize(
-    ("filter_func", "filter_name"), zip(filter_funcs, filter_names), ids=filter_names
-)
-@pytest.mark.parametrize("p", [1, 5], ids=["univariate (p=1)", "multivariate (p=5)"])
-def test_missing_data(filter_func, filter_name, p, rng):
-    m, r, n = 5, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng, missing_data=1)
-
-    outputs = filter_func(*inputs)
-    for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
-
-
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-@pytest.mark.parametrize("output_idx", [(0, 2), (3, 5)], ids=["smoothed_states", "smoothed_covs"])
-def test_last_smoother_is_last_filtered(filter_func, output_idx, rng):
-    p, m, r, n = 1, 5, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-    outputs = filter_func(*inputs)
-
-    filtered = outputs[output_idx[0]]
-    smoothed = outputs[output_idx[1]]
-
-    assert_allclose(filtered[-1], smoothed[-1])
-
-
-@pytest.mark.parametrize(
-    "filter_func, filter_name", zip(filter_funcs, filter_names), ids=filter_names
-)
-@pytest.mark.parametrize("n_missing", [0, 5], ids=["n_missing=0", "n_missing=5"])
-@pytest.mark.skipif(floatX == "float32", reason="Tests are too sensitive for float32")
-def test_filters_match_statsmodel_output(filter_func, filter_name, n_missing, rng):
-    fit_sm_mod, [data, a0, P0, c, d, T, Z, R, H, Q] = nile_test_test_helper(rng, n_missing)
-    if filter_name == "CholeskyFilter":
-        P0 = np.linalg.cholesky(P0)
-    inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
-    outputs = filter_func(*inputs)
-
-    for output_idx, name in enumerate(output_names):
-        ref_val = get_sm_state_from_output_name(fit_sm_mod, name)
-        val_to_test = outputs[output_idx].squeeze()
-
-        if name == "smoothed_covs":
-            # TODO: The smoothed covariance matrices have large errors (1e-2) ONLY in the first few states -- no idea why.
-            assert_allclose(
-                val_to_test[5:],
-                ref_val[5:],
-                atol=ATOL,
-                rtol=RTOL,
-                err_msg=f"{name} does not match statsmodels",
-            )
-        elif name.startswith("predicted"):
-            # statsmodels doesn't throw away the T+1 forecast in the predicted states like we do
-            assert_allclose(
-                val_to_test,
-                ref_val[:-1],
-                atol=ATOL,
-                rtol=RTOL,
-                err_msg=f"{name} does not match statsmodels",
-            )
-        else:
-            # Need atol = 1e-7 for smoother tests to pass
-            assert_allclose(
-                val_to_test,
-                ref_val,
-                atol=ATOL,
-                rtol=RTOL,
-                err_msg=f"{name} does not match statsmodels",
-            )
-
-
-@pytest.mark.parametrize(
-    "filter_func, filter_name", zip(filter_funcs[:-1], filter_names[:-1]), ids=filter_names[:-1]
-)
-@pytest.mark.parametrize("n_missing", [0, 5], ids=["n_missing=0", "n_missing=5"])
-@pytest.mark.parametrize("obs_noise", [True, False])
-def test_all_covariance_matrices_are_PSD(filter_func, filter_name, n_missing, obs_noise, rng):
-    if (floatX == "float32") & (filter_name == "UnivariateFilter"):
-        # TODO: These tests all pass locally for me with float32 but they fail on the CI, so i'm just disabling them.
-        pytest.skip("Univariate filter not stable at half precision without measurement error")
-
-    fit_sm_mod, [data, a0, P0, c, d, T, Z, R, H, Q] = nile_test_test_helper(rng, n_missing)
-    if filter_name == "CholeskyFilter":
-        P0 = np.linalg.cholesky(P0)
-
-    H *= int(obs_noise)
-    inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
-    outputs = filter_func(*inputs)
-
-    for output_idx, name in zip([3, 4, 5], output_names[3:-2]):
-        cov_stack = outputs[output_idx]
-        w, v = np.linalg.eig(cov_stack)
-
-        assert_array_less(0, w, err_msg=f"Smallest eigenvalue of {name}: {min(w.ravel())}")
-        assert_allclose(
-            cov_stack,
-            np.swapaxes(cov_stack, -2, -1),
-            rtol=RTOL,
-            atol=ATOL,
-            err_msg=f"{name} is not symmetrical",
-        )
-
-
-@pytest.mark.parametrize(
-    "filter",
-    [StandardFilter, SquareRootFilter],
-    ids=["standard", "cholesky"],
-)
-def test_kalman_filter_jax(filter):
-    pytest.importorskip("jax")
-    from pymc.sampling.jax import get_jaxified_graph
-
-    # TODO: Add UnivariateFilter to test; need to figure out the broadcasting issue when 2nd data dim is defined
-
-    p, m, r, n = 1, 5, 1, 10
-    inputs, outputs = initialize_filter(filter(), mode="JAX", p=p, m=m, r=r, n=n)
-    inputs_np = make_test_inputs(p, m, r, n, rng)
-
-    f_jax = get_jaxified_graph(inputs, outputs)
-    f_pt = pytensor.function(inputs, outputs, mode="FAST_COMPILE")
-
-    jax_outputs = f_jax(*inputs_np)
-    pt_outputs = f_pt(*inputs_np)
-
-    for name, jax_res, pt_res in zip(output_names, jax_outputs, pt_outputs):
-        assert_allclose(jax_res, pt_res, atol=ATOL, rtol=RTOL, err_msg=f"{name} failed!")

tests/statespace/test_representation.py

@@ -1,175 +0,0 @@
-import unittest
-
-import numpy as np
-import pytensor
-import pytensor.tensor as pt
-
-from numpy.testing import assert_allclose
-
-from pymc_extras.statespace.core.representation import PytensorRepresentation
-from tests.statespace.utilities.shared_fixtures import TEST_SEED
-from tests.statespace.utilities.test_helpers import fast_eval, make_test_inputs
-
-floatX = pytensor.config.floatX
-atol = 1e-12 if floatX == "float64" else 1e-6
-
-
-def unpack_ssm_dims(ssm):
-    p = ssm.k_endog
-    m = ssm.k_states
-    r = ssm.k_posdef
-
-    return p, m, r
-
-
-class BasicFunctionality(unittest.TestCase):
-    def setUp(self):
-        self.rng = np.random.default_rng(TEST_SEED)
-
-    def test_numpy_to_pytensor(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        X = np.eye(5)
-        X_pt = ssm._numpy_to_pytensor("transition", X)
-        self.assertTrue(isinstance(X_pt, pt.TensorVariable))
-        assert_allclose(ssm["transition"].type.shape, X.shape)
-
-        assert ssm["transition"].name == "transition"
-
-    def test_default_shapes_full_rank(self):
-        ssm = PytensorRepresentation(k_endog=5, k_states=5, k_posdef=5)
-        p, m, r = unpack_ssm_dims(ssm)
-
-        assert_allclose(ssm["design"].type.shape, (p, m))
-        assert_allclose(ssm["transition"].type.shape, (m, m))
-        assert_allclose(ssm["selection"].type.shape, (m, r))
-        assert_allclose(ssm["state_cov"].type.shape, (r, r))
-        assert_allclose(ssm["obs_cov"].type.shape, (p, p))
-
-    def test_default_shapes_low_rank(self):
-        ssm = PytensorRepresentation(k_endog=5, k_states=5, k_posdef=2)
-        p, m, r = unpack_ssm_dims(ssm)
-
-        assert_allclose(ssm["design"].type.shape, (p, m))
-        assert_allclose(ssm["transition"].type.shape, (m, m))
-        assert_allclose(ssm["selection"].type.shape, (m, r))
-        assert_allclose(ssm["state_cov"].type.shape, (r, r))
-        assert_allclose(ssm["obs_cov"].type.shape, (p, p))
-
-    def test_matrix_assignment(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=2)
-
-        ssm["design", 0, 0] = 3.0
-        ssm["transition", 0, :] = 2.7
-        ssm["selection", -1, -1] = 9.9
-
-        assert_allclose(fast_eval(ssm["design"][0, 0]), 3.0, atol=atol)
-        assert_allclose(fast_eval(ssm["transition"][0, :]), 2.7, atol=atol)
-        assert_allclose(fast_eval(ssm["selection"][-1, -1]), 9.9, atol=atol)
-
-        assert ssm["design"].name == "design"
-        assert ssm["transition"].name == "transition"
-        assert ssm["selection"].name == "selection"
-
-    def test_build_representation_from_data(self):
-        p, m, r, n = 3, 6, 1, 10
-        inputs = [data, a0, P0, c, d, T, Z, R, H, Q] = make_test_inputs(
-            p, m, r, n, self.rng, missing_data=0
-        )
-
-        ssm = PytensorRepresentation(
-            k_endog=p,
-            k_states=m,
-            k_posdef=r,
-            design=Z,
-            transition=T,
-            selection=R,
-            state_cov=Q,
-            obs_cov=H,
-            initial_state=a0,
-            initial_state_cov=P0,
-            state_intercept=c,
-            obs_intercept=d,
-        )
-
-        names = [
-            "initial_state",
-            "initial_state_cov",
-            "state_intercept",
-            "obs_intercept",
-            "transition",
-            "design",
-            "selection",
-            "obs_cov",
-            "state_cov",
-        ]
-
-        for name, X in zip(names, inputs[1:]):
-            assert_allclose(X, fast_eval(ssm[name]), err_msg=name)
-
-        for name, X in zip(names, inputs[1:]):
-            assert ssm[name].name == name
-            assert_allclose(ssm[name].type.shape, X.shape, err_msg=f"{name} shape test")
-
-    def test_assign_time_varying_matrices(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=2)
-        n = 10
-
-        ssm["design", 0, 0] = 3.0
-        ssm["transition", 0, :] = 2.7
-        ssm["selection", -1, -1] = 9.9
-
-        ssm["state_intercept"] = np.zeros((n, 5))
-        ssm["state_intercept", :, 0] = np.arange(n)
-
-        assert_allclose(fast_eval(ssm["design"][0, 0]), 3.0, atol=atol)
-        assert_allclose(fast_eval(ssm["transition"][0, :]), 2.7, atol=atol)
-        assert_allclose(fast_eval(ssm["selection"][-1, -1]), 9.9, atol=atol)
-        assert_allclose(fast_eval(ssm["state_intercept"][:, 0]), np.arange(n), atol=atol)
-
-    def test_invalid_key_name_raises(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        with self.assertRaises(IndexError) as e:
-            X = ssm["invalid_key"]
-        msg = str(e.exception)
-        self.assertEqual(msg, "invalid_key is an invalid state space matrix name")
-
-    def test_non_string_key_raises(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        with self.assertRaises(IndexError) as e:
-            X = ssm[0]
-        msg = str(e.exception)
-        self.assertEqual(msg, "First index must the name of a valid state space matrix.")
-
-    def test_invalid_key_tuple_raises(self):
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        with self.assertRaises(IndexError) as e:
-            X = ssm[0, 1, 1]
-        msg = str(e.exception)
-        self.assertEqual(msg, "First index must the name of a valid state space matrix.")
-
-    def test_slice_statespace_matrix(self):
-        T = np.eye(5)
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1, transition=T)
-        T_out = ssm["transition", :3, :]
-        assert_allclose(T[:3], fast_eval(T_out))
-
-    def test_update_matrix_via_key(self):
-        T = np.eye(5)
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        ssm["transition"] = T
-
-        assert_allclose(T, fast_eval(ssm["transition"]))
-
-    def test_update_matrix_with_invalid_shape_raises(self):
-        T = np.eye(10)
-        ssm = PytensorRepresentation(k_endog=3, k_states=5, k_posdef=1)
-        with self.assertRaises(ValueError) as e:
-            ssm["transition"] = T
-        msg = str(e.exception)
-        self.assertEqual(
-            msg, "The last two dimensions of transition must be (5, 5), found (10, 10)"
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()