diff --git a/examples/hovercraft_example.jl b/examples/hovercraft_example.jl index f0db39b..07d10f3 100644 --- a/examples/hovercraft_example.jl +++ b/examples/hovercraft_example.jl @@ -28,4 +28,3 @@ constant_over_collocation.(u, t) # needed for collocation # Solve optimize!(im) - diff --git a/src/InfiniteExaModels.jl b/src/InfiniteExaModels.jl index a423aba..5af3463 100644 --- a/src/InfiniteExaModels.jl +++ b/src/InfiniteExaModels.jl @@ -6,6 +6,7 @@ import InfiniteOpt.TranscriptionOpt as _TO include("infiniteopt_backend.jl") include("operators.jl") +include("grouped_patterns.jl") include("transform.jl") export ExaMappingData, ExaTranscriptionBackend diff --git a/src/grouped_patterns.jl b/src/grouped_patterns.jl new file mode 100644 index 0000000..45e4395 --- /dev/null +++ b/src/grouped_patterns.jl @@ -0,0 +1,287 @@ +# Integer alias for types of InfiniteOpt modelling objects to use in hashing expressions +const _VariableTypeHashingInt = Dict( + InfiniteOpt.FiniteParameterIndex => -2, + InfiniteOpt.ParameterFunctionIndex => -3, + InfiniteOpt.InfiniteVariableIndex => -4, + InfiniteOpt.DerivativeIndex => -4, + InfiniteOpt.SemiInfiniteVariableIndex => -5, + InfiniteOpt.PointVariableIndex => -6, + InfiniteOpt.FiniteVariableIndex => -7, + InfiniteOpt.IndependentParameterIndex => -8, + InfiniteOpt.DependentParameterIndex => -9, + InfiniteOpt.MeasureIndex => -10, +) + +# Appropriately encode an InfiniteOpt variable such that grouping variables can be appropriately assigned +function _encode_variable(v::InfiniteOpt.GeneralVariableRef) + if v.index_type == InfiniteOpt.PointVariableIndex + group_idxs = InfiniteOpt.parameter_group_int_indices(InfiniteOpt.infinite_variable_ref(v)) + elseif v.index_type == InfiniteOpt.SemiInfiniteVariableIndex + group_idxs = InfiniteOpt.parameter_group_int_indices(InfiniteOpt.infinite_variable_ref(v)) + group_idxs = vcat(group_idxs, InfiniteOpt.parameter_group_int_indices(v)) # append the semi-infinite index group indices distinguish y(0, x) from y(t, -1) for example + else + group_idxs = InfiniteOpt.parameter_group_int_indices(v) + end + return _VariableTypeHashingInt[v.index_type], group_idxs +end + +## Extract the following from an expression: +# 1. A hash of the symbolic expression structure +# 2. A list of all variable references in the expression in the order they appear +# 3. A list of all constant values in the expression in the order they appear +function _encode_expr(expr::JuMP.AbstractJuMPScalar) + return _encode_expr(expr, hash(:+), InfiniteOpt.GeneralVariableRef[], Float64[]) +end +function _encode_expr(c::Real, h::UInt, refs, consts) + return hash(-1, h), refs, push!(consts, c) # -1 indicates a symbolic constant +end +function _encode_expr(v::InfiniteOpt.GeneralVariableRef, h::UInt, refs, consts) + return hash(_encode_variable(v), h), push!(refs, v), consts +end +function _encode_expr( + expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, + h::UInt, + refs, + consts + ) where {C, V} + return _encode_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), h, refs, consts) +end +function _encode_expr(expr::JuMP.GenericNonlinearExpr, h::UInt, refs, consts) # TODO remove recursion + h = hash((expr.head, length(expr.args)), h) + for arg in expr.args + h, _, _ = _encode_expr(arg, h, refs, consts) + end + return h, refs, consts +end + +# Traverse expression in same order as _encode_expr and exafy it +function _exafy_grouped_expr( + ::Real, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return popfirst!(consts) +end +function _exafy_grouped_expr( + ::InfiniteOpt.GeneralVariableRef, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return popfirst!(vrefs) +end +function _exafy_grouped_expr( + expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, + vrefs::Vector{Any}, + consts::Vector{Any} + ) where {C, V} + return _exafy_grouped_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), vrefs, consts) +end +function _exafy_grouped_expr( + expr::JuMP.GenericNonlinearExpr, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return _nl_op(expr.head)((_exafy_grouped_expr(a, vrefs, consts) for a in expr.args)...) +end + +# Print a message about a group +function _group_info_msg(group, msg) + idxs = [JuMP.index(cref).value for cref in group] + @info "$msg constraint group with indices: $(idxs)" + return +end + +# Get the grouped index of a variable based on its direct exaified variable reference +function _get_grouped_idx(em_var::ExaModels.Var, grouped_var::ExaModels.Variable) + idx = em_var.i - grouped_var.offset + @assert 1 <= idx <= grouped_var.size[end] && length(grouped_var.size) == 1 + return idx +end +function _get_grouped_idx( + em_var::Union{ExaModels.Variable, ExaModels.Parameter}, + grouped_var::Union{ExaModels.Variable, ExaModels.Parameter} + ) + idx = (em_var.offset - grouped_var.offset) ÷ em_var.length + 1 + @assert 1 <= idx <= grouped_var.size[end] + return idx +end +function _get_grouped_idx(vref::InfiniteOpt.GeneralVariableRef, data::ExaMappingData) + if vref.index_type in (InfiniteOpt.SemiInfiniteVariableIndex, InfiniteOpt.PointVariableIndex) + vref = InfiniteOpt.infinite_variable_ref(vref) + end + em_var = data[vref] + grouped_var = data.var_to_grouped_var[vref] + return _get_grouped_idx(em_var, grouped_var) +end + +# Given the lists from _encode_expr, create the exafied expression and the finite iterator for the grouped algebraic pattern +# TODO: possible take in idx counters as input to avoid clashing (for sums) +function _process_grouped_expression( + expr::JuMP.AbstractJuMPScalar, + vref_lists,#::Vector{Vector{InfiniteOpt.GeneralVariableRef}}, + const_lists,#::Vector{Vector{Float64}}, + data::ExaMappingData + ) + # determine which vrefs and consts change across the array + vrefs1 = vref_lists[1] + is_grouped_var = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] + consts1 = const_lists[1] + is_grouped_data = [any(l -> l[i] != consts1[i], const_lists) for i in eachindex(consts1)] + # exafy the vrefs + exafied_vrefs = Vector{Any}(undef, length(vrefs1)) + var_itr = Any[(;) for _ in 1:length(vref_lists)] + group_var_idx = 1 + restricted_idx = 1 + for (i, vref) in enumerate(vrefs1) + if is_grouped_var[i] + @assert haskey(data.var_to_grouped_var, vref) + base_idxs = Tuple(_index_params(vref, data)) + itr_alias = Symbol("grouped_vidx$group_var_idx") + data_src = ExaModels.DataSource() + alias_map = Dict{Int, Symbol}() + var_idxs = (begin + if k > length(base_idxs) + data_src[itr_alias] + elseif base_idxs[k] isa Int # for restricted variables + alias_map[k] = Symbol("restricted_idx$restricted_idx") + restricted_idx += 1 + data_src[alias_map[k]] + else + base_idxs[k] + end + end for k in 1:length(base_idxs)+1) + src_var = data.var_to_grouped_var[vref] + exafied_vrefs[i] = src_var[var_idxs...] + for j in 1:length(vref_lists) + infvar = vref_lists[j][i] + @assert data.var_to_grouped_var[infvar] == src_var + var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(infvar, data)) + if !isempty(alias_map) # add in restricted variables indices if they exist + ridxs = Tuple(_index_params(infvar, data)) + var_itr[j] = merge(var_itr[j], NamedTuple(alias => ridxs[k] for (k, alias) in alias_map)) + end + end + group_var_idx += 1 + else + exafied_vrefs[i] = _exafy(vref, data) + end + end + # exafy the consts + exafied_consts = Vector{Any}(undef, length(consts1)) + const_itr = Any[(;) for _ in 1:length(const_lists)] + grouped_const_idx = 1 + for (i, c) in enumerate(consts1) + if is_grouped_data[i] + itr_alias = Symbol("grouped_const$grouped_const_idx") + exafied_consts[i] = ExaModels.DataSource()[itr_alias] + for j in 1:length(const_lists) + const_itr[j] = (; const_itr[j]..., itr_alias => const_lists[j][i]) + end + grouped_const_idx += 1 + else + exafied_consts[i] = c + end + end + # build the ExaModels graph and the finite iterator for the algebraic pattern + em_expr = _finalize_expr(_exafy_grouped_expr(expr, exafied_vrefs, exafied_consts)) + finite_itr = [merge(var_itr[i], const_itr[i]) for i in 1:length(vref_lists)] + return em_expr, finite_itr +end + +# Given a candidate group of constraint, seek to merge together and add as a single constraint pattern to `core` +function _process_candidate_constraint_group( + core::ExaModels.ExaCore, + data::ExaMappingData, + crefs::Vector{InfiniteOpt.InfOptConstraintRef}, + vref_lists::Vector{Vector{InfiniteOpt.GeneralVariableRef}}, + const_lists::Vector{Vector{Float64}}, + sets::Vector{_MOI.AbstractSet} + ) + # build the expression graph and finite iterator for the algebraic pattern + raw_expr = JuMP.jump_function(JuMP.constraint_object(first(crefs))) + em_expr, finite_itr = _process_grouped_expression(raw_expr, vref_lists, const_lists, data) + # process the iterator + infinite_itr = _get_constraint_iterator(first(crefs), data) + itr = vec([merge(i...) for i in Iterators.product(infinite_itr, finite_itr)]) + # add the constraints to the core + lbs = Vector{Float64}(undef, length(crefs)) + ubs = Vector{Float64}(undef, length(crefs)) + for (i, s) in enumerate(sets) + lbs[i], ubs[i] = _get_constr_bounds(s) + end + full_lbs = repeat(lbs, inner = length(infinite_itr)) + full_ubs = repeat(ubs, inner = length(infinite_itr)) + core, con = ExaModels.add_con(core, em_expr, itr, lcon = full_lbs, ucon = full_ubs) + # save the constraint mappings + inf_len = length(infinite_itr) + for (i, cref) in enumerate(crefs) + base_idx = (i - 1) * inf_len + 1 + sliced_itr = itr[base_idx:base_idx + inf_len - 1] + offset = con.offset + base_idx - 1 + data.constraint_mappings[cref] = ExaModels.Constraint(con.f, sliced_itr, offset, (inf_len,), nothing) + end + return core +end + +# Iterate over constraints in the InfiniteOpt model, group by algebraic pattern, and add to the ExaModels core +function _group_and_add_constraints( + core::ExaModels.ExaCore, + data::ExaMappingData, + inf_model::InfiniteOpt.InfiniteModel + ) + # set up dictionaries for tracking patterns + hash_to_patterns = Dict{UInt, Tuple{Vector{Vector{InfiniteOpt.GeneralVariableRef}}, Vector{Vector{Float64}}, Vector{_MOI.AbstractSet}}}() + hash_to_constrs = Dict{UInt, Vector{InfiniteOpt.InfOptConstraintRef}}() + # iterate over constraints and group by hashed algebraic pattern + for cref in JuMP.all_constraints(inf_model) + InfiniteOpt.is_variable_domain_constraint(cref) && continue + isempty(JuMP.owner_model(cref).constraints[JuMP.index(cref)].measure_indices) || continue # TODO: temporary restriction + expr = JuMP.jump_function(JuMP.constraint_object(cref)) + expr isa JuMP.AbstractJuMPScalar || continue + h, vrefs, consts = _encode_expr(expr) + if haskey(hash_to_patterns, h) + push!(hash_to_patterns[h][1], vrefs) + push!(hash_to_patterns[h][2], consts) + push!(hash_to_patterns[h][3], JuMP.moi_set(JuMP.constraint_object(cref))) + push!(hash_to_constrs[h], cref) + else + hash_to_patterns[h] = ([vrefs], [consts], [JuMP.moi_set(JuMP.constraint_object(cref))]) + hash_to_constrs[h] = [cref] + end + end + # process each grouped pattern (requiring at least 2 constraints to be grouped) + for (h, crefs) in hash_to_constrs + if length(crefs) < 2 + continue + end + core = _process_candidate_constraint_group(core, data, crefs, hash_to_patterns[h]...) + _group_info_msg(crefs, "Successfully added") + end + return core +end + +## Given an objective expression, see if it can be expressed as a finite sum of grouped terms +# NonlinearExpr +function _process_candidate_sum_group( + expr::JuMP.GenericNonlinearExpr, + data::ExaMappingData + ) + expr.head == :+ || return _exafy(expr, data), [(;)] # TODO: check for other sum-like operations + length(expr.args) == 1 && _process_candidate_sum_group(expr.args[1], data) + flat_expr = JuMP.flatten!(JuMP.GenericNonlinearExpr(expr.head, copy(expr.args))) + vref_lists = Vector{Vector{InfiniteOpt.GeneralVariableRef}}(undef, length(flat_expr.args)) + const_lists = Vector{Vector{Float64}}(undef, length(flat_expr.args)) + hs = Vector{UInt}(undef, length(flat_expr.args)) + for (i, arg) in enumerate(flat_expr.args) + hs[i], vref_lists[i], const_lists[i] = _encode_expr(arg) + end + all(hs[1] == h for h in hs) || return _exafy(expr, data), [(;)] # TODO: perhaps we can break this up + return _process_grouped_expression(flat_expr.args[1], vref_lists, const_lists, data) +end +# Fallback for other expressions +function _process_candidate_sum_group( + expr::JuMP.AbstractJuMPScalar, + data::ExaMappingData + ) + return _exafy(expr, data), [(;)] +end diff --git a/src/infiniteopt_backend.jl b/src/infiniteopt_backend.jl index 7f761c0..dd0f289 100644 --- a/src/infiniteopt_backend.jl +++ b/src/infiniteopt_backend.jl @@ -31,6 +31,10 @@ struct ExaMappingData Vector{Any} } } + # Point variable info + point_indicies::Dict{InfiniteOpt.GeneralVariableRef, Tuple} + # Finite template constraint metadata + var_to_grouped_var::Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}} # Default constructor function ExaMappingData() @@ -52,6 +56,8 @@ struct ExaMappingData Vector{Any} } }(), + Dict{InfiniteOpt.GeneralVariableRef, Tuple}(), + Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}}(), ) end end @@ -109,7 +115,7 @@ mutable struct ExaTranscriptionBackend{B} <: InfiniteOpt.AbstractTransformationB end # Constructors -function ExaTranscriptionBackend(; backend = nothing) +function ExaTranscriptionBackend(; backend = nothing, ) return ExaTranscriptionBackend( nothing, nothing, diff --git a/src/transform.jl b/src/transform.jl index f1ab2fe..60425d4 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -51,8 +51,7 @@ _process_value(pf::InfiniteOpt.ParameterFunction, supp) = pf(supp) ## Determine the bounds of an InfiniteOpt variable # Real bound and start value function _get_variable_bounds_and_start( - info::JuMP.VariableInfo{<:Real, <:Real, <:Real, <:Real}, - itrs = nothing + info::JuMP.VariableInfo{<:Real, <:Real, <:Real, <:Real} ) lb = -Inf ub = Inf @@ -76,25 +75,26 @@ end function _get_variable_bounds_and_start(info::JuMP.VariableInfo, itrs) # set up the collection arrays dims = Tuple(length(itr) for itr in itrs) - lb = fill(-Inf, dims) - ub = fill(Inf, dims) - start = fill(0.0, dims) + lin_idxs = LinearIndices(dims) + lb = fill(-Inf, length(lin_idxs)) + ub = fill(Inf, length(lin_idxs)) + start = fill(0.0, length(lin_idxs)) # iterate over all support combinations and fill in the arrays for i in Iterators.product(itrs...) supp = [s for nt in i for s in Iterators.drop(values(nt), 1)] if info.has_fix val = _process_value(info.fixed_value, supp) - lb[first.(i)...] = val - ub[first.(i)...] = val + lb[lin_idxs[first.(i)...]] = val + ub[lin_idxs[first.(i)...]] = val end if info.has_lb - lb[first.(i)...] = _process_value(info.lower_bound, supp) + lb[lin_idxs[first.(i)...]] = _process_value(info.lower_bound, supp) end if info.has_ub - ub[first.(i)...] = _process_value(info.upper_bound, supp) + ub[lin_idxs[first.(i)...]] = _process_value(info.upper_bound, supp) end if info.has_start - start[first.(i)...] = _process_value(info.start, supp) + start[lin_idxs[first.(i)...]] = _process_value(info.start, supp) end end return lb, ub, start @@ -104,7 +104,7 @@ end function _get_name(vref::InfiniteOpt.GeneralVariableRef, default_name = "var") raw_name = JuMP.name(vref) sym_name = isempty(raw_name) ? Symbol("$(default_name)$(vref.raw_index)") : Symbol(raw_name) - return Val(sym_name) + return sym_name end # Add all the finite variables from an InfiniteModel to a ExaCore @@ -113,13 +113,18 @@ function _add_finite_variables( data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel ) - for vref in JuMP.all_variables(inf_model, InfiniteOpt.FiniteVariable) + vrefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteVariable) + core, ex_vars = ExaModels.add_var(core, length(vrefs), name = Val(:finvar)) + for (i, vref) in enumerate(vrefs) info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo _ensure_continuous(info) lb, ub, start = _get_variable_bounds_and_start(info) - vname = _get_name(vref, "finvar") - core, new_var = ExaModels.add_var(core, 1, start = start, lvar = lb, uvar = ub, name = vname) - data.finvar_mappings[vref] = new_var[1] + ex_var = ex_vars[i] + data.finvar_mappings[vref] = ex_var + core.lvar[ex_var.i] = lb + core.uvar[ex_var.i] = ub + core.x0[ex_var.i] = start + data.var_to_grouped_var[vref] = ex_vars end return core end @@ -130,10 +135,13 @@ function _add_finite_parameters( data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel ) - for pref in JuMP.all_variables(inf_model, InfiniteOpt.FiniteParameter) - param_val = InfiniteOpt.parameter_value(pref) - core, new_par = ExaModels.add_par(core, [param_val]) - data.param_mappings[pref] = new_par + prefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteParameter) + core, ex_pars = ExaModels.add_par(core, length(prefs)) + offset = ex_pars.offset + for (i, pref) in enumerate(prefs) + data.param_mappings[pref] = ExaModels.Parameter((1,), 1, offset + i - 1, nothing) + core.θ[offset + i] = InfiniteOpt.parameter_value(pref) + data.var_to_grouped_var[pref] = ex_pars end return core end @@ -144,24 +152,44 @@ function _add_infinite_variables( data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel ) - # Get the raw variables + # get the raw variables ivrefs = JuMP.all_variables(inf_model, InfiniteOpt.InfiniteVariable) InfiniteOpt.reformulate_high_order_derivatives!(inf_model) drefs = InfiniteOpt.all_derivatives(inf_model) - # now process and add each infinite variable - for vref in append!(ivrefs, drefs) - # retrieve basic information - info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo - _ensure_continuous(info) - group_idxs = InfiniteOpt.parameter_group_int_indices(vref) - # prepare the bounds and start values + vrefs = append!(ivrefs, drefs) + # sort the variables by parameter groups + if length(data.base_itrs) > 1 + group_to_vrefs = Dict{Vector{Int}, Vector{InfiniteOpt.GeneralVariableRef}}() + for vref in vrefs + group_idxs = InfiniteOpt.parameter_group_int_indices(vref) + if !haskey(group_to_vrefs, group_idxs) + group_to_vrefs[group_idxs] = [vref] + else + push!(group_to_vrefs[group_idxs], vref) + end + end + else + group_idxs = InfiniteOpt.parameter_group_int_indices(first(vrefs)) + group_to_vrefs = Dict(group_idxs => vrefs) + end + # add each group of variables in group_to_vrefs to the ExaCore + for (group_idxs, vrefs) in group_to_vrefs itrs = map(i -> data.base_itrs[i], group_idxs) - lb, ub, start = _get_variable_bounds_and_start(info, itrs) - # create the ExaModels variable dims = Tuple(length(itr) for itr in itrs) - vname = _get_name(vref, vref in drefs ? "deriv" : "infvar") - core, new_var = ExaModels.add_var(core, dims...; start = start, lvar = lb, uvar = ub, name = vname) - data.infvar_mappings[vref] = new_var + core, ex_vars = ExaModels.add_var(core, dims..., length(vrefs), name = Val(:infvar)) + offset = ex_vars.offset + for vref in vrefs + info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo + _ensure_continuous(info) + lb, ub, start = _get_variable_bounds_and_start(info, itrs) + vname = _get_name(vref, vref in drefs ? "deriv" : "infvar") + data.infvar_mappings[vref] = ExaModels.Variable(dims, length(lb), offset, vname, nothing) + copyto!(@view(core.lvar[offset+1:offset+length(lb)]), lb) + copyto!(@view(core.uvar[offset+1:offset+length(ub)]), ub) + copyto!(@view(core.x0[offset+1:offset+length(start)]), start) + offset += length(lb) + data.var_to_grouped_var[vref] = ex_vars + end end return core end @@ -172,21 +200,42 @@ function _add_parameter_functions( data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel ) - for pfref in InfiniteOpt.all_parameter_functions(inf_model) - # gather the basic information - group_idxs = InfiniteOpt.parameter_group_int_indices(pfref) - pfunc = InfiniteOpt.core_object(pfref) - # compute the value for each support combination and store + pfrefs = InfiniteOpt.all_parameter_functions(inf_model) + iszero(length(pfrefs)) && return core + # sort the parameter functions by parameter groups + if length(data.base_itrs) > 1 + group_to_pfrefs = Dict{Vector{Int}, Vector{InfiniteOpt.GeneralVariableRef}}() + for pfref in pfrefs + group_idxs = InfiniteOpt.parameter_group_int_indices(pfref) + if !haskey(group_to_pfrefs, group_idxs) + group_to_pfrefs[group_idxs] = [pfref] + else + push!(group_to_pfrefs[group_idxs], pfref) + end + end + else + group_idxs = InfiniteOpt.parameter_group_int_indices(first(pfrefs)) + group_to_pfrefs = Dict(group_idxs => pfrefs) + end + # add each group of parameter functions to the ExaCore + for (group_idxs, group_pfrefs) in group_to_pfrefs itrs = map(i -> data.base_itrs[i], group_idxs) dims = Tuple(length(itr) for itr in itrs) - vals = Array{Float64}(undef, dims...) - for i in Iterators.product(itrs...) - supp = [s for nt in i for s in Iterators.drop(values(nt), 1)] - vals[first.(i)...] = pfunc(supp) + core, ex_pars = ExaModels.add_par(core, dims..., length(group_pfrefs)) + offset = ex_pars.offset + for pfref in group_pfrefs + pfunc = InfiniteOpt.core_object(pfref) + lin_idxs = LinearIndices(dims) + vals = Vector{Float64}(undef, length(lin_idxs)) + for i in Iterators.product(itrs...) + supp = [s for nt in i for s in Iterators.drop(values(nt), 1)] + vals[lin_idxs[first.(i)...]] = pfunc(supp) + end + copyto!(@view(core.θ[offset+1:offset+length(vals)]), vals) + data.param_mappings[pfref] = ExaModels.Parameter(dims, length(vals), offset, nothing) + offset += length(vals) + data.var_to_grouped_var[pfref] = ex_pars end - # Register the parameter function values in the ExaCore & mapping data - core, new_par = ExaModels.add_par(core, vals) - data.param_mappings[pfref] = new_par end return core end @@ -218,6 +267,9 @@ function _process_semi_infinite_var(vref, data) else mapped_var = data.infvar_mappings[ivref] end + if haskey(data.var_to_grouped_var, ivref) + data.var_to_grouped_var[vref] = data.var_to_grouped_var[ivref] + end return data.semivar_info[vref] = (mapped_var, indexing) end @@ -275,7 +327,12 @@ function _process_point_var(vref, data) end group_idxs = InfiniteOpt.parameter_group_int_indices(ivref) idxs = Tuple(data.support_to_index[i, s] for (i, s) in zip(group_idxs, supp)) - return data.infvar_mappings[ivref][idxs...] + pt = data.infvar_mappings[ivref][idxs...] + if haskey(data.var_to_grouped_var, ivref) + data.var_to_grouped_var[vref] = data.var_to_grouped_var[ivref] + end + data.point_indicies[vref] = idxs + return data.finvar_mappings[vref] = pt end # Add all the point variables from an InfiniteModel to a ExaCore @@ -287,7 +344,6 @@ function _add_point_variables( for vref in JuMP.all_variables(inf_model, InfiniteOpt.PointVariable) # store the index mapping for the point variable pt = _process_point_var(vref, data) - data.finvar_mappings[vref] = pt # update the bounds and start value if needed info = InfiniteOpt.core_object(vref).info # InfiniteOpt.RestrictedDomainInfo _update_bounds_and_start(core, info, pt) @@ -295,6 +351,52 @@ function _add_point_variables( return end +# Get the index parameters for a variable reference (used by `_map_variable`) +function _index_params( + vref::InfiniteOpt.GeneralVariableRef, + data::ExaMappingData + ) + _index_params(vref, vref.index_type, data) +end +function _index_params( + vref::InfiniteOpt.GeneralVariableRef, + ::Type{V}, + data::ExaMappingData + ) where V <: Union{InfiniteOpt.InfiniteVariableIndex, InfiniteOpt.DerivativeIndex, InfiniteOpt.ParameterFunctionIndex} + group_idxs = InfiniteOpt.parameter_group_int_indices(vref) + data_src = ExaModels.DataSource() + return (data_src[data.group_alias[i]] for i in group_idxs) +end +function _index_params( + vref::InfiniteOpt.GeneralVariableRef, + ::Type{InfiniteOpt.SemiInfiniteVariableIndex}, + data::ExaMappingData + ) + if !haskey(data.semivar_info, vref) + _process_semi_infinite_var(vref, data) + end + _, inds = data.semivar_info[vref] + data_src = ExaModels.DataSource() + return (i isa Int ? i : data_src[i] for i in inds) +end +function _index_params( + vref::InfiniteOpt.GeneralVariableRef, + ::Type{InfiniteOpt.PointVariableIndex}, + data::ExaMappingData + ) + if !haskey(data.finvar_mappings, vref) + _process_point_var(vref, data) + end + return data.point_indicies[vref] +end +function _index_params( + vref::InfiniteOpt.GeneralVariableRef, + type, + data::ExaMappingData + ) + return () +end + # Add user-defined operators to ExaModels function _add_user_operators(inf_model::InfiniteOpt.InfiniteModel) for op in InfiniteOpt.added_nonlinear_operators(inf_model) @@ -316,68 +418,59 @@ function _add_user_operators(inf_model::InfiniteOpt.InfiniteModel) end # Map variable references based on their underlying type (used by `_exafy`) -function _map_variable(vref, ::Type{InfiniteOpt.FiniteVariableIndex}, data_src, data) +function _map_variable(vref, ::Type{InfiniteOpt.FiniteVariableIndex}, data) return data.finvar_mappings[vref] end -function _map_variable(vref, ::Type{InfiniteOpt.PointVariableIndex}, data_src, data) +function _map_variable(vref, ::Type{InfiniteOpt.PointVariableIndex}, data) if haskey(data.finvar_mappings, vref) return data.finvar_mappings[vref] else - var = _process_point_var(vref, data) - data.finvar_mappings[vref] = var - return var + return _process_point_var(vref, data) end end function _map_variable( vref, ::Type{V}, - data_src, data ) where V <: Union{InfiniteOpt.InfiniteVariableIndex, InfiniteOpt.DerivativeIndex} - group_idxs = InfiniteOpt.parameter_group_int_indices(vref) - idx_pars = (data_src[data.group_alias[i]] for i in group_idxs) + idx_pars = _index_params(vref, V, data) return data.infvar_mappings[vref][idx_pars...] end -function _map_variable(vref, ::Type{InfiniteOpt.SemiInfiniteVariableIndex}, data_src, data) - if !haskey(data.semivar_info, vref) - _process_semi_infinite_var(vref, data) - end - ivar, inds = data.semivar_info[vref] - idx_pars = (i isa Int ? i : data_src[i] for i in inds) +function _map_variable(vref, ::Type{InfiniteOpt.SemiInfiniteVariableIndex}, data) + idx_pars = _index_params(vref, InfiniteOpt.SemiInfiniteVariableIndex, data) + ivar, _ = data.semivar_info[vref] return ivar[idx_pars...] end -function _map_variable(vref, ::Type{<:InfiniteOpt.InfiniteParameterIndex}, data_src, data) - return data_src[data.param_alias[vref]] +function _map_variable(vref, ::Type{<:InfiniteOpt.InfiniteParameterIndex}, data) + return ExaModels.DataSource()[data.param_alias[vref]] end -function _map_variable(vref, ::Type{InfiniteOpt.FiniteParameterIndex}, data_src, data) +function _map_variable(vref, ::Type{InfiniteOpt.FiniteParameterIndex}, data) return data.param_mappings[vref][1] end -function _map_variable(vref, ::Type{InfiniteOpt.ParameterFunctionIndex}, data_src, data) - group_idxs = InfiniteOpt.parameter_group_int_indices(vref) - idx_pars = (data_src[data.group_alias[i]] for i in group_idxs) +function _map_variable(vref, ::Type{InfiniteOpt.ParameterFunctionIndex}, data) + idx_pars = _index_params(vref, InfiniteOpt.ParameterFunctionIndex, data) return data.param_mappings[vref][idx_pars...] end -function _map_variable(vref, IdxType, data_src, data) +function _map_variable(vref, IdxType, data) error("Unable to add `$vref` to an ExaModel, it's index type `$IdxType`" * " is not yet supported by InfiniteExaModels.") end -# Convert as InfiniteOpt expression into a ExaModel expression using the DataIndexed `data_src` -function _exafy(vref::InfiniteOpt.GeneralVariableRef, data_src, data) - return _map_variable(vref, vref.index_type, data_src, data) +# Convert as InfiniteOpt expression into a ExaModel expression +function _exafy(vref::InfiniteOpt.GeneralVariableRef, data) + return _map_variable(vref, vref.index_type, data) end -function _exafy(c::Number, data_src, data) +function _exafy(c::Number, data) return c end function _exafy( - aff::JuMP.GenericAffExpr{C, InfiniteOpt.GeneralVariableRef}, - data_src, + aff::JuMP.GenericAffExpr{C, InfiniteOpt.GeneralVariableRef}, data ) where {C} c = JuMP.constant(aff) if !isempty(aff.terms) ex = sum(begin - v_ex = _exafy(v, data_src, data) + v_ex = _exafy(v, data) isone(c) ? v_ex : c * v_ex end for (c, v) in JuMP.linear_terms(aff) ) @@ -387,19 +480,18 @@ function _exafy( end end function _exafy( - quad::JuMP.GenericQuadExpr{C, InfiniteOpt.GeneralVariableRef}, - data_src, + quad::JuMP.GenericQuadExpr{C, InfiniteOpt.GeneralVariableRef}, data ) where {C} - aff = _exafy(quad.aff, data_src, data) + aff = _exafy(quad.aff, data) if !isempty(quad.terms) ex = sum(begin if v1 == v2 - v_ex = _exafy(v1, data_src, data) + v_ex = _exafy(v1, data) isone(c) ? abs2(v_ex) : c * abs2(v_ex) else - v1_ex = _exafy(v1, data_src, data) - v2_ex = _exafy(v2, data_src, data) + v1_ex = _exafy(v1, data) + v2_ex = _exafy(v2, data) isone(c) ? v1_ex * v2_ex : c * v1_ex * v2_ex end end for (c, v1, v2) in JuMP.quad_terms(quad) @@ -411,10 +503,35 @@ function _exafy( end function _exafy( nl::JuMP.GenericNonlinearExpr{InfiniteOpt.GeneralVariableRef}, - data_src, data ) - return _nl_op(nl.head)((_exafy(a, data_src, data) for a in nl.args)...) + return _nl_op(nl.head)((_exafy(a, data) for a in nl.args)...) +end + +# Check if NamedTuple iterator respects the restriction +function _support_in_restriction(restriction, itr, data) + supp = [itr[data.param_alias[p]] for p in restriction.parameter_refs] + return restriction(supp) +end + +# Prepare the constraint iterator over the infinite parameters +function _get_constraint_iterator(cref, data) + group_idxs = InfiniteOpt.parameter_group_int_indices(cref) + # prepare the iterator of NamedTuples (contains support values, iterator values, and constants from parameter functions) + if isempty(group_idxs) # we have a finite constraint + itr = [(;)] + elseif length(group_idxs) == 1 # we only depend on one independent infinite parameter + itr = data.base_itrs[first(group_idxs)] + else # we depend on multiple independent infinite parameters + itrs = map(i -> data.base_itrs[i], group_idxs) + itr = vec([merge(i...) for i in Iterators.product(itrs...)]) + end + # Remove any elements of the iterator that violate the domain restriction + if InfiniteOpt.has_domain_restriction(cref) + restriction = InfiniteOpt.domain_restriction(cref) + itr = filter(i -> _support_in_restriction(restriction, i, data), itr) + end + return itr end # Finalize exafied expressions to avoid scalars @@ -439,12 +556,6 @@ function _get_constr_bounds(set) "if you need support for this constraint type, please open an issue.") end -# Check if NamedTuple iterator respects the restriction -function _support_in_restriction(restriction, itr, data) - supp = [itr[data.param_alias[p]] for p in restriction.parameter_refs] - return restriction(supp) -end - # Add all the constraints from an InfiniteModel to an ExaCore function _add_constraints( core::ExaModels.ExaCore, @@ -452,8 +563,9 @@ function _add_constraints( inf_model::InfiniteOpt.InfiniteModel ) for cref in JuMP.all_constraints(inf_model) - # skip if the constraint is a variable bound or type + # skip if the constraint is a variable bound or already added (as a grouped constraint) InfiniteOpt.is_variable_domain_constraint(cref) && continue + haskey(data.constraint_mappings, cref) && continue # parse the basic information constr = JuMP.constraint_object(cref) if isempty(inf_model.constraints[JuMP.index(cref)].measure_indices) @@ -463,24 +575,10 @@ function _add_constraints( expr = InfiniteOpt.expand_measures(JuMP.jump_function(constr), inf_model) end set = JuMP.moi_set(constr) - group_idxs = InfiniteOpt.parameter_group_int_indices(cref) - # prepare the iterator of NamedTuples (contains support values, iterator values, and constants from parameter functions) - if isempty(group_idxs) # we have a finite constraint - itr = [(;)] - elseif length(group_idxs) == 1 # we only depend on one independent infinite parameter - itr = data.base_itrs[first(group_idxs)] - else # we depend on multiple independent infinite parameters - itrs = map(i -> data.base_itrs[i], group_idxs) - itr = vec([merge(i...) for i in Iterators.product(itrs...)]) - end - # Remove any elements of the iterator that violate the domain restriction - if InfiniteOpt.has_domain_restriction(cref) - restriction = InfiniteOpt.domain_restriction(cref) - itr = filter(i -> _support_in_restriction(restriction, i, data), itr) - end + # prepare the constraint iterator + itr = _get_constraint_iterator(cref, data) # create the ExaModels expression tree based on expr - data_src = ExaModels.DataSource() - em_expr = _finalize_expr(_exafy(expr, data_src, data)) + em_expr = _finalize_expr(_exafy(expr, data)) # get the constraint bounds lb, ub = _get_constr_bounds(set) # create the ExaModels constraint @@ -493,7 +591,7 @@ end # Make dispatch type to pass the data needed by `make_reduced_expr` struct _DerivReductionBackendInfo <: InfiniteOpt.AbstractTransformationBackend data::ExaMappingData - data_src::ExaModels.DataSource + alias_map::Union{Nothing, Dict{InfiniteOpt.GeneralVariableRef, Symbol}} end # Extend make_reduced_expr to create an ExaModel expression @@ -506,7 +604,7 @@ function InfiniteOpt.make_reduced_expr( ) group_idx = InfiniteOpt.parameter_group_int_index(pref) data = dispatch_data.data - data_src = dispatch_data.data_src + data_src = ExaModels.DataSource() alias = data.group_alias[group_idx] if vref.index_type == InfiniteOpt.SemiInfiniteVariableIndex @assert haskey(data.semivar_info, vref) @@ -520,7 +618,8 @@ function InfiniteOpt.make_reduced_expr( data_src[i] end end for i in inds) - return ivar[idx_pars...] + grouped_var = data.var_to_grouped_var[vref] + return grouped_var[idx_pars..., data_src[dispatch_data.alias_map[vref]]] else # either an infinite variable or a derivative variable group_idxs = InfiniteOpt.parameter_group_int_indices(vref) idx_pars = (begin @@ -531,9 +630,9 @@ function InfiniteOpt.make_reduced_expr( data_src[g_alias] end end for i in group_idxs) - return data.infvar_mappings[vref][idx_pars...] + grouped_var = data.var_to_grouped_var[vref] + return grouped_var[idx_pars..., data_src[dispatch_data.alias_map[vref]]] end - return end # Add the approximation equations for each derivative variable @@ -542,14 +641,27 @@ function _add_derivative_approximations( data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel ) + # group all the derivatives of the same order, method, and infinite parameter dependencies + signature_to_derivs = Dict{ + Tuple{InfiniteOpt.GeneralVariableRef, Int, Vector{Int}, DataType}, + Tuple{Vector{InfiniteOpt.GeneralVariableRef}, Vector{InfiniteOpt.GeneralVariableRef}} + }() for dref in InfiniteOpt.all_derivatives(inf_model) - # gather the derivative information vref = InfiniteOpt.derivative_argument(dref) pref = InfiniteOpt.operator_parameter(dref) order = InfiniteOpt.derivative_order(dref) - method = InfiniteOpt.derivative_method(dref) - # gather the needed infinite parameter data group_idxs = InfiniteOpt.parameter_group_int_indices(vref) + if !haskey(signature_to_derivs, (pref, order, group_idxs, vref.index_type)) + signature_to_derivs[pref, order, group_idxs, vref.index_type] = + (InfiniteOpt.GeneralVariableRef[], InfiniteOpt.GeneralVariableRef[]) + end + push!(signature_to_derivs[pref, order, group_idxs, vref.index_type][1], dref) + push!(signature_to_derivs[pref, order, group_idxs, vref.index_type][2], vref) + end + # iterate over each group of derivatives and add the approximation equations + for ((pref, order, group_idxs, _), (drefs, vrefs)) in signature_to_derivs + # gather basic info + method = InfiniteOpt.derivative_method(drefs[1]) pref_group = InfiniteOpt.parameter_group_int_index(pref) # sort the base support iterator p_alias = data.param_alias[pref] @@ -561,30 +673,32 @@ function _add_derivative_approximations( end # collect the expression data supps = map(p -> p[p_alias], srt_itr) - idxs, arg_itrs... = InfiniteOpt.derivative_expr_data(dref, order, supps, method) + idxs, arg_itrs... = InfiniteOpt.derivative_expr_data(drefs[1], order, supps, method) # make the iterator aliases = Tuple(Symbol("d_arg$i") for i in eachindex(arg_itrs)) pref_itr = [(; srt_itr[i]..., zip(aliases, args)...) for (i, args...) in zip(idxs, arg_itrs...)] - if length(group_idxs) > 1 - itrs = [g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] - itr = [merge(i...) for i in Iterators.product(itrs...)] - else - itr = pref_itr - end - # make the ExaModel expression tree + itrs = Any[g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] + push!( + itrs, + [(; :grouped_didx => _get_grouped_idx(dref, data), + :grouped_vidx => _get_grouped_idx(vref, data)) + for (dref, vref) in zip(drefs, vrefs) + ]) + itr = length(itrs) > 1 ? vec([merge(i...) for i in Iterators.product(itrs...)]) : pref_itr + # make the ExaModel expression tree and add the constraint(s) data_src = ExaModels.DataSource() + alias_map = Dict(drefs[1] => :grouped_didx, vrefs[1] => :grouped_vidx) em_expr = InfiniteOpt.make_indexed_derivative_expr( - dref, - vref, - pref, - order, - data_src[data.group_alias[pref_group]], - supps, - _DerivReductionBackendInfo(data, data_src), + drefs[1], + vrefs[1], + pref, + order, + data_src[data.group_alias[pref_group]], + supps, + _DerivReductionBackendInfo(data, alias_map), method, (data_src[a] for a in aliases)... - ) - # add the constraint + ) core, _ = ExaModels.add_con(core, em_expr, itr) end return core @@ -611,18 +725,31 @@ function _add_collocation_restrictions( ubs = repeat(2+num_nodes:num_nodes+1:num_supps, inner = num_nodes) pts = filter(i -> !(i in ubs), 2:num_supps-1) pref_itr = [(i1 = ub, i2 = pt) for (ub, pt) in zip(ubs, pts)] - # make the constraints for each infinite variable + # group the variables by their input infinite parameters + group_idxs_to_vrefs = Dict{Vector{Int}, Vector{InfiniteOpt.GeneralVariableRef}}() for vidx in vidxs vref = InfiniteOpt.GeneralVariableRef(inf_model, vidx) group_idxs = InfiniteOpt.parameter_group_int_indices(vref) + if !haskey(group_idxs_to_vrefs, group_idxs) + group_idxs_to_vrefs[group_idxs] = InfiniteOpt.GeneralVariableRef[] + end + push!(group_idxs_to_vrefs[group_idxs], vref) + end + # add the constraints for each group of variables + for (group_idxs, vrefs) in group_idxs_to_vrefs + # prepare the iterator aliases = (data.group_alias[g] for g in group_idxs) itrs = (g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs) - itr = vec([merge(i...) for i in Iterators.product(itrs...)]) + finite_itr = [(; :grouped_vidx => _get_grouped_idx(vref, data)) for vref in vrefs] + itr = vec([merge(i...) for i in Iterators.product(itrs..., finite_itr)]) + # prepare the variable indices data_src = ExaModels.DataSource() - idx_pars1 = (a == pref_alias ? data_src[:i1] : data_src[a] for a in aliases) - idx_pars2 = (a == pref_alias ? data_src[:i2] : data_src[a] for a in aliases) - ivar = data.infvar_mappings[vref] - em_expr = ivar[idx_pars1...] - ivar[idx_pars2...] + alias_tuple = (aliases..., :grouped_vidx) + idx_pars1 = (a == pref_alias ? data_src[:i1] : data_src[a] for a in alias_tuple) + idx_pars2 = (a == pref_alias ? data_src[:i2] : data_src[a] for a in alias_tuple) + # create the ExaModel expression tree and add the constraint + grouped_var = data.var_to_grouped_var[vrefs[1]] + em_expr = grouped_var[idx_pars1...] - grouped_var[idx_pars2...] core, _ = ExaModels.add_con(core, em_expr, itr) end end @@ -639,7 +766,7 @@ const _ObjMeasureExpansionWarn = string( # Write a finite expression `expr` in a single objective term (this is a generic fallback) function _add_generic_objective_term(core, expr, data) - em_expr = _finalize_expr(_exafy(expr, (;), data)) + em_expr = _finalize_expr(_exafy(expr, data)) return ExaModels.add_obj(core, em_expr, [(;)])[1] end @@ -716,29 +843,44 @@ end # Helper function for adding "affine" terms as independent objective terms # Note the `coef` doesn't have to be a constant, it can be an expression that doesn't contain measures -function _add_objective_aff_term(core, coef, vref, data) - return _add_objective_aff_term(core, coef, vref, vref.index_type, data) +function _add_objective_aff_term(core, coef, vref, data, group_repeated_sums = false) + return _add_objective_aff_term(core, coef, vref, vref.index_type, data, group_repeated_sums) end -function _add_objective_aff_term(core, coef, vref, ::Type{InfiniteOpt.MeasureIndex}, data) +function _add_objective_aff_term(core, coef, vref, ::Type{InfiniteOpt.MeasureIndex}, data, group_repeated_sums) # process the measure structure recursively as needed mexpr, itr = _process_measure_sum(vref, data) + # form the exafied expression and iterator + c = ExaModels.DataSource()[:c] + if group_repeated_sums + exafied_expr, finite_itr = _process_candidate_sum_group(mexpr, data) + if length(finite_itr) > 1 + @info "Successfully grouped $(length(finite_itr)) finite terms together into a single objective pattern." + final_itr = vec([merge(i...) for i in Iterators.product(itr, finite_itr)]) + else + final_itr = itr + end + else + exafied_expr = _exafy(mexpr, data) + final_itr = itr + end # prepare the examodel expression tree - data_src = ExaModels.DataSource() - em_expr = data_src.c * _exafy(coef * mexpr, data_src, data) + em_expr = isone(coef) ? c * exafied_expr : _exafy(coef, data) * (c * exafied_expr) # add the term to the objective - core, _ = ExaModels.add_obj(core, _finalize_expr(em_expr), itr) + core, _ = ExaModels.add_obj(core, _finalize_expr(em_expr), final_itr) return core end -function _add_objective_aff_term(core, coef, vref, _, data) - return _add_generic_objective_term(core, coef * vref, data) +function _add_objective_aff_term(core, coef, vref, _, data, group_repeated_sums) + expr = isone(coef) ? vref : coef * vref + return _add_generic_objective_term(core, expr, data) end # Add the objective from an InfiniteModel to an ExaCore function _add_objective( core::ExaModels.ExaCore, - expr::JuMP.AbstractJuMPScalar, + expr::JuMP.AbstractJuMPScalar, # generic fallback (heuristics fail to find a summed measure structure) data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + group_repeated_sums::Bool = false ) vrefs = InfiniteOpt.all_expression_variables(expr) if any(v.index_type == InfiniteOpt.MeasureIndex for v in vrefs) @@ -751,19 +893,21 @@ function _add_objective( core::ExaModels.ExaCore, vref::InfiniteOpt.GeneralVariableRef, # can be finite var, point var, finite param, or measure that fully evaluates the measures inside data::ExaMappingData, - ::InfiniteOpt.InfiniteModel + ::InfiniteOpt.InfiniteModel; + group_repeated_sums::Bool = false ) - return _add_objective_aff_term(core, 1.0, vref, data) + return _add_objective_aff_term(core, 1.0, vref, data, group_repeated_sums) end function _add_objective( core::ExaModels.ExaCore, aff::JuMP.GenericAffExpr, data::ExaMappingData, - ::InfiniteOpt.InfiniteModel + ::InfiniteOpt.InfiniteModel; + group_repeated_sums::Bool = false ) - # TODO should we check if there are a lot of terms? + # TODO should we check if there are a lot of terms? (use group_repeated_sums) for (coef, vref) in JuMP.linear_terms(aff) - core = _add_objective_aff_term(core, coef, vref, data) + core = _add_objective_aff_term(core, coef, vref, data, group_repeated_sums) end c = JuMP.constant(aff) if !iszero(c) @@ -775,7 +919,8 @@ function _add_objective( core::ExaModels.ExaCore, quad::InfiniteOpt.GenericQuadExpr, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + group_repeated_sums::Bool = false ) # process the quadratic terms for (coef, vref1, vref2) in JuMP.quad_terms(quad) @@ -785,9 +930,9 @@ function _add_objective( new_expr = InfiniteOpt.expand_measures(coef * vref1 * vref2, inf_model) core = _add_generic_objective_term(core, new_expr, data) elseif vref1.index_type == InfiniteOpt.MeasureIndex - core = _add_objective_aff_term(core, coef * vref2, vref1, data) + core = _add_objective_aff_term(core, coef * vref2, vref1, data, group_repeated_sums) else - core = _add_objective_aff_term(core, coef * vref1, vref2, data) + core = _add_objective_aff_term(core, coef * vref1, vref2, data, group_repeated_sums) end end # add the affine terms @@ -800,28 +945,48 @@ end function build_exa_core!( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + group_repeated_algebraic_patterns = false ) # initial setup _build_base_iterators(data, inf_model) # add the variables and appropriate mappings core = _add_finite_parameters(core, data, inf_model) core = _add_finite_variables(core, data, inf_model) - core = _add_infinite_variables(core, data, inf_model) # includes derivatives + core = _add_infinite_variables(core, data, inf_model) core = _add_parameter_functions(core, data, inf_model) _add_semi_infinite_variables(core, data, inf_model) _add_point_variables(core, data, inf_model) # account for user-defined nonlinear operators _add_user_operators(inf_model) # add the constraints + if group_repeated_algebraic_patterns + core = _group_and_add_constraints(core, data, inf_model) # TODO: can eventually replace `_add_constraints` if it works well + num_grouped_constraints = length(core.cons) + end core = _add_constraints(core, data, inf_model) + if group_repeated_algebraic_patterns + num_ungrouped_constraints = length(core.cons) - num_grouped_constraints + end core = _add_derivative_approximations(core, data, inf_model) core = _add_collocation_restrictions(core, data, inf_model) # add the objective if there is one expr = JuMP.objective_function(inf_model) sense = JuMP.objective_sense(inf_model) if sense != _MOI.FEASIBILITY_SENSE - core = _add_objective(core, expr, data, inf_model) + core = _add_objective( + core, + expr, + data, + inf_model, + group_repeated_sums = group_repeated_algebraic_patterns + ) + end + if group_repeated_algebraic_patterns + num_con_patterns = length(core.cons) + num_grouped_constraints = num_con_patterns - num_ungrouped_constraints + @info "In total, $num_con_patterns constraint pattern(s) was/were added of which $num_grouped_constraints are grouped constraints." + @info "In total, $(length(core.obj)) objective sum pattern(s) was/were added. Check the logs to determine how many were grouped." end return core end @@ -831,24 +996,33 @@ end inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; [backend = nothing, - concrete_core::Bool = false] + concrete_core::Bool = false, + group_repeated_algebraic_patterns = false] # experimental )::ExaModels.ExaCore Create `ExaModels.ExaCore` from `inf_model` using the provided -`ExaMappingData` to store the variable and constraint mappings. +`ExaMappingData` to store the variable and constraint mappings. The setting `concrete_core = true` will create a concrete `ExaModels.ExaCore` type, which is useful for performance in some cases. +Optionally, try to aggregate common algebraic constraint and objective patterns +by setting `group_repeated_algebraic_patterns = true`. This is an +experimental feature that may encounter issues and may be removed/modified in the future. """ function ExaModels.ExaCore( inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; backend = nothing, - concrete_core::Bool = false + concrete_core::Bool = false, + group_repeated_algebraic_patterns = false ) # TODO add support for other float types once InfiniteOpt does minimize = JuMP.objective_sense(inf_model) == _MOI.MIN_SENSE core = ExaModels.ExaCore(; backend = backend, minimize = minimize, concrete = Val(concrete_core)) - return build_exa_core!(core, data, inf_model) + return build_exa_core!( + core, + data, + inf_model; group_repeated_algebraic_patterns = group_repeated_algebraic_patterns + ) end """ @@ -856,23 +1030,45 @@ end inf_model::InfiniteOpt.InfiniteModel, [data::ExaMappingData]; [backend = nothing, - concrete_core::Bool = false] + concrete_core::Bool = false, + group_repeated_algebraic_patterns = false] # experimental )::ExaModels.ExaModel Create an `ExaModels.ExaModel` from `inf_model` and store the mappings in `data`. If `data` is not provided, the mappings cannot be readily extracted. The `concrete_core` setting will create a concrete `ExaModels.ExaCore` type, which is useful for performance in some cases. +Optionally, try to aggregate common algebraic constraint/objective patterns +by setting `group_repeated_algebraic_patterns = true`. This is an +experimental feature that may encounter issues and may be removed/modified in the future. """ function ExaModels.ExaModel( inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; backend = nothing, - concrete_core::Bool = false + concrete_core::Bool = false, + group_repeated_algebraic_patterns = false + ) + core = ExaModels.ExaCore( + inf_model, + data; + backend = backend, + concrete_core = concrete_core, + group_repeated_algebraic_patterns = group_repeated_algebraic_patterns ) - core = ExaModels.ExaCore(inf_model, data; backend = backend, concrete_core = concrete_core) return ExaModels.ExaModel(core) end -function ExaModels.ExaModel(inf_model::InfiniteOpt.InfiniteModel; backend = nothing) - return ExaModels.ExaModel(inf_model, ExaMappingData(), backend = backend) +function ExaModels.ExaModel( + inf_model::InfiniteOpt.InfiniteModel; + backend = nothing, + group_repeated_algebraic_patterns = false, + concrete_core::Bool = false +) + return ExaModels.ExaModel( + inf_model, + ExaMappingData(); + backend = backend, + group_repeated_algebraic_patterns = group_repeated_algebraic_patterns, + concrete_core = concrete_core + ) end diff --git a/test/solve.jl b/test/solve.jl index 20e6b59..69a9c84 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -16,7 +16,7 @@ tol = 1E-6 yval = value(y) zval = value(z) dyval = value(∂(y, t)) - # test with ExaTranscriptionBackend + # test with ExaTranscriptionBackend @test set_transformation_backend(m, ExaTranscriptionBackend(IpoptSolver)) isa Nothing @test set_silent(m) isa Nothing @test optimize!(m).status == :first_order @@ -92,6 +92,71 @@ end end end +@testset "Dynamic Repeated Constraint Patterns" begin + m = InfiniteModel(Ipopt.Optimizer) + @infinite_parameter(m, t in [0, 1], num_supports = 11, derivative_method = OrthogonalCollocation(3)) + @variable(m, x[1:3], Infinite(t)) + @variable(m, u, Infinite(t)) + @finite_parameter(m, p == 1.0) + @parameter_function(m, pf == t->2*t) + @objective(m, Min, ∫(sum(@force_nonlinear((x[i] - p)^2) for i in 1:3), t)) + @constraint(m, c1[i in 1:3], ∂(x[i], t) == u^(i+2) - pf) + @constraint(m, c2[i in 1:3], x[i] <= i) + @constraint(m, c3[i in 1:3], x[i](0) == 0.1*i) + constant_over_collocation(u, t) + set_silent(m) + optimize!(m) + obj = objective_value(m) + xval = value.(x) + uval = value(u) + c2_duals = dual.(c2) + @test set_transformation_backend(m, ExaTranscriptionBackend(IpoptSolver)) isa Nothing + @test set_silent(m) isa Nothing + @test optimize!(m, group_repeated_algebraic_patterns = true).status == :first_order + @test isapprox(obj, objective_value(m), atol = tol) + for i in 1:3 + @test all(isapprox.(xval[i], value(x[i]), atol = tol)) + @test all(isapprox.(c2_duals[i], dual(c2[i]), atol = tol)) + end + @test all(isapprox.(uval, value(u), atol = tol)) +end + +@testset "PDE Repeated Patterns" begin + m = InfiniteModel(Ipopt.Optimizer) + @infinite_parameter(m, t in [0, 10], num_supports = 15, derivative_method = OrthogonalCollocation(3)) + @infinite_parameter(m, x in [0, 1], num_supports = 8, derivative_method = FiniteDifference(Central())) + @variable(m, y[1:3], Infinite(t, x)) + @variable(m, q, Infinite(x, t)) + @variable(m, u, Infinite(t), start = 1.0) + @variable(m, 0 <= z[1:3] <= 0.0) + @finite_parameter(m, p == 1.0) + @parameter_function(m, pf == sin(t)) + @finite_parameter(m, p2 == 3.0) + @objective(m, Min, ∫(∫((q - p2)^2, t), x)) + @constraint(m, c1[i in 1:3], ∂(y[i], t) == p*∂(y[i], x, x) + u * pf) + @constraint(m, c2[i in 1:3], y[i](0, x) == 0.1*i) + @constraint(m, c3[i in 1:3], y[i](t, 0) == z[i], DomainRestriction(s -> s != 0, t)) + @constraint(m, c4, sum(y) == q) + constant_over_collocation(u, t) + set_silent(m) + optimize!(m) + obj = objective_value(m) + yval = value.(y) + qval = value(q) + uval = value(u) + zval = value.(z) + @test set_transformation_backend(m, ExaTranscriptionBackend(IpoptSolver)) isa Nothing + @test set_silent(m) isa Nothing + @test optimize!(m, group_repeated_algebraic_patterns = true).status == :first_order + @test isapprox(obj, objective_value(m), atol = tol) + for i in 1:3 + @test all(isapprox.(yval[i], value(y[i]), atol = 1e-2)) # this is fairly tempermental depending on the hardware and OS + end + @test all(isapprox.(qval, value(q), atol = 1e-4)) + @test all(isapprox.(uval, value(u), atol = 1e-4)) + @test all(isapprox.(zval, value.(z), atol = tol)) +end + @testset "User-Defined Operators" begin model = InfiniteModel(Ipopt.Optimizer) set_silent(model) diff --git a/test/transcription.jl b/test/transcription.jl index 7407ccf..664fc52 100644 --- a/test/transcription.jl +++ b/test/transcription.jl @@ -117,9 +117,9 @@ end @test transformation_variable(x, exaBackend) == xMapping @test transformation_variable(y[1], exaBackend) == y1Mapping @test transformation_variable(y[2], exaBackend) == y2Mapping - @test InfiniteExaModels._map_variable(x, x.index_type, 0, exaData) isa ExaModels.ParameterNode{Int} - @test InfiniteExaModels._map_variable(y[1], y[1].index_type, 0, exaData) isa ExaModels.ParameterNode{Int} - @test InfiniteExaModels._map_variable(y[2], y[2].index_type, 0, exaData) isa ExaModels.ParameterNode{Int} + @test InfiniteExaModels._map_variable(x, x.index_type, exaData) isa ExaModels.ParameterNode{Int} + @test InfiniteExaModels._map_variable(y[1], y[1].index_type, exaData) isa ExaModels.ParameterNode{Int} + @test InfiniteExaModels._map_variable(y[2], y[2].index_type, exaData) isa ExaModels.ParameterNode{Int} @test length(exaModel.θ) == 3 @test exaModel.θ[xMapping.offset + 1] == 42 @test exaModel.θ[y1Mapping.offset + 1] == yVals[1]