#########################################################################
## Particle filtering and smoothing                                   ###
#########################################################################

##### Example of Andrade-Netto #######################
##### (nonlinear AR(1), observation =squared states + noise)

##### Functions to be loaded ###################################

f.andrade <- function(x,u,i)
{
  ## Purpose: simulates one state transition   
    ##          x = vector of starting values
    ##          u = vector of innovations
    ##          i = time index
  0.5*x + 25*x/(1+x^2) + 8*cos(1.2*i) + sqrt(10)*u
} 

g.square <- function(x,y) {dnorm(y-0.05*x^2)}
## computes the likelihood for observation y and vector x of state values

bal.sample <- function(size=length(prob),prob)
{
  ## Purpose: generating a balanced sample from (1,2,..,length(prob))
  ## ----------------------------------------------------------------------
  ## Arguments: size=sample size, prob=vector of probabilites
  ## Output: Vector of multiplicites and indices of resampled particles
  ## ----------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date:  9 Jun 2005, 14:22
  w <- floor(size*cumsum(prob) + runif(1))
  w <- w-c(0,w[-length(prob)])
  list(mult=w,index=rep(1:length(prob),w))
}


f.part.smooth <- function(x0,y, f.state, g, ...)
{
  ## Purpose: Computing the Kitagawa particle smoother (extending the state)
  ## ----------------------------------------------------------------------
  ## Arguments: x0=sample from the starting density
  ##            y=observed values
  ##            f.state=function describing the evolution of states
  ##                    x_t = f(x_{t-1},u_t,t) with u_t standard normal
  ##            g=observation density (time homogeneous)
  ## Output: x[t,s,] = ensemble for state at time s-1 based on y_{1:t-1}
  ##         (array of dimension c(n+1,n+1,n.rep)
  ## ----------------------------------------------------------------------
  ## Author: H. R. Kuensch, Date: 13 Mar 2005, 22:11
  n.rep <- length(x0)
  n <- length(y)
    x <- array(NA,dim=c(n+1,n+1,n.rep))
         # 
  x[1,1,] <- x0
  for (i in (1:n)){
      x.p <- f.state(x[i,i,],rnorm(n.rep),i, ...) #propagation
      w <- g(x.p,y[i], ...)
      w <- w/sum(w) #reweighting
      index <- bal.sample(n.rep,w)$index
      x[i+1,1:i,] <- x[i,1:i,index]
      x[i+1,i+1,] <- x.p[index] #resampling
  }
  x
}



###########  End functions ########################################

###########  Simulation of the model and plotting ##################
n.t <- 100 #number of time points 
x <- rnorm(n.t+1)
x[1] <- 5*x[1]
for (i in (1:n.t)) x[i+1] <- f.andrade(x[i],x[i+1],i)
y <- 0.05*x[-1]^2 + rnorm(n.t)
plot(x[-1],type="l",xlab="time",ylab="state/observations",main=
      "Simulated states and observations of Andrade-Netto")
lines(y,col=2)



##### Filter and smoothing intervals for all observation times ######
## The particle approximation of the full smoothing distribution (obtained
## by enlarging the state) is computed and the 80% credible intervals
## at time t based on y_{1:t} (filter) and on y_{1:t+L} (smoother) are
## plotted for all time points.

n.t <- length(y)
n.rep <- 200 #number of particles
x0 <- 5*rnorm(n.rep) #starting values

x.path <- f.part.smooth(x0,y,f.state=f.andrade,g=g.square) #full smoother

par(mfrow=c(1,2)) # computing and plotting quantiles
x.filtq <- matrix(0,nrow=3,ncol=n.t+1)
for (i in (1:(n.t+1))) x.filtq[,i] <- quantile(x.path[i,i,],
                                               probs=c(0.1,0.5,0.9))

plot(0:n.t,x,type="l",xlab="time",ylab="state",lwd=3,
     main="True state, filter median and 80% filter interval")
polygon(c(0:n.t,n.t:0),c(x.filtq[1,],rev(x.filtq[3,])),col="pink",border=NA)
lines(0:n.t,x,type="l",lwd=3)
lines(0:n.t,x.filtq[2,],col=2,lwd=2)    

L <- 2
x.smoothq <- matrix(0,nrow=3,ncol=n.t+1-L)
for (i in (1:(n.t+1-L))) x.smoothq[,i] <- quantile(x.path[i+L,i,],
                                               probs=c(0.1,0.5,0.9))

plot(0:n.t,x,type="l",xlab="time",ylab="state",lwd=3,
     main="True state, smoother median and 80% smoothing interval")
polygon(c(0:(n.t-L),(n.t-L):0),c(x.smoothq[1,],rev(x.smoothq[3,])),
        col="pink",border=NA)
lines(0:n.t,x,type="l",lwd=3)
lines(0:(n.t-L),x.smoothq[2,],col=2,lwd=2)