I would suggest you do everything in long form, so instead of having a matrix of dimensions (observation, patient), you just have a vector of length n_obs * n_patients with basically a multi-index. I’m not sure if you have a balanced panel or not, but I find it’s easier to index into long-form data.
I’m also leery of the fact that you’re expanding the a variable with patient_idx just to take the mean of it. That doesn’t seem correct.