I am tracking public opinion over time using a Dirichlet where the public opion in any month is informed by polling and opinion in the previous month.
My question is how do I vectorise this.
model = pm.Model()
with model:
# temporal model for monthly voting intention
vi = []
for i in range(P["n_months"]):
if not i:
a = P["start"]
else:
a = vi[i - 1]
vi.append(
pm.Dirichlet(f"vi_{i:02}", a=a * P["innovation"], shape=(P["n_parties"],))
)
vi_stack = pm.Deterministic("vi_stack", pm.math.stack(vi, axis=1))
print(vi_stack.ndim)
# likelihood / observation model
for i in range(P["n_polls"]):
P["y"][i] = pm.Dirichlet(
f"y_{i:03}",
a=vi_stack[:, P["month"][i]] * P["innovation"],
observed=P["y"][i],
)
Full notebook here: https://github.com/bpalmer4/Australian-Federal-Election-2025/blob/main/notebooks/_poll_agg_states_pymc.ipynb