Learning Poisson State Space Model parameters for large number of samples

Questions version agnostic
5

Thanks a lot, this is really useful. I spend some time making it run, apparently, there is some issue with numpy version and PyMC version. I am wondering how the model runs. If I look at the below lines then Poisson uses the ground truth hidden states, not the hidden_states_pt.

mod.register_rv(hidden_states_pt, name='hidden_states', initval=pt.zeros((T, 2)))
obs = pm.Poisson('obs', pt.exp(hidden_states[:, 0]), observed=data)

Also, I am trying to convert it to a case where we have more than one number of samples-

T = 100
num_samples = 200
num_features = 10
latent_variables = 2

# Local level model
A = np.array([[1, 1],
                   [0, 1]])

Z = np.array([[1, 0]])
R = np.eye(2)

true_Q = np.diag(np.random.normal(scale=1e-2, size=2) ** 2)
hidden_states = np.zeros((num_samples, T, 2))
# hidden_states[0] = true_x0

data = np.zeros((num_samples, T))

fig, ax = plt.subplots()
for t in range(1, T):
    innov_t = np.random.multivariate_normal(mean=np.zeros(2), cov=true_Q)
    hidden_states[:,t,:] =  hidden_states[:,t-1,:]@A + R @ innov_t
    data_mu = np.exp(hidden_states[:,t,:]@Z.T)
    data[:,t] = np.squeeze(np.random.poisson(lam=data_mu))

ax.plot(data[2,:]);

with pm.Model() as mod:
    A = pt.as_tensor(A)
    R = pt.eye(2)
    Z = pt.as_tensor(Z)
    
    sigmas_Q = pm.HalfNormal('sigmas_Q', sigma=1000, size=2)
    Q = pt.diag(sigmas_Q)
    x0 = pt.zeros((num_samples,2))
    
    def step(x, A, R, Z, Q):
        innov = pm.MvNormal.dist(mu=0, tau=Q)
        next_x = x.dot(A) + innov
        return next_x, collect_default_updates([x,A,R,Z,Q], [next_x])
    
    hidden_states_pt, updates = pytensor.scan(step, 
                                           outputs_info=x0, 
                                           non_sequences=[A, R, Z, Q],
                                           n_steps=T, 
                                           strict=True)
    
    mod.register_rv(hidden_states_pt, name='hidden_states', initval=pt.zeros((num_samples, T, 2)))
    obs = pm.Poisson('obs', np.squeeze(pt.exp(hidden_states_pt[:, :, 0])), observed=data.T)

There is some assertion error assert “x.ndim == 2, return Apply(self, , [x.type()])”. I am not sure what is the issue in the multivariate case. I am new to PyMC and trying to learn it. Any help is greatly appreciated. Thanks a lot again!