This repository has been archived by the owner on Jun 9, 2023. It is now read-only.
/
gromov_hausdorff.R
176 lines (158 loc) · 8.56 KB
/
gromov_hausdorff.R
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
#' gromov hausdorff
#' @description This function implements the discrete, extended version of gromov hausdorff distance proposed by
#' Memoli in section 7 of (1) for comparing two measure-metric spaces.
#' @param d_X either a dist object, or (n x n) numeric metric distance matrix over \code{X}.
#' @param d_Y either a dist object, or (m x m) numeric metric distance matrix over \code{Y}.
#' @param mu_X (n-length) numeric vector giving the probability measure at each point of \code{X}.
#' @param mu_Y (m-length) numeric vector giving the probability measure at each point of \code{Y}.
#' @param p The norm to compute.
#' @param flb_solver Which LOP ROI plugin solver to use to solve the lower bound approximation.
#' @param flb_only If only crude approximation is needed, the lower bound can be returned. See details.
#' @param options Options to \code{\link[nloptr]{auglag}}. Ignored if \code{flb_only=TRUE}.
#' @param control Control options to pass \code{\link[nloptr]{auglag}}'s control.
#' @param return_optimizer Whether to return the optimizer, or just the results.
## @param roi_opt Outer options to pass to ROI_solve
#' @references 1. Mémoli, Facundo. "On the use of Gromov-Hausdorff distances for shape comparison." (2007).
#' @details This function provides an implementation of the computational objective \emph{(Pp)} listed in Section 7 of (1), which
#' calculates an Lp approximition of Gromov-Hausdorff (i.e. wasserstein relaxation) distance between two measure-metric spaces. The first lower bound (FLB)
#' is first computed via an LOP, and then the solution is used as the initialization point for the corresponding quadratic optimization
#' routine. This code utilizes the \emph{R Optimization Infrastructure} (\pkg{ROI}), and relies on having available at least one plugin
#' available for to solve the initial lower bound (a linear program). If more than the lower bound is requested, the optimization
#' is cast as a generic nonlinear optimization.
#' @return If \code{flb_only} is set to TRUE, a list with components:
#' \itemize{
#' \item gh The gromov-wasserstein distance.
#' \item mu The result of the optimization.
#' \item correspondences A surjective matching from X to Y and from Y to X.
#' }
#' If \code{flb_only} is set to FALSE, a list of both the LOP and QOP optimization results are returned in the
#' above format.
#'
#' \strong{NOTE:} The quadratic optimization problem (QOP) is non-convex, and requires solving for (n^2 x m^2) variables,
#' which may be very computationally expensive to run. The FLB is an LOP of the same size.
#' @export
gromov_hausdorff <- function(d_X, d_Y, mu_X, mu_Y, p = 1,
flb_solver="glpk", flb_only=TRUE,
options = list(localsolver="COBYLA", localtol = 1e-7),
control = list(maxeval=300),
return_optimizer = FALSE){
roi_installed <- requireNamespace("ROI", quietly = TRUE)
if (!roi_installed){ stop("This function requires the R Optimization Infrastructure (ROI) package to be installed.") }
## Preprocessing
stopifnot((sum(mu_X) - 1.0) < sqrt(.Machine$double.eps))
stopifnot((sum(mu_Y) - 1.0) < sqrt(.Machine$double.eps))
if (is.matrix(d_X)) { stopifnot(dim(d_X)[[1]] == dim(d_X)[[2]]) }
else { stopifnot(is(d_X, "dist")); d_X <- as.matrix(d_X) }
if (is.matrix(d_Y)) { stopifnot(dim(d_Y)[[1]] == dim(d_Y)[[2]]) }
else { stopifnot(is(d_Y, "dist")); d_Y <- as.matrix(d_Y) }
## Wrap in requireNamesapce to dismiss dependency warnings
if ( requireNamespace("ROI", quietly = TRUE) ){
## Useful variables
{ n_x <- nrow(d_X); n_y <- nrow(d_Y) }
idx <- matrix(seq(n_x * n_y), nrow = n_x, ncol = n_y)
## Make constraints
A <- gh_make_A(idx-1L)
constraints <- ROI::L_constraint(L = A, dir = rep("==", n_x + n_y), rhs = c(mu_X, mu_Y))
## Make objective for FLB
s <- function(d_X, p, lambda){ apply(d_X, 1, function(d_i){ sum((d_i^p) * lambda)^(1/p) }) }
sp_X <- s(d_X, p = 1, lambda = mu_X)
sp_Y <- s(d_Y, p = 1, lambda = mu_Y)
flb_objective <- ROI::L_objective(as.vector(outer(1L:n_x, 1L:n_y, FUN = Vectorize(function(i,j){ abs(sp_X[i] - sp_Y[j]) }))))
## Bounds on mu
mu_bnds <- ROI::V_bound(li = seq(n_x*n_y), ui = seq(n_x*n_y), lb = rep(0, n_x*n_y), ub = rep(1,n_x*n_y))
## The lower-bound -- as a LOP
lop <- ROI::OP(objective = flb_objective, constraints = constraints, bounds = mu_bnds)
## Need to have a solver available
available_solvers <- ROI::ROI_applicable_solvers(lop)
if (is.null(available_solvers)){
stop("No applicable solvers found with ROI! The FLB is a linear optimization.")
}
if (!missing(flb_solver)){ stopifnot(flb_solver %in% available_solvers) }
## Solve the first lower bound
flb <- ROI::ROI_solve(lop, solver = flb_solver)
if (flb_only) { return(wrap_roi_solution(flb, idx, return_optimizer)) }
else {
nloptr_installed <- requireNamespace("nloptr", quietly = TRUE)
if (!nloptr_installed){ stop("The quadratic optimization requires 'nloptr' to be installed.") }
## Make the Q matrix of distances
Q <- gh_make_Q(d_X, d_Y)
## Setup objective function(s), gradients, etc.
N <- n_x * n_y
f <- function(mu){
as.vector({ structure(mu, dim=c(1L, N)) %*% Q %*% structure(mu, dim=c(N, 1L)) })
}
grad_f <- function(mu){ Q %*% matrix(mu, ncol = 1) }
heq <- function(mu){ (A %*% mu) - c(mu_X, mu_Y) } # rowSums(relist(mu, idx)) - c(mu_X, mu_Y)
# heqjac <- function(x) { nloptr::nl.jacobian(x, heq) }
## Minor fix to the FLB bounds to make optimization more stable. This usually isn't necessary.
flb$solution[flb$solution < 0] <- 0.0
flb$solution[flb$solution > 1] <- 1.0
## Set solver options
{ lb <- rep(0, n_x*n_y); ub <- rep(1, n_x*n_y) }
default_opts <- list(x0=flb$solution, fn = f, lower = lb, upper = ub, heq = heq)
default_opts$control <- modifyList(list(maxeval=300, stopval=flb$objval), control)
options <- modifyList(default_opts, val = options)
## Attach gradient function if COBYLA isn't used
if (options$localsolver %in% c("LBFGS", "MMA", "SLSQP")){
options$gr <- grad_f
}
## Use nloptr augmented lagrangian approach
al_res <- do.call(nloptr::auglag, options)
## Return results
return(list(lop_res = wrap_roi_solution(flb, idx, return_optimizer),
qop_res = wrap_nloptr_solution(al_res, idx, return_optimizer)))
}
}
}
wrap_nloptr_solution <- function(res, idx, return_opt){
tmp <- list(
gh = res$value,
mu = res$par,
correspondences = list(
xy = apply(utils::relist(res$par, idx), 1, which.max),
yx = apply(utils::relist(res$par, idx), 2, which.max)
)
)
if (return_opt){ tmp$optim_status <- res[c("global_solver", "local_solver", "convergence", "message")] }
return(tmp)
}
wrap_roi_solution <- function(res, idx, return_opt){
tmp <- list(
gh = res$objval,
mu = res$solution,
correspondences = list(
xy = apply(utils::relist(res$solution, idx), 1, which.max),
yx = apply(utils::relist(res$solution, idx), 2, which.max)
)
)
if (return_opt){ tmp$optim_status <- res[c("status", "message")] }
return(tmp)
}
## Extra code
# A <- outer(X = seq(n_x + n_y), Y = seq(n_x * n_y), FUN = Vectorize(function(i, j){
# if (i <= n_x){ ifelse(j %in% idx[i,], 1L, 0L) }
# else { ifelse(j %in% idx[,(i-n_x)], 1L, 0L) }
# }))
## Make the control
# control <- list(
# algorithm = "NLOPT_LN_AUGLAG_EQ",
# maxeval = 10L,
# local_opts = list(algorithm = "NLOPT_LD_SLSQP", xtol_abs = 1e-4, stopval = 0, heq = heq, heqjac = heqjac),
# print_level = 0,
# check_derivatives = FALSE
# )
# The indices to use to compute the distances, filling Q row-wise
# rhs_idx <- cbind(as.vector(row(idx)), as.vector(col(idx)))
# master_idx <- do.call(rbind, lapply(seq(nrow(rhs_idx)), function(i){
# cbind( t(replicate(n = nrow(rhs_idx), rhs_idx[i,])), rhs_idx)
# }))
## The distances to use
# Q <- matrix(apply(master_idx, 1, function(ijkl){
# { i <- ijkl[1]; j <- ijkl[2]; k <- ijkl[3]; l <- ijkl[4] }
# abs(d_X[i,k] - d_Y[j, l])
# }), nrow = n_x*n_y, ncol = n_x*n_y, byrow = TRUE)
## Solve the QOP -- using general non-linear optimizer, since there is no non-convex solvers available
## -- Unfortunately, the interface isn't quite as flexible as one would hope
# q_obj <- ROI::Q_objective(Q = Q)
# qop <- ROI::OP(objective = q_obj, constraints = constraints, bounds = mu_bnds)
# qop_res <- ROI::ROI_solve(qop, solver = "nloptr", x0 = flb$solution, control = control)