Allow exogenous variables in BayesianVARMAX

Author: jessegrabowski

pymc_extras/statespace/models/VARMAX.py

@@ -14,11 +14,13 @@
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     AR_PARAM_DIM,
+    EXOGENOUS_DIM,
     MA_PARAM_DIM,
     OBS_STATE_AUX_DIM,
     OBS_STATE_DIM,
     SHOCK_AUX_DIM,
     SHOCK_DIM,
+    TIME_DIM,
 )
 
 floatX = pytensor.config.floatX
@@ -264,6 +266,14 @@ def param_names(self):
             names.remove("ar_params")
         if self.q == 0:
             names.remove("ma_params")
+
+        # Add exogenous regression coefficents rather than remove, since we might have to handle
+        # several (if self.exog_state_names is a dict)
+        if isinstance(self.exog_state_names, list):
+            names.append("beta_exog")
+        elif isinstance(self.exog_state_names, dict):
+            names.extend([f"beta_{name}" for name in self.exog_state_names.keys()])
+
         return names
 
     @property
@@ -295,11 +305,49 @@ def param_info(self) -> dict[str, dict[str, Any]]:
             },
         }
 
+        if isinstance(self.exog_state_names, list):
+            k_exog = len(self.exog_state_names)
+            info["beta_exog"] = {
+                "shape": (self.k_endog, k_exog),
+                "constraints": "None",
+            }
+
+        elif isinstance(self.exog_state_names, dict):
+            for name, exog_names in self.exog_state_names.items():
+                k_exog = len(exog_names)
+                info[f"beta_{name}"] = {
+                    "shape": (k_exog,),
+                    "constraints": "None",
+                }
+
         for name in self.param_names:
             info[name]["dims"] = self.param_dims[name]
 
         return {name: info[name] for name in self.param_names}
 
+    @property
+    def data_info(self) -> dict[str, dict[str, Any]]:
+        info = None
+
+        if isinstance(self.exog_state_names, list):
+            info = {
+                "exogenous_data": {
+                    "dims": (TIME_DIM, EXOGENOUS_DIM),
+                    "shape": (None, self.k_exog),
+                }
+            }
+
+        elif isinstance(self.exog_state_names, dict):
+            info = {
+                f"{endog_state}_exogenous_data": {
+                    "dims": (TIME_DIM, f"{EXOGENOUS_DIM}_{endog_state}"),
+                    "shape": (None, len(exog_names)),
+                }
+                for endog_state, exog_names in self.exog_state_names.items()
+            }
+
+        return info
+
     @property
     def data_names(self) -> list[str]:
         if isinstance(self.exog_state_names, list):
@@ -312,10 +360,10 @@ def data_names(self) -> list[str]:
     def state_names(self):
         state_names = self.endog_names.copy()
         state_names += [
-            f"L{i + 1}.{state}" for i in range(self.p - 1) for state in self.endog_names
+            f"L{i + 1}_{state}" for i in range(self.p - 1) for state in self.endog_names
         ]
         state_names += [
-            f"L{i + 1}.{state}_innov" for i in range(self.q) for state in self.endog_names
+            f"L{i + 1}_{state}_innov" for i in range(self.q) for state in self.endog_names
         ]
 
         return state_names
@@ -340,6 +388,12 @@ def coords(self) -> dict[str, Sequence]:
         if self.q > 0:
             coords.update({MA_PARAM_DIM: list(range(1, self.q + 1))})
 
+        if isinstance(self.exog_state_names, list):
+            coords[EXOGENOUS_DIM] = self.exog_state_names
+        elif isinstance(self.exog_state_names, dict):
+            for name, exog_names in self.exog_state_names.items():
+                coords[f"{EXOGENOUS_DIM}_{name}"] = exog_names
+
         return coords
 
     @property
@@ -363,6 +417,14 @@ def param_dims(self):
             del coord_map["P0"]
             del coord_map["x0"]
 
+        if isinstance(self.exog_state_names, list):
+            coord_map["beta_exog"] = (OBS_STATE_DIM, EXOGENOUS_DIM)
+        elif isinstance(self.exog_state_names, dict):
+            # If each state has its own exogenous variables, each parameter needs it own dim, since we expect the
+            # dim labels to all be different (otherwise we'd be in the list case).
+            for name in self.exog_state_names.keys():
+                coord_map[f"beta_{name}"] = (f"{EXOGENOUS_DIM}_{name}",)
+
         return coord_map
 
     def add_default_priors(self):
@@ -450,6 +512,61 @@ def make_symbolic_graph(self) -> None:
         )
         self.ssm["state_cov", :, :] = state_cov
 
+        if self.exog_state_names is not None:
+            if isinstance(self.exog_state_names, list):
+                beta_exog = self.make_and_register_variable(
+                    "beta_exog", shape=(self.k_posdef, self.k_exog), dtype=floatX
+                )
+                exog_data = self.make_and_register_data(
+                    "exogenous_data", shape=(None, self.k_exog), dtype=floatX
+                )
+
+                obs_intercept = exog_data @ beta_exog.T
+
+            elif isinstance(self.exog_state_names, dict):
+                obs_components = []
+                for i, name in enumerate(self.endog_names):
+                    if name in self.exog_state_names:
+                        k_exog = len(self.exog_state_names[name])
+                        beta_exog = self.make_and_register_variable(
+                            f"beta_{name}", shape=(k_exog,), dtype=floatX
+                        )
+                        exog_data = self.make_and_register_data(
+                            f"{name}_exogenous_data", shape=(None, k_exog), dtype=floatX
+                        )
+                        obs_components.append(pt.expand_dims(exog_data @ beta_exog, axis=-1))
+                    else:
+                        obs_components.append(pt.zeros((1, 1), dtype=floatX))
+
+                # TODO: Replace all of this with pt.concat_with_broadcast once PyMC works with pytensor >= 2.32
+
+                # If there were any zeros, they need to be broadcast against the non-zeros.
+                # Core shape is the last dim, the time dim is always broadcast
+                non_concat_shape = [1, None]
+
+                # Look for the first non-zero component to get the shape from
+                for tensor_inp in obs_components:
+                    for i, (bcast, sh) in enumerate(
+                        zip(tensor_inp.type.broadcastable, tensor_inp.shape)
+                    ):
+                        if bcast or i == 1:
+                            continue
+                        non_concat_shape[i] = sh
+
+                assert non_concat_shape.count(None) == 1
+
+                bcast_tensor_inputs = []
+                for tensor_inp in obs_components:
+                    non_concat_shape[1] = tensor_inp.shape[1]
+                    bcast_tensor_inputs.append(pt.broadcast_to(tensor_inp, non_concat_shape))
+
+                obs_intercept = pt.join(1, *bcast_tensor_inputs)
+
+            else:
+                raise NotImplementedError()
+
+            self.ssm["obs_intercept"] = obs_intercept
+
         if self.stationary_initialization:
             # Solve for matrix quadratic for P0
             T = self.ssm["transition"]

tests/statespace/models/test_VARMAX.py

@@ -9,6 +9,7 @@
 import statsmodels.api as sm
 
 from numpy.testing import assert_allclose, assert_array_less
+from pymc.model.transform.optimization import freeze_dims_and_data
 
 from pymc_extras.statespace import BayesianVARMAX
 from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG
@@ -203,6 +204,10 @@ def test_create_varmax_with_exogenous(data):
     assert mod.k_exog == 2
     assert mod.exog_state_names == ["exogenous_0", "exogenous_1"]
     assert mod.data_names == ["exogenous_data"]
+    assert mod.param_dims["beta_exog"] == ("observed_state", "exogenous")
+    assert mod.coords["exogenous"] == ["exogenous_0", "exogenous_1"]
+    assert mod.param_info["beta_exog"]["shape"] == (mod.k_endog, 2)
+    assert mod.param_info["beta_exog"]["dims"] == ("observed_state", "exogenous")
 
     # Case 2: exog_state_names as list, k_exog is None
     mod = BayesianVARMAX(
@@ -216,6 +221,10 @@ def test_create_varmax_with_exogenous(data):
     assert mod.k_exog == 2
     assert mod.exog_state_names == ["foo", "bar"]
     assert mod.data_names == ["exogenous_data"]
+    assert mod.param_dims["beta_exog"] == ("observed_state", "exogenous")
+    assert mod.coords["exogenous"] == ["foo", "bar"]
+    assert mod.param_info["beta_exog"]["shape"] == (mod.k_endog, 2)
+    assert mod.param_info["beta_exog"]["dims"] == ("observed_state", "exogenous")
 
     # Case 3: k_exog as int, exog_state_names as list (matching)
     mod = BayesianVARMAX(
@@ -230,6 +239,10 @@ def test_create_varmax_with_exogenous(data):
     assert mod.k_exog == 2
     assert mod.exog_state_names == ["a", "b"]
     assert mod.data_names == ["exogenous_data"]
+    assert mod.param_dims["beta_exog"] == ("observed_state", "exogenous")
+    assert mod.coords["exogenous"] == ["a", "b"]
+    assert mod.param_info["beta_exog"]["shape"] == (mod.k_endog, 2)
+    assert mod.param_info["beta_exog"]["dims"] == ("observed_state", "exogenous")
 
     # Case 4: k_exog as dict, exog_state_names is None
     k_exog = {"observed_0": 2, "observed_1": 1, "observed_2": 0}
@@ -252,6 +265,25 @@ def test_create_varmax_with_exogenous(data):
         "observed_1_exogenous_data",
         "observed_2_exogenous_data",
     ]
+    assert mod.param_dims["beta_observed_0"] == ("exogenous_observed_0",)
+    assert mod.param_dims["beta_observed_1"] == ("exogenous_observed_1",)
+    assert (
+        "beta_observed_2" not in mod.param_dims
+        or mod.param_info.get("beta_observed_2") is None
+        or mod.param_info.get("beta_observed_2", {}).get("shape", (0,))[0] == 0
+    )
+
+    assert mod.coords["exogenous_observed_0"] == [
+        "observed_0_exogenous_0",
+        "observed_0_exogenous_1",
+    ]
+    assert mod.coords["exogenous_observed_1"] == ["observed_1_exogenous_0"]
+    assert "exogenous_observed_2" in mod.coords and mod.coords["exogenous_observed_2"] == []
+
+    assert mod.param_info["beta_observed_0"]["shape"] == (2,)
+    assert mod.param_info["beta_observed_0"]["dims"] == ("exogenous_observed_0",)
+    assert mod.param_info["beta_observed_1"]["shape"] == (1,)
+    assert mod.param_info["beta_observed_1"]["dims"] == ("exogenous_observed_1",)
 
     # Case 5: exog_state_names as dict, k_exog is None
     exog_state_names = {"observed_0": ["a", "b"], "observed_1": ["c"], "observed_2": []}
@@ -270,6 +302,22 @@ def test_create_varmax_with_exogenous(data):
         "observed_1_exogenous_data",
         "observed_2_exogenous_data",
     ]
+    assert mod.param_dims["beta_observed_0"] == ("exogenous_observed_0",)
+    assert mod.param_dims["beta_observed_1"] == ("exogenous_observed_1",)
+    assert (
+        "beta_observed_2" not in mod.param_dims
+        or mod.param_info.get("beta_observed_2") is None
+        or mod.param_info.get("beta_observed_2", {}).get("shape", (0,))[0] == 0
+    )
+
+    assert mod.coords["exogenous_observed_0"] == ["a", "b"]
+    assert mod.coords["exogenous_observed_1"] == ["c"]
+    assert "exogenous_observed_2" in mod.coords and mod.coords["exogenous_observed_2"] == []
+
+    assert mod.param_info["beta_observed_0"]["shape"] == (2,)
+    assert mod.param_info["beta_observed_0"]["dims"] == ("exogenous_observed_0",)
+    assert mod.param_info["beta_observed_1"]["shape"] == (1,)
+    assert mod.param_info["beta_observed_1"]["dims"] == ("exogenous_observed_1",)
 
     # Case 6: k_exog as dict, exog_state_names as dict (matching)
     k_exog = {"observed_0": 2, "observed_1": 1}
@@ -286,6 +334,14 @@ def test_create_varmax_with_exogenous(data):
     assert mod.k_exog == k_exog
     assert mod.exog_state_names == exog_state_names
     assert mod.data_names == ["observed_0_exogenous_data", "observed_1_exogenous_data"]
+    assert mod.param_dims["beta_observed_0"] == ("exogenous_observed_0",)
+    assert mod.param_dims["beta_observed_1"] == ("exogenous_observed_1",)
+    assert mod.coords["exogenous_observed_0"] == ["a", "b"]
+    assert mod.coords["exogenous_observed_1"] == ["c"]
+    assert mod.param_info["beta_observed_0"]["shape"] == (2,)
+    assert mod.param_info["beta_observed_0"]["dims"] == ("exogenous_observed_0",)
+    assert mod.param_info["beta_observed_1"]["shape"] == (1,)
+    assert mod.param_info["beta_observed_1"]["dims"] == ("exogenous_observed_1",)
 
     # Error: k_exog as int, exog_state_names as list (length mismatch)
     with pytest.raises(
@@ -348,3 +404,108 @@ def test_create_varmax_with_exogenous(data):
             measurement_error=False,
             stationary_initialization=False,
         )
+
+
+@pytest.mark.parametrize(
+    "k_exog, exog_state_names",
+    [
+        (2, None),
+        (None, ["foo", "bar"]),
+        (None, {"y1": ["a", "b"], "y2": ["c"]}),
+    ],
+    ids=["k_exog_int", "exog_state_names_list", "exog_state_names_dict"],
+)
+@pytest.mark.filterwarnings("ignore::UserWarning")
+def test_varmax_with_exog(rng, k_exog, exog_state_names):
+    endog_names = ["y1", "y2", "y3"]
+    n_obs = 50
+    time_idx = pd.date_range(start="2020-01-01", periods=n_obs, freq="D")
+
+    y = rng.normal(size=(n_obs, len(endog_names)))
+    df = pd.DataFrame(y, columns=endog_names, index=time_idx).astype(floatX)
+
+    if isinstance(exog_state_names, dict):
+        exog_data = {
+            f"{name}_exogenous_data": pd.DataFrame(
+                rng.normal(size=(n_obs, len(exog_names))).astype(floatX),
+                columns=exog_names,
+                index=time_idx,
+            )
+            for name, exog_names in exog_state_names.items()
+        }
+    else:
+        exog_names = exog_state_names or [f"exogenous_{i}" for i in range(k_exog)]
+        exog_data = {
+            "exogenous_data": pd.DataFrame(
+                rng.normal(size=(n_obs, k_exog or len(exog_state_names))).astype(floatX),
+                columns=exog_names,
+                index=time_idx,
+            )
+        }
+
+    mod = BayesianVARMAX(
+        endog_names=endog_names,
+        order=(1, 0),
+        k_exog=k_exog,
+        exog_state_names=exog_state_names,
+        verbose=True,
+        measurement_error=False,
+        stationary_initialization=False,
+        mode="JAX",
+    )
+
+    with pm.Model(coords=mod.coords) as m:
+        for var_name, data in exog_data.items():
+            pm.Data(var_name, data, dims=mod.data_info[var_name]["dims"])
+
+        x0 = pm.Deterministic("x0", pt.zeros(mod.k_states), dims=mod.param_dims["x0"])
+        P0_diag = pm.Exponential("P0_diag", 1.0, dims=mod.param_dims["P0"][0])
+        P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=mod.param_dims["P0"])
+
+        ar_params = pm.Normal("ar_params", mu=0, sigma=1, dims=mod.param_dims["ar_params"])
+        state_cov_diag = pm.Exponential("state_cov_diag", 1.0, dims=mod.param_dims["state_cov"][0])
+        state_cov = pm.Deterministic(
+            "state_cov", pt.diag(state_cov_diag), dims=mod.param_dims["state_cov"]
+        )
+
+        # Exogenous priors
+        if isinstance(mod.exog_state_names, list):
+            beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=mod.param_dims["beta_exog"])
+        elif isinstance(mod.exog_state_names, dict):
+            for name in mod.exog_state_names:
+                if mod.exog_state_names.get(name):
+                    pm.Normal(f"beta_{name}", mu=0, sigma=1, dims=mod.param_dims[f"beta_{name}"])
+
+        mod.build_statespace_graph(data=df)
+
+    with freeze_dims_and_data(m):
+        prior = pm.sample_prior_predictive(
+            draws=10, random_seed=rng, compile_kwargs={"mode": "JAX"}
+        )
+
+    prior_cond = mod.sample_conditional_prior(prior, mvn_method="eigh")
+    beta_dot_data = prior_cond.filtered_prior_observed.values - prior_cond.filtered_prior.values
+
+    if isinstance(exog_state_names, list) or k_exog is not None:
+        beta = prior.prior.beta_exog
+        assert beta.shape == (1, 10, 3, 2)
+
+        np.testing.assert_allclose(
+            beta_dot_data,
+            np.einsum("tx,...sx->...ts", exog_data["exogenous_data"].values, beta),
+            atol=1e-2,
+        )
+
+    elif isinstance(exog_state_names, dict):
+        assert prior.prior.beta_y1.shape == (1, 10, 2)
+        assert prior.prior.beta_y2.shape == (1, 10, 1)
+
+        obs_intercept = [
+            np.einsum("tx,...x->...t", exog_data[f"{name}_exogenous_data"].values, beta)
+            for name, beta in zip(["y1", "y2"], [prior.prior.beta_y1, prior.prior.beta_y2])
+        ]
+
+        # y3 has no exogenous variables
+        obs_intercept.append(np.zeros_like(obs_intercept[0]))
+
+        np.testing.assert_allclose(beta_dot_data, np.stack(obs_intercept, axis=-1), atol=1e-2)