lindermanlab / ssm

Bayesian learning and inference for state space models
MIT License
540 stars 196 forks source link

Constraining AR-HMM emission parameters #158

Open JustFineNeuro opened 1 year ago

JustFineNeuro commented 1 year ago

Assuming an AR-HMM, I want to know if it's possible to constrain the AR model parameters. Specifically, is it possible to set the 'A' matrix to fixed parameters and only estimate parameters for the input matrix (labeled as 'V' in the ssm package)?

Thanks!

slinderman commented 1 year ago

Good question! I don’t think you can do that with the built in code. You could make a custom observation class and only include the updates for the other parameters, via a custom M step. It should end up looking pretty similar to the existing AR observations class. Does that make sense?

JustFineNeuro commented 1 year ago

Thanks for the rapid response! Yes, this makes sense. I see the AR class in https://github.com/lindermanlab/ssm/blob/master/ssm/observations.py. Can you point me to where the M step is handled, I assume by HMM.fit(). But I couldn't find it easily.

slinderman commented 1 year ago

First let me just make sure I understand the model you'd like to fit. I think you're describing an autoregressive observation model,

$$ p(xt \mid x{t-1}, z_t, u_t) = \mathcal{N}(xt \mid A x{t-1} + B_{z_t} ut, Q{z_t}) $$

where the free parameters are ${B_k, Qk}{k=1}^K$. In your case, the $A$ matrix is fixed (and shared by all discrete states). Is that right?

If so, I think you could equivalently define an HMM on the residuals $r_t = xt - A x{t-1}$, then have a simple HMM with "linear regression" observations,

$$ p(r_t \mid z_t, u_t) = \mathcal{N}(rt \mid B{z_t} ut, Q{z_t}) $$

This is also called a "switching linear regression model." I don't think we have such an observation model built into SSM, but we do have one in dynamax (see here)

JustFineNeuro commented 1 year ago

Yes, that should work well. Namely, I know the feedforward matrix transition matrix of the system dynamics "A". What I am trying to do is parse behavioral states by assuming they only differ by projection into a control subspace "B". Side question, then: in dynamax, are there options for computing the HMM with multiple emission types? For example, a gaussian for $p(r_t|z_t,u_t)$ and Poisson for firing rates?