Support quadratic constraints in MathOptNLSModel

Author: amontoison

src/moi_nls_model.jl

@@ -9,6 +9,7 @@ mutable struct MathOptNLSModel <: AbstractNLSModel{Float64, Vector{Float64}}
   linequ::LinearEquations
   nlequ::NonLinearStructure
   lincon::LinearConstraints
+  quadcon::QuadraticConstraints
   nlcon::NonLinearStructure
   counters::NLSCounters
 end
@@ -39,11 +40,11 @@ function MathOptNLSModel(cmodel::JuMP.Model, F; hessian::Bool = true, name::Stri
   Fnnzj = linequ.nnzj + nlequ.nnzj
   Fnnzh = nlequ.nnzh
 
-  ncon = nlin + nnln
-  lcon = vcat(lin_lcon, nl_lcon)
-  ucon = vcat(lin_ucon, nl_ucon)
-  cnnzj = lincon.nnzj + nlcon.nnzj
-  cnnzh = lls.nnzh + nlcon.nnzh
+  ncon = nlin + quadcon.nquad + nnln
+  lcon = vcat(lin_lcon, quad_lcon, nl_lcon)
+  ucon = vcat(lin_ucon, quad_ucon, nl_ucon)
+  cnnzj = lincon.nnzj + quadcon.nnzj + nlcon.nnzj
+  cnnzh = lls.nnzh + quadcon.nnzh + nlcon.nnzh
 
   meta = NLPModelMeta(
     nvar,
@@ -58,7 +59,7 @@ function MathOptNLSModel(cmodel::JuMP.Model, F; hessian::Bool = true, name::Stri
     nnzh = cnnzh,
     lin = collect(1:nlin),
     lin_nnzj = lincon.nnzj,
-    nln_nnzj = nlcon.nnzj,
+    nln_nnzj = quadcon.nnzj + nlcon.nnzj,
     minimize = objective_sense(cmodel) == MOI.MIN_SENSE,
     islp = false,
     name = name,
@@ -73,6 +74,7 @@ function MathOptNLSModel(cmodel::JuMP.Model, F; hessian::Bool = true, name::Stri
     linequ,
     nlequ,
     lincon,
+    quadcon,
     nlcon,
     NLSCounters(),
   )
@@ -253,7 +255,13 @@ end
 
 function NLPModels.cons_nln!(nls::MathOptNLSModel, x::AbstractVector, c::AbstractVector)
   increment!(nls, :neval_cons_nln)
-  MOI.eval_constraint(nls.ceval, c, x)
+  for i = 1:(nls.quadcon.nquad)
+    qcon = nls.quadcon.constraints[i]
+    c[i] = 0.5 * coo_sym_dot(qcon.A.rows, qcon.A.cols, qcon.A.vals, x, x) + dot(qcon.b, x)
+  end
+  if nls.meta.nnln > nls.quadcon.nquad
+    MOI.eval_constraint(nls.ceval, view(c, (nls.quadcon.nquad + 1):(nls.meta.nnln)), x)
+  end
   return c
 end
 
@@ -272,8 +280,20 @@ function NLPModels.jac_nln_structure!(
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
 )
-  view(rows, 1:(nls.nlcon.nnzj)) .= nls.nlcon.jac_rows
-  view(cols, 1:(nls.nlcon.nnzj)) .= nls.nlcon.jac_cols
+  if nls.quadcon.nquad > 0
+    index = 0
+    for i = 1:(nls.quadcon.nquad)
+      qcon = nls.quadcon.constraints[i]
+      view(rows, index+1:index+qcon.nnzg) .= i
+      view(cols, index+1:index+qcon.nnzg) .= qcon.g
+      index += qcon.nnzg
+    end
+  end
+  if nls.meta.nnln > nls.quadcon.nquad
+    ind_nnln = (nls.quadcon.nnzj + 1):(nls.quadcon.nnzj + nls.nlcon.nnzj)
+    view(rows, ind_nnln) .= nls.nlcon.jac_rows
+    view(cols, ind_nnln) .= nls.nlcon.jac_cols
+  end
   return rows, cols
 end
 
@@ -285,7 +305,33 @@ end
 
 function NLPModels.jac_nln_coord!(nls::MathOptNLSModel, x::AbstractVector, vals::AbstractVector)
   increment!(nls, :neval_jac_nln)
-  MOI.eval_constraint_jacobian(nls.ceval, view(vals, 1:(nls.nlcon.nnzj)), x)
+  if nls.quadcon.nquad > 0
+    index = 0
+    view(vals, 1:nls.quadcon.nnzj) .= 0.0
+    for i = 1:(nls.quadcon.nquad)
+      qcon = nls.quadcon.constraints[i]
+      for (j, ind) in enumerate(qcon.b.nzind)
+        k = qcon.dg[ind]
+        vals[index+k] += qcon.b.nzval[j]
+      end
+      for j = 1:qcon.nnzh
+        row = qcon.A.rows[j]
+        col = qcon.A.cols[j]
+        val = qcon.A.vals[j]
+        k1 = qcon.dg[row]
+        vals[index+k1] += val * x[col]
+        if row != col
+          k2 = qcon.dg[col]
+          vals[index+k2] += val * x[row]
+        end
+      end
+      index += qcon.nnzg
+    end
+  end
+  if nls.meta.nnln > nls.quadcon.nquad
+    ind_nnln = (nls.quadcon.nnzj + 1):(nls.quadcon.nnzj + nls.nlcon.nnzj)
+    MOI.eval_constraint_jacobian(nls.ceval, view(vals, ind_nnln), x)
+  end
   return vals
 end
 
@@ -313,7 +359,18 @@ function NLPModels.jprod_nln!(
   Jv::AbstractVector,
 )
   increment!(nls, :neval_jprod_nln)
-  MOI.eval_constraint_jacobian_product(nls.ceval, Jv, x, v)
+  if nls.meta.nnln > nls.quadcon.nquad
+    ind_nnln = (nls.quadcon.nquad + 1):(nls.meta.nnln)
+    MOI.eval_constraint_jacobian_product(nls.ceval, view(Jv, ind_nnln), x, v)
+  end
+  (nls.meta.nnln == nls.quadcon.nquad) && (Jv .= 0.0)
+  if nls.quadcon.nquad > 0
+    for i = 1:(nls.quadcon.nquad)
+      # Jv[i] = (Aᵢ * x + bᵢ)ᵀ * v
+      qcon = nls.quadcon.constraints[i]
+      v[i] += coo_sym_dot(qcon.A.rows, qcon.A.cols, qcon.A.vals, x, v) + dot(qcon.b, v)
+    end
+  end
   return Jv
 end
 
@@ -341,7 +398,19 @@ function NLPModels.jtprod_nln!(
   Jtv::AbstractVector,
 )
   increment!(nls, :neval_jtprod_nln)
-  MOI.eval_constraint_jacobian_transpose_product(nls.ceval, Jtv, x, v)
+  if nls.meta.nnln > nls.quadcon.nquad
+    ind_nnln = (nls.quadcon.nquad + 1):(nls.meta.nnln)
+    MOI.eval_constraint_jacobian_transpose_product(nls.ceval, Jtv, x, view(v, ind_nnln))
+  end
+  (nls.meta.nnln == nls.quadcon.nquad) && (Jtv .= 0.0)
+  if nls.quadcon.nquad > 0
+    for i = 1:(nls.quadcon.nquad)
+      # Jtv += v[i] * (Aᵢ * x + bᵢ)
+      qcon = nls.quadcon.constraints[i]
+      coo_sym_add_mul!(rows, cols, vals, x, Jtv, v[i])
+      Jtv .+= v[i] .* qcon.b
+    end
+  end
   return Jtv
 end
 
@@ -354,9 +423,18 @@ function NLPModels.hess_structure!(
     view(rows, 1:(nls.lls.nnzh)) .= nls.lls.hessian.rows
     view(cols, 1:(nls.lls.nnzh)) .= nls.lls.hessian.cols
   end
-  if nls.nls_meta.nnln > 0
-    view(rows, (nls.lls.nnzh + 1):(nls.meta.nnzh)) .= nls.nlcon.hess_rows
-    view(cols, (nls.lls.nnzh + 1):(nls.meta.nnzh)) .= nls.nlcon.hess_cols
+  if (nls.nls_meta.nnln > 0) || (nlp.meta.nnln > nlp.quadcon.nquad)
+    view(rows, (nls.lls.nnzh + nlp.quadcon.nnzh + 1):(nls.meta.nnzh)) .= nls.nlcon.hess_rows
+    view(cols, (nls.lls.nnzh + nlp.quadcon.nnzh + 1):(nls.meta.nnzh)) .= nls.nlcon.hess_cols
+  end
+  if nls.quadcon.nquad > 0
+    index = nls.lls.nnzh
+    for i = 1:(nls.quadcon.nquad)
+      qcon = nls.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
   return rows, cols
 end
@@ -372,15 +450,24 @@ function NLPModels.hess_coord!(
   if nls.nls_meta.nlin > 0
     view(vals, 1:(nls.lls.nnzh)) .= obj_weight .* nls.lls.hessian.vals
   end
-  if (nls.nls_meta.nnln > 0) || (nls.meta.nnln > 0)
+  if (nls.nls_meta.nnln > 0) || (nls.meta.nnln > nls.quadcon.nquad)
+    λ = view(y, (nls.meta.nlin + nls.quadcon.nquad + 1):(nls.meta.ncon))
     MOI.eval_hessian_lagrangian(
       nls.ceval,
-      view(vals, (nls.lls.nnzh + 1):(nls.meta.nnzh)),
+      view(vals, (nls.lls.nnzh + nls.quadcon.nnzh + 1):(nls.meta.nnzh)),
       x,
       obj_weight,
-      view(y, nls.meta.nln),
+      λ
     )
   end
+  if nls.quadcon.nquad > 0
+    index = nls.lls.nnzh
+    for i = 1:(nls.quadcon.nquad)
+      qcon = nls.quadcon.constraints[i]
+      view(vals, (index + 1):(index + qcon.nnzh)) .= y[i] .* qcon.A.vals
+      index += qcon.nnzh
+    end
+  end
   return vals
 end
 
@@ -394,16 +481,18 @@ function NLPModels.hess_coord!(
   if nls.nls_meta.nlin > 0
     view(vals, 1:(nls.lls.nnzh)) .= obj_weight .* nls.lls.hessian.vals
   end
+  view(vals, (nls.lls.nnzh + 1):(nls.lls.nnzh + nls.quadcon.nnzh)) .= 0.0
   if nls.nls_meta.nnln > 0
+    λ = zeros(nlp.meta.nnln - nlp.quadcon.nquad)  # Should be stored in the structure MathOptNLSModel
     MOI.eval_hessian_lagrangian(
       nls.ceval,
-      view(vals, (nls.lls.nnzh + 1):(nls.meta.nnzh)),
+      view(vals, (nls.lls.nnzh + nls.quadcon.nnzh + 1):(nls.meta.nnzh)),
       x,
       obj_weight,
-      zeros(nls.meta.nnln),
+      λ,
     )
   else
-    view(vals, (nls.lls.nnzh + 1):(nls.meta.nnzh)) .= 0.0
+    view(vals, (nls.lls.nnzh + nls.quadcon.nnzh + 1):(nls.meta.nnzh)) .= 0.0
   end
   return vals
 end
@@ -417,11 +506,12 @@ function NLPModels.hprod!(
   obj_weight::Float64 = 1.0,
 )
   increment!(nls, :neval_hprod)
-  if (nls.nls_meta.nnln > 0) || (nls.meta.nnln > 0)
-    MOI.eval_hessian_lagrangian_product(nls.ceval, hv, x, v, obj_weight, view(y, nls.meta.nln))
+  if (nls.nls_meta.nnln > 0) || (nls.meta.nnln > nls.quadcon.nquad)
+    λ = view(y, (nls.meta.nlin + nls.quadcon.nquad + 1):(nls.meta.ncon))
+    MOI.eval_hessian_lagrangian_product(nls.ceval, hv, x, v, obj_weight, λ)
   end
+  (nls.nls_meta.nnln == 0) && (nls.meta.nnln == nls.quadcon.nquad) && (hv .= 0.0)
   if nls.nls_meta.nlin > 0
-    (nls.nls_meta.nnln == 0) && (nls.meta.nnln == 0) && (hv .= 0.0)
     coo_sym_add_mul!(
       nls.lls.hessian.rows,
       nls.lls.hessian.cols,
@@ -431,6 +521,12 @@ function NLPModels.hprod!(
       obj_weight,
     )
   end
+  if nls.quadcon.nquad > 0
+    for i = 1:(nls.quadcon.nquad)
+      qcon = nls.quadcon.constraints[i]
+      coo_sym_add_mul!(qcon.A.rows, qcon.A.cols, qcon.A.vals, v, hv, y[nls.meta.nlin + i])
+    end
+  end
   return hv
 end
 
@@ -443,7 +539,8 @@ function NLPModels.hprod!(
 )
   increment!(nls, :neval_hprod)
   if nls.nls_meta.nnln > 0
-    MOI.eval_hessian_lagrangian_product(nls.ceval, hv, x, v, obj_weight, zeros(nls.meta.nnln))
+    λ = zeros(nls.meta.nnln - nls.quadcon.nquad)  # Should be stored in the structure MathOptNLSModel
+    MOI.eval_hessian_lagrangian_product(nls.ceval, hv, x, v, obj_weight, λ)
   end
   if nls.nls_meta.nlin > 0
     (nls.nls_meta.nnln == 0) && (hv .= 0.0)