R/ame2.R

In ianmtaylor1/amen2: Alternate MCMC Algorithm for AME Models

ame2 <- function(Y, Xdyad=NULL, Xrow=NULL, Xcol=NULL, intercept=TRUE,
                    model="nrm", rvar=TRUE, cvar=TRUE,
                    prior="ig",
                    burn=500, n.iter=10000, save.every=25, print=TRUE) {
  # Check for appropriate model
  stopifnot(model %in% c("nrm", "bin"))

  # Check for appropriate prior
  stopifnot(prior %in% c("ig", "hs", "tiny"))
  
  # Set up data as appropriate
  stopifnot(length(dim(Y)) == 2)
  stopifnot(dim(Y)[1] == dim(Y)[2])
  n <- dim(Y)[1]
  X <- amen::design_array(Xrow, Xcol, Xdyad, intercept, n)
  p <- dim(X)[3]
  diag(Y) <- NA

  # Initial values
  a <- rep(0, n)
  b <- rep(0, n)
  beta <- rep(0, p)
  sigma.squared <- 1
  sigma.ab <- diag(2)
  if (model == "nrm") {
    Z <- Y
  } else if (model == "bin") {
    Z <- matrix(rep(0,n*n), ncol=n)
    diag(Z) <- NA
  }

  # Lists to save the samples in temporarily
  samples.a <- list()
  samples.b <- list()
  samples.beta <- list()
  samples.sigma.squared <- list()
  samples.sigma.ab <- list()

  # Iterate through the gibbs sampler
  for (i in 1:(n.iter + burn)) {
    # Update latent normal variables (if probit model)
    if (model == "nrm") {
      Z <- Y
    } else if (model == "bin") {
      Z <- update_Z_bin_fc(Y=Y, X=X, beta=beta, sigma.squared=sigma.squared, a=a, b=b, sigma.ab=sigma.ab)
    }
    # Update sigma.squared
    if (model == "nrm") {
      sigma.squared <- update_sigma.squared_fc(Y=Y, Z=Z, X=X, beta=beta, a=a, b=b, sigma.ab=sigma.ab)
    } else if (model == "bin") {
      sigma.squared <- 1
    }
    # Update Sigma.ab
    if (rvar || cvar) {
      if (prior == "ig") {
        sigma.ab <- update_sigma.ab_fc(Y=Y, Z=Z, X=X, beta=beta, sigma.squared=sigma.squared, a=a, b=b)
      } else if (prior == "hs") {
        sigma.ab <- update_sigma.ab_horseshoe_fc(Y=Y, Z=Z, X=X, beta=beta, sigma.squared=sigma.squared, a=a, b=b, sigma.ab=sigma.ab)
      } else if (prior == "tiny") {
        sigma.ab <- diag(.0001, nrow=2, ncol=2)
      }
      
    }
    # Update beta
    beta <- update_beta_fc(Y=Y, Z=Z, X=X, sigma.squared=sigma.squared, a=a, b=b, sigma.ab=sigma.ab)
    # Update a,b
    if (rvar) {
      a <- update_a_fc(Y=Y, Z=Z, X=X, beta=beta, sigma.squared=sigma.squared, b=b, sigma.ab=sigma.ab)
    }
    if (cvar) {
      b <- update_b_fc(Y=Y, Z=Z, X=X, beta=beta, sigma.squared=sigma.squared, a=a, sigma.ab=sigma.ab)
    }

    # Save if necessary
    if ((i > burn) && ((i-burn) %% save.every == 0)) {
      if (print) {
        cat("Saving...",i,"\n")
      }
      samples.a <- append(samples.a, list(a))
      samples.b <- append(samples.b, list(b))
      samples.beta <- append(samples.beta, list(beta))
      samples.sigma.squared <- append(samples.sigma.squared, list(sigma.squared))
      samples.sigma.ab <- append(samples.sigma.ab, list(sigma.ab))
    }
  }

  # Create form in which to return results
  results.beta <- matrix(unlist(samples.beta), ncol=p, byrow=TRUE)
  colnames(results.beta) <- dimnames(X)[[3]]
  results.a <- matrix(unlist(samples.a), ncol=n, byrow=TRUE)
  results.b <- matrix(unlist(samples.b), ncol=n, byrow=TRUE)
  vcdata <- c()
  for (i in 1:length(samples.sigma.ab)) {
    vcdata <- c(vcdata,
                samples.sigma.ab[[i]][1,1], #va
                samples.sigma.ab[[i]][2,2], #vb
                samples.sigma.ab[[i]][1,2], #cab
                0, #rho
                samples.sigma.squared[[i]] #ve
    )
  }
  results.vc <- matrix(vcdata, ncol=5, byrow=TRUE)
  colnames(results.vc) <- c("va","vb","cab","rho","ve")

  # Return Samples
  return(
    list(
      BETA=results.beta,
      A=results.a,
      B=results.b,
      VC=results.vc
    )
  )
}