Update of DFM draft implementation

In the notebook a comparison between the custom DFM and the implemented DFM (which has an hardcoded version of make_symbolic_graph, that work just in this case)

Author: andreacate

notebooks/Making a Custom DFM.ipynb

@@ -4,9 +4,11 @@
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
 from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     AR_PARAM_DIM,
     MA_PARAM_DIM,
+    OBS_STATE_DIM,
     SHOCK_DIM,
 )
 
@@ -28,6 +30,8 @@ class BayesianDynamicFactor(PyMCStateSpace):
 
     exog : array_like, optional
         Array of exogenous regressors for the observation equation (nobs x k_exog).
+        Default is None, meaning no exogenous regressors.
+        Not implemented yet.
 
     error_order : int, optional
         Order of the AR process for the observation error component.
@@ -58,6 +62,9 @@ class BayesianDynamicFactor(PyMCStateSpace):
     the observed time series are driven by a set of latent factors that evolve
     according to a VAR process, possibly along with an autoregressive error term.
 
+    Up to now just a draft implementation to test the working of the class and comparing
+    with the Custom model done in the Notebook (notebook/Making a Custom DFM.ipynb).
+    The model work just with two observations and one factor (k_endog=2, k_factors=1).
 
 
     """
@@ -113,9 +120,15 @@ def __init__(
 
     @property
     def param_names(self):
-        names = ["factor_loadings", "factor_ar", "error_ar", "error_sigma"]
-
-        # factor_sigma is fixed and equal to the identity matrix
+        names = [
+            "x0",
+            "P0",
+            "factor_loadings",
+            "factor_ar",
+            "factor_sigma",
+            "error_ar",
+            "error_sigma",
+        ]
 
         # Handle cases where parameters should be excluded based on model settings
         if self.factor_order == 0:
@@ -130,6 +143,14 @@ def param_names(self):
     @property
     def param_info(self) -> dict[str, dict[str, Any]]:
         info = {
+            "x0": {
+                "shape": (self.k_factors,),
+                "constraints": None,
+            },
+            "P0": {
+                "shape": (self.k_factors, self.k_factors),
+                "constraints": "Positive Semi-definite",
+            },
             "factor_loadings": {
                 "shape": (self.k_endog, self.k_factors),
                 "constraints": None,
@@ -138,6 +159,10 @@ 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, self.k_endog)
                 if self.error_var
@@ -167,9 +192,7 @@ def param_info(self) -> dict[str, dict[str, Any]]:
 
     @property
     def state_names(self):
-        # Initialize state names based on the endogenous variables
-        state_names = self.endog_names.copy()
-
+        state_names = []
         # Add names for the factor loadings (one per observation and factor)
         for i in range(self.k_endog):
             for j in range(self.k_factors):
@@ -216,10 +239,6 @@ def state_names(self):
 
         return state_names
 
-    @property
-    def observed_states(self):
-        return self.endog_names
-
     @property
     def shock_names(self):
         shock_names = []
@@ -235,39 +254,86 @@ def shock_names(self):
 
         return shock_names
 
+    @property
+    def param_dims(self):
+        """
+        Define parameter dimensions for the Dynamic Factor Model (DFM).
+
+        Returns
+        -------
+        dict
+            Dictionary mapping parameter names to their respective dimensions.
+        """
+        coord_map = {
+            "x0": (ALL_STATE_DIM,),  # Initial state dimension
+            "P0": (ALL_STATE_DIM, ALL_STATE_DIM),  # Initial state covariance dimension
+            "factor_loadings": (OBS_STATE_DIM, ALL_STATE_AUX_DIM),  # Factor loadings dimension
+            "factor_sigma": (ALL_STATE_DIM,),  # Factor variances dimension
+        }
 
-@property
-def param_dims(self):
-    """
-    Define parameter dimensions for the Dynamic Factor Model (DFM).
+        # Factor AR coefficients if applicable
+        if self.factor_order > 0:
+            coord_map["factor_ar"] = (AR_PARAM_DIM, SHOCK_DIM, SHOCK_DIM)
 
-    Returns
-    -------
-    dict
-        Dictionary mapping parameter names to their respective dimensions.
-    """
-    coord_map = {
-        "factor_loadings": (ALL_STATE_DIM, SHOCK_DIM),  # Factor loadings dimension
-        "factor_sigma": (SHOCK_DIM,),  # Factor shocks (one per factor)
-    }
-
-    # Factor AR coefficients if applicable
-    if self.factor_order > 0:
-        coord_map["factor_ar"] = (AR_PARAM_DIM, SHOCK_DIM, SHOCK_DIM)
-
-    # Error AR coefficients and variances
-    if self.error_order > 0:
-        if self.error_cov_type == "diagonal":
-            coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM)  # AR for errors
-            coord_map["error_sigma"] = (SHOCK_DIM,)  # One variance for each observed variable
-        elif self.error_cov_type == "scalar":
-            coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM)
-            coord_map["error_sigma"] = None  # Single scalar for error variance
-        elif self.error_cov_type == "unstructured":
-            coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM, SHOCK_DIM)  # AR for errors
-            coord_map["error_cov_L"] = (SHOCK_DIM, SHOCK_DIM)  # Lower triangular Cholesky factor
-            coord_map["error_cov_sd"] = (SHOCK_DIM,)  # Standard deviations for diagonal
-        else:
-            raise ValueError("Invalid error covariance type.")
-
-    return coord_map
+        # Error AR coefficients and variances
+        if self.error_order > 0:
+            if self.error_cov_type == "diagonal":
+                coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM)  # AR for errors
+                coord_map["error_sigma"] = (SHOCK_DIM,)  # One variance for each observed variable
+            elif self.error_cov_type == "scalar":
+                coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM)
+                coord_map["error_sigma"] = None  # Single scalar for error variance
+            elif self.error_cov_type == "unstructured":
+                coord_map["error_ar"] = (MA_PARAM_DIM, SHOCK_DIM, SHOCK_DIM)  # AR for errors
+                coord_map["error_cov_L"] = (
+                    SHOCK_DIM,
+                    SHOCK_DIM,
+                )  # Lower triangular Cholesky factor
+                coord_map["error_cov_sd"] = (SHOCK_DIM,)  # Standard deviations for diagonal
+            else:
+                raise ValueError("Invalid error covariance type.")
+
+        return coord_map
+
+    # def make_symbolic_graph(self):
+    # We will implement this in a moment. For now, we need to overwrite it with nothing to avoid a NotImplementedError
+    # when we initialize a class instance.
+    #    pass
+
+    def make_symbolic_graph(self):
+        """
+        Create the symbolic graph for the Dynamic Factor Model (DFM).
+        This method sets up the state space model, including the design, transition,
+        selection, and initial state matrices, as well as the parameters for the model.
+
+
+        Up to know just a draft implementation to test the working of the class and comparing
+        with the Custom model done in the Notebook (notebook/Making a Custom DFM.ipynb).
+        """
+
+        # Create symbolic variables for 1D state
+        x0 = self.make_and_register_variable("x0", shape=(1,))
+        P0 = self.make_and_register_variable("P0", shape=(1, 1))
+        factor_loading = self.make_and_register_variable("factor_loadings", shape=(2, 1))
+
+        factor_ar = self.make_and_register_variable("factor_ar", shape=())
+        sigma_f = self.make_and_register_variable("factor_sigma", shape=())
+
+        # Initialize matrices with correct dimensions
+        self.ssm["design", :, :] = np.array([[0.0], [0.0]])  # 2x1 matrix
+        self.ssm["transition", :, :] = np.array([[0.0]])  # 1x1 matrix
+        self.ssm["selection", :, :] = np.array([[1.0]])  # 1x1 matrix
+
+        # Set initial state and covariance
+        self.ssm["initial_state", :] = x0
+        self.ssm["initial_state_cov", :, :] = P0
+
+        # Set design matrix parameters
+        self.ssm["design", 0, 0] = factor_loading[0, 0]  # First observation loading
+        self.ssm["design", 1, 0] = factor_loading[1, 0]  # Second observation loading
+
+        # Set transition parameter (AR coefficient)
+        self.ssm["transition", 0, 0] = factor_ar
+
+        # Set state covariance
+        self.ssm["state_cov", 0, 0] = sigma_f