kalman.R

### Libraries
library(pscl)
library(rjags)
library(rstan)
data(AustralianElectionPolling,package="pscl")
dat <- AustralianElectionPolling



dat$startDate <- as.Date(dat$startDate)
dat$endDate <- as.Date(dat$endDate)
### Start at 2004 elex.
orig.date <- as.Date("2004-10-09")
### End at 2007 elex.
end.date <- as.Date("2007-11-24")

dat$startDate.num <- julian(dat$startDate,origin=orig.date)
dat$endDate.num <- julian(dat$endDate,origin=orig.date)
dat$fieldDate.num <- floor((dat$startDate.num + dat$endDate.num) / 2)

dat$y <- dat$ALP / 100

## Jags will use precision
### Stan will use variance
dat$var <- dat$y*(1-dat$y)/dat$sampleSize
dat$sd <- sqrt(dat$var)
dat$prec <- 1 / dat$var

forJags <- list(y = dat$y,
	date = dat$fieldDate.num,
	nPolls = nrow(dat),
	kappa = 0.01,
	myvar = dat$var,
	mysd = dat$sd,
	prec = dat$prec,
	xi = c(0.376,rep(NA,end.date - orig.date - 2),0.434),
	nPeriods = as.numeric(end.date - orig.date))

jags_code <- '
	model {
		for (i in 1:nPolls) {
			mu[i] <- xi[date[i]]
			y[i] ~ dnorm(mu[i],prec[i])
		}
		for (t in 2:(nPeriods-1)) {
			xi[t] ~ dnorm(xi[t-1],tau)
		}

		omega ~ dunif(0,kappa)
		tau <- 1/pow(omega,2)
	}
'

writeLines(jags_code,con="kalman.bug")


system.time(jags.mod <- jags.model("kalman.bug",
	data = forJags))
update(jags.mod,n.iter = 10000)
system.time(out <- coda.samples(jags.mod,variable.names = c("xi"),
	n.iter = 10000,
	n.thin = 10))

holder <- summary(out)
plot(1:forJags$nPeriods,holder$statistics[,1],
	xlab = "Day",
	ylab = "ALP vote intention",type="l")

stan_code <- '
	data {
		int<lower=0> nPolls;
		int<lower=0> nPeriods; //
		int<lower=0> date[nPolls];
		real y[nPolls];
		real kappa;
		real mysd[nPolls];
	} 
	parameters {
		real<lower=0,upper=1> xi[nPeriods];
		real<lower=0,upper=.01>tau;
	}
	model {
		// observation model
		for (i in 1:nPolls) {
			y[i]  ~ normal(xi[date[i]],mysd[i]);
		}

		// transition model
		// First period: take from known value 
		xi[1] ~ normal(0.376,0.00001);
		for (t in 2:(nPeriods-1)) {
			xi[t] ~ normal(xi[t-1],tau);
		}
		xi[nPeriods] ~ normal(0.434,0.00001);
	}
'

## Stan will complain if we leave partial data in
forJags$xi <- NULL

## Init func with dramatic step change in latent support
my.initfunc <- function() {
	xi <- c(rep(0.376,floor(forJags$nPeriods/2)),
		rep(0.434,ceiling(forJags$nPeriods/2)))
	xi <- rep(0.5,forJags$nPeriods)
	tau <- runif(1,0,.01)
	list(xi = xi,tau=tau)
}

system.time(fit <- stan(model_code = stan_code,
	data = forJags,
	chains = 1,
	init = my.initfunc,
	iter = 1000,
	thin = 1))

samples <- extract(fit, c("xi"))
plot(1:forJags$nPeriods,colMeans(samples$xi),
	xlab = "Day",
	ylab = "ALP vote intention",type="l")