myFKF: Fast Kalman Filtering and Smoothing

#' Generate parameters for various univariate state-space models.
#'
#' @param ... Named parameters for univariate state space models. These can be
#'   any of the following (defaults given)
#'
#'   \code{sig2epsilon = 1, sig2xi = 1, sig2eta = 1, sig2omega = 1, sig2cycle =
#'   1,d = 2, r = 2, constSlope = 0, nu = FALSE, phi1 = 1, phi2 = 1, period =
#'   12, damp = 1, ccycle = double(2), obsConst = 0, timepoints = 1:100)}
#'
#'   Generally, these are set by other functions (\code{\link{trendModelArb},
#'   \link{unstrucSeasonal}, \link{trigSeasonal}, \link{tvarCycle},
#'   \link{constCycle}, or \link{generateSSmodel}}.
#'
#' @return A list of parameters for use in generating certain kinds of
#'   univariate state space models.
#' @export
parmList <- function(...){
  if(is.vector(data)) data = matrix(data, nrow=1)
  chg = list(...)
  defaults = list(sig2epsilon = 1, sig2xi = 1, sig2eta = 1, sig2omega = 1, sig2cycle = 1,
                  d = 2, r = 2, constSlope = 0, nu = FALSE, phi1 = 1, phi2 = 1,
                  period = 12, damp = 1, ccycle = double(2),
                  obsConst = 0, timepoints = 1:100)
  m = match(names(chg), names(defaults))
  keep = !is.na(m)
  m = m[keep]
  defaults[m] = chg[keep]
  return(defaults)
}

#' Univariate trend components
#'
#' This function can generate the appropriate state and transition matrices
#' (\eqn{Tt}, \eqn{Zt}, \eqn{GGt}, \eqn{HHt}) for different possible trend
#' components. Possible univariate trends are given with: \deqn{y(t) = mu(t) +
#' extra stuff} \deqn{(1 - phi1 L)^d mu(t) = nu(t-1) + xi(t-1)} \deqn{(1 - phi2
#' L)^r nu(t) = eta(t)} Noise terms \eqn{eta} and \eqn{xi} are iid normal mean 0
#' with variances \code{sig2eta} etc. This function generates all possible
#' combinations of these. See Chapter 3 in (Durbin and Koopman) for more
#' information. "Extra stuff" can be simply an error variance or may include
#' seasonal, cyclical, or other components.
#'
#' @param parmlist
#'
#'   The input contains named elements \code{d, phi1, r, phi2, sig2eta, sig2xi}.
#'   All are scalars parameterizing the above model. It also contains a logical
#'   for \code{nu}, which determines whether to include the \eqn{nu} component
#'   (default is \code{FALSE}). All of these can be set with
#'   \code{\link{parmList}}.
#'
#'   The default values (generated by \code{\link{parmList}}) are
#'
#'   \code{d} = 2 \cr \code{phi1} = 1\cr \code{r} = 2\cr \code{phi2} = 1\cr
#'   \code{sig2eta} = 1\cr \code{sig2xi} = 1\cr \code{nu} = \code{FALSE}
#'
#' @return A list of matrices with components \code{Z}, \code{Tt}, \code{Q}, and
#'   \code{R}.
#' @export
trendModelArb <- function(parmlist){
  with(parmlist, {
    ## First build mu contribution
    muMats = buildContrib(d, phi1)
    ## Now build nu contribution
    if(!nu){
      Rmu = muMats$Rt
      Tmu = muMats$Tt
      Qmu = matrix(sig2xi)
    }else{
      nuMats = buildContrib(r, phi2)
      Tmu = as.matrix(Matrix::bdiag(muMats$Tt, nuMats$Tt))
      Tmu[1,ncol(muMats$Tt)+1] = 1
      Qmu = diag(c(sig2xi, sig2eta))
      Rmu = rbind(cbind(muMats$Rt, 0), cbind(0, nuMats$Rt))
    }
    Zmu = matrix(0,1,nrow(Tmu))
    Zmu[1] = 1
    return(list(Z = Zmu, Tt = Tmu, Q=Qmu, R=Rmu))
  }
  )
}



#' Specific trend models
#'
#' This is simply a wrapper for \code{\link{trendModelArb}} which produces
#' specific trend models discussed in Chapter 3 of Durbin and Koopman.
#'
#' @param trendType The name of a specific trend type. Options are
#'   \code{LocalLevel}, \code{LocalLinTrend}, \code{SmoothSlope},
#'   \code{ConstantSlope}, and \code{IntRandWalk}. Input should be a quoted
#'   string. This updates the parmlist to reflect the specific trend choice. Any
#'   other string simply applies \code{\link{trendModelArb}} to the parmlist.
#'
#' @details The named trend models have the following forms:
#'
#' 1. LocalLevel: \deqn{y_t = mu_t + more} \deqn{mu_t+1 = mu_t + xi_t} 2.
#' LocalLinTrend: \deqn{y_t = mu_t + more} \deqn{mu_t+1 = mu_t + nu_t + xi_t}
#' \deqn{nu_t+1 = nu_t + eta_t} 3. SmoothSlope: \deqn{y_t = mu_t + more}
#' \deqn{mu_t+1 = mu_t + nu_t} \deqn{nu_t+1 = nu_t + eta_t} 4. ConstantSlope has
#' trend = constSlope, gives linear trend \deqn{y_t = mu_t + more} \deqn{mu_t+1
#' = constSlope + mu_t} 5. IRW(r), r > 1,default is r=2 \deqn{y_t = mu_t + more}
#' \deqn{mu_t+1 = 2mu_t - mu_t-1 + xi_t}
#'
#' @param parmlist A list of parameters as generated by \code{\link{parmList}}.
#'
#' @return A list of matrices with components \code{Z}, \code{Tt}, \code{Q}, and
#'   \code{R}.
#'
#' @references Durbin, James, and Koopman, Siem Jan. \emph{Time Series Analysis
#' by State Space Methods} (2nd Edition). Cambridge, GBR: OUP Oxford, 2012.
#' @export
namedTrends <- function(trendType, parmlist){
  pars = within(parmlist, {
    switch(trendType,
           'LocalLevel' = {
             nu <- FALSE
             d <- 1; phi1 <- 1
             },
           'LocalLinTrend' = {
             nu <- TRUE
             d <- 1; phi1 <- 1; r <- 1; phi2 <- 1
             },
           'SmoothSlope' = {
             nu <- TRUE
             d <- 1; phi1 <- 1; r <- 1; phi2 <- 1; sig2xi <- 0
             },
           'ConstantSlope' = {
             nu <- FALSE
             d <- 1; phi1 <- 1; sig2xi <- 0
             },
           'IntRandWalk' = {
             nu <- FALSE
             d <- ifelse(d>1, d, 2)
             },
           ## any other name gives arbitrary model
           TRUE
           )
  }
  )
  return(trendModelArb(pars))
}

buildContrib <- function(d, phi){
  if(d == 1){
    Tt = phi * matrix(1)
    Rt = matrix(1)
  }else{
    Tt.lower = cbind(diag(d-1),0)
    Tt = rbind(lagCoef(d) * phi^(1:d), Tt.lower)
    Rt = matrix(0, ncol=1, nrow=dim(Tt)[1])
    Rt[1] = 1
  }
  return(list(Tt=Tt, Rt = Rt))
}



#' Unstructured seasonal component
#'
#' @param parmlist A list of parameters as generated by \code{\link{parmList}}.
#'
#' @details This function generates state space model matrices related to an an
#'   unstructured seasonal component. The model has the form: \deqn{y(t) = Z_gam
#'   * gam(t) + extra stuff} \deqn{gam(t+1) = Tt_gam * gam(t) + R * sig_omega}
#'
#'   For a particular period (\code{parmlist$period}), this function produces
#'   the appropriate model matrices above. Note that the size of the state
#'   vector (related to the seasonal) is equal to the period. This can get
#'   large.
#'
#' @return A list of matrices with components \code{Z}, \code{Tt}, \code{Q}, and
#'   \code{R}.
#' @export
unstrucSeasonal <- function(parmlist){
  sig2omega = parmlist$sig2omega
  s = parmlist$period
  Tgam1 = rep(-1, s-1)
  Tgam = rbind(Tgam1, cbind(Matrix::Diagonal(s-2),0))
  if(s > 2){
    zeros = rep(0, s-2)
  } else {
    zeros = NULL
  }
  Zgam = matrix(c(1, zeros), nrow=1)
  Qgam = matrix(sig2omega)
  Rgam = matrix(0, nrow=dim(Tgam)[1], ncol=1)
  Rgam[1] = 1
  return(list(Z = Zgam, Tt = Tgam, Q = Qgam, R = Rgam))
}

#' Trigonometric seasonal component
#'
#' @param parmlist A list of parameters as generated by \code{\link{parmList}}.
#'
#' @details This function generates state space model matrices related to an a
#'   trigonometric seasonal component. The model has the form: \deqn{y(t) =
#'   Z_gam * gam(t) + extra stuff} \deqn{gam(t+1) = Tt_gam * gam(t) + R *
#'   sig_omega}
#'
#'   For a particular period (\code{parmlist$period}), this function produces
#'   the appropriate model matrices above. Note that the size of the state
#'   vector (related to the seasonal) is equal to the period. This can get
#'   large.
#'
#' @return A list of matrices with components \code{Z}, \code{Tt}, \code{Q}, and
#'   \code{R}.
#' @export
trigSeasonal <- function(parmlist){
  s = parmlist$period
  sig2omega = parmlist$sig2omega
  sstar = floor(s/2)
  Rgam = diag(s-1)
  Qgam = sig2omega * Rgam
  Zgam = matrix(rep(c(1,0), length = s-1), nrow=1)
  Cjs = lapply(1:sstar, Cj, s=s)
  if(s %% 2 == 0) Cjs[[sstar+1]] = matrix(-1)
  Tgam = as.matrix(Matrix::bdiag(Cjs))
  return(list(Z = Zgam, Tt = Tgam, Q = Qgam, R = Rgam))
}

Cj <- function(j, s){
  lam = 2*pi*j/s
  Cj = matrix(c(cos(lam), -sin(lam), sin(lam), cos(lam)), 2, 2)
  return(Cj)
}


#' Time-varying cycle component
#'
#' @param parmlist A list of parameters as generated by \code{\link{parmList}}.
#'
#' @details This function generates state space model matrices related to an a
#'   time-varying cyclic component. Such a specification is usually more slowly
#'   varying than a seasonal component, say for a yearly cycle with daily data.
#'   The model has the form: \deqn{y(t) = Z_c * c(t) + extra stuff} \deqn{c(t+1)
#'   = Tt_c * c(t) + R * sig_cycle}
#'
#'   In particular, the vector c(t) evolves as \deqn{c_1(t+1) = c_1(t) *
#'   cos(lam) + c_2(t) * sin(lam) + omega_1} \deqn{c_2(t+1) = -c_1(t) * sin(lam)
#'   + c_2(t) * cos(lam) + omega_2} where \eqn{omega_1} and \eqn{omega_2} are
#'   iid normal with mean 0 and variance sig2cycle, and \eqn{lam = 2\pi/period}.
#'
#'   For a particular period (\code{parmlist$period}), this function produces
#'   the appropriate model matrices above.
#'
#' @return A list of matrices with components \code{Z}, \code{Tt}, \code{Q}, and
#'   \code{R}.
#' @export
tvarCycle <- function(parmlist){
  sig2cycle = parmlist$sig2cycle
  damp = parmlist$damp
  Zc = matrix(c(1,0), 1)
  Rc = diag(2)
  Qc = sig2cycle * diag(2)
  Tc = damp*Cj(1, parmlist$period)
  return(list(Z=Zc, Tt=Tc, Q=Qc, R=Rc))
}

#' Constant cycle component
#'
#' @param parmlist A list of parameters as generated by \code{\link{parmList}}.
#'
#' @details This function generates state space model matrices related to an a
#'   constant cyclic component. Such a specification is usually more slowly
#'   varying than a seasonal component, say for a yearly cycle with daily data.
#'   The model has the form: \deqn{y(t) = c(t) + extra stuff} where \deqn{c(t) =
#'   c_1 * cos(lam*t) + c_2 * sin(lam*t)} with \eqn{lam=2\pi/period} and \eqn{t}
#'   simply being the times at which the observations occurred.
#'
#'   Note that this really is no more than a time varying intercept in the
#'   observation equation. For a particular period (\code{parmlist$period}),
#'   this function produces the appropriate model matrices above.
#'
#' @return A vector \code{ct} of the same length as the data \code{yt}.
#' @export
constCycle <- function(parmlist){
  cc = parmlist$ccycle
  period = parmlist$period
  tt = parmlist$timepoints
  ct = cc %*% rbind(cos(2*pi*tt/period), sin(2*pi*tt/period))
  return(ct)
}

#' Generate (univariate) state space model matrices.
#'
#' This function generates all necessary matrices for various types of
#' univariate structural state space models. Not particularly useful on its own,
#' it serves as input for the \code{\link{fkf}} function.
#'
#' @param trendModel The name of a specific trend type. Options are
#'   \code{LocalLevel}, \code{LocalLinTrend}, \code{SmoothSlope},
#'   \code{ConstantSlope}, and \code{IntRandWalk}. Input should be a quoted
#'   string. This updates the parmlist to reflect the specific trend choice. Any
#'   other string simply applies \code{\link{trendModelArb}} to the parmlist.
#' @param seasonalModel The name of a seasonal type, either \code{trigSeasonal}
#'   or \code{unstrucSeasonal}. The default is \code{'none'}.
#' @param cycleModel The name of a cycle type, either \code{tvarCycle} or
#'   \code{constCycle}. The default is \code{'none'}.
#' @param parmlist A list of parameters. Likely the output of the function
#'   \code{\link{parmList}}. These will be used to parameterize the model
#'   matrices.
#'
#'
#' @return A list of components
#'
#'   \code{dt}	A matrix giving the intercept of the transition equation.\cr
#'   \code{ct}	A matrix giving the intercept of the measurement equation.\cr
#'   \code{Tt}	An array giving the factor of the transition equation.\cr
#'   \code{Zt}	An array giving the factor of the measurement equation.\cr
#'   \code{HHt}	An array giving the variance of the innovations of the
#'   transition equation.\cr \code{GGt}	An array giving the variance of the
#'   disturbances of the measurement equation.\cr \code{parmlist} The parameters
#'   used to generate the model (may modify the input).
#' @examples ssMod = generateSSmodel('LocalLinTrend',cycleModel = 'tvarCycle',
#'    parmlist=parmList(sig2xi=.5,sig2eta=.1))
#' simOut = simSS(250, ssMod)
#' plot(1:250,simOut$y,ty='l',lwd=2)
#' ## Plot the trend
#' lines(1:250,c(1,1)%*%simOut$a[1:2,],col=2)
#' lines(1:250,c(1,1)%*%simOut$a[3:4,]+ #the cycle added to the trend
#'    c(1,1)%*%simOut$a[1:2,],col=4)
#' @export
generateSSmodel <- function(trendModel, seasonalModel='none', cycleModel='none',
                            parmlist = parmList()){
  MuMats = namedTrends(trendModel, parmlist)
  GamMats = switch(seasonalModel,
                   trigSeasonal = trigSeasonal(parmlist),
                   unstrucSeasonal = unstrucSeasonal(parmlist),
                   list(Z = NULL, Tt = NULL, Q = NULL, R = NULL)
  )
  cMats = switch(cycleModel,
                 tvarCycle = tvarCycle(parmlist),
                 list(NULL, Z = NULL, Tt = NULL, Q = NULL, R = NULL)
  )
  Zt = cbind(MuMats$Z, GamMats$Z, cMats$Z)
  Tt = list(MuMats$Tt, GamMats$Tt, cMats$Tt)
  Qt = list(MuMats$Q, GamMats$Q, cMats$Q)
  Rt = list(MuMats$R, GamMats$R, cMats$R)
  Tt = as.matrix(Matrix::.bdiag(Tt[!sapply(Tt,is.null)]))
  Qt = as.matrix(Matrix::.bdiag(Qt[!sapply(Qt,is.null)]))
  Rt = as.matrix(Matrix::.bdiag(Rt[!sapply(Rt,is.null)]))
  if(cycleModel=='constCycle'){
    ct = constCycle(parmlist)
  } else {
    ct = matrix(0, nrow = 1)
  }
  dt = matrix(0, nrow = nrow(Tt))
  dt[1] = parmlist$constSlope
  GGt = matrix(parmlist$sig2epsilon)
  HHt = Rt %*% Qt %*% t(Rt)
  return(list(dt=dt,ct=ct+parmlist$obsConst,Tt=Tt,Zt=Zt,HHt=HHt,GGt = GGt,parms=parmlist))
}



lagCoef <- function(d){
  x = 0:d
  x = -choose(d, x) * (-1)^(x)
  return(x[-1])
}