Dirichlet error: updating 2 patch gaussian model to 4 patch

Hi. I am trying to write a gaussian mixture model to estimate values of parameters. I am able to write the code for 2 gaussian patches, but when I try to run the code with 4 gaussian patches I get the following error:

Unrecoverable error evaluating the log probability at the initial value.
Exception: dirichlet_lpmf: prior sample sizes has dimension = 2, expecting dimension = 4; a function was called with arguments of different scalar, array, vector, or matrix types, and they were not consistently sized; all arguments must be scalars or multidimensional values of the same shape. (in ‘unknown file name’ at line 52)

I know the error has to do with the dirichlet function and I’ve tried multiple different ways of fixing it. For example, if I put:

theta ~ dirichlet(rep_vector(4.0, K)

I get a similar error.
Here is my whole code:

from future import division
import pystan
from scipy import stats
import matplotlib.pyplot as plt
import theano.tensor as tt
import numpy as np
import os
import pandas as pd
x=np.array([366.43,368.44,370.06,366.43,368.43,370.05,367.59,358.35,352.57])
y=np.array([3169.77,3175.17,3176.19,3169.78,3175.17,3176.19,3178.73,3177.84,3175.35])

x=np.array(zip(x,y))
c=np.array([2.27,0.92,0.44,4.32,0.59,3.45,52.11,6.66,1.68])
s = np.array([366.61,3179.73])/100
st=np.ones(9)
pt = np.ones(3)

model=’’’
data {
int D; //number of dimensions
int K; //number of gaussians
int N; //number of data
int M; // prediction day
vector[D] y[N]; // observation data
real con[N]; //concentration
vector[D] s;//oil spill location
real st[N]; // sample time
real pt[M]; // prediction time
}

parameters {
simplex[K] theta; //mixing proportions
vector[D] v[K];
vector<lower=0>[D] Dif[K];
cholesky_factor_corr[D] L[K]; //cholesky factor of correlation matrix
}

transformed parameters {
cholesky_factor_cov[D,D] cov[K,N];
vector<lower=0>[D] sigma[K,N]; // standard deviations
vector[D] mu[K,N];
real ro[K];
matrix[D,D] Omega[K];
matrix[D,D] Sigma[K,N];
vector[N] lamba;

for (k in 1:K) {
Omega[k] = multiply_lower_tri_self_transpose(L[k]);
for (n in 1:N){
sigma[k,n] = 0.05 + sqrt(2*st[n]*Dif[k]);
mu[k,n] = s+v[k]*st[n];
cov[k,n] = diag_pre_multiply(sigma[k,n],L[k]);
Sigma[k,n] = quad_form_diag(Omega[k], sigma[k,n]);
}
ro[k]=Omega[k,2,1];
}

for (i in 1 : N) {lamba[i] = 1/(theta[1](1./2./3.1415926/sqrt (Sigma[1,i, 1, 1])/sqrt (Sigma[1,i, 2, 2])/sqrt (1 - ro[1]ro[1]))exp (-1./2./(1 - ro[1]ro[1])(-(y[i, 1] - mu[1,i, 1])(y[i, 1] - mu[1,i, 1])/Sigma[1, i,1, 1] - (y[i, 2] - mu[1, i,2])(y[i, 2] - mu[1, i,2])/Sigma[1,i, 2, 2] + 2.ro[1](y[i, 1] - mu[1,i, 1])(y[i, 2] - mu[1,i, 2])/sqrt (Sigma[1, i,1, 1])/sqrt (Sigma[1,i, 2, 2]))) +
theta[2](1./2./3.1415926/sqrt (Sigma[2, i,1, 1])/sqrt (Sigma[2,i, 2, 2])/sqrt (1 - ro[2]ro[2]))exp (-1./2./(1 - ro[2]ro[2])(-(y[i, 1] - mu[2, i,1])(y[i, 1] - mu[2, i,1])/Sigma[2, i,1, 1] - (y[i, 2] - mu[2,i, 2])(y[i, 2] - mu[2, i,2])/Sigma[2,i, 2, 2] + 2.ro[2](y[i, 1] - mu[2, i,1])(y[i, 2] - mu[2, i,2])/sqrt (Sigma[2, i,1, 1])/sqrt (Sigma[2, i,2, 2]))) +
theta[3](1./2./3.1415926/sqrt (Sigma[3, i,1, 1])/sqrt (Sigma[3,i, 2, 2])/sqrt (1 - ro[3]ro[3]))exp (-1./2./(1 - ro[3]ro[3])(-(y[i, 1] - mu[3, i,1])(y[i, 1] - mu[3, i,1])/Sigma[3, i,1, 1] - (y[i, 2] - mu[3,i, 2])(y[i, 2] - mu[3, i,2])/Sigma[3,i, 2, 2] + 2.ro[3](y[i, 1] - mu[3, i,1])(y[i, 2] - mu[3, i,2])/sqrt (Sigma[3, i,1, 1])/sqrt (Sigma[3, i,2, 2]))) +
theta[4](1./2./3.1415926/sqrt (Sigma[4, i,1, 1])/sqrt (Sigma[4,i, 2, 2])/sqrt (1 - ro[4]ro[4]))exp (-1./2./(1 - ro[4]ro[4])(-(y[i, 1] - mu[4, i,1])(y[i, 1] - mu[4, i,1])/Sigma[4, i,1, 1] - (y[i, 2] - mu[4,i, 2])(y[i, 2] - mu[4, i,2])/Sigma[4,i, 2, 2] + 2.ro[4](y[i, 1] - mu[4, i,1])(y[i, 2] - mu[4, i,2])/sqrt (Sigma[4, i,1, 1])/sqrt (Sigma[4, i,2, 2]))));}
}

model {
real ps[K];
theta ~ dirichlet(rep_vector(2.0, K));
for(k in 1:K){
v[k,1] ~ normal(0.0,4.1);// uniform(340/100,380/100);//
v[k,2]~ normal(0.0,4.1);//uniform(3160/100,3190/100);//
Dif[k] ~ normal(0.5,0.2);//exponential(0.05);//beta(2,5);//uniform(0.1,0.89); //////normal(1,2);
L[k] ~ lkj_corr_cholesky(4);// contain rho
con ~ exponential(lamba);
}

for (n in 1:N){
for (k in 1:K){
ps[k] = log(theta[k])+multi_normal_cholesky_lpdf(y[n] | mu[k,N], cov[k,N]); //increment log probability of the gaussian
}
target += log_sum_exp(ps);
}
for(i in 1:N){
target += - lamba[i]*con[i]+log(lamba[i]);
}
}
}
‘’’
dat={‘D’:2, ‘K’:4, ‘N’:9, ‘y’:x, ‘con’:c, ‘s’:s, ‘st’:st, ‘pt’:pt, ‘M’:3}
fit = pystan.stan(model_code=model,data=dat,iter=10000,warmup=500, chains=1,init_r=0.5)
print(fit)

If anyone has any suggestions, let me know!