R/bddp_ps.R

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

bddp_ps <-
function(y, X, K, alpha.fixed = TRUE, alpha, aalpha, balpha, a, b, atau, btau, m, S, Psi, nu, nsim, L, standardise = TRUE) {	
	multinom <- function(prob) {
        probs <- t(apply(prob,1,cumsum)) 
        res <- rowSums(probs - runif(nrow(probs)) < 0) + 1 
        return(res)  
    }

    yt <- y  
	if (standardise == TRUE) {
		yt <- (y - mean(y))/sd(y)
	}

	n <-  length(y)
	k <-  ncol(X)
	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.
	np.e <- cumsum(dim.comp)
	np.s <- np.e - dim.comp + 1

	# Fixed effects
	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)
    S_inv <- solve(S)
    mu <- mvrnorm(1, mu = m, Sigma = S)
    Sigmainv <- rWishart(1, df = nu, solve(nu*Psi))[,,1]

    # Complete
    mu_c <- rep(0, k)
	Sigmainv_c <- matrix(0, ncol = k, nrow = k)
	mu_c[ind_fixed] <- mu
	Sigmainv_c[ind_fixed,ind_fixed] <- Sigmainv
	
	if(!alpha.fixed) {
		alpha <- numeric(nsim)
		alpha[1] <- 1
	}
	
	p <- ns <- rep(0,L)
	v <- rep(1/L,L)
	v[L] <- 1

	beta <- matrix(0, nrow = L, ncol = k)
	aux <- try(solve(t(X)%*%X)%*%t(X)%*%yt, silent = TRUE)
	if(!inherits(aux, "try-error")) {
		for(l in 1:L) {
		  beta[l,] <- aux   
		}
	}

	tau_s <- rep(1/var(yt), L)			 			# residual variance
	tau_f <- matrix(1, nrow = L, ncol = ncomp)  	# smoothing parameter
	tau_f[, rk == 0] <- 0                       	# tau_f is zero for unpenalized coefficients
	prop <- prob <- matrix(0, nrow = n, ncol = L)

	inv_error <- P <- Tau_s <- Sigma2 <- matrix(0, nrow = nsim, ncol = L)
	Tau_f <- array(0, c(nsim, L, ncomp))
	Beta <- array(0, c(nsim, L, k))
	
	Beta[1,,] <- beta
	Tau_s[1,] <- tau_s
	Tau_f[1,,] <- tau_f
	
	z <- matrix(0, nrow = n, ncol = nsim)
	z[,1] <- rep(1,n)
  
	for(i in 2:nsim) {
		verbose <- FALSE
		if (verbose) {
            if (i%%100 == 0) {
                cat(paste("MCMC iteration: ", i, "\n", sep = ""))
            }
        }
		if(L > 1) {
			cumv <- cumprod(1-v)
			p[1] <- v[1]
			p[2:L] <- v[2:L]*cumv[1:(L-1)]

	        for(l in 1:L) {
	            prop[,l] <- p[l]*dnorm(yt, mean = X%*%beta[l,], sd = sqrt(1/tau_s[l]))
	        }
	        prob <- prop/rowSums(prop)
			z_tmp <- multinom(prob)
			ns <- sapply(1:L, function(x, v) sum(v == x), v = z_tmp)

			if(alpha.fixed) {
				v[1:(L-1)] <- rbeta(L-1, 1 + ns[1:(L-1)], alpha + rev(cumsum(rev(ns[-1]))))
			} else {
		    	v[1:(L-1)] <- rbeta(L-1, 1 + ns[1:(L-1)], alpha[i-1] + rev(cumsum(rev(ns[-1]))))
		    	alpha[i] <- rgamma(1, shape = aalpha + L - 1, balpha - sum(log(1 - v[1:(L-1)])))
	    	}

		} else {
			p <- 1
			z_tmp <- rep(1, n)
			ns <- sapply(1:L, function(x, v) sum(v == x), v = z_tmp)
		}

		Sigmainv_mu_c <- Sigmainv_c%*%mu_c

		for(l in 1:L) {
			if(ncomp == 1) {
				Kaux <- tau_f[l,]*K[[1]]
			} else {
				Kaux <- as.matrix(Matrix::bdiag(mapply("*", K, tau_f[l,], SIMPLIFY = FALSE)))
			}
			Xl  <- matrix(X[z_tmp == l, ], ncol = k, nrow = ns[l])
			yl <- yt[z_tmp == l]
			
			# Cholesky
			Vinv <- tau_s[l]*crossprod(Xl) + Kaux + Sigmainv_c
			cholVinv <- try(chol(Vinv), silent = TRUE)
			if(inherits(cholVinv, "try-error")) { 
				mu1 <- rep(0, length = k)
				V <- MASS::ginv(Kaux)
				beta[l,] <- MASS::mvrnorm(1, mu = mu1, Sigma = V)
			} else {
				beta_aux <- rnorm(k)
				mu1 <- backsolve(cholVinv, forwardsolve(t(cholVinv), Sigmainv_mu_c + tau_s[l]*crossprod(Xl, yl)))
				err <- backsolve(cholVinv, beta_aux)
				beta[l,] <- mu1 + err
			}

			yXB <- yl - Xl%*%beta[l,]
			tau_s[l] <- rgamma(1, shape = a + ns[l]/2, rate = b + 0.5*crossprod(yXB))
			for (nc in 1:ncomp) {
				if(rk[nc] == 0) {
					Tau_f[i,l,nc] <- tau_f[l,nc] <- 0
				} else {
					if(!is.null(attr(K[[nc]], "atau")) & !is.null(attr(K[[nc]], "btau"))) {
						atau_tmp <- attr(K[[nc]], "atau")
						btau_tmp <- attr(K[[nc]], "btau")
						Tau_f[i,l,nc] <- tau_f[l,nc] <- rgamma(1, shape = atau_tmp + rk[nc]/2, rate = btau_tmp + 0.5*(t(beta[l,np.s[nc]:np.e[nc]])%*%K[[nc]]%*%(beta[l,np.s[nc]:np.e[nc]])))
					} else { 
						Tau_f[i,l,nc] <- tau_f[l,nc] <- rgamma(1, shape = atau + rk[nc]/2, rate = btau + 0.5*(t(beta[l,np.s[nc]:np.e[nc]])%*%K[[nc]]%*%(beta[l,np.s[nc]:np.e[nc]])))
					}
				}
			}
		}
		# ---------------------------------------
		# Fixed/parametric effects
		# ---------------------------------------
		Vaux_inv <- S_inv + L*Sigmainv
		eig <- eigen(Vaux_inv, symmetric = TRUE)
		Vaux <- crossprod(t(eig$vector)/sqrt(eig$values))
        if(k == 1) {
            meanmu <- Vaux%*%(S_inv%*%m + Sigmainv%*%sum(beta[,ind_fixed,drop = FALSE]))
        } else {
            meanmu <- Vaux%*%(S_inv%*%m + Sigmainv%*%colSums(beta[,ind_fixed,drop = FALSE]))
        }
        mu <- mvrnorm(1, mu = meanmu, Sigma = Vaux)
        
        Vaux1 <- 0
        for(l in 1:L) {         
            Vaux1 <- Vaux1 + crossprod(beta[l,ind_fixed,drop = FALSE] - mu)
        }
        Sigmainv <- rWishart(1, nu + L, solve(nu*Psi + Vaux1))[,,1]

        mu_c[ind_fixed] <- mu
		Sigmainv_c[ind_fixed,ind_fixed] <- Sigmainv
		# ---------------------------------------

		P[i,] <- p
        z[,i] <- z_tmp
        Beta[i,,] <- beta
        Tau_s[i,] <- tau_s
	}
	if (standardise == TRUE) {
        Beta[,,1] <- sd(y)*Beta[,,1] + mean(y)
        if(k > 1) {
            Beta[,,2:k] <- sd(y)*Beta[,,2:k]
        }
        Sigma2 <- var(y)*(1/Tau_s)
    } else {
        Sigma2 <- 1/Tau_s
    }
 	return(list(P = P, Beta = Beta, Sigma2 = Sigma2, Tau_f = Tau_f, Z = z))
}