Add support of custom data values for JuMP model elements

src/MOIwrapper.jl

@@ -14,15 +14,17 @@ mutable struct _VarInfo
     index::MOI.VariableIndex
     name::String
     var::Variable
+    data::Union{Nothing, BlockDecomposition.AbstractCustomData}
 end
-_VarInfo(var::Variable) = _VarInfo(_NONE, _NONE, _CONT, MOI.VariableIndex(0), "", var)
+_VarInfo(var::Variable) = _VarInfo(_NONE, _NONE, _CONT, MOI.VariableIndex(0), "", var, nothing)
 
 mutable struct _ConstrInfo
     name::String
     index::Union{Nothing, MOI.ConstraintIndex}
     constr::Constraint
+    data::Union{Nothing, BlockDecomposition.AbstractCustomData}
 end
-_ConstrInfo(constr::Constraint) = _ConstrInfo("", nothing, constr)
+_ConstrInfo(constr::Constraint) = _ConstrInfo("", nothing, constr, nothing)
 
 mutable struct Optimizer <: MOI.AbstractOptimizer
     env::Env
@@ -166,6 +168,8 @@ MOI.supports(::Optimizer, ::MOI.ConstraintDualStart) = false
 MOI.supports(::Optimizer, ::BlockDecomposition.ConstraintDecomposition) = true
 MOI.supports(::Optimizer, ::BlockDecomposition.VariableDecomposition) = true
 MOI.supports(::Optimizer, ::BlockDecomposition.RepresentativeVar) = true
+MOI.supports(::Optimizer, ::BlockDecomposition.CustomVarValue) = true   
+MOI.supports(::Optimizer, ::BlockDecomposition.CustomConstrValue) = true
 
 # Parameters
 function MOI.set(model::Optimizer, param::MOI.RawOptimizerAttribute, val)
@@ -856,6 +860,15 @@ function MOI.set(
     return
 end
 
+function MOI.set(
+    model::Optimizer, ::BD.CustomVarValue, varid::MOI.VariableIndex, custom_data
+)
+    MOI.throw_if_not_valid(model, varid)
+    var = _info(model, varid).var
+    var.custom_data = custom_data
+    return
+end
+
 function MOI.get(model::Optimizer, ::BD.VarBranchingPriority, varid::MOI.VariableIndex)
     var = _info(model, varid).var
     return var.branching_priority
@@ -909,6 +922,16 @@ function MOI.set(
     return
 end
 
+function MOI.set(
+    model::Optimizer, ::BlockDecomposition.CustomConstrValue, constrid::MOI.ConstraintIndex,
+    custom_data
+)
+    MOI.throw_if_not_valid(model, constrid)
+    constr = _info(model, constrid).constr
+    constr.custom_data = custom_data
+    return
+end
+
 function MOI.get(model::Optimizer, ::BD.ConstraintDecomposition, index::MOI.ConstraintIndex)
     MOI.throw_if_not_valid(model, index)
     constrinfo = _info(model, index)

src/MathProg/MathProg.jl

@@ -91,7 +91,8 @@ export Variable, Constraint, VarId, ConstrId, VarMembership, ConstrMembership,
     getperenub, getcurub, setcurub!, getperenrhs, setperenrhs!, getcurrhs, setcurrhs!, getperensense, setperensense!,
     getcursense, setcursense!, getperenkind, getcurkind, setcurkind!, getperenincval,
     getcurincval, setcurincval!, isperenactive, iscuractive, activate!, deactivate!,
-    isexplicit, getname, getbranchingpriority, reset!, getreducedcost, setperenkind!, isfixed, fix!, unfix!
+    isexplicit, getname, getbranchingpriority, reset!, getreducedcost, setperenkind!, isfixed, fix!, unfix!,
+    getcustomdata
 
 # Types & methods related to solutions & bounds
 export PrimalBound, DualBound, AbstractSolution, PrimalSolution, DualSolution, ActiveBound, ObjValues,

src/MathProg/clone.jl

@@ -14,7 +14,8 @@ function clonevar!(
     is_active::Bool = isperenactive(originform, var),
     is_explicit::Bool = isexplicit(originform, var),
     branching_priority::Float64 = getbranchingpriority(originform, var),
-    members::Union{ConstrMembership,Nothing} = nothing
+    members::Union{ConstrMembership,Nothing} = nothing,
+    custom_data = getcustomdata(originform, var)
 )
     id_of_clone = VarId(
         getid(var);
@@ -26,7 +27,7 @@ function clonevar!(
         cost = cost, lb = lb, ub = ub, kind = kind,
         inc_val = inc_val, is_active = is_active, is_explicit = is_explicit, 
         branching_priority = branching_priority, members = members, 
-        id = id_of_clone
+        id = id_of_clone, custom_data = custom_data
     )
 end
 
@@ -44,7 +45,8 @@ function cloneconstr!(
     is_active::Bool = isperenactive(originform, constr),
     is_explicit::Bool = isexplicit(originform, constr),
     members::Union{VarMembership,Nothing}  = nothing,
-    loc_art_var_abs_cost::Float64 = 0.0
+    loc_art_var_abs_cost::Float64 = 0.0,
+    custom_data = getcustomdata(originform, constr)
 )
     id_of_clone = ConstrId(
         getid(constr);
@@ -56,7 +58,7 @@ function cloneconstr!(
         rhs = rhs, kind = kind, sense = sense, inc_val = inc_val,
         is_active = is_active, is_explicit = is_explicit, members = members,
         loc_art_var_abs_cost = loc_art_var_abs_cost, 
-        id = id_of_clone
+        id = id_of_clone, custom_data = custom_data
     )
 end
 

src/MathProg/varconstr.jl

@@ -704,6 +704,21 @@ Return the branching priority of a variable
 getbranchingpriority(form::Formulation, varid::VarId) = getvar(form, varid).branching_priority
 getbranchingpriority(::Formulation, var::Variable) = var.branching_priority
 
+
+"""
+    getcustomdata(formulation, var)
+    getcustomdata(formulation, varid)
+    getcustomdata(formulation, constr)
+    getcustomdata(formulation, constrid)
+
+Return the custom data of a variable or a constraint in a formulation.
+"""
+getcustomdata(form::Formulation, varid::VarId) = getcustomdata(form, getvar(form, varid))
+getcustomdata(::Formulation, var::Variable) = var.custom_data
+getcustomdata(form::Formulation, constrid::ConstrId) = getcustomdata(form, getconstr(form, constrid))
+getcustomdata(::Formulation, constr::Constraint) = constr.custom_data
+
+
 # Reset (this method is used only in tests... @guimarqu doesn't know if we should keep it)
 """
     reset!(form, var)

test/integration/custom_data/attach_custom_data.jl

@@ -0,0 +1,50 @@
+struct TestAttachCustomDataAlgorithm end
+
+struct CustomVarData <: BD.AbstractCustomData
+    var_value::Int
+end
+
+struct CustomConstrData <: BD.AbstractCustomData
+    constr_value::Int
+end
+
+function Coluna.Algorithm.run!(::TestAttachCustomDataAlgorithm, _, form, _)
+    vars = Dict{String, Coluna.MathProg.Variable}()
+    for (_, var) in Coluna.MathProg.getvars(form)
+        vars[getname(form, var)] = var
+    end
+
+    constrs = Dict{String, Coluna.MathProg.Constraint}()
+    for (_, constr) in Coluna.MathProg.getconstrs(form)
+        constrs[getname(form, constr)] = constr
+    end
+
+    @test Coluna.MathProg.getcustomdata(form, vars["x[1]"]).var_value == 1
+    @test Coluna.MathProg.getcustomdata(form, vars["x[2]"]).var_value == 2
+    @test Coluna.MathProg.getcustomdata(form, constrs["c"]).constr_value == 3
+    return Coluna.Algorithm.OptimizationState(form)
+end
+
+function attach_custom_data()
+    coluna = JuMP.optimizer_with_attributes(
+        Coluna.Optimizer,
+        "params" => CL.Params(
+            solver = TestAttachCustomDataAlgorithm()
+        ),
+        "default_optimizer" => GLPK.Optimizer,
+    )
+
+    model = BlockModel(coluna)
+    @variable(model, x[1:2], Bin)
+    @constraint(model, c, x[1] + x[2] <= 1)
+    @objective(model, Max, 2x[1] + 3x[2])
+
+    customvars!(model, CustomVarData)
+    customconstrs!(model, CustomConstrData)
+
+    customdata!(x[1], CustomVarData(1))
+    customdata!(x[2], CustomVarData(2))
+    customdata!(c, CustomConstrData(3))
+    optimize!(model)
+end
+register!(integration_tests, "attach_custom_data", attach_custom_data)

test/integration/run.jl

@@ -0,0 +1,12 @@
+dirs = [
+    "custom_data"
+]
+
+for dir in dirs
+    dirpath = joinpath(@__DIR__, dir)
+    for filename in readdir(dirpath)
+        include(joinpath(dirpath, filename))
+    end
+end
+
+run_integration_tests() = run_tests(integration_tests)

test/revise.jl

@@ -63,9 +63,11 @@ for file in all_test_files
 end
 
 include("unit/run.jl")
+include("integration/run.jl")
 include("e2e/run.jl")
 
 listen_to_tests([
-    run_unit_tests,
+    #run_unit_tests,
+    run_integration_tests,
     #run_e2e_tests
 ])

test/runtests.jl

@@ -28,6 +28,7 @@ using .TestRegistry
 unit_tests = Registry()
 include("parser.jl")
 
+integration_tests = Registry()
 e2e_tests = Registry()
 
 const MODULES = [
@@ -46,8 +47,11 @@ if !isempty(ARGS)
     include("revise.jl")
 else
     include("unit/run.jl")
+    include("integration/run.jl")
     include("e2e/run.jl")
+
     run_unit_tests()
+    run_integration_tests()
 
     @testset "MOI integration" begin
         include("MathOptInterface/MOI_wrapper.jl")