From 6833b66de1cfc2e2f6feefc1ac5bcd80a491acff Mon Sep 17 00:00:00 2001 From: pulsipher Date: Thu, 13 Aug 2026 15:58:57 -0400 Subject: [PATCH 01/27] init commit --- src/InfiniteExaModels.jl | 1 + src/infiniteopt_backend.jl | 3 + src/template_constraints.jl | 144 ++++++++++++++++++++++++++++++++++++ src/transform.jl | 103 ++++++++++++++++---------- 4 files changed, 212 insertions(+), 39 deletions(-) create mode 100644 src/template_constraints.jl diff --git a/src/InfiniteExaModels.jl b/src/InfiniteExaModels.jl index a423aba..5d6475f 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("template_constraints.jl") include("transform.jl") export ExaMappingData, ExaTranscriptionBackend diff --git a/src/infiniteopt_backend.jl b/src/infiniteopt_backend.jl index 5b18604..16d9ad3 100644 --- a/src/infiniteopt_backend.jl +++ b/src/infiniteopt_backend.jl @@ -31,6 +31,8 @@ struct ExaMappingData Vector{Any} } } + # Finite template constraint metadata + indexed_var_to_exa_var::Dict{InfiniteOpt.GeneralVariableRef, ExaModels.Variable} # Default constructor function ExaMappingData() @@ -52,6 +54,7 @@ struct ExaMappingData Vector{Any} } }(), + Dict{InfiniteOpt.GeneralVariableRef, ExaModels.Variable}() ) end end diff --git a/src/template_constraints.jl b/src/template_constraints.jl new file mode 100644 index 0000000..24f492f --- /dev/null +++ b/src/template_constraints.jl @@ -0,0 +1,144 @@ +## 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(-2, h), push!(refs, v), consts # -2 indicates a symbolic variable +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 + +function _exafy_iterated_expr( + expr::Real, + data::ExaMappingData + ) + # TODO: finish +end +function _exafy_iterated_expr( + expr::InfiniteOpt.GeneralVariableRef, + data::ExaMappingData + ) + # TODO: finish +end +function _exafy_iterated_expr( + expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, + data::ExaMappingData + ) where {C, V} + return _exafy_iterated_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), data) +end +function _exafy_iterated_expr( + expr::JuMP.GenericNonlinearExpr, + data::ExaMappingData, + + ) + return _nl_op(expr.head)(_exafy_iterated_expr(a, data) for a in expr.args...) +end + +# Given the 1st and last variable in an array of variables, create a new ExaModels.Variable +# that represents the entire array +function _make_exa_variable(v1::ExaModels.Var, vf::ExaModels.Var, num_idxs::Int) + num_vars = vf.i - v1.i + 1 + @assert num_vars == num_idxs + return ExaModels.Variable((1:num_vars,), num_vars, v1.i, :x, nothing) +end +function _make_exa_variable(v1::ExaModels.Variable, vf::ExaModels.Variable, num_idxs::Int) + @assert v1.size == vf.size + num_vars = v1.length * num_idxs + @assert vf.offset + vf.length - v1.offset == num_vars + return ExaModels.Variable((v1.size..., 1:num_idxs), num_vars, v1.offset, :x, nothing) +end + +function _process_array_object( + obj::Vector{InfiniteOpt.GeneralVariableRef}, # TODO generalize to other arrays + data::ExaMappingData + ) + v_first = first(obj) + v_last = last(obj) + # TODO: what about infinite parameter arrays that are indexed across constraints? + v_first.index_type in (InfiniteOpt.DependentParametersIndex, InfiniteOpt.IndependentParameterIndex) && return + indexed_var = _make_exa_variable(data[v_first], data[v_last], length(obj)) + for v in obj + @assert !haskey(data.indexed_var_to_exa_var, v) + data.indexed_var_to_exa_var[v] = indexed_var + end + return +end +function _process_array_object( + obj::Vector{InfiniteOpt.InfOptConstraintRef}, # TODO generalize to other arrays + data::ExaMappingData + ) + # encode each constraint + h_ids = Vector{UInt}(undef, length(obj)) + vref_lists = Vector{Vector{InfiniteOpt.GeneralVariableRef}}(undef, length(obj)) + const_lists = Vector{Vector{Float64}}(undef, length(obj)) + for (i, cref) in enumerate(obj) + con = JuMP.constraint_object(cref) + expr = JuMP.jump_function(con) + h_id, vrefs, consts = _encode_expr(expr) + h_ids[i] = h_id + vref_lists[i] = vrefs + const_lists[i] = consts + end + # ensure we only have one unique constraint pattern in the array + length(unique!(h_ids)) == 1 || return + # determine which vrefs and consts change across the array + vrefs1 = vref_lists[1] + indexed_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] + consts1 = const_lists[1] + indexed_data_idxs = [any(l -> l[i] != consts1[i], const_lists) for i in eachindex(consts1)] + # exafy the vrefs + exafied_vrefs = Vector{Any}(undef, length(vrefs1)) + iterated_var_idx = 1 + for (i, vref) in enumerate(exafied_vrefs) + if indexed_var_idxs[i] + haskey(data.indexed_var_to_exa_var, vref) || return + base_idxs = _index_params(vref, data) + itr_alias = Symbol("itrvar$iterated_var_idx") + var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) + exafied_vrefs[i] = data.indexed_var_to_exa_var[vref][var_idxs...] + iterated_var_idx += 1 + else + exafied_vrefs[i] = _exafy(vref, data) + end + end + # exafy the consts + exafied_consts = Vector{Any}(undef, length(consts1)) + + # build the ExaModels graph for the constraint pattern + + # TODO: finish + return +end +function _process_array_object(obj, data::ExaMappingData) + return # fallback to ignore the unsupported/irrelevant object type +end + +# Extract repeated constraint patterns using the object dictionary of the InfiniteModel +function _process_array_objects( + core::ExaModels.ExaCore, + data::ExaMappingData, + inf_model::InfiniteOpt.InfiniteModel + ) + for (name, obj) in JuMP.object_dictionary(inf_model) + _process_array_object(obj, data) + end +end \ No newline at end of file diff --git a/src/transform.jl b/src/transform.jl index b6b7e70..ca66e2e 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -286,11 +286,47 @@ 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, + data::ExaMappingData + ) + return () +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 @@ -302,53 +338,46 @@ 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) ) @@ -358,19 +387,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) @@ -382,10 +410,9 @@ 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 # Finalize exafied expressions to avoid scalars @@ -450,8 +477,7 @@ function _add_constraints( itr = filter(i -> _support_in_restriction(restriction, i, data), itr) end # 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 @@ -464,7 +490,6 @@ end # Make dispatch type to pass the data needed by `make_reduced_expr` struct _DerivReductionBackendInfo <: InfiniteOpt.AbstractTransformationBackend data::ExaMappingData - data_src::ExaModels.DataSource end # Extend make_reduced_expr to create an ExaModel expression @@ -477,7 +502,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) @@ -551,7 +576,7 @@ function _add_derivative_approximations( order, data_src[data.group_alias[pref_group]], supps, - _DerivReductionBackendInfo(data, data_src), + _DerivReductionBackendInfo(data), method, (data_src[a] for a in aliases)... ) @@ -695,7 +720,7 @@ function _add_objective_aff_term(core, coef, vref, ::Type{InfiniteOpt.MeasureInd mexpr, itr = _process_measure_sum(vref, data) # prepare the examodel expression tree data_src = ExaModels.DataSource() - em_expr = data_src.c * _exafy(coef * mexpr, data_src, data) + em_expr = data_src.c * _exafy(coef * mexpr, data) # add the term to the objective core, _ = ExaModels.add_obj(core, _finalize_expr(em_expr), itr) return core From 2386e1c3fec58d5f2e89b3aac403ce7064cac924 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 16:19:57 -0400 Subject: [PATCH 02/27] experimental finite index support --- src/InfiniteExaModels.jl | 2 +- src/infiniteopt_backend.jl | 11 +- src/iterated_constraints.jl | 228 ++++++++++++++++++++++++++++++++++++ src/template_constraints.jl | 144 ----------------------- src/transform.jl | 104 +++++++++++----- 5 files changed, 308 insertions(+), 181 deletions(-) create mode 100644 src/iterated_constraints.jl delete mode 100644 src/template_constraints.jl diff --git a/src/InfiniteExaModels.jl b/src/InfiniteExaModels.jl index 5d6475f..21c6687 100644 --- a/src/InfiniteExaModels.jl +++ b/src/InfiniteExaModels.jl @@ -6,7 +6,7 @@ import InfiniteOpt.TranscriptionOpt as _TO include("infiniteopt_backend.jl") include("operators.jl") -include("template_constraints.jl") +include("iterated_constraints.jl") include("transform.jl") export ExaMappingData, ExaTranscriptionBackend diff --git a/src/infiniteopt_backend.jl b/src/infiniteopt_backend.jl index 16d9ad3..fd9f4aa 100644 --- a/src/infiniteopt_backend.jl +++ b/src/infiniteopt_backend.jl @@ -32,7 +32,7 @@ struct ExaMappingData } } # Finite template constraint metadata - indexed_var_to_exa_var::Dict{InfiniteOpt.GeneralVariableRef, ExaModels.Variable} + indexed_var_to_exa_var::Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}} # Default constructor function ExaMappingData() @@ -54,7 +54,7 @@ struct ExaMappingData Vector{Any} } }(), - Dict{InfiniteOpt.GeneralVariableRef, ExaModels.Variable}() + Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}}(), ) end end @@ -112,7 +112,7 @@ mutable struct ExaTranscriptionBackend{B} <: InfiniteOpt.AbstractTransformationB end # Constructors -function ExaTranscriptionBackend(; backend = nothing) +function ExaTranscriptionBackend(; backend = nothing, ) return ExaTranscriptionBackend( nothing, nothing, @@ -152,10 +152,11 @@ InfiniteOpt.transformation_data(backend::ExaTranscriptionBackend) = backend.data # Build out the backend function InfiniteOpt.build_transformation_backend!( model::InfiniteOpt.InfiniteModel, - backend::ExaTranscriptionBackend + backend::ExaTranscriptionBackend; + kwargs... ) empty!(backend) - backend.core = ExaModels.ExaCore(model, backend.data; backend = backend.backend) + backend.core = ExaModels.ExaCore(model, backend.data; backend = backend.backend, kwargs...) backend.model = ExaModels.ExaModel(backend.core) end diff --git a/src/iterated_constraints.jl b/src/iterated_constraints.jl new file mode 100644 index 0000000..da976c3 --- /dev/null +++ b/src/iterated_constraints.jl @@ -0,0 +1,228 @@ +## 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(-2, h), push!(refs, v), consts # -2 indicates a symbolic variable +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_iterated_expr( + ::Real, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return popfirst!(consts) +end +function _exafy_iterated_expr( + ::InfiniteOpt.GeneralVariableRef, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return popfirst!(vrefs) +end +function _exafy_iterated_expr( + expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, + vrefs::Vector{Any}, + consts::Vector{Any} + ) where {C, V} + return _exafy_iterated_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), vrefs, consts) +end +function _exafy_iterated_expr( + expr::JuMP.GenericNonlinearExpr, + vrefs::Vector{Any}, + consts::Vector{Any} + ) + return _nl_op(expr.head)((_exafy_iterated_expr(a, vrefs, consts) for a in expr.args)...) +end + +# Given the 1st and last variable in an array of variables, create a new ExaModels.Variable +# that represents the entire array +function _make_exa_variable(v1::ExaModels.Var, vf::ExaModels.Var, num_idxs::Int) + num_vars = vf.i - v1.i + 1 + @assert num_vars == num_idxs + return ExaModels.Variable((1:num_vars,), num_vars, v1.i, :x, nothing) +end +function _make_exa_variable(v1::ExaModels.Variable, vf::ExaModels.Variable, num_idxs::Int) + @assert v1.size == vf.size + num_vars = v1.length * num_idxs + @assert vf.offset + vf.length - v1.offset == num_vars + return ExaModels.Variable((v1.size..., 1:num_idxs), num_vars, v1.offset, :x, nothing) +end +function _make_exa_variable(v1::ExaModels.Parameter, vf::ExaModels.Parameter, num_idxs::Int) + @assert v1.size == vf.size + num_vars = v1.length * num_idxs + @assert vf.offset + vf.length - v1.offset == num_vars + return ExaModels.Parameter((v1.size..., 1:num_idxs), num_vars, v1.offset, nothing) +end + +# get the repeated index of a variable based on its direct exaified variable reference +function _get_repeated_idx(em_var::ExaModels.Var, repeated_var::ExaModels.Variable) + idx = em_var.i - repeated_var.offset + 1 + @assert idx in repeated_var.size[end] + return idx +end +function _get_repeated_idx(em_var::ExaModels.Variable, repeated_var::ExaModels.Variable) + idx = (em_var.offset - repeated_var.offset) ÷ em_var.length + 1 + @assert idx in repeated_var.size[end] + return idx +end + +function _process_variable_array( + obj::Vector{InfiniteOpt.GeneralVariableRef}, # TODO generalize to other arrays + data::ExaMappingData + ) + v_first = first(obj) + v_last = last(obj) + # TODO: what about infinite parameter arrays that are indexed across constraints? + v_first.index_type in (InfiniteOpt.DependentParametersIndex, InfiniteOpt.IndependentParameterIndex) && return false + indexed_var = _make_exa_variable(data[v_first], data[v_last], length(obj)) + for v in obj + @assert !haskey(data.indexed_var_to_exa_var, v) + data.indexed_var_to_exa_var[v] = indexed_var + end + return true +end +function _process_variable_array(obj, data::ExaMappingData) + return false +end + +function _process_constraint_array( + core::ExaModels.ExaCore, + obj::Vector{InfiniteOpt.InfOptConstraintRef}, # TODO generalize to other arrays + data::ExaMappingData + ) + # encode each constraint + h_ids = Vector{UInt}(undef, length(obj)) + vref_lists = Vector{Vector{InfiniteOpt.GeneralVariableRef}}(undef, length(obj)) + const_lists = Vector{Vector{Float64}}(undef, length(obj)) + sets = Vector{_MOI.AbstractSet}(undef, length(obj)) + for (i, cref) in enumerate(obj) + isempty(JuMP.owner_model(cref).constraints[JuMP.index(cref)].measure_indices) || return false, core + con = JuMP.constraint_object(cref) + expr = JuMP.jump_function(con) + h_id, vrefs, consts = _encode_expr(expr) + h_ids[i] = h_id + vref_lists[i] = vrefs + const_lists[i] = consts + sets[i] = JuMP.moi_set(con) + end + # ensure we only have one unique constraint pattern in the array + length(unique!(h_ids)) == 1 || return false + # determine which vrefs and consts change across the array + vrefs1 = vref_lists[1] + indexed_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] + consts1 = const_lists[1] + indexed_data_idxs = [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(obj)] + iterated_var_idx = 1 + for (i, vref) in enumerate(vrefs1) + if indexed_var_idxs[i] + haskey(data.indexed_var_to_exa_var, vref) || return false, core + base_idxs = _index_params(vref, data) + itr_alias = Symbol("itrvar$iterated_var_idx") + var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) + src_var = data.indexed_var_to_exa_var[vref] + exafied_vrefs[i] = src_var[var_idxs...] + for j in 1:length(obj) + em_var = data[vref_lists[j][i]] + var_itr[j] = (; var_itr[j]..., itr_alias => _get_repeated_idx(em_var, src_var)) + end + iterated_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(obj)] + iteratered_const_idx = 1 + for (i, c) in enumerate(consts1) + if indexed_data_idxs[i] + itr_alias = Symbol("itrconst$iteratered_const_idx") + exafied_consts[i] = ExaModels.DataSource()[itr_alias] + for j in 1:length(obj) + const_itr[j] = (; const_itr[j]..., itr_alias => const_lists[j][i]) + end + iteratered_const_idx += 1 + else + exafied_consts[i] = c + end + end + # build the ExaModels graph for the constraint pattern + raw_expr = JuMP.jump_function(JuMP.constraint_object(first(obj))) + em_expr = _finalize_expr(_exafy_iterated_expr(raw_expr, exafied_vrefs, exafied_consts)) + # process the iterator + infinite_itr = _get_constraint_iterator(first(obj), data) + finite_itr = [merge(var_itr[i], const_itr[i]) for i in 1:length(obj)] + itr = vec([merge(i...) for i in Iterators.product(infinite_itr, finite_itr)]) + # add the constraints to the core + lbs = Vector{Float64}(undef, length(obj)) + ubs = Vector{Float64}(undef, length(obj)) + 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(obj) + 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, (1:inf_len,), nothing) + end + return true, core +end +function _process_constraint_array(core::ExaModels.ExaCore, obj, data::ExaMappingData) + return false, core # fallback to ignore the unsupported/irrelevant object type +end + +# Extract repeated constraint patterns using the object dictionary of the InfiniteModel +function _process_array_objects( + core::ExaModels.ExaCore, + data::ExaMappingData, + inf_model::InfiniteOpt.InfiniteModel; + print_info::Bool = true + ) + # create ExaModels.Variable for each array of variables in the model + # have to handle variables first before constraints since dictionary is unordered + for (name, obj) in JuMP.object_dictionary(inf_model) + success = _process_variable_array(obj, data) + if success && print_info + @info "Successfully processed finite indices for variable array: $name" + end + end + # create the necessary metadata for each array of constraints in the model + for (name, obj) in JuMP.object_dictionary(inf_model) + success, core = _process_constraint_array(core, obj, data) + if success && print_info + @info "Successfully added finite indices for constraint array: $name" + end + end + return core +end diff --git a/src/template_constraints.jl b/src/template_constraints.jl deleted file mode 100644 index 24f492f..0000000 --- a/src/template_constraints.jl +++ /dev/null @@ -1,144 +0,0 @@ -## 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(-2, h), push!(refs, v), consts # -2 indicates a symbolic variable -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 - -function _exafy_iterated_expr( - expr::Real, - data::ExaMappingData - ) - # TODO: finish -end -function _exafy_iterated_expr( - expr::InfiniteOpt.GeneralVariableRef, - data::ExaMappingData - ) - # TODO: finish -end -function _exafy_iterated_expr( - expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, - data::ExaMappingData - ) where {C, V} - return _exafy_iterated_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), data) -end -function _exafy_iterated_expr( - expr::JuMP.GenericNonlinearExpr, - data::ExaMappingData, - - ) - return _nl_op(expr.head)(_exafy_iterated_expr(a, data) for a in expr.args...) -end - -# Given the 1st and last variable in an array of variables, create a new ExaModels.Variable -# that represents the entire array -function _make_exa_variable(v1::ExaModels.Var, vf::ExaModels.Var, num_idxs::Int) - num_vars = vf.i - v1.i + 1 - @assert num_vars == num_idxs - return ExaModels.Variable((1:num_vars,), num_vars, v1.i, :x, nothing) -end -function _make_exa_variable(v1::ExaModels.Variable, vf::ExaModels.Variable, num_idxs::Int) - @assert v1.size == vf.size - num_vars = v1.length * num_idxs - @assert vf.offset + vf.length - v1.offset == num_vars - return ExaModels.Variable((v1.size..., 1:num_idxs), num_vars, v1.offset, :x, nothing) -end - -function _process_array_object( - obj::Vector{InfiniteOpt.GeneralVariableRef}, # TODO generalize to other arrays - data::ExaMappingData - ) - v_first = first(obj) - v_last = last(obj) - # TODO: what about infinite parameter arrays that are indexed across constraints? - v_first.index_type in (InfiniteOpt.DependentParametersIndex, InfiniteOpt.IndependentParameterIndex) && return - indexed_var = _make_exa_variable(data[v_first], data[v_last], length(obj)) - for v in obj - @assert !haskey(data.indexed_var_to_exa_var, v) - data.indexed_var_to_exa_var[v] = indexed_var - end - return -end -function _process_array_object( - obj::Vector{InfiniteOpt.InfOptConstraintRef}, # TODO generalize to other arrays - data::ExaMappingData - ) - # encode each constraint - h_ids = Vector{UInt}(undef, length(obj)) - vref_lists = Vector{Vector{InfiniteOpt.GeneralVariableRef}}(undef, length(obj)) - const_lists = Vector{Vector{Float64}}(undef, length(obj)) - for (i, cref) in enumerate(obj) - con = JuMP.constraint_object(cref) - expr = JuMP.jump_function(con) - h_id, vrefs, consts = _encode_expr(expr) - h_ids[i] = h_id - vref_lists[i] = vrefs - const_lists[i] = consts - end - # ensure we only have one unique constraint pattern in the array - length(unique!(h_ids)) == 1 || return - # determine which vrefs and consts change across the array - vrefs1 = vref_lists[1] - indexed_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] - consts1 = const_lists[1] - indexed_data_idxs = [any(l -> l[i] != consts1[i], const_lists) for i in eachindex(consts1)] - # exafy the vrefs - exafied_vrefs = Vector{Any}(undef, length(vrefs1)) - iterated_var_idx = 1 - for (i, vref) in enumerate(exafied_vrefs) - if indexed_var_idxs[i] - haskey(data.indexed_var_to_exa_var, vref) || return - base_idxs = _index_params(vref, data) - itr_alias = Symbol("itrvar$iterated_var_idx") - var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) - exafied_vrefs[i] = data.indexed_var_to_exa_var[vref][var_idxs...] - iterated_var_idx += 1 - else - exafied_vrefs[i] = _exafy(vref, data) - end - end - # exafy the consts - exafied_consts = Vector{Any}(undef, length(consts1)) - - # build the ExaModels graph for the constraint pattern - - # TODO: finish - return -end -function _process_array_object(obj, data::ExaMappingData) - return # fallback to ignore the unsupported/irrelevant object type -end - -# Extract repeated constraint patterns using the object dictionary of the InfiniteModel -function _process_array_objects( - core::ExaModels.ExaCore, - data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel - ) - for (name, obj) in JuMP.object_dictionary(inf_model) - _process_array_object(obj, data) - end -end \ No newline at end of file diff --git a/src/transform.jl b/src/transform.jl index ca66e2e..84f7480 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -415,6 +415,32 @@ function _exafy( 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 _finalize_expr(expr) = expr _finalize_expr(c::Real) = ExaModels.Null(c) @@ -437,12 +463,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,6 +472,7 @@ function _add_constraints( for cref in JuMP.all_constraints(inf_model) # skip if the constraint is a variable bound or type 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) @@ -461,21 +482,8 @@ 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 em_expr = _finalize_expr(_exafy(expr, data)) # get the constraint bounds @@ -796,7 +804,8 @@ end function build_exa_core!( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + process_finite_indices_of_named_constraints = false ) # initial setup _build_base_iterators(data, inf_model) @@ -808,6 +817,9 @@ function build_exa_core!( _add_semi_infinite_variables(core, data, inf_model) _add_point_variables(core, data, inf_model) # add the constraints + if process_finite_indices_of_named_constraints + core = _process_array_objects(core, data, inf_model) + end core = _add_constraints(core, data, inf_model) core = _add_derivative_approximations(core, data, inf_model) core = _add_collocation_restrictions(core, data, inf_model) @@ -824,41 +836,71 @@ end ExaModels.ExaCore( inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; - [backend = nothing] + [backend = nothing, + process_finite_indices_of_named_constraints = 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. +Optionally, try to aggregate common algebraic constraint structures +over a named constraint array (e.g., `@constraint(model, con[i=1:N], ...)`) +by setting `process_finite_indices_of_named_constraints = 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 + backend = nothing, + process_finite_indices_of_named_constraints = 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(true)) - return build_exa_core!(core, data, inf_model) + return build_exa_core!( + core, + data, + inf_model; process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints + ) end """ ExaModels.ExaModel( inf_model::InfiniteOpt.InfiniteModel, [data::ExaMappingData]; - [backend = nothing] + [backend = nothing, + process_finite_indices_of_named_constraints = 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. +Optionally, try to aggregate common algebraic constraint structures +over a named constraint array (e.g., `@constraint(model, con[i=1:N], ...)`) +by setting `process_finite_indices_of_named_constraints = 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 + backend = nothing, + process_finite_indices_of_named_constraints = false + ) + core = ExaModels.ExaCore( + inf_model, + data; + backend = backend, + process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints ) - core = ExaModels.ExaCore(inf_model, data; backend = backend) 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, + process_finite_indices_of_named_constraints = false +) + return ExaModels.ExaModel( + inf_model, + ExaMappingData(); + backend = backend, + process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints + ) end From 98f7ac931ad09ba2d1a4d45a1c51d489fdf96c94 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 16:51:36 -0400 Subject: [PATCH 03/27] Support AbstractArray constraint groups --- src/iterated_constraints.jl | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/iterated_constraints.jl b/src/iterated_constraints.jl index da976c3..7fa3b86 100644 --- a/src/iterated_constraints.jl +++ b/src/iterated_constraints.jl @@ -110,7 +110,7 @@ end function _process_constraint_array( core::ExaModels.ExaCore, - obj::Vector{InfiniteOpt.InfOptConstraintRef}, # TODO generalize to other arrays + obj::Vector{InfiniteOpt.InfOptConstraintRef}, data::ExaMappingData ) # encode each constraint @@ -198,6 +198,20 @@ function _process_constraint_array( end return true, core end +function _process_constraint_array( + core::ExaModels.ExaCore, + obj::AbstractArray{<:InfiniteOpt.InfOptConstraintRef}, + data::ExaMappingData + ) + return _process_constraint_array(core, vec(collect(values(obj))), data) +end +function _process_constraint_array( + core::ExaModels.ExaCore, + obj::JuMP.Containers.SparseAxisArray{<:InfiniteOpt.InfOptConstraintRef}, + data::ExaMappingData + ) + return _process_constraint_array(core, collect(values(obj.data)), data) +end function _process_constraint_array(core::ExaModels.ExaCore, obj, data::ExaMappingData) return false, core # fallback to ignore the unsupported/irrelevant object type end From 40f6c68163c206998a0974120e5a973516e840b6 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 17:03:10 -0400 Subject: [PATCH 04/27] test fix --- test/transcription.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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] From 2aab7206931ce6169a80f197b7252f7fd471d9e1 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 17:12:51 -0400 Subject: [PATCH 05/27] Add support for variables stored in multi-dimensional dense arrays --- src/iterated_constraints.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/iterated_constraints.jl b/src/iterated_constraints.jl index 7fa3b86..829373f 100644 --- a/src/iterated_constraints.jl +++ b/src/iterated_constraints.jl @@ -90,7 +90,7 @@ function _get_repeated_idx(em_var::ExaModels.Variable, repeated_var::ExaModels.V end function _process_variable_array( - obj::Vector{InfiniteOpt.GeneralVariableRef}, # TODO generalize to other arrays + obj::Array{<:InfiniteOpt.GeneralVariableRef}, data::ExaMappingData ) v_first = first(obj) @@ -104,6 +104,12 @@ function _process_variable_array( end return true end +function _process_variable_array( + obj::JuMP.Containers.DenseAxisArray{<:InfiniteOpt.GeneralVariableRef}, + data::ExaMappingData + ) + return _process_variable_array(obj.data, data) +end function _process_variable_array(obj, data::ExaMappingData) return false end From 615ae4e418d298601514b733b56d4afd26deefae Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 17:17:55 -0400 Subject: [PATCH 06/27] bug fix --- src/transform.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transform.jl b/src/transform.jl index 1364cba..58b479b 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -372,7 +372,7 @@ function _map_variable(vref, ::Type{InfiniteOpt.SemiInfiniteVariableIndex}, data return ivar[idx_pars...] end function _map_variable(vref, ::Type{<:InfiniteOpt.InfiniteParameterIndex}, data) - return ExaModels.DataSource[data.param_alias[vref]] + return ExaModels.DataSource()[data.param_alias[vref]] end function _map_variable(vref, ::Type{InfiniteOpt.FiniteParameterIndex}, data) return data.param_mappings[vref][1] From ba6485bef949b314df4a19a0fcf31fd4e096961c Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 17:25:45 -0400 Subject: [PATCH 07/27] bug fix --- src/transform.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transform.jl b/src/transform.jl index 58b479b..7e897e3 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -666,7 +666,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 From 67ef4e48e0ee7f37ac83307f262ffe41d9443410 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 14 Aug 2026 19:26:52 -0400 Subject: [PATCH 08/27] support sparse array variables --- src/iterated_constraints.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/iterated_constraints.jl b/src/iterated_constraints.jl index 829373f..88ebda8 100644 --- a/src/iterated_constraints.jl +++ b/src/iterated_constraints.jl @@ -110,6 +110,12 @@ function _process_variable_array( ) return _process_variable_array(obj.data, data) end +function _process_variable_array( + obj::JuMP.Containers.SparseAxisArray{<:InfiniteOpt.GeneralVariableRef}, + data::ExaMappingData + ) + return _process_variable_array(collect(values(obj.data)), data) +end function _process_variable_array(obj, data::ExaMappingData) return false end From dac2755c37742a6971ddba0471155b53fdb54168 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Tue, 18 Aug 2026 16:43:15 -0400 Subject: [PATCH 09/27] general refactor --- src/InfiniteExaModels.jl | 2 +- src/grouped_constraints.jl | 216 ++++++++++++++++++++++++++++++ src/infiniteopt_backend.jl | 2 +- src/iterated_constraints.jl | 254 ------------------------------------ src/transform.jl | 161 +++++++++++++++++++---- 5 files changed, 351 insertions(+), 284 deletions(-) create mode 100644 src/grouped_constraints.jl delete mode 100644 src/iterated_constraints.jl diff --git a/src/InfiniteExaModels.jl b/src/InfiniteExaModels.jl index 21c6687..cbfa5f0 100644 --- a/src/InfiniteExaModels.jl +++ b/src/InfiniteExaModels.jl @@ -6,7 +6,7 @@ import InfiniteOpt.TranscriptionOpt as _TO include("infiniteopt_backend.jl") include("operators.jl") -include("iterated_constraints.jl") +include("grouped_constraints.jl") include("transform.jl") export ExaMappingData, ExaTranscriptionBackend diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl new file mode 100644 index 0000000..c2c1bd0 --- /dev/null +++ b/src/grouped_constraints.jl @@ -0,0 +1,216 @@ +# 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.FiniteVariableIndex => -5, + InfiniteOpt.SemiInfiniteVariableIndex => -6, + InfiniteOpt.PointVariableIndex => -7, + InfiniteOpt.IndependentParameterIndex => -8, + InfiniteOpt.DependentParameterIndex => -9, + InfiniteOpt.MeasureIndex => -10, +) + +## 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(_VariableTypeHashingInt[v.index_type], 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 idx in grouped_var.size[end] + return idx +end +function _get_grouped_idx(em_var::ExaModels.Variable, grouped_var::ExaModels.Variable) + idx = (em_var.offset - grouped_var.offset) ÷ em_var.length + 1 + @assert idx in grouped_var.size[end] + return idx +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} + ) + # determine which vrefs and consts change across the array + vrefs1 = vref_lists[1] + grouped_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] + consts1 = const_lists[1] + grouped_data_idxs = [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(crefs)] + group_var_idx = 1 + for (i, vref) in enumerate(vrefs1) + if grouped_var_idxs[i] + if !haskey(data.var_to_grouped_var, vref) + _group_info_msg(crefs, "Failure: $(vref) is not a grouped variable/parameter which prevents adding") + return false, core + end + base_idxs = _index_params(vref, data) + itr_alias = Symbol("grouped_var$group_var_idx") + var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) + src_var = data.var_to_grouped_var[vref] + exafied_vrefs[i] = src_var[var_idxs...] + for j in 1:length(crefs) + infvar = vref_lists[j][i] + if data.var_to_grouped_var[infvar] != src_var + _group_info_msg(crefs, "Failure: Unable to correctly set up variable grouping for $(infvar) which prevents adding") + return false, core + end + em_var = data[infvar] + var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(em_var, src_var)) + 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(crefs)] + grouped_const_idx = 1 + for (i, c) in enumerate(consts1) + if grouped_data_idxs[i] + itr_alias = Symbol("grouped_const$grouped_const_idx") + exafied_consts[i] = ExaModels.DataSource()[itr_alias] + for j in 1:length(crefs) + 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 for the constraint pattern + raw_expr = JuMP.jump_function(JuMP.constraint_object(first(crefs))) + em_expr = _finalize_expr(_exafy_grouped_expr(raw_expr, exafied_vrefs, exafied_consts)) + # process the iterator + infinite_itr = _get_constraint_iterator(first(crefs), data) + finite_itr = [merge(var_itr[i], const_itr[i]) for i in 1:length(crefs)] + 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, (1:inf_len,), nothing) + end + return true, 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 + success = _process_candidate_constraint_group(core, data, crefs, hash_to_patterns[h]...) + if success + _group_info_msg(crefs, "Successfully added") + end + end + return core +end \ No newline at end of file diff --git a/src/infiniteopt_backend.jl b/src/infiniteopt_backend.jl index fd9f4aa..8be4919 100644 --- a/src/infiniteopt_backend.jl +++ b/src/infiniteopt_backend.jl @@ -32,7 +32,7 @@ struct ExaMappingData } } # Finite template constraint metadata - indexed_var_to_exa_var::Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}} + var_to_grouped_var::Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}} # Default constructor function ExaMappingData() diff --git a/src/iterated_constraints.jl b/src/iterated_constraints.jl deleted file mode 100644 index 88ebda8..0000000 --- a/src/iterated_constraints.jl +++ /dev/null @@ -1,254 +0,0 @@ -## 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(-2, h), push!(refs, v), consts # -2 indicates a symbolic variable -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_iterated_expr( - ::Real, - vrefs::Vector{Any}, - consts::Vector{Any} - ) - return popfirst!(consts) -end -function _exafy_iterated_expr( - ::InfiniteOpt.GeneralVariableRef, - vrefs::Vector{Any}, - consts::Vector{Any} - ) - return popfirst!(vrefs) -end -function _exafy_iterated_expr( - expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, - vrefs::Vector{Any}, - consts::Vector{Any} - ) where {C, V} - return _exafy_iterated_expr(convert(JuMP.GenericNonlinearExpr{V}, expr), vrefs, consts) -end -function _exafy_iterated_expr( - expr::JuMP.GenericNonlinearExpr, - vrefs::Vector{Any}, - consts::Vector{Any} - ) - return _nl_op(expr.head)((_exafy_iterated_expr(a, vrefs, consts) for a in expr.args)...) -end - -# Given the 1st and last variable in an array of variables, create a new ExaModels.Variable -# that represents the entire array -function _make_exa_variable(v1::ExaModels.Var, vf::ExaModels.Var, num_idxs::Int) - num_vars = vf.i - v1.i + 1 - @assert num_vars == num_idxs - return ExaModels.Variable((1:num_vars,), num_vars, v1.i, :x, nothing) -end -function _make_exa_variable(v1::ExaModels.Variable, vf::ExaModels.Variable, num_idxs::Int) - @assert v1.size == vf.size - num_vars = v1.length * num_idxs - @assert vf.offset + vf.length - v1.offset == num_vars - return ExaModels.Variable((v1.size..., 1:num_idxs), num_vars, v1.offset, :x, nothing) -end -function _make_exa_variable(v1::ExaModels.Parameter, vf::ExaModels.Parameter, num_idxs::Int) - @assert v1.size == vf.size - num_vars = v1.length * num_idxs - @assert vf.offset + vf.length - v1.offset == num_vars - return ExaModels.Parameter((v1.size..., 1:num_idxs), num_vars, v1.offset, nothing) -end - -# get the repeated index of a variable based on its direct exaified variable reference -function _get_repeated_idx(em_var::ExaModels.Var, repeated_var::ExaModels.Variable) - idx = em_var.i - repeated_var.offset + 1 - @assert idx in repeated_var.size[end] - return idx -end -function _get_repeated_idx(em_var::ExaModels.Variable, repeated_var::ExaModels.Variable) - idx = (em_var.offset - repeated_var.offset) ÷ em_var.length + 1 - @assert idx in repeated_var.size[end] - return idx -end - -function _process_variable_array( - obj::Array{<:InfiniteOpt.GeneralVariableRef}, - data::ExaMappingData - ) - v_first = first(obj) - v_last = last(obj) - # TODO: what about infinite parameter arrays that are indexed across constraints? - v_first.index_type in (InfiniteOpt.DependentParametersIndex, InfiniteOpt.IndependentParameterIndex) && return false - indexed_var = _make_exa_variable(data[v_first], data[v_last], length(obj)) - for v in obj - @assert !haskey(data.indexed_var_to_exa_var, v) - data.indexed_var_to_exa_var[v] = indexed_var - end - return true -end -function _process_variable_array( - obj::JuMP.Containers.DenseAxisArray{<:InfiniteOpt.GeneralVariableRef}, - data::ExaMappingData - ) - return _process_variable_array(obj.data, data) -end -function _process_variable_array( - obj::JuMP.Containers.SparseAxisArray{<:InfiniteOpt.GeneralVariableRef}, - data::ExaMappingData - ) - return _process_variable_array(collect(values(obj.data)), data) -end -function _process_variable_array(obj, data::ExaMappingData) - return false -end - -function _process_constraint_array( - core::ExaModels.ExaCore, - obj::Vector{InfiniteOpt.InfOptConstraintRef}, - data::ExaMappingData - ) - # encode each constraint - h_ids = Vector{UInt}(undef, length(obj)) - vref_lists = Vector{Vector{InfiniteOpt.GeneralVariableRef}}(undef, length(obj)) - const_lists = Vector{Vector{Float64}}(undef, length(obj)) - sets = Vector{_MOI.AbstractSet}(undef, length(obj)) - for (i, cref) in enumerate(obj) - isempty(JuMP.owner_model(cref).constraints[JuMP.index(cref)].measure_indices) || return false, core - con = JuMP.constraint_object(cref) - expr = JuMP.jump_function(con) - h_id, vrefs, consts = _encode_expr(expr) - h_ids[i] = h_id - vref_lists[i] = vrefs - const_lists[i] = consts - sets[i] = JuMP.moi_set(con) - end - # ensure we only have one unique constraint pattern in the array - length(unique!(h_ids)) == 1 || return false - # determine which vrefs and consts change across the array - vrefs1 = vref_lists[1] - indexed_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] - consts1 = const_lists[1] - indexed_data_idxs = [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(obj)] - iterated_var_idx = 1 - for (i, vref) in enumerate(vrefs1) - if indexed_var_idxs[i] - haskey(data.indexed_var_to_exa_var, vref) || return false, core - base_idxs = _index_params(vref, data) - itr_alias = Symbol("itrvar$iterated_var_idx") - var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) - src_var = data.indexed_var_to_exa_var[vref] - exafied_vrefs[i] = src_var[var_idxs...] - for j in 1:length(obj) - em_var = data[vref_lists[j][i]] - var_itr[j] = (; var_itr[j]..., itr_alias => _get_repeated_idx(em_var, src_var)) - end - iterated_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(obj)] - iteratered_const_idx = 1 - for (i, c) in enumerate(consts1) - if indexed_data_idxs[i] - itr_alias = Symbol("itrconst$iteratered_const_idx") - exafied_consts[i] = ExaModels.DataSource()[itr_alias] - for j in 1:length(obj) - const_itr[j] = (; const_itr[j]..., itr_alias => const_lists[j][i]) - end - iteratered_const_idx += 1 - else - exafied_consts[i] = c - end - end - # build the ExaModels graph for the constraint pattern - raw_expr = JuMP.jump_function(JuMP.constraint_object(first(obj))) - em_expr = _finalize_expr(_exafy_iterated_expr(raw_expr, exafied_vrefs, exafied_consts)) - # process the iterator - infinite_itr = _get_constraint_iterator(first(obj), data) - finite_itr = [merge(var_itr[i], const_itr[i]) for i in 1:length(obj)] - itr = vec([merge(i...) for i in Iterators.product(infinite_itr, finite_itr)]) - # add the constraints to the core - lbs = Vector{Float64}(undef, length(obj)) - ubs = Vector{Float64}(undef, length(obj)) - 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(obj) - 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, (1:inf_len,), nothing) - end - return true, core -end -function _process_constraint_array( - core::ExaModels.ExaCore, - obj::AbstractArray{<:InfiniteOpt.InfOptConstraintRef}, - data::ExaMappingData - ) - return _process_constraint_array(core, vec(collect(values(obj))), data) -end -function _process_constraint_array( - core::ExaModels.ExaCore, - obj::JuMP.Containers.SparseAxisArray{<:InfiniteOpt.InfOptConstraintRef}, - data::ExaMappingData - ) - return _process_constraint_array(core, collect(values(obj.data)), data) -end -function _process_constraint_array(core::ExaModels.ExaCore, obj, data::ExaMappingData) - return false, core # fallback to ignore the unsupported/irrelevant object type -end - -# Extract repeated constraint patterns using the object dictionary of the InfiniteModel -function _process_array_objects( - core::ExaModels.ExaCore, - data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel; - print_info::Bool = true - ) - # create ExaModels.Variable for each array of variables in the model - # have to handle variables first before constraints since dictionary is unordered - for (name, obj) in JuMP.object_dictionary(inf_model) - success = _process_variable_array(obj, data) - if success && print_info - @info "Successfully processed finite indices for variable array: $name" - end - end - # create the necessary metadata for each array of constraints in the model - for (name, obj) in JuMP.object_dictionary(inf_model) - success, core = _process_constraint_array(core, obj, data) - if success && print_info - @info "Successfully added finite indices for constraint array: $name" - end - end - return core -end diff --git a/src/transform.jl b/src/transform.jl index 7e897e3..4d708aa 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -107,14 +107,22 @@ end function _add_finite_variables( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + create_variable_groups::Bool = false ) - for vref in JuMP.all_variables(inf_model, InfiniteOpt.FiniteVariable) + vrefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteVariable) + if create_variable_groups + ex_var = ExaModels.Variable((1:length(vrefs),), length(vrefs), core.nvar, :x, nothing) + end + for vref in vrefs info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo _ensure_continuous(info) lb, ub, start = _get_variable_bounds_and_start(info) core, new_var = ExaModels.add_var(core, 1, start = start, lvar = lb, uvar = ub) data.finvar_mappings[vref] = new_var[1] + if create_variable_groups + data.var_to_grouped_var[vref] = ex_var + end end return core end @@ -123,12 +131,20 @@ end function _add_finite_parameters( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + create_parameter_groups::Bool = false ) - for pref in JuMP.all_variables(inf_model, InfiniteOpt.FiniteParameter) + prefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteParameter) + if create_parameter_groups + ex_par = ExaModels.Parameter((1:length(prefs),), length(prefs), core.npar, nothing) + end + for pref in prefs param_val = InfiniteOpt.parameter_value(pref) core, new_par = ExaModels.add_par(core, [param_val]) data.param_mappings[pref] = new_par + if create_parameter_groups + data.var_to_grouped_var[pref] = ex_par + end end return core end @@ -137,14 +153,28 @@ end function _add_infinite_variables( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + create_variable_groups::Bool = false ) # 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) + vrefs = append!(ivrefs, drefs) + # sort by infinite parameter groups (required to capture repeated finite patterns) + if create_variable_groups && length(data.base_itrs) > 1 + all_group_idxs = map(v -> InfiniteOpt.parameter_group_int_indices(v), vrefs) + sorted_int_idxs = sortperm(all_group_idxs) + permute!(vrefs, sorted_int_idxs) + permute!(all_group_idxs, sorted_int_idxs) + group_types = unique(all_group_idxs) + vref_groups = Dict(t => InfiniteOpt.GeneralVariableRef[] for t in group_types) + for (vref, group_idx) in zip(vrefs, all_group_idxs) + push!(vref_groups[group_idx], vref) + end + end # now process and add each infinite variable - for vref in append!(ivrefs, drefs) + for vref in vrefs # retrieve basic information info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo _ensure_continuous(info) @@ -157,6 +187,25 @@ function _add_infinite_variables( core, new_var = ExaModels.add_var(core, dims...; start = start, lvar = lb, uvar = ub) data.infvar_mappings[vref] = new_var end + # process and store a grouped variable for each group of variables with common infinite parameter groups + if create_variable_groups && length(data.base_itrs) > 1 + for group_idx in group_types + group_vrefs = vref_groups[group_idx] + v1 = first(group_vrefs) + num_vars = v1.length * length(group_vrefs) + ex_var = ExaModels.Variable((v1.size..., 1:length(group_vrefs)), num_vars, v1.offset, :x, nothing) + for vref in group_vrefs + data.var_to_grouped_var[vref] = ex_var + end + end + elseif create_variable_groups + v1 = first(vrefs) + num_vars = v1.length * length(vrefs) + ex_var = ExaModels.Variable((v1.size..., 1:length(vrefs)), num_vars, v1.offset, :x, nothing) + for vref in vrefs + data.var_to_grouped_var[vref] = ex_var + end + end return core end @@ -164,9 +213,24 @@ end function _add_parameter_functions( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + create_parameter_groups::Bool = false ) - for pfref in InfiniteOpt.all_parameter_functions(inf_model) + pfrefs = InfiniteOpt.all_parameter_functions(inf_model) + # sort by infinite parameter groups (required to capture repeated finite patterns) + if create_parameter_groups && length(data.base_itrs) > 1 + all_group_idxs = map(v -> InfiniteOpt.parameter_group_int_indices(v), pfrefs) + sorted_int_idxs = sortperm(all_group_idxs) + permute!(pfrefs, sorted_int_idxs) + permute!(all_group_idxs, sorted_int_idxs) + group_types = unique(all_group_idxs) + pfref_groups = Dict(t => InfiniteOpt.GeneralVariableRef[] for t in group_types) + for (pfref, group_idx) in zip(pfrefs, all_group_idxs) + push!(pfref_groups[group_idx], pfref) + end + end + # iterate and add each parameter function + for pfref in pfrefs # gather the basic information group_idxs = InfiniteOpt.parameter_group_int_indices(pfref) pfunc = InfiniteOpt.core_object(pfref) @@ -182,6 +246,25 @@ function _add_parameter_functions( core, new_par = ExaModels.add_par(core, vals) data.param_mappings[pfref] = new_par end + # process and store a grouped variable for each group of variables with common infinite parameter groups + if create_parameter_groups && length(data.base_itrs) > 1 + for group_idx in group_types + group_pfrefs = pfref_groups[group_idx] + pf1 = first(group_pfrefs) + num_pars = pf1.length * length(group_pfrefs) + ex_var = ExaModels.Parameter((pf1.size..., 1:length(group_pfrefs)), num_pars, pf1.offset, nothing) + for pfref in group_pfrefs + data.var_to_grouped_var[pfref] = ex_var + end + end + elseif create_parameter_groups + pf1 = first(pfrefs) + num_pars = pf1.length * length(pfrefs) + ex_par = ExaModels.Parameter((pf1.size..., 1:length(pfrefs)), num_pars, pf1.offset, nothing) + for pfref in pfrefs + data.var_to_grouped_var[pfref] = ex_par + end + end return core end @@ -493,7 +576,7 @@ 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 @@ -564,6 +647,7 @@ function InfiniteOpt.make_reduced_expr( end # Add the approximation equations for each derivative variable +# TODO: group together all the derivatives that share derivative method, order, and infinite parameters function _add_derivative_approximations( core::ExaModels.ExaCore, data::ExaMappingData, @@ -618,6 +702,7 @@ function _add_derivative_approximations( end # Add the constraints needed for piecewise constant control variables +# TODO: group together all the piecewise variables that share the same infinite parameter(s) function _add_collocation_restrictions( core::ExaModels.ExaCore, data::ExaMappingData, @@ -828,22 +913,42 @@ function build_exa_core!( core::ExaModels.ExaCore, data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel; - process_finite_indices_of_named_constraints = false + group_repeated_constraint_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_parameter_functions(core, data, inf_model) - _add_semi_infinite_variables(core, data, inf_model) - _add_point_variables(core, data, inf_model) + core = _add_finite_parameters( + core, + data, + inf_model, + create_parameter_groups = group_repeated_constraint_patterns + ) + core = _add_finite_variables( + core, + data, + inf_model, + create_variable_groups = group_repeated_constraint_patterns + ) + core = _add_infinite_variables( # includes derivatives + core, + data, + inf_model, + create_variable_groups = group_repeated_constraint_patterns + ) + core = _add_parameter_functions( + core, + data, + inf_model, + create_parameter_groups = group_repeated_constraint_patterns + ) + _add_semi_infinite_variables(core, data, inf_model) # TODO: add support for grouping semi-infinite variables + _add_point_variables(core, data, inf_model) # TODO: add support for grouping point variables # account for user-defined nonlinear operators _add_user_operators(inf_model) # add the constraints - if process_finite_indices_of_named_constraints - core = _process_array_objects(core, data, inf_model) + if group_repeated_constraint_patterns + core = _group_and_add_constraints(core, data, inf_model) # TODO: can eventually replace `_add_constraints` if it works well end core = _add_constraints(core, data, inf_model) core = _add_derivative_approximations(core, data, inf_model) @@ -862,21 +967,21 @@ end inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; [backend = nothing, - process_finite_indices_of_named_constraints = false] # experimental + group_repeated_constraint_patterns = false] # experimental )::ExaModels.ExaCore Create `ExaModels.ExaCore` from `inf_model` using the provided `ExaMappingData` to store the variable and constraint mappings. Optionally, try to aggregate common algebraic constraint structures over a named constraint array (e.g., `@constraint(model, con[i=1:N], ...)`) -by setting `process_finite_indices_of_named_constraints = true`. This is an +by setting `group_repeated_constraint_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, - process_finite_indices_of_named_constraints = false + group_repeated_constraint_patterns = false ) # TODO add support for other float types once InfiniteOpt does minimize = JuMP.objective_sense(inf_model) == _MOI.MIN_SENSE @@ -884,7 +989,7 @@ function ExaModels.ExaCore( return build_exa_core!( core, data, - inf_model; process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints + inf_model; group_repeated_constraint_patterns = group_repeated_constraint_patterns ) end @@ -893,39 +998,39 @@ end inf_model::InfiniteOpt.InfiniteModel, [data::ExaMappingData]; [backend = nothing, - process_finite_indices_of_named_constraints = false] # experimental + group_repeated_constraint_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. Optionally, try to aggregate common algebraic constraint structures over a named constraint array (e.g., `@constraint(model, con[i=1:N], ...)`) -by setting `process_finite_indices_of_named_constraints = true`. This is an +by setting `group_repeated_constraint_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, - process_finite_indices_of_named_constraints = false + group_repeated_constraint_patterns = false ) core = ExaModels.ExaCore( inf_model, data; backend = backend, - process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints + group_repeated_constraint_patterns = group_repeated_constraint_patterns ) return ExaModels.ExaModel(core) end function ExaModels.ExaModel( inf_model::InfiniteOpt.InfiniteModel; backend = nothing, - process_finite_indices_of_named_constraints = false + group_repeated_constraint_patterns = false ) return ExaModels.ExaModel( inf_model, ExaMappingData(); backend = backend, - process_finite_indices_of_named_constraints = process_finite_indices_of_named_constraints + group_repeated_constraint_patterns = group_repeated_constraint_patterns ) end From f4306f92141e9e705a365a7c2086e9de2f65ff0f Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 09:02:57 -0400 Subject: [PATCH 10/27] add support for restricted variables --- src/grouped_constraints.jl | 42 +++++++++++++++++++++++++------------- src/infiniteopt_backend.jl | 3 +++ src/transform.jl | 26 ++++++++++++++++++----- 3 files changed, 52 insertions(+), 19 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index c2c1bd0..6738098 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -4,12 +4,12 @@ const _VariableTypeHashingInt = Dict( InfiniteOpt.ParameterFunctionIndex => -3, InfiniteOpt.InfiniteVariableIndex => -4, InfiniteOpt.DerivativeIndex => -4, + InfiniteOpt.SemiInfiniteVariableIndex => -4, + InfiniteOpt.PointVariableIndex => -4, InfiniteOpt.FiniteVariableIndex => -5, - InfiniteOpt.SemiInfiniteVariableIndex => -6, - InfiniteOpt.PointVariableIndex => -7, - InfiniteOpt.IndependentParameterIndex => -8, - InfiniteOpt.DependentParameterIndex => -9, - InfiniteOpt.MeasureIndex => -10, + InfiniteOpt.IndependentParameterIndex => -6, + InfiniteOpt.DependentParameterIndex => -7, + InfiniteOpt.MeasureIndex => -8, ) ## Extract the following from an expression: @@ -78,17 +78,32 @@ function _group_info_msg(group, msg) return end +# TODO: Make work for parameters, semi-infinite variables, and point variables # 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) +function _get_grouped_idx( + em_var::Union{ExaModels.Var, ExaModels.Par}, + grouped_var::Union{ExaModels.Variable, ExaModels.Parameter} + ) idx = em_var.i - grouped_var.offset - @assert idx in grouped_var.size[end] + @assert idx in grouped_var.size[end] && length(grouped_var.size) == 1 return idx end -function _get_grouped_idx(em_var::ExaModels.Variable, grouped_var::ExaModels.Variable) +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 idx in 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 a candidate group of constraint, seek to merge together and add as a single constraint pattern to `core` function _process_candidate_constraint_group( @@ -101,15 +116,15 @@ function _process_candidate_constraint_group( ) # determine which vrefs and consts change across the array vrefs1 = vref_lists[1] - grouped_var_idxs = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] + is_grouped_var = [any(l -> l[i] != vrefs1[i], vref_lists) for i in eachindex(vrefs1)] consts1 = const_lists[1] - grouped_data_idxs = [any(l -> l[i] != consts1[i], const_lists) for i in eachindex(consts1)] + 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(crefs)] group_var_idx = 1 for (i, vref) in enumerate(vrefs1) - if grouped_var_idxs[i] + if is_grouped_var[i] if !haskey(data.var_to_grouped_var, vref) _group_info_msg(crefs, "Failure: $(vref) is not a grouped variable/parameter which prevents adding") return false, core @@ -125,8 +140,7 @@ function _process_candidate_constraint_group( _group_info_msg(crefs, "Failure: Unable to correctly set up variable grouping for $(infvar) which prevents adding") return false, core end - em_var = data[infvar] - var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(em_var, src_var)) + var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(infvar, data)) end group_var_idx += 1 else @@ -138,7 +152,7 @@ function _process_candidate_constraint_group( const_itr = Any[(;) for _ in 1:length(crefs)] grouped_const_idx = 1 for (i, c) in enumerate(consts1) - if grouped_data_idxs[i] + 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(crefs) diff --git a/src/infiniteopt_backend.jl b/src/infiniteopt_backend.jl index 8be4919..9b90fe0 100644 --- a/src/infiniteopt_backend.jl +++ b/src/infiniteopt_backend.jl @@ -31,6 +31,8 @@ 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}} @@ -54,6 +56,7 @@ struct ExaMappingData Vector{Any} } }(), + Dict{InfiniteOpt.GeneralVariableRef, Tuple}(), Dict{InfiniteOpt.GeneralVariableRef, Union{ExaModels.Variable, ExaModels.Parameter}}(), ) end diff --git a/src/transform.jl b/src/transform.jl index 4d708aa..3f57e0c 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -295,6 +295,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] = mapped_var + end return data.semivar_info[vref] = (mapped_var, indexing) end @@ -352,7 +355,13 @@ 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...] + mapped_var = data.infvar_mappings[ivref] + pt = mapped_var[idxs...] + if haskey(data.var_to_grouped_var, ivref) + data.var_to_grouped_var[vref] = mapped_var + end + data.point_indicies[vref] = idxs + return data.finvar_mappings[vref] = pt end # Add all the point variables from an InfiniteModel to a ExaCore @@ -364,7 +373,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) @@ -400,6 +408,16 @@ function _index_params( 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, @@ -436,9 +454,7 @@ 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( From 46702089b2bc7cadb25d2a6d21af6f22aa10b2e6 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 09:14:55 -0400 Subject: [PATCH 11/27] bug fix --- src/grouped_constraints.jl | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index 6738098..b0a22b5 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -78,12 +78,8 @@ function _group_info_msg(group, msg) return end -# TODO: Make work for parameters, semi-infinite variables, and point variables # Get the grouped index of a variable based on its direct exaified variable reference -function _get_grouped_idx( - em_var::Union{ExaModels.Var, ExaModels.Par}, - grouped_var::Union{ExaModels.Variable, ExaModels.Parameter} - ) +function _get_grouped_idx(em_var::ExaModels.Var, grouped_var::ExaModels.Variable) idx = em_var.i - grouped_var.offset @assert idx in grouped_var.size[end] && length(grouped_var.size) == 1 return idx From 7d223f54831d430b93721643230d188ed9f60248 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 15:16:44 -0400 Subject: [PATCH 12/27] add support for constant over collocation --- src/transform.jl | 42 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index 3f57e0c..a5bf09b 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -722,7 +722,8 @@ end function _add_collocation_restrictions( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + group_constraints::Bool = false ) for (pidx, vidxs) in inf_model.piecewise_vars # gather the basic information @@ -739,19 +740,48 @@ 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...)]) + if group_constraints + finite_itr = [(; :group_var => _get_grouped_idx(vref, data)) for vref in vrefs] + itr = vec([merge(i...) for i in Iterators.product(itrs..., finite_itr)]) + else + itr = vec([merge(i...) for i in Iterators.product(itrs...)]) + end + # 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...] - core, _ = ExaModels.add_con(core, em_expr, itr) + if group_constraints + idx_pars1 = (idx_pars1..., data_src[:grouped_var]) + idx_pars2 = (idx_pars2..., data_src[:grouped_var]) + end + # create the ExaModel expression tree and add the constraint + if group_constraints + 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) + else + for vref in vrefs + ivar = data.infvar_mappings[vref] + em_expr = ivar[idx_pars1...] - ivar[idx_pars2...] + core, _ = ExaModels.add_con(core, em_expr, itr) + end + end + end end return core From a8777cb170e19f2f98a8383f0440d64e4939b659 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 15:18:27 -0400 Subject: [PATCH 13/27] minor updates --- src/transform.jl | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index a5bf09b..fde65ba 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -718,7 +718,6 @@ function _add_derivative_approximations( end # Add the constraints needed for piecewise constant control variables -# TODO: group together all the piecewise variables that share the same infinite parameter(s) function _add_collocation_restrictions( core::ExaModels.ExaCore, data::ExaMappingData, @@ -988,8 +987,8 @@ function build_exa_core!( inf_model, create_parameter_groups = group_repeated_constraint_patterns ) - _add_semi_infinite_variables(core, data, inf_model) # TODO: add support for grouping semi-infinite variables - _add_point_variables(core, data, inf_model) # TODO: add support for grouping point variables + _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 @@ -998,7 +997,12 @@ function build_exa_core!( end core = _add_constraints(core, data, inf_model) core = _add_derivative_approximations(core, data, inf_model) - core = _add_collocation_restrictions(core, data, inf_model) + core = _add_collocation_restrictions( + core, + data, + inf_model, + group_constraints = group_repeated_constraint_patterns + ) # add the objective if there is one expr = JuMP.objective_function(inf_model) sense = JuMP.objective_sense(inf_model) From 241d6368de2a8b264f5c28c822929c98f4610f79 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 15:28:57 -0400 Subject: [PATCH 14/27] bug fix --- src/transform.jl | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index fde65ba..5e475d3 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -667,7 +667,8 @@ end function _add_derivative_approximations( core::ExaModels.ExaCore, data::ExaMappingData, - inf_model::InfiniteOpt.InfiniteModel + inf_model::InfiniteOpt.InfiniteModel; + group_constraints::Bool = false ) for dref in InfiniteOpt.all_derivatives(inf_model) # gather the derivative information @@ -755,19 +756,16 @@ function _add_collocation_restrictions( 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) if group_constraints - finite_itr = [(; :group_var => _get_grouped_idx(vref, data)) for vref in vrefs] + finite_itr = [(; :grouped_var => _get_grouped_idx(vref, data)) for vref in vrefs] itr = vec([merge(i...) for i in Iterators.product(itrs..., finite_itr)]) else itr = vec([merge(i...) for i in Iterators.product(itrs...)]) end # 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) - if group_constraints - idx_pars1 = (idx_pars1..., data_src[:grouped_var]) - idx_pars2 = (idx_pars2..., data_src[:grouped_var]) - end + alias_tuple = group_constraints ? (aliases..., :grouped_var) : Tuple(aliases) + 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 if group_constraints grouped_var = data.var_to_grouped_var[vrefs[1]] From eeb04dd2ee2a1d08bd628a837f7e3727a86cdee6 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 16:05:27 -0400 Subject: [PATCH 15/27] support grouped derivative constraints --- src/transform.jl | 91 +++++++++++++++++++++++++++++++++++------------- 1 file changed, 66 insertions(+), 25 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index 5e475d3..f5f5240 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -620,6 +620,7 @@ end # Make dispatch type to pass the data needed by `make_reduced_expr` struct _DerivReductionBackendInfo <: InfiniteOpt.AbstractTransformationBackend data::ExaMappingData + grouped_constraints::Bool end # Extend make_reduced_expr to create an ExaModel expression @@ -646,7 +647,12 @@ function InfiniteOpt.make_reduced_expr( data_src[i] end end for i in inds) - return ivar[idx_pars...] + if dispatch_data.grouped_constraints + grouped_var = data.var_to_grouped_var[vref] + return grouped_var[idx_pars..., data_src[:grouped_var]] + else + return ivar[idx_pars...] + end else # either an infinite variable or a derivative variable group_idxs = InfiniteOpt.parameter_group_int_indices(vref) idx_pars = (begin @@ -657,27 +663,42 @@ function InfiniteOpt.make_reduced_expr( data_src[g_alias] end end for i in group_idxs) - return data.infvar_mappings[vref][idx_pars...] + if dispatch_data.grouped_constraints + grouped_var = data.var_to_grouped_var[vref] + return grouped_var[idx_pars..., data_src[:grouped_var]] + else + return data.infvar_mappings[vref][idx_pars...] + end end return end # Add the approximation equations for each derivative variable -# TODO: group together all the derivatives that share derivative method, order, and infinite parameters function _add_derivative_approximations( core::ExaModels.ExaCore, data::ExaMappingData, inf_model::InfiniteOpt.InfiniteModel; group_constraints::Bool = false ) + # group all the derivatives of the same order, method, and infinite parameter dependencies + signature_to_derivs = Dict{ + Tuple{InfiniteOpt.GeneralVariableRef, Int, Vector{Int}}, + 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)) + signature_to_derivs[pref, order, group_idxs] = InfiniteOpt.GeneralVariableRef[] + end + push!(signature_to_derivs[pref, order, group_idxs], dref) + end + # iterate over each group of derivatives and add the approximation equations + for ((pref, order, group_idxs), drefs) 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] @@ -693,27 +714,42 @@ function _add_derivative_approximations( # 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 + itrs = [g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] + if group_constraints + push!(itrs, [(; :grouped_var => _get_grouped_idx(dref, data)) for dref in drefs]) end - # make the ExaModel expression tree + 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() - em_expr = InfiniteOpt.make_indexed_derivative_expr( - dref, - vref, - pref, - order, - data_src[data.group_alias[pref_group]], - supps, - _DerivReductionBackendInfo(data), - method, - (data_src[a] for a in aliases)... + if grouped_constraints + # TODO: make this handle the grouped indexing + em_expr = InfiniteOpt.make_indexed_derivative_expr( + drefs[1], + vref, + pref, + order, + data_src[data.group_alias[pref_group]], + supps, + _DerivReductionBackendInfo(data, true), + method, + (data_src[a] for a in aliases)... ) - # add the constraint - core, _ = ExaModels.add_con(core, em_expr, itr) + else + for dref in drefs + em_expr = InfiniteOpt.make_indexed_derivative_expr( + dref, + vref, + pref, + order, + data_src[data.group_alias[pref_group]], + supps, + _DerivReductionBackendInfo(data, false), + method, + (data_src[a] for a in aliases)... + ) + core, _ = ExaModels.add_con(core, em_expr, itr) + end + end end return core end @@ -994,7 +1030,12 @@ function build_exa_core!( core = _group_and_add_constraints(core, data, inf_model) # TODO: can eventually replace `_add_constraints` if it works well end core = _add_constraints(core, data, inf_model) - core = _add_derivative_approximations(core, data, inf_model) + core = _add_derivative_approximations( + core, + data, + inf_model, + group_constraints = group_repeated_constraint_patterns + ) core = _add_collocation_restrictions( core, data, From 6231c2f51ec71d860116c82c6f268e027fb68814 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 16:09:30 -0400 Subject: [PATCH 16/27] minor fixes --- src/transform.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index f5f5240..bd9c313 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -620,7 +620,7 @@ end # Make dispatch type to pass the data needed by `make_reduced_expr` struct _DerivReductionBackendInfo <: InfiniteOpt.AbstractTransformationBackend data::ExaMappingData - grouped_constraints::Bool + group_constraints::Bool end # Extend make_reduced_expr to create an ExaModel expression @@ -647,7 +647,7 @@ function InfiniteOpt.make_reduced_expr( data_src[i] end end for i in inds) - if dispatch_data.grouped_constraints + if dispatch_data.group_constraints grouped_var = data.var_to_grouped_var[vref] return grouped_var[idx_pars..., data_src[:grouped_var]] else @@ -663,7 +663,7 @@ function InfiniteOpt.make_reduced_expr( data_src[g_alias] end end for i in group_idxs) - if dispatch_data.grouped_constraints + if dispatch_data.group_constraints grouped_var = data.var_to_grouped_var[vref] return grouped_var[idx_pars..., data_src[:grouped_var]] else @@ -721,8 +721,7 @@ function _add_derivative_approximations( 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() - if grouped_constraints - # TODO: make this handle the grouped indexing + if group_constraints em_expr = InfiniteOpt.make_indexed_derivative_expr( drefs[1], vref, From c4354061124612c2723716d294a60085fa38a365 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 16:13:04 -0400 Subject: [PATCH 17/27] bug fix --- src/transform.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index bd9c313..d6ba4ca 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -710,7 +710,7 @@ 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...)] @@ -724,12 +724,12 @@ function _add_derivative_approximations( if group_constraints em_expr = InfiniteOpt.make_indexed_derivative_expr( drefs[1], - vref, - pref, - order, - data_src[data.group_alias[pref_group]], - supps, - _DerivReductionBackendInfo(data, true), + vref, + pref, + order, + data_src[data.group_alias[pref_group]], + supps, + _DerivReductionBackendInfo(data, group_constraints), method, (data_src[a] for a in aliases)... ) @@ -742,7 +742,7 @@ function _add_derivative_approximations( order, data_src[data.group_alias[pref_group]], supps, - _DerivReductionBackendInfo(data, false), + _DerivReductionBackendInfo(data, group_constraints), method, (data_src[a] for a in aliases)... ) From b7fe851774761ede58da82f8e2ed128a3257b0bf Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 16:24:22 -0400 Subject: [PATCH 18/27] bug fix --- src/transform.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/transform.jl b/src/transform.jl index d6ba4ca..243f77e 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -682,7 +682,7 @@ function _add_derivative_approximations( ) # group all the derivatives of the same order, method, and infinite parameter dependencies signature_to_derivs = Dict{ - Tuple{InfiniteOpt.GeneralVariableRef, Int, Vector{Int}}, + Tuple{InfiniteOpt.GeneralVariableRef, Int, Vector{Int}, DataType}, Vector{InfiniteOpt.GeneralVariableRef} }() for dref in InfiniteOpt.all_derivatives(inf_model) @@ -690,13 +690,13 @@ function _add_derivative_approximations( pref = InfiniteOpt.operator_parameter(dref) order = InfiniteOpt.derivative_order(dref) group_idxs = InfiniteOpt.parameter_group_int_indices(vref) - if !haskey(signature_to_derivs, (pref, order, group_idxs)) - signature_to_derivs[pref, order, group_idxs] = InfiniteOpt.GeneralVariableRef[] + if !haskey(signature_to_derivs, (pref, order, group_idxs, vref.index_type)) + signature_to_derivs[pref, order, group_idxs, vref.index_type] = InfiniteOpt.GeneralVariableRef[] end - push!(signature_to_derivs[pref, order, group_idxs], dref) + push!(signature_to_derivs[pref, order, group_idxs, vref.index_type], dref) end # iterate over each group of derivatives and add the approximation equations - for ((pref, order, group_idxs), drefs) in signature_to_derivs + for ((pref, order, group_idxs, _), drefs) in signature_to_derivs # gather basic info method = InfiniteOpt.derivative_method(drefs[1]) pref_group = InfiniteOpt.parameter_group_int_index(pref) @@ -724,7 +724,7 @@ function _add_derivative_approximations( if group_constraints em_expr = InfiniteOpt.make_indexed_derivative_expr( drefs[1], - vref, + InfiniteOpt.derivative_argument(drefs[1]), pref, order, data_src[data.group_alias[pref_group]], @@ -737,7 +737,7 @@ function _add_derivative_approximations( for dref in drefs em_expr = InfiniteOpt.make_indexed_derivative_expr( dref, - vref, + InfiniteOpt.derivative_argument(dref), pref, order, data_src[data.group_alias[pref_group]], From fdd748848307705b1bca892b3e2676df9ea9b0ad Mon Sep 17 00:00:00 2001 From: pulsipher Date: Wed, 19 Aug 2026 22:54:06 -0400 Subject: [PATCH 19/27] bug fixes --- src/grouped_constraints.jl | 6 +++--- src/transform.jl | 32 ++++++++++++++++---------------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index b0a22b5..5728242 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -81,7 +81,7 @@ 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 idx in grouped_var.size[end] && length(grouped_var.size) == 1 + @assert 1 <= idx <= grouped_var.size[end] && length(grouped_var.size) == 1 return idx end function _get_grouped_idx( @@ -89,7 +89,7 @@ function _get_grouped_idx( grouped_var::Union{ExaModels.Variable, ExaModels.Parameter} ) idx = (em_var.offset - grouped_var.offset) ÷ em_var.length + 1 - @assert idx in grouped_var.size[end] + @assert 1 <= idx <= grouped_var.size[end] return idx end function _get_grouped_idx(vref::InfiniteOpt.GeneralVariableRef, data::ExaMappingData) @@ -217,7 +217,7 @@ function _group_and_add_constraints( if length(crefs) < 2 continue end - success = _process_candidate_constraint_group(core, data, crefs, hash_to_patterns[h]...) + success, core = _process_candidate_constraint_group(core, data, crefs, hash_to_patterns[h]...) if success _group_info_msg(crefs, "Successfully added") end diff --git a/src/transform.jl b/src/transform.jl index 243f77e..51fc2d5 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -111,8 +111,8 @@ function _add_finite_variables( create_variable_groups::Bool = false ) vrefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteVariable) - if create_variable_groups - ex_var = ExaModels.Variable((1:length(vrefs),), length(vrefs), core.nvar, :x, nothing) + if create_variable_groups && length(vrefs) > 1 + ex_var = ExaModels.Variable((length(vrefs),), length(vrefs), core.nvar, :x, nothing) end for vref in vrefs info = InfiniteOpt.core_object(vref).info # JuMP.VariableInfo @@ -135,8 +135,8 @@ function _add_finite_parameters( create_parameter_groups::Bool = false ) prefs = JuMP.all_variables(inf_model, InfiniteOpt.FiniteParameter) - if create_parameter_groups - ex_par = ExaModels.Parameter((1:length(prefs),), length(prefs), core.npar, nothing) + if create_parameter_groups && length(prefs) > 1 + ex_par = ExaModels.Parameter((length(prefs),), length(prefs), core.npar, nothing) end for pref in prefs param_val = InfiniteOpt.parameter_value(pref) @@ -191,17 +191,17 @@ function _add_infinite_variables( if create_variable_groups && length(data.base_itrs) > 1 for group_idx in group_types group_vrefs = vref_groups[group_idx] - v1 = first(group_vrefs) + v1 = data[first(group_vrefs)] num_vars = v1.length * length(group_vrefs) - ex_var = ExaModels.Variable((v1.size..., 1:length(group_vrefs)), num_vars, v1.offset, :x, nothing) + ex_var = ExaModels.Variable((v1.size..., length(group_vrefs)), num_vars, v1.offset, :x, nothing) for vref in group_vrefs data.var_to_grouped_var[vref] = ex_var end end elseif create_variable_groups - v1 = first(vrefs) + v1 = data[first(vrefs)] num_vars = v1.length * length(vrefs) - ex_var = ExaModels.Variable((v1.size..., 1:length(vrefs)), num_vars, v1.offset, :x, nothing) + ex_var = ExaModels.Variable((v1.size..., length(vrefs)), num_vars, v1.offset, :x, nothing) for vref in vrefs data.var_to_grouped_var[vref] = ex_var end @@ -217,6 +217,7 @@ function _add_parameter_functions( create_parameter_groups::Bool = false ) pfrefs = InfiniteOpt.all_parameter_functions(inf_model) + length(pfrefs) == 0 && return core # sort by infinite parameter groups (required to capture repeated finite patterns) if create_parameter_groups && length(data.base_itrs) > 1 all_group_idxs = map(v -> InfiniteOpt.parameter_group_int_indices(v), pfrefs) @@ -250,17 +251,17 @@ function _add_parameter_functions( if create_parameter_groups && length(data.base_itrs) > 1 for group_idx in group_types group_pfrefs = pfref_groups[group_idx] - pf1 = first(group_pfrefs) + pf1 = data[first(group_pfrefs)] num_pars = pf1.length * length(group_pfrefs) - ex_var = ExaModels.Parameter((pf1.size..., 1:length(group_pfrefs)), num_pars, pf1.offset, nothing) + ex_var = ExaModels.Parameter((pf1.size..., length(group_pfrefs)), num_pars, pf1.offset, nothing) for pfref in group_pfrefs data.var_to_grouped_var[pfref] = ex_var end end elseif create_parameter_groups - pf1 = first(pfrefs) + pf1 = data[first(pfrefs)] num_pars = pf1.length * length(pfrefs) - ex_par = ExaModels.Parameter((pf1.size..., 1:length(pfrefs)), num_pars, pf1.offset, nothing) + ex_par = ExaModels.Parameter((pf1.size..., length(pfrefs)), num_pars, pf1.offset, nothing) for pfref in pfrefs data.var_to_grouped_var[pfref] = ex_par end @@ -355,10 +356,9 @@ 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)) - mapped_var = data.infvar_mappings[ivref] - pt = mapped_var[idxs...] + pt = data.infvar_mappings[ivref][idxs...] if haskey(data.var_to_grouped_var, ivref) - data.var_to_grouped_var[vref] = mapped_var + data.var_to_grouped_var[vref] = data.var_to_grouped_var[ivref] end data.point_indicies[vref] = idxs return data.finvar_mappings[vref] = pt @@ -714,7 +714,7 @@ function _add_derivative_approximations( # 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...)] - itrs = [g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] + itrs = Any[g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] if group_constraints push!(itrs, [(; :grouped_var => _get_grouped_idx(dref, data)) for dref in drefs]) end From 3a989298e2ce4140b882e64a4aabc93a76a1d0b4 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Thu, 20 Aug 2026 08:50:19 -0400 Subject: [PATCH 20/27] bug fix --- src/transform.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transform.jl b/src/transform.jl index 51fc2d5..1e0f1c3 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -733,6 +733,7 @@ function _add_derivative_approximations( method, (data_src[a] for a in aliases)... ) + core, _ = ExaModels.add_con(core, em_expr, itr) else for dref in drefs em_expr = InfiniteOpt.make_indexed_derivative_expr( From 7130421c24f72e5ad8524e466335e5e53b7ebc23 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Thu, 20 Aug 2026 22:36:27 -0400 Subject: [PATCH 21/27] bug fix --- src/grouped_constraints.jl | 4 ++-- src/transform.jl | 35 ++++++++++++++++++++++------------- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index 5728242..cb0555f 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -126,7 +126,7 @@ function _process_candidate_constraint_group( return false, core end base_idxs = _index_params(vref, data) - itr_alias = Symbol("grouped_var$group_var_idx") + itr_alias = Symbol("grouped_vidx$group_var_idx") var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) src_var = data.var_to_grouped_var[vref] exafied_vrefs[i] = src_var[var_idxs...] @@ -181,7 +181,7 @@ function _process_candidate_constraint_group( 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, (1:inf_len,), nothing) + data.constraint_mappings[cref] = ExaModels.Constraint(con.f, sliced_itr, offset, (inf_len,), nothing) end return true, core end diff --git a/src/transform.jl b/src/transform.jl index 1e0f1c3..327ca48 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -621,6 +621,7 @@ end struct _DerivReductionBackendInfo <: InfiniteOpt.AbstractTransformationBackend data::ExaMappingData group_constraints::Bool + alias_map::Union{Nothing, Dict{InfiniteOpt.GeneralVariableRef, Symbol}} end # Extend make_reduced_expr to create an ExaModel expression @@ -649,7 +650,7 @@ function InfiniteOpt.make_reduced_expr( end for i in inds) if dispatch_data.group_constraints grouped_var = data.var_to_grouped_var[vref] - return grouped_var[idx_pars..., data_src[:grouped_var]] + return grouped_var[idx_pars..., data_src[dispatch_data.alias_map[vref]]] else return ivar[idx_pars...] end @@ -665,7 +666,7 @@ function InfiniteOpt.make_reduced_expr( end for i in group_idxs) if dispatch_data.group_constraints grouped_var = data.var_to_grouped_var[vref] - return grouped_var[idx_pars..., data_src[:grouped_var]] + return grouped_var[idx_pars..., data_src[dispatch_data.alias_map[vref]]] else return data.infvar_mappings[vref][idx_pars...] end @@ -683,7 +684,7 @@ function _add_derivative_approximations( # group all the derivatives of the same order, method, and infinite parameter dependencies signature_to_derivs = Dict{ Tuple{InfiniteOpt.GeneralVariableRef, Int, Vector{Int}, DataType}, - Vector{InfiniteOpt.GeneralVariableRef} + Tuple{Vector{InfiniteOpt.GeneralVariableRef}, Vector{InfiniteOpt.GeneralVariableRef}} }() for dref in InfiniteOpt.all_derivatives(inf_model) vref = InfiniteOpt.derivative_argument(dref) @@ -691,12 +692,14 @@ function _add_derivative_approximations( order = InfiniteOpt.derivative_order(dref) 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[] + 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], dref) + 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) in signature_to_derivs + 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) @@ -716,34 +719,40 @@ function _add_derivative_approximations( pref_itr = [(; srt_itr[i]..., zip(aliases, args)...) for (i, args...) in zip(idxs, arg_itrs...)] itrs = Any[g == pref_group ? pref_itr : data.base_itrs[g] for g in group_idxs] if group_constraints - push!(itrs, [(; :grouped_var => _get_grouped_idx(dref, data)) for dref in drefs]) + push!( + itrs, + [(; :grouped_didx => _get_grouped_idx(dref, data), + :grouped_vidx => _get_grouped_idx(vref, data)) + for (dref, vref) in zip(drefs, vrefs) + ]) end 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() if group_constraints + alias_map = Dict(drefs[1] => :grouped_didx, vrefs[1] => :grouped_vidx) em_expr = InfiniteOpt.make_indexed_derivative_expr( drefs[1], - InfiniteOpt.derivative_argument(drefs[1]), + vrefs[1], pref, order, data_src[data.group_alias[pref_group]], supps, - _DerivReductionBackendInfo(data, group_constraints), + _DerivReductionBackendInfo(data, group_constraints, alias_map), method, (data_src[a] for a in aliases)... ) core, _ = ExaModels.add_con(core, em_expr, itr) else - for dref in drefs + for (dref, vref) in zip(drefs, vrefs) em_expr = InfiniteOpt.make_indexed_derivative_expr( dref, - InfiniteOpt.derivative_argument(dref), + vref, pref, order, data_src[data.group_alias[pref_group]], supps, - _DerivReductionBackendInfo(data, group_constraints), + _DerivReductionBackendInfo(data, group_constraints, nothing), method, (data_src[a] for a in aliases)... ) @@ -799,7 +808,7 @@ function _add_collocation_restrictions( end # prepare the variable indices data_src = ExaModels.DataSource() - alias_tuple = group_constraints ? (aliases..., :grouped_var) : Tuple(aliases) + alias_tuple = group_constraints ? (aliases..., :grouped_vidx) : Tuple(aliases) 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 From 689752dfa8da83fa338dc5eb4351e95f713a3a62 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 15:13:58 -0400 Subject: [PATCH 22/27] bug fix --- src/grouped_constraints.jl | 40 +++++++++++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index cb0555f..8fa709a 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -4,12 +4,12 @@ const _VariableTypeHashingInt = Dict( InfiniteOpt.ParameterFunctionIndex => -3, InfiniteOpt.InfiniteVariableIndex => -4, InfiniteOpt.DerivativeIndex => -4, - InfiniteOpt.SemiInfiniteVariableIndex => -4, - InfiniteOpt.PointVariableIndex => -4, - InfiniteOpt.FiniteVariableIndex => -5, - InfiniteOpt.IndependentParameterIndex => -6, - InfiniteOpt.DependentParameterIndex => -7, - InfiniteOpt.MeasureIndex => -8, + InfiniteOpt.SemiInfiniteVariableIndex => -5, + InfiniteOpt.PointVariableIndex => -6, + InfiniteOpt.FiniteVariableIndex => -7, + InfiniteOpt.IndependentParameterIndex => -8, + InfiniteOpt.DependentParameterIndex => -9, + InfiniteOpt.MeasureIndex => -10, ) ## Extract the following from an expression: @@ -23,7 +23,12 @@ 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(_VariableTypeHashingInt[v.index_type], h), push!(refs, v), consts + if v.index_type in (InfiniteOpt.SemiInfiniteVariableIndex, InfiniteOpt.PointVariableIndex) + group_idxs = InfiniteOpt.parameter_group_int_indices(InfiniteOpt.infinite_variable_ref(v)) + else + group_idxs = InfiniteOpt.parameter_group_int_indices(v) + end + return hash((_VariableTypeHashingInt[v.index_type], group_idxs), h), push!(refs, v), consts end function _encode_expr( expr::Union{JuMP.GenericAffExpr{C, V}, JuMP.GenericQuadExpr{C, V}}, @@ -119,15 +124,28 @@ function _process_candidate_constraint_group( exafied_vrefs = Vector{Any}(undef, length(vrefs1)) var_itr = Any[(;) for _ in 1:length(crefs)] group_var_idx = 1 + restricted_idx = 1 for (i, vref) in enumerate(vrefs1) if is_grouped_var[i] if !haskey(data.var_to_grouped_var, vref) _group_info_msg(crefs, "Failure: $(vref) is not a grouped variable/parameter which prevents adding") return false, core end - base_idxs = _index_params(vref, data) + base_idxs = collect(_index_params(vref, data)) itr_alias = Symbol("grouped_vidx$group_var_idx") - var_idxs = (base_idxs..., ExaModels.DataSource()[itr_alias]) + 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(crefs) @@ -137,6 +155,10 @@ function _process_candidate_constraint_group( return false, core end var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(infvar, data)) + if !isempty(alias_map) + ridxs = _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 From b0e09094fc297539aff68793ea235b1b6bf4c4a5 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 15:20:40 -0400 Subject: [PATCH 23/27] bug_fix --- src/transform.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transform.jl b/src/transform.jl index 327ca48..c7a8b81 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -801,7 +801,7 @@ function _add_collocation_restrictions( 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) if group_constraints - finite_itr = [(; :grouped_var => _get_grouped_idx(vref, data)) for vref in vrefs] + 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)]) else itr = vec([merge(i...) for i in Iterators.product(itrs...)]) From 655e3ec6eabbe9741898a8b7b9074e84cd633fed Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 15:26:52 -0400 Subject: [PATCH 24/27] update --- src/grouped_constraints.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index 8fa709a..e31a655 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -155,7 +155,7 @@ function _process_candidate_constraint_group( return false, core end var_itr[j] = (; var_itr[j]..., itr_alias => _get_grouped_idx(infvar, data)) - if !isempty(alias_map) + if !isempty(alias_map) # add in restricted variables indices if they exist ridxs = _index_params(infvar, data) var_itr[j] = merge(var_itr[j], NamedTuple(alias => ridxs[k] for (k, alias) in alias_map)) end From 749b2f841419a6da86bd9731363a1d4d8e67fdf2 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 16:04:59 -0400 Subject: [PATCH 25/27] add tests --- test/solve.jl | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/test/solve.jl b/test/solve.jl index 20e6b59..b089972 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,31 @@ 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)) + @objective(m, Min, ∫(sum((x[i] - 1)^2 for i in 1:3), t)) + @constraint(m, c1[i in 1:3], ∂(x[i], t) == u^(i+2)) + @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) + @test set_transformation_backend(m, ExaTranscriptionBackend(IpoptSolver)) isa Nothing + @test set_silent(m) isa Nothing + @test optimize!(m, group_repeated_constraint_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)) + end + @test all(isapprox.(uval, value(u), atol = tol)) +end + @testset "User-Defined Operators" begin model = InfiniteModel(Ipopt.Optimizer) set_silent(model) From 724fb56bec5352d713db34091f9ae88df7345cc7 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 16:54:13 -0400 Subject: [PATCH 26/27] typo fix --- src/transform.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transform.jl b/src/transform.jl index 8449851..bdd3ad9 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -1082,7 +1082,7 @@ function ExaModels.ExaCore( inf_model::InfiniteOpt.InfiniteModel, data::ExaMappingData; backend = nothing, - concrete_core::Bool = false + concrete_core::Bool = false, group_repeated_constraint_patterns = false ) # TODO add support for other float types once InfiniteOpt does From 807c615c2f9f1321477a6a09a0094a3ea0fe9178 Mon Sep 17 00:00:00 2001 From: pulsipher Date: Fri, 21 Aug 2026 22:46:31 -0400 Subject: [PATCH 27/27] semi-infinite bug fix --- src/grouped_constraints.jl | 4 ++-- src/transform.jl | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/grouped_constraints.jl b/src/grouped_constraints.jl index e31a655..81d28d8 100644 --- a/src/grouped_constraints.jl +++ b/src/grouped_constraints.jl @@ -131,7 +131,7 @@ function _process_candidate_constraint_group( _group_info_msg(crefs, "Failure: $(vref) is not a grouped variable/parameter which prevents adding") return false, core end - base_idxs = collect(_index_params(vref, data)) + base_idxs = Tuple(_index_params(vref, data)) itr_alias = Symbol("grouped_vidx$group_var_idx") data_src = ExaModels.DataSource() alias_map = Dict{Int, Symbol}() @@ -156,7 +156,7 @@ function _process_candidate_constraint_group( end 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 = _index_params(infvar, data) + 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 diff --git a/src/transform.jl b/src/transform.jl index bdd3ad9..8f80a48 100644 --- a/src/transform.jl +++ b/src/transform.jl @@ -297,7 +297,7 @@ function _process_semi_infinite_var(vref, data) mapped_var = data.infvar_mappings[ivref] end if haskey(data.var_to_grouped_var, ivref) - data.var_to_grouped_var[vref] = mapped_var + data.var_to_grouped_var[vref] = data.var_to_grouped_var[ivref] end return data.semivar_info[vref] = (mapped_var, indexing) end