R/DDPstar.R

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

DDPstar <-
function(formula, data, standardise = TRUE, 
    compute.lpml = FALSE, compute.WAIC = FALSE, compute.DIC = FALSE, 
    prior = priorcontrol(), mcmc = mcmccontrol(), verbose = FALSE) {	
	mcmc <- do.call("mcmccontrol", mcmc)
    prior <- do.call("priorcontrol", prior)

	if(inherits(formula, "character")) {		  
		formula <- as.formula(formula)
	}
	# Marker variable
    tf <- terms.formula(formula, specials = c("f","rae"))
    if (attr(tf, "response") > 0) {
        response <- as.character(attr(tf, "variables")[2])
    } else {
        stop("The formula should include the response variable (left hand side)")
    }
    
    # Variables in the model
    names.cov <- get_vars_formula(formula)
    
    if (inherits(data, what = 'data.frame')) {
        data <- as.data.frame(data)
    } else {
        stop("The object specified in argument 'data' is not a data frame")
    }

    if(sum(is.na(match(c(response, names.cov), names(data))))) {
        stop("Not all needed variables are supplied in data")
    }

    # New data, removing missing values
    data.new <- data[,c(response,names.cov)]
    omit <- apply(data.new[, c(response,names.cov)], 1, anyNA)    
    data.new <- data.new[!omit,,drop = FALSE]
        
    # Hyperparameters
    alpha.fixed <- prior$alpha.fixed 
    alpha <- prior$alpha
	aalpha <- prior$aalpha
    balpha <- prior$balpha
	m <- prior$m0
    S <- prior$S0
    nu <- prior$nu
    Psi <- prior$Psi
	a <- prior$a
	b <- prior$b
	L <- prior$L

    # Lang & Brezger
    atau <- prior$atau
    btau <- prior$btau
    
    # MCMC specifications
	nburn <- mcmc$nburn
    nsave <- mcmc$nsave
    nskip <- mcmc$nskip

    # Design and precision matrix
	MM <- design.matrix.DDPstar(formula, data.new, standardise)
	X <- MM$X
	K <- MM$K # This is a list with as many elements as model components (separated by fixed and random). For "fixed" components, K is a matrix of zeroes of the same size as the number of coefficients

	# Fixed effects (to obtain/check prior distributions)
	ncomp <- length(K) 						# Number of components
	dim.comp <- unlist(lapply(K, ncol)) 	# Number of coefficients per component
	rk <- unlist(lapply(K, function(x) qr(x)$rank)) # Rank of the "penalty" matrix. A value of zero will indicate fixed (unpenalized) coefficients.
	ind_fixed <- vector()
	for(i in 1:ncomp) {
		if(rk[i] == 0) {
			ind_fixed <- c(ind_fixed, rep(TRUE, dim.comp[i]))
		} else {
			ind_fixed <- c(ind_fixed, rep(FALSE, dim.comp[i]))
		}
	}
	k_fixed <- sum(ind_fixed)
	if(!standardise & (anyNA(prior$m0) | anyNA(prior$S0) | anyNA(prior$Psi) | is.na(prior$b))) {
		X_fixed <- X[,ind_fixed]
		response_val <- data.new[,response]
		n <- nrow(X_fixed)
        res <- ols.function(X_fixed, response_val, vcov = TRUE)
        coefs <- res$coeff
        var <- sum((response_val - X_fixed %*% coefs)^2)/(n - ncol(X_fixed))
        cov <- res$vcov*var
    }
    if(is.na(L)) {
        L <- 10 
    } else { 
        if(length(L) != 1) {
            stop(paste0("L must be a constant"))
        }
    }
    if(anyNA(m)) {
        if(standardise) m <- rep(0, k_fixed)
        else m <- coefs
    } else { 
        if(length(m) != k_fixed) {
            stop(paste0("'m' must be a vector of length ", k_fixed))
        }
    }
    if(anyNA(S)) {
        if(standardise) S <- 10*diag(k_fixed)
        else S <- cov
    } else { 
        if(!is.matrix(S) | !all(dim(S) == c(k_fixed, k_fixed))) {
            stop(paste0("'S' must be a matrix of dimension ", k_fixed, "x", k_fixed))
        }
    }

    if(anyNA(Psi)) {
        if(standardise) Psi <- diag(k_fixed)
        else Psi <- 30*cov
    } else { 
        if(!is.matrix(Psi) | !all(dim(Psi) == c(k_fixed, k_fixed))) {
            stop(paste0("'Psi' must be a matrix of dimension ", k_fixed, "x", k_fixed))
        }
    }

    if(is.na(nu)) {
        nu <- k_fixed + 2
    } else { 
        if(nu < k_fixed + 2){
            stop(paste0("'nu' must be larger than ", k_fixed + 2))
        }
    }

    if(is.na(a)) {
        a <- 2
    } else { 
        if(length(a) != 1) {
            stop(paste0("'a' must be a constant"))
        }
    }

    if(is.na(b)) {
        if(standardise) {
            b <- 0.5
        } else b <- var/2
    } else { 
        if(length(b) != 1) {
            stop(paste0("'b' must be a constant"))
        }
    }

	fit <- bddp_ps_steps(y = data.new[,response], X = X, K = K,
        alpha.fixed = alpha.fixed, alpha = alpha,
        aalpha = aalpha, balpha = balpha, 
        a = a, b = b, atau = atau, btau = btau, 
		m = m, S = S, Psi = Psi, nu = nu,
		nsim = nburn + nsave, L = L, 
		standardise = standardise, verbose = verbose)

    alpha.fit <- if(alpha.fixed) {
        fit$alpha

    } else {
        fit$alpha[seq(nburn+1, nburn + nsave, by = nskip)]
    }

	res <- list()
    res$call <- match.call()
    res$data <- data
    res$data_model <- list(y = data.new[,response], X = X)
    res$missing.ind <- omit
    res$mcmc <- mcmc
	res$fit <- list(formula = formula, names_cov = names.cov, mm = MM, 
        beta = fit$Beta[seq(nburn+1, nburn + nsave, by = nskip),,,drop = FALSE], 
		sd = sqrt(fit$Sigma2[seq(nburn+1, nburn + nsave, by = nskip),,drop = FALSE]), 
		probs = fit$P[seq(nburn+1, nburn + nsave, by = nskip),,drop = FALSE],
		tau_sm = fit$Tau_f[seq(nburn+1, nburn + nsave, by = nskip),,,drop = FALSE],
        S = fit$Z[,seq(nburn+1, nburn + nsave, by = nskip), drop = FALSE], alpha = alpha.fit)

    if(L > 1) {
        res$fit$moves <- fit$moves[seq(nburn+1, nburn + nsave, by = nskip),]
    }
	if(compute.lpml | compute.WAIC | compute.DIC) {
		term <- inf_criteria(y = data.new[,response], X = X, res = res$fit)        
	}

	if(compute.lpml) {
		res$lpml <- lpml_new(y = data.new[,response], X = X, res = res$fit, termsum = term)
	}
	if(compute.WAIC) {
		res$WAIC <- waicnp_new(y = data.new[,response], X = X, res = res$fit, termsum = term)
	}
	if(compute.DIC) {
		res$DIC <- dic(y = data.new[,response], X = X, res = res$fit, termsum = term)
	}


	class(res) <- "DDPstar"
	res
}