## Examples of Applied Time Series Analysis, By He Shuyuan
## Chpater 3, MA and ARMA


## example data of chemical reaction time series, two-hourly data

demo.chemical <- function(){
  x <- scan('data-b8.txt')
  x <- ts(x)
  oldpar <- par(mfrow=c(2,2)); on.exit(par(oldpar))
  plot(x, main="Chemical reaction time series")
  y <- diff(x)
  plot(y, main="First order difference")
  acf(y)
}
## demo.chemical()


## MA theoretical spectrum given MA coefficients
ma.true.spectrum <- function(a, ngrid=256, sigma=1,
                             tit="True MA Spectral Density",
                             plot.it=TRUE){
  p <- length(a)
  freqs <- seq(from=0, to=pi, length=ngrid)
  spec <- numeric(ngrid)
  for(ii in seq(ngrid)){
    spec[ii] <- 1 + sum(complex(mod=a, arg=freqs[ii]*seq(p)))
  }
  spec = sigma^2 / (2*pi) * abs(spec)^2
  if(plot.it){
    plot(freqs, spec, type='l',
         main=tit,
         xlab="frequency", ylab="spectrum",
         axes=FALSE)
    axis(2)
    axis(1, at=(0:6)/6*pi,
         labels=c(0, expression(pi/6),
           expression(pi/3), expression(pi/2),
           expression(2*pi/3), expression(5*pi/6), expression(pi)))
  }
  box()
  list(frequencies=freqs, spectrum=spec,
       ma.coefficients=a, sigma=sigma)
}

## simulate MA(q) model.
## a is vector a_1, \dots, a_q
## Model is X_t = eps_t + a_1 eps_{t-1} + \dots + a_q eps_{t-q}
## \Var(\epsilon_t) = sigma^2
## This version uses the filter function.
ma.gen <- function(n, a, sigma=1.0, by.roots=FALSE,
                   plot.it=FALSE){
  if(by.roots){
    require(polynom)
    cf <- Re(c(poly.calc(a)))
    cf <- cf / cf[1]
    a <- cf
  }
  q <- length(a)
  n2 <- n+q
  eps <- rnorm(n2, 0, sigma)
  x2 <- filter(eps, c(1,a), method="convolution", side=1)
  x <- x2[(q+1):n2]
  x <- ts(x)
  attr(x, "model") <- "MA"
  attr(x, "b") <- a
  attr(x, "sigma") <- sigma
  if(plot.it) plot(x)
  x
}


demo.ma1 <- function(){
  oldpar <- par(mfrow=c(2,1)); on.exit(par(oldpar))
  n <- 100
  x1 <- ma.gen(n, a=0.6)
  plot(x1, main="MA(1) Series: b=0.6")
  ma.true.spectrum(a=0.6, tit="MA(1) Spectral Density: b=0.6")
  locator(1)
  x1 <- ma.gen(n, a=-0.6)
  plot(x1, main="MA(1) Series: b=-0.6")
  ma.true.spectrum(a=-0.6, tit="MA(1) Spectral Density: b=-0.6")
  invisible()
}


demo.ma2 <- function(){
  oldpar <- par(mfrow=c(2,1)); on.exit(par(oldpar))
  n <- 120
  sigma <- 2
  b <- c(-0.36, 0.85)
  x <- ma.gen(n, sigma=2, a=b)
  plot(x, main="MA(2) Series: b=(-0.36,0.85)")
  ma.true.spectrum(a=b, tit="MA(2) Spectral Density: b=(-0.36,0.85)")

  gms <- numeric(3)
  gms[1] <- sigma^2 * (1 + b[1]^2 + b[2]^2)
  gms[2] <- sigma^2 * (b[1] + b[1]*b[2])
  gms[3] <- sigma^2 * b[2]
  names(gms) <- 0:2
  cat("======== Autocovariances ==========\n")
  print(gms)
  ks <- c(6, 12, 20, 30, 40, 50, 60, 100)
  nks <- length(ks)
  results <- matrix(0, nrow=nks, ncol=3)
  rownames(results) <- ks
  for(ii in seq(along=ks)){
    k <- ks[ii]
    res <- ma.solve(gms, k=k)
    results[ii,] <- c(res$b[1], res$b[2], res$s2)
  }
  cat("======== Solving MA coefficients ========\n")
  cat("True coefficients: ", b, "\n")
  cat("True white noise variance =", sigma^2, "\n")
  cat("Coefficients estimates using true autocovariance and k=6,12,...:\n")
  print(results)
  invisible()
}

## Given \gamma_0, \gamma_1, \dots, \gamma_q,
## Solve MA(q) coefficients b_1, \dots, b_q, \sigma^2
## Using Li Lei's algorithm.
## Input: gms -- \gamma_0, \gamma_1, \dots, \gamma_q
ma.solve <- function(gms, k=100){
  q <- length(gms)-1
  if(q==1){
    rho1 <- gms[2] / gms[1]
    b <- (1 - sqrt(1 - 4*rho1^2))/(2*rho1)
    s2 <- gms[1] / (1 + b^2)
    return(list(b=b, s2=s2))
  }
  
  A <- matrix(0, nrow=q, ncol=q)
  for(j in seq(2,q)){
    A[j-1,j] <- 1
  }
  cc <- numeric(q); cc[1] <- 1
  gamma0 <- gms[1]
  gammas <- numeric(q+k)
  gammas[1:(q+1)] <- gms
  gamq <- gms[-1]
  Gammak <- matrix(0, nrow=k, ncol=k)
  for(ii in seq(k)){
    for(jj in seq(k)){
      Gammak[ii,jj] <- gammas[abs(ii-jj)+1]
    }
  }
  
  Omk <- matrix(0, nrow=q, ncol=k)
  for(ii in seq(q)){
    for(jj in seq(k)){
      Omk[ii,jj] <- gammas[ii+jj-1+1]
    }
  }
  PI <- Omk %*% solve(Gammak, t(Omk))

  s2 <- gamma0 - c(t(cc) %*% PI %*% cc)
  b <- 1/s2 * c(gamq - A %*% PI %*% cc)
  return(list(b=b, s2=s2))
}


## ARMA theoretical spectrum given ARMA coefficients
arma.true.spectrum <- function(a, b, ngrid=256, sigma=1,
                               tit="True ARMA Spectral Density",
                               plot.it=TRUE){
  p <- length(a)
  q <- length(b)
  freqs <- seq(from=0, to=pi, length=ngrid)
  spec1 <- numeric(ngrid)
  spec2 <- numeric(ngrid)
  for(ii in seq(ngrid)){
    spec1[ii] <- 1 + sum(complex(mod=b, arg=freqs[ii]*seq(q)))
    spec2[ii] <- 1 - sum(complex(mod=a, arg=freqs[ii]*seq(p)))
  }
  spec = sigma^2 / (2*pi) * abs(spec1)^2 / abs(spec2)^2
  if(plot.it){
    plot(freqs, spec, type='l',
         main=tit,
         xlab="frequency", ylab="spectrum",
         axes=FALSE)
    axis(2)
    axis(1, at=(0:6)/6*pi,
         labels=c(0, expression(pi/6),
           expression(pi/3), expression(pi/2),
           expression(2*pi/3), expression(5*pi/6), expression(pi)))
  }
  box()
  list(frequencies=freqs, spectrum=spec,
       ar.coefficients=a, ma.coefficients=b,
       sigma=sigma)
}


## simulate ARMA(p, q) model.
## a is vector a_1, \dots, a_p
## b is vector b_1, \dots, b_q
## Model is
##   X_t = a_1 X_{t-1} + \dots + a_p X_{t-p}
##         + \epsilon_t + b_1 \epsilon_{t-1} + \dots + b_q \epsilon_{t-q}
## \Var(\epsilon_t) = sigma^2
arma.gen <- function(n, a, b, sigma=1.0,
                     by.roots.ar=FALSE, by.roots.ma=FALSE,
                     plot.it=FALSE, n0=1000,
                     x0=numeric(length(a))){
  n2 <- n0 + n
  p <- length(a)
  ## first generate n0+n MA(q) series
  eps <- ma.gen(n2, b, sigma=sigma,
                by.roots=by.roots.ma,
                plot.it=FALSE)
  b <- attr(eps, "b")
  if(by.roots.ar){
    require(polynom)
    cf <- Re(c(poly.calc(a)))
    cf <- cf / cf[1]
    a <- -cf[-1]
  }

  ##set.seed(1)
  x2 <- filter(eps, a, method="recursive", side=1, init=x0)
  x <- x2[(n0+1):n2]
  x <- ts(x)
  attr(x, "model") <- "ARMA"
  attr(x, "a") <- a
  attr(x, "b") <- b
  attr(x, "sigma") <- sigma
  if(plot.it) plot(x)
  x
}

## Wold coefficients for the ARMA model
arma.Wold <- function(n, a, b=numeric(0)){
  p <- length(a)
  q <- length(b)
  arev <- rev(a)
  psi <- numeric(n)
  psi[1] <- 1
  for(j in seq(n-1)){
    if(j <= q) bj=b[j]
    else bj=0
    psis <- psi[max(1, j+1-p):j]
    np <- length(psis)
    if(np < p) psis <- c(rep(0,p-np), psis)
    psi[j+1] <- bj + sum(arev * psis)
  }
  
  psi
}

## Calculate theoretical autocovariance function
## of ARMA model using Wold expansion
arma.gamma.by.Wold <- function(n, a, b=numeric(0), sigma=1){
  nn <- n + 100
  psi <- arma.Wold(nn, a, b)
  gam <- numeric(n)
  for(ii in seq(0, n-1)){
    gam[ii+1] <- sum(psi[1:(nn-ii)] * psi[(ii+1):nn])
  }
  gam <- (sigma^2) * gam
  gam
}
arma.gamma <- arma.gamma.by.Wold

## Solve ARMA parameters given autocovariance functions
arma.solve <- function(gms, p, q){
  Gpq <- matrix(0, nrow=p, ncol=p)
  for(ii in seq(p)) for(jj in seq(p)){
    Gpq[ii,jj] <- gms[abs(q + ii - jj)+1]
  }
  gs <- gms[(q+1+1):(q+p+1)]
  a <- solve(Gpq, gs)
  aa <- c(-1, a)

  gys <- numeric(q+1)
  for(k in seq(0, q)){
    gys[k+1] <- sum(c(outer(0:p,0:p,
                            function(ii,jj) aa[ii+1] * aa[jj+1]
                            * gms[abs(k+jj-ii)+1])))
  }

  res <- ma.solve(gys)
  b <- res$b
  sigma <- sqrt(res$s2)

  list(a=a, b=b, sigma=sigma)
}


demo.arma42 <- function(n=80){
  a <- c(-0.9, -1.4, -0.7, -0.6)
  b <- c(0.5, -0.4)

  x <- arma.gen(n, a, b, plot.it=F)

  oldpar <- par(mfrow=c(2,2)); on.exit(par(oldpar))
  ts.plot(x, main="ARMA(4,2) Series")
  acf(x)
  arma.true.spectrum(a, b)
  ng <- 21
  gams <- arma.gamma(ng, a, b, sigma=1)
  plot(0:(ng-1), gams, type="h",
       main="Theoretical gamma by Wold",
       xlab="k", ylab=expression(gamma[k]))
  abline(h=0)
  cat("\n====== Autocovariances ========\n")
  names(gams) <- 0:(ng-1)
  print(gams)

  cat("\n===== Solve ARMA from ACV ======\n")
  res <- arma.solve(gams, p=4, q=2)
  print(res)
#  psi <- arma.Wold(31, a, b)
#   ff <- seq(0, pi, length=200)
#   yy <- numeric(200)
#   for(ii in seq(200)){
#     fr <- ff[ii]
#     yy[ii] <- sum(complex(mod=psi, arg=(0:30)*fr))
#   }
#   yy <- Mod(yy)^2 / (2*pi)
#   plot(ff, yy, type="l", main="Wold spectrum")

  cat("\n\n====== Example 3.2.2 Solve ARMA parameters\n")
  gams2 <- c(4.61, -1.06, 0.29, 0.69, -0.12)
  res2 <- arma.solve(gams2, p=2, q=2)
  print(res2)
  
  invisible()
}

demo03 <- function(){
  demos <- list("化学反应时间序列" = demo.chemical,
                "MA(1)演示" = demo.ma1,
                "MA(2)演示" = demo.ma2,
                "ARMA(4,2)演示" = demo.arma42
                )
  ndemos <- length(demos)
  prompts <- paste("   ", seq(ndemos), ". ", names(demos), "\n", sep="")
  cat("第三章演示:\n",
      prompts)
  sel <- scan(n=1)
  demos[[sel]]()
}