# How to estimate state specific parameters in PyMC3 model

I have the following bayesian regression

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!

1 Like

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?

1 Like

Hey Martin,

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.

1 Like

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.

4 Likes

Works fine.
Thanks a lot!

2 Likes