softKmeans_stan.R

#library(parallel)
#THe model is from Stan User's Guide and Reference Manual
#https://github.com/stan-dev/stan/releases/download/v2.0.1/stan-reference-2.0.1.pdf

library(rstan)
model <- "
data {
       int<lower=0> N;  // number of data points
       int<lower=1> D;  // number of dimensions
       int<lower=1> K;  // number of clusters
       vector[D] y[N];  // observations
}

transformed data {
       real<upper=0> neg_log_K;
       neg_log_K <- -log(K);
}

parameters {
       vector[D] mu[K]; // cluster means
}

transformed parameters {
  real<upper=0> soft_z[N,K]; // log unnormalized clusters
  for (n in 1:N)
    for (k in 1:K)
      soft_z[n,k] <- neg_log_K - 0.5 * dot_self(mu[k] - y[n]);
}

model {
  // prior
  for (k in 1:K)
       mu[k] ~ normal(0,1);
  // likelihood
  for (n in 1:N)
       increment_log_prob(log_sum_exp(soft_z[n]));
}
"

d1<-cbind(rnorm(100,0,10),rnorm(100,0,5))
d2<-cbind(rnorm(50,25,5),rnorm(50,25,10))
#plot(rbind(c,d))
d <- rbind(d1,d2)

fit <- stan(model_code = model,
            pars=c("mu"),
            data = list(N=nrow(d),D=ncol(d),K=2,y=d),
           chains = 1, iter = 200)
traceplot(fit)
print(fit)