Does anyone know how I could fit a spline in a multilevel model? In other words, how can I define a separate spline basis for each group?
Here is what I have tried. The list of
B matrices is
B_list and the group of each data point is in the
"group" column of the data frame
d. The problem is that I cannot figure out a way to successfully index a collection of the spline bases. They are different sizes (i.e. different number of data points per group) so I cannot stack them into a 3D array.
group_idx = d["group"].cat.codes.to_list()
w_shape = (NUM_KNOTS + 2, len(np.unique(d["group"])))
with pm.Model() as m:
a = pm.Normal("a", 0, 2)
w = pm.Normal("w", 0, 1, shape=w_shape)
# Linear model
mu = pm.Deterministic("mu", a + pm.math.dot(B_list[group_idx], w[:, group_idx]))
sgima = pm.Exponential("sigma", 1)
D = pm.Normal("D", mu, sigma, observed=d["observed_D"])
Any help is greatly appreciated!
(Edit: To state what was not working.)