[WIP] Support MOI.VectorNonlinearOracle

src/moi_nlp_model.jl

@@ -36,6 +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, noracle, oracle_lcon, oracle_ucon, nnzj_oracle, nnzh_oracle = parser_oracles(moimodel, index_map)
   oracles = Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache}[]
   counters = Counters()
   λ = zeros(Float64, nnln)  # Lagrange multipliers for hess_coord! and hprod! without y
@@ -47,12 +48,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 + noracle
+  lcon = vcat(lin_lcon, quad_lcon, nl_lcon, oracle_lcon)
+  ucon = vcat(lin_ucon, quad_ucon, nl_ucon, oracle_ucon)
+  nnzj = lincon.nnzj + quadcon.nnzj + nlcon.nnzj + nnzj_oracle
+  nnzh = obj.nnzh + quadcon.nnzh + nlcon.nnzh + nnzh_oracle
 
   meta = NLPModelMeta(
     nvar,
@@ -67,7 +68,7 @@ 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 + nnzj_oracle,
     minimize = MOI.get(moimodel, MOI.ObjectiveSense()) == MOI.MIN_SENSE,
     islp = (obj.type == "LINEAR") && (nnln == 0) && (quadcon.nquad == 0),
     name = name,
@@ -204,14 +205,24 @@ 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)
-    view(rows, ind_nnln) .= nlp.quadcon.nquad .+ nlp.nlcon.jac_rows
+    # structure of oracle Jacobians
+    for (k, (f, s)) in enumerate(nlp.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[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
+    # non-oracle nonlinear constraints
+    ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
+    view(rows, ind_nnln) .= nlp.quadcon.nquad .+ sum(s.set.output_dimension for (_, s) in nlp.oracles) .+ nlp.nlcon.jac_rows
     view(cols, ind_nnln) .= nlp.nlcon.jac_cols
   end
   return rows, cols
@@ -290,15 +301,15 @@ 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)
+    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):(offset + nnz_oracle)), s.x)
+      offset += nnz_oracle
+    end
+    ind_nnln = (offset + 1):(offset + nlp.nlcon.nnzj)
     MOI.eval_constraint_jacobian(nlp.eval, view(vals, ind_nnln), x)
   end
   return vals
@@ -337,14 +348,14 @@ 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
+      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):(offset + nnz_oracle)), s.x)
+        offset += nnz_oracle
+      end
       ind_nnln = (offset + 1):(nlp.meta.nnzj)
       MOI.eval_constraint_jacobian(nlp.eval, view(vals, ind_nnln), x)
     end
@@ -516,26 +527,28 @@ 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 !isempty(nlp.oracles)
+    for (f, s) in nlp.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
+  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
+  end
   return rows, cols
 end
 

src/utils.jl

@@ -84,6 +84,17 @@ mutable struct QuadraticConstraints
   nnzh::Int
 end
 
+"""
+    NonLinearStructure
+
+Structure containing Jacobian and Hessian structures of nonlinear constraints:
+- jac_rows: row indices of the Jacobian in Coordinate format (COO) format
+- jac_cols: column indices of the Jacobian in COO format
+- nnzj: number of non-zero entries in the Jacobian
+- hess_rows: row indices of the Hessian in COO format
+- hess_cols: column indices of the Hessian in COO format
+- nnzh: number of non-zero entries in the Hessian
+"""
 mutable struct NonLinearStructure
   jac_rows::Vector{Int}
   jac_cols::Vector{Int}
@@ -417,6 +428,10 @@ function parser_MOI(moimodel, index_map, nvar)
 
   contypes = MOI.get(moimodel, MOI.ListOfConstraintTypesPresent())
   for (F, S) in contypes
+    # Ignore VectorNonlinearOracle here, we'll parse it separately
+    if F == MOI.VectorOfVariables && S <: MOI.VectorNonlinearOracle{Float64}
+        continue
+    end
     (F == VNF) && error(
       "The function $F is not supported. Please use `.<=`, `.==`, and `.>=` in your constraints to ensure compatibility with ScalarNonlinearFunction.",
     )
@@ -538,6 +553,12 @@ end
     parser_NL(nlp_data; hessian)
 
 Parse nonlinear constraints of an `nlp_data`.
+
+Returns:
+- nnln: number of nonlinear constraints
+- nlcon: NonLinearStructure containing Jacobian and Hessian structures
+- nl_lcon: lower bounds of nonlinear constraints
+- nl_ucon: upper bounds of nonlinear constraints
 """
 function parser_NL(nlp_data; hessian::Bool = true)
   nnln = length(nlp_data.constraint_bounds)
@@ -560,6 +581,47 @@ function parser_NL(nlp_data; hessian::Bool = true)
   return nnln, nlcon, nl_lcon, nl_ucon
 end
 
+"""
+    parser_oracles(moimodel, index_map)
+
+Parse nonlinear oracles of a `MOI.ModelLike`.
+"""
+function parser_oracles(moimodel, index_map)
+    oracles = Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache}[]
+    l_oracle = Float64[]
+    u_oracle = Float64[]
+
+    # We know this pair exists from ListOfConstraintTypesPresent
+    for ci in MOI.get(
+        moimodel,
+        MOI.ListOfConstraintIndices{MOI.VectorOfVariables, MOI.VectorNonlinearOracle{Float64}}(),
+    )
+        f   = MOI.get(moimodel, MOI.ConstraintFunction(), ci)  # ::MOI.VectorOfVariables
+        set = MOI.get(moimodel, MOI.ConstraintSet(), ci)       # ::MOI.VectorNonlinearOracle{Float64}
+
+        cache = _VectorNonlinearOracleCache(set)
+        push!(oracles, (f, cache))
+
+        # Bounds: MOI.VectorNonlinearOracle stores them internally (l, u)
+        append!(l_oracle, set.l)
+        append!(u_oracle, set.u)
+    end
+
+    # Sizes: number of scalar constraints represented by all oracles
+    noracle = length(l_oracle)
+
+    # Sparsity:
+    nnzj_oracle = 0
+    nnzh_oracle = 0
+    for (_, cache) in oracles
+        nnzj_oracle += length(cache.set.jacobian_structure)
+        # there may or may not be Hessian info
+        nnzh_oracle += length(cache.set.hessian_lagrangian_structure)
+    end
+
+    return oracles, noracle, l_oracle, u_oracle, nnzj_oracle, nnzh_oracle
+end
+
 """
     parser_variables(model)