removed commented code and tests, updated validate names test to be reused in both endog and some exog cases

Author: Dekermanjian

pymc_extras/statespace/models/DFM.py

@@ -5,7 +5,7 @@
 import pytensor.tensor as pt
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
-from pymc_extras.statespace.models.utilities import _validate_endog_names, make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -279,7 +279,7 @@ def __init__(
 
         """
 
-        k_endog = _validate_endog_names(endog_names)
+        k_endog = validate_names(endog_names, var_name="endog_names", optional=False)
         self.endog_names = endog_names
         self.k_endog = k_endog
         self.k_factors = k_factors
@@ -291,7 +291,7 @@ def __init__(
         if exog_names is not None:
             self.shared_exog_states = shared_exog_states
             self.exog_innovations = exog_innovations
-            k_exog = len(exog_names) if exog_names is not None else 0
+            k_exog = validate_names(exog_names, var_name="exog_names", optional=True) or 0
             self.k_exog = k_exog
             self.exog_names = exog_names
         else:

pymc_extras/statespace/models/ETS.py

@@ -9,7 +9,7 @@
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
-from pymc_extras.statespace.models.utilities import _validate_endog_names, make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -261,7 +261,7 @@ def __init__(
         if self.seasonal and self.seasonal_periods is None:
             raise ValueError("If seasonal is True, seasonal_periods must be provided.")
 
-        k_endog = _validate_endog_names(endog_names)
+        k_endog = validate_names(endog_names, var_name="endog_names", optional=False)
         self.endog_names = list(endog_names)
 
         if dense_innovation_covariance and k_endog == 1:

pymc_extras/statespace/models/SARIMAX.py

@@ -12,6 +12,7 @@
     make_default_coords,
     make_harvey_state_names,
     make_SARIMA_transition_matrix,
+    validate_names,
 )
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
@@ -207,10 +208,7 @@ def __init__(
         if seasonal_order is None:
             seasonal_order = (0, 0, 0, 0)
 
-        if exog_state_names is None:
-            k_exog = 0
-        else:
-            k_exog = len(exog_state_names)
+        k_exog = validate_names(exog_state_names, var_name="exog_state_names", optional=True) or 0
 
         self.exog_state_names = exog_state_names
         self.k_exog = k_exog

pymc_extras/statespace/models/VARMAX.py

@@ -9,7 +9,7 @@
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
-from pymc_extras.statespace.models.utilities import _validate_endog_names, make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -152,7 +152,7 @@ def __init__(
 
         """
 
-        k_endog = _validate_endog_names(endog_names)
+        k_endog = validate_names(endog_names, var_name="endog_names", optional=False)
 
         needs_exog_data = False
 

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

@@ -3,6 +3,7 @@
 from pytensor import tensor as pt
 
 from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.models.utilities import validate_names
 from pymc_extras.statespace.utils.constants import TIME_DIM
 
 
@@ -107,7 +108,6 @@ class RegressionComponent(Component):
 
     def __init__(
         self,
-        # k_exog: int | None = None,
         name: str | None = "regression",
         state_names: list[str] | None = None,
         observed_state_names: list[str] | None = None,
@@ -120,7 +120,7 @@ def __init__(
             observed_state_names = ["data"]
 
         self.innovations = innovations
-        k_exog = self._handle_input_data(state_names, name)
+        k_exog = self._handle_input_data(state_names)
 
         k_states = k_exog
         k_endog = len(observed_state_names)
@@ -140,17 +140,8 @@ def __init__(
             obs_state_idxs=np.ones(k_states),
         )
 
-    @staticmethod
-    def _get_state_names(state_names: list[str] | None, name: str):
-        if state_names is None:
-            raise ValueError("Must specify state_names")
-        else:
-            k_exog = len(state_names)
-
-        return k_exog, state_names
-
-    def _handle_input_data(self, state_names: list[str] | None, name) -> int:
-        k_exog, state_names = self._get_state_names(state_names, name)
+    def _handle_input_data(self, state_names: list[str] | None) -> int:
+        k_exog = validate_names(state_names, var_name="state_names", optional=False)
         self.state_names = state_names
 
         return k_exog

pymc_extras/statespace/models/utilities.py

@@ -672,9 +672,9 @@ def get_exog_dims_from_idata(exog_name, idata):
     return exog_dims
 
 
-def _validate_endog_names(endog_names) -> int:
-    if endog_names is None:
-        raise ValueError("Must specify endog_names")
-    else:
-        k_endog = len(endog_names)
-    return k_endog
+def validate_names(names: list[str], var_name: str, optional: bool = True) -> int | None:
+    if names is None:
+        if optional:
+            return None
+        raise ValueError(f"Must specify {var_name}")
+    return len(names)

tests/statespace/models/test_SARIMAX.py

@@ -422,9 +422,6 @@ def test_invalid_order_raises(order, name):
         BayesianSARIMAX(order=(p, 0, q), seasonal_order=(P, 0, Q, 4))
 
 
-@pytest.mark.filterwarnings(
-    "ignore::RuntimeWarning"
-)  # Needed this due to RuntimeWarning: divide by zero encountered in matmul
 def test_SARIMA_with_exogenous(rng, mock_sample):
     # Note something odd is happening with stationary_initialization where the matrix is not singular
     ss_mod = BayesianSARIMAX(

tests/statespace/models/test_VARMAX.py

@@ -201,23 +201,6 @@ def test_forecast(varma_mod, idata, rng):
 
 
 class TestVARMAXWithExogenous:
-    # def test_create_varmax_with_exogenous_k_exog_int(self, data):
-    #     mod = BayesianVARMAX(
-    #         endog_names=["realgdp", "realcons", "realinv"],
-    #         order=(1, 0),
-    #         exog_state_names=["exogenous_0", "exogenous_1"],
-    #         verbose=False,
-    #         measurement_error=False,
-    #         stationary_initialization=False,
-    #     )
-    #     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")
-
     def test_create_varmax_with_exogenous_list_of_names(self, data):
         mod = BayesianVARMAX(
             endog_names=["realgdp", "realcons", "realinv"],
@@ -252,47 +235,6 @@ def test_create_varmax_with_exogenous_both_defined_correctly(self, data):
         assert mod.param_info["beta_exog"]["shape"] == (mod.k_endog, 2)
         assert mod.param_info["beta_exog"]["dims"] == ("observed_state", "exogenous")
 
-    # def test_create_varmax_with_exogenous_k_exog_dict(self, data):
-    #     k_exog = {"observed_0": 2, "observed_1": 1, "observed_2": 0}
-    #     mod = BayesianVARMAX(
-    #         endog_names=["observed_0", "observed_1", "observed_2"],
-    #         order=(1, 0),
-    #         exog_state_names=k_exog,
-    #         verbose=False,
-    #         measurement_error=False,
-    #         stationary_initialization=False,
-    #     )
-    #     assert mod.k_exog == k_exog
-    #     assert mod.exog_state_names == {
-    #         "observed_0": ["observed_0_exogenous_0", "observed_0_exogenous_1"],
-    #         "observed_1": ["observed_1_exogenous_0"],
-    #         "observed_2": [],
-    #     }
-    #     assert mod.data_names == [
-    #         "observed_0_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",)
-
     def test_create_varmax_with_exogenous_exog_names_dict(self, data):
         exog_state_names = {"observed_0": ["a", "b"], "observed_1": ["c"], "observed_2": []}
         mod = BayesianVARMAX(
@@ -327,30 +269,6 @@ def test_create_varmax_with_exogenous_exog_names_dict(self, data):
         assert mod.param_info["beta_observed_1"]["shape"] == (1,)
         assert mod.param_info["beta_observed_1"]["dims"] == ("exogenous_observed_1",)
 
-    # def test_create_varmax_with_exogenous_both_dict_correct(self, data):
-    #     k_exog = {"observed_0": 2, "observed_1": 1}
-    #     exog_state_names = {"observed_0": ["a", "b"], "observed_1": ["c"]}
-    #     mod = BayesianVARMAX(
-    #         endog_names=["observed_0", "observed_1"],
-    #         order=(1, 0),
-    #         k_exog=k_exog,
-    #         exog_state_names=exog_state_names,
-    #         verbose=False,
-    #         measurement_error=False,
-    #         stationary_initialization=False,
-    #     )
-    #     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",)
-
     def test_create_varmax_with_exogenous_dict_converts_to_list(self, data):
         exog_state_names = {
             "observed_0": ["a", "b"],
@@ -374,76 +292,6 @@ def test_create_varmax_with_exogenous_dict_converts_to_list(self, data):
         assert mod.param_info["beta_exog"]["shape"] == (mod.k_endog, 2)
         assert mod.param_info["beta_exog"]["dims"] == ("observed_state", "exogenous")
 
-    # def test_create_varmax_with_exogenous_raises_if_args_disagree(self, data):
-    #     # List case
-    #     with pytest.raises(
-    #         ValueError, match="Length of exog_state_names does not match provided k_exog"
-    #     ):
-    #         BayesianVARMAX(
-    #             k_endog=2,
-    #             order=(1, 0),
-    #             k_exog=3,
-    #             exog_state_names=["a", "b"],
-    #             verbose=False,
-    #             measurement_error=False,
-    #             stationary_initialization=False,
-    #         )
-
-    #     # Dict case
-    #     with pytest.raises(
-    #         ValueError,
-    #         match="If k_exog is an int, exog_state_names must be a list of the same length",
-    #     ):
-    #         BayesianVARMAX(
-    #             k_endog=2,
-    #             order=(1, 0),
-    #             k_exog=2,
-    #             exog_state_names={"observed_0": ["a"], "observed_1": ["b"]},
-    #             verbose=False,
-    #             measurement_error=False,
-    #             stationary_initialization=False,
-    #         )
-
-    #     # dict + list
-    #     with pytest.raises(
-    #         ValueError, match="If k_exog is a dict, exog_state_names must be a dict as well"
-    #     ):
-    #         BayesianVARMAX(
-    #             endog_names=["observed_0", "observed_1"],
-    #             order=(1, 0),
-    #             k_exog={"observed_0": 1, "observed_1": 1},
-    #             exog_state_names=["a", "b"],
-    #             verbose=False,
-    #             measurement_error=False,
-    #             stationary_initialization=False,
-    #         )
-
-    #     # Dict/dict, key mismatch
-    #     with pytest.raises(
-    #         ValueError, match="Keys of k_exog and exog_state_names dicts must match"
-    #     ):
-    #         BayesianVARMAX(
-    #             endog_names=["observed_0", "observed_1"],
-    #             order=(1, 0),
-    #             k_exog={"observed_0": 1, "observed_1": 1},
-    #             exog_state_names={"observed_0": ["a"], "observed_2": ["b"]},
-    #             verbose=False,
-    #             measurement_error=False,
-    #             stationary_initialization=False,
-    #         )
-
-    #     # Dict/dict, length mismatch
-    #     with pytest.raises(ValueError, match="lengths of exog_state_names lists must match"):
-    #         BayesianVARMAX(
-    #             endog_names=["observed_0", "observed_1"],
-    #             order=(1, 0),
-    #             k_exog={"observed_0": 2, "observed_1": 1},
-    #             exog_state_names={"observed_0": ["a"], "observed_1": ["b"]},
-    #             verbose=False,
-    #             measurement_error=False,
-    #             stationary_initialization=False,
-    #         )
-
     def _build_varmax(self, df, exog_state_names, exog_data):
         endog_names = df.columns.values.tolist()
 
@@ -514,11 +362,10 @@ def test_varmax_with_exog(self, rng, exog_state_names):
                 for name, exog_names in exog_state_names.items()
             }
         else:
-            exog_names = exog_state_names
             exog_data = {
                 "exogenous_data": pd.DataFrame(
                     rng.normal(size=(n_obs, len(exog_state_names))).astype(floatX),
-                    columns=exog_names,
+                    columns=exog_state_names,
                     index=time_idx,
                 )
             }