title: Reproducing Hamilton 1989 Markov Switching Model in PyMC3
body: <p>I am trying to understand of Hamilton's 1989 Markov-Switching Autoregression model. I attempt to reproduce the results with a Bayesian twist. I wrote a number of PyMC3 models using <a href="https://ericmjl.github.io/essays-on-data-science/machine-learning/markov-models/" rel="nofollow noreferrer">Eric Ma's tutorial about HMM's in PyMC3</a> and the latest iteration can be found below.</p>
  <p>Without autoregression, the model converges to mu values close to Hamilton's (1.16 and -0.36) and realistic transition probabilities.
  However, when autoregression is added, the model fails to converge and fit coefficients close to Hamilton's results. The transition probabilities are fit especially poorly.</p>
  <p>What am I missing?</p>
  <pre><code># %%
  import pymc3 as pm
  import theano.tensor as tt
  import theano.tensor.slinalg as sla  # theano-wrapped scipy linear algebra
  import theano.tensor.nlinalg as nla  # theano-wrapped numpy linear algebra
  import theano
  
  theano.config.gcc.cxxflags = &quot;-Wno-c++11-narrowing&quot;
  
  import pandas as pd
  dta = pd.read_stata('https://www.stata-press.com/data/r14/rgnp.dta').iloc[1:]
  dta.index = pd.DatetimeIndex(dta.date, freq='QS')
  dta_hamilton = dta.rgnp
  
  # Plot the data
  dta_hamilton.plot(title='Growth rate of Real GNP', figsize=(12,3))
  
  # %%
  fig, ax = plt.subplots(figsize=(12, 4))
  #plt.plot(np.round(trace[&quot;hmm_states&quot;].mean(axis=0)), label=&quot;inferred&quot;)
  plt.plot(dta_hamilton.values, label=&quot;true&quot;)
  
  # %%
  def solve_equilibrium(n_states, p_transition):
      A = tt.dmatrix('A')
      A = tt.eye(n_states) - p_transition + tt.ones(shape=(n_states, n_states))
      p_equilibrium = pm.Deterministic(&quot;p_equilibrium&quot;, sla.solve(A.T, tt.ones(shape=(n_states))))
      return p_equilibrium
  
  class HMMStates(pm.Categorical):
      def __init__(self, p_transition, p_equilibrium, n_states, *args, **kwargs):
          &quot;&quot;&quot;You can ignore this section for the time being.&quot;&quot;&quot;
          super(pm.Categorical, self).__init__(*args, **kwargs)
          self.p_transition = p_transition
          self.p_equilibrium = p_equilibrium
          # This is needed
          self.k = n_states
          # This is only needed because discrete distributions must define a mode.
          self.mode = tt.cast(0,dtype='int64')
  
      def logp(self, x):
          &quot;&quot;&quot;Focus your attention here!&quot;&quot;&quot;
          p_eq = self.p_equilibrium
          # Broadcast out the transition probabilities,
          # so that we can broadcast the calculation
          # of log-likelihoods
          p_tr = self.p_transition[x[:-1]]
  
          # the logp of the initial state evaluated against the equilibrium probabilities
          initial_state_logp = pm.Categorical.dist(p_eq).logp(x[0])
  
          # the logp of the rest of the states.
          x_i = x[1:]
          ou_like = pm.Categorical.dist(p_tr).logp(x_i)
          transition_logp = tt.sum(ou_like)
          return initial_state_logp + transition_logp
  
  # %%
  class HamiltonEmissions(pm.Continuous):
      def __init__(self, states, phi, sigma, mu, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.states = states
          self.phi = phi
          self.mu = mu
          self.sigma = sigma  # variance
  
      def logp(self, x):
          &quot;&quot;&quot;
          x: observations
          &quot;&quot;&quot;
          states = self.states
          sigma = self.sigma[states]
          mu = self.mu[states]
          phi = self.phi
  
          z = x - mu  # Centered version of x
  
          ar_mean = \
              phi[0] * z[0:-4] + \
                  phi[1] * z[1:-3] + \
                      phi[2] * z[2:-2] + \
                          phi[3] * z[3:-1]
          
          ar_like = tt.sum(pm.Normal.dist(mu= ar_mean + mu[4:], sigma=sigma[4:]).logp(x[4:]))
  
          boundary_like = pm.Normal.dist(mu=0, sigma=sigma[:4]).logp(x[:4])
          return ar_like + boundary_like
  # %%
  n_states = 2
  with pm.Model() as model:
      # Priors for transition matrix
      p_transition = pm.Dirichlet(&quot;p_transition&quot;,
          a=tt.ones((n_states, n_states)),
          shape=(n_states, n_states))
  
      # Solve for the equilibrium state
      p_equilibrium = solve_equilibrium(n_states, p_transition)
  
      # HMM state
      hmm_states = HMMStates(
          &quot;hmm_states&quot;,
          p_transition=p_transition,
          p_equilibrium=p_equilibrium,
          n_states=n_states,
          shape=(len(dta_hamilton),)
      )
  
      # Prior for mu and sigma
      mu = pm.Normal(&quot;mu&quot;, mu=0, sigma=1, shape=(n_states,))
      sigma = pm.Exponential(&quot;sigma&quot;, lam=2, shape=(n_states,))
      phi = pm.Normal(&quot;phi&quot;, 0, 0.5, shape=(4, ))
  
      # Observed emission likelihood
      obs = HamiltonEmissions(
          &quot;emission&quot;,
          states=hmm_states,
          mu=mu,
          sigma=sigma,
          phi=phi,
          observed=dta_hamilton
      )
  # %%
  with model:
      start = pm.find_MAP()
      step1 = pm.Metropolis(vars=[mu, sigma, phi, p_transition, emission])
      step2 = pm.BinaryGibbsMetropolis(vars=[hmm_states])
      trace = pm.sample(2500, cores=1, chains=2, step=[step1, step2], tune=1500)
  
  # %%
  import arviz as az
  az.plot_trace(trace, var_names=[&quot;p_transition&quot;])
  </code></pre>