In Junzheshao5959/ssmimputedemo: SSM Imputation DEMO(Title Case)

knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
library(stringr)

plot.KFS=function(KFS_out,selection=NULL,range=NULL){
  if(class(KFS_out)!="KFS"){stop("The result is not a KFS object.")}
  if(is.null(selection)){selection=1:ncol(KFS_out$att)}
  if(is.null(range)){range=1:nrow(KFS_out$att)}
  if(!all(selection %in% 1:ncol(KFS_out$att))){stop("The selected column overflow the #variables")}
  for(id in selection){
    plot.KFS_each(KFS_out,id,range=range)
  }
}
plot.KFS_each=function(KFS_out,k,range, print = F){
  sd_temp=sqrt(KFS_out$Ptt[k,k,])
  fit = KFS_out$att[range,k]
  upper = KFS_out$att[range,k]+1.65*sd_temp[range]
  lower = KFS_out$att[range,k]-1.65*sd_temp[range]
  min_y = min(fit,upper,lower)
  max_y = max(fit,upper,lower)
  x = seq(1,length(fit),1)
  plot(upper,ylim = range(min_y-0.1*abs(min_y),max_y+0.1*abs(max_y)),ylab=colnames(KFS_out$model$T)[k],type="l", xlab = "Time", col = "grey")
  lines(upper,col="grey")
  polygon(c(x, rev(x)), c(lower, rev(upper)),col=gray(0.9), border = NA)
  lines(fit,col=gray(0.4),lwd = 1.5)
  abline(h = 0, col=gray(0.3), lwd=1, lty=3)
  if(print){
    text(length(fit)*0.73, max_y-0.3*abs(max_y),expression(hat(beta)[t==T[0]]))
    text(length(fit)*0.8, max_y-0.3*abs(max_y),"=")
    text(length(fit)*0.9, max_y-0.3*abs(max_y),round(fit[length(fit)],3), cex = 0.9)
    arrows(x0 = length(fit)*0.8,
           y0 = max_y-0.4*abs(max_y),
           x1 = length(fit),
           y1 = round(fit[length(fit)],3), lty = 1, length = 0.07, lwd = 1.4, col = "#2B4593")
  }
  cat(colnames(KFS_out$att)[k],": ",KFS_out$att[max(range),k]," (",KFS_out$att[max(range),k]-1.65*sd_temp[max(range)],",",KFS_out$att[max(range),k]+1.65*sd_temp[max(range)],").\n",sep="")
}

library(KFAS)
library(ssmimputedemo)
head(data_stationary)
# function for exploration
run_exploration =function(formula = "y ~ x + c", data.input = data_stationary, method.input = "BFGS"){
  resp = str_split(formula, " ~ ")[[1]][1]
  #print(resp)
  f = str_split(formula, " ~ ")[[1]][2]
  #print(f)
  n_variables = length(str_split(f, " \\+ ")[[1]])
  print(n_variables)
  col_list = colnames(data.input)
  #print(col_list)
  var_list = c(resp,str_split(f, " \\+ ")[[1]])
  #print(var_list)
  #print(mean(var_list %in% col_list))
  if (mean(var_list %in% col_list)!=1){
    stop("variables in the formula are not found in the input data")
  }else{
    model2_c_test=SSModel(data.input[,c(resp)] ~ -1 + 
                        SSMregression(as.formula(paste("~ ",f, sep = " ")),data=data.input,Q=diag(rep(NA,n_variables))),
                        data.input,H=NA)
    model2_c_fit=fitSSM(model2_c_test, inits =rep(0,n_variables+1), method = method.input) # n_variables + one NA in H
    model2_c_out=KFS(model2_c_fit$model)
    plot.KFS(model2_c_out,range=30:708)
  }

  return(model2_c_out)
}
run_exploration(formula = "y ~ x + c", data.input = data_stationary)

n_variables=2
model2_c_test=SSModel(y ~ -1 + SSMregression(~ x + c,
                                                           data=data.input,Q=diag(rep(NA,n_variables))),
                      data,H=NA) # baseline variance constant
model2_c_fit=fitSSM(model2_c_test, inits =rep(0,n_variables+1), method = method.input) # n_variables + one NA in H
model2_c_out=KFS(model2_c_fit$model)
plot.KFS(model2_c_out,range=30:708)
model2_c_out$model$H;model2_c_out$model$Q # check intercept(0.012), logit_TAM_phone_1(0.0005), keycontacts_text_outdegree_merged (0.002)
formula = "y ~ x + c"
str_split(formula, "~")[[1]]