JuMP front end: make On(geom, mask) eager, dropping the region union

On stored its region as Union{Vector{Tuple{Int,Int}},Vector{Bool}} and
resolved a Bool mask to (vertex, element) pairs lazily at add_constraint time,
because On(mask) alone lacked the geometry's V. That deferral was the front
end's one non-eager corner. Take the geometry explicitly instead --
On(geom, mask) converts to pairs at construction -- so On.pairs is a single
concrete field again and the extension's deferred-resolution path (_region_pairs,
the unresolved RegionConstraint) is gone. Mask length errors now surface at the
On(...) call rather than inside @constraint.

Author: Sébastien Loisel

docs/src/jump.md

@@ -85,7 +85,7 @@ accepts it directly:
 - `Coef(m, vals::AbstractVector)`
 - `set_start(u, vals::AbstractVector)`
 - `EpiPower(pvals::AbstractVector)`
-- `On(mask::AbstractVector{Bool})` — region membership, one Bool per node
+- `On(geom, mask::AbstractVector{Bool})` — region membership, one Bool per node
 
 Node `i` is vertex `v` of element `e` with `i = v + (e-1)V` where
 `V = size(geom.x, 1)`: the rows of `reshape(geom.x, :, d)` are the node
@@ -109,13 +109,12 @@ say) drops in directly, without wrapping it in an interpolating closure.
 
 ## Regions: constraints on part of the domain
 
-A constraint holds everywhere by default; adding `On(region)` restricts it to
-a node set. The region is either a Bool mask over the nodal vectors (one
-entry per broken node, in the ordering of
-[the data model](@ref jump-data-model)) or a vector of `(vertex, element)`
-pairs — the format of [`find_boundary`](@ref) and the low-level
-`dirichlet_nodes` API; a mask is sugar that resolves to the pair set when the
-constraint is added. Equality + `On` is a Dirichlet condition;
+A constraint holds everywhere by default; adding `On(...)` restricts it to a
+node set of `(vertex, element)` pairs — the format of [`find_boundary`](@ref)
+and the low-level `dirichlet_nodes` API. `On(geom, mask)` is grid-level sugar
+for the same thing: a Bool mask over the nodal vectors (one entry per broken
+node, in the ordering of [the data model](@ref jump-data-model)), converted
+eagerly to the pair set. Equality + `On` is a Dirichlet condition;
 inequality/cone + `On` becomes a piecewise barrier, active only on the
 region. Region *selection* is ordinary data preparation — a comparison on the
 node coordinates gives the mask.
@@ -133,7 +132,7 @@ set_attribute(m2, "verbose", false)
 set_start(s2, 100.0)
 @constraint(m2, u2 == Coef(m2, 0.0), On(find_boundary(geom2)))
 @constraint(m2, [deriv(u2, :dx); deriv(u2, :dy); s2] in EpiPower(2.0))
-@constraint(m2, u2 >= Coef(m2, x -> 0.25 - x[1]^2 - x[2]^2), On(left))
+@constraint(m2, u2 >= Coef(m2, x -> 0.25 - x[1]^2 - x[2]^2), On(geom2, left))
 @objective(m2, Min, integral(Coef(m2, -1.0) * u2 + s2))
 optimize!(m2)
 plot(mgb_solution(m2)); savefig("jump_obstacle.svg"); nothing  # hide

ext/MultiGridBarrierJuMPExt/MultiGridBarrierJuMPExt.jl

@@ -389,32 +389,17 @@ end
 MOI.dimension(s::MOIEpiPower) = s.dim
 JuMP.moi_set(s::EpiPower, dim::Int) = MOIEpiPower(s.p, dim)
 
-# region-restricted constraint: wrap the underlying constraint. Carries the On
-# unresolved because a Bool mask needs the geometry (V) to become pairs, and
-# build_constraint has no model access; add_constraint resolves it.
+# region-restricted constraint: wrap the underlying constraint. On is already
+# resolved to (vertex, element) pairs (masks convert eagerly in On(geom, mask)).
 struct RegionConstraint{C} <: JuMP.AbstractConstraint
     con::C
-    region::On
+    pairs::Vector{Tuple{Int,Int}}
 end
 JuMP.jump_function(rc::RegionConstraint) = JuMP.jump_function(rc.con)
 JuMP.moi_set(rc::RegionConstraint) = JuMP.moi_set(rc.con)
 
-# resolve an On region to (vertex, element) pairs — the package's canonical
-# node-set format (find_boundary, dirichlet_nodes). A Bool mask is grid-level
-# sugar for it: mask entry i = v + (e-1)V selects vertex v of element e.
-function _region_pairs(m::MGBModel, on::On)
-    r = on.region
-    r isa Vector{Bool} || return r
-    length(r) == m.nnodes || _argerror(
-        "region mask has length $(length(r)) but the geometry has $(m.nnodes) " *
-        "broken nodes; entry i is vertex v of element e with i = v + (e-1)V, " *
-        "V = $(size(m.geometry.x, 1))")
-    V = size(m.geometry.x, 1)
-    Tuple{Int,Int}[(mod1(i, V), cld(i, V)) for i in findall(r)]
-end
-
 JuMP.build_constraint(err::Function, f, set, on::On) =
-    RegionConstraint(JuMP.build_constraint(err, f, set), on)
+    RegionConstraint(JuMP.build_constraint(err, f, set), on.pairs)
 
 # JuMP's macro maps generic scalar comparisons (non-Number rhs) to its
 # Zeros/Nonnegatives/Nonpositives shortcut sets through VARIADIC
@@ -424,19 +409,19 @@ function JuMP.build_constraint(err::Function, f, set::JuMP.Zeros, on::On)
     inner = f isa AbstractVector ?
         JuMP.build_constraint(err, f, MOI.Zeros(length(f))) :
         JuMP.build_constraint(err, f, MOI.EqualTo(0.0))
-    RegionConstraint(inner, on)
+    RegionConstraint(inner, on.pairs)
 end
 function JuMP.build_constraint(err::Function, f, set::JuMP.Nonnegatives, on::On)
     inner = f isa AbstractVector ?
         JuMP.build_constraint(err, f, MOI.Nonnegatives(length(f))) :
         JuMP.build_constraint(err, f, MOI.GreaterThan(0.0))
-    RegionConstraint(inner, on)
+    RegionConstraint(inner, on.pairs)
 end
 function JuMP.build_constraint(err::Function, f, set::JuMP.Nonpositives, on::On)
     inner = f isa AbstractVector ?
         JuMP.build_constraint(err, f, MOI.Nonpositives(length(f))) :
         JuMP.build_constraint(err, f, MOI.LessThan(0.0))
-    RegionConstraint(inner, on)
+    RegionConstraint(inner, on.pairs)
 end
 # Newer JuMP versions route scalar comparisons with non-Number sides through
 # internal *Zero marker sets, again variadically; cover them when they exist.
@@ -446,7 +431,7 @@ for (marker, moiset) in ((:GreaterThanZero, :(MOI.GreaterThan(0.0))),
     if isdefined(JuMP, marker)
         @eval function JuMP.build_constraint(err::Function, f,
                                              set::JuMP.$marker, on::On)
-            RegionConstraint(JuMP.build_constraint(err, f, $moiset), on)
+            RegionConstraint(JuMP.build_constraint(err, f, $moiset), on.pairs)
         end
     end
 end
@@ -553,11 +538,10 @@ function JuMP.add_constraint(m::MGBModel, c::JuMP.VectorConstraint,
 end
 function JuMP.add_constraint(m::MGBModel, rc::RegionConstraint, name::String = "")
     c = rc.con
-    pairs = _region_pairs(m, rc.region)
     if c isa JuMP.ScalarConstraint
-        _add_scalar(m, JuMP.jump_function(c), JuMP.moi_set(c), pairs, name)
+        _add_scalar(m, JuMP.jump_function(c), JuMP.moi_set(c), rc.pairs, name)
     elseif c isa JuMP.VectorConstraint
-        _add_vector(m, JuMP.jump_function(c), JuMP.moi_set(c), pairs, name)
+        _add_vector(m, JuMP.jump_function(c), JuMP.moi_set(c), rc.pairs, name)
     else
         _argerror("unsupported constraint type $(typeof(c)) with On(...)")
     end

src/jump_frontend.jl

@@ -106,22 +106,30 @@ function solver_log end
 
 """
     On(pairs::Vector{Tuple{Int,Int}})
-    On(mask::AbstractVector{Bool})
+    On(geom::Geometry, mask::AbstractVector{Bool})
 
-Constraint region for a [`MGBModel`](@ref): a set of broken nodes, as
+Constraint region for a [`MGBModel`](@ref): a set of broken nodes as
 `(vertex, element)` pairs — the same format as [`find_boundary`](@ref) and the
-low-level `dirichlet_nodes` API — or as a Bool mask over the nodal vectors of
-[`Coef`](@ref), where `mask[v + (e-1)V]` selects vertex `v` of element `e`
-(the mask is resolved to the pair set when the constraint is added). The mask
-form composes directly with grid-level data: `On(reshape(geom.x, :, d)[:, 1] .< 0)`
-is the left half of the domain. `@constraint(m, u == g, On(pairs))` is a
-Dirichlet condition; an inequality/cone with `On` holds only on those nodes.
+low-level `dirichlet_nodes` API. The second form is grid-level sugar: a Bool
+mask over the nodal vectors of [`Coef`](@ref), converted eagerly to the pair
+set (`mask[v + (e-1)V]` selects vertex `v` of element `e`; the geometry
+supplies `V`). A mask composes directly with grid-level data:
+`On(geom, reshape(geom.x, :, d)[:, 1] .< 0)` is the left half of the domain.
+`@constraint(m, u == g, On(pairs))` is a Dirichlet condition; an
+inequality/cone with `On` holds only on those nodes.
 """
 struct On
-    region::Union{Vector{Tuple{Int,Int}},Vector{Bool}}
+    pairs::Vector{Tuple{Int,Int}}
     On(pairs::AbstractVector{<:Tuple{Integer,Integer}}) =
         new(Tuple{Int,Int}[(Int(v), Int(e)) for (v, e) in pairs])
-    On(mask::AbstractVector{Bool}) = new(collect(Bool, mask))
+end
+function On(geom::Geometry, mask::AbstractVector{Bool})
+    V, N = size(geom.x, 1), size(geom.x, 2)
+    length(mask) == V * N || throw(ArgumentError(
+        "region mask has length $(length(mask)) but the geometry has $(V * N) " *
+        "broken nodes; entry i is vertex v of element e with i = v + (e-1)V, " *
+        "V = $V"))
+    On(Tuple{Int,Int}[(mod1(i, V), cld(i, V)) for i in findall(mask)])
 end
 
 """

test/test_jump.jl

@@ -146,18 +146,18 @@ _maxdiff(ref, cols...) =
         left = [(v, e) for e in 1:Nn for v in 1:Vn if geom3.x[v, e, 1] < 0]
         mask = reshape(geom3.x, :, 2)[:, 1] .< 0
 
-        function obstacle(region)
+        function obstacle(region::On)
             m = MGBModel(geom3); _quiet!(m)
             @variable(m, u); @variable(m, s, Broken())
             set_start(s, 100.0)
             @constraint(m, u == Coef(m, 0.0), On(find_boundary(geom3)))
             @constraint(m, [deriv(u, :dx); deriv(u, :dy); s] in EpiPower(2.0))
-            @constraint(m, u >= Coef(m, x -> 0.25 - x[1]^2 - x[2]^2), On(region))
+            @constraint(m, u >= Coef(m, x -> 0.25 - x[1]^2 - x[2]^2), region)
             @objective(m, Min, integral(Coef(m, -1.0) * u + s))
             optimize!(m)
             m, u
         end
-        m, u = obstacle(left)
+        m, u = obstacle(On(left))
         @test termination_status(m) == JuMP.MOI.LOCALLY_SOLVED
         @test mgb_solution(m).SOL_feasibility !== nothing   # phase 1 ran
         phi = value(Coef(m, x -> 0.25 - x[1]^2 - x[2]^2))
@@ -167,11 +167,11 @@ _maxdiff(ref, cols...) =
         @test count(<(1e-4), gapL) > 0           # ... with actual contact
         @test minimum(zu[.!mask] .- phi[.!mask]) < 0 # ... and is absent off it
 
-        # the mask is sugar for the same node set
-        m2, u2 = obstacle(mask)
+        # On(geom, mask) is eager sugar for the same node set
+        @test On(geom3, mask).pairs == On(left).pairs
+        m2, u2 = obstacle(On(geom3, mask))
         @test value(u2) == zu
-        @test_throws ArgumentError @constraint(m2, u2 >= Coef(m2, 0.0),
-                                               On(trues(5)))
+        @test_throws ArgumentError On(geom3, trues(5))
     end
 
     @testset "nodal vectors are the fundamental data form" begin