R/design.matrix.DDPstar.R

In DDPstar: Density Regression via Dirichlet Process Mixtures of Normal Structured Additive Regression Models

design.matrix.DDPstar <-
function(formula, data, standardise = TRUE) {
	iform <- interpret.DDPstarformula(formula, data)
	data.cov <- iform$data.cov
    data.cov.std <- iform$data.cov.std    
	X <- NULL			# Design matrices
	K <- list() 		# Penalty matrices
	K.length <- NULL	# Penalty matrices per component
	K[[1]] <-  diag(x = 0L, ncol = 1, nrow = 1) # Intercept
	K.length <- c(K.length, 1)
	Xterms <- list()
	paracoeff <- TRUE # First "coefficients" is parametric (Intercept)
	j <- 1

	if(iform$npartial == 0) { # Only the intercept
        X <- matrix(1, ncol = 1, nrow = nrow(data))
        colnames(X) <- "(Intercept)"
        res <- list()
        res$X <- X
        res$K <- K
		res$K.length <- K.length
        res$terms <- NULL
        res$iformula <- iform
	} else {
		for(i in 1:iform$npartial) {
			if(any(iform$II[,i] == -1)) {
				if(iform$h[i] == 0) { # Parametric components
					if(standardise) {
                        mf <- model.frame(paste0("~", iform$II[2,i]), data.cov.std, drop.unused.levels = TRUE)
                    } else {
                        mf <- model.frame(paste0("~", iform$II[2,i]), data.cov, drop.unused.levels = TRUE)
                    }                    
					mt <- terms(mf)				
					MM <- model.matrix(mt, mf)[,-1, drop = FALSE] # Here we delete the intercept
					paracoeff <- c(paracoeff, rep(TRUE, ncol(MM)))
					X <- cbind(X, MM)
					attr(mt, "contrast") <- attr(MM,"contrast")
					attr(mt, "xlev") <- .getXlevels(mt, mf)
					Xterms[[i]] <- mt
				
					K[[j+1]] <- diag(x = rep(0L, ncol(MM)), ncol = ncol(MM), nrow = ncol(MM))
					K.length <- c(K.length, 1)
					j <- j + 1
				} else if(iform$h[i] == -1) { # Smooth effects
					MM <- bbase.bs(data.cov[,iform$II[2,i]], ndx = iform$nseg[[i]], bdeg = iform$bdeg[[i]], pord = iform$pord[[i]], intercept = FALSE)
					Xterms[[i]] <- MM
					Bs <- MM$B
					colnames(Bs) <- paste0(iform$partial[i],".", 1:ncol(Bs))
                    paracoeff <- c(paracoeff, rep(FALSE, ncol(Bs)))
					X <- cbind(X, Bs)
					for(k in 1: length(MM$K)) {
						K[[j+1]] <- MM$K[[k]]
						attr(K[[j+1]], "atau") <- iform$atau[[i]]
						attr(K[[j+1]], "btau") <- iform$btau[[i]]
						j <- j + 1
					}
					K.length <- c(K.length, length(MM$K))
				} else { # Random effects
					mf <- model.frame(paste0("~", iform$II[2,i]), data.cov, drop.unused.levels = TRUE, na.action = NULL)               
					mt <- terms(mf)
					f.terms <- attr(mt, "term.labels")[attr(mt,"dataClasses") == "factor"]
  					MM <- model.matrix(mt, data = mf, contrasts.arg = lapply(mf[,f.terms, drop = FALSE], contrasts, contrasts=FALSE))
  					MM[is.na(MM)] <- 0
  					attr(mt, "contrast") <- attr(MM,"contrast")
  					attr(mt, "xlev") <- .getXlevels(mt, mf)
  					MM <- MM[,-1, drop = FALSE]
					paracoeff <- c(paracoeff, rep(FALSE, ncol(MM)))
					X <- cbind(X, MM)
					K[[j+1]] <- diag(x = rep(1L, ncol(MM)), ncol = ncol(MM), nrow = ncol(MM))
					attr(K[[j+1]], "atau") <- iform$atau[[i]]
					attr(K[[j+1]], "btau") <- iform$btau[[i]]
					j <- j + 1
					K.length <- c(K.length, 1) # TODO
					Xterms[[i]] <- mt
				}
			} else { # Factor by curve, varying coefficient or 2D
				if(!iform$by.var[i]) { # 2D (no by variable)
					MM <- bbase.psanova.bs(x1 = data.cov[,iform$II[1,i]], x2 = data.cov[,iform$II[2,i]], 
						ndx = iform$nseg[[i]], bdeg = iform$bdeg[[i]], pord = iform$pord[[i]], prior.2D = iform$prior.2D[i])
					Xterms[[i]] <- MM$terms
					Bs <- MM$B
					colnames(Bs) <- paste0(iform$partial[i],".", 1:ncol(Bs))
                    paracoeff <- c(paracoeff, rep(FALSE, ncol(Bs)))
					X <- cbind(X, Bs)
					for(k in 1:length(MM$K)) {
						K[[j+1]] <- MM$K[[k]]
						attr(K[[j+1]], "atau") <- iform$atau[[i]]
						attr(K[[j+1]], "btau") <- iform$btau[[i]]
						j <- j + 1
					}
					K.length <- c(K.length, length(MM$K))
				} else { # Factor by curve or varying coefficient
					if(is.factor(data.cov[,iform$II[1,i]])) { # Factor by curve
						MM <- bbase.interaction.factor.by.curve.bs(x = data.cov[,iform$II[2,i]], factor = data.cov[,iform$II[1,i]], ndx = iform$nseg[[i]], bdeg = iform$bdeg[[i]], pord = iform$pord[[i]])
						Xterms[[i]] <- MM$terms
						Bs <- MM$B
						colnames(Bs) <- paste0(iform$partial[i],".", 1:ncol(Bs))
					    paracoeff <- c(paracoeff, rep(FALSE, ncol(Bs)))
						X <- cbind(X, Bs)
						for(k in 1:length(MM$K)) {
							for(l in 1:length(MM$K[[k]])) {
								K[[j+1]] <- MM$K[[k]][[l]]
								attr(K[[j+1]], "atau") <- iform$atau[[i]]
								attr(K[[j+1]], "btau") <- iform$btau[[i]]
								j <- j + 1
							}
						}
						K.length <- c(K.length, sum(unlist(lapply(MM$K, length))))
					} else { # Varying coefficient
						if(standardise) {
							MM <- bbase.interaction.vc.bs(x1 = data.cov[,iform$II[2,i]], x2 = data.cov.std[,iform$II[1,i]], ndx = iform$nseg[[i]], bdeg = iform$bdeg[[i]], pord = iform$pord[[i]])
						} else {
							MM <- bbase.interaction.vc.bs(x1 = data.cov[,iform$II[2,i]], x2 = data.cov[,iform$II[1,i]], ndx = iform$nseg[[i]], bdeg = iform$bdeg[[i]], pord = iform$pord[[i]])
						}
						Xterms[[i]] <- MM$terms
						Bs <- MM$B
						colnames(Bs) <- paste0(iform$partial[i],".", 1:ncol(Bs))
						paracoeff <- c(paracoeff, rep(FALSE, ncol(Bs)))
						X <- cbind(X, Bs)
						for(k in 1: length(MM$K)) {
							K[[j+1]] <- MM$K[[k]]
							attr(K[[j+1]], "atau") <- iform$atau[[i]]
							attr(K[[j+1]], "btau") <- iform$btau[[i]]
							j <- j + 1
						}
						K.length <- c(K.length, length(MM$K))
					}
				}
			}
		}
		# Add the intercept
		X <- cbind(1, X)
		res <- list()
		res$X <- X
		res$K <- K
		res$K.length <- K.length
		res$terms <- Xterms
		res$iformula <- iform
	}
	res$paracoeff <- paracoeff
    res$standardise <- standardise
	class(res) <- "design.matrix.DDPstar"
	res
}