Extensions

JuMP provides a variety of ways to extend the basic modeling functionality.

Compatibility

When writing JuMP extensions, you should carefully consider the compatibility guarantees that JuMP makes. In particular:

• All functions, structs, and constants which do not begin with an underscore (_) are public. These are always safe to use, and they should all have corresponding documentation.
• All identifiers beginning with an underscore (_) are private. These are not safe to use, because they may break in any JuMP release, including patch releases.
• Unless explicitly mentioned in the documentation, all fields of a struct are private. These are not safe to use, because they may break in any JuMP release, including patch releases. An example of a field which is safe to use is the model.ext extension dictionary, which is documented in The extension dictionary.

In general, we strongly encourage you to use only the public API of JuMP. If you are missing a feature, please open a GitHub issue.

However, if you do use the private API (for example, because your feature request has not been implemented yet), then you must carefully restrict the versions of JuMP that your package is compatible with in the Project.toml file. The easiest way to do this is via the hyphen specifiers. For example, if your package supports all JuMP versions between v1.0.0 and v1.1.1, do:

JuMP = "1.0.0 - 1.1.1"

Then, whenever JuMP releases a new version, you should check if your package is still compatible and update the bound accordingly.

Define a new set

To define a new set for JuMP, subtype MOI.AbstractScalarSet or MOI.AbstractVectorSet and implement Base.copy for the set.

julia> struct NewMOIVectorSet <: MOI.AbstractVectorSet
           dimension::Int
       end

julia> Base.copy(x::NewMOIVectorSet) = x

julia> model = Model();

julia> @variable(model, x[1:2]);

julia> @constraint(model, x in NewMOIVectorSet(2))
[x[1], x[2]] ∈ NewMOIVectorSet(2)

However, for vector-sets, this requires the user to specify the dimension argument to their set, even though we could infer it from the length of x!

You can make a more user-friendly set by subtyping AbstractVectorSet and implementing moi_set.

julia> struct NewVectorSet <: JuMP.AbstractVectorSet end

julia> JuMP.moi_set(::NewVectorSet, dim::Int) = NewMOIVectorSet(dim)

julia> @constraint(model, x in NewVectorSet())
[x[1], x[2]] ∈ NewMOIVectorSet(2)

Extend @variable

Just as Bin and Int create binary and integer variables, you can extend the @variable macro to create new types of variables. Here is an explanation by example, where we create a AddTwice type, that creates a tuple of two JuMP variables instead of a single variable.

First, create a new struct. This can be anything. Our struct holds a VariableInfo object that stores bound information, and whether the variable is binary or integer.

julia> struct AddTwice
           info::JuMP.VariableInfo
       end

Second, implement build_variable, which takes ::Type{AddTwice} as an argument, and returns an instance of AddTwice. Note that you can also receive keyword arguments.

julia> function JuMP.build_variable(
           _err::Function,
           info::JuMP.VariableInfo,
           ::Type{AddTwice};
           kwargs...
       )
           println("Can also use $kwargs here.")
           return AddTwice(info)
       end

Third, implement add_variable, which takes the instance of AddTwice from the previous step, and returns something. Typically, you will want to call add_variable here. For example, our AddTwice call is going to add two JuMP variables.

julia> function JuMP.add_variable(
           model::JuMP.Model,
           duplicate::AddTwice,
           name::String,
       )
           a = JuMP.add_variable(
               model,
               JuMP.ScalarVariable(duplicate.info),
               "$(name)_a",
            )
           b = JuMP.add_variable(
               model,
               JuMP.ScalarVariable(duplicate.info),
               "$(name)_b",
            )
           return (a, b)
       end

Now AddTwice can be passed to @variable similar to Bin or Int, or through the variable_type keyword. However, now it adds two variables instead of one.

julia> model = Model();

julia> @variable(model, x[i=1:2], variable_type = AddTwice, kw = i)
Can also use Base.Pairs(:kw => 1) here.
Can also use Base.Pairs(:kw => 2) here.
2-element Vector{Tuple{VariableRef, VariableRef}}:
 (x[1]_a, x[1]_b)
 (x[2]_a, x[2]_b)

julia> num_variables(model)
4

julia> first(x[1])
x[1]_a

julia> last(x[2])
x[2]_b

Extend @constraint

The @constraint macro has three steps that can be intercepted and extended: parse time, build time, and add time.

Parse

To extend the @constraint macro at parse time, implement one of the following methods:

• parse_constraint_head
• parse_constraint_call

parse_constraint_head should be implemented to intercept an expression based on the .head field of Base.Expr. For example:

julia> using JuMP

julia> const MutableArithmetics = JuMP._MA;

julia> model = Model(); @variable(model, x);

julia> function JuMP.parse_constraint_head(
           error_fn::Function,
           ::Val{:≔},
           lhs,
           rhs,
       )
           println("Rewriting ≔ as ==")
           new_lhs, parse_code = MutableArithmetics.rewrite(lhs)
           build_code = :(
               build_constraint($(error_fn), $(new_lhs), MOI.EqualTo($(rhs)))
           )
           return false, parse_code, build_code
       end

julia> @constraint(model, x + x ≔ 1.0)
Rewriting ≔ as ==
2 x = 1

parse_constraint_call should be implemented to intercept an expression of the form Expr(:call, op, args...). For example:

julia> using JuMP

julia> const MutableArithmetics = JuMP._MA;

julia> model = Model(); @variable(model, x);

julia> function JuMP.parse_constraint_call(
           error_fn::Function,
           is_vectorized::Bool,
           ::Val{:my_equal_to},
           lhs,
           rhs,
       )
           println("Rewriting my_equal_to to ==")
           new_lhs, parse_code = MutableArithmetics.rewrite(lhs)
           build_code = if is_vectorized
               :(build_constraint($(error_fn), $(new_lhs), MOI.EqualTo($(rhs)))
           )
           else
               :(build_constraint.($(error_fn), $(new_lhs), MOI.EqualTo($(rhs))))
           end
           return parse_code, build_code
       end

julia> @constraint(model, my_equal_to(x + x, 1.0))
Rewriting my_equal_to to ==
2 x = 1

To prevent JuMP from promoting the set to the same value type as the model, use SkipModelConvertScalarSetWrapper.

Build

To extend the @constraint macro at build time, implement a new build_constraint method.

This may mean implementing a method for a specific function or set created at parse time, or it may mean implementing a method which handles additional positional arguments.

build_constraint must return an AbstractConstraint, which can either be an AbstractConstraint already supported by JuMP, for example, ScalarConstraint or VectorConstraint, or a custom AbstractConstraint with a corresponding add_constraint method (see Add).

Here is an example of adding extra arguments to build_constraint:

julia> model = Model(); @variable(model, x);

julia> struct MyConstrType end

julia> function JuMP.build_constraint(
            error_fn::Function,
            f::JuMP.GenericAffExpr,
            set::MOI.EqualTo,
            extra::Type{MyConstrType};
            d = 0,
       )
            new_set = MOI.LessThan(set.value + d)
            return JuMP.build_constraint(error_fn, f, new_set)
       end

julia> @constraint(model, my_con, x == 0, MyConstrType, d = 2)
my_con : x ≤ 2

Add

build_constraint returns an AbstractConstraint object. To extend @constraint at add time, define a subtype of AbstractConstraint, implement build_constraint to return an instance of the new type, and then implement add_constraint.

Here is an example:

julia> model = Model(); @variable(model, x);

julia> struct MyTag
           name::String
       end

julia> struct MyConstraint{S} <: AbstractConstraint
           name::String
           f::AffExpr
           s::S
       end

julia> function JuMP.build_constraint(
            error_fn::Function,
            f::AffExpr,
            set::MOI.AbstractScalarSet,
            extra::MyTag,
       )
            return MyConstraint(extra.name, f, set)
       end

julia> function JuMP.add_constraint(
            model::Model,
            con::MyConstraint,
            name::String,
       )
            return add_constraint(
                model,
                ScalarConstraint(con.f, con.s),
                "$(con.name)[$(name)]",
            )
       end

julia> @constraint(model, my_con, 2x <= 1, MyTag("my_prefix"))
my_prefix[my_con] : 2 x - 1 ≤ 0

The extension dictionary

Every JuMP model has a field .ext::Dict{Symbol,Any} that can be used by extensions. This is useful if your extensions to @variable and @constraint need to store information between calls.

The most common way to initialize a model with information in the .ext dictionary is to provide a new constructor:

julia> function MyModel()
           model = Model()
           model.ext[:MyModel] = 1
           return model
       end
MyModel (generic function with 1 method)

julia> model = MyModel()
A JuMP Model
├ solver: none
├ objective_sense: FEASIBILITY_SENSE
├ num_variables: 0
├ num_constraints: 0
└ Names registered in the model: none

julia> model.ext
Dict{Symbol, Any} with 1 entry:
  :MyModel => 1

If you define extension data, implement copy_extension_data to support copy_model.

Defining new JuMP models

If extending individual calls to @variable and @constraint is not sufficient, it is possible to implement a new model via a subtype of AbstractModel. You can also define new AbstractVariableRefs to create different types of JuMP variables.

If you define new types, you will need to implement a considerable number of methods, and doing so will require a detailed understanding of the JuMP internals. Therefore, the list of methods to implement is currently undocumented.

The easiest way to extend JuMP by defining a new model type is to follow an existing example. A simple example to follow is the JuMPExtension module in the JuMP test suite. The best example of an external JuMP extension that implements an AbstractModel is InfiniteOpt.jl.

Testing JuMP extensions

The JuMP test suite contains a large number of tests for JuMP extensions. You can run these tests by copying the MIT-licensed Kokako.jl file in the JuMP tests into your /test folder, and then adding this snippet to your /test/runtests.jl file:

using MyJuMPExtension
import JuMP
include("Kokako.jl")
const MODULES_TO_TEST = Kokako.include_modules_to_test(JuMP)
Kokako.run_tests(
    MODULES_TO_TEST,
    MyJuMPExtension.MyModel,
    MyJuMPExtension.MyVariableRef;
    test_prefix = "test_extension_",
)

Set an optimize! hook

Some extensions require modification to the problem after the user has finished constructing the problem, but before optimize! is called. For these situations, JuMP provides set_optimize_hook, which lets you intercept the optimize! call.

Here's a simple example of adding an optimize hook that extends optimize! to take a keyword argument silent:

julia> using JuMP, HiGHS

julia> model = Model(HiGHS.Optimizer);

julia> @variable(model, x >= 1.5, Int);

julia> @objective(model, Min, x);

julia> function silent_hook(model; silent::Bool)
           if silent
               set_silent(model)
           else
               unset_silent(model)
           end
           ## Make sure you set ignore_optimize_hook = true, or we'll
           ## recursively enter the optimize hook!
           return optimize!(model; ignore_optimize_hook = true)
       end
silent_hook (generic function with 1 method)

julia> set_optimize_hook(model, silent_hook)
silent_hook (generic function with 1 method)

julia> optimize!(model; silent = true)

julia> optimize!(model; silent = false)
Coefficient ranges:
  Cost   [1e+00, 1e+00]
  Bound  [2e+00, 2e+00]
Assessing feasibility of MIP using primal feasibility and integrality tolerance of       1e-06
Solution has               num          max          sum
Col     infeasibilities      0            0            0
Integer infeasibilities      0            0            0
Row     infeasibilities      0            0            0
Row     residuals            0            0            0
Presolving model
0 rows, 0 cols, 0 nonzeros  0s
0 rows, 0 cols, 0 nonzeros  0s
Presolve: Optimal

        Nodes      |    B&B Tree     |            Objective Bounds              |  Dynamic Constraints |       Work
     Proc. InQueue |  Leaves   Expl. | BestBound       BestSol              Gap |   Cuts   InLp Confl. | LpIters     Time

         0       0         0   0.00%   2               2                  0.00%        0      0      0         0     0.0s

Solving report
  Status            Optimal
  Primal bound      2
  Dual bound        2
  Gap               0% (tolerance: 0.01%)
  Solution status   feasible
                    2 (objective)
                    0 (bound viol.)
                    0 (int. viol.)
                    0 (row viol.)
  Timing            0.00 (total)
                    0.00 (presolve)
                    0.00 (postsolve)
  Nodes             0
  LP iterations     0 (total)
                    0 (strong br.)
                    0 (separation)
                    0 (heuristics)

Creating new container types

JuMP macros (for example, @variable) accept a container keyword argument to force the type of container that is chosen. By default, JuMP supports container = Array, container = DenseAxisArray, container = SparseAxisArray and container = Auto. You can extend support to user-defined types by implementing Containers.container.

For example, here is a container that reverses the order of the indices:

julia> struct Foo end

julia> function Containers.container(f::Function, indices, ::Type{Foo})
           return reverse([f(i...) for i in indices])
       end

julia> model = Model();

julia> @variable(model, x[1:3], container = Foo)
3-element Vector{VariableRef}:
 x[3]
 x[2]
 x[1]

julia> x[1]
x[3]

julia> @variable(model, y[1:3, 1:2], container = Foo)
3×2 Matrix{VariableRef}:
 y[3,2]  y[3,1]
 y[2,2]  y[2,1]
 y[1,2]  y[1,1]

julia> y[1, 1]
y[3,2]

julia> @variable(model, z[i=1:3; isodd(i)], container = Foo)
2-element Vector{VariableRef}:
 z[3]
 z[1]

julia> z[2]
z[1]

julia> model = Model();

julia> @variable(model, x[1:3])
3-element Vector{VariableRef}:
 x[1]
 x[2]
 x[3]

julia> y = reverse(x)
3-element Vector{VariableRef}:
 x[3]
 x[2]
 x[1]

Performance tips for extensions

The function-in-set design of MathOptInterface causes type stability issues in Julia if you try to iterate over all of the constraints in a model. The easiest way to fix this is to use a function barrier.

For example, instead of:

function all_names_slow(model)
    names = Set{String}()
    for ci in all_constraints(model)
        push!(names, name(ci))
    end
    return names
end

use:

function _function_barrier(names, model, ::Type{F}, ::Type{S}) where {F,S}
    for ci in all_constraints(model, F, S)
        push!(names, name(ci))
    end
    return
end

function all_names_fast(model)
    names = Set{String}()
    for (F, S) in list_of_constraint_types(model)
        _function_barrier(names, model, F, S)
    end
    return names
end