Error in sample_conditional_posterior with PyMCStateSpace model

[rrebane]

Related PyMC Discourse post: https://discourse.pymc.io/t/stochastic-volatility-with-pymcstatespace/16354

As noted in one of the answers, the problem is that the model needs to record the shape of h after done fitting, but it does not.

After fitting the model, when running sample_conditional_posterior I get the following error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[300], line 1
----> 1 cond_post = vkm.sample_conditional_posterior(idata)

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py:1345, in PyMCStateSpace.sample_conditional_posterior(self, idata, random_seed, **kwargs)
   1318 def sample_conditional_posterior(
   1319     self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs
   1320 ):
   1321     """
   1322     Sample from the conditional posterior; that is, given parameter draws from the posterior distribution,
   1323     compute Kalman filtered trajectories. Trajectories are drawn from a single multivariate normal with mean and
   (...)
   1342          "predicted_posterior", and "smoothed_posterior".
   1343     """
-> 1345     return self._sample_conditional(idata, "posterior", random_seed, **kwargs)

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py:1117, in PyMCStateSpace._sample_conditional(self, idata, group, random_seed, data, **kwargs)
   1101 group_idata = getattr(idata, group)
   1103 with pm.Model(coords=self._fit_coords) as forward_model:
   1104     (
   1105         [
   1106             x0,
   1107             P0,
   1108             c,
   1109             d,
   1110             T,
   1111             Z,
   1112             R,
   1113             H,
   1114             Q,
   1115         ],
   1116         grouped_outputs,
-> 1117     ) = self._kalman_filter_outputs_from_dummy_graph(data=data)
   1119     for name, (mu, cov) in zip(FILTER_OUTPUT_TYPES, grouped_outputs):
   1120         dummy_ll = pt.zeros_like(mu)

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py:1001, in PyMCStateSpace._kalman_filter_outputs_from_dummy_graph(self, data, data_dims, scenario)
    998     scenario = dict()
   1000 pm_mod = modelcontext(None)
-> 1001 self._build_dummy_graph()
   1002 self._insert_random_variables()
   1004 for name in self.data_names:

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc_extras/statespace/core/statespace.py:965, in PyMCStateSpace._build_dummy_graph(self)
    954 """
    955 Build a dummy computation graph for the state space model matrices.
    956 
   (...)
    962     A list of pm.Flat variables representing all parameters estimated by the model.
    963 """
    964 for name in self.param_names:
--> 965     pm.Flat(
    966         name,
    967         shape=self._name_to_variable[name].type.shape,
    968         dims=self._fit_dims.get(name, None),
    969     )

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc/distributions/distribution.py:511, in Distribution.__new__(cls, name, rng, dims, initval, observed, total_size, transform, default_transform, *args, **kwargs)
    508     elif observed is not None:
    509         kwargs["shape"] = tuple(observed.shape)
--> 511 rv_out = cls.dist(*args, **kwargs)
    513 rv_out = model.register_rv(
    514     rv_out,
    515     name,
   (...)
    521     initval=initval,
    522 )
    524 # add in pretty-printing support

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc/distributions/continuous.py:369, in Flat.dist(cls, **kwargs)
    367 @classmethod
    368 def dist(cls, **kwargs):
--> 369     res = super().dist([], **kwargs)
    370     return res

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pymc/distributions/distribution.py:580, in Distribution.dist(cls, dist_params, shape, **kwargs)
    577     ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp
    579 create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
--> 580 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
    582 _add_future_warning_tag(rv_out)
    583 return rv_out

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/random/op.py:315, in RandomVariable.__call__(self, size, name, rng, dtype, *args, **kwargs)
    313     props["dtype"] = dtype
    314     new_op = type(self)(**props)
--> 315     return new_op.__call__(
    316         *args, size=size, name=name, rng=rng, dtype=dtype, **kwargs
    317     )
    319 res = super().__call__(rng, size, *args, **kwargs)
    321 if name is not None:

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/random/op.py:319, in RandomVariable.__call__(self, size, name, rng, dtype, *args, **kwargs)
    314     new_op = type(self)(**props)
    315     return new_op.__call__(
    316         *args, size=size, name=name, rng=rng, dtype=dtype, **kwargs
    317     )
--> 319 res = super().__call__(rng, size, *args, **kwargs)
    321 if name is not None:
    322     res.name = name

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/graph/op.py:293, in Op.__call__(self, name, return_list, *inputs, **kwargs)
    249 def __call__(
    250     self, *inputs: Any, name=None, return_list=False, **kwargs
    251 ) -> Variable | list[Variable]:
    252     r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    253 
    254     This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    291 
    292     """
--> 293     node = self.make_node(*inputs, **kwargs)
    294     if name is not None:
    295         if len(node.outputs) == 1:

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/random/op.py:349, in RandomVariable.make_node(self, rng, size, *dist_params)
    326 def make_node(self, rng, size, *dist_params):
    327     """Create a random variable node.
    328 
    329     Parameters
   (...)
    347 
    348     """
--> 349     size = normalize_size_param(size)
    351     dist_params = tuple(
    352         as_tensor_variable(p) if not isinstance(p, Variable) else p
    353         for p in dist_params
    354     )
    356     if rng is None:

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/random/utils.py:190, in normalize_size_param(shape)
    186     raise TypeError(
    187         "Parameter size must be None, an integer, or a sequence with integers."
    188     )
    189 else:
--> 190     shape = cast(as_tensor_variable(shape, ndim=1, dtype="int64"), "int64")
    192     if not isinstance(shape, Constant):
    193         # This should help ensure that the length of non-constant `size`s
    194         # will be available after certain types of cloning (e.g. the kind
    195         # `Scan` performs)
    196         shape = specify_shape(shape, (get_vector_length(shape),))

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/__init__.py:50, in as_tensor_variable(x, name, ndim, **kwargs)
     18 def as_tensor_variable(
     19     x: TensorLike, name: str | None = None, ndim: int | None = None, **kwargs
     20 ) -> "TensorVariable":
     21     """Convert `x` into an equivalent `TensorVariable`.
     22 
     23     This function can be used to turn ndarrays, numbers, `ScalarType` instances,
   (...)
     48 
     49     """
---> 50     return _as_tensor_variable(x, name, ndim, **kwargs)

File ~/miniconda3/envs/py312/lib/python3.12/functools.py:907, in singledispatch.<locals>.wrapper(*args, **kw)
    903 if not args:
    904     raise TypeError(f'{funcname} requires at least '
    905                     '1 positional argument')
--> 907 return dispatch(args[0].__class__)(*args, **kw)

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/basic.py:177, in _as_tensor_Sequence(x, name, ndim, dtype, **kwargs)
    172     # In this case, we have at least one non-`Constant` term, so we
    173     # couldn't get an underlying non-symbolic sequence of objects and we to
    174     # symbolically join terms.
    175     return stack(x)
--> 177 return constant(x, name=name, ndim=ndim, dtype=dtype)

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/tensor/basic.py:223, in constant(x, name, ndim, dtype)
    220     else:
    221         x = x.data
--> 223 x_ = ps.convert(x, dtype=dtype)
    225 if ndim is not None:
    226     if x_.ndim < ndim:

File ~/miniconda3/envs/py312/lib/python3.12/site-packages/pytensor/scalar/basic.py:249, in convert(x, dtype)
    247     if dtype == "floatX":
    248         dtype = config.floatX
--> 249     x_ = np.asarray(x).astype(dtype)
    250 else:
    251     # In this case, this function should infer the dtype according to the
    252     # autocasting rules. See autocasting above.
    253     x_ = None

TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'

The model definition and the model sampling block looks like this:

import pymc as pm
import pytensor.tensor as pt

from pymc_extras.statespace.utils.constants import (
    ALL_STATE_DIM,
    ALL_STATE_AUX_DIM,
    OBS_STATE_DIM,
    SHOCK_DIM,
    TIME_DIM,
)
from pymc_extras.statespace.core.statespace import PyMCStateSpace
from pymc_extras.statespace.models.utilities import make_default_coords

class VlekkeKoopmanMellens2021(PyMCStateSpace):
    def __init__(self):
        k_states = 5 # size of the state vector x
        k_posdef = 5 # number of shocks (size of the state covariance matrix Q)
        k_endog = 1  # number of observed states

        super().__init__(k_endog=k_endog, k_states=k_states, k_posdef=k_posdef)

    def make_symbolic_graph(self):
        exog_data = self.make_and_register_data('exog_data', shape=(None, self.k_states))
        x0 = self.make_and_register_variable('x0', shape=(self.k_states,))
        P0 = self.make_and_register_variable('P0', shape=(self.k_states, self.k_states))
        h = self.make_and_register_variable('h', shape=(None, self.k_endog))
        sigma_sq_eta = self.make_and_register_variable('sigma_sq_eta', shape=(self.k_posdef,))

        self.ssm['transition', :, :] = np.eye(self.k_states)
        assert self.k_states == self.k_posdef
        self.ssm['selection', :, :] = np.eye(self.k_states, self.k_posdef)

        self.ssm['initial_state', :] = x0
        self.ssm['initial_state_cov', :, :] = P0
        self.ssm['design'] = pt.expand_dims(exog_data, 1)  
        self.ssm['obs_cov'] = pt.vectorize(pt.diag, '(a)->(a,b)')(pt.exp(h))
        self.ssm['state_cov', *np.diag_indices(self.k_posdef)] = sigma_sq_eta

    @property
    def param_names(self):
        return ['x0', 'P0', 'h', 'sigma_sq_eta']

    @property
    def state_names(self):
        return ['beta', 'theta', 'phi', 'gamma_1', 'gamma_2']

    @property
    def shock_names(self):
        return [f'{state}_shock' for state in self.state_names]

    @property
    def observed_states(self):
        return ['pi']

    @property
    def param_dims(self):
        return {
            'x0': (ALL_STATE_DIM,),
            'P0': (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
            'h': (TIME_DIM, OBS_STATE_DIM),
            'sigma_sq_eta': (SHOCK_DIM,),
        }

    @property
    def coords(self):
        coords = make_default_coords(self)
        coords.update({'exog_data_dim': ['u_gap', 'pi_p', 'pi_h', 'rel_i_p', 'oil_p']})

        return coords

    @property
    def param_info(self):
        info = {
            'x0': {
                'shape': (self.k_states,),
                'constraints': 'None',
            },
            'P0': {
                'shape': (self.k_states, self.k_states),
                'constraints': 'Positive Semi-definite',
            },
            'h': {
                'shape': (None, self.k_endog),
                'constraints': 'Time-varying',
            },
            'sigma_sq_eta': {
                'shape': (self.k_posdef,),
                'constraints': 'Positive',
            },
        }

        for name in self.param_names:
            info[name]['dims'] = self.param_dims[name]

        return info

    @property
    def data_names(self):
        return ['exog_data']
    
    @property
    def data_info(self):
        return {
            'exog_data': {
                'shape': (None, self.k_states),
                'dims': (TIME_DIM, 'exog_data_dim')
            }
        }

# Note: model_df is omitted for brevity
exog_data_df = model_df[['u_gap', 'e_p_cpi_1_meb', 'e_h_cpi_1_meb', 'rel_i_p', 'oil_p']]
obs_df = model_df[['cpi_meb']]

vkm = VlekkeKoopmanMellens2021()

b_40 = np.zeros(5, dtype="float")
P_40 = np.eye(5) * 0.1
h0_mu, h0_sigma = 0.0, 0.25

with pm.Model(coords=vkm.coords) as mod:
    x0 = pm.Deterministic('x0', pt.as_tensor(b_40), dims=[ALL_STATE_DIM])
    P0 = pm.Deterministic('P0', pt.as_tensor(P_40), dims=[ALL_STATE_DIM, ALL_STATE_AUX_DIM])
    h0 = pm.Normal.dist(mu=h0_mu, sigma=h0_sigma)

    # For sigma_sq_nu and sigma_sq_eta priors, see table 2 in appendix C in the refrence paper
    sigma_sq_nu = pm.InverseGamma('sigma_sq_nu', alpha=(vkm.k_states + 2) / 2, beta=0.454 / 2)
    sigma_sq_eta = pm.InverseGamma('sigma_sq_eta', alpha=(vkm.k_states + 2) / 2, beta=0.268 / 2, dims=[SHOCK_DIM])

    h = pm.GaussianRandomWalk('h', sigma=pt.sqrt(sigma_sq_nu), init_dist=h0, shape=obs_df.shape, dims=[TIME_DIM, OBS_STATE_DIM])
    sigma_sq_eps = pm.Deterministic('sigma_sq_eps', pt.exp(h), dims=[TIME_DIM, OBS_STATE_DIM])
    
    exog_data = pm.Data('exog_data', exog_data_df, dims=[TIME_DIM, 'exog_data_dim'])

    vkm.build_statespace_graph(data=obs_df, mode='JAX')
    idata = pm.sample(nuts_sampler='numpyro', chains=4, draws=1000)