# R/bkmr_parameter_update_functions.R In bkmr: Bayesian Kernel Machine Regression

```beta.update <- function(X, Vinv, y, sigsq.eps) {
XVinv <- crossprod(X, Vinv)
Vbeta <- chol2inv(chol(XVinv %*% X))
cholVbeta <- chol(Vbeta)
betahat <- Vbeta %*% XVinv %*% y
n01 <- rnorm(ncol(X))
betahat + crossprod(sqrt(sigsq.eps)*cholVbeta, n01)
}

sigsq.eps.update <- function(y, X, beta, Vinv, a.eps=1e-3, b.eps=1e-3) {
mu <- y - X%*%beta
prec.y <- rgamma(1, shape=a.eps + nrow(X)/2, rate=b.eps + 1/2*crossprod(mu, Vinv)%*%mu)
1/prec.y
}

ystar.update <- function(y, X, beta, h) {
mu <-  drop(h + X %*% beta)
lower <- ifelse(y == 1, 0, -Inf)
upper <- ifelse(y == 0, 0,  Inf)
samp <- truncnorm::rtruncnorm(1, a = lower, b = upper, mean = mu, sd = 1)
drop(samp)
}
#' @importFrom tmvtnorm rtmvnorm
ystar.update.noh <- function(y, X, beta, Vinv, ystar) {
mu <-  drop(X %*% beta)
lower <- ifelse(y == 1, 0, -Inf)
upper <- ifelse(y == 0, 0,  Inf)
samp <- tmvtnorm::rtmvnorm(1, mean = mu, H = Vinv, lower = lower, upper = upper, algorithm = "gibbs", start.value = ystar)
#samp <- truncnorm::rtruncnorm(1, a = lower, b = upper, mean = mu, sd = 1)
drop(samp)
}

r.update <- function(r, whichcomp, delta, lambda, y, X, beta, sigsq.eps, Vcomps, Z, data.comps, control.params, rprop.gen, rprop.logdens, rprior.logdens, ...) {
# r.params <- set.r.params(r.prior = control.params\$r.prior, comp = whichcomp, r.params = control.params\$r.params)
r.params <- make_r_params_comp(control.params\$r.params, whichcomp)
rcomp <- unique(r[whichcomp])
if(length(rcomp) > 1) stop("rcomp should only be 1-dimensional")

## generate a proposal
rcomp.star <- rprop.gen(current = rcomp, r.params = r.params)
lambda.star <- lambda
delta.star <- delta
move.type <- NA

## part of M-H ratio that depends on the proposal distribution
negdifflogproposal <- -rprop.logdens(rcomp.star, rcomp, r.params = r.params) + rprop.logdens(rcomp, rcomp.star, r.params = r.params)

## prior distribution
diffpriors <- rprior.logdens(rcomp.star, r.params = r.params) - rprior.logdens(rcomp, r.params = r.params)

r.star <- r
r.star[whichcomp] <- rcomp.star

## M-H step
return(MHstep(r=r, lambda=lambda, lambda.star=lambda.star, r.star=r.star, delta=delta, delta.star=delta.star, y=y, X=X, Z=Z, beta=beta, sigsq.eps=sigsq.eps, diffpriors=diffpriors, negdifflogproposal=negdifflogproposal, Vcomps=Vcomps, move.type=move.type, data.comps=data.comps))
}

rdelta.comp.update <- function(r, delta, lambda, y, X, beta, sigsq.eps, Vcomps, Z, ztest, data.comps, control.params, rprop.gen2, rprop.logdens1, rprior.logdens, rprior.logdens2, rprop.logdens2, rprop.gen1, ...) { ## individual variable selection
r.params <- control.params\$r.params
a.p0 <- control.params\$a.p0
b.p0 <- control.params\$b.p0
delta.star <- delta
r.star <- r

move.type <- ifelse(all(delta[ztest] == 0), 1, sample(c(1,2),1))
move.prob <- ifelse(all(delta[ztest] == 0), 1, 1/2)
if(move.type == 1) {
comp <- ifelse(length(ztest) == 1, ztest, sample(ztest, 1))
r.params <- set.r.params(r.prior = control.params\$r.prior, comp = comp, r.params = r.params)

delta.star[comp] <- 1 - delta[comp]
move.prob.star <- ifelse(all(delta.star[ztest] == 0), 1, 1/2)
r.star[comp] <- ifelse(delta.star[comp] == 0, 0, rprop.gen1(r.params = r.params))

diffpriors <- (lgamma(sum(delta.star[ztest]) + a.p0) + lgamma(length(ztest) - sum(delta.star[ztest]) + b.p0) - lgamma(sum(delta[ztest]) + a.p0) - lgamma(length(ztest) - sum(delta[ztest]) + b.p0)) + ifelse(delta[comp] == 1, -1, 1)*with(list(r.sel = ifelse(delta[comp] == 1, r[comp], r.star[comp])), rprior.logdens(x = r.sel, r.params = r.params))

negdifflogproposal <- -log(move.prob.star) + log(move.prob) - ifelse(delta[comp] == 1, -1, 1)*with(list(r.sel = ifelse(delta[comp] == 1, r[comp], r.star[comp])), rprop.logdens1(x = r.sel, r.params = r.params))

} else if(move.type == 2) {
comp <- ifelse(length(which(delta == 1)) == 1, which(delta == 1), sample(which(delta == 1), 1))
r.params <- set.r.params(r.prior = control.params\$r.prior, comp = comp, r.params = r.params)

r.star[comp] <- rprop.gen2(current = r[comp], r.params = r.params)

diffpriors <- rprior.logdens(r.star[comp], r.params = r.params) - rprior.logdens(r[comp], r.params = r.params)

negdifflogproposal <- -rprop.logdens2(r.star[comp], r[comp], r.params = r.params) + rprop.logdens2(r[comp], r.star[comp], r.params = r.params)
}

lambda.star <- lambda

## M-H step
return(MHstep(r=r, lambda=lambda, lambda.star=lambda.star, r.star=r.star, delta=delta, delta.star=delta.star, y=y, X=X, Z=Z, beta=beta, sigsq.eps=sigsq.eps, diffpriors=diffpriors, negdifflogproposal=negdifflogproposal, Vcomps=Vcomps, move.type=move.type, data.comps=data.comps))
}

rdelta.group.update <- function(r, delta, lambda, y, X, beta, sigsq.eps, Vcomps, Z, ztest, data.comps, control.params, rprop.gen1, rprior.logdens, rprop.logdens1, rprop.gen2, rprop.logdens2, ...) { ## grouped variable selection
r.params <- control.params\$r.params
a.p0 <- control.params\$a.p0
b.p0 <- control.params\$b.p0
groups <- control.params\$group.params\$groups
sel.groups <- control.params\$group.params\$sel.groups
neach.group <- control.params\$group.params\$neach.group
delta.star <- delta
r.star <- r

# if(length(mu.r) == 1) mu.r <- rep(mu.r, nz)
# if(length(sigma.r) == 1) sigma.r <- rep(sigma.r, nz)

delta.source <- sapply(sel.groups, function(x) ifelse(any(delta[which(groups == groups[x])] == 1), 1, 0))
delta.source.star <- delta.source

## randomly select move type
if(all(delta.source == 0)) {
move.type <- 1
move.prob <- 1
} else if(length(which(neach.group > 1 & delta.source == 1)) == 0) {
move.type <- sample(c(1, 3), 1)
move.prob <- 1/2
} else {
move.type <- sample(1:3, 1)
move.prob <- 1/3
}
# move.type <- ifelse(all(delta.source == 0), 1, ifelse(length(which(neach.group > 1 & delta.source == 1)) == 0, sample(c(1, 3), 1), sample(1:3, 1)))

# print(move.type)

if(move.type == 1) { ## randomly select a source and change its state (e.g., from being in the model to not being in the model)

source <- sample(seq_along(delta.source), 1)
source.comps <- which(groups == source)

# r.params <- set.r.params(r.prior = control.params\$r.prior, comp = source.comps, r.params = r.params)

delta.source.star[source] <- 1 - delta.source[source]
delta.star[source.comps] <- rmultinom(1, delta.source.star[source], rep(1/length(source.comps), length(source.comps)))
move.prob.star <- ifelse(all(delta.source.star == 0), 1, ifelse(length(which(neach.group > 1 & delta.source.star == 1)) == 0, 1/2, 1/3))

## which component got switched
comp <- ifelse(delta.source[source] == 1, source.comps[which(delta[source.comps] == 1)], source.comps[which(delta.star[source.comps] == 1)])
r.params <- set.r.params(r.prior = control.params\$r.prior, comp = comp, r.params = r.params)

r.star[comp] <- ifelse(delta.star[comp] == 0, 0, rprop.gen1(r.params = r.params))

# diffpriors <- ifelse(delta.source[source] == 1, log(length(sel.groups) - sum(delta.source) + b.p0) - log(sum(delta.source.star) + a.p0), log(sum(delta.source) + a.p0) - log(length(sel.groups) - sum(delta.source.star) + b.p0)) + ifelse(delta.source[source] == 1, 1, -1)*log(length(source.comps)) + ifelse(delta.source[source] == 1, -1, 1)*with(list(r.sel = ifelse(delta.source[source] == 1, r[source.comps][which(delta[source.comps] == 1)], r.star[source.comps][which(delta.star[source.comps] == 1)])), rprior.logdens(x = r.sel, r.params = r.params))
diffpriors <- ifelse(delta.source[source] == 1, log(length(sel.groups) - sum(delta.source) + b.p0) - log(sum(delta.source.star) + a.p0), log(sum(delta.source) + a.p0) - log(length(sel.groups) - sum(delta.source.star) + b.p0)) + ifelse(delta.source[source] == 1, 1, -1)*log(length(source.comps)) + ifelse(delta.source[source] == 1, -1, 1)*with(list(r.sel = ifelse(delta.source[source] == 1, r[comp], r.star[comp])), rprior.logdens(x = r.sel, r.params = r.params))

# negdifflogproposal <- -log(move.prob.star) + log(move.prob) -ifelse(delta.source[source] == 1, 1, -1)*(log(length(source.comps)) - with(list(r.sel = ifelse(delta.source[source] == 1, r[source.comps][which(delta[source.comps] == 1)], r.star[source.comps][which(delta.star[source.comps] == 1)])), rprop.logdens1(x = r.sel, r.params = r.params)))
negdifflogproposal <- -log(move.prob.star) + log(move.prob) -ifelse(delta.source[source] == 1, 1, -1)*(log(length(source.comps)) - with(list(r.sel = ifelse(delta.source[source] == 1, r[comp], r.star[comp])), rprop.logdens1(x = r.sel, r.params = r.params)))

} else if(move.type == 2) { ## randomly select a multi-component source that is in the model and change which component is included

tmp <- which(neach.group > 1 & delta.source == 1)
source <- ifelse(length(tmp) == 1, tmp, sample(tmp, 1))
source.comps <- which(groups == source)

oldcomp <- source.comps[delta[source.comps] == 1]
tmp <- source.comps[delta[source.comps] == 0]
comp <- ifelse(length(tmp) == 1, tmp, sample(tmp, 1))

r.params.oldcomp <- set.r.params(r.prior = control.params\$r.prior, comp = oldcomp, r.params = r.params)
r.params <- set.r.params(r.prior = control.params\$r.prior, comp = comp, r.params = r.params)

delta.star[oldcomp] <- 0
delta.star[comp] <- 1

r.star[oldcomp] <- 0
r.star[comp] <- rprop.gen1(r.params = r.params)

diffpriors <- rprior.logdens(r.star[comp], r.params = r.params) - rprior.logdens(r[oldcomp], r.params = r.params.oldcomp)

negdifflogproposal <- -rprop.logdens1(r.star[comp], r.params = r.params) + rprop.logdens1(r[oldcomp], r.params = r.params.oldcomp)

} else if(move.type == 3) { ## randomly select a component that is in the model and update it
tmp <- which(delta == 1)
comp <- ifelse(length(tmp) == 1, tmp, sample(tmp, 1))

r.params <- set.r.params(r.prior = control.params\$r.prior, comp = comp, r.params = r.params)

r.star[comp] <- rprop.gen2(current = r[comp], r.params = r.params)

diffpriors <- rprior.logdens(r.star[comp], r.params = r.params) - rprior.logdens(r[comp], r.params = r.params)

negdifflogproposal <- -rprop.logdens2(r.star[comp], r[comp], r.params = r.params) + rprop.logdens2(r[comp], r.star[comp], r.params = r.params)
}

lambda.star <- lambda

## M-H step
return(MHstep(r=r, lambda=lambda, lambda.star=lambda.star, r.star=r.star, delta=delta, delta.star=delta.star, y=y, X=X, Z=Z, beta=beta, sigsq.eps=sigsq.eps, diffpriors=diffpriors, negdifflogproposal=negdifflogproposal, Vcomps=Vcomps, move.type=move.type, data.comps=data.comps))
}

lambda.update <- function(r, delta, lambda, whichcomp=1, y, X, Z = Z, beta, sigsq.eps, Vcomps, data.comps, control.params) {
lambda.jump <- control.params\$lambda.jump[whichcomp]
mu.lambda <- control.params\$mu.lambda[whichcomp]
sigma.lambda <- control.params\$sigma.lambda[whichcomp]
lambdacomp <- lambda[whichcomp]

## generate a proposal
lambdacomp.star <- rgamma(1, shape=lambdacomp^2/lambda.jump^2, rate=lambdacomp/lambda.jump^2)
r.star <- r
delta.star <- delta
move.type <- NA

## part of M-H ratio that depends on the proposal distribution
negdifflogproposal <- -dgamma(lambdacomp.star, shape=lambdacomp^2/lambda.jump^2, rate=lambdacomp/lambda.jump^2, log=TRUE) + dgamma(lambdacomp, shape=lambdacomp.star^2/lambda.jump^2, rate=lambdacomp.star/lambda.jump^2, log=TRUE)

## prior distribution
diffpriors <- dgamma(lambdacomp.star, shape=mu.lambda^2/sigma.lambda^2, rate=mu.lambda/sigma.lambda^2, log=TRUE) - dgamma(lambdacomp, shape=mu.lambda^2/sigma.lambda^2, rate=mu.lambda/sigma.lambda^2, log=TRUE)

lambda.star <- lambda
lambda.star[whichcomp] <- lambdacomp.star

## M-H step
return(MHstep(r=r, lambda=lambda, lambda.star=lambda.star, r.star=r.star, delta=delta, delta.star=delta.star, y=y, X=X, Z=Z, beta=beta, sigsq.eps=sigsq.eps, diffpriors=diffpriors, negdifflogproposal=negdifflogproposal, Vcomps=Vcomps, move.type=move.type, data.comps=data.comps))
}

MHstep <- function(r, lambda, lambda.star, r.star, delta, delta.star, y, X, Z, beta, sigsq.eps, diffpriors, negdifflogproposal, Vcomps, move.type, data.comps) {
## compute log M-H ratio
Vcomps.star <- makeVcomps(r.star, lambda.star, Z, data.comps)
mu <- y - X%*%beta
diffliks <- 1/2*Vcomps.star\$logdetVinv - 1/2*Vcomps\$logdetVinv - 1/2/sigsq.eps*crossprod(mu, Vcomps.star\$Vinv - Vcomps\$Vinv)%*%mu
logMHratio <- diffliks + diffpriors + negdifflogproposal
logalpha <- min(0,logMHratio)

## return value
acc <- FALSE
if( log(runif(1)) <= logalpha ) {
r <- r.star
delta <- delta.star
lambda <- lambda.star
Vcomps <- Vcomps.star
acc <- TRUE
}
return(list(r=r, lambda=lambda, delta=delta, acc=acc, Vcomps=Vcomps, move.type=move.type))
}

h.update <- function(lambda, Vcomps, sigsq.eps, y, X, beta, r, Z, data.comps) {
if (is.null(Vcomps)) {
Vcomps <- makeVcomps(r = r, lambda = lambda, Z = Z, data.comps = data.comps)
}
if(is.null(Vcomps\$Q)) {
Kpart <- makeKpart(r, Z)
K <- exp(-Kpart)
Vinv <- Vcomps\$Vinv
lambda <- lambda[1] ## in case with random intercept (randint==TRUE), where lambda is 2-dimensional
lamKVinv <- lambda*K%*%Vinv
h.postmean <- lamKVinv%*%(y-X%*%beta)
##h.postvar <- sigsq.eps*lamKVinv
h.postvar <- sigsq.eps*lambda*(K - lamKVinv%*%K)
h.postvar.sqrt <- try(chol(h.postvar), silent=TRUE)
if(inherits(h.postvar.sqrt, "try-error")) {
sigsvd <- svd(h.postvar)
h.postvar.sqrt <- t(sigsvd\$v %*% (t(sigsvd\$u) * sqrt(sigsvd\$d)))
}
hsamp <- h.postmean + crossprod(h.postvar.sqrt, rnorm(length(h.postmean)))
hcomps <- list(hsamp = hsamp)
} else {
h.star.postvar.sqrt <- sqrt(sigsq.eps*lambda)*forwardsolve(t(Vcomps\$cholR), Vcomps\$Q)
h.star.postmean <- lambda[1]*Vcomps\$Q %*% Vcomps\$Rinv %*% Vcomps\$K10 %*% (y - X %*% beta)
hsamp.star <- h.star.postmean + crossprod(h.star.postvar.sqrt, rnorm(length(h.star.postmean)))
hsamp <- t(Vcomps\$K10) %*% Vcomps\$Qinv %*% hsamp.star
hcomps <- list(hsamp = hsamp, hsamp.star = hsamp.star)
}
hcomps
}

newh.update <- function(Z, Znew, Vcomps, lambda, sigsq.eps, r, y, X, beta, data.comps) {

if(is.null(data.comps\$knots)) {
n0 <- nrow(Z)
n1 <- nrow(Znew)
nall <- n0 + n1
# Kpartall <- makeKpart(r, rbind(Z, Znew))
# Kmat <- exp(-Kpartall)
# Kmat0 <- Kmat[1:n0,1:n0 ,drop=FALSE]
# Kmat1 <- Kmat[(n0+1):nall,(n0+1):nall ,drop=FALSE]
# Kmat10 <- Kmat[(n0+1):nall,1:n0 ,drop=FALSE]
Kmat1 <- exp(-makeKpart(r, Znew))
Kmat10 <- exp(-makeKpart(r, Znew, Z))

if(is.null(Vcomps)) {
Vcomps <- makeVcomps(r = r, lambda = lambda, Z = Z, data.comps = data.comps)
}
Vinv <- Vcomps\$Vinv

lamK10Vinv <- lambda[1]*Kmat10 %*% Vinv
Sigma.hnew <- lambda[1]*sigsq.eps*(Kmat1 - lamK10Vinv %*% t(Kmat10))
mu.hnew <- lamK10Vinv %*% (y - X%*%beta)
root.Sigma.hnew <- try(chol(Sigma.hnew), silent=TRUE)
if(inherits(root.Sigma.hnew, "try-error")) {
sigsvd <- svd(Sigma.hnew)
root.Sigma.hnew <- t(sigsvd\$v %*% (t(sigsvd\$u) * sqrt(sigsvd\$d)))
}
hsamp <- mu.hnew + crossprod(root.Sigma.hnew, rnorm(n1))
} else {
n0 <- nrow(data.comps\$knots)
n1 <- nrow(Znew)
nall <- n0 + n1
# Kpartall <- makeKpart(r, rbind(data.comps\$knots, Znew))
# Kmat <- exp(-Kpartall)
# Kmat0 <- Kmat[1:n0,1:n0 ,drop=FALSE]
# Kmat1 <- Kmat[(n0+1):nall,(n0+1):nall ,drop=FALSE]
# Kmat10 <- Kmat[(n0+1):nall,1:n0 ,drop=FALSE]
Kmat10 <- exp(-makeKpart(r, Znew, data.comps\$knots))

if(is.null(Vcomps)) {
Vcomps <- makeVcomps(r = r, lambda = lambda, Z = Z, data.comps = data.comps)
h.star.postvar.sqrt <- sqrt(sigsq.eps*lambda[1])*forwardsolve(t(Vcomps\$cholR), Vcomps\$Q)
h.star.postmean <- lambda[1]*Vcomps\$Q %*% Vcomps\$Rinv %*% Vcomps\$K10 %*% (y - X %*% beta)
Vcomps\$hsamp.star <- h.star.postmean + crossprod(h.star.postvar.sqrt, rnorm(length(h.star.postmean)))
}
hsamp <- Kmat10 %*% Vcomps\$Qinv %*% Vcomps\$hsamp.star
}

hsamp
}

## function to obtain posterior samples of h(znew) from fit of Bayesian kernel machine regression
predz.samps <- function(fit, Znew, verbose = TRUE) {
if(is.null(dim(Znew))) Znew <- matrix(Znew, nrow=1)
if(inherits(Znew, "data.frame")) Znew <- data.matrix(Znew)
Z <- fit\$Z
if(ncol(Z) != ncol(Znew)) {
stop("Znew must have the same number of columns as Z")
}

hnew.samps <- sapply(1:fit\$nsamp, function(s) {
if(s%%(fit\$nsamp/10)==0 & verbose) print(s)
newh.update(Z = Z, Znew = Znew, Vcomps = NULL, lambda = fit\$lambda[s], sigsq.eps = fit\$sigsq.eps[s], r = fit\$r[s,], y = fit\$y, X = fit\$X, beta = fit\$beta[s,], data.comps = fit\$data.comps)
})
rownames(hnew.samps) <- rownames(Znew)
t(hnew.samps)
}

## function to approximate the posterior mean and variance as a function of the estimated tau, lambda, beta, and sigsq.eps
newh.postmean <- function(fit, Znew, sel) {
if(is.null(dim(Znew))) Znew <- matrix(Znew, nrow=1)
if(inherits(Znew, "data.frame")) Znew <- data.matrix(Znew)

Z <- fit\$Z
X <- fit\$X
y <- fit\$y
data.comps <- fit\$data.comps
lambda <- colMeans(fit\$lambda[sel, ,drop = FALSE])
sigsq.eps <- mean(fit\$sigsq.eps[sel])
r <- colMeans(fit\$r[sel,])
beta <- colMeans(fit\$beta[sel, ,drop=FALSE])

if(is.null(data.comps\$knots)) {
n0 <- nrow(Z)
n1 <- nrow(Znew)
nall <- n0 + n1
Kpartall <- makeKpart(r, rbind(Z, Znew))
Kmat <- exp(-Kpartall)
Kmat0 <- Kmat[1:n0,1:n0 ,drop=FALSE]
Kmat1 <- Kmat[(n0+1):nall,(n0+1):nall ,drop=FALSE]
Kmat10 <- Kmat[(n0+1):nall,1:n0 ,drop=FALSE]

Vcomps <- makeVcomps(r = r, lambda = lambda, Z = Z, data.comps = data.comps)
Vinv <- Vcomps\$Vinv

lamK10Vinv <- lambda[1]*Kmat10 %*% Vinv
Sigma.hnew <- lambda[1]*sigsq.eps*(Kmat1 - lamK10Vinv %*% t(Kmat10))
mu.hnew <- lamK10Vinv %*% (y - X%*%beta)
} else {
n0 <- nrow(data.comps\$knots)
n1 <- nrow(Znew)
nall <- n0 + n1
Kpartall <- makeKpart(r, rbind(data.comps\$knots, Znew))
# Kmat <- exp(-Kpartall)
# Kmat0 <- Kmat[1:n0,1:n0 ,drop=FALSE]
# Kmat1 <- Kmat[(n0+1):nall,(n0+1):nall ,drop=FALSE]
# Kmat10 <- Kmat[(n0+1):nall,1:n0 ,drop=FALSE]
Kmat1 <- exp(-makeKpart(r, Znew))
Kmat10 <- exp(-makeKpart(r, Znew, data.comps\$knots))

Vcomps <- makeVcomps(r = r, lambda = lambda, Z = Z, data.comps = data.comps)

Sigma.hnew <- lambda[1]*sigsq.eps*Kmat10 %*% Vcomps\$Rinv %*% t(Kmat10)
mu.hnew <- lambda[1]*Kmat10 %*% Vcomps\$Rinv %*% Vcomps\$K10 %*% (y - X%*%beta)
}

ret <- list(postmean = drop(mu.hnew), postvar = Sigma.hnew)
ret
}
```

## Try the bkmr package in your browser

Any scripts or data that you put into this service are public.

bkmr documentation built on March 28, 2022, 9:11 a.m.