MBM: optional GP-based M sampling on InfiniteGDP models

Project.toml

@@ -8,14 +8,19 @@ 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"
+KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 
 [extensions]
 InfiniteDisjunctiveProgramming = "InfiniteOpt"
+InfiniteGPDisjunctiveProgramming = ["InfiniteOpt", "AbstractGPs", "KernelFunctions"]
 
 [compat]
+AbstractGPs = "0.5"
 Aqua = "0.8"
 JuMP = "1.18"
+KernelFunctions = "0.10"
 Reexport = "1"
 julia = "1.10"
 Juniper = "0.9.3"
@@ -30,4 +35,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", "KernelFunctions"]

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).
+    - `M_sampler`: Strategy for computing M values across the supports of an infinite model. Default: `:auto`, which uses a Gaussian-process sampler (`GPSampler`) when AbstractGPs is loaded and otherwise solves an M subproblem at every support (`:exact`). 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`) activates 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. The filled values are heuristic upper estimates rather than certificates; see the `GPSampler` docstring for the tuning knobs (`kappa`, `budget`) and use `MBM(optimizer, M_sampler = :exact)` to force exact solves.
+
 ## 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/InfiniteDisjunctiveProgramming.jl

@@ -263,31 +263,72 @@ 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 an M sampler needs it.
+function _support_grids(sub::DP.GDPSubmodel, mini_expr)
+    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
+
+# Resolve `:auto` to the GP sampler when the GP extension is loaded
+function _resolve_M_sampler(sampler)
+    sampler === :auto || return sampler
+    gp = Base.get_extension(DP, :InfiniteGPDisjunctiveProgramming)
+    return isnothing(gp) ? :exact : DP.GPSampler()
+end
+
+# Solve the M subproblem exactly at every support
+function DP.sample_M_values(
+    sampler::Symbol,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids
+    )
+    sampler === :exact || error(
+        "Unrecognized `M_sampler` `$(repr(sampler))` for MBM on an " *
+        "infinite model. Use `:auto`, `:exact`, or a `GPSampler`.")
+    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
+# resolved M 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}}())
+    sampler = _resolve_M_sampler(method.M_sampler)
+    M_vals = DP.sample_M_values(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)

ext/InfiniteGPDisjunctiveProgramming.jl

@@ -0,0 +1,133 @@
+module InfiniteGPDisjunctiveProgramming
+
+import InfiniteOpt, JuMP
+import AbstractGPs, KernelFunctions
+import DisjunctiveProgramming as DP
+
+################################################################################
+#                                 GP SAMPLER
+################################################################################
+# See the `GPSampler` docstring in `src/extension_api.jl`
+struct GPSampler{K}
+    kappa::Float64
+    budget::Float64
+    min_solves::Int
+    kernel::K
+    detect_uniform_M::Bool
+end
+
+function DP.GPSampler(;
+    kappa::Real = 2.5,
+    budget::Real = 0.25,
+    min_solves::Int = 6,
+    kernel = nothing,
+    detect_uniform_M::Bool = true
+    )
+    kappa >= 0 || error("`kappa` must be nonnegative.")
+    0 < budget <= 1 || error("`budget` must be in `(0, 1]`.")
+    min_solves >= 1 || error("`min_solves` must be at least 1.")
+    return GPSampler(Float64(kappa), Float64(budget), min_solves,
+        kernel, detect_uniform_M)
+end
+
+################################################################################
+#                                 GP FITTING
+################################################################################
+const _LENGTHSCALES = (0.05, 0.1, 0.2, 0.4, 0.8)
+const _JITTER = 1e-8
+
+# Normalized to [0, 1]^d so one lengthscale works across dimensions
+function _support_coords(grids)
+    idxs = CartesianIndices(length.(grids))
+    los = [minimum(g) for g in grids]
+    rng = [max(maximum(g) - minimum(g), eps()) for g in grids]
+    return [[(grids[d][I[d]] - los[d]) / rng[d] for d in 1:length(grids)]
+            for I in vec(idxs)]
+end
+
+# Lengthscale by marginal likelihood unless the user gave a kernel
+function _fit_posterior(sampler::GPSampler, X, y)
+    isnothing(sampler.kernel) || return AbstractGPs.posterior(
+        AbstractGPs.GP(sampler.kernel)(X, _JITTER), y)
+    best_post, best_lp = nothing, -Inf
+    for ls in _LENGTHSCALES
+        kern = KernelFunctions.with_lengthscale(
+            KernelFunctions.SqExponentialKernel(), ls)
+        fx = AbstractGPs.GP(kern)(X, _JITTER)
+        lp = AbstractGPs.logpdf(fx, y)
+        if lp > best_lp
+            best_post, best_lp = AbstractGPs.posterior(fx, y), lp
+        end
+    end
+    return best_post
+end
+
+function _mean_sd(sampler::GPSampler, X, solved)
+    lis = collect(keys(solved))
+    y = [solved[li] for li in lis]
+    ybar = sum(y) / length(y)
+    # floored so near-equal solved values still cushion the filled ones
+    ystd = max(sqrt(sum(abs2, y .- ybar) / max(length(y) - 1, 1)),
+        1e-2 * abs(ybar), 1e-8)
+    post = _fit_posterior(sampler, X[lis], (y .- ybar) ./ ystd)
+    mz = AbstractGPs.mean(post, X)
+    vz = max.(AbstractGPs.var(post, X), 0.0)
+    return mz .* ystd .+ ybar, sqrt.(vz) .* ystd
+end
+
+################################################################################
+#                              M VALUE SAMPLING
+################################################################################
+# Solve M at max-UCB selected supports, fill the rest with the bound
+function DP.sample_M_values(
+    sampler::GPSampler,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids
+    )
+    idxs = collect(CartesianIndices(objectives))
+    n = length(idxs)
+    solved = Dict{Int, Float64}()
+    solve_at(li) = begin
+        m = DP.raw_M(sub, objectives[idxs[li]], method)
+        m === nothing && return false
+        solved[li] = m
+        return true
+    end
+    # golden fractions; even spacing aliases with a periodic M
+    for s in unique([1, n, 1 + floor(Int, 0.618 * (n - 1)),
+                     1 + floor(Int, 0.382 * (n - 1))])
+        solve_at(s) || 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 = clamp(
+        ceil(Int, sampler.budget * n), min(sampler.min_solves, n), n)
+    X = _support_coords(support_grids())
+    while length(solved) < budget
+        ms, ss = _mean_sd(sampler, X, solved)
+        acq = ms .+ sampler.kappa .* ss
+        for li in keys(solved)
+            acq[li] = -Inf
+        end
+        solve_at(argmax(acq)) || return nothing
+    end
+    M_vals = Array{Float64}(undef, size(objectives))
+    if length(solved) == n # nothing left to estimate
+        for (li, I) in enumerate(idxs)
+            M_vals[I] = solved[li]
+        end
+        return M_vals
+    end
+    ms, ss = _mean_sd(sampler, X, solved)
+    for (li, I) in enumerate(idxs) # exact M values are nonnegative
+        M_vals[I] = get(solved, li, max(ms[li] + sampler.kappa * ss[li], 0.0))
+    end
+    return M_vals
+end
+
+end

src/datatypes.jl

@@ -368,28 +368,36 @@ struct BigM{T} <: AbstractReformulationMethod
 end
 
 """
-    MBM{O, T, L <: LogicalVariableRef} <: AbstractReformulationMethod
+    MBM{O, T} <: AbstractReformulationMethod
 
 A type for using the multiple big-M reformulation approach for disjunctive constraints.
 
 **Fields**
 - `optimizer::O`: Optimizer to use when solving mini-models (required).
 - `default_M::T`: Default big-M value to use if no big-M is specified for a logical variable (1e9).
+- `M_sampler::Any`: Strategy for computing M values across the supports
+  of an infinite model (`:auto`). `:auto` uses [`GPSampler`](@ref) when
+  the AbstractGPs extension is loaded and `:exact` otherwise; `:exact`
+  solves an M subproblem at every support. Ignored for finite models.
 """
 mutable struct MBM{O, T} <: AbstractReformulationMethod
     optimizer::O
     default_M::T
-
+    M_sampler::Any
+
     # Constructor with optimizer (required) and optional default_M
-    function MBM(optimizer::O, default_M::T = 1e9) where {O, T}
-        new{O, T}(optimizer, default_M)
+    function MBM(
+        optimizer::O, default_M::T = 1e9; M_sampler = :auto
+        ) where {O, T}
+        new{O, T}(optimizer, default_M, M_sampler)
     end
 end
 
 mutable struct _MBM{O, T, M <: JuMP.AbstractModel} <: AbstractReformulationMethod
     optimizer::O
     M::Dict{LogicalVariableRef{M}, Any}
     default_M::T
+    M_sampler::Any
     subproblem_indicators::Vector{LogicalVariableRef{M}}
     # Cached submodels: indicator => GDPSubmodel.
     # Typed Any so extensions can store different types.
@@ -400,6 +408,7 @@ mutable struct _MBM{O, T, M <: JuMP.AbstractModel} <: AbstractReformulationMetho
             method.optimizer,
             Dict{LogicalVariableRef{M}, Any}(),
             method.default_M,
+            method.M_sampler,
             Vector{LogicalVariableRef{M}}(),
             Dict{LogicalVariableRef{M}, Any}()
         )

src/extension_api.jl

@@ -38,3 +38,64 @@ Y(t, x)
 ```
 """
 function InfiniteLogical end
+
+"""
+    GPSampler(; kappa = 2.5, budget = 0.25, min_solves = 6,
+              kernel = nothing, detect_uniform_M = true)
+
+Creates a Gaussian-process M sampler for [`MBM`](@ref) on infinite
+models. Instead of solving an M subproblem at every support of the
+infinite parameters, the sampler solves a subset of the supports
+selected by an upper-confidence-bound acquisition and fills the
+remaining supports with the posterior upper confidence bound
+`mean + kappa * sd`. The filled values are heuristic upper estimates
+of the exact M values, not certificates; increase `kappa` for more
+conservative estimates or use `MBM(...; M_sampler = :exact)` to solve
+every support. This requires that InfiniteOpt and AbstractGPs be
+imported first, in which case it is also the default M sampler (see
+the `M_sampler` field of [`MBM`](@ref)).
+
+**Keyword Arguments**
+- `kappa::Real`: Upper-confidence-bound multiplier used to select the
+  next support to solve and to fill unsolved supports (2.5).
+- `budget::Real`: Fraction of the supports to solve exactly, in
+  `(0, 1]` (0.25).
+- `min_solves::Int`: Minimum number of exactly solved supports (6).
+- `kernel`: Covariance kernel for the GP fit. Defaults to a squared
+  exponential kernel whose lengthscale is selected by maximizing the
+  marginal likelihood; pass a `KernelFunctions` kernel to override.
+- `detect_uniform_M::Bool`: If `true` (the default), M values that
+  agree at the first few supports are taken to be uniform and used
+  for every support. This is cheap and is what makes infinite
+  parameters without a support grid (e.g. dependent ones) workable,
+  but it assumes M does not vary elsewhere. Set it to `false` to
+  always fit the GP, which leaves the usual `kappa * sd` cushion on
+  the unsolved supports at the cost of the extra solves, and which
+  requires that every infinite parameter have a support grid.
+
+**Example**
+```julia
+julia> using DisjunctiveProgramming, InfiniteOpt, AbstractGPs, HiGHS
+
+julia> method = MBM(HiGHS.Optimizer, M_sampler = GPSampler(kappa = 4.0))
+```
+"""
+function GPSampler end
+
+"""
+    sample_M_values(sampler, objectives, sub, method, support_grids)
+
+Compute the MBM M values at the transcription supports of an infinite
+model. `objectives` is the array of per-support objective expressions,
+`sub` is the transcribed submodel wrapped as a `GDPSubmodel`, `method`
+is the `_MBM` data, and `support_grids` is a function returning the
+support vectors of the infinite parameters. It is a function because
+the supports are only well defined once M is known to vary over them,
+so samplers that return early (or never need coordinates) must not
+call it. Returns an array of M values shaped like
+`objectives`, a scalar when M is uniform across the supports, or
+`nothing` if an M subproblem is infeasible. Extensions implement
+methods for their sampler types (e.g. [`GPSampler`](@ref)); the
+`:exact` sampler solves every support.
+"""
+function sample_M_values end

src/mbm.jl

@@ -87,7 +87,8 @@ function reformulate_disjunct_constraint(
         },
     method::_MBM
     )
-    ref_cons = reformulate_disjunction(model, con, MBM(method.optimizer))
+    ref_cons = reformulate_disjunction(model, con, MBM(
+        method.optimizer, method.default_M, M_sampler = method.M_sampler))
     new_ref_cons = Vector{JuMP.AbstractConstraint}()
     for ref_con in ref_cons
         append!(new_ref_cons,