Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor JuMP macros #3513

Closed
odow opened this issue Sep 18, 2023 · 15 comments · Fixed by #3629
Closed

Refactor JuMP macros #3513

odow opened this issue Sep 18, 2023 · 15 comments · Fixed by #3629
Labels
Category: Internals Related to JuMP internals

Comments

@odow
Copy link
Member

odow commented Sep 18, 2023

We have a roadmap item for this: https://github.com/jump-dev/JuMP.jl/blob/master/docs/src/developers/roadmap.md, but it'd be useful to have an issue tracking it (x-ref #3500 (comment)).

I don't have any suggestions for what a "good" macros.jl file would look like, but it would be cleaner, with less code being run in the macros, and no user-facing changes.

@odow odow added the Category: Internals Related to JuMP internals label Sep 18, 2023
@chriscoey
Copy link
Contributor

Maybe precluded by the "no user-facing changes", but it would be great if macro changes could make it easier to use JuMP "programmatically". Possibly that could be a separate roadmap item. I don't know what needs to be done really because I'm no expert on macros etc, but I have had to resort to generating low level MOI code to build complex models programmatically.

@odow
Copy link
Member Author

odow commented Sep 18, 2023

I think the non-macro version is just a documentation issue. We already have all of the machinery available in JuMP.

Given:

using JuMP
model = Model()
@variable(model, 0 <= x[1:I, 1:T] <= 100)
@constraint(model, [i in 1:I, t in 2:T], x[i, t] - x[i, t-1] <= 10)
@constraint(model, [i in 1:I, t in 2:T], x[i, t] - x[i, t-1] >= -10)
@objective(model, Min, sum(x))

One version of the non-macro version is

model = Model()
x = [VariableRef(model) for i in 1:I, t in 1:T]
set_lower_bound.(x, 0)
set_upper_bound.(x, 1000)
for i in 1:I, t in 2:T
    add_constraint(model, ScalarConstraint(x[i, t] - x[i, t-1], MOI.LessThan(10.0)))
    add_constraint(model, ScalarConstraint(x[i, t] - x[i, t-1], MOI.GreaterThan(-10.0)))
end
set_objective(model, MOI.MIN_SENSE, sum(x))

It's also possible to minimally use the macros:

model = Model()
x = [@variable(model) for i in 1:I, t in 1:T]
set_lower_bound.(x, 0)
set_upper_bound.(x, 1000)
for i in 1:I, t in 2:T
    f, set = x[i, t] - x[i, t-1], MOI.LessThan(10.0)
    @constraint(model, f in set)
    f, set = x[i, t] - x[i, t-1], MOI.GreaterThan(-10.0)
    @constraint(model, f in set)
end
sense, f = MOI.MIN_SENSE, sum(x)
@objective(model, sense, f)

But these don't give good names, and it's a bit more complicated if you want to construct containers that are not arrays, etc.

There's the much more complicated build_variable add_variable API, but that's really just for the macros. I don't think users gain much by using it.

But I also don't think users gain much by using the non-macro JuMP syntax. If you're doing it programmatically, what's wrong with MOI?

@blegat
Copy link
Member

blegat commented Sep 19, 2023

I think JuMP is still a bit more convenient, even for programatic use, mainly because the model is stored inside a VariableRef or ConstraintRef while a MOI.Index does not store the model.
I agree it's mostly a documentation thing to demystify the JuMP's macro. Maybe the main painpoint when using it it programatically is to build the VariableInfo struct.

@pulsipher
Copy link
Contributor

From the extension side of things, it would be nice to have some of the common macro tooling be accessible. Specifically, I encounter macros of the form @myobject(model, name[inds...], expr_to_parse, extra_args...; base_name = "", container = Automatic, kw_args...) which creates a container or singleton of myobjects which are constructed with appropriate build_myobject and add_myobject functions. With this a lot of the internal methods and workflow used in the JuMP macros have to be duplicated. Examples include:

Currently, JuMP provides JuMP.Containers.build_ref_sets to help with these macros, but that is it. It would be ideal to have a workflow where only the build_myobject and add_myobject functions have to be provided along with code to specify how expr_to_parse should be processed.

@odow
Copy link
Member Author

odow commented Dec 9, 2023

I have been making reasonable progress. The @variable macro is...complicated.

Once I have things under control, I'll think about how we could expose some public API for building related calls.

@odow
Copy link
Member Author

odow commented Dec 11, 2023

@pulsipher interested in your feedback on these changes.

The two most useful new additions are:

args, kwargs = Containers.parse_macro_arguments(error_fn, input_args)

and

Containers.container_name(ref_expr)

I haven't made any larger structural changes, or introduced a "simple" way to write macros. There's still a lot of fluff around error checking the number of arguments, and dealing with optional keyword arguments. But I don't know if there is a good universal design.

The @variable and @constraint macros are pretty complicated, but @objective and @expression are pretty simple.

@odow
Copy link
Member Author

odow commented Dec 11, 2023

Actually. I've simplified to parse_ref_sets:

https://jump.dev/JuMP.jl/previews/PR3620/api/JuMP.Containers/#Containers.container_code

@odow
Copy link
Member Author

odow commented Dec 11, 2023

Perhaps the best way to look at the new code is https://github.com/jump-dev/JuMP.jl/tree/od/container-name/src/macros

Here's the @expression macro, which uses everything without being overly complicated like @constraint or @variable:

macro expression(input_args...)
error_fn(str...) = _macro_error(:expression, input_args, __source__, str...)
args, kwargs = Containers.parse_macro_arguments(error_fn, input_args)
if !(2 <= length(args) <= 3)
error_fn("expected 2 or 3 positional arguments, got $(length(args)).")
elseif Meta.isexpr(args[2], :block)
error_fn("Invalid syntax. Did you mean to use `@expressions`?")
elseif !isempty(kwargs)
for key in keys(kwargs)
if key != :container
error_fn("unsupported keyword argument `$key`.")
end
end
end
name_expr = length(args) == 3 ? args[2] : nothing
name, index_vars, indices = Containers.parse_ref_sets(error_fn, name_expr)
if args[1] in index_vars
error_fn(
"Index $(args[1]) is the same symbol as the model. Use a " *
"different name for the index.",
)
end
model = esc(args[1])
expr, build_code = _rewrite_expression(args[end])
code = quote
$build_code
# Don't leak a `_MA.Zero` if the expression is an empty summation, or
# other structure that returns `_MA.Zero()`.
_replace_zero($model, $expr)
end
container = get(kwargs, :container, :Auto)
return _finalize_macro(
model,
Containers.container_code(index_vars, indices, code, container),
__source__;
register_name = name,
wrap_let = true,
)
end

Here's what it used to look like

JuMP.jl/src/macros.jl

Lines 1923 to 1975 in 419c334

macro expression(args...)
_error(str...) = _macro_error(:expression, args, __source__, str...)
args, kw_args, requested_container = Containers._extract_kw_args(args)
if length(args) == 3
m = esc(args[1])
c = args[2]
x = args[3]
elseif length(args) == 2
m = esc(args[1])
c = gensym()
x = args[2]
else
_error("needs at least two arguments.")
end
length(kw_args) == 0 || _error("unrecognized keyword argument")
if isexpr(args[2], :block)
_error("Invalid syntax. Did you mean to use `@expressions`?")
end
anonvar = isexpr(c, :vect) || isexpr(c, :vcat) || length(args) == 2
variable = gensym()
idxvars, indices = Containers.build_ref_sets(_error, c)
if args[1] in idxvars
_error(
"Index $(args[1]) is the same symbol as the model. Use a " *
"different name for the index.",
)
end
expr_var, build_code = _rewrite_expression(x)
code = quote
$build_code
# Don't leak a `_MA.Zero` if the expression is an empty summation, or
# other structure that returns `_MA.Zero()`.
_replace_zero($m, $expr_var)
end
code =
Containers.container_code(idxvars, indices, code, requested_container)
# Wrap the entire code block in a let statement to make the model act as
# a type stable local variable.
code = _wrap_let(m, code)
# don't do anything with the model, but check that it's valid anyway
if anonvar
macro_code = code
else
macro_code = _macro_assign_and_return(
code,
variable,
Containers._get_name(c);
model_for_registering = m,
)
end
return _finalize_macro(m, macro_code, __source__)
end

@pulsipher
Copy link
Contributor

Hi @odow, you've been busy! I definitely like the direction these are going in. Some suggestions to improve things further:

  • Have parse_macro_arguments accept an valid_kwargs argument that defaults to something like (:container,) such that it errors if some unexpected keyword argument is given.
  • Have parse_ref_sets do the check whether index is the same symbol as the model, this is used by every macro.
  • Perhaps have container_code accept the kwargs output from parse_macro_arguments extract the requested container since every macro will follow the same workflow
  • It would be really really nice to make _finalize_macro public since this brings everything together and its current implementation is already general enough to handle extensions.
  • Perhaps make _macro_error public, it is a small function, but every extension macro I have made ends up repeating it.

To guide refining the API, perhaps you could make a PR in JuMP that I could use to make an InfiniteOpt PR that tests out the new API.

@odow
Copy link
Member Author

odow commented Dec 11, 2023

  1. Good idea
  2. Something like forbidden_indices::Vector{Symbol}?
  3. I don't really know how this would work. Do you have an example of the code you'd like the macro to generate?
  4. Yeah I was thinking like extension_finalize_macro_helper or something
  5. Good idea

@pulsipher
Copy link
Contributor

Something like forbidden_indices::Vector{Symbol}?

Yes

I don't really know how this would work. Do you have an example of the code you'd like the macro to generate?

For instance, right now in @expression you have:

args, kwargs = Containers.parse_macro_arguments(error_fn, input_args) 
# PUT CODE HERE to derive the code to define `index_vars`, `indices`, `code`
container = get(kwargs, :container, :Auto)
container_code = Containers.container_code(index_vars, indices, code, container)

By allowing the 4th argument of container_code to additionally accept the type of kwargs (which I believe is a Dict) then we don't have to do any processing of containers outside the API. Then we could have:

args, kwargs = Containers.parse_macro_arguments(error_fn, input_args) 
# PUT CODE HERE to derive the code to define `index_vars`, `indices`, `code`
container_code = Containers.container_code(index_vars, indices, code, kwargs)

Yeah I was thinking like extension_finalize_macro_helper or something

Awesome!

@pulsipher
Copy link
Contributor

One last suggestion that I forgot would be to make _add_additional_args public as well.

@odow
Copy link
Member Author

odow commented Dec 11, 2023

Try #3620. I haven't made _finalize_macro public yet.

@odow
Copy link
Member Author

odow commented Dec 12, 2023

@pulsipher I've been thinking about _finalize_macro.

It seems like a a weird function to export from JuMP. But it also can't go in Containers like the rest of the macro stuff. The code is also pretty trivial, so perhaps it's not the end of the world if you have to copy-paste just this one part.

@pulsipher
Copy link
Contributor

While I would prefer for _finalize_macro to be public, it is not too much trouble to copy it.

I tried out #3620 in InfiniteOpt.jl, see #3620 for details.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Category: Internals Related to JuMP internals
4 participants