Updating test following test_ETS.py and small adjustment for exog variables in DFM.py

Author: andreacate

pymc_extras/statespace/models/DFM.py

@@ -47,10 +47,11 @@ class BayesianDynamicFactor(PyMCStateSpace):
     endog_names : Sequence[str], optional
         Names of the observed time series. If not provided, default names will be generated as `endog_1`, `endog_2`, ..., `endog_k`.
 
-    exog : array_like, optional
-        Array of exogenous variables for the observation equation (nobs x k_exog).
-        Default is None, meaning no exogenous variables.
-        Not implemented yet.
+    k_exog : int
+        Number of exogenous variables, optional. If not provided, model will not have exogenous variables.
+
+    exog_names : Sequence[str], optional
+        Names of the exogenous variables. If not provided, but `k_exog` is specified, default names will be generated as `exog_1`, `exog_2`, ..., `exog_k`.
 
     error_order : int, optional
         Order of the AR process for the observation error component.
@@ -240,9 +241,6 @@ class BayesianDynamicFactor(PyMCStateSpace):
             # AR coefficients for factor dynamics: shape (k_factors, factor_order)
             factor_ar = pm.Normal("factor_ar", sigma=1, dims=["k_factors", "factor_order"])
 
-            # Innovation std dev of factors: shape (k_factors,)
-            # factor_sigma = pm.Deterministic("factor_sigma", pt.constant([1.0], dtype=float)) TODO could be removed
-
             # AR coefficients for observation noise: shape (k_endog, error_order)
             error_ar = pm.Normal("error_ar", sigma=1, dims=["k_endog", "error_order"])
 
@@ -271,7 +269,8 @@ def __init__(
         factor_order: int,
         k_endog: int | None = None,
         endog_names: Sequence[str] | None = None,
-        exog: np.ndarray | None = None,
+        k_exog: int | None = None,
+        exog_names: Sequence[str] | None = None,
         error_order: int = 0,
         error_var: bool = False,
         error_cov_type: str = "diagonal",
@@ -290,7 +289,7 @@ def __init__(
                 "Joint error modeling (error_var=True) is not yet implemented."
             )
 
-        if exog is not None:
+        if k_exog is not None or exog_names is not None:
             raise NotImplementedError("Exogenous variables (exog) are not yet implemented.")
 
         self.endog_names = endog_names
@@ -300,10 +299,8 @@ def __init__(
         self.error_order = error_order
         self.error_var = error_var
         self.error_cov_type = error_cov_type
-        self.exog = exog
         # TODO add exogenous variables support
         # TODO add error_var support
-        # TODO understanding if the factor_sigma matrix is the identity?
 
         # Determine the dimension for the latent factor states.
         # For static factors, one use k_factors.
@@ -341,7 +338,6 @@ def param_names(self):
             "P0",
             "factor_loadings",
             "factor_ar",
-            # "factor_sigma",
             "error_ar",
             "error_sigma",
             "sigma_obs",
@@ -379,10 +375,6 @@ def param_info(self) -> dict[str, dict[str, Any]]:
                 "shape": (self.k_factors, self.factor_order * self.k_factors),
                 "constraints": None,
             },
-            # "factor_sigma": {
-            #     "shape": (self.k_factors,),
-            #     "constraints": "Positive",
-            # },
             "error_ar": {
                 "shape": (self.k_endog, self.error_order),
                 "constraints": None,
@@ -470,7 +462,6 @@ def param_dims(self):
             "x0": (ALL_STATE_DIM,),
             "P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
             "factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
-            # "factor_sigma": (FACTOR_DIM,),
         }
         if self.factor_order > 0:
             coord_map["factor_ar"] = (FACTOR_DIM, AR_PARAM_DIM)
@@ -601,11 +592,6 @@ def build_independent_var_block_matrix(ar_coeffs, k_series, p):
                 col = self.k_factors + i
                 self.ssm["selection", row, col] = 1.0
 
-        # # State covariance matrix
-        # factor_sigma = self.make_and_register_variable(
-        #     "factor_sigma", shape=(self.k_factors,), dtype=floatX
-        # )
-        # factor_cov = pt.diag(factor_sigma)
         factor_cov = pt.eye(self.k_factors, dtype=floatX)
 
         # Handle error_sigma and error_cov depending on error_cov_type

tests/statespace/models/test_DFM.py

@@ -1,3 +1,5 @@
+from itertools import product
+
 import numpy as np
 import pandas as pd
 import pymc as pm
@@ -7,10 +9,11 @@
 import statsmodels.api as sm
 
 from numpy.testing import assert_allclose
+from pytensor.graph.basic import explicit_graph_inputs
 from statsmodels.tsa.statespace.dynamic_factor import DynamicFactor
 
 from pymc_extras.statespace.models.DFM import BayesianDynamicFactor
-from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG
+from pymc_extras.statespace.utils.constants import LONG_MATRIX_NAMES
 from tests.statespace.shared_fixtures import rng
 
 floatX = pytensor.config.floatX
@@ -27,95 +30,65 @@ def data():
     return df
 
 
-@pytest.mark.parametrize("k_factors", [1, 2])
-@pytest.mark.parametrize("factor_order", [0, 1, 2])
-@pytest.mark.parametrize("error_order", [0, 1, 2])
-def test_dfm_parameter_and_matrix_match(data, k_factors, factor_order, error_order):
-    # --- Statsmodels DFM ---
-    sm_dfm = DynamicFactor(
-        endog=data,
-        k_factors=k_factors,
-        factor_order=factor_order,
-        error_order=error_order,
-    )
-
-    # Use deterministic small parameters for reproducibility
-    param_array = np.full(len(sm_dfm.param_names), 0.5)
-    sm_dfm.update(param_array)
-
-    # Only request matrices that actually exist in ssm.__getitem__
-    valid_names = ["design", "obs_cov", "transition", "state_cov", "selection"]
-    sm_matrices = {name: np.array(sm_dfm.ssm[name]) for name in valid_names}
+def create_sm_test_values_mapping(test_values, data, k_factors, factor_order, error_order):
+    """Convert PyMC test values to statsmodels parameter format"""
+    sm_test_values = {}
 
-    # --- PyMC DFM ---
-    mod = BayesianDynamicFactor(
-        k_factors=k_factors,
-        factor_order=factor_order,
-        k_endog=data.shape[1],
-        error_order=error_order,
-        measurement_error=False,
-        verbose=False,
+    # 1. Factor loadings: PyMC shape (n_endog, k_factors) -> statsmodels individual params
+    factor_loadings = test_values["factor_loadings"]
+    all_pairs = product(data.columns, range(1, k_factors + 1))
+    sm_test_values.update(
+        {
+            f"loading.f{factor_idx}.{endog_name}": value
+            for (endog_name, factor_idx), value in zip(all_pairs, factor_loadings.ravel())
+        }
     )
 
-    coords = mod.coords
-    with pm.Model(coords=coords):
-        k_endog = data.shape[1]
-        factor_part = max(1, factor_order) * k_factors
-        error_part = error_order * k_endog if error_order > 0 else 0
-        k_states = factor_part + error_part
-
-        pm.Deterministic("x0", pt.constant(np.full((k_states,), 0.5), dtype=floatX))
-        pm.Deterministic("P0", pt.constant(np.full((k_states, k_states), 0.5), dtype=floatX))
-        pm.Deterministic(
-            "factor_loadings", pt.constant(np.full((k_endog, k_factors), 0.5), dtype=floatX)
+    # 2. Factor AR coefficients: PyMC shape (k_factors, factor_order*k_factors) -> L{lag}.f{to}.f{from}
+    if factor_order > 0 and "factor_ar" in test_values:
+        factor_ar = test_values["factor_ar"]
+        triplets = product(
+            range(1, k_factors + 1), range(1, factor_order + 1), range(1, k_factors + 1)
+        )
+        sm_test_values.update(
+            {
+                f"L{lag}.f{to_factor}.f{from_factor}": factor_ar[
+                    from_factor - 1, (lag - 1) * k_factors + (to_factor - 1)
+                ]
+                for from_factor, lag, to_factor in triplets
+            }
         )
 
-        if factor_order > 0:
-            pm.Deterministic(
-                "factor_ar",
-                pt.constant(np.full((k_factors, factor_order * k_factors), 0.5), dtype=floatX),
-            )
-        if error_order > 0:
-            pm.Deterministic(
-                "error_ar", pt.constant(np.full((k_endog, error_order), 0.5), dtype=floatX)
-            )
-        # pm.Deterministic("factor_sigma", pt.constant(np.full((k_factors,), 0.5), dtype=floatX))
-        pm.Deterministic("error_sigma", pt.constant(np.full((k_endog,), 0.5), dtype=floatX))
-        pm.Deterministic("sigma_obs", pt.constant(np.full((k_endog,), 0.5), dtype=floatX))
-
-        mod._insert_random_variables()
-
-        pymc_matrices = pm.draw(mod.subbed_ssm)
-        pymc_matrices = dict(zip(SHORT_NAME_TO_LONG.values(), pymc_matrices))
-
-    # --- Compare ---
-    for mat_name in valid_names:
-        assert_allclose(
-            pymc_matrices[mat_name],
-            sm_matrices[mat_name],
-            atol=1e-12,
-            err_msg=f"Matrix mismatch: {mat_name} (k_factors={k_factors}, factor_order={factor_order}, error_order={error_order})",
+    # 3. Error AR coefficients: PyMC shape (n_endog, error_order) -> L{lag}.e(var).e(var)
+    if error_order > 0 and "error_ar" in test_values:
+        error_ar = test_values["error_ar"]
+        pairs = product(enumerate(data.columns), range(1, error_order + 1))
+        sm_test_values.update(
+            {
+                f"L{lag}.e({endog_name}).e({endog_name})": error_ar[endog_idx, lag - 1]
+                for (endog_idx, endog_name), lag in pairs
+            }
         )
 
+    # 4. Observation error variances:
+    if "error_sigma" in test_values:
+        error_sigma = test_values["error_sigma"]
+        sm_test_values.update(
+            {
+                f"sigma2.{endog_name}": error_sigma[endog_idx]
+                for endog_idx, endog_name in enumerate(data.columns)
+            }
+        )
+
+    return sm_test_values
+
 
 @pytest.mark.parametrize("k_factors", [1, 2])
 @pytest.mark.parametrize("factor_order", [0, 1, 2])
-@pytest.mark.parametrize("error_order", [1, 2, 3])
+@pytest.mark.parametrize("error_order", [0, 1, 2])
 @pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.EstimationWarning")
 @pytest.mark.filterwarnings("ignore::FutureWarning")
 def test_DFM_update_matches_statsmodels(data, k_factors, factor_order, error_order, rng):
-    # --- Fit Statsmodels DynamicFactor with random small params ---
-    sm_dfm = DynamicFactor(
-        endog=data,
-        k_factors=k_factors,
-        factor_order=factor_order,
-        error_order=error_order,
-    )
-    param_names = sm_dfm.param_names
-    param_dict = {param: getattr(np, floatX)(rng.normal(scale=0.1) ** 2) for param in param_names}
-    sm_res = sm_dfm.fit_constrained(param_dict)
-
-    # --- Setup BayesianDynamicFactor ---
     mod = BayesianDynamicFactor(
         k_factors=k_factors,
         factor_order=factor_order,
@@ -124,96 +97,51 @@ def test_DFM_update_matches_statsmodels(data, k_factors, factor_order, error_ord
         measurement_error=False,
         verbose=False,
     )
+    sm_dfm = DynamicFactor(
+        endog=data,
+        k_factors=k_factors,
+        factor_order=factor_order,
+        error_order=error_order,
+    )
+
+    # Generate test values for PyMC model
+    test_values = {}
+    test_values["x0"] = rng.normal(size=mod.k_states)
+    test_values["P0"] = np.eye(mod.k_states)  # Use identity for stability
+    test_values["factor_loadings"] = rng.normal(size=(data.shape[1], k_factors))
+
+    if factor_order > 0:
+        test_values["factor_ar"] = rng.normal(size=(k_factors, factor_order * k_factors))
+
+    if error_order > 0:
+        test_values["error_ar"] = rng.normal(size=(data.shape[1], error_order))
 
-    # Convert flat param dict to PyTensor variables as needed
-    # Reshape factor_ar and error_ar parameters according to model expected shapes
-    factor_ar_shape = (k_factors, factor_order * k_factors)
-    error_ar_shape = (data.shape[1], error_order) if error_order > 0 else (0,)
+    test_values["error_sigma"] = rng.beta(1, 1, size=data.shape[1])
 
-    # Prepare parameter arrays to set as deterministic
-    # Extract each group of parameters by name pattern (simplified)
-    factor_loadings = np.array([param_dict[p] for p in param_names if "loading" in p]).reshape(
-        (data.shape[1], k_factors)
+    # Convert to statsmodels format
+    sm_test_values = create_sm_test_values_mapping(
+        test_values, data, k_factors, factor_order, error_order
     )
 
-    # Handle factor_ar parameters - need to account for different factor orders
-    factor_ar_params = []
+    # Initialize and constrain statsmodels model
+    x0 = test_values["x0"]
+    P0 = test_values["P0"]
 
-    for factor_idx in range(1, k_factors + 1):
-        for lag in range(1, factor_order + 1):
-            for factor_idx2 in range(1, k_factors + 1):
-                param_pattern = f"L{lag}.f{factor_idx2}.f{factor_idx}"
-                if param_pattern in param_names:
-                    factor_ar_params.append(param_pattern)
+    sm_dfm.initialize_known(initial_state=x0, initial_state_cov=P0)
+    sm_dfm.fit_constrained({name: sm_test_values[name] for name in sm_dfm.param_names})
 
-    if len(factor_ar_params) > 0:
-        factor_ar_values = [param_dict[p] for p in factor_ar_params]
-        factor_ar = np.array(factor_ar_values).reshape(factor_ar_shape)
-    else:
-        factor_ar = np.zeros(factor_ar_shape)
+    # Get PyMC matrices using the same pattern as ETS test
+    matrices = mod._unpack_statespace_with_placeholders()
+    inputs = list(explicit_graph_inputs(matrices))
+    input_names = [x.name for x in inputs]
 
-    # factor_sigma = np.array([param_dict[p] for p in param_names if "factor.sigma" in p])
+    f_matrices = pytensor.function(inputs, matrices)
+    test_values_subset = {name: test_values[name] for name in input_names}
 
-    # Handle error AR parameters - need to account for different error orders and variables
-    if error_order > 0:
-        error_ar_params = []
-        var_names = [col for col in data.columns]  # Get variable names from data
-
-        # Order parameters by variable first, then by lag to match expected shape (n_vars, n_lags)
-        for var_name in var_names:
-            for lag in range(1, error_order + 1):
-                param_pattern = f"L{lag}.e({var_name}).e({var_name})"
-                if param_pattern in param_names:
-                    error_ar_params.append(param_pattern)
-
-        if len(error_ar_params) > 0:
-            error_ar_values = [param_dict[p] for p in error_ar_params]
-            error_ar = np.array(error_ar_values).reshape(error_ar_shape)
-        else:
-            error_ar = np.zeros(error_ar_shape)
-
-    # Handle observation error variances - look for sigma2 pattern
-    sigma_obs_params = [p for p in param_names if "sigma2." in p]
-    sigma_obs = np.array([param_dict[p] for p in sigma_obs_params])
-
-    # Handle error variances (if needed separately from sigma_obs)
-    if error_order > 0:
-        error_sigma = sigma_obs  # In this case, error_sigma is the same as sigma_obs
+    pymc_matrices = f_matrices(**test_values_subset)
 
-    coords = mod.coords
-    with pm.Model(coords=coords) as model:
-        k_states = k_factors * max(1, factor_order) + (
-            error_order * data.shape[1] if error_order > 0 else 0
-        )
-        pm.Deterministic("x0", pt.zeros(k_states, dtype=floatX))
-        pm.Deterministic("P0", pt.eye(k_states, dtype=floatX))
-        # Set deterministic variables with constrained parameter values
-        pm.Deterministic("factor_loadings", pt.as_tensor_variable(factor_loadings))
-        if factor_order > 0:
-            pm.Deterministic("factor_ar", pt.as_tensor_variable(factor_ar))
-        # pm.Deterministic("factor_sigma", pt.as_tensor_variable(factor_sigma))
-        if error_order > 0:
-            pm.Deterministic("error_ar", pt.as_tensor_variable(error_ar))
-        pm.Deterministic("error_sigma", pt.as_tensor_variable(error_sigma))
-        pm.Deterministic("sigma_obs", pt.as_tensor_variable(sigma_obs))
-
-        mod._insert_random_variables()
-
-        # Draw the substituted state-space matrices from PyMC model
-        matrices = pm.draw(mod.subbed_ssm)
-        matrix_dict = dict(zip(SHORT_NAME_TO_LONG.values(), matrices))
-
-    # Matrices to check
-    matrices_to_check = ["transition", "selection", "state_cov", "obs_cov", "design"]
-
-    # Compare matrices from PyMC and Statsmodels
-    for mat_name in matrices_to_check:
-        sm_mat = np.array(sm_dfm.ssm[mat_name])
-        pm_mat = matrix_dict[mat_name]
-
-        assert_allclose(
-            pm_mat,
-            sm_mat,
-            atol=1e-10,
-            err_msg=f"Matrix mismatch: {mat_name} (k_factors={k_factors}, factor_order={factor_order}, error_order={error_order})",
-        )
+    sm_matrices = [sm_dfm.ssm[name] for name in LONG_MATRIX_NAMES[2:]]
+
+    # Compare matrices (skip x0 and P0)
+    for matrix, sm_matrix, name in zip(pymc_matrices[2:], sm_matrices, LONG_MATRIX_NAMES[2:]):
+        assert_allclose(matrix, sm_matrix, err_msg=f"{name} does not match")