From a930770b7ff25a45693b8c03a58797cc7a6a7c65 Mon Sep 17 00:00:00 2001
From: d227nguyen <d227nguy@uwaterloo.ca>
Date: Wed, 29 Jul 2026 12:33:36 -0400
Subject: [PATCH 1/4] Initial

---
 Project.toml                          |   2 +
 ext/InfiniteDisjunctiveProgramming.jl | 104 +++++++++++++++++++++++---
 2 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/Project.toml b/Project.toml
index bdec9be..5de25a7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,7 +4,9 @@ authors = ["hdavid16 <hperez16@gmail.com>"]
 version = "0.6.1"
 
 [deps]
+AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [weakdeps]
diff --git a/ext/InfiniteDisjunctiveProgramming.jl b/ext/InfiniteDisjunctiveProgramming.jl
index 5f77dce..a251067 100644
--- a/ext/InfiniteDisjunctiveProgramming.jl
+++ b/ext/InfiniteDisjunctiveProgramming.jl
@@ -3,6 +3,7 @@ module InfiniteDisjunctiveProgramming
 import JuMP.MOI as _MOI
 import InfiniteOpt, JuMP
 import DisjunctiveProgramming as DP
+import AbstractGPs, KernelFunctions
 
 ################################################################################
 #                                   MODEL
@@ -263,9 +264,95 @@ 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.
+# ------ GP active-learning M(d) (hard-coded for the mbm_gp experiments) ------
+# Coordinates of every support, normalized to [0,1]^d so one isotropic
+# lengthscale works across dimensions.
+function _gp_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]
+    X = [[(grids[d][I[d]] - los[d]) / rng[d] for d in 1:length(grids)]
+         for I in idxs]
+    return X, collect(idxs)
+end
+
+# Posterior mean and sd at all coords, given solved (index => M) samples.
+# Outputs are standardized so the k*sd term scales with the M spread.
+function _gp_mean_sd(X, solved)
+    lis = collect(keys(solved))
+    yt = [solved[li] for li in lis]
+    ybar = sum(yt) / length(yt)
+    ystd = max(sqrt(sum(abs2, yt .- ybar) / max(length(yt) - 1, 1)), 1e-8)
+    kern = KernelFunctions.with_lengthscale(
+        KernelFunctions.SqExponentialKernel(), 0.1)
+    post = AbstractGPs.posterior(
+        AbstractGPs.GP(kern)(X[lis], 1e-8), (yt .- ybar) ./ ystd)
+    mz = AbstractGPs.mean(post, X)
+    vz = max.(AbstractGPs.var(post, X), 0.0)
+    return mz .* ystd .+ ybar, sqrt.(vz) .* ystd
+end
+
+# Count of per-support M subproblems solved (for the grid-vs-GP comparison).
+const _M_SOLVE_COUNT = Ref(0)
+
+# Original workflow: solve M at every support (grid). Kept for the mbm_gp
+# vs grid comparison, gated by ENV["DP_MBM_GRID"].
+function _grid_M_vals(objectives, inner_sub, method)
+    M_vals = Array{Float64}(undef, size(objectives))
+    for I in eachindex(objectives)
+        _M_SOLVE_COUNT[] += 1
+        m = DP.raw_M(inner_sub, objectives[I], method)
+        m === nothing && return nothing
+        M_vals[I] = m
+    end
+    return M_vals
+end
+
+# Solve M at actively-selected supports (max-UCB acquisition) and fill the
+# rest with the UCB (mean + k*sd), a valid over-estimate. Returns a scalar
+# when M is uniform (e.g. dependent multi-dim parameters where M does not
+# vary), matching the grid workflow's early return.
+function _gp_active_M_vals(objectives, inner_sub, method, sub, mini_expr)
+    idxs = collect(CartesianIndices(objectives))
+    n = length(idxs)
+    k = 2.5
+    budget = min(n, max(6, cld(n, 4)))
+    solved = Dict{Int, Float64}()
+    solve_at!(li) = begin
+        _M_SOLVE_COUNT[] += 1
+        m = DP.raw_M(inner_sub, objectives[idxs[li]], method)
+        m === nothing && return false
+        solved[li] = m
+        true
+    end
+    for s in unique([1, cld(n + 1, 2), n])
+        solve_at!(s) || return nothing
+    end
+    seed = collect(values(solved))
+    all(==(first(seed)), seed) && return first(seed)   # uniform M -> scalar
+    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)
+    grids = Tuple(InfiniteOpt.supports(p) for p in prefs)
+    X, _ = _gp_support_coords(grids)
+    while length(solved) < budget
+        ms, ss = _gp_mean_sd(X, solved)
+        acq = ms .+ k .* ss
+        for li in keys(solved)
+            acq[li] = -Inf
+        end
+        solve_at!(argmax(acq)) || return nothing
+    end
+    ms, ss = _gp_mean_sd(X, solved)
+    M_vals = Array{Float64}(undef, size(objectives))
+    for (li, I) in enumerate(idxs)
+        M_vals[I] = get(solved, li, ms[li] + k * ss[li])
+    end
+    return M_vals
+end
+
+# Transcribe mini_expr, then approximate M(d) with an actively-sampled GP
+# instead of solving at every support; aggregate to a scalar if uniform.
 function DP.raw_M(
     sub::DP.GDPSubmodel{<:InfiniteOpt.InfiniteModel},
     mini_expr::JuMP.AbstractJuMPScalar,
@@ -276,12 +363,11 @@ function DP.raw_M(
     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
-    end
+    M_vals = haskey(ENV, "DP_MBM_GRID") ?
+        _grid_M_vals(objectives, inner_sub, method) :
+        _gp_active_M_vals(objectives, inner_sub, method, sub, mini_expr)
+    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)

From ed08522cbfc1466d71845ffcda6eb85b07174166 Mon Sep 17 00:00:00 2001
From: d227nguyen <d227nguy@uwaterloo.ca>
Date: Sat, 1 Aug 2026 21:54:27 -0400
Subject: [PATCH 2/4] Simplify and code coverage

---
 Project.toml                                  |   9 +-
 README.md                                     |   3 +
 ext/InfiniteDisjunctiveProgramming.jl         | 124 ++++---------
 ext/InfiniteGPDisjunctiveProgramming.jl       | 135 +++++++++++++++
 src/datatypes.jl                              |  17 +-
 src/extension_api.jl                          |  50 ++++++
 src/mbm.jl                                    |   3 +-
 .../InfiniteDisjunctiveProgramming.jl         |  33 +++-
 .../InfiniteGPDisjunctiveProgramming.jl       | 163 ++++++++++++++++++
 test/runtests.jl                              |   1 +
 10 files changed, 438 insertions(+), 100 deletions(-)
 create mode 100644 ext/InfiniteGPDisjunctiveProgramming.jl
 create mode 100644 test/extensions/InfiniteGPDisjunctiveProgramming.jl

diff --git a/Project.toml b/Project.toml
index 5de25a7..ba1a492 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,20 +4,23 @@ authors = ["hdavid16 <hperez16@gmail.com>"]
 version = "0.6.1"
 
 [deps]
-AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
-KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 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"
@@ -32,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"]
diff --git a/README.md b/README.md
index 543d1f7..c3cd037 100644
--- a/README.md
+++ b/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).
diff --git a/ext/InfiniteDisjunctiveProgramming.jl b/ext/InfiniteDisjunctiveProgramming.jl
index a251067..2055422 100644
--- a/ext/InfiniteDisjunctiveProgramming.jl
+++ b/ext/InfiniteDisjunctiveProgramming.jl
@@ -3,7 +3,6 @@ module InfiniteDisjunctiveProgramming
 import JuMP.MOI as _MOI
 import InfiniteOpt, JuMP
 import DisjunctiveProgramming as DP
-import AbstractGPs, KernelFunctions
 
 ################################################################################
 #                                   MODEL
@@ -264,116 +263,61 @@ function _interpolate_at(
         )
 end
 
-# ------ GP active-learning M(d) (hard-coded for the mbm_gp experiments) ------
-# Coordinates of every support, normalized to [0,1]^d so one isotropic
-# lengthscale works across dimensions.
-function _gp_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]
-    X = [[(grids[d][I[d]] - los[d]) / rng[d] for d in 1:length(grids)]
-         for I in idxs]
-    return X, collect(idxs)
+# 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
 
-# Posterior mean and sd at all coords, given solved (index => M) samples.
-# Outputs are standardized so the k*sd term scales with the M spread.
-function _gp_mean_sd(X, solved)
-    lis = collect(keys(solved))
-    yt = [solved[li] for li in lis]
-    ybar = sum(yt) / length(yt)
-    ystd = max(sqrt(sum(abs2, yt .- ybar) / max(length(yt) - 1, 1)), 1e-8)
-    kern = KernelFunctions.with_lengthscale(
-        KernelFunctions.SqExponentialKernel(), 0.1)
-    post = AbstractGPs.posterior(
-        AbstractGPs.GP(kern)(X[lis], 1e-8), (yt .- ybar) ./ ystd)
-    mz = AbstractGPs.mean(post, X)
-    vz = max.(AbstractGPs.var(post, X), 0.0)
-    return mz .* ystd .+ ybar, sqrt.(vz) .* ystd
-end
-
-# Count of per-support M subproblems solved (for the grid-vs-GP comparison).
-const _M_SOLVE_COUNT = Ref(0)
-
-# Original workflow: solve M at every support (grid). Kept for the mbm_gp
-# vs grid comparison, gated by ENV["DP_MBM_GRID"].
-function _grid_M_vals(objectives, inner_sub, method)
+# Solve the M subproblem exactly at every support
+function DP.sample_M_values(
+    sampler::Symbol,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    grids::Tuple
+    )
+    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_SOLVE_COUNT[] += 1
-        m = DP.raw_M(inner_sub, objectives[I], method)
+        m = DP.raw_M(sub, objectives[I], method)
         m === nothing && return nothing
         M_vals[I] = m
     end
     return M_vals
 end
 
-# Solve M at actively-selected supports (max-UCB acquisition) and fill the
-# rest with the UCB (mean + k*sd), a valid over-estimate. Returns a scalar
-# when M is uniform (e.g. dependent multi-dim parameters where M does not
-# vary), matching the grid workflow's early return.
-function _gp_active_M_vals(objectives, inner_sub, method, sub, mini_expr)
-    idxs = collect(CartesianIndices(objectives))
-    n = length(idxs)
-    k = 2.5
-    budget = min(n, max(6, cld(n, 4)))
-    solved = Dict{Int, Float64}()
-    solve_at!(li) = begin
-        _M_SOLVE_COUNT[] += 1
-        m = DP.raw_M(inner_sub, objectives[idxs[li]], method)
-        m === nothing && return false
-        solved[li] = m
-        true
-    end
-    for s in unique([1, cld(n + 1, 2), n])
-        solve_at!(s) || return nothing
-    end
-    seed = collect(values(solved))
-    all(==(first(seed)), seed) && return first(seed)   # uniform M -> scalar
-    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)
-    grids = Tuple(InfiniteOpt.supports(p) for p in prefs)
-    X, _ = _gp_support_coords(grids)
-    while length(solved) < budget
-        ms, ss = _gp_mean_sd(X, solved)
-        acq = ms .+ k .* ss
-        for li in keys(solved)
-            acq[li] = -Inf
-        end
-        solve_at!(argmax(acq)) || return nothing
-    end
-    ms, ss = _gp_mean_sd(X, solved)
-    M_vals = Array{Float64}(undef, size(objectives))
-    for (li, I) in enumerate(idxs)
-        M_vals[I] = get(solved, li, ms[li] + k * ss[li])
-    end
-    return M_vals
-end
-
-# Transcribe mini_expr, then approximate M(d) with an actively-sampled GP
-# instead of solving at every support; aggregate to a scalar if uniform.
+# 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)
+    # 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 = haskey(ENV, "DP_MBM_GRID") ?
-        _grid_M_vals(objectives, inner_sub, method) :
-        _gp_active_M_vals(objectives, inner_sub, method, sub, mini_expr)
-    M_vals === nothing && return nothing
-    M_vals isa Number && return M_vals
-    all(==(first(M_vals)), M_vals) && return first(M_vals)
+    inner_sub = DP.GDPSubmodel(transcribed, JuMP.VariableRef[],
+        Dict{JuMP.VariableRef, Vector{JuMP.VariableRef}}())
     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)
+    sampler = _resolve_M_sampler(method.M_sampler)
+    M_vals = DP.sample_M_values(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)
+    main = JuMP.owner_model(first(prefs))
     param_func = InfiniteOpt.build_parameter_function(
         error, _interpolate(grids, M_vals), prefs)
     return InfiniteOpt.add_parameter_function(main, param_func)
diff --git a/ext/InfiniteGPDisjunctiveProgramming.jl b/ext/InfiniteGPDisjunctiveProgramming.jl
new file mode 100644
index 0000000..0e9e136
--- /dev/null
+++ b/ext/InfiniteGPDisjunctiveProgramming.jl
@@ -0,0 +1,135 @@
+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
+end
+
+function DP.GPSampler(;
+    kappa::Real = 2.5,
+    budget::Real = 0.25,
+    min_solves::Int = 6,
+    kernel = nothing
+    )
+    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)
+end
+
+################################################################################
+#                                 GP FITTING
+################################################################################
+# Lengthscale candidates on the [0, 1]-normalized support coordinates
+const _LENGTHSCALES = (0.05, 0.1, 0.2, 0.4, 0.8)
+
+# Observation jitter for the GP fit
+const _JITTER = 1e-8
+
+# Coordinates of every support in linear index order, normalized to
+# [0, 1]^d so one isotropic 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
+
+# Fit the GP posterior on the solved coordinates; `y` is standardized
+# by the caller. With no user kernel, select the lengthscale of a
+# squared exponential kernel by maximizing the marginal likelihood.
+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
+
+# Posterior mean and sd at all coords given the solved (index => M)
+# samples, destandardized back to M units
+function _mean_sd(sampler::GPSampler, X, solved)
+    lis = collect(keys(solved))
+    y = [solved[li] for li in lis]
+    ybar = sum(y) / length(y)
+    ystd = max(sqrt(sum(abs2, y .- ybar) / max(length(y) - 1, 1)), 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 actively-selected supports (max-UCB acquisition) and fill
+# the rest with the upper confidence bound mean + kappa * sd, a
+# heuristic over-estimate. Returns a scalar when the seed M values are
+# uniform (e.g. M does not vary over the supports).
+function DP.sample_M_values(
+    sampler::GPSampler,
+    objectives::AbstractArray,
+    sub::DP.GDPSubmodel,
+    method::DP._MBM,
+    grids::Tuple
+    )
+    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
+    for s in unique([1, cld(n + 1, 2), n])
+        solve_at(s) || return nothing
+    end
+    seed = collect(values(solved))
+    all(==(first(seed)), seed) && return first(seed)
+    budget = clamp(
+        ceil(Int, sampler.budget * n), min(sampler.min_solves, n), n)
+    X = _support_coords(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, so the fill is too
+        M_vals[I] = get(solved, li, max(ms[li] + sampler.kappa * ss[li], 0.0))
+    end
+    return M_vals
+end
+
+end
diff --git a/src/datatypes.jl b/src/datatypes.jl
index bdcd4c2..2bb9038 100644
--- a/src/datatypes.jl
+++ b/src/datatypes.jl
@@ -368,21 +368,28 @@ 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
 
@@ -390,6 +397,7 @@ mutable struct _MBM{O, T, M <: JuMP.AbstractModel} <: AbstractReformulationMetho
     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}()
         )
diff --git a/src/extension_api.jl b/src/extension_api.jl
index ad3566f..6bfa129 100644
--- a/src/extension_api.jl
+++ b/src/extension_api.jl
@@ -38,3 +38,53 @@ Y(t, x)
 ```
 """
 function InfiniteLogical end
+
+"""
+    GPSampler(; kappa = 2.5, budget = 0.25, min_solves = 6,
+              kernel = nothing)
+
+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.
+
+**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, 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 `grids` are the support vectors of the
+infinite parameters. 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
diff --git a/src/mbm.jl b/src/mbm.jl
index 23c5096..4441d95 100644
--- a/src/mbm.jl
+++ b/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,
diff --git a/test/extensions/InfiniteDisjunctiveProgramming.jl b/test/extensions/InfiniteDisjunctiveProgramming.jl
index 47bb884..a2daf8e 100644
--- a/test/extensions/InfiniteDisjunctiveProgramming.jl
+++ b/test/extensions/InfiniteDisjunctiveProgramming.jl
@@ -376,7 +376,7 @@ function test_raw_M_infinite_scalar()
     @constraint(model, con, x >= 5, Disjunct(Y[1]))
     @constraint(model, con2, x <= 3, Disjunct(Y[2]))
     @disjunction(model, Y)
-    mbm = DP._MBM(MBM(HiGHS.Optimizer), model)
+    mbm = DP._MBM(MBM(HiGHS.Optimizer, M_sampler = :exact), model)
     sub = DP.copy_model_with_constraints(
         model, DP.DisjunctConstraintRef[con2], mbm)
     obj = DP.prepare_max_M_objective(
@@ -399,7 +399,7 @@ function test_raw_M_infinite_param_function()
     @constraint(model, con, x <= f, Disjunct(Y[1]))
     @constraint(model, con2, x >= 0.5, Disjunct(Y[2]))
     @disjunction(model, Y)
-    mbm = DP._MBM(MBM(HiGHS.Optimizer), model)
+    mbm = DP._MBM(MBM(HiGHS.Optimizer, M_sampler = :exact), model)
     sub = DP.copy_model_with_constraints(
         model, DP.DisjunctConstraintRef[con2], mbm)
     obj = DP.prepare_max_M_objective(
@@ -413,6 +413,34 @@ function test_raw_M_infinite_param_function()
     end
 end
 
+# raw_M over two infinite parameters with different support counts.
+# Transcription orders the objective dimensions by parameter group,
+# which need not be the ascending order of the grids, so the M values
+# must be permuted to line up. Setup: x(t, s) in [0, 10],
+# disj1: x <= t + s, disj2: x >= 0.5. Slack r(x) = x - t - s
+# maximized over x in [0.5, 10]: 10 - t - s.
+function test_raw_M_infinite_two_params()
+    model = InfiniteGDPModel()
+    @infinite_parameter(model, t ∈ [0, 1], supports = [0.0, 0.5, 1.0])
+    @infinite_parameter(model, s ∈ [0, 1], supports = [0.0, 1.0])
+    @variable(model, 0 <= x <= 10, Infinite(t, s))
+    @variable(model, Y[1:2], InfiniteLogical(t, s))
+    @constraint(model, con, x <= t + s, Disjunct(Y[1]))
+    @constraint(model, con2, x >= 0.5, Disjunct(Y[2]))
+    @disjunction(model, Y)
+    mbm = DP._MBM(MBM(HiGHS.Optimizer, M_sampler = :exact), model)
+    sub = DP.copy_model_with_constraints(
+        model, DP.DisjunctConstraintRef[con2], mbm)
+    obj = DP.prepare_max_M_objective(
+        model, JuMP.constraint_object(con), sub)
+    M = DP.raw_M(sub, obj, mbm)
+    @test M isa InfiniteOpt.GeneralVariableRef
+    raw_fn = InfiniteOpt.raw_function(M)
+    for t_val in [0.0, 0.5, 1.0], s_val in [0.0, 1.0]
+        @test raw_fn(t_val, s_val) >= 10.0 - t_val - s_val - 1e-6
+    end
+end
+
 # Piecewise-constant max-of-corners: returns the maximum value over
 # the 2^n corners of the cell containing the query.
 function test_interpolate()
@@ -834,6 +862,7 @@ end
         test_interpolate()
         test_raw_M_infinite_scalar()
         test_raw_M_infinite_param_function()
+        test_raw_M_infinite_two_params()
         test_mbm_finite_and_integer_var()
         test_mbm_infinite_simple()
         test_mbm_infinite_param_dependent()
diff --git a/test/extensions/InfiniteGPDisjunctiveProgramming.jl b/test/extensions/InfiniteGPDisjunctiveProgramming.jl
new file mode 100644
index 0000000..8b702df
--- /dev/null
+++ b/test/extensions/InfiniteGPDisjunctiveProgramming.jl
@@ -0,0 +1,163 @@
+using InfiniteOpt, HiGHS, AbstractGPs, KernelFunctions
+import DisjunctiveProgramming as DP
+
+# Helpers to access internal functions of the two extensions
+const IGDP = Base.get_extension(DP, :InfiniteGPDisjunctiveProgramming)
+const IDP = Base.get_extension(DP, :InfiniteDisjunctiveProgramming)
+
+function test_gp_sampler_creation()
+    sampler = GPSampler()
+    @test sampler isa IGDP.GPSampler
+    @test sampler.kappa == 2.5
+    @test sampler.budget == 0.25
+    @test sampler.min_solves == 6
+    @test isnothing(sampler.kernel)
+    kern = with_lengthscale(SqExponentialKernel(), 0.3)
+    sampler = GPSampler(
+        kappa = 4.0, budget = 0.1, min_solves = 3, kernel = kern)
+    @test sampler.kappa == 4.0
+    @test sampler.budget == 0.1
+    @test sampler.min_solves == 3
+    @test sampler.kernel === kern
+    @test_throws ErrorException GPSampler(kappa = -1)
+    @test_throws ErrorException GPSampler(budget = 0)
+    @test_throws ErrorException GPSampler(budget = 1.5)
+    @test_throws ErrorException GPSampler(min_solves = 0)
+end
+
+function test_gp_sampler_resolution()
+    # the GP extension is loaded, so :auto resolves to a GPSampler
+    @test IDP._resolve_M_sampler(:auto) isa IGDP.GPSampler
+    @test IDP._resolve_M_sampler(:exact) === :exact
+    sampler = GPSampler(kappa = 3.0)
+    @test IDP._resolve_M_sampler(sampler) === sampler
+    @test MBM(HiGHS.Optimizer).M_sampler === :auto
+    @test MBM(HiGHS.Optimizer, M_sampler = :exact).M_sampler === :exact
+end
+
+# Mirror of test_raw_M_infinite_scalar: uniform seed M values collapse
+# to the exactly-solved scalar under the GP sampler
+function test_gp_raw_M_scalar()
+    model = InfiniteGDPModel()
+    @infinite_parameter(model, t ∈ [0, 1], num_supports = 5)
+    @variable(model, 0 <= x <= 10, Infinite(t))
+    @variable(model, Y[1:2], InfiniteLogical(t))
+    @constraint(model, con, x >= 5, Disjunct(Y[1]))
+    @constraint(model, con2, x <= 3, Disjunct(Y[2]))
+    @disjunction(model, Y)
+    mbm = DP._MBM(MBM(HiGHS.Optimizer), model)
+    sub = DP.copy_model_with_constraints(
+        model, DP.DisjunctConstraintRef[con2], mbm)
+    obj = DP.prepare_max_M_objective(
+        model, JuMP.constraint_object(con), sub)
+    @test DP.raw_M(sub, obj, mbm) == 5.0
+end
+
+# With few supports the budget floor covers every support, so the GP
+# sampler solves all of them exactly and must reproduce the exact
+# grid parameter function
+function test_gp_raw_M_matches_exact()
+    function pfunc_values(M_sampler, supports)
+        model = InfiniteGDPModel()
+        @infinite_parameter(model, t ∈ [0, 1], supports = supports)
+        @variable(model, 0 <= x <= 10, Infinite(t))
+        @variable(model, Y[1:2], InfiniteLogical(t))
+        @parameter_function(model, f == t -> 2*t)
+        @constraint(model, con, x <= f, Disjunct(Y[1]))
+        @constraint(model, con2, x >= 0.5, Disjunct(Y[2]))
+        @disjunction(model, Y)
+        mbm = DP._MBM(
+            MBM(HiGHS.Optimizer, M_sampler = M_sampler), model)
+        sub = DP.copy_model_with_constraints(
+            model, DP.DisjunctConstraintRef[con2], mbm)
+        obj = DP.prepare_max_M_objective(
+            model, JuMP.constraint_object(con), sub)
+        M = DP.raw_M(sub, obj, mbm)
+        @test M isa InfiniteOpt.GeneralVariableRef
+        return [InfiniteOpt.raw_function(M)(t_val) for t_val in supports]
+    end
+    supports = [0.0, 0.25, 0.5, 0.75, 1.0]
+    exact_vals = pfunc_values(:exact, supports)
+    @test pfunc_values(GPSampler(), supports) == exact_vals
+    # a user kernel skips the lengthscale fit but solves the same supports
+    kern = with_lengthscale(SqExponentialKernel(), 0.2)
+    @test pfunc_values(GPSampler(kernel = kern), supports) == exact_vals
+end
+
+# an empty disjunct region makes the M subproblems infeasible; both
+# samplers propagate that up to the reformulation error
+function test_gp_infeasible_disjunct()
+    function build()
+        model = InfiniteGDPModel(HiGHS.Optimizer)
+        set_silent(model)
+        @infinite_parameter(model, t ∈ [0, 1], num_supports = 5)
+        @variable(model, 0 <= x <= 10, Infinite(t))
+        @variable(model, Y[1:2], InfiniteLogical(t))
+        @parameter_function(model, f == t -> 2*t)
+        @constraint(model, x <= f, Disjunct(Y[1]))
+        @constraint(model, x >= 8, Disjunct(Y[2]))
+        @constraint(model, x <= 3, Disjunct(Y[2]))
+        @disjunction(model, Y)
+        @objective(model, Max, 𝔼(x, t))
+        return model
+    end
+    for sampler in (:exact, GPSampler())
+        model = build()
+        @test_throws ErrorException optimize!(model,
+            gdp_method = MBM(HiGHS.Optimizer, M_sampler = sampler))
+    end
+end
+
+# optimum (10) needs M(t) >= 10 - 2t pointwise; the GP fill is heuristic
+function test_gp_mbm_solve_equivalence()
+    function solve_with(M_sampler)
+        model = InfiniteGDPModel(HiGHS.Optimizer)
+        set_silent(model)
+        @infinite_parameter(model, t ∈ [0, 1], num_supports = 20)
+        @variable(model, 0 <= x <= 10, Infinite(t))
+        @variable(model, Y[1:2], InfiniteLogical(t))
+        @parameter_function(model, f == t -> 2*t)
+        @constraint(model, x <= f, Disjunct(Y[1]))
+        @constraint(model, x >= 0.5, Disjunct(Y[2]))
+        @disjunction(model, Y)
+        @objective(model, Max, 𝔼(x, t))
+        optimize!(model,
+            gdp_method = MBM(HiGHS.Optimizer, M_sampler = M_sampler))
+        @test termination_status(model) == MOI.OPTIMAL
+        return objective_value(model)
+    end
+    obj_exact = solve_with(:exact)
+    obj_auto = solve_with(:auto)
+    obj_gp = solve_with(GPSampler(kappa = 4.0, budget = 0.2))
+    @test obj_exact ≈ 10.0 atol = 1e-4
+    # over-M can't raise the optimum, under-M can only shave it a bit
+    @test obj_auto <= obj_exact + 1e-6
+    @test obj_auto ≈ obj_exact atol = 1e-2
+    @test obj_gp <= obj_exact + 1e-6
+    @test obj_gp ≈ obj_exact atol = 1e-2
+end
+
+function test_gp_unknown_sampler_error()
+    model = InfiniteGDPModel(HiGHS.Optimizer)
+    set_silent(model)
+    @infinite_parameter(model, t ∈ [0, 1], num_supports = 5)
+    @variable(model, 0 <= x <= 10, Infinite(t))
+    @variable(model, Y[1:2], InfiniteLogical(t))
+    @parameter_function(model, f == t -> 2*t)
+    @constraint(model, x <= f, Disjunct(Y[1]))
+    @constraint(model, x >= 0.5, Disjunct(Y[2]))
+    @disjunction(model, Y)
+    @objective(model, Max, 𝔼(x, t))
+    @test_throws ErrorException optimize!(model,
+        gdp_method = MBM(HiGHS.Optimizer, M_sampler = :grid))
+end
+
+@testset "InfiniteGPDisjunctiveProgramming" begin
+    test_gp_sampler_creation()
+    test_gp_sampler_resolution()
+    test_gp_raw_M_scalar()
+    test_gp_raw_M_matches_exact()
+    test_gp_mbm_solve_equivalence()
+    test_gp_unknown_sampler_error()
+    test_gp_infeasible_disjunct()
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 06e8813..569a12b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -24,4 +24,5 @@ include("constraints/disjunction.jl")
 include("print.jl")
 include("solve.jl")
 include("extensions/InfiniteDisjunctiveProgramming.jl")
+include("extensions/InfiniteGPDisjunctiveProgramming.jl")
 

From fa4a7ca7530537442e8d56802986b8f5102701a5 Mon Sep 17 00:00:00 2001
From: d227nguyen <d227nguy@uwaterloo.ca>
Date: Sat, 1 Aug 2026 23:39:33 -0400
Subject: [PATCH 3/4] Don't compute grids until we need them

---
 ext/InfiniteDisjunctiveProgramming.jl         | 20 ++++++++++++-----
 ext/InfiniteGPDisjunctiveProgramming.jl       |  4 ++--
 src/extension_api.jl                          |  9 +++++---
 .../InfiniteDisjunctiveProgramming.jl         | 22 +++++++++++++++++++
 4 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/ext/InfiniteDisjunctiveProgramming.jl b/ext/InfiniteDisjunctiveProgramming.jl
index 2055422..37eb8ff 100644
--- a/ext/InfiniteDisjunctiveProgramming.jl
+++ b/ext/InfiniteDisjunctiveProgramming.jl
@@ -263,6 +263,16 @@ function _interpolate_at(
         )
 end
 
+# 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(reverse_map[p]
+                  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
@@ -276,7 +286,7 @@ function DP.sample_M_values(
     objectives::AbstractArray,
     sub::DP.GDPSubmodel,
     method::DP._MBM,
-    grids::Tuple
+    support_grids
     )
     sampler === :exact || error(
         "Unrecognized `M_sampler` `$(repr(sampler))` for MBM on an " *
@@ -308,15 +318,13 @@ function DP.raw_M(
     transcribed = InfiniteOpt.transformation_model(sub.model)
     inner_sub = DP.GDPSubmodel(transcribed, JuMP.VariableRef[],
         Dict{JuMP.VariableRef, Vector{JuMP.VariableRef}}())
-    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)
-    grids = Tuple(InfiniteOpt.supports(p) for p in prefs)
     sampler = _resolve_M_sampler(method.M_sampler)
-    M_vals = DP.sample_M_values(sampler, objectives, inner_sub, method, grids)
+    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)
+    prefs, grids = _support_grids(sub, mini_expr)
     main = JuMP.owner_model(first(prefs))
     param_func = InfiniteOpt.build_parameter_function(
         error, _interpolate(grids, M_vals), prefs)
diff --git a/ext/InfiniteGPDisjunctiveProgramming.jl b/ext/InfiniteGPDisjunctiveProgramming.jl
index 0e9e136..4e93519 100644
--- a/ext/InfiniteGPDisjunctiveProgramming.jl
+++ b/ext/InfiniteGPDisjunctiveProgramming.jl
@@ -90,7 +90,7 @@ function DP.sample_M_values(
     objectives::AbstractArray,
     sub::DP.GDPSubmodel,
     method::DP._MBM,
-    grids::Tuple
+    support_grids
     )
     idxs = collect(CartesianIndices(objectives))
     n = length(idxs)
@@ -108,7 +108,7 @@ function DP.sample_M_values(
     all(==(first(seed)), seed) && return first(seed)
     budget = clamp(
         ceil(Int, sampler.budget * n), min(sampler.min_solves, n), n)
-    X = _support_coords(grids)
+    X = _support_coords(support_grids())
     while length(solved) < budget
         ms, ss = _mean_sd(sampler, X, solved)
         acq = ms .+ sampler.kappa .* ss
diff --git a/src/extension_api.jl b/src/extension_api.jl
index 6bfa129..38273ba 100644
--- a/src/extension_api.jl
+++ b/src/extension_api.jl
@@ -75,13 +75,16 @@ julia> method = MBM(HiGHS.Optimizer, M_sampler = GPSampler(kappa = 4.0))
 function GPSampler end
 
 """
-    sample_M_values(sampler, objectives, sub, method, grids)
+    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 `grids` are the support vectors of the
-infinite parameters. Returns an array of M values shaped like
+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
diff --git a/test/extensions/InfiniteDisjunctiveProgramming.jl b/test/extensions/InfiniteDisjunctiveProgramming.jl
index a2daf8e..3526d62 100644
--- a/test/extensions/InfiniteDisjunctiveProgramming.jl
+++ b/test/extensions/InfiniteDisjunctiveProgramming.jl
@@ -441,6 +441,27 @@ function test_raw_M_infinite_two_params()
     end
 end
 
+# Dependent parameters have no per-parameter support grid, so raw_M
+# must not need one when M does not vary. Setup as in
+# test_raw_M_infinite_scalar, over a dependent parameter array.
+function test_raw_M_infinite_dependent_params()
+    model = InfiniteGDPModel()
+    @infinite_parameter(model, ξ[1:2] ∈ [0, 1], num_supports = 4)
+    @variable(model, 0 <= x <= 10, Infinite(ξ))
+    @variable(model, Y[1:2], InfiniteLogical(ξ))
+    @constraint(model, con, x >= 5, Disjunct(Y[1]))
+    @constraint(model, con2, x <= 3, Disjunct(Y[2]))
+    @disjunction(model, Y)
+    for sampler in (:exact, :auto)
+        mbm = DP._MBM(MBM(HiGHS.Optimizer, M_sampler = sampler), model)
+        sub = DP.copy_model_with_constraints(
+            model, DP.DisjunctConstraintRef[con2], mbm)
+        obj = DP.prepare_max_M_objective(
+            model, JuMP.constraint_object(con), sub)
+        @test DP.raw_M(sub, obj, mbm) == 5.0
+    end
+end
+
 # Piecewise-constant max-of-corners: returns the maximum value over
 # the 2^n corners of the cell containing the query.
 function test_interpolate()
@@ -863,6 +884,7 @@ end
         test_raw_M_infinite_scalar()
         test_raw_M_infinite_param_function()
         test_raw_M_infinite_two_params()
+        test_raw_M_infinite_dependent_params()
         test_mbm_finite_and_integer_var()
         test_mbm_infinite_simple()
         test_mbm_infinite_param_dependent()

From 5a010118e6327aebf71a209d5c0ebea20f05ae09 Mon Sep 17 00:00:00 2001
From: d227nguyen <d227nguy@uwaterloo.ca>
Date: Sun, 2 Aug 2026 13:26:48 -0400
Subject: [PATCH 4/4] detect_uniform_M arg

---
 ext/InfiniteDisjunctiveProgramming.jl         |  7 +-
 ext/InfiniteGPDisjunctiveProgramming.jl       | 42 +++++-----
 src/extension_api.jl                          | 10 ++-
 .../InfiniteGPDisjunctiveProgramming.jl       | 82 ++++++++++++++++++-
 4 files changed, 114 insertions(+), 27 deletions(-)

diff --git a/ext/InfiniteDisjunctiveProgramming.jl b/ext/InfiniteDisjunctiveProgramming.jl
index 37eb8ff..38d0a82 100644
--- a/ext/InfiniteDisjunctiveProgramming.jl
+++ b/ext/InfiniteDisjunctiveProgramming.jl
@@ -268,8 +268,11 @@ end
 # 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(reverse_map[p]
-                  for p in InfiniteOpt.parameter_refs(mini_expr))
+    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
 
diff --git a/ext/InfiniteGPDisjunctiveProgramming.jl b/ext/InfiniteGPDisjunctiveProgramming.jl
index 4e93519..887f837 100644
--- a/ext/InfiniteGPDisjunctiveProgramming.jl
+++ b/ext/InfiniteGPDisjunctiveProgramming.jl
@@ -13,31 +13,30 @@ struct GPSampler{K}
     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
+    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)
+    return GPSampler(Float64(kappa), Float64(budget), min_solves,
+        kernel, detect_uniform_M)
 end
 
 ################################################################################
 #                                 GP FITTING
 ################################################################################
-# Lengthscale candidates on the [0, 1]-normalized support coordinates
 const _LENGTHSCALES = (0.05, 0.1, 0.2, 0.4, 0.8)
-
-# Observation jitter for the GP fit
 const _JITTER = 1e-8
 
-# Coordinates of every support in linear index order, normalized to
-# [0, 1]^d so one isotropic lengthscale works across dimensions
+# 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]
@@ -46,9 +45,7 @@ function _support_coords(grids)
             for I in vec(idxs)]
 end
 
-# Fit the GP posterior on the solved coordinates; `y` is standardized
-# by the caller. With no user kernel, select the lengthscale of a
-# squared exponential kernel by maximizing the marginal likelihood.
+# 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)
@@ -65,13 +62,13 @@ function _fit_posterior(sampler::GPSampler, X, y)
     return best_post
 end
 
-# Posterior mean and sd at all coords given the solved (index => M)
-# samples, destandardized back to M units
 function _mean_sd(sampler::GPSampler, X, solved)
     lis = collect(keys(solved))
     y = [solved[li] for li in lis]
     ybar = sum(y) / length(y)
-    ystd = max(sqrt(sum(abs2, y .- ybar) / max(length(y) - 1, 1)), 1e-8)
+    # 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)
@@ -81,10 +78,7 @@ end
 ################################################################################
 #                              M VALUE SAMPLING
 ################################################################################
-# Solve M at actively-selected supports (max-UCB acquisition) and fill
-# the rest with the upper confidence bound mean + kappa * sd, a
-# heuristic over-estimate. Returns a scalar when the seed M values are
-# uniform (e.g. M does not vary over the supports).
+# Solve M at max-UCB selected supports, fill the rest with the bound
 function DP.sample_M_values(
     sampler::GPSampler,
     objectives::AbstractArray,
@@ -101,11 +95,16 @@ function DP.sample_M_values(
         solved[li] = m
         return true
     end
-    for s in unique([1, cld(n + 1, 2), n])
+    # 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
-    seed = collect(values(solved))
-    all(==(first(seed)), seed) && return first(seed)
+    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())
@@ -125,8 +124,7 @@ function DP.sample_M_values(
         return M_vals
     end
     ms, ss = _mean_sd(sampler, X, solved)
-    for (li, I) in enumerate(idxs)
-        # exact M values are nonnegative, so the fill is too
+    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
diff --git a/src/extension_api.jl b/src/extension_api.jl
index 38273ba..7c220b2 100644
--- a/src/extension_api.jl
+++ b/src/extension_api.jl
@@ -41,7 +41,7 @@ function InfiniteLogical end
 
 """
     GPSampler(; kappa = 2.5, budget = 0.25, min_solves = 6,
-              kernel = nothing)
+              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
@@ -64,6 +64,14 @@ the `M_sampler` field of [`MBM`](@ref)).
 - `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
diff --git a/test/extensions/InfiniteGPDisjunctiveProgramming.jl b/test/extensions/InfiniteGPDisjunctiveProgramming.jl
index 8b702df..796507b 100644
--- a/test/extensions/InfiniteGPDisjunctiveProgramming.jl
+++ b/test/extensions/InfiniteGPDisjunctiveProgramming.jl
@@ -12,13 +12,15 @@ function test_gp_sampler_creation()
     @test sampler.budget == 0.25
     @test sampler.min_solves == 6
     @test isnothing(sampler.kernel)
+    @test sampler.detect_uniform_M
     kern = with_lengthscale(SqExponentialKernel(), 0.3)
-    sampler = GPSampler(
-        kappa = 4.0, budget = 0.1, min_solves = 3, kernel = kern)
+    sampler = GPSampler(kappa = 4.0, budget = 0.1, min_solves = 3,
+        kernel = kern, detect_uniform_M = false)
     @test sampler.kappa == 4.0
     @test sampler.budget == 0.1
     @test sampler.min_solves == 3
     @test sampler.kernel === kern
+    @test !sampler.detect_uniform_M
     @test_throws ErrorException GPSampler(kappa = -1)
     @test_throws ErrorException GPSampler(budget = 0)
     @test_throws ErrorException GPSampler(budget = 1.5)
@@ -137,6 +139,79 @@ function test_gp_mbm_solve_equivalence()
     @test obj_gp ≈ obj_exact atol = 1e-2
 end
 
+# A periodic M must not read as uniform. With f(t) = 2|cos(2*pi*t)|
+# on these supports, M = 10 - f is 8 at supports 1, 3, 5 and 10 at
+# supports 2, 4, so evenly spaced probes alias and collapse M to 8,
+# which caps x at 8 and cuts the optimum from 10 down to 9.
+function test_gp_periodic_M_not_uniform()
+    supports = [0.0, 0.25, 0.5, 0.75, 1.0]
+    function solve_with(M_sampler)
+        model = InfiniteGDPModel(HiGHS.Optimizer)
+        set_silent(model)
+        @infinite_parameter(model, t ∈ [0, 1], supports = supports)
+        @variable(model, 0 <= x <= 10, Infinite(t))
+        @variable(model, Y[1:2], InfiniteLogical(t))
+        @parameter_function(model, f == t -> 2 * abs(cos(2 * pi * t)))
+        @constraint(model, x <= f, Disjunct(Y[1]))
+        @constraint(model, x >= 0.5, Disjunct(Y[2]))
+        @disjunction(model, Y)
+        @objective(model, Max, 𝔼(x, t))
+        optimize!(model, gdp_method = MBM(
+            HiGHS.Optimizer, M_sampler = M_sampler))
+        return objective_value(model)
+    end
+    @test solve_with(:exact) ≈ 10.0 atol = 1e-6
+    @test solve_with(GPSampler()) ≈ 10.0 atol = 1e-6
+end
+
+# With detection off the uniform M is not collapsed to a scalar: the
+# GP is fit and the unsolved supports keep their kappa * sd cushion,
+# which must sit above the M that detection would have returned.
+function test_gp_detect_uniform_M_off()
+    function raw_M_with(detect)
+        model = InfiniteGDPModel()
+        @infinite_parameter(model, t ∈ [0, 1], num_supports = 20)
+        @variable(model, 0 <= x <= 10, Infinite(t))
+        @variable(model, Y[1:2], InfiniteLogical(t))
+        @constraint(model, con, x >= 5, Disjunct(Y[1]))
+        @constraint(model, con2, x <= 3, Disjunct(Y[2]))
+        @disjunction(model, Y)
+        mbm = DP._MBM(MBM(HiGHS.Optimizer,
+            M_sampler = GPSampler(detect_uniform_M = detect)), model)
+        sub = DP.copy_model_with_constraints(
+            model, DP.DisjunctConstraintRef[con2], mbm)
+        obj = DP.prepare_max_M_objective(
+            model, JuMP.constraint_object(con), sub)
+        return DP.raw_M(sub, obj, mbm)
+    end
+    @test raw_M_with(true) == 5.0
+    M = raw_M_with(false)
+    @test M isa InfiniteOpt.GeneralVariableRef
+    raw_fn = InfiniteOpt.raw_function(M)
+    vals = [raw_fn(t) for t in range(0, 1, length = 20)]
+    @test all(vals .>= 5.0 - 1e-6)
+    @test maximum(vals) > 5.0
+end
+
+# Dependent parameters have no support grid, so turning detection off
+# leaves the GP with nothing to fit over
+function test_gp_detect_uniform_M_off_dependent()
+    model = InfiniteGDPModel()
+    @infinite_parameter(model, ξ[1:2] ∈ [0, 1], num_supports = 4)
+    @variable(model, 0 <= x <= 10, Infinite(ξ))
+    @variable(model, Y[1:2], InfiniteLogical(ξ))
+    @constraint(model, con, x >= 5, Disjunct(Y[1]))
+    @constraint(model, con2, x <= 3, Disjunct(Y[2]))
+    @disjunction(model, Y)
+    mbm = DP._MBM(MBM(HiGHS.Optimizer,
+        M_sampler = GPSampler(detect_uniform_M = false)), model)
+    sub = DP.copy_model_with_constraints(
+        model, DP.DisjunctConstraintRef[con2], mbm)
+    obj = DP.prepare_max_M_objective(
+        model, JuMP.constraint_object(con), sub)
+    @test_throws ErrorException DP.raw_M(sub, obj, mbm)
+end
+
 function test_gp_unknown_sampler_error()
     model = InfiniteGDPModel(HiGHS.Optimizer)
     set_silent(model)
@@ -158,6 +233,9 @@ end
     test_gp_raw_M_scalar()
     test_gp_raw_M_matches_exact()
     test_gp_mbm_solve_equivalence()
+    test_gp_periodic_M_not_uniform()
+    test_gp_detect_uniform_M_off()
+    test_gp_detect_uniform_M_off_dependent()
     test_gp_unknown_sampler_error()
     test_gp_infeasible_disjunct()
 end