State Space: Exogenous Vars - Conflicting dimensions for time (bug?)

[rOybOiii]

Hey PyMC & @jessegrabowski

I’m trying to use state space’s forecast method with exogenous variables, but I keep getting this error:

ValueError: conflicting sizes for dimension ‘time’: length 1107 on the data but length 8 on coordinate ‘time’

My scenario: I’m passing exogenous indicator variables totaling 18 columns and 8 rows long, for an 8 day forecast. The period I’m forecasting is not at the -very- end of my time series. I’m testing my exogenous variable’s ability to reduce variance on events and holidays.

Time series info:

Length: 1107 data points
Start: 2022-01-19
End: 2025-01-29

Relevant code:

// Define start date and forecast period
start_date, n_periods = pd.to_datetime("2024-1-25"), 8

// Extract exogenous indicator variables for the forecast period
scenario = {'data_exog': pd.DataFrame(df_holidays.loc[start_date:].iloc[:n_periods].to_numpy(dtype=float), 
                                      columns=df_holidays.columns)}

// Generate the forecast
forecasts = ss_mod.forecast(trace, start=start_date, periods=n_periods, scenario=scenario)

My second line of code is extracting the relevant portion of my exogenous variables and assigning it to the scenario.

Am I getting this error because I’m not forecasting at the very end of my time series? (edit: dumb question, I’ll test this later.)

Conda Environment-
PyMC Version: 5.20.1
PyMC Extras Version: 0.2.3
Python Version: Python 3.11.11

Thanks,
Roy

[rOybOiii]

Can offer data and Jupyter file. Lemme know.

-Roy

[jessegrabowski]

Can you provide a very simple example that reproduces the error?

[rOybOiii]

Back at the other forum, I do have my model and some data available. I have not had a moment yet to pare it down to something a little simpler, but Jonathan was able to run it in 41 minutes, after putting in some column names into the data.

I'll try to provide something even simpler soon - but work calls me now.

-Roy

[rOybOiii]

Hey @jessegrabowski :

I modified my model so it's a little simpler, with only 2 exogenous variables, and I deleted about 25-30% of the data points. Total run time is about 26 minutes for me, but with your M4 processor, I imagine it'll be closer to 8 minutes. Do you need something even simpler?

I ran this model using a fresh environment with pymc-extras 0.2.3.

Here's the Jupyter file

Here's the data

-Mike

[Dekermanjian]

Hey @jessegrabowski,

Here is a runnable example that I believe is simpler than the above. This runnable will result in the conflicting dimension on 'time':

from pymc_extras.statespace.utils.constants import (
    ALL_STATE_DIM,
    ALL_STATE_AUX_DIM,
    SHOCK_DIM,
    TIME_DIM
)
from pymc_extras.statespace.models.utilities import make_default_coords
from pymc_extras.statespace.core.statespace import PyMCStateSpace
import pytensor.tensor as pt
import pymc as pm
import numpy as np

data = np.random.normal(size=100)[:, np.newaxis]
exog = np.random.normal(size=100)[:, np.newaxis]

class AutoRegressiveThree(PyMCStateSpace):
    def __init__(self, k_exog=1):
        k_states = 3 + k_exog # size of the state vector x
        k_posdef = 1  # number of shocks (size of the state covariance matrix Q)
        k_endog = 1  # number of observed states
        self.k_exog = k_exog
        super().__init__(k_endog=k_endog, k_states=k_states, k_posdef=k_posdef)

    def make_symbolic_graph(self):
        exogenous_data = self.make_and_register_data("exogenous_data", shape=(None, self.k_exog))
        beta_exog = self.make_and_register_variable("beta_exog", shape=(self.k_exog,))
        
        n_obs = exogenous_data.shape[0]

        x0 = self.make_and_register_variable("x0", shape=(3,))
        P0 = self.make_and_register_variable("P0", shape=(4, 4))
        ar_params = self.make_and_register_variable("ar_params", shape=(3,))
        sigma_x = self.make_and_register_variable("sigma_x", shape=(1,))

        a = np.eye(4, k=-1)
        a[3, 2] = 0.
        a[3, 3] = 1.
        self.ssm["transition", :, :] = a
        self.ssm["transition", 0, :3] = ar_params

        self.ssm["selection", 0, 0] = 1.
        
        self.ssm["design"] = pt.concatenate(
                                                (
                                                    pt.expand_dims(pt.repeat(pt.as_tensor([[1.,0.,0.]]), n_obs, axis=0), 1),
                                                    pt.expand_dims(exogenous_data, 1)
                                                ), axis=2
                                            )

        self.ssm["initial_state", :] = pt.concatenate((x0, beta_exog), axis=0)
        self.ssm["initial_state_cov", :, :] = P0
        self.ssm["state_cov", :, :] = sigma_x

    @property
    def data_names(self) -> list[str]:
        """
        Names of data variables expected by the model.

        This does not include the observed data series, which is automatically handled by PyMC. This property only
        needs to be implemented for models that expect exogenous data.
        """
        return ["exogenous_data"]
    
    @property
    def data_info(self):
        """
        Information about Data variables that need to be declared in the PyMC model block.

        Returns a dictionary of data_name: dictionary of property-name:property description pairs. The return value is
        used by the ``_print_data_requirements`` method, to print a message telling users how to define the necessary
        data for the model. Each dictionary should have the following key-value pairs:
            * key: "shape", value: a tuple of integers
            * key: "dims", value: tuple of strings
        """
        return {"exogenous_data": {"shape": (None, self.k_exog), "dims": (TIME_DIM, "exog_dims")}}

    @property
    def param_names(self):
        return ["x0", "P0", "ar_params", "sigma_x", "beta_exog"]

    @property
    def state_names(self):
        # Since the three states are lags of the data, i'll call them L1, L2 L3
        return ["L1.data", "L2.data", "L3.data", "beta"]

    @property
    def shock_names(self):
        # There is one shock, called the "innovations" in the literature, so i'll go with that
        return ["innovations"]

    @property
    def observed_states(self):
        # Inspired, I know
        return ["data"]

    @property
    def param_dims(self):
        # This needs to map the 4 parameters to the names of the coords.
        # There are special standardized names to use here. You can import them from
        # pymc_experimental.statespace.utils.constants

        # Not the best system. Something to improve on in the future.
        return {
            "x0": (ALL_STATE_DIM,),
            "P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
            "ar_params": ("ar_lags",),
            "sigma_x": (SHOCK_DIM,),
            "beta_exog": ("exog_dims",)
        }

    @property
    def coords(self):
        # This function puts coords on all those statespace matrices (x0, P0, c, d, T, Z, R, H, Q)
        # and also on the different filter outputs so you don't have to worry about it. You only need to set
        # the coords for the dims unique to your model. In this case, it's just "ar_lags"
        coords = make_default_coords(self)
        coords.update({"ar_lags": [1, 2, 3], "exog_dims": ['beta0']}) # need to define n_obs upfront

        return coords

    @property
    def param_info(self):
        # This needs to return a dictionary where the keys are the parameter names, and the values are a
        # dictionary. The value dictionary should have the following keys: "shape", "constraints", and "dims".

        info = {
            "x0": {
                "shape": (self.k_states - 1,),
                "constraints": "None",
            },
            "P0": {
                "shape": (self.k_states, self.k_states),
                "constraints": "Positive Semi-definite",
            },
            "sigma_x": {
                "shape": (self.k_posdef,),
                "constraints": "Positive",
            },
            "ar_params": {
                "shape": (3,),
                "constraints": "None",
            },
            "beta_exog": {
                "shape": (self.k_exog,),
                "constraints": "None",
            },
        }

        # Lazy way to add the dims without making any typos
        for name in self.param_names:
            info[name]["dims"] = self.param_dims[name]

        return info
    
ar3 = AutoRegressiveThree()
with pm.Model(coords=ar3.coords) as pymc_mod:
    exogenous_data = pm.Data("exogenous_data", exog, dims=[TIME_DIM, 'exog_dims'])
    beta_exog = pm.Normal("beta_exog", 0, 1, dims=["exog_dims"])

    x0 = pm.Deterministic(
        "x0",
        pt.as_tensor([0., 0., 0.]),
        
    )
    P0 = pm.Deterministic("P0", pt.eye(4) * 10., dims=["state", "state_aux"])

    ar_params = pm.Normal("ar_params", sigma=0.25, dims=["ar_lags"])
    sigma_x = pm.Exponential("sigma_x", 1, dims=["shock"])
    

    ar3.build_statespace_graph(data=data, mode="JAX")
    idata = pm.sample(tune=50, draws=50, nuts_sampler="numpyro")

forecasts = ar3.forecast(idata, scenario={"exogenous_data": exog[:4]}, start=99, periods=4)

And this runnable avoids the error by giving the exogenous data a different name than 'time' for the coord. In this case you also need to provide the length of the coord upfront. This works without any error:

from pymc_extras.statespace.utils.constants import (
    ALL_STATE_DIM,
    ALL_STATE_AUX_DIM,
    SHOCK_DIM,
    TIME_DIM
)
from pymc_extras.statespace.models.utilities import make_default_coords
from pymc_extras.statespace.core.statespace import PyMCStateSpace
import pytensor.tensor as pt
import pymc as pm
import numpy as np

data = np.random.normal(size=100)[:, np.newaxis]
exog = np.random.normal(size=100)[:, np.newaxis]

class AutoRegressiveThree(PyMCStateSpace):
    def __init__(self, k_exog=1):
        k_states = 3 + k_exog # size of the state vector x
        k_posdef = 1  # number of shocks (size of the state covariance matrix Q)
        k_endog = 1  # number of observed states
        self.k_exog = k_exog
        super().__init__(k_endog=k_endog, k_states=k_states, k_posdef=k_posdef)

    def make_symbolic_graph(self):
        exogenous_data = self.make_and_register_data("exogenous_data", shape=(None, self.k_exog))
        beta_exog = self.make_and_register_variable("beta_exog", shape=(self.k_exog,))
        
        n_obs = exogenous_data.shape[0]

        x0 = self.make_and_register_variable("x0", shape=(3,))
        P0 = self.make_and_register_variable("P0", shape=(4, 4))
        ar_params = self.make_and_register_variable("ar_params", shape=(3,))
        sigma_x = self.make_and_register_variable("sigma_x", shape=(1,))

        a = np.eye(4, k=-1)
        a[3, 2] = 0.
        a[3, 3] = 1.
        self.ssm["transition", :, :] = a
        self.ssm["transition", 0, :3] = ar_params

        self.ssm["selection", 0, 0] = 1.
        
        self.ssm["design"] = pt.concatenate(
                                                (
                                                    pt.expand_dims(pt.repeat(pt.as_tensor([[1.,0.,0.]]), n_obs, axis=0), 1),
                                                    pt.expand_dims(exogenous_data, 1)
                                                ), axis=2
                                            )

        self.ssm["initial_state", :] = pt.concatenate((x0, beta_exog), axis=0)
        self.ssm["initial_state_cov", :, :] = P0
        self.ssm["state_cov", :, :] = sigma_x

    @property
    def data_names(self) -> list[str]:
        """
        Names of data variables expected by the model.

        This does not include the observed data series, which is automatically handled by PyMC. This property only
        needs to be implemented for models that expect exogenous data.
        """
        return ["exogenous_data"]
    
    @property
    def data_info(self):
        """
        Information about Data variables that need to be declared in the PyMC model block.

        Returns a dictionary of data_name: dictionary of property-name:property description pairs. The return value is
        used by the ``_print_data_requirements`` method, to print a message telling users how to define the necessary
        data for the model. Each dictionary should have the following key-value pairs:
            * key: "shape", value: a tuple of integers
            * key: "dims", value: tuple of strings
        """
        return {"exogenous_data": {"shape": (None, self.k_exog), "dims": ('n_obs', "exog_dims")}}

    @property
    def param_names(self):
        return ["x0", "P0", "ar_params", "sigma_x", "beta_exog"]

    @property
    def state_names(self):
        # Since the three states are lags of the data, i'll call them L1, L2 L3
        return ["L1.data", "L2.data", "L3.data", "beta"]

    @property
    def shock_names(self):
        # There is one shock, called the "innovations" in the literature, so i'll go with that
        return ["innovations"]

    @property
    def observed_states(self):
        # Inspired, I know
        return ["data"]

    @property
    def param_dims(self):
        # This needs to map the 4 parameters to the names of the coords.
        # There are special standardized names to use here. You can import them from
        # pymc_experimental.statespace.utils.constants

        # Not the best system. Something to improve on in the future.
        return {
            "x0": (ALL_STATE_DIM,),
            "P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
            "ar_params": ("ar_lags",),
            "sigma_x": (SHOCK_DIM,),
            "beta_exog": ("exog_dims",)
        }

    @property
    def coords(self):
        # This function puts coords on all those statespace matrices (x0, P0, c, d, T, Z, R, H, Q)
        # and also on the different filter outputs so you don't have to worry about it. You only need to set
        # the coords for the dims unique to your model. In this case, it's just "ar_lags"
        coords = make_default_coords(self)
        coords.update({"ar_lags": [1, 2, 3], "exog_dims": ['beta0'], "n_obs": np.arange(100)}) # need to define n_obs upfront

        return coords

    @property
    def param_info(self):
        # This needs to return a dictionary where the keys are the parameter names, and the values are a
        # dictionary. The value dictionary should have the following keys: "shape", "constraints", and "dims".

        info = {
            "x0": {
                "shape": (self.k_states - 1,),
                "constraints": "None",
            },
            "P0": {
                "shape": (self.k_states, self.k_states),
                "constraints": "Positive Semi-definite",
            },
            "sigma_x": {
                "shape": (self.k_posdef,),
                "constraints": "Positive",
            },
            "ar_params": {
                "shape": (3,),
                "constraints": "None",
            },
            "beta_exog": {
                "shape": (self.k_exog,),
                "constraints": "None",
            },
        }

        # Lazy way to add the dims without making any typos
        for name in self.param_names:
            info[name]["dims"] = self.param_dims[name]

        return info
    
ar3 = AutoRegressiveThree()
with pm.Model(coords=ar3.coords) as pymc_mod:
    exogenous_data = pm.Data("exogenous_data", exog, dims=['n_obs', 'exog_dims'])
    beta_exog = pm.Normal("beta_exog", 0, 1, dims=["exog_dims"])

    x0 = pm.Deterministic(
        "x0",
        pt.as_tensor([0., 0., 0.]),
        
    )
    P0 = pm.Deterministic("P0", pt.eye(4) * 10., dims=["state", "state_aux"])

    ar_params = pm.Normal("ar_params", sigma=0.25, dims=["ar_lags"])
    sigma_x = pm.Exponential("sigma_x", 1, dims=["shock"])
    

    ar3.build_statespace_graph(data=data, mode="JAX")
    idata = pm.sample(tune=50, draws=50, nuts_sampler="numpyro")

forecasts = ar3.forecast(idata, scenario={"exogenous_data": exog[:4]}, start=99, periods=4)

I guess this could either be a bug with updating coords which is avoided when defining a different name on the exogenous coords. I am still trying to find exactly where this bug could be occurring in the forecast() method.

EDIT:
this is the case for version 0.2.5. With the fix where you pass in the scenario to _kalman_filter_outputs_from_dummy_graph() both the above will work.

[jessegrabowski]

This is still a pretty beefy example, can you reproduce it just using the statespace.structural.ExogenousRegression component?

When I say "minimial", i'm hoping for like ~10 lines

[Dekermanjian]

Hey @jessegrabowski, what about this:

@pytest.fixture(scope="session")
def exog_data(rng):
    # simulate data
    index = pd.date_range(start="2023-05-01", end="2023-05-10", freq="D")
    T, k = len(index), 1
    data = np.zeros((T, k))
    idx = rng.choice(T, size=3, replace=False)
    cols = rng.choice(k, size=3, replace=True)

    data[idx, cols] = 1

    df_holidays = pd.DataFrame(data, index=index, columns=["Holiday 1"])

    data = rng.normal(size=(T, 1))
    nan_locs = rng.choice(T, size=3, replace=False)
    data[nan_locs] = np.nan
    y = pd.DataFrame(data, index=index, columns=["sales"])
    return pd.concat((df_holidays, y), axis=1)


@pytest.fixture(scope="session")
def exog_ss_mod(exog_data):
    level_trend = st.LevelTrendComponent(order=1, innovations_order=[0])
    exog = st.RegressionComponent(
        name="exog",  # Name of this exogenous variable component
        k_exog=1,  # Only one exogenous variable now
        innovations=False,  # Typically fixed effect (no stochastic evolution)
        state_names=exog_data[['Holiday 1']].columns.tolist(),
    )

    combined_model = level_trend + exog
    return combined_model.build()

@pytest.fixture(scope="session")
def exog_pymc_mod(exog_ss_mod, exog_data):
    # define pymc model
    with pm.Model(coords=exog_ss_mod.coords) as struct_model:
        P0_diag = pm.Gamma("P0_diag", alpha=2, beta=4, dims=["state"])
        P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=["state", "state_aux"])

        initial_trend = pm.Normal("initial_trend", mu=[0], sigma=[0.005], dims=["trend_state"])

        data_exog = pm.Data("data_exog", exog_data['Holiday 1'].values[:, None], dims=["time", "exog_state"])
        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["exog_state"])

        exog_ss_mod.build_statespace_graph(exog_data['sales'], mode="JAX")

    return struct_model

@pytest.mark.filterwarnings("ignore:Provided data contains missing values.")
@pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
@pytest.mark.filterwarnings("ignore:No time index found on the supplied data. A simple range index will be automatically generated.")
@pytest.mark.filterwarnings("ignore:Skipping `CheckAndRaise` Op")
def test_foreacast_valid_index(exog_pymc_mod, exog_ss_mod, exog_data):
    # Regression test for issue reported at  https://github.com/pymc-devs/pymc-extras/issues/424
    with exog_pymc_mod:
        idata = pm.sample_prior_predictive()

    # Define start date and forecast period
    start_date, n_periods = pd.to_datetime("2023-05-05"), 5

    # Extract exogenous data for the forecast period
    scenario = {
        "data_exog": pd.DataFrame(
            exog_data[['Holiday 1']].loc[start_date:].iloc[:n_periods], columns=exog_data[['Holiday 1']].columns
        )
    }

    # Generate the forecast
    forecasts = exog_ss_mod.forecast(idata.prior, scenario=scenario, use_scenario_index=True)

It is more than 10 lines but I broke it out into smaller pieces and this passes tests on the last push to the PR you did a couple of days ago.

[jessegrabowski]

Yeah this is great! Do you want to check out my branch and just push the new test to it? No is an acceptable answer -- I'll do it in that case :)

[Dekermanjian]

@jessegrabowski I'd be happy to do that. I am still new to contributing to open source and I am not really sure what the etiquette is like. So, I want to make sure I don't do anything that I'm not supposed to be doing. I will push everything later today!

[jessegrabowski]

Great! I'll look for the push and give some feedback