Skip to content

[WIP] Support MOI.VectorNonlinearOracle #3

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
amontoison merged 9 commits into from
Dec 16, 2025
Merged

[WIP] Support MOI.VectorNonlinearOracle #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
dd41b6e
skip support check oracle parser_MOI
andrewrosemberg Dec 12, 2025
97bec14
add parser oracles
andrewrosemberg Dec 12, 2025
167b38b
count number of elements considering oracles
andrewrosemberg Dec 12, 2025
6321fe7
update NLPModels.jac_nln_structure!
andrewrosemberg Dec 12, 2025
dbd5cf1
update NLPModels.jac_nln_coord!
andrewrosemberg Dec 12, 2025
ba11908
update NLPModels.jac_coord!
andrewrosemberg Dec 12, 2025
4ad64b1
update NLPModels.hess_structure!
andrewrosemberg Dec 12, 2025
b493185
rearrange order nnl and oracle
andrewrosemberg Dec 15, 2025
6908d38
fix bug
andrewrosemberg Dec 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
250 changes: 171 additions & 79 deletions src/moi_nlp_model.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ mutable struct MathOptNLPModel <: AbstractNLPModel{Float64, Vector{Float64}}
lincon::LinearConstraints
quadcon::QuadraticConstraints
nlcon::NonLinearStructure
oracles::Vector{Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache}}
oracles_data::OraclesData
λ::Vector{Float64}
hv::Vector{Float64}
obj::Objective
Expand Down Expand Up @@ -36,7 +36,7 @@ function nlp_model(moimodel::MOI.ModelLike; hessian::Bool = true, name::String =

nlp_data = _nlp_block(moimodel)
nnln, nlcon, nl_lcon, nl_ucon = parser_NL(nlp_data, hessian = hessian)
oracles = Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache}[]
oracles_data = parser_oracles(moimodel)
counters = Counters()
λ = zeros(Float64, nnln) # Lagrange multipliers for hess_coord! and hprod! without y
hv = zeros(Float64, nvar) # workspace for ghjvprod!
Expand All @@ -47,13 +47,12 @@ function nlp_model(moimodel::MOI.ModelLike; hessian::Bool = true, name::String =
obj = parser_objective_MOI(moimodel, nvar, index_map)
end

# Update ncon, lcon, ucon, nnzj and nnzh for the oracles
ncon = nlin + quadcon.nquad + nnln
lcon = vcat(lin_lcon, quad_lcon, nl_lcon)
ucon = vcat(lin_ucon, quad_ucon, nl_ucon)
nnzj = lincon.nnzj + quadcon.nnzj + nlcon.nnzj
nnzh = obj.nnzh + quadcon.nnzh + nlcon.nnzh

# Total counts
ncon = nlin + quadcon.nquad + nnln + oracles_data.noracle
lcon = vcat(lin_lcon, quad_lcon, nl_lcon, oracles_data.l_oracle)
ucon = vcat(lin_ucon, quad_ucon, nl_ucon, oracles_data.u_oracle)
nnzj = lincon.nnzj + quadcon.nnzj + nlcon.nnzj + oracles_data.nnzj_oracle
nnzh = obj.nnzh + quadcon.nnzh + nlcon.nnzh + oracles_data.nnzh_oracle
meta = NLPModelMeta(
nvar,
x0 = x0,
Expand All @@ -67,17 +66,17 @@ function nlp_model(moimodel::MOI.ModelLike; hessian::Bool = true, name::String =
nnzh = nnzh,
lin = collect(1:nlin),
lin_nnzj = lincon.nnzj,
nln_nnzj = quadcon.nnzj + nlcon.nnzj,
nln_nnzj = quadcon.nnzj + nlcon.nnzj + oracles_data.nnzj_oracle,
minimize = MOI.get(moimodel, MOI.ObjectiveSense()) == MOI.MIN_SENSE,
islp = (obj.type == "LINEAR") && (nnln == 0) && (quadcon.nquad == 0),
name = name,
jprod_available = isempty(oracles),
jtprod_available = isempty(oracles),
hprod_available = isempty(oracles) && hessian,
jprod_available = isempty(oracles_data.oracles),
jtprod_available = isempty(oracles_data.oracles),
hprod_available = isempty(oracles_data.oracles) && hessian,
hess_available = hessian,
)

return MathOptNLPModel(meta, nlp_data.evaluator, lincon, quadcon, nlcon, oracles, λ, hv, obj, counters),
return MathOptNLPModel(meta, nlp_data.evaluator, lincon, quadcon, nlcon, oracles_data, λ, hv, obj, counters),
index_map
end

Expand Down Expand Up @@ -128,8 +127,9 @@ function NLPModels.cons_nln!(nlp::MathOptNLPModel, x::AbstractVector, c::Abstrac
end
end
if nlp.meta.nnln > nlp.quadcon.nquad
offset = nlp.quadcon.nquad
for (f, s) in nlp.oracles
offset = nlp.quadcon.nquad + nlp.nlcon.nnln
MOI.eval_constraint(nlp.eval, view(c, (nlp.quadcon.nquad + 1):(offset)), x)
for (f, s) in nlp.oracles_data.oracles
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end
Expand All @@ -138,8 +138,6 @@ function NLPModels.cons_nln!(nlp::MathOptNLPModel, x::AbstractVector, c::Abstrac
s.set.eval_f(c_oracle, s.x)
offset += s.set.output_dimension
end
index_nnln = (offset + 1):(nlp.meta.nnln)
MOI.eval_constraint(nlp.eval, view(c, index_nnln), x)
end
return c
end
Expand All @@ -158,8 +156,9 @@ function NLPModels.cons!(nlp::MathOptNLPModel, x::AbstractVector, c::AbstractVec
end
end
if nlp.meta.nnln > nlp.quadcon.nquad
offset = nlp.meta.nlin + nlp.quadcon.nquad
for (f, s) in nlp.oracles
offset = nlp.meta.nlin + nlp.quadcon.nquad + nlp.nlcon.nnln
MOI.eval_constraint(nlp.eval, view(c, (nlp.meta.nlin + nlp.quadcon.nquad + 1):(offset)), x)
for (f, s) in nlp.oracles_data.oracles
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end
Expand All @@ -168,8 +167,6 @@ function NLPModels.cons!(nlp::MathOptNLPModel, x::AbstractVector, c::AbstractVec
s.set.eval_f(c_oracle, s.x)
offset += s.set.output_dimension
end
index_nnln = (offset + 1):(nlp.meta.ncon)
MOI.eval_constraint(nlp.eval, view(c, index_nnln), x)
end
end
return c
Expand Down Expand Up @@ -204,15 +201,26 @@ function NLPModels.jac_nln_structure!(
end
if nlp.meta.nnln > nlp.quadcon.nquad
offset = nlp.quadcon.nnzj
# for (f, s) in nlp.oracles
# for (i, j) in s.set.jacobian_structure
# ...
# offset += 1
# end
# end
ind_nnln = (offset + 1):(nlp.quadcon.nnzj + nlp.nlcon.nnzj)
# non-oracle nonlinear constraints
ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
view(rows, ind_nnln) .= nlp.quadcon.nquad .+ nlp.nlcon.jac_rows
view(cols, ind_nnln) .= nlp.nlcon.jac_cols
offset += nlp.nlcon.nnzj
# structure of oracle Jacobians
for (k, (f, s)) in enumerate(nlp.oracles_data.oracles)
# Shift row index by quadcon.nquad
# plus previous oracle outputs.
row_offset = nlp.quadcon.nquad
for i in 1:(k - 1)
row_offset += nlp.oracles_data.oracles[i][2].set.output_dimension
end

for (r, c) in s.set.jacobian_structure
offset += 1
rows[offset] = row_offset + r
cols[offset] = f.variables[c].value
end
end
end
return rows, cols
end
Expand Down Expand Up @@ -240,16 +248,22 @@ function NLPModels.jac_structure!(
end
end
if nlp.meta.nnln > nlp.quadcon.nquad
# non-oracle nonlinear constraints
offset = nlp.lincon.nnzj + nlp.quadcon.nnzj
# for (f, s) in nlp.oracles
# for (i, j) in s.set.jacobian_structure
# ...
# offset += 1
# end
# end
ind_nnln = (offset + 1):(nlp.meta.nnzj)
ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
view(rows, ind_nnln) .= nlp.meta.nlin .+ nlp.quadcon.nquad .+ nlp.nlcon.jac_rows
view(cols, ind_nnln) .= nlp.nlcon.jac_cols
offset += nlp.nlcon.nnzj
# structure of oracle Jacobians
row_offset = nlp.meta.nlin + nlp.quadcon.nquad + nlp.nlcon.nnln
for (f, s) in nlp.oracles_data.oracles
for (i, j) in s.set.jacobian_structure
offset += 1
rows[offset] = row_offset + i
cols[offset] = f.variables[j].value
end
row_offset += s.set.output_dimension
end
end
end
return rows, cols
Expand Down Expand Up @@ -290,16 +304,17 @@ function NLPModels.jac_nln_coord!(nlp::MathOptNLPModel, x::AbstractVector, vals:
end
if nlp.meta.nnln > nlp.quadcon.nquad
offset = nlp.quadcon.nnzj
# for (f, s) in nlp.oracles
# for i in 1:s.set.input_dimension
# s.x[i] = x[f.variables[i].value]
# end
# nnz_oracle = length(s.set.jacobian_structure)
# s.set.eval_jacobian(view(values, (offset + 1):(oracle + nnz_oracle)), s.x)
# offset += nnz_oracle
# end
ind_nnln = (offset + 1):(nlp.quadcon.nnzj + nlp.nlcon.nnzj)
ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
MOI.eval_constraint_jacobian(nlp.eval, view(vals, ind_nnln), x)
offset += nlp.nlcon.nnzj
for (f, s) in nlp.oracles_data.oracles
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end
nnz_oracle = length(s.set.jacobian_structure)
s.set.eval_jacobian(view(values, (offset + 1):(offset + nnz_oracle)), s.x)
offset += nnz_oracle
end
end
return vals
end
Expand Down Expand Up @@ -337,18 +352,20 @@ function NLPModels.jac_coord!(nlp::MathOptNLPModel, x::AbstractVector, vals::Abs
end
if nlp.meta.nnln > nlp.quadcon.nquad
offset = nlp.lincon.nnzj + nlp.quadcon.nnzj
# for (f, s) in nlp.oracles
# for i in 1:s.set.input_dimension
# s.x[i] = x[f.variables[i].value]
# end
# nnz_oracle = length(s.set.jacobian_structure)
# s.set.eval_jacobian(view(values, (offset + 1):(oracle + nnz_oracle)), s.x)
# offset += nnz_oracle
# end
ind_nnln = (offset + 1):(nlp.meta.nnzj)
ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
MOI.eval_constraint_jacobian(nlp.eval, view(vals, ind_nnln), x)
offset += nlp.nlcon.nnzj
for (f, s) in nlp.oracles_data.oracles
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end
nnz_oracle = length(s.set.jacobian_structure)
s.set.eval_jacobian(view(values, (offset + 1):(offset + nnz_oracle)), s.x)
offset += nnz_oracle
end
end
end
@assert offset == nlp.meta.nnzj
return vals
end

Expand Down Expand Up @@ -516,26 +533,29 @@ function NLPModels.hess_structure!(
view(rows, 1:(nlp.obj.nnzh)) .= nlp.obj.hessian.rows
view(cols, 1:(nlp.obj.nnzh)) .= nlp.obj.hessian.cols
end
if (nlp.obj.type == "NONLINEAR") || (nlp.meta.nnln > nlp.quadcon.nquad)
view(rows, (nlp.obj.nnzh + nlp.quadcon.nnzh + 1):(nlp.meta.nnzh)) .= nlp.nlcon.hess_rows
view(cols, (nlp.obj.nnzh + nlp.quadcon.nnzh + 1):(nlp.meta.nnzh)) .= nlp.nlcon.hess_cols
end
index = nlp.obj.nnzh
if nlp.quadcon.nquad > 0
index = nlp.obj.nnzh
for i = 1:(nlp.quadcon.nquad)
qcon = nlp.quadcon.constraints[i]
view(rows, (index + 1):(index + qcon.nnzh)) .= qcon.A.rows
view(cols, (index + 1):(index + qcon.nnzh)) .= qcon.A.cols
index += qcon.nnzh
end
end
# if !isempty(nlp.oracles)
# for (f, s) in nlp.oracles
# for (i, j) in s.set.hessian_lagrangian_structure
# ...
# end
# end
# end
if (nlp.obj.type == "NONLINEAR") || (nlp.meta.nnln > nlp.quadcon.nquad)
view(rows, (index + 1):(index + nlp.nlcon.nnzh)) .= nlp.nlcon.hess_rows
view(cols, (index + 1):(index + nlp.nlcon.nnzh)) .= nlp.nlcon.hess_cols
index += nlp.nlcon.nnzh
end
if !isempty(nlp.oracles_data.oracles)
for (f, s) in nlp.oracles_data.oracles
for (i_local, j_local) in s.set.hessian_lagrangian_structure
index += 1
rows[index] = f.variables[i_local].value
cols[index] = f.variables[j_local].value
end
end
end
return rows, cols
end

Expand All @@ -547,9 +567,16 @@ function NLPModels.hess_coord!(
obj_weight::Float64 = 1.0,
)
increment!(nlp, :neval_hess)

# Running index over Hessian nonzeros we've filled so far.
index = 0

# Objective Hessian block
if nlp.obj.type == "QUADRATIC"
view(vals, 1:(nlp.obj.nnzh)) .= obj_weight .* nlp.obj.hessian.vals
index += nlp.obj.nnzh
end
# Nonlinear objective Hessian block
if (nlp.obj.type == "NONLINEAR") || (nlp.meta.nnln > nlp.quadcon.nquad)
λ = view(y, (nlp.meta.nlin + nlp.quadcon.nquad + 1):(nlp.meta.ncon))
MOI.eval_hessian_lagrangian(
Expand All @@ -560,23 +587,42 @@ function NLPModels.hess_coord!(
λ,
)
end
# Quadratic constraint Hessian blocks
if nlp.quadcon.nquad > 0
index = nlp.obj.nnzh
for i = 1:(nlp.quadcon.nquad)
qcon = nlp.quadcon.constraints[i]
view(vals, (index + 1):(index + qcon.nnzh)) .= y[nlp.meta.nlin + i] .* qcon.A.vals
index += qcon.nnzh
end
end
# if !isempty(nlp.oracles)
# for (f, s) in nlp.oracles
# for i in 1:s.set.input_dimension
# s.x[i] = x[f.variables[i].value]
# end
# H_nnz = length(s.set.hessian_lagrangian_structure)
# ...
# end
# end
# Oracle Hessian blocks
if !isempty(nlp.oracles_data.oracles)
λ_oracle_all = view(
y,
(nlp.meta.nlin + nlp.quadcon.nquad + nlp.nlcon.nnln + 1):
(nlp.meta.nlin + nlp.quadcon.nquad + nlp.nlcon.nnln + nlp.oracles_data.noracle),
)

λ_offset = 0
for (f, s) in nlp.oracles_data.oracles
# build local x for this oracle
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end

nout = s.set.output_dimension
μ = view(λ_oracle_all, (λ_offset + 1):(λ_offset + nout))

H_nnz = length(s.set.hessian_lagrangian_structure)
ind = (index + 1):(index + H_nnz)
s.set.eval_hessian_lagrangian(view(vals, ind), s.x, μ)

index += H_nnz
λ_offset += nout
end
end

@assert index == nlp.meta.nnzh
return vals
end

Expand Down Expand Up @@ -631,9 +677,55 @@ function NLPModels.jth_hess_coord!(
)
nlp.λ[j - nlp.meta.nlin - nlp.quadcon.nquad] = 0.0
end
# if !isempty(nlp.oracles)
# ...
# end
# check for oracle constraints
if !isempty(nlp.oracles_data.oracles)
n_nl = nlp.nlcon.nnln
noracle = nlp.oracles_data.noracle

# check if j is an oracle constraint index
first_oracle = nlp.meta.nlin + nlp.quadcon.nquad + n_nl + 1
last_oracle = first_oracle + noracle - 1

if first_oracle ≤ j ≤ last_oracle
# local oracle constraint index k ∈ 1: noracle
k = j - (nlp.meta.nlin + nlp.quadcon.nquad + n_nl)

# starting index of this oracle's Hessian block
index = nlp.obj.nnzh + nlp.quadcon.nnzh + nlp.nlcon.nnzh

# walk through oracles to find which one owns component k
offset_outputs = 0
for (f, s) in nlp.oracles_data.oracles
nout = s.set.output_dimension
H_nnz = length(s.set.hessian_lagrangian_structure)

if k > offset_outputs + nout
# skip this oracle's block
index += H_nnz
offset_outputs += nout
continue
end

# this oracle owns constraint k; local index ℓ
ℓ = k - offset_outputs

# build local x
for i in 1:s.set.input_dimension
s.x[i] = x[f.variables[i].value]
end

# local multipliers μ: one-hot at ℓ
μ = zeros(eltype(x), nout)
μ[ℓ] = 1.0

# fill this oracle's Hessian block
ind = (index + 1):(index + H_nnz)
s.set.eval_hessian_lagrangian(view(vals, ind), s.x, μ)

break
end
end
end
return vals
end

Expand Down
Loading
Loading