pymc-extras 0.2.0__py3-none-any.whl

tests/statespace/test_representation.py

@@ -0,0 +1,175 @@
+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()