Building a custom "variable in set" constraint

Hi,

I implemented my own type of sets (let’s call it MySetType) and now I want to override Julia’s in functionality, similar to what is already supported for SecondOrderCones.

So, I would like to make this work with JuMP:

model = Model(...)
@variable(model, x)
@constraint(model, x in X)

where X is of type MySetType.

Currently, I’m using something like this:

# (...)
function Base::in(x, X::T) where {T <: MyAbstractSetType}
   return @build_constraint x .<= X.x_max
end

my_constraints = x in X
add_constraint.(model, my_constraints)
# (...)

This works, but it would be great if I didn’t have to use the add_constraint function and simply go for the @constraint macro. I know this works for MOI.SecondOrderCone, where you can just do the following:

@constraint(model, x in my_second_order_cone)

I tried to find the source code behind this, but failed.
Has someone done this before? Do I somehow need to override the @constraint macro?

Thanks for your help!
Best,
Jannes

Hi @janneshb,

See https://jump.dev/JuMP.jl/stable/developers/extensions/#Define-a-new-set

Hi @odow

thanks for the super quick response.

I looked into the link you send, that was already really helpful.

However, I don’t really understand yet what function I need to override to produce the constraints for my set.
I see that @constraint(model, x in X) calls add_constraint. I that the one I should override? Or moi_set maybe?
Unfortunately, the documentation is a bit sparse on this…

Thanks
Jannes

Sorry, I thought by “implemented my own type of sets” you meant a MOI.AbstractSet. Note that JuMP does not intercept Base.in, nor does it rewrite in SecondOrderCone to be something different. It just calls moi_set(SecondOrderCone(), dim) to get the correct MOI set to pass to the solver.

Can you give an actual minimal working example of the set you want to add?

Why is this not sufficient:

add_constraint.(model, x in X)

or even

function add_my_constraint(model, x, set::MyAbstractSetType)
    return @constraint(model, x .<= set.x_max)
end
add_my_constraint(model, x, set)

I’d recommend that you don’t (ab)use JuMP’s macro code when writing your own function is simpler to implement and maintain.

1 Like

Thanks for following up

Yes,

add_constraint.(model, x in X)

does work when I do something like

function Base::in(x, X::mySet)
    return @build_constraint # ...
end

But if JuMP does not do anything with the in operator, why can’t I overload it to work with @constraint?
I like the @constraint(model, x in X) macro more than the add_constraint function.
So basically my question comes down to: how is @constraint(model, x in SecondOrderCone()) implemented? What does Base::in need to return in order for it to work with @constraint?

The sets I implemented are not something entirely new. Mainly specific polyhedra. I know there are alternatives out there to implementing my own, but I want to do it nonetheless :slight_smile:

Best
Jannes

Because JuMP rewrites it to something else. We don’t call Base.in. We see the symbol :in and rewrite it.

how is @constraint(model, x in SecondOrderCone()) implemented

See @macroexpand

julia> @macroexpand @constraint(model, x in SecondOrderCone())
quote
    #= REPL[2]:1 =#
    JuMP._valid_model(model, :model)
    begin
        #= /Users/oscar/.julia/packages/JuMP/6RAQ9/src/macros.jl:392 =#
        let model = model
            #= /Users/oscar/.julia/packages/JuMP/6RAQ9/src/macros.jl:393 =#
            begin
                #= /Users/oscar/.julia/packages/JuMP/6RAQ9/src/macros/@constraint.jl:171 =#
                var"#1###225" = (MutableArithmetics).copy_if_mutable(x)
                #= /Users/oscar/.julia/packages/JuMP/6RAQ9/src/macros/@constraint.jl:172 =#
                var"#2#build" = JuMP.model_convert(
                    model, 
                    JuMP.build_constraint(
                        JuMP.Containers.var"#error_fn#98"{String}("At REPL[2]:1: `@constraint(model, x in SecondOrderCone())`: "),
                        var"#1###225",
                        SecondOrderCone(),
                    ),
                )
                #= /Users/oscar/.julia/packages/JuMP/6RAQ9/src/macros/@constraint.jl:173 =#
                JuMP.add_constraint(model, var"#2#build", "")
            end
        end
    end
end

Strip away all the line numbers etc, and you get something like:

JuMP._valid_model(model, :model)
lhs = JuMP.MutableArithmetics.copy_if_mutable(x)
constraint = JuMP.build_constraint(error, lhs, SecondOrderCone())
constraint = JuMP.model_convert(model, constraint)
JuMP.add_constraint(model, constraint, "")

This transformation happens because we have implemented parse_constraint_call to intercept the :in and :∈ symbols:

What does Base::in need to return in order for it to work with @constraint ?

We do not call Base.in. The @constraint macro will not call it.

Ah, okay, this makes a lot more sense now.

So, if I want to make this work, I would have to directly override build_constraint with something like this?

function build_constraint(err, x, set::MySet)
    return @build_constraint x .<= set.x_max
end

Thanks a lot for your help, I don’t know how long it would have taken me to figure out the in-intercept.

Jannes

build_constraint will only work if you return an object that JuMP already knows how to add_constraint.

If you get an error, you might have to also implement add_constraint. See Extensions · JuMP and Extensions · JuMP. The fallback error messages should also point you the right direction.

I’m not convinced that this is simpler than just writing a function though.

Perhaps I’ll also add: the JuMP extension hooks could be better documented, but it is also somewhat purposeful. They should really be used only by expert JuMP users. I would strongly encourage people reading this to think about other ways they could structure their code. Messing with the JuMP macros is non-trivial, because the code you see at the top-level like @constraint(model, x in set) has nothing to do with the code that is actually being evaluated.

Yeah, no, I think this creates a lot more overhead than just using add_constraints(model, x, my_set).

Probably, I will just go with that.

Thanks for your help, I really appreciate it :slight_smile:

Best
Jannes

1 Like