R/bddp_ps_steps.R

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

bddp_ps_steps <-
function(y, X, K, alpha.fixed = TRUE, alpha, aalpha, balpha, a, b, atau, btau, m, S, Psi, nu, nsim, L, standardise = TRUE, verbose = FALSE) {	
	multinom <- function(prob) {
        probs <- t(apply(prob,1,cumsum)) 
        res <- rowSums(probs - runif(nrow(probs)) < 0) + 1 
        return(res)  
    }
    swap <- function(vec, from, to) {
  		tmp <- to[ match(vec, from) ]
  		tmp[is.na(tmp)] <- vec[is.na(tmp)]
  		return(tmp)
	}


    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]

    v <- p <- ns <- rep(0, L)
    if(!alpha.fixed) {
		alpha <- numeric(nsim)
		# Random
		#alpha[1] <- rgamma(1, shape = aalpha, balpha)
		#v[1:(L-1)] <- rbeta(L-1, 1, alpha[1])
		#v[L] <- 1
		
		# Fixed (all equal)
		# alpha[1] <- 1
		# v <- 1/rev(1:L)

		# Fixed (decreasing)
		#alpha[1] <- 1
		#v <- rep(1/L, L)
		#v[L] <- 1

		# Fixed (increasing)
		alpha[1] <- 1
		p <- L:1/sum(L:1)
		p_sum <- cumsum(p)
		v[1] <- p[1]
		v[2:L] <- p[2:L]/(1-p_sum[1:(L-1)])

	} else {
		# Random
		#v[1:(L-1)] <- rbeta(L-1, 1, alpha)
		#v[L] <- 1

		# Fixed (all equal)
		# v <- 1/rev(1:L)

		# Fixed (decreasing)
		#v <- rep(1/L, L)
		#v[L] <- 1

		# Fixed (increasing)
		p <- L:1/sum(L:1)
		p_sum <- cumsum(p)
		v[1] <- p[1]
		v[2:L] <- p[2:L]/(1-p_sum[1:(L-1)])
	}

	if(L > 1) {
		# Moves
		moves <- matrix(0, ncol = 2, nrow = nsim)
		colnames(moves) <- c("Move 1", "Move 2")
	}
	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)

	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))
	
	z <- matrix(0, nrow = n, ncol = nsim)
	z[,1] <- rep(1,n)

	# ---------------------------------------
	# Initial beta
	# ---------------------------------------
	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
	Sigmainv_mu_c <- Sigmainv_c%*%mu_c

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

	aux <- rep(-1, nsim)
	aux[1] <- sum(log(1 - v[1:(L-1)]))
	V <- matrix(0, nrow = nsim, ncol = L)
	V[1,] <- v
	for(i in 2:nsim) {
		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]))))
				v <- replace(v, v > 1, 1)
			} else {
		    	v[1:(L-1)] <- rbeta(L-1, 1 + ns[1:(L-1)], alpha[i-1] + rev(cumsum(rev(ns[-1]))))
		    	v <- replace(v, v > 1, 1)
		    	alpha[i] <- rgamma(1, shape = aalpha + L - 1, balpha - sum(log(1 - v[1:(L-1)])))
		    	aux[i] <- 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)
		}
		# For fixed effects
		Sigmainv_mu <- Sigmainv%*%mu
		for(l in 1:L) {
			Xl  <- matrix(X[z_tmp == l, ], ncol = k, nrow = ns[l])
			yl <- yt[z_tmp == l]

			# Fixed/Parametric effects
			Xl_fixed  <- Xl[,ind_fixed, drop = FALSE]
			Xl_random <- Xl[,!ind_fixed, drop = FALSE]
			b_random <- beta[l,!ind_fixed]
			Vinv <- tau_s[l]*crossprod(Xl_fixed) + Sigmainv
			cholVinv <- try(chol(Vinv), silent = TRUE)
			mu1 <- backsolve(cholVinv, forwardsolve(t(cholVinv), Sigmainv_mu + tau_s[l]*crossprod(Xl_fixed, yl - Xl_random%*%b_random)))
			err <- backsolve(cholVinv, rnorm(sum(ind_fixed)))
			beta[l,ind_fixed] <- mu1 + err

			# Random/Smooth effects
			for (nc in 1:ncomp) {
				if(rk[nc] == 0) {
					Tau_f[i,l,nc] <- tau_f[l,nc] <- 0
				} else {
					# Coefficients
					Xl_comp  <- Xl[,np.s[nc]:np.e[nc], drop = FALSE]
					Xl_ncomp  <- Xl[,-(np.s[nc]:np.e[nc]), drop = FALSE]
					beta_ncomp <- beta[l,-(np.s[nc]:np.e[nc])]
					Vinv <- tau_s[l]*crossprod(Xl_comp) + tau_f[l,nc]*K[[nc]]
					cholVinv <- try(chol(Vinv), silent = TRUE)
					mu1 <- backsolve(cholVinv, forwardsolve(t(cholVinv), tau_s[l]*crossprod(Xl_comp, yl - Xl_ncomp%*%beta_ncomp)))
					err <- backsolve(cholVinv, rnorm(dim.comp[nc]))
					beta[l,np.s[nc]:np.e[nc]] <- mu1 + err

					# Variance components
					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[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[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]])))
					}
				}
			}
			# 1/sigma_2
			yXB <- yl - Xl%*%beta[l,]
			tau_s[l] <- rgamma(1, shape = a + ns[l]/2, rate = b + 0.5*crossprod(yXB))
		}
		#---------------------------------------
		# Moves (Paper Omiros 2008)
		#---------------------------------------
		if(L > 1) {
	        # Move 1
	        act_comp <- (1:L)[ns != 0]
	        if(length(act_comp) > 1) {
		        comp_sel <- sample(act_comp, 2, replace = FALSE)
		        prob_change <- min(1, (p[comp_sel[1]]/p[comp_sel[2]])^(ns[comp_sel[2]]-ns[comp_sel[1]]))
		        if(!is.na(prob_change) && runif(1) <= prob_change) {
		        	beta[comp_sel,]  <- beta[rev(comp_sel),]
		        	tau_s[comp_sel]  <- tau_s[rev(comp_sel)]
		        	tau_f[comp_sel,] <- tau_f[rev(comp_sel),]
		        	p[comp_sel] <- p[rev(comp_sel)]	        	
		        	#v[comp_sel] <- v[rev(comp_sel)]
		        	z_tmp <- swap(z_tmp, comp_sel, rev(comp_sel))
		        	moves[i,1] <- 1
		        }
		    }
	        # Move 2
	        # comp_sel <- sample(1:(L-2), 1, replace = FALSE)
	        # prob_change <- min(1, ((1-v[comp_sel + 1])^ns[comp_sel])/((1-v[comp_sel])^ns[comp_sel+1]))
	        # if(!is.na(prob_change) && runif(1) <= prob_change) {	        	
	        # 	comp_sel <- c(comp_sel, comp_sel + 1)
	        # 	beta[comp_sel,]  <- beta[rev(comp_sel),]
	        # 	tau_s[comp_sel]  <- tau_s[rev(comp_sel)]
	        # 	tau_f[comp_sel,] <- tau_f[rev(comp_sel),]
	        # 	p[comp_sel] <- p[rev(comp_sel)]
	        # 	v[comp_sel] <- v[rev(comp_sel)]
	        # 	z_tmp <- swap(z_tmp, comp_sel, rev(comp_sel))
	        # 	moves[i,2] <- 1
	        # }

	        comp_sel <- sample(2:L, 1, replace = FALSE)
	        prob_change <- min(1, ((1-v[comp_sel])^ns[comp_sel - 1])/((1-v[comp_sel - 1])^ns[comp_sel]))
	        if(!is.na(prob_change) && runif(1) <= prob_change) {	        	
	        	comp_sel <- c(comp_sel, comp_sel - 1)
	         	beta[comp_sel,]  <- beta[rev(comp_sel),]
	         	tau_s[comp_sel]  <- tau_s[rev(comp_sel)]
	         	tau_f[comp_sel,] <- tau_f[rev(comp_sel),]
	         	p[comp_sel] <- p[rev(comp_sel)]
	         	v[comp_sel] <- v[rev(comp_sel)]
	         	z_tmp <- swap(z_tmp, comp_sel, rev(comp_sel))
	         	moves[i,2] <- 1
	        }
	    }
		# ---------------------------------------
		# ---------------------------------------
		# 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]
		# ---------------------------------------
		
		P[i,] <- p
		V[i,] <- v
        z[,i] <- z_tmp
        Beta[i,,] <- beta
        Tau_s[i,] <- tau_s
        Tau_f[i,,] <- tau_f
	}
	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, alpha = alpha))
    res <- list()
    res$P <- P
    res$Beta <- Beta
    res$Sigma2 <- Sigma2
    res$Tau_f <- Tau_f
    res$Z <- z
    res$alpha <- alpha
    if(L > 1) {
    	res$moves <- moves
    }
    res 	
}