# Modular implementation of optimization problems with JuMP

**URL:** <https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261>\
**Category:** Optimization (Mathematical)\
**Tags:** jump, energy\
**Created:** [June 26, 2024, 2:06pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261 "2024-06-26T14:06:27Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [June 26, 2024, 2:06pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/1 "2024-06-26T14:06:27Z")

</div>

Hello,

This is a follow-up of my question on the thread about dynamic JuMP variable names [Using anonymous and or symbols to create dynamic variable names? - #3 by pierre-haessig](https://discourse.julialang.org/t/using-anonymous-and-or-symbols-to-create-dynamic-varible-names/83095/3).

I’m sharing here a JuMP code design study on how to implement single/multi stage variants of the same optimization problem in a **modular way** , i.e. sharing as much common model code as possible.

I’m interested in your thoughts on this code design.

## Motivation

In my use case, modularity is welcomed because each stage _has more or less the same variables/constraints_. This differs from many exemplar two stage programs, such as the farmer example implemented in [InfiniteOpt.jl/Two-Stage Stochastic Program](https://infiniteopt.github.io/InfiniteOpt.jl/stable/examples/Stochastic%20Optimization/farmer/#Two-Stage-Stochastic-Program), where first and second stage variables are different (and named as such x vs y(\xi)).

## Study notebook

It would be a bit long to paste the example, so here is the link to the Jupyter notebook in a Gist. Notes:

- The example is a toy energy system optimization problem, i.e. with optimal sizing and operation jointly solved. Since there is only one power source, the trivial solution is that `Pgen` should be the maximum of the requested `load` time series, i.e. the peak load.
- The two-stage aspect comes from considering two subsequent periods, that is allowing a mid-project generator resize.
- The stochastic aspect comes from the second stage load being uncertain, so different resize values are to be considered.

Also notice that I arrived at two different code designs (“A” and “B”, named chronologically) and in the end I prefer the simplicity of option “B”.

> <https://gist.github.com/pierre-haessig/79358e7c45e7ab9954938188162d92b6>

## Summary of “Modular implementation B”

The key building block is the inner stage problem wrapped in a function (`m` is JuMP `Model`, `dat` is a `Dict{String,Any}` to store variables)

```julia
function optim_model_stage_B!(m, dat, load; stage="")
    K = length(load)

    # Variables
    Pgen = add_stage_variable_B!(m, dat, "Pgen"; stage=stage)
    gen = add_stage_variable_B!(m, dat, "gen"; K=K, stage=stage)
    
    # Constraints:
    dat["balance"] = @constraint(m, [k=1:K], gen[k] == load[k], base_name="balance$stage")
    dat["Pgen_cons"] = @constraint(m, [k=1:K], gen[k] <= Pgen, base_name="Pgen_cons$stage")

    # scenario cost
    dat["cost"] = Pgen + sum(gen)
end

```

This uses an `add_stage_variable_B!` wrapper of `@variable` macro to avoid having name conflict in the JuMP `Model`, in the spirit of [Using anonymous and or symbols to create dynamic variable names? - #3 by pierre-haessig](https://discourse.julialang.org/t/using-anonymous-and-or-symbols-to-create-dynamic-varible-names/83095/3). Instead of using the JuMP `Model`, variables are stored in the `Dict`.

### About _not_ wrapping `@constraint`

Notice that in the above example, I use a wrapper for `@variable`, but not for `@constraint`.  
I find that the code looks a bit unbalanced (asymetrical) in that aspect, but I was unable to create a constraint wrapper allowing the large formulation flexibility permitted by `@constraint`. As an example consider “Pgen\_cons” above which is a vector constraint mixing vector and scalar variables.

### Single stage usage

This stage problem can be used as a single stage (i.e. a classical deterministic energy system sizing), after setting the objective (notice `@objective` macro is not called inside `optim_model_stage_B!`, but rather defines `dat["cost"]` which is only a fragment of the global cost in a multi-stage context):

```julia
m = Model(HiGHS.Optimizer)
data = Dict{String,Any}()
load = [1., 1.5, 1.2, 0.9] # dummy problem data

optim_model_stage_B!(m, data, load)

@objective(m, Min, data["cost"])
optimize!(m)
Pgen_val = value(data["Pgen"])
println("Pgen=$(Pgen_val)") # Pgen=1.5

```

### Two-stage stochastic program usage

Remember that the goal of this code design study is that the _inner stage can be reused for a larger two-stage problem_. And indeed, 1+n\_s stage instances are created, that is 1 for the first stage, and 1 for each of the n\_s stochastic scenarios for the second stage.

```julia
# second stage data: 3 equally probable load changes
load_scenar = [0.5, 1.0, 2.0] # (decrease, constant, increase)
ns = length(load_scenar)
proba_scenar = [1., 1., 1.]/ns

m = Model(HiGHS.Optimizer)

data1 = Dict{String,Any}()
optim_model_stage_B!(m, data1, load; stage=1)

data2 = [Dict{String,Any}() for s=1:ns]
for s=1:ns
    load_s = load*load_scenar[s]
    optim_model_stage_B!(m, data2[s], load_s; stage="2$s")
end

```

Then, there is more work to define the global problem by joining the different stages, either in the objective and/or constraints. Here are examples of both:

- cost is the sum of expected stage costs
- a coupling constraints forbids having 2nd stage generator smaller than at the 1st stage

```julia
cost2 = [data2[s]["cost"] for s=1:ns]
cost2_exp = cost2' * proba_scenar
@objective(m, Min, data1["cost"] + cost2_exp)

@constraint(m, [s=1:ns], data1["Pgen"] <= data2[s]["Pgen"])

```

Optimization results are retrieved by digging into the `Dict` containers of each stage:

```julia
optimize!(m)
Pgen1_val = value(data1["Pgen"])
println("Pgen1=$(Pgen1_val)") # Pgen1=1.5
Pgen2_val = [value(data2[s]["Pgen"]) for s=1:ns]
println("Pgen2=$(Pgen2_val)")
# Pgen2=[1.5, 1.5, 3.0] and not [0.75, 1.5, 3.0]
# due to the stage coupling constraint

```

That’s it, any thoughts welcome!

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [June 26, 2024, 10:09pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/2 "2024-06-26T22:09:41Z")

</div>

Take a read of [Design patterns for larger models · JuMP](https://jump.dev/JuMP.jl/stable/tutorials/getting_started/design_patterns_for_larger_models/)

---

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [June 27, 2024, 8:42am UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/3 "2024-06-27T08:42:20Z")

</div>

Yes, this doc section was part of my “literature review” since you had already mentioned it on Stack Overflow or here, so it’s in the “References” section of my notebook.

Most parts (wrapping the model in a function, having subfunctions to create variables or constraints) do apply (and are applied) to my use case. Only I didn’t need dedicated data structures like [`KnapsackObject`](https://jump.dev/JuMP.jl/stable/tutorials/getting_started/design_patterns_for_larger_models/#Create-better-data-structures).

So the only aspect I didn’t find in that doc (and that was my motivation for making this code design study and posting it here) was the ability to _instantiate several replicas of more or less the same inner problem inside one larger problem_. Thus my interest in [dynamic variable naming](https://discourse.julialang.org/t/using-anonymous-and-or-symbols-to-create-dynamic-varible-names/83095). I the end, I decided to use anonynous variables, identified by their `String` key in a `Dict`. Having one `Dict` per inner problem is like having a dedicated namespace for each inner subproblem, so this solves the duplicate name clashes. Still, as a bonus, I append an optional “stage” suffix to the `base_name` of each variable, so that the String/Latex representation is less ambiguous.

---

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [July 4, 2024, 12:33pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/4 "2024-07-04T12:33:14Z")

</div>

I just came across the 2019 post [Pyomo style "blocks" in JuMP](https://discourse.julialang.org/t/pyomo-style-blocks-in-jump/31849) by @jonmat because the more I dig into it, the more I realize that I’m indeed in search of the **“namespace”** feature mentioned here (and I didn’t know that Pyomo provides this). Indeed, for some other project, I’m may be looking for a Modelica-style object/component oriented structuring of JuMP model.

Now, I understand that in the present state of JuMP it’s indeed possible to implement such things using anonymous variables. However, looking at my “add\_variable” wrapper of the classical `@variable` macro (see below), I feel like I’m just creating a more cumbersome version of `@variable` (see e.g. the presence of `lb` and `ub` keyword args since I cannot use the nice `lb <= x <= ub` syntax without a macro), for the only reason that I want to store the `VariableRef` in my container of choice (see `model_data::Dict{String,Any}`) rather than in the existing global `Model` container of registered variables.

So my question is: what would be the difficulty to bake this “reference storage hijacking” directly in the `@variable` macro? Or should I copy-paste the [`@variable` source code](https://github.com/jump-dev/JuMP.jl/blob/master/src/macros/%40variable.jl) and try to adapt it for my use case?

## my “add\_variable” wrapper

This is the present state of my custom “add\_variable” wrapper around `@variable` (early version discussed in [Using anonymous and or symbols to create dynamic variable names? - #3 by pierre-haessig](https://discourse.julialang.org/t/using-anonymous-and-or-symbols-to-create-dynamic-variable-names/83095/3)). Main features are:

- uses anonymous variables to avoid registration in the global namespace of `Model`
- but instead store the `VariableRef` in a `Dict`
- and then returns the `VariableRef` to make it easy to bind it to a similarly name Julia variable in the caller.

but to get these features, as said above, I have to pay the price of a much worse API compared to the original `@variable`.

```julia-auto
"""
    add_var!(model_data::Dict{String,Any}, name::String;
             lb=nothing, ub=nothing, fix=nothing, K::Int=1, stage="")

add a variable to the JuMP `Model` `model_data["model"]` and store the reference
to that variable in `model_data[name]`.

Returns the variable reference.

`name` is also used as the `base_name` of the JuMP variable (used for printing),
with the optional `stage` argument (`Int`, `String`...) used as a suffix.

# Other optional arguments:
- lower and/or upper bounds can be set with `lb` and `ub`
- equal bounds can be set with `fix`
- the variable can be a Vector of length `K` if `K` >1 (timeseries).
"""
function add_var!(model_data::Dict{String,Any}, name::String;
                  lb=nothing, ub=nothing, fix=nothing, K::Int=1, stage="")
    model = model_data["model"]
    name_suffix = "$name$stage"
    
    if K==1
        v = @variable(model, base_name=name_suffix)
    elseif K>1
        v = @variable(model, [1:K], base_name=name_suffix)
    else
        throw(ArgumentError("`K` array size parameter should be >=1"))
    end
    model_data[name] = v

    # Bounds:
    if (lb!==nothing || ub!==nothing) && fix!==nothing
        throw(ArgumentError("`lb` and `ub` bounds cannot be used in conjunction with `fix`"))
    end
    if fix !==nothing
        fix.(v)
    end
    if lb !==nothing
        set_lower_bound.(v, lb)
    end
    if ub !==nothing
        set_upper_bound.(v, ub)
    end
    
    return v
end

```

which is to be used as

```julia-auto
# One shot model setup
md = Dict{String,Any}()
md["model"] = Model(HiGHS.Optimizer)
# add piles of variables:
x = add_var!(model_data, "x", lb=0.0, ub=power_rated_gen_max)
y = add_var!(model_data, "y", lb=0.0, ub=power_rated_gen_max)

```

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [July 13, 2024, 2:29am UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/5 "2024-07-13T02:29:30Z")

</div>

In almost all cases, writing your own macro is the wrong thing to do.

You could also use the regular `@variable` macro, and then call `unregister` to free up the game again: [JuMP · JuMP](https://jump.dev/JuMP.jl/stable/api/JuMP/#unregister)

---

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [July 15, 2024, 7:49am UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/6 "2024-07-15T07:49:48Z")

</div>

> [@odow](#):
>
> In almost all cases, writing your own macro is the wrong thing to do.

Yes I agree, it would be a lot of work to make a worse/more fragile copy of the original one.

So perhaps a variant of my question is: is it imaginable that a future version of the regular `@variable` macro would accept an optional “model container/data” keyword that would store the `VariableRef` instead of the model itself.

The idea would be to condense the pattern of the last three line into just one line:

```julia
data = Dict{String,Any}()
model = Model(HiGHS.Optimizer)

@variable(model, 0 <= x <= 1)
unregister(mode, :x)
data["x"] = x

```

(and same for constraints)

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [July 15, 2024, 10:58pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/7 "2024-07-15T22:58:40Z")

</div>

Just do:

```Julia
data = Dict{String,Any}()
model = Model(HiGHS.Optimizer)
data["x"] = @variable(model, 0 <= x <= 1)
unregister(model, :x)

```

From a code-readability standard, I think that writing something slightly more verbose that uses the standard JuMP features than trying to writing something custom that saves a line of typing but which is non-standard.

---

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [July 16, 2024, 9:23am UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/8 "2024-07-16T09:23:10Z")

</div>

You’re right, it occurred to me afterwards that the Dict registration line could be merged with the `@variable` declaration line. So the final “code writing cost” is:

- writing the variable name three times, but at least these duplications are contiguous, so little risk of getting out of sync
- one extra line for `unregister`

About you’re final comment, I’m assuming that you meant something like

> “From a code-readability standard, I think that writing something slightly more verbose that uses the standard JuMP features _[is better/safer/simpler to read and maintain…?]_ than trying to writing something custom that saves a line of typing but which is non-standard.”

and I agree.

Thanks for your many feedbacks.

---

<div class="post-metadata">

**Author:** ![ericphanson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ericphanson/32/215186_2.png) [@ericphanson](https://discourse.julialang.org/u/ericphanson)\
**Post date:** [July 16, 2024, 12:44pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/9 "2024-07-16T12:44:39Z")

</div>

The duplication of the variable name and the unregister line could be eliminated with anonymous variables, at the cost of uglier syntax for the upper/lower bounds: [Variables · JuMP](https://jump.dev/JuMP.jl/stable/manual/variables/#anonymous_variables)

---

<div class="post-metadata">

**Author:** ![pierre-haessig](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pierre-haessig/32/217129_2.png) [@pierre-haessig](https://discourse.julialang.org/u/pierre-haessig)\
**Post date:** [July 17, 2024, 2:59pm UTC](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/10 "2024-07-17T14:59:31Z")

</div>

Yes, with anonymous variables come the need to use the `lower_bound`/`upper_bound` keywords. Then this starts to look like my custom `add_var!` wrapper mentioned [above](https://discourse.julialang.org/t/modular-implementation-of-optimization-problems-with-jump/116261/4) that I’d like to avoid.

But the main argument against anonymous variables is the need to add the variable to the current namespace (I say “need” in the sense that it’s handy for writing subsequent model lines). Plus, using anonymous variables needs setting the `base_name` attribute (here “need” means only that is useful when printing expressions or constraints)

In the end, it’s about trade-offs:

```julia
data["x"] = @variable(model, 0 <= x <= 1)
unregister(model, :x)

```

versus

```julia
x = @variable(model, lower_bound=0, upper_bound=1, base_name="x") # base_name is optional, only for prints
data["x"] = x # and this may be compressed into data["x"] = x = @var...

```

Now, thinking a bit more, if I’m sure I want to unregister _all_ registered symbols, it’s possible to do it in one loop at the end using [`object_dictionary(model)`](https://jump.dev/JuMP.jl/stable/api/JuMP/#object_dictionary) mentioned in [How to get all the registered variable names of a `JuMP` model?](https://discourse.julialang.org/t/how-to-get-all-the-registered-variable-names-of-a-jump-model/113837). The model definition would then look like this:

```julia
data["x"] = @variable(model, 0 <= x <= 1)
data["y"] = @variable(model, 0 <= y <= 1)
@constraint(model, x+y <= 1) # here comes the interest of having x and y in the namespace
# ...

for var in keys(object_dictionary(model))
    unregister(model, var)
end

```