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: `ExhaustiveSampler()`, 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,173 @@
+module AbstractGPsDisjunctiveProgramming
+
+import AbstractGPs
+import AbstractGPs.KernelFunctions
+import DisjunctiveProgramming as DP
+
+################################################################################
+#                               SAMPLER CONFIG
+################################################################################
+# The concrete sampler behind DP.GPSampler; the base package only
+# carries the function stub, so constructing one requires AbstractGPs.
+struct _GPSampler{F} <: DP.AbstractMBMSampler
+    f::F
+    std_dev_margin::Float64
+    frac_supports::Float64
+    detect_uniform_M::Bool
+    initial_supports::Union{Int, Vector{Float64}}
+
+    function _GPSampler(
+        f::F;
+        std_dev_margin::Real = 2.5,
+        frac_supports::Real = 0.25,
+        detect_uniform_M::Bool = true,
+        initial_supports = 4
+        ) where {F}
+        f isa Union{Nothing, AbstractGPs.GP} || error(
+            "`f` must be an `AbstractGPs.GP` prior, e.g. " *
+            "`GP(Matern52Kernel())`.")
+        std_dev_margin >= 0 || error("`std_dev_margin` must be nonnegative.")
+        0 < frac_supports <= 1 || error("`frac_supports` must be in `(0, 1]`.")
+        if initial_supports isa Int
+            initial_supports >= 2 ||
+                error("`initial_supports` must be at least 2.")
+        else
+            initial_supports = collect(Float64, initial_supports)
+            (!isempty(initial_supports) &&
+                all(frac -> 0 <= frac <= 1, initial_supports)) ||
+                error("`initial_supports` must be fractions in `[0, 1]`.")
+        end
+        new{F}(f, Float64(std_dev_margin), Float64(frac_supports),
+            detect_uniform_M, initial_supports)
+    end
+end
+
+DP.GPSampler(f = nothing; kwargs...) = _GPSampler(f; kwargs...)
+
+################################################################################
+#                                 GP FITTING
+################################################################################
+# Normalized to [0, 1]^d so one lengthscale works across dimensions.
+# An independent parameter contributes one coordinate; a dependent
+# group contributes its joint-support column.
+function _support_coords(grids::Tuple)::Vector{Vector{Float64}}
+    axis_coords = map(grids) do g
+        g isa AbstractMatrix ? [g[:, j] for j in axes(g, 2)] :
+            [[v] for v in g]
+    end
+    coords = [reduce(vcat, getindex.(axis_coords, Tuple(I)))
+              for I in vec(CartesianIndices(length.(axis_coords)))]
+    dims = eachindex(first(coords))
+    mins = [minimum(c[d] for c in coords) for d in dims]
+    ranges = [max(maximum(c[d] for c in coords) - mins[d], eps())
+              for d in dims]
+    return [[(c[d] - mins[d]) / ranges[d] for d in dims]
+            for c in coords]
+end
+
+# Defaults behind GPSampler(): candidate lengthscales for the
+# marginal-likelihood fit (relative to the unit box) and the
+# observation-noise nugget.
+const _LENGTHSCALES = [0.05, 0.1, 0.2, 0.4, 0.8]
+const _JITTER = 1e-8
+
+# A user prior is used as given; the default squared exponential has
+# its lengthscale selected by marginal likelihood
+function _fit_posterior(
+    sampler::_GPSampler,
+    X::Vector{Vector{Float64}},
+    y::Vector{Float64}
+    )
+    sampler.f === nothing ||
+        return AbstractGPs.posterior(sampler.f(X, _JITTER), y)
+    best_posterior, best_log_prob = nothing, -Inf
+    for lengthscale in _LENGTHSCALES
+        kernel = KernelFunctions.with_lengthscale(
+            KernelFunctions.SqExponentialKernel(), lengthscale)
+        finite_gp = AbstractGPs.GP(kernel)(X, _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::_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)
+    posterior = _fit_posterior(
+        sampler, X[solved_indices], (y .- y_mean) ./ y_scale)
+    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::_GPSampler,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids::Tuple
+    )
+    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 initial count, or user-given fractions
+    fractions = sampler.initial_supports isa Int ?
+        range(0, 1, length = sampler.initial_supports) :
+        sampler.initial_supports
+    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
+        probes = collect(values(solved))
+        all(==(first(probes)), probes) && return first(probes)
+    end
+    solve_target = min(ceil(Int, sampler.frac_supports * n), n)
+    X = _support_coords(support_grids)
+    while length(solved) < solve_target
+        means, sds = _mean_sd(sampler, X, solved)
+        acquisition = means .+ sampler.std_dev_margin .* 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.std_dev_margin * sds[index], 0.0))
+    end
+    return M_vals
+end
+
+end

ext/InfiniteDisjunctiveProgramming.jl

@@ -238,56 +238,84 @@ function DP.prepare_max_M_objective(
     return obj.set.lower - obj_func
 end
 
-# Constant interpolation
-function _interpolate(
-    grids::NTuple{N, AbstractVector{<:Real}},
-    values::AbstractArray{<:Real, N}
-    ) where {N}
+# Candidate indices along one axis of the M array: the corners of the
+# grid cell bracketing a scalar query (independent parameter), or the
+# column matching a joint-support query (dependent group; every column
+# when the query is off-support, so the estimate stays conservative)
+function _axis_candidates(grid::AbstractVector{<:Real}, arg::Real)
+    lo = clamp(searchsortedlast(grid, arg), 1, length(grid) - 1)
+    return lo:(lo + 1)
+end
+function _axis_candidates(grid::AbstractMatrix{<:Real}, arg)
+    j = findfirst(k -> isapprox(view(grid, :, k), arg, atol = 1e-10),
+                  axes(grid, 2))
+    return isnothing(j) ? axes(grid, 2) : (j:j)
+end
+
+# Constant interpolation: max of `values` over the candidate indices
+function _interpolate(grids::Tuple, values::AbstractArray{<:Real})
     # mimic the call form of Interpolations.jl's interpolation
-    return (args...) -> _interpolate_at(grids, values, args)
+    return (args...) -> maximum(
+        values[I...] for I in Iterators.product(
+            map(_axis_candidates, grids, args)...))
+end
+
+# The infinite parameters of `mini_expr` and their supports, in the
+# ascending order of `parameter_refs`. An independent parameter gives
+# its sorted support vector; a dependent group gives the matrix whose
+# columns are its joint supports. Grids are read off the mini model so
+# their column order matches the transcription axes; the returned
+# prefs are the main-model parameters.
+function _support_grids(
+    sub::DP.GDPSubmodel, mini_expr::JuMP.AbstractJuMPScalar)
+    reverse_map = Dict(ws[1] => v for (v, ws) in sub.fwd_map)
+    mini_prefs = InfiniteOpt.parameter_refs(mini_expr)
+    prefs = Tuple(getindex.(Ref(reverse_map), p) for p in mini_prefs)
+    return prefs, Tuple(InfiniteOpt.supports(p) for p in mini_prefs)
 end
 
-function _interpolate_at(
-    grids::NTuple{N, AbstractVector{<:Real}},
-    values::AbstractArray{<:Real, N},
-    args::NTuple{N, <:Real}
-    ) where {N}
-    # lower-corner cell index per dimension
-    idx_lo = ntuple(d -> 
-        clamp(searchsortedlast(grids[d], args[d]),1, length(grids[d]) - 1), N
+# Solve the M subproblem exactly at every support
+function DP.sample_M_values(
+    sampler::DP.ExhaustiveSampler,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    support_grids::Tuple
     )
-    # max over the 2^N corners; bit d of k picks lower or upper
-    return maximum(
-        values[ntuple(d -> idx_lo[d] +((k >> (d - 1)) & 1), N)...]
-        for k in 0:(2^N - 1)
-        )
+    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, solve per support on the transcribed JuMP
-# model, and aggregate to a scalar if uniform, else to a parameter
-# function on main.
+# 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}}())
+    prefs, grids = _support_grids(sub, mini_expr)
+    M_vals = DP.sample_M_values(method.sampler, objectives,
+        inner_sub, method, grids)
+    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)
-    main = JuMP.owner_model(first(prefs))
-    grids = Tuple(InfiniteOpt.supports(p) for p in prefs)
+    main = JuMP.owner_model(first(keys(sub.fwd_map)))
     param_func = InfiniteOpt.build_parameter_function(
         error, _interpolate(grids, M_vals), prefs)
     return InfiniteOpt.add_parameter_function(main, param_func)