Refactor JuMP macros

[odow]

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.

[chriscoey]

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]

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]

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]

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:

• @finite_parameter, @parameter_function: https://github.com/infiniteopt/InfiniteOpt.jl/blob/master/src/macros.jl
• @disjunction: https://github.com/hdavid16/DisjunctiveProgramming.jl/blob/d501c77dbc859c7594600f429a9c6cb068a88c88/src/macros.jl

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]

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]

@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]

Actually. I've simplified to parse_ref_sets:

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

[odow]

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:

JuMP.jl/src/macros/@expression.jl

Lines 66 to 104 in ea44fe4

	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]

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]

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]

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]

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

[odow]

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

[odow]

@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]

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.