Added test_DFM_update_matches_statsmodels and small corrections to DFM.py

Author: andreacate

pymc_extras/statespace/models/DFM.py

@@ -241,7 +241,7 @@ class BayesianDynamicFactor(PyMCStateSpace):
             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))
+            # 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"])
@@ -302,6 +302,8 @@ def __init__(
         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.
@@ -533,13 +535,9 @@ def build_var_block_matrix(ar_coeffs, k_factors, p):
 
             # Sub-diagonal identity blocks (shift structure)
             if p > 1:
-                for i in range(1, p):
-                    row_start = i * k_factors
-                    col_start = (i - 1) * k_factors
-                    block = pt.set_subtensor(
-                        block[row_start : row_start + k_factors, col_start : col_start + k_factors],
-                        pt.eye(k_factors, dtype=floatX),
-                    )
+                # Create the identity pattern for all sub-diagonal blocks
+                identity_pattern = pt.eye(k_factors * (p - 1), dtype=floatX)
+                block = pt.set_subtensor(block[k_factors:, : k_factors * (p - 1)], identity_pattern)
 
             return block
 
@@ -559,26 +557,15 @@ def build_independent_var_block_matrix(ar_coeffs, k_series, p):
             # First block row: AR coefficients per series (block diagonal)
             for j in range(k_series):
                 for lag in range(p):
-                    col_idx = lag * k_series + j  # interleaved layout
+                    col_idx = lag * k_series + j
                     block = pt.set_subtensor(block[j, col_idx], ar_coeffs[j, lag])
 
             # Sub-diagonal identity blocks (shift)
             if p > 1:
-                for lag in range(1, p):
-                    row_start = lag * k_series
-                    col_start = (lag - 1) * k_series
-                    block = pt.set_subtensor(
-                        block[row_start : row_start + k_series, col_start : col_start + k_series],
-                        pt.eye(k_series, dtype=floatX),
-                    )
-
+                identity_pattern = pt.eye(k_series * (p - 1), dtype=floatX)
+                block = pt.set_subtensor(block[k_series:, : k_series * (p - 1)], identity_pattern)
             return block
 
-        # def build_ar_block_matrix(ar_coeffs):
-        #     # ar_coeffs: (p,)
-        #     p = ar_coeffs.shape[0]
-        #     return pt.eye(p, k=-1)[0].set(ar_coeffs)
-
         transition_blocks = []
 
         if self.factor_order > 0:

tests/statespace/models/test_DFM.py

@@ -79,7 +79,7 @@ def test_dfm_parameter_and_matrix_match(data, k_factors, factor_order, error_ord
             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("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))
 
@@ -96,3 +96,124 @@ def test_dfm_parameter_and_matrix_match(data, k_factors, factor_order, error_ord
             atol=1e-12,
             err_msg=f"Matrix mismatch: {mat_name} (k_factors={k_factors}, factor_order={factor_order}, error_order={error_order})",
         )
+
+
+@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.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,
+        error_order=error_order,
+        k_endog=data.shape[1],
+        measurement_error=False,
+        verbose=False,
+    )
+
+    # 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,)
+
+    # 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)
+    )
+
+    # Handle factor_ar parameters - need to account for different factor orders
+    factor_ar_params = []
+
+    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)
+
+    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)
+
+    # factor_sigma = np.array([param_dict[p] for p in param_names if "factor.sigma" in p])
+
+    # 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
+
+    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})",
+        )