Posterior predictive on testing data in State Space Models

I have been working on modeling the below State Space Model

X_{i,t+1} \sim N(\mu_t, Q)\\ \mu_{i,t} = Ax_{i,t} + BU_{i}\\ y_{i,t} \sim Poisson(\lambda_{i,t}) \\\lambda_{i,t} = C \exp(X_{i,t})

where X_{i,t+1} is the hidden state for ith sample at time t, U_i is known exogenous variables, \lambda_{i,t} is the prediction for ith sample at time t. Q, A, B and C are parameters of the model. Below is the code for the model

with pm.Model() as mod:

    sd_dist = pm.Exponential.dist(1.0, shape=latent_variables)
    mu1 = np.zeros(latent_variables)
    chol1, corr, stds = pm.LKJCholeskyCov('chol_cov1', n=latent_variables, eta=2,
    sd_dist=sd_dist, compute_corr=True)
    pred_Z = pm.MvNormal('pred_Z', mu=mu1, chol=chol1, shape=(latent_variables))
    mu2 = np.zeros(num_features)
    chol2, corr, stds = pm.LKJCholeskyCov('chol_cov2', n=num_features, eta=2,
    sd_dist=sd_dist, compute_corr=True)
    pred_B = pm.MvNormal('B', mu=mu2, chol=chol2, shape=(num_features))
    A_pt = pt.as_tensor_variable(A_block)

    sigmas_Q = pm.HalfNormal('sigmas_Q', sigma=1, shape= (latent_variables))
    Q = pt.diag(sigmas_Q)
    u = pm.MutableData("input", u_scale)
    #x0 = pm.MutableData("x_init", x_0)
    data = pm.MutableData("data", data_train)
    samples = u.shape[0]
    x0 = pt.zeros((num_samples , latent_variables))
    def step(x, A, Q, samples):
        innov = pm.MvNormal.dist(mu=0, tau=Q, size = (samples))
        next_x = pt.nlinalg.matrix_dot(x,A) + innov
        return next_x, collect_default_updates([x, A, Q, samples], [next_x])
    hidden_states_pt, updates = pytensor.scan(step, 
                                              non_sequences=[A_pt, Q, samples],
    mod.register_rv(hidden_states_pt, name='hidden_states', initval=pt.zeros((T, samples, latent_variables)))
    #pred_Z = pm.HalfNormal('pred_Z', sigma=1, size= latent_variables)
    temp1 = pt.transpose(pt.nlinalg.matrix_dot(u,pred_B))
    #hidden_states_pt = pt.reshape(hidden_states_pt,(T,num_samples,latent_variables))
    #lambdas = pt.exp(hidden_states_pt[:, np.arange(0, num_samples * latent_variables, 2)]+temp1)
    lambdas = pt.exp(pt.squeeze(pt.nlinalg.matrix_dot(hidden_states_pt,pred_Z))+temp1)

    obs = pm.Poisson('obs', lambdas, observed=data)

I estimate the model parameters using the training data \lambda and U. Now I want to do some prediction on test data U_{test}. I was thinking of using posterior predictive samples but that won’t work because the hidden states X need to be estimated again for the new sample. What will be the best way to get predictions for the new samples and also estimate the hidden states?

Do these Utest happen after the U in your original model? In that case you can create a new model which starts where U “left off” and call posterior predictive there.

I have an example in this WIP notebook with a single timeseries, but the same concept applies. You can sample new unobserved variables in posterior predictive.

1 Like

Thanks! This is really helpful!

1 Like