Nothing
#' p-Wasserstein distance projections using simulated annealing
#'
#' @param X Covariate vector
#' @param Y Predictions
#' @param theta Optional matrix of parameters for generating predictions
#' @param power Power of the Wasserstein distance
#' @param force Any covariates to force into the model?
#' @param model.size Maximum number of coefficients
#' @param nvars The number of variables to explore. Should be an integer vector of model sizes. Default is NULL which will explore all models from 1 to `model.size`.
#' @param maxit Maximum number of iterations
#' @param temps Number of temperatures
#' @param max.time Maximum time in seconds to run
#' @param const Maximum value for simulated annealing distance
#' @param proposal Proposal function. There is a default method but can provide your own with parameters `xty`, `cur`, `idx`, `force`, `d`, `method`
#' @param options Options for simulated annealing
#' @param display.progress Whether to display solver progress. TRUE or FALSE. Default is FALSE.
#' @param parallel A [foreach::foreach()] backend
#' @param calc.theta Should the model save the linear coefficients? TRUE or FALSE. Default is TRUE
#' @param xtx precomputed crossproduct \code{crossprod(X,X)}
#' @param xty precomputed \code{crossprod(X, Y)}
#'
#' @return An object of class `WpProj`
#' @keywords internal
#'
# @examples
# if(rlang::is_installed("stats")) {
# n <- 128
# p <- 10
# s <- 99
# x <- matrix( stats::rnorm( p * n ), nrow = n, ncol = p )
# x_ <- t(x)
# beta <- (1:10)/10
# y <- x %*% beta + rnorm(n)
# post_beta <- matrix(beta, nrow=p, ncol=s) + rnorm(p*s, 0, 0.1)
# post_mu <- x %*% post_beta
# sv <- WPSA(X=x, Y=t(post_mu), theta = t(post_beta),
# force = 1, p = 2, nvars = 2:(p-1),
# maxit=5, temps = 5,
# max.time = 100,
# options = list(method = c("selection.variable"),
# energy.distribution = "boltzman",
# transport.method = "exact",
# cooling.schedule = "Geman-Geman",
# proposal.method = "covariance")
# )
# }
WPSA <- function(X, Y=NULL, theta,
power = 2, force = NULL,
model.size = 3,
nvars = NULL,
# groups = NULL,
maxit=1, temps = 1000,
max.time = 3600, const = NULL,
proposal = proposal.fun,
options = list(method = c("selection.variable","scale","projection"),
transport.method = transport_options(),
energy.distribution = "boltzman",
cooling.schedule = "Geman-Geman",
proposal.method = "covariance",
epsilon = 0.05,
OTmaxit = 100),
display.progress = FALSE,
parallel = NULL,
calc.theta = TRUE,
xtx = NULL,
xty = NULL,
...
)
{
# groups <- as.integer(groups)
# unique.groups <- as.integer(unique.groups)
this.call <- as.list(match.call()[-1])
if(!is.matrix(X)) X <- as.matrix(X)
if(!is.matrix(theta)) theta <- as.matrix(theta)
time.exceed <- function(max.time, start.time) {
elapsed <- proc.time() - start.time
return(elapsed[3] > max.time)
}
get.theta <- calc.theta
iter <- maxit
p <- ground_p <- power
if (is.null(model.size)) {
model.size <- ncol(X)-1L
}
if (is.null(nvars)) {
nvars <- 1:model.size
}
names(options) <- match.arg(names(options), c("energy.distribution",
"method",
"transport.method",
"cooling.schedule",
"proposal.method"),
several.ok = TRUE)
energy.dist <- match.arg(options$energy.distribution, c("boltzman","bose-einstein"))
meth <- match.arg(options$method, c("selection.variable","scale","projection"))
transmeth <- match.arg(options$transport.method, transport_options() )
cool.sched <- match.arg(options$cooling.schedule, c("Geman-Geman","exponential"))
obs.direction <- "colwise"
prop.meth <- match.arg(options$proposal.method, c("covariance","random"))
epsilon <- options$epsilon
OTmaxit <- options$OTmaxit
if(is.null(epsilon)) epsilon <- 0.05
if(is.null(OTmaxit)) OTmaxit <- 100
if(ncol(theta) == ncol(X)){
theta_ <- t(theta)
} else {
theta_ <- theta
}
#transpose X
X_ <- t(X)
if(is.null(Y)) {
same <- TRUE
Y_ <- crossprod(X_, theta_)
} else{
Y <- as.matrix(Y)
if(!any(dim(Y) %in% dim(X_))) stop("dimensions of Y must match X")
if(!is.matrix(Y)) Y <- as.matrix(Y)
if(nrow(Y) != ncol(X_)){
Y_ <- t(Y)
} else{
Y_ <- Y
}
same <- FALSE
if(all(Y_==crossprod(X_, theta_))) same <- TRUE
}
OToptions <- list(same = as.logical(same),
method = as.character(meth),
transport.method = as.character(transmeth),
epsilon = as.double(epsilon),
niter = as.integer(OTmaxit))
d <- nrow(X_)
n <- ncol(X_)
if (is.null(xtx) | is.null(xty) ) {
augmat <- sufficientStatistics(X, Y_, t(theta_), OToptions)
xtx <- augmat$XtX
xty <- augmat$XtY
}
if(!is.null(parallel)){
if(!inherits(parallel, "cluster")) {
stop("parallel must be a registered cluster backend")
}
doParallel::registerDoParallel(parallel)
display.progress <- FALSE
} else{
foreach::registerDoSEQ()
}
message <- "completed"
start.time <- proc.time()
idx <- 1:d
if (!is.null(force) & is.null(const)) {
beta_temp <- calc.beta(xtx=xtx, xty=xty, force, meth, OToptions = OToptions, x=X_, theta_, Y_, niter=500)
if(meth != "projection"){
temp_mu <- selVarMeanGen(X_, theta_, beta_temp)
} else {
temp_mu <- crossprod(X_, beta_temp)
}
const <- WpProj::wasserstein(temp_mu,
Y_, p, ground_p,
obs.direction, transmeth)
} else if (is.null(const)) {
if (meth == "projection") {
mu_calc <- function(X, theta, beta) {
return(crossprod(X, beta))
}
} else {
mu_calc <- function(X, theta, beta) {
return(selVarMeanGen(X, theta, beta))
}
}
w2s <- foreach::foreach(i = idx, .combine = "unlist") %dopar%
{
beta_temp <- calc.beta(xtx=xtx, xty=xty, i, meth, OToptions = OToptions, x=X_, theta_, Y_, niter=500)
temp_mu <- mu_calc(X_, theta_, beta_temp)
return(WpProj::wasserstein(temp_mu, Y_, p, ground_p, obs.direction, transmeth) )
}
const <- max(w2s)
}
if(length(nvars) == 1) {
if(length(force) == max(nvars)) {
# warning("Number of variables forced into model the same as desired model size")
old.dist <- best.dist <- dist.fun(X_, force, theta_, Y_, xtx, xty, OToptions, obs.direction, p, ground_p)
final.beta <- best.beta <- calc.beta(xtx,xty,force, meth, OToptions = OToptions,
x=X_, theta= theta_, Y = Y_, niter = 500)
output <- list(final = list(index = force,
distance = old.dist,
beta = final.beta),
optimal = list(index = force,
distance = best.dist,
beta = best.beta),
force = force,
history = old.dist,
temps = NULL,
accept.pct = NULL,
call = this.call,
message = message,
method = meth)
class(output) <- c("WpProj", "annealing")
if(get.theta == TRUE){
extract <- extractTheta(output, theta_)
output$theta <- extract$theta
output$nzero <- extract$nzero
output$eta <- lapply(output$theta, function(tt) crossprod(X_, tt))
} else {
output$theta <- NULL
output$nzero <- nvars
output$eta <- NULL
}
return(output)
}
if (iter <= 0) {
stop("iter should be greater than 0")
}
if (temps <= 0) {
stop("temps should be greater than 0")
}
# tY <- t(Y)
# if (grepl("univariate", transmeth) ) {
# obs.direction <- "rowwise"
# # X_ <- t(X_)
# # Y_ <- t(Y_)
# } else {
obs.direction <- "colwise"
# }
if(is.null(proposal)) proposal <- proposal.fun
# wt <- pseudo.obs/(n+pseudo.obs)
# theta_norm <- rowMeans(theta_^2)
# w_xtx <- wt * diag(theta_norm) + (1-wt) * xtx
# w_xty <- wt * theta_norm + (1-wt) * xty
if(d == nvars) {
warning("Model is same size as number of parameters")
old.dist <- best.dist <- dist.fun(X_, 1:d, theta_, Y_, xtx, xty, OToptions, obs.direction, p, ground_p)
final.beta <- best.beta <- calc.beta(xtx,xty,1:d, meth, OToptions = OToptions,
x=X_, theta= theta_, Y = Y_, niter = 500)
# transport.method
output <- list(final = list(index = 1:d,
distance = old.dist,
beta = final.beta),
optimal = list(index = 1:d,
distance = best.dist,
beta = best.beta),
force = force,
history = old.dist,
temps = NULL,
accept.pct = NULL,
call = this.call,
message = message,
method = meth)
class(output) <- c("WpProj", "annealing")
if(get.theta == TRUE){
extract <- extractTheta(output, theta_)
output$theta <- extract$theta
output$nzero <- extract$nzero
output$eta <- lapply(output$theta, function(tt) crossprod(X_, tt))
} else {
output$theta <- NULL
output$nzero <- nvars
output$eta <- NULL
}
return(output)
}
temp_sched <- coolingSchedule(1:temps, const, cool.sched)
previous <- const
active.set <- idx[!(idx %in% force)]
mod.size <- nvars - length(force)
active.idx <- sort(c(force, sample(active.set, mod.size)))
old.dist <- dist.fun(X_, active.idx, theta_, Y_, xtx, xty, OToptions, obs.direction, p,ground_p)
dist.hist <- matrix(NA, nrow=iter, ncol=temps)
colnames(dist.hist) <- as.character(temp_sched)
rownames(dist.hist) <- as.character(1:iter)
best.idx <- active.idx
best.dist <- old.dist
acpct <- rep(NA, temps) #accept percentage
accept_cnt <- zero_counter <- 0
max.save <- min(c(choose(d, nvars), length(temp_sched)*iter, 1e6))
saved.dist <- rep(Inf, max.save)
idx.key <- rep(-Inf, max.save)
# pot.comb <- rbind(force, combn(idx[!(idx %in% force)], nvars-length(force)))
# names(saved.dist) <- apply(pot.comb,2,paste, collapse="")
if(display.progress) pb <- utils::txtProgressBar(min = 0, max = length(temp_sched), style = 3)
for (i in seq_along(temp_sched) ){
# print(i)
if(display.progress) utils::setTxtProgressBar(pb, i)
for (j in 1:iter) {
#get proposal
proposed <- proposal(xty, active.idx, idx, force, d, prop.meth)
# calculate proposed distance
idx.switch <- proposed$proposal
idx.name <- calc_key(idx.switch$active)
if( !(idx.name %in% idx.key )) {
proposed.dist <- dist.fun(X_, idx.switch$active, theta_, Y_, xtx, xty,OToptions, obs.direction, p, ground_p)
if(proposed.dist < saved.dist[max.save]){
# names(proposed.dist) <- idx.name
ords <- order(c(saved.dist, proposed.dist))
saved.dist <- c(saved.dist, proposed.dist)[ords[1:max.save]]
idx.key <- c(idx.key, idx.name)[ords[1:max.save]]
}
# saved.dist[idx.name] <- proposed.dist
} else {
proposed.dist <- saved.dist[idx.key == idx.name]
}
# get transition probs
log.proposal.probs <- log(proposed$probs)
#calculate MH prob
# if(all(idx.switch$active %in% c(1,3,6,7))) browser()
log.MHprob <- energy.fun(temp_sched[i], proposed.dist, old.dist, energy.dist) + log.proposal.probs[1] - log.proposal.probs[2]
# if(is.nan(log.MHprob)) browser()
# if(is.na(log.MHprob)) browser()
# if(all(idx.switch$active %in% c(1,6,7))) browser()
if(log(stats::runif(1)) < log.MHprob){
active.idx <- idx.switch$active
old.dist <- proposed.dist
accept_cnt <- accept_cnt + 1
if(old.dist < best.dist) {
best.dist <- old.dist
best.idx <- active.idx
}
}
# if(length(old.dist)>1) browser()
dist.hist[j,i] <- old.dist
}
acpct[i] <- accept_cnt/iter
if(accept_cnt == 0) {
zero_counter <- zero_counter + 1
} else {
zero_counter <- 0
}
accept_cnt <- 0
if(zero_counter > 2) {
temp_sched[is.na(acpct)] <- NA
break
}
if(j %% 10) {
if(time.exceed(max.time, start.time)) {
messages <- "Hit max time. Did not reach final temperature"
break
}
}
}
if(display.progress) close(pb)
final.beta <- calc.beta(xtx, xty, active.idx, meth, OToptions = OToptions, X_, theta_)
best.beta <- calc.beta(xtx, xty, best.idx, meth, OToptions = OToptions, X_, theta_)
output <- list(final = list(index = active.idx,
distance = old.dist,
beta = final.beta),
optimal = list(index = best.idx,
distance = best.dist,
beta = best.beta),
force = force,
history = dist.hist,
temps = temp_sched,
accept.pct = acpct,
call = this.call,
message = message,
method = meth)
}
else if (length(nvars) > 1) {
if(display.progress) pb <- utils::txtProgressBar(min = 0, max = length(nvars), style = 3)
# vector("list", length(nvars))
out <- foreach::foreach( i = seq_along(nvars),
.inorder = FALSE) %dorng% {
mm <- nvars[i]
if(display.progress) utils::setTxtProgressBar(pb, i)
if ( time.exceed(max.time, start.time) ) { # doesn't work in parallel
# message <- "Hit max time exploring model sizes"
# if (i < length(out)) out[(i + 1):length(out)] <- NULL
return(NULL)
}
result <- WPSA(X=X, Y=Y_, theta=theta_,
force = force, power = p, nvars = mm,
maxit=iter, temps = temps,
max.time = max.time/length(nvars), const = const,
proposal = proposal,
options = options,
display.progress = FALSE,
parallel = NULL,
calc.theta = FALSE,
xtx = xtx,
xty = xty)
return(result)
}
if( time.exceed(max.time, start.time) ) {
message <- "Hit max time exploring model sizes"
}
if(display.progress) close(pb)
# final <- lapply(out, function(o) o$final)
# optimal <- lapply(out, function(o) o$optimal)
# dist.hist <- lapply(out, function(o) o$history)
# temp_sched <- lapply(out, function(o) o$temps)
# acpct <- lapply(out, function(o) o$accept.pct)
# rm(out)
call <- formals(WPSA)
call[names(this.call)] <- this.call
output <-list(final = lapply(out, function(o) o$final),
optimal = lapply(out, function(o) o$optimal),
force = force,
history = lapply(out, function(o) o$history),
temps = lapply(out, function(o) o$temps),
accept.pct = lapply(out, function(o) o$accept.pct),
call = this.call,
message = message,
method = meth)
rm(out)
}
class(output) <- c("WpProj","annealing")
if(get.theta == TRUE){
extract <- extractTheta(output, theta_)
output$theta <- extract$theta
output$nzero <- extract$nzero
output$eta <- lapply(output$theta, function(tt) X %*% tt)
} else {
output$theta <- NULL
output$nzero <- nvars
output$eta <- NULL
}
return(output)
}
dist.fun <- function(X=NULL, active.idx=NULL, theta = NULL, Y = NULL, xtx=NULL, xty = NULL, OToptions = NULL, obs.direction ="colwise",
p = 2, ground_p = 2){
if(is.null(OToptions)) {
OToptions <- list(same = TRUE,
method = "selection.variable",
transport.method = "hilbert",
epsilon = 0.05,
niter = 100)
} else {
if(is.null(OToptions$same)) OToptions$same <- FALSE
if(is.null(OToptions$method)) OToptions$method <- "selection.variable"
if(is.null(OToptions$transport.method)) OToptions$transport.method <- "hilbert"
if(is.null(OToptions$epsilon)) OToptions$epsilon <- 0.05
if(is.null(OToptions$niter)) OToptions$niter <- 100
}
beta <- calc.beta(xtx, xty, active.idx = active.idx,
method = OToptions$method, OToptions = OToptions,
x = X, theta =theta, Y = Y)
if(OToptions$method == "projection") {
mu <- crossprod(X, beta)
# } else if (method == "selection.variable") {
# # wp <- WpProj::wasserstein(mu, Y, p = p, ground_p = p, "colwise", "exact")
# mu <- crossprod(X[active.idx, , drop = FALSE], theta[active.idx,, drop = FALSE])
} else {
mu <- selVarMeanGen(X,theta,beta)
}
shortcut <- (OToptions$method == "projection") & (OToptions$transport.method == "exact")
wp <- if(shortcut){
(((sum((mu - Y)^ground_p)^(1/ground_p))^p)^(1/p))/ncol(Y)
} else {
WpProj::wasserstein(mu, Y, p = p, ground_p = ground_p, obs.direction, OToptions$transport.method)
}
return(wp)
}
energy.fun <- function(temp, proposed.dist, old.dist, method="boltzman") {
dist <- proposed.dist - old.dist
if(method == "boltzman"){
return(-dist/temp)
}
if(method == "bose-einstein") {
return(log(expm1(old.dist/temp)) - log(expm1(proposed.dist/temp)))
# return(log(expm1(dist/temp)))
}
}
proposal.fun <- function(xty, cur, idx, force, d, method = "covariance") {
forced <- idx %in% force
choice.new <- !( forced | (idx %in% cur) )
choice.old <- cur[!(cur %in% force)]
not.act.idx <- idx[choice.new]
act.idx <- idx[choice.old]
n.new <- length(not.act.idx)
n.old <- length(act.idx) + sum(forced)
if (method == "covariance") {
# xty_new <- xty_old <-
potentials <- abs(xty)
# xty_new[choice] <- 0
# xty_old[!choice | forced] <- 0
# if(!all(is.finite(probs_rem))) browser()
# if(!all(is.finite(probs_add))) browser()
# if(length(probs_rem)==0) browser()
}
if (method == "random") {
potentials <- rep(1, n.old + n.new)
# probs_rem <- probs_rev_rem <- rep(1/n.old, n.old)
# probs_add <- probs_rev_add <- rep(1/n.new, n.new)
}
probs_add <- potentials[not.act.idx]/sum(potentials[not.act.idx])
probs_rem <- 1-potentials[act.idx]/sum(potentials[act.idx])
if(all(probs_rem == 0)) probs_rem <- 1/n.old
if(all(probs_add == 0)) probs_add <- 1/n.new
proposed.removal <- which(stats::rmultinom(n=1, size=1, prob=probs_rem)==1)
proposed.addition <- which(stats::rmultinom(n=1, size=1, prob=probs_add)==1)
new.act <- c(not.act.idx[proposed.addition], act.idx[-(proposed.removal)])
new.not.act <- c(act.idx[proposed.removal], not.act.idx[-(proposed.addition)])
probs_rev_rem <- 1-potentials[new.act]/sum(potentials[new.act])
probs_rev_add <- potentials[new.not.act]/sum(potentials[new.not.act])
if(all(probs_rev_rem == 0)) probs_rev_rem <- rep(1/n.old, n.old)
if(all(probs_rev_add == 0)) probs_rev_add <-rep(1/n.new, n.new)
new.probs <- probs_add[proposed.addition] * probs_rem[proposed.removal]
old.probs <- probs_rev_rem[1] * probs_rev_add[1] #make sure this makes sense
new.act.idx <- sort(c(new.act, force)) #c(new.act, force) #
new.not.act.idx <- sort(new.not.act) # new.not.act #
return(list(proposal=list(active=new.act.idx, not.active=new.not.act.idx),
probs=c(new.probs, old.probs)))
}
coolingSchedule <- function(t,c, method="Geman-Geman") {
if(method == "Geman-Geman"){
return(c/log(t+1))
}
if(method == "exponential") {
Ti <- c/log(2)
Tf <- c/log(.Machine$double.xmax)
alpha <- -(log(Tf) - log(Ti))/length(t)
return(exp(-(t-1) *alpha) * Ti)
}
}
calc_key <- function(b){
a <- log(sort(b) + 1)
s <- length(a)
val <- 2^a + pi * a
out <- sum(val)
return(out)
}
Any scripts or data that you put into this service are public.
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.