Add shared_state argument to CycleComponent

Author: jessegrabowski

pymc_extras/statespace/models/structural/components/cycle.py

@@ -43,6 +43,11 @@ class CycleComponent(Component):
         Names of the observed state variables. For univariate time series, defaults to ``["data"]``.
         For multivariate time series, specify a list of names for each endogenous variable.
 
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states. This argument has no effect if `k_endog` is 1.
+
     Notes
     -----
     The cycle component is very similar in implementation to the frequency domain seasonal component, expect that it
@@ -155,6 +160,7 @@ def __init__(
         dampen: bool = False,
         innovations: bool = True,
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
         if observed_state_names is None:
             observed_state_names = ["data"]
@@ -167,6 +173,7 @@ def __init__(
             cycle = int(cycle_length) if cycle_length is not None else "Estimate"
             name = f"Cycle[s={cycle}, dampen={dampen}, innovations={innovations}]"
 
+        self.share_states = share_states
         self.estimate_cycle_length = estimate_cycle_length
         self.cycle_length = cycle_length
         self.innovations = innovations
@@ -175,8 +182,8 @@ def __init__(
 
         k_endog = len(observed_state_names)
 
-        k_states = 2 * k_endog
-        k_posdef = 2 * k_endog
+        k_states = 2 if share_states else 2 * k_endog
+        k_posdef = 2 if share_states else 2 * k_endog
 
         obs_state_idx = np.zeros(k_states)
         obs_state_idx[slice(0, k_states, 2)] = 1
@@ -193,18 +200,22 @@ def __init__(
         )
 
     def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
         Z = np.array([1.0, 0.0]).reshape((1, -1))
-        design_matrix = block_diag(*[Z for _ in range(self.k_endog)])
+        design_matrix = block_diag(*[Z for _ in range(k_endog_effective)])
         self.ssm["design", :, :] = pt.as_tensor_variable(design_matrix)
 
         # selection matrix R defines structure of innovations (always identity for cycle components)
         # when innovations=False, state cov Q=0, hence R @ Q @ R.T = 0
         R = np.eye(2)  # 2x2 identity for each cycle component
-        selection_matrix = block_diag(*[R for _ in range(self.k_endog)])
+        selection_matrix = block_diag(*[R for _ in range(k_endog_effective)])
         self.ssm["selection", :, :] = pt.as_tensor_variable(selection_matrix)
 
         init_state = self.make_and_register_variable(
-            f"{self.name}", shape=(self.k_endog, 2) if self.k_endog > 1 else (self.k_states,)
+            f"{self.name}",
+            shape=(k_endog_effective, 2) if k_endog_effective > 1 else (self.k_states,),
         )
         self.ssm["initial_state", :] = init_state.ravel()
 
@@ -219,37 +230,45 @@ def make_symbolic_graph(self) -> None:
             rho = 1
 
         T = rho * _frequency_transition_block(lamb, j=1)
-        transition = block_diag(*[T for _ in range(self.k_endog)])
+        transition = block_diag(*[T for _ in range(k_endog_effective)])
         self.ssm["transition"] = pt.specify_shape(transition, (self.k_states, self.k_states))
 
         if self.innovations:
-            if self.k_endog == 1:
+            if k_endog_effective == 1:
                 sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=())
                 self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
             else:
                 sigma_cycle = self.make_and_register_variable(
-                    f"sigma_{self.name}", shape=(self.k_endog,)
+                    f"sigma_{self.name}", shape=(k_endog_effective,)
                 )
                 state_cov = block_diag(
-                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(self.k_endog)]
+                    *[pt.eye(2) * sigma_cycle[i] ** 2 for i in range(k_endog_effective)]
                 )
                 self.ssm["state_cov"] = pt.specify_shape(state_cov, (self.k_states, self.k_states))
         else:
             # explicitly set state cov to 0 when no innovations
             self.ssm["state_cov", :, :] = pt.zeros((self.k_posdef, self.k_posdef))
 
     def populate_component_properties(self):
-        self.state_names = [
-            f"{f}_{self.name}[{var_name}]" if self.k_endog > 1 else f"{f}_{self.name}"
-            for var_name in self.observed_state_names
-            for f in ["Cos", "Sin"]
-        ]
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        base_names = [f"{f}_{self.name}" for f in ["Cos", "Sin"]]
+
+        if self.share_states:
+            self.state_names = [f"{name}[shared]" for name in base_names]
+        else:
+            self.state_names = [
+                f"{name}[{var_name}]" if k_endog_effective > 1 else name
+                for var_name in self.observed_state_names
+                for name in base_names
+            ]
 
         self.param_names = [f"{self.name}"]
 
-        if self.k_endog == 1:
+        if k_endog_effective == 1:
             self.param_dims = {self.name: (f"state_{self.name}",)}
-            self.coords = {f"state_{self.name}": self.state_names}
+            self.coords = {f"state_{self.name}": base_names}
             self.param_info = {
                 f"{self.name}": {
                     "shape": (2,),
@@ -265,7 +284,7 @@ def populate_component_properties(self):
             }
             self.param_info = {
                 f"{self.name}": {
-                    "shape": (self.k_endog, 2),
+                    "shape": (k_endog_effective, 2),
                     "constraints": None,
                     "dims": (f"endog_{self.name}", f"state_{self.name}"),
                 }
@@ -274,22 +293,22 @@ def populate_component_properties(self):
         if self.estimate_cycle_length:
             self.param_names += [f"length_{self.name}"]
             self.param_info[f"length_{self.name}"] = {
-                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
                 "constraints": "Positive, non-zero",
-                "dims": None if self.k_endog == 1 else f"endog_{self.name}",
+                "dims": None if k_endog_effective == 1 else f"endog_{self.name}",
             }
 
         if self.dampen:
             self.param_names += [f"dampening_factor_{self.name}"]
             self.param_info[f"dampening_factor_{self.name}"] = {
-                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "shape": () if k_endog_effective == 1 else (k_endog_effective,),
                 "constraints": "0 < x ≤ 1",
-                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+                "dims": None if k_endog_effective == 1 else (f"endog_{self.name}",),
             }
 
         if self.innovations:
             self.param_names += [f"sigma_{self.name}"]
-            if self.k_endog == 1:
+            if k_endog_effective == 1:
                 self.param_info[f"sigma_{self.name}"] = {
                     "shape": (),
                     "constraints": "Positive",
@@ -298,7 +317,7 @@ def populate_component_properties(self):
             else:
                 self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
                 self.param_info[f"sigma_{self.name}"] = {
-                    "shape": (self.k_endog,),
+                    "shape": (k_endog_effective,),
                     "constraints": "Positive",
                     "dims": (f"endog_{self.name}",),
                 }

tests/statespace/models/structural/components/test_cycle.py

@@ -3,6 +3,8 @@
 
 from numpy.testing import assert_allclose
 from pytensor import config
+from pytensor.graph.basic import explicit_graph_inputs
+from scipy import linalg
 
 from pymc_extras.statespace.models import structural as st
 from pymc_extras.statespace.models.structural.utils import _frequency_transition_block
@@ -105,6 +107,27 @@ def test_cycle_multivariate_deterministic(rng):
     np.testing.assert_allclose(R, expected_R)
 
 
+def test_multivariate_cycle_with_shared(rng):
+    cycle = st.CycleComponent(
+        name="cycle",
+        cycle_length=12,
+        estimate_cycle_length=False,
+        innovations=False,
+        observed_state_names=["data_1", "data_2", "data_3"],
+        share_states=True,
+    )
+
+    assert cycle.state_names == ["Cos_cycle[shared]", "Sin_cycle[shared]"]
+    assert cycle.shock_names == []
+    assert cycle.param_names == ["cycle"]
+
+    params = {"cycle": np.array([1.0, 2.0], dtype=config.floatX)}
+    x, y = simulate_from_numpy_model(cycle, rng, params, steps=12 * 12)
+
+    np.testing.assert_allclose(y[:, 0], y[:, 1], atol=ATOL, rtol=RTOL)
+    np.testing.assert_allclose(y[:, 0], y[:, 2], atol=ATOL, rtol=RTOL)
+
+
 def test_cycle_multivariate_with_dampening(rng):
     """Test multivariate cycle component with dampening."""
     cycle = st.CycleComponent(
@@ -286,3 +309,90 @@ def test_add_multivariate_cycle_components_with_different_observed():
     for i in range(4):
         expected_R[2 * i : 2 * i + 2, 2 * i : 2 * i + 2] = np.eye(2)
     assert_allclose(R, expected_R)
+
+
+def test_add_multivariate_shared_and_not_shared():
+    cycle_shared = st.CycleComponent(
+        name="shared_cycle",
+        cycle_length=12,
+        estimate_cycle_length=False,
+        innovations=True,
+        observed_state_names=["gdp", "inflation", "unemployment"],
+        share_states=True,
+    )
+    cycle_individual = st.CycleComponent(
+        name="individual_cycle",
+        estimate_cycle_length=True,
+        innovations=False,
+        observed_state_names=["gdp", "inflation", "unemployment"],
+        dampen=True,
+    )
+    mod = (cycle_shared + cycle_individual).build(verbose=False)
+
+    assert mod.k_endog == 3
+    assert mod.k_states == 2 + 3 * 2
+    assert mod.k_posdef == 2 + 3 * 2
+
+    expected_states = [
+        "Cos_shared_cycle[shared]",
+        "Sin_shared_cycle[shared]",
+        "Cos_individual_cycle[gdp]",
+        "Sin_individual_cycle[gdp]",
+        "Cos_individual_cycle[inflation]",
+        "Sin_individual_cycle[inflation]",
+        "Cos_individual_cycle[unemployment]",
+        "Sin_individual_cycle[unemployment]",
+    ]
+
+    assert mod.state_names == expected_states
+    assert mod.shock_names == expected_states[:2]
+
+    assert mod.param_names == [
+        "shared_cycle",
+        "sigma_shared_cycle",
+        "individual_cycle",
+        "length_individual_cycle",
+        "dampening_factor_individual_cycle",
+        "P0",
+    ]
+
+    assert "endog_shared_cycle" not in mod.coords
+    assert mod.coords["state_shared_cycle"] == ["Cos_shared_cycle", "Sin_shared_cycle"]
+    assert mod.coords["state_individual_cycle"] == ["Cos_individual_cycle", "Sin_individual_cycle"]
+    assert mod.coords["endog_individual_cycle"] == ["gdp", "inflation", "unemployment"]
+
+    assert mod.param_info["shared_cycle"]["dims"] == ("state_shared_cycle",)
+    assert mod.param_info["shared_cycle"]["shape"] == (2,)
+
+    assert mod.param_info["sigma_shared_cycle"]["dims"] is None
+    assert mod.param_info["sigma_shared_cycle"]["shape"] == ()
+
+    assert mod.param_info["individual_cycle"]["dims"] == (
+        "endog_individual_cycle",
+        "state_individual_cycle",
+    )
+    assert mod.param_info["individual_cycle"]["shape"] == (3, 2)
+
+    params = {
+        "length_individual_cycle": 12.0,
+        "dampening_factor_individual_cycle": 0.95,
+    }
+    outputs = [mod.ssm["transition"], mod.ssm["design"], mod.ssm["selection"]]
+    T, Z, R = pytensor.function(
+        list(explicit_graph_inputs(outputs)),
+        outputs,
+        mode="FAST_COMPILE",
+    )(**params)
+
+    lamb = 2 * np.pi / 12  # dampening factor for individual cycle
+    transition_block = np.array(
+        [[np.cos(lamb), np.sin(lamb)], [-np.sin(lamb), np.cos(lamb)]], dtype=config.floatX
+    )
+    T_expected = linalg.block_diag(transition_block, *[0.95 * transition_block] * 3)
+    np.testing.assert_allclose(T, T_expected)
+
+    np.testing.assert_allclose(
+        Z, np.array([[1, 0, 1, 0, 0, 0, 0, 0], [1, 0, 0, 0, 1, 0, 0, 0], [1, 0, 0, 0, 0, 0, 1, 0]])
+    )
+
+    np.testing.assert_allclose(R, np.eye(8, dtype=config.floatX))