From ea49afd8f4a808f372f007a7fb5c33e01d170dfc Mon Sep 17 00:00:00 2001 From: pulsipher Date: Mon, 6 Nov 2023 12:41:34 -0500 Subject: [PATCH 1/3] add support for optimze hooks --- src/TranscriptionOpt/optimize.jl | 11 ++++++--- src/datatypes.jl | 14 +++++++---- src/optimize.jl | 41 +++++++++++++++++++++++++++++--- test/optimizer.jl | 21 ++++++++++++++++ 4 files changed, 76 insertions(+), 11 deletions(-) diff --git a/src/TranscriptionOpt/optimize.jl b/src/TranscriptionOpt/optimize.jl index 89889dbfa..6094ac221 100644 --- a/src/TranscriptionOpt/optimize.jl +++ b/src/TranscriptionOpt/optimize.jl @@ -4,7 +4,7 @@ Return the transcription model stored in `model` if that is what is stored in `model.optimizer_model`. """ -function transcription_model(model::InfiniteOpt.InfiniteModel)::JuMP.Model +function transcription_model(model::InfiniteOpt.InfiniteModel) trans_model = InfiniteOpt.optimizer_model(model) if !is_transcription_model(trans_model) error("The model does not contain a transcription model.") @@ -26,8 +26,13 @@ using [`build_transcription_model!`](@ref). function InfiniteOpt.build_optimizer_model!( model::InfiniteOpt.InfiniteModel, key::Val{:TransData}; - check_support_dims::Bool = true - )::Nothing + check_support_dims::Bool = true, + extra_kwargs... + ) + # throw error for extra keywords + for (kw, _) in extra_kwargs + error("Unrecognized keyword argument `$kw` for building transcription models.") + end # clear the optimzier model contents trans_model = InfiniteOpt.clear_optimizer_model_build!(model) # build the transcription model based on model diff --git a/src/datatypes.jl b/src/datatypes.jl index fcca42bc9..3c8a49468 100644 --- a/src/datatypes.jl +++ b/src/datatypes.jl @@ -1316,6 +1316,7 @@ mutable struct InfiniteModel <: JuMP.AbstractModel # Extensions ext::Dict{Symbol, Any} + optimize_hook::Any end """ @@ -1402,7 +1403,8 @@ function InfiniteModel(; # Optimize data nothing, OptimizerModel(; kwargs...), false, # Extensions - Dict{Symbol, Any}() + Dict{Symbol, Any}(), + nothing ) end @@ -1411,13 +1413,13 @@ end function _set_optimizer_constructor( model::InfiniteModel, constructor::MOI.OptimizerWithAttributes - )::Nothing + ) model.optimizer_constructor = constructor.optimizer_constructor return end # No attributes -function _set_optimizer_constructor(model::InfiniteModel, constructor)::Nothing +function _set_optimizer_constructor(model::InfiniteModel, constructor) model.optimizer_constructor = constructor return end @@ -1427,7 +1429,7 @@ function InfiniteModel( optimizer_constructor; OptimizerModel::Function = TranscriptionModel, kwargs... - )::InfiniteModel + ) model = InfiniteModel() model.optimizer_model = OptimizerModel(optimizer_constructor; kwargs...) _set_optimizer_constructor(model, optimizer_constructor) @@ -1436,6 +1438,7 @@ end # Define basic InfiniteModel extension functions Base.broadcastable(model::InfiniteModel) = Ref(model) +JuMP.variable_ref_type(::Type{InfiniteModel}) = GeneralVariableRef """ JuMP.object_dictionary(model::InfiniteModel)::Dict{Symbol, Any} @@ -1446,7 +1449,7 @@ registered to a specific symbol in the macros. For example, `@variable(model, x[1:2, 1:2])` registers the array of variables `x` to the symbol `:x`. """ -JuMP.object_dictionary(model::InfiniteModel)::Dict{Symbol, Any} = model.obj_dict +JuMP.object_dictionary(model::InfiniteModel) = model.obj_dict """ Base.empty!(model::InfiniteModel)::InfiniteModel @@ -1492,6 +1495,7 @@ function Base.empty!(model::InfiniteModel) empty!(model.optimizer_model) model.ready_to_optimize = false empty!(model.ext) + model.optimize_hook = nothing return model end diff --git a/src/optimize.jl b/src/optimize.jl index 11bda01cc..01d0d84a3 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -903,13 +903,42 @@ end ################################################################################ # OPTIMIZATION METHODS ################################################################################ +""" + JuMP.set_optimize_hook( + model::InfiniteModel, + hook::Union{Function, Nothing} + )::Nothing + +Set the function `hook` as the optimize hook for `model` where `hook` should +have be of the form `hook(model::GenericModel; hook_specfic_kwargs..., kwargs...)`. +The `kwargs` are those passed to [`optimize!`](@ref). The `hook_specifc_kwargs` +are passed as additional keywords by the user when they call [`optimize!`](@ref). + +## Notes + +* The optimize hook should generally modify the model, or some external state +in some way, and then call `optimize!(model; ignore_optimize_hook = true)` to +optimize the problem, bypassing the hook. +* Use `set_optimize_hook(model, nothing)` to unset an optimize hook. +""" +function JuMP.set_optimize_hook( + model::InfiniteModel, + hook::Union{Function, Nothing} + ) + model.optimize_hook = hook + set_optimizer_model_ready(model, false) + return +end + """ JuMP.optimize!(model::InfiniteModel; [kwargs...]) Extend `JuMP.optimize!` to optimize infinite models using the internal -optimizer model. Will call [`build_optimizer_model!`](@ref) if the optimizer +optimizer model. Calls [`build_optimizer_model!`](@ref) if the optimizer model isn't up to date. The `kwargs` correspond to keyword arguments passed to -[`build_optimizer_model!`](@ref) if any are defined. +[`build_optimizer_model!`](@ref) if any are defined. The `kwargs` can also +include arguments that are passed to an optimize hook if one was set with +[`JuMP.set_optimize_hook`](@ref). **Example** ```julia-repl @@ -919,7 +948,13 @@ julia> has_values(model) true ``` """ -function JuMP.optimize!(model::InfiniteModel; kwargs...) +function JuMP.optimize!( + model::InfiniteModel; + ignore_optimize_hook = isnothing(model.optimize_hook), + kwargs...) + if !ignore_optimize_hook + return model.optimize_hook(model; kwargs...) + end if !optimizer_model_ready(model) build_optimizer_model!(model; kwargs...) end diff --git a/test/optimizer.jl b/test/optimizer.jl index f33c82a53..0db9b98dc 100644 --- a/test/optimizer.jl +++ b/test/optimizer.jl @@ -21,6 +21,8 @@ @constraint(m, c4, meas2 - 2y0 + x <= 1, DomainRestrictions(par => [0.5, 1])) @constraint(m, c5, meas2 == 0) @objective(m, Min, x0 + meas1) + # test extra keywords + @test_throws ErrorException build_optimizer_model!(m, bad = 42) # test normal usage @test isa(build_optimizer_model!(m), Nothing) @test optimizer_model_ready(m) @@ -181,6 +183,25 @@ end @test isa(optimize!(m, check_support_dims = false), Nothing) @test optimizer_model_ready(m) @test num_variables(optimizer_model(m)) == 8 + # test optimize hook + function myhook(model; n = "", ub = 2, kwargs...) + if !isempty(n) + var = variable_by_name(model, n) + set_upper_bound(var, ub) + end + optimize!(model; ignore_optimize_hook = true, kwargs...) + return + end + @test set_optimize_hook(m, myhook) isa Nothing + @test optimize!(m, n = "x", check_support_dims = false) isa Nothing + @test optimizer_model_ready(m) + @test num_variables(optimizer_model(m)) == 8 + @test upper_bound(x) == 2 + @test set_optimize_hook(m, nothing) isa Nothing + @test isnothing(m.optimize_hook) + @test_throws ErrorException optimize!(m, n = "x") + @test optimize!(m) isa Nothing + @test optimizer_model_ready(m) end # Test JuMP.result_count From 7decf01a93e0a0140faa412159fef43f9a1aad0d Mon Sep 17 00:00:00 2001 From: pulsipher Date: Mon, 6 Nov 2023 12:43:44 -0500 Subject: [PATCH 2/3] add docstring to docs --- docs/src/manual/model.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/manual/model.md b/docs/src/manual/model.md index 5d9a47083..67a0bcd8f 100644 --- a/docs/src/manual/model.md +++ b/docs/src/manual/model.md @@ -9,6 +9,7 @@ InfiniteModel() JuMP.object_dictionary(::InfiniteModel) has_internal_supports Base.empty!(::InfiniteModel) +JuMP.set_optimize_hook(::InfiniteModel, ::Union{Function, Nothing}) ``` ## Abstract Dependencies From a914d98398dc61a438fe8dd9b6434a09d277635f Mon Sep 17 00:00:00 2001 From: pulsipher Date: Mon, 6 Nov 2023 14:26:24 -0500 Subject: [PATCH 3/3] fix coverage --- test/datatypes.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/datatypes.jl b/test/datatypes.jl index 0165a7eb7..0cc633035 100644 --- a/test/datatypes.jl +++ b/test/datatypes.jl @@ -198,6 +198,8 @@ end @test InfiniteOpt._param_object_indices(m) isa Vector{Union{IndependentParameterIndex, DependentParametersIndex}} # test other methods @test empty!(InfiniteModel(mockoptimizer)).optimizer_constructor == mockoptimizer + @test variable_ref_type(InfiniteModel) == GeneralVariableRef + @test variable_ref_type(InfiniteModel()) == GeneralVariableRef end # Test reference variable datatypes