How to estimate state specific parameters in PyMC3 model

Questions
1

I have the following bayesian regression

Capture
Capture696×168 3.65 KB

where I want to estimate the parameters beta which depend on the state s=[1, 2, 3] which are given by an array S. y_it is of shape 100x90 and X_it is of shape 100x90x20 and beta is of shape 20x1 while N=90 refers to the number of individuals, T=90 to the time increments and K=20 to the number of predictors.
For the priors I stick to weakly informative priors such a standard normal distributions.
In absence of the state variable, my model would like as follows:

with pm.Model() as model:
    alpha = pm.Normal(name="alpha", mu=0.0, sigma=1.0, shape=())
    beta = pm.Normal(name="beta", mu=0.0, sigma=1.0, shape=K)
    eps = pm.Gamma(name="eps", alpha=9.0, beta=4.0, shape=())
    y_hat = pm.Deterministic(name="y_hat", var=pm.math.dot(X, beta) + alpha)
    y_like = pm.Normal(name="y_like", mu=y_hat, sigma=eps, observed=y)

How can I accommodate the state variable in the model at hand?

Thank you very much!

2

Hi there,

Just to clarify: do I understand you correctly that the state s_{it} is known for each observation? Or do you need to estimate which state each observation belongs to, too?

3

Hey Martin,

thanks for your reply.
In first place the state is observed for every observation.

However, I would also be keen how the model looks like if one has to estimate the states, too.

4

Hi again,

I had a go at this. Here’s the code:

# Make up some random data
K = 20
N = 100
T = 90
S = 3

X = np.random.randn(N, T, K)
y = np.random.randn(N, T)

# Number of states
# States, zero-indexed
s = np.random.randint(0, S, size=(N, T))

with pm.Model() as model:
    
    # Note the added dimension for alpha and beta
    alpha = pm.Normal(name="alpha", mu=0.0, sigma=1.0, shape=(S))
    beta = pm.Normal(name="beta", mu=0.0, sigma=1.0, shape=(S, K))
    eps = pm.Gamma(name="eps", alpha=9.0, beta=4.0, shape=())
    
    # This picks out the relevant alphas for each observation
    # It should have shape (N, T)
    rel_alpha = alpha[s]
    
    # This picks out the relevant betas for each observation
    # It should have shape (N, T, K)
    rel_beta = beta[s]
        
    # The sum here computes the dot product over the last axis; maybe there is a nicer/faster way.
    y_hat = pm.Deterministic(name="y_hat", var=pm.math.sum(X * rel_beta, axis=2) + rel_alpha)
    
    # This is the same as before
    y_like = pm.Normal(name="y_like", mu=y_hat, sigma=eps, observed=y)

The trick is basically to use the matrix of states s to index alpha and beta and pull out the relevant coefficients. A slight complication is that we can’t seem to use pm.math.dot with this because there’s an extra axis now. I got around this using the pm.math.sum expression, but there’s probably a nicer way, maybe using tensordot. This should work though.

I’ve tried to comment each step; let me know if it makes sense! To work out the shapes, it helped me to play around with things in numpy first:

sample_beta = np.random.randn(S, K)
sample_alpha = np.random.randn(S)
sample_beta[s].shape, X.shape, sample_alpha[s].shape

This returns:

((100, 90, 20), (100, 90, 20), (100, 90))

If you want to estimate the states, too, things get a bit more complicated. Basically, you’ll want to fit something quite similar to a mixture model, so maybe looking at some tutorials for that (maybe e.g. here) could give you some initial pointers.

5

Works fine.
Thanks a lot!