Added support for joint VAR modelling (error_var=True)

Author: andreacate

pymc_extras/statespace/models/DFM.py

@@ -201,8 +201,7 @@ class BayesianDynamicFactor(PyMCStateSpace):
         These models are only identified up to a sign flip in the factor loadings. Proper prior specification is crucial
         for good estimation and inference.
 
-    Currently, the implementation assumes same factor order for all the factors,
-    does not yet support exogenous variables and joint (VAR) error modeling.
+    Currently, the implementation does not yet support exogenous variables
 
     Examples
     --------
@@ -284,11 +283,6 @@ def __init__(
         if endog_names is None:
             endog_names = [f"endog_{i}" for i in range(k_endog)]
 
-        if error_var:
-            raise NotImplementedError(
-                "Joint error modeling (error_var=True) is not yet implemented."
-            )
-
         if k_exog is not None or exog_names is not None:
             raise NotImplementedError("Exogenous variables (exog) are not yet implemented.")
 
@@ -300,7 +294,6 @@ def __init__(
         self.error_var = error_var
         self.error_cov_type = error_cov_type
         # TODO add exogenous variables support
-        # TODO add error_var support
 
         # Determine the dimension for the latent factor states.
         # For static factors, one use k_factors.
@@ -376,7 +369,10 @@ def param_info(self) -> dict[str, dict[str, Any]]:
                 "constraints": None,
             },
             "error_ar": {
-                "shape": (self.k_endog, self.error_order),
+                "shape": (
+                    self.k_endog,
+                    self.error_order * self.k_endog if self.error_var else self.error_order,
+                ),
                 "constraints": None,
             },
             "error_sigma": {
@@ -437,7 +433,10 @@ def coords(self) -> dict[str, Sequence]:
 
         # If error_order > 0
         if self.error_order > 0:
-            coords[ERROR_AR_PARAM_DIM] = list(range(1, self.error_order + 1))
+            if self.error_var:
+                coords[ERROR_AR_PARAM_DIM] = list(range(1, (self.error_order * self.k_endog) + 1))
+            else:
+                coords[ERROR_AR_PARAM_DIM] = list(range(1, self.error_order + 1))
 
         return coords
 
@@ -509,26 +508,26 @@ def make_symbolic_graph(self):
 
         # Transition matrix
         # auxiliary function to build transition matrix block
-        def build_var_block_matrix(ar_coeffs, k_factors, p):
+        def build_var_block_matrix(ar_coeffs, k_series, p):
             """
             Build the VAR(p) companion matrix for the factors.
 
-            ar_coeffs: PyTensor matrix of shape (k_factors, p * k_factors)
+            ar_coeffs: PyTensor matrix of shape (k_series, p * k_series)
                     [A1 | A2 | ... | Ap] horizontally concatenated.
-            k_factors: number of factors
+            k_series: number of series
             p: lag order
             """
-            size = k_factors * p
+            size = k_series * p
             block = pt.zeros((size, size), dtype=floatX)
 
             # First block row: the AR coefficient matrices for each lag
-            block = pt.set_subtensor(block[0:k_factors, 0 : k_factors * p], ar_coeffs)
+            block = pt.set_subtensor(block[0:k_series, 0 : k_series * p], ar_coeffs)
 
             # Sub-diagonal identity blocks (shift structure)
             if p > 1:
                 # 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)
+                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
 
@@ -571,7 +570,14 @@ def build_independent_var_block_matrix(ar_coeffs, k_series, p):
         else:
             transition_blocks.append(pt.zeros((self.k_factors, self.k_factors), dtype=floatX))
 
-        if self.error_order > 0:
+        if self.error_order > 0 and self.error_var:
+            error_ar = self.make_and_register_variable(
+                "error_ar", shape=(self.k_endog, self.error_order * self.k_endog), dtype=floatX
+            )
+            transition_blocks.append(
+                build_var_block_matrix(error_ar, self.k_endog, self.error_order)
+            )
+        elif self.error_order > 0 and not self.error_var:
             error_ar = self.make_and_register_variable(
                 "error_ar", shape=(self.k_endog, self.error_order), dtype=floatX
             )

tests/statespace/models/test_DFM.py

@@ -18,6 +18,9 @@
 
 floatX = pytensor.config.floatX
 
+# TODO: check test for error_var=True, since there are problems with statsmodels, the matrices looks the same by some experiments done in notebooks
+# (FAILED tests/statespace/models/test_DFM.py::test_DFM_update_matches_statsmodels[True-2-2-2] - numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite)
+
 
 @pytest.fixture(scope="session")
 def data():
@@ -30,7 +33,9 @@ def data():
     return df
 
 
-def create_sm_test_values_mapping(test_values, data, k_factors, factor_order, error_order):
+def create_sm_test_values_mapping(
+    test_values, data, k_factors, factor_order, error_order, error_var
+):
     """Convert PyMC test values to statsmodels parameter format"""
     sm_test_values = {}
 
@@ -59,8 +64,8 @@ def create_sm_test_values_mapping(test_values, data, k_factors, factor_order, er
             }
         )
 
-    # 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:
+    # 3a. Error AR coefficients: PyMC shape (n_endog, error_order) -> L{lag}.e(var).e(var)
+    if error_order > 0 and not error_var 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(
@@ -70,6 +75,24 @@ def create_sm_test_values_mapping(test_values, data, k_factors, factor_order, er
             }
         )
 
+    # 3b. Error AR coefficients: PyMC shape (n_endog, error_order * n_endog) -> L{lag}.e(var).e(var)
+    elif error_order > 0 and error_var and "error_ar" in test_values:
+        error_ar = test_values["error_ar"]
+        triplets = product(
+            enumerate(data.columns), range(1, error_order + 1), enumerate(data.columns)
+        )
+        sm_test_values.update(
+            {
+                f"L{lag}.e({from_endog_name}).e({to_endog_name})": error_ar[
+                    from_endog_idx, (lag - 1) * data.shape[1] + to_endog_idx
+                ]
+                for (from_endog_idx, from_endog_name), lag, (
+                    to_endog_idx,
+                    to_endog_name,
+                ) in triplets
+            }
+        )
+
     # 4. Observation error variances:
     if "error_sigma" in test_values:
         error_sigma = test_values["error_sigma"]
@@ -86,22 +109,25 @@ def create_sm_test_values_mapping(test_values, data, k_factors, factor_order, er
 @pytest.mark.parametrize("k_factors", [1, 2])
 @pytest.mark.parametrize("factor_order", [0, 1, 2])
 @pytest.mark.parametrize("error_order", [0, 1, 2])
+@pytest.mark.parametrize("error_var", [False])
 @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):
+def test_DFM_update_matches_statsmodels(data, k_factors, factor_order, error_order, error_var, rng):
     mod = BayesianDynamicFactor(
         k_factors=k_factors,
         factor_order=factor_order,
         error_order=error_order,
         k_endog=data.shape[1],
         measurement_error=False,
+        error_var=error_var,
         verbose=False,
     )
     sm_dfm = DynamicFactor(
         endog=data,
         k_factors=k_factors,
         factor_order=factor_order,
         error_order=error_order,
+        error_var=error_var,
     )
 
     # Generate test values for PyMC model
@@ -113,14 +139,16 @@ def test_DFM_update_matches_statsmodels(data, k_factors, factor_order, error_ord
     if factor_order > 0:
         test_values["factor_ar"] = rng.normal(size=(k_factors, factor_order * k_factors))
 
-    if error_order > 0:
+    if error_order > 0 and error_var:
+        test_values["error_ar"] = rng.normal(size=(data.shape[1], error_order * data.shape[1]))
+    elif error_order > 0 and not error_var:
         test_values["error_ar"] = rng.normal(size=(data.shape[1], error_order))
 
     test_values["error_sigma"] = rng.beta(1, 1, size=data.shape[1])
 
     # Convert to statsmodels format
     sm_test_values = create_sm_test_values_mapping(
-        test_values, data, k_factors, factor_order, error_order
+        test_values, data, k_factors, factor_order, error_order, error_var
     )
 
     # Initialize and constrain statsmodels model