MBM: optional GP-based M sampling on InfiniteGDP models

Project.toml

@@ -8,12 +8,15 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [weakdeps]
+AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 
 [extensions]
+AbstractGPsDisjunctiveProgramming = "AbstractGPs"
 InfiniteDisjunctiveProgramming = "InfiniteOpt"
 
 [compat]
+AbstractGPs = "0.5"
 Aqua = "0.8"
 JuMP = "1.18"
 Reexport = "1"
@@ -30,4 +33,4 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Juniper = "2ddba703-00a4-53a7-87a5-e8b9971dde84"
 
 [targets]
-test = ["Aqua", "HiGHS", "Test", "Juniper", "Ipopt", "InfiniteOpt"]
+test = ["Aqua", "HiGHS", "Test", "Juniper", "Ipopt", "InfiniteOpt", "AbstractGPs"]

README.md

@@ -176,6 +176,7 @@ The following reformulation methods are currently supported:
 
     - `optimizer`: Optimizer to use when solving subproblems to determine M values. This is a required value.
     - `default_M`: Default big-M value to use if no big-M is specified for a logical variable (1e9).
+    - `sampler`: M-value sampler for infinite models. Default: `nothing`, which solves an M subproblem at every support. Pass a `GPSampler` to instead solve a subset of the supports and fill the rest with a conservative Gaussian-process estimate. Ignored for finite models.
 
 5. [P-Split](https://arxiv.org/abs/2202.05198): This method reformulates each disjunct constraint into P constraints, each with a partitioned group defined by the user. This method requires that terms in the constraint be convex additively seperable with respect to each variable. The `PSplit` struct is created with the following required arguments:
 
@@ -223,6 +224,8 @@ optimize!(model, gdp_method = Hull())
 value(W)
 ```
 
+When the `MBM` reformulation is used on an infinite model, an M subproblem is solved at every support by default. Loading [AbstractGPs.jl](https://github.com/JuliaGaussianProcesses/AbstractGPs.jl) (`using AbstractGPs`) enables an additional extension that instead solves M at a subset of the supports and fills the rest with a conservative Gaussian-process estimate, which can substantially reduce the number of subproblem solves. To opt in, pass a `GPSampler` via the `sampler` keyword, e.g. `MBM(optimizer, sampler = GPSampler())` (squared exponential kernel with the lengthscale selected by marginal likelihood), `GPSampler(Matern52Kernel())` (a custom kernel, lengthscale selected the same way), or `GPSampler(Matern52Kernel(), lengthscales = [0.2])` (lengthscale pinned). The filled values are heuristic upper estimates rather than certificates; see the `GPSampler` docstring for the tuning keywords (`kappa`, `budget`, `detect_uniform_M`, `lengthscales`, `jitter`, `seeds`).
+
 ## Release Notes
 
 Prior to `v0.4.0`, the package did not leverage the JuMP extension capabilities and was not as robust. For these earlier releases, refer to [Perez, Joshi, and Grossmann, 2023](https://arxiv.org/abs/2304.10492v1) and the following [JuliaCon 2022 Talk](https://www.youtube.com/watch?v=AMIrgTTfUkI).

ext/AbstractGPsDisjunctiveProgramming.jl

@@ -0,0 +1,119 @@
+module AbstractGPsDisjunctiveProgramming
+
+import AbstractGPs
+import AbstractGPs.KernelFunctions
+import DisjunctiveProgramming as DP
+
+################################################################################
+#                                 GP FITTING
+################################################################################
+# Normalized to [0, 1]^d so one lengthscale works across dimensions
+function _support_coords(
+    grids::Tuple{Vararg{Vector{Float64}}}
+    )::Vector{Vector{Float64}}
+    indices = CartesianIndices(length.(grids))
+    mins = [minimum(g) for g in grids]
+    ranges = [max(maximum(g) - minimum(g), eps()) for g in grids]
+    return [[(grids[d][I[d]] - mins[d]) / ranges[d]
+             for d in 1:length(grids)] for I in vec(indices)]
+end
+
+# Lengthscale selected by marginal likelihood over the candidates
+function _fit_posterior(
+    kernel::KernelFunctions.Kernel,
+    X::Vector{Vector{Float64}},
+    y::Vector{Float64},
+    sampler::DP.GPSampler
+    )
+    best_posterior, best_log_prob = nothing, -Inf
+    for lengthscale in sampler.lengthscales
+        scaled_kernel = KernelFunctions.with_lengthscale(
+            kernel, lengthscale)
+        finite_gp = AbstractGPs.GP(scaled_kernel)(X, sampler.jitter)
+        log_prob = AbstractGPs.logpdf(finite_gp, y)
+        if log_prob > best_log_prob
+            best_posterior = AbstractGPs.posterior(finite_gp, y)
+            best_log_prob = log_prob
+        end
+    end
+    return best_posterior
+end
+
+function _mean_sd(
+    sampler::DP.GPSampler,
+    X::Vector{Vector{Float64}},
+    solved::Dict{Int, Float64}
+    )
+    solved_indices = collect(keys(solved))
+    y = [solved[i] for i in solved_indices]
+    y_mean = sum(y) / length(y)
+    # floored so near-equal solved values still cushion the filled ones
+    y_scale = max(sqrt(sum(abs2, y .- y_mean) / max(length(y) - 1, 1)),
+        1e-2 * abs(y_mean), 1e-8)
+    kernel = something(sampler.kernel,
+        KernelFunctions.SqExponentialKernel())
+    posterior = _fit_posterior(
+        kernel, X[solved_indices], (y .- y_mean) ./ y_scale, sampler)
+    posterior_mean = AbstractGPs.mean(posterior, X)
+    posterior_var = max.(AbstractGPs.var(posterior, X), 0.0)
+    return posterior_mean .* y_scale .+ y_mean,
+        sqrt.(posterior_var) .* y_scale
+end
+
+################################################################################
+#                              M VALUE SAMPLING
+################################################################################
+# Solve M at max-UCB selected supports, fill the rest with the bound
+function DP.sample_M_values(
+    sampler::DP.GPSampler,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids::Function
+    )
+    indices = collect(CartesianIndices(objectives))
+    n = length(indices)
+    solved = Dict{Int, Float64}()
+    solve_at(index::Int) = begin
+        M_val = DP.raw_M(sub, objectives[indices[index]], method)
+        M_val === nothing && return false
+        solved[index] = M_val
+        return true
+    end
+    # an evenly spaced seed count, or user-given seed fractions
+    fractions = sampler.seeds isa Int ?
+        range(0, 1, length = sampler.seeds) : sampler.seeds
+    for index in unique(1 .+ round.(Int, fractions .* (n - 1)))
+        solve_at(index) || return nothing
+    end
+    if sampler.detect_uniform_M
+        # a uniform M needs no fit, and so no support grid either
+        probes = collect(values(solved))
+        all(==(first(probes)), probes) && return first(probes)
+    end
+    budget = min(ceil(Int, sampler.budget * n), n)
+    X = _support_coords(support_grids())
+    while length(solved) < budget
+        means, sds = _mean_sd(sampler, X, solved)
+        acquisition = means .+ sampler.kappa .* sds
+        for index in keys(solved)
+            acquisition[index] = -Inf
+        end
+        solve_at(argmax(acquisition)) || return nothing
+    end
+    M_vals = Array{Float64}(undef, size(objectives))
+    if length(solved) == n # nothing left to estimate
+        for (index, I) in enumerate(indices)
+            M_vals[I] = solved[index]
+        end
+        return M_vals
+    end
+    means, sds = _mean_sd(sampler, X, solved)
+    for (index, I) in enumerate(indices) # exact M values are nonnegative
+        M_vals[I] = get(solved, index,
+            max(means[index] + sampler.kappa * sds[index], 0.0))
+    end
+    return M_vals
+end
+
+end

ext/InfiniteDisjunctiveProgramming.jl

@@ -38,6 +38,67 @@ function DP.requires_disaggregation(vref::InfiniteOpt.GeneralVariableRef)
     return !_is_parameter(vref)
 end
 
+# Bound info for parameter refs in disjunct constraints: parameter
+# functions report their support extrema, finite parameters a point,
+# other parameters (-Inf, Inf). Returns nothing for variables.
+function _parameter_bound_info(
+    vref::InfiniteOpt.GeneralVariableRef
+    )::Union{Nothing, Tuple{Float64, Float64}}
+    _is_parameter(vref) || return nothing
+    dvref = InfiniteOpt.dispatch_variable_ref(vref)
+    if dvref isa InfiniteOpt.ParameterFunctionRef
+        prefs = InfiniteOpt.parameter_list(dvref)
+        length(prefs) == 1 || return (-Inf, Inf)
+        supports = InfiniteOpt.supports(first(prefs))
+        isempty(supports) && return (-Inf, Inf)
+        func = InfiniteOpt.raw_function(dvref)
+        func_values = [func(s) for s in supports]
+        return (minimum(func_values), maximum(func_values))
+    elseif dvref isa InfiniteOpt.FiniteParameterRef
+        value = InfiniteOpt.parameter_value(dvref)
+        return (value, value)
+    end
+    return (-Inf, Inf)
+end
+
+function DP.set_variable_bound_info(
+    vref::InfiniteOpt.GeneralVariableRef, ::DP.BigM)
+    info = _parameter_bound_info(vref)
+    info === nothing || return info
+    lb = JuMP.has_lower_bound(vref) ? JuMP.lower_bound(vref) : -Inf
+    ub = JuMP.has_upper_bound(vref) ? JuMP.upper_bound(vref) : Inf
+    return lb, ub
+end
+
+# Hull and PSplit require finite bounds that include 0
+function _zero_inclusive_bounds(
+    vref::InfiniteOpt.GeneralVariableRef,
+    info::Union{Nothing, Tuple{Float64, Float64}},
+    method_name::String
+    )::Tuple{Float64, Float64}
+    info === nothing || return (min(0, info[1]), max(0, info[2]))
+    if !JuMP.has_lower_bound(vref) || !JuMP.has_upper_bound(vref)
+        error("Variable $vref must have both lower and upper " *
+              "bounds defined when using the $method_name " *
+              "reformulation.")
+    end
+    return (min(0, JuMP.lower_bound(vref)),
+            max(0, JuMP.upper_bound(vref)))
+end
+
+function DP.set_variable_bound_info(
+    vref::InfiniteOpt.GeneralVariableRef, ::DP.Hull)
+    return _zero_inclusive_bounds(vref,
+        _parameter_bound_info(vref), "Hull")
+end
+
+function DP.set_variable_bound_info(
+    vref::InfiniteOpt.GeneralVariableRef,
+    ::Union{DP.PSplit, DP._PSplit})
+    return _zero_inclusive_bounds(vref,
+        _parameter_bound_info(vref), "PSplit")
+end
+
 function DP.VariableProperties(vref::InfiniteOpt.GeneralVariableRef)
     info = DP.get_variable_info(vref)
     name = JuMP.name(vref)
@@ -263,31 +324,62 @@ function _interpolate_at(
         )
 end
 
-# Transcribe mini_expr, solve per support on the transcribed JuMP
-# model, and aggregate to a scalar if uniform, else to a parameter
-# function on main.
+# The infinite parameters of `mini_expr` and their supports, in the
+# ascending order of `parameter_refs`. Only defined when M varies over
+# the supports, so it is deferred until sample_M_values needs it.
+function _support_grids(
+    sub::DP.GDPSubmodel, mini_expr::JuMP.AbstractJuMPScalar)
+    reverse_map = Dict(ws[1] => v for (v, ws) in sub.fwd_map)
+    prefs = Tuple(get(reverse_map, p) do
+            error("MBM cannot build a support grid over `$p`, which " *
+                  "is a group of dependent infinite parameters, so M " *
+                  "must not vary over its supports.")
+        end for p in InfiniteOpt.parameter_refs(mini_expr))
+    return prefs, Tuple(InfiniteOpt.supports(p) for p in prefs)
+end
+
+# Solve the M subproblem exactly at every support
+function DP.sample_M_values(
+    sampler::Nothing,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids::Function
+    )
+    M_vals = Array{Float64}(undef, size(objectives))
+    for I in eachindex(objectives)
+        m = DP.raw_M(sub, objectives[I], method)
+        m === nothing && return nothing
+        M_vals[I] = m
+    end
+    return M_vals
+end
+
+# Transcribe mini_expr, compute the per-support M values with the
+# method's sampler, and aggregate to a scalar if uniform, else to a
+# parameter function on main.
 function DP.raw_M(
     sub::DP.GDPSubmodel{<:InfiniteOpt.InfiniteModel},
     mini_expr::JuMP.AbstractJuMPScalar,
     method::DP._MBM
     )
     objectives = InfiniteOpt.transformation_expression(mini_expr)
-    transcribed = InfiniteOpt.transformation_model(sub.model)
-    inner_sub = DP.GDPSubmodel(transcribed,JuMP.VariableRef[],
-        Dict{JuMP.VariableRef, Vector{JuMP.VariableRef}}()
-        )
-    M_vals = Array{typeof(method.default_M)}(undef, size(objectives))
-    for I in eachindex(objectives)
-        m = DP.raw_M(inner_sub, objectives[I], method)
-        m === nothing && return nothing
-        M_vals[I] = m
+    # transcription orders the dimensions by parameter group, which is
+    # not the ascending order `parameter_refs` gives the grids below
+    group_idxs = InfiniteOpt.parameter_group_int_indices(mini_expr)
+    if length(group_idxs) > 1 && ndims(objectives) == length(group_idxs)
+        objectives = permutedims(objectives, sortperm(group_idxs))
     end
+    transcribed = InfiniteOpt.transformation_model(sub.model)
+    inner_sub = DP.GDPSubmodel(transcribed, JuMP.VariableRef[],
+        Dict{JuMP.VariableRef, Vector{JuMP.VariableRef}}())
+    M_vals = DP.sample_M_values(method.sampler, objectives,
+        inner_sub, method, () -> _support_grids(sub, mini_expr)[2])
+    M_vals === nothing && return nothing
+    M_vals isa Number && return M_vals
     all(==(first(M_vals)), M_vals) && return first(M_vals)
-    mini_prefs = InfiniteOpt.parameter_refs(mini_expr)
-    reverse_map = Dict(ws[1] => v for (v, ws) in sub.fwd_map)
-    prefs = Tuple(reverse_map[p] for p in mini_prefs)
+    prefs, grids = _support_grids(sub, mini_expr)
     main = JuMP.owner_model(first(prefs))
-    grids = Tuple(InfiniteOpt.supports(p) for p in prefs)
     param_func = InfiniteOpt.build_parameter_function(
         error, _interpolate(grids, M_vals), prefs)
     return InfiniteOpt.add_parameter_function(main, param_func)