Did state-space models change the way they handle time-varying matrices?

Questions v5
,
1

I recently installed pymc on my windows computer using the pip installer. The following code works on my linux machine, which has an older version installed via conda. It doesn’t work on my windows laptop.

import numpy as np
import pandas as pd
import pytensor
import pytensor.tensor as pt
import pymc as pm

from pymc_extras.statespace.core import PyMCStateSpace
import matplotlib.pyplot as plt

# version = [0,0]
# version = [0,1]
# version = [1,0]
version = [1,1]

class SSM(PyMCStateSpace):
   def __init__(self, N):
      self.k_states = 1 
      self.k_posdef = 1
      self.k_endog = 1

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

   def make_symbolic_graph(self):
      y_0 = self.make_and_register_variable("y_0", shape=(self.k_states,))
      self.ssm["initial_state",:] = y_0

      # define transition matrix
      M = self.make_and_register_variable('M', shape=(self.N,))
      T = pt.zeros((self.N, self.k_states, self.k_states))
      T = T[:,0,0].set( M )
      self.ssm["transition"] = T[:,:,:]

      if version[0] == 0:
         # define observation matrix
         self.ssm['design',0,0] = 1

      if version[0] == 1:
         # define observation matrix
         Z = pt.zeros((self.N, self.k_endog, self.k_states))
         Z = Z[:,0,0].set( 1 )
         self.ssm['design'] = Z[:,:,:]

      if version[1] == 0:
         # define observation noise
         sigma = self.make_and_register_variable('sigma', shape=())
         self.ssm['obs_cov',0,0] = sigma

      if version[1] == 1:
         # define observation noise
         sigma = self.make_and_register_variable('sigma', shape=())
         N_o = pt.zeros((self.N, self.k_endog, self.k_endog))
         N_o = N_o[:,0,0].set( sigma )
         self.ssm['obs_cov'] = N_o[:,:,:]

   @property
   def param_names(self):
      return ["y_0", "M", "sigma"]


t = [i for i in range(100)]
omega_true = 0.002
sigma_true = 0.05
X_0 = 2
X = []
X_prev = X_0
for i in t:
   X_next = (1-omega_true*i)*X_prev
   X += [X_next]
   X_prev = X_next

Y = X + np.random.normal(0, sigma_true, size=len(t))

sm = SSM(len(t))

with pm.Model() as pymc_mod:
   tt = pm.Data('tt', t)
   y_0 = pm.Normal("y_0", mu=3, sigma=1, shape=(1,))
   omega = pm.Normal('omega', mu=0, sigma=1)
   sigma = pm.Normal('sigma', mu=0, sigma=1)
   M = pm.Deterministic('M', 1-omega*tt)

   sm.build_statespace_graph(data=np.asmatrix(Y).T)
   idata = pm.sample()

The code fails with the following errors:

version [0,0]

ValueError: When compiling the inner function of scan (the function called by scan in each of its iterations) the following error has been encountered: The initial state (outputs_info in scan nomenclature) of variable SetSubtensor{:stop}.0 (argument number 1) has 2 dimension(s), while the corresponding variable in the result of the inner function of scan (fn) has 2 dimension(s) (it should be one less than the initial state). For example, if the inner function of scan returns a vector of size d and scan uses the values of the previous time-step, then the initial state in scan should be a matrix of shape (1, d). The first dimension of this matrix corresponds to the number of previous time-steps that scan uses in each of its iterations. In order to solve this issue if the two variable currently have the same dimensionality, you can increase the dimensionality of the variable in the initial state of scan by using dimshuffle or shape_padleft.

version [0,1]

ValueError: Incompatible shared dimension for dot product: (1, 1), (100, 1)

versions [1,0] and [1,1]

ValueError: Incompatible shared dimension for dot product: (1, 100), (1, 100)

2

I think so, CC @jessegrabowski

3

Yes, I changed it in this PR. I’m trying to get to full batch support for statespace. Previously, a matrix/vector was time varying if it had exactly 1 extra dimension. If we allow batch support, that’s not enough to pin that down – the extra dimension could be some arbitrary batch dimension. So now you need to explicitly tell the model that an object is time varying with the declare_time_varying method. Your code should now look like this:

class SSM(PyMCStateSpace):
   def __init__(self, N):
      self.k_states = 1
      self.k_posdef = 1
      self.k_endog = 1

      self.N = N

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

   def make_symbolic_graph(self):
      y_0 = self.make_and_register_variable("y_0", shape=(self.k_states,))
      self.ssm["initial_state",:] = y_0

      # define transition matrix
      M = self.make_and_register_variable('M', shape=(self.N,))
      T = pt.zeros((self.N, self.k_states, self.k_states))
      T = T[:,0,0].set( M )
      self.ssm["transition"] = T[:,:,:]

      # Big change here!
      self.ssm.declare_time_varying("transition")

      if version[0] == 0:
         # define observation matrix
         self.ssm['design',0,0] = 1

      if version[0] == 1:
         # define observation matrix
         Z = pt.zeros((self.N, self.k_endog, self.k_states))
         Z = Z[:,0,0].set( 1 )
         self.ssm['design'] = Z[:,:,:]
         self.ssm.declare_time_varying("design")

      if version[1] == 0:
         # define observation noise
         sigma = self.make_and_register_variable('sigma', shape=())
         self.ssm['obs_cov',0,0] = sigma

      if version[1] == 1:
         # define observation noise
         sigma = self.make_and_register_variable('sigma', shape=())
         N_o = pt.zeros((self.N, self.k_endog, self.k_endog))
         N_o = N_o[:,0,0].set( sigma )
         self.ssm['obs_cov'] = N_o[:,:,:]
         self.ssm.declare_time_varying("obs_cov")

   @property
   def param_names(self):
      return ["y_0", "M", "sigma"]

Unrelated, but you also have a prior on sigma that is allowed to be negative. That makes the samper’s job harder, because it has to learn to discard all the negative covariances.

4

Thanks! That seems to have fixed everything.