Add cons_h_weighted! to handle σ=0 in Lagrangian Hessian

When computing the Lagrangian Hessian with lag_h!, the case σ=0 requires
special handling since it reduces to just the weighted sum of constraint
Hessians (Σᵢ λᵢ∇²cᵢ) without the objective contribution.

Previously, this case would fail when cons_h was not explicitly requested
but lag_h was, because the constraint Hessian preparations were not created.

This commit:
- Always creates constraint Hessian preparations when either cons_h or lag_h is true
- Adds cons_h_weighted!(H, θ, λ) function to compute the weighted sum directly into H
- Updates lag_h! to use cons_h_weighted! when σ=0

This fixes the edge case in OptimizationMadNLP where the solver could hit
σ=0 during iterations, particularly with exact Hessian and sparse KKT systems.

Applies to both OptimizationDIExt and OptimizationZygoteExt.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: SebastianM-C

lib/OptimizationBase/ext/OptimizationZygoteExt.jl

@@ -186,20 +186,54 @@ function OptimizationBase.instantiate_function(
 
     conshess_sparsity = f.cons_hess_prototype
     conshess_colors = f.cons_hess_colorvec
-    if cons !== nothing && cons_h == true && f.cons_h === nothing
+
+    # Prepare constraint Hessian preparations if needed by lag_h or cons_h
+    if cons !== nothing && f.cons_h === nothing && (cons_h == true || lag_h == true)
         prep_cons_hess = [prepare_hessian(
                               cons_oop, soadtype, x, Constant(i), strict = Val(false))
                           for i in 1:num_cons]
+    else
+        prep_cons_hess = nothing
+    end
+
+    # Generate cons_h! functions
+    if cons !== nothing && f.cons_h === nothing && prep_cons_hess !== nothing
+        # Standard cons_h! that returns array of matrices
+        if cons_h == true
+            cons_h! = function (H, θ)
+                for i in 1:num_cons
+                    hessian!(cons_oop, H[i], prep_cons_hess[i], soadtype, θ, Constant(i))
+                end
+            end
+        else
+            cons_h! = nothing
+        end
+
+        # Weighted sum dispatch for cons_h! (always created if prep_cons_hess exists)
+        # This is used by lag_h! when σ=0
+        cons_h_weighted! = function (H::AbstractMatrix, θ, λ)
+            # Compute weighted sum: H = Σᵢ λᵢ∇²cᵢ
+            H .= zero(eltype(H))
+
+            # Create a single temporary matrix to reuse for all constraints
+            Hi = similar(H)
 
-        function cons_h!(H, θ)
             for i in 1:num_cons
-                hessian!(cons_oop, H[i], prep_cons_hess[i], soadtype, θ, Constant(i))
+                if λ[i] != zero(eltype(λ))
+                    # Compute constraint's Hessian into temporary matrix
+                    hessian!(cons_oop, Hi, prep_cons_hess[i], soadtype, θ, Constant(i))
+                    # Add weighted Hessian to result using in-place operation
+                    # H += λ[i] * Hi
+                    @. H += λ[i] * Hi
+                end
             end
         end
     elseif cons !== nothing && cons_h == true
         cons_h! = (res, θ) -> f.cons_h(res, θ, p)
+        cons_h_weighted! = nothing
     else
         cons_h! = nothing
+        cons_h_weighted! = nothing
     end
 
     lag_hess_prototype = f.lag_hess_prototype
@@ -212,8 +246,8 @@ function OptimizationBase.instantiate_function(
 
         function lag_h!(H::AbstractMatrix, θ, σ, λ)
             if σ == zero(eltype(θ))
-                cons_h!(H, θ)
-                H *= λ
+                # When σ=0, use the weighted sum function
+                cons_h_weighted!(H, θ, λ)
             else
                 hessian!(lagrangian, H, lag_extras, soadtype, θ,
                     Constant(σ), Constant(λ), Constant(p))
@@ -512,8 +546,8 @@ function OptimizationBase.instantiate_function(
 
         function lag_h!(H::AbstractMatrix, θ, σ, λ)
             if σ == zero(eltype(θ))
-                cons_h!(H, θ)
-                H *= λ
+                # When σ=0, use the weighted sum function
+                cons_h_weighted!(H, θ, λ)
             else
                 hessian!(lagrangian, H, lag_extras, soadtype, θ,
                     Constant(σ), Constant(λ), Constant(p))

lib/OptimizationBase/src/OptimizationDIExt.jl

@@ -177,20 +177,53 @@ function instantiate_function(
 
     conshess_sparsity = f.cons_hess_prototype
     conshess_colors = f.cons_hess_colorvec
-    # Generate cons_h! if explicitly requested OR if lag_h needs it
+
+    # Prepare constraint Hessian preparations if needed by lag_h or cons_h
     if f.cons !== nothing && f.cons_h === nothing && (cons_h == true || lag_h == true)
         prep_cons_hess = [prepare_hessian(cons_oop, soadtype, x, Constant(i))
                           for i in 1:num_cons]
+    else
+        prep_cons_hess = nothing
+    end
+
+    # Generate cons_h! functions
+    if f.cons !== nothing && f.cons_h === nothing && prep_cons_hess !== nothing
+        # Standard cons_h! that returns array of matrices
+        if cons_h == true
+            cons_h! = function (H, θ)
+                for i in 1:num_cons
+                    hessian!(cons_oop, H[i], prep_cons_hess[i], soadtype, θ, Constant(i))
+                end
+            end
+        else
+            cons_h! = nothing
+        end
+
+        # Weighted sum dispatch for cons_h! (always created if prep_cons_hess exists)
+        # This is used by lag_h! when σ=0
+        cons_h_weighted! = function (H::AbstractMatrix, θ, λ)
+            # Compute weighted sum: H = Σᵢ λᵢ∇²cᵢ
+            H .= zero(eltype(H))
+
+            # Create a single temporary matrix to reuse for all constraints
+            Hi = similar(H)
 
-        function cons_h!(H, θ)
             for i in 1:num_cons
-                hessian!(cons_oop, H[i], prep_cons_hess[i], soadtype, θ, Constant(i))
+                if λ[i] != zero(eltype(λ))
+                    # Compute constraint's Hessian into temporary matrix
+                    hessian!(cons_oop, Hi, prep_cons_hess[i], soadtype, θ, Constant(i))
+                    # Add weighted Hessian to result using in-place operation
+                    # H += λ[i] * Hi
+                    @. H += λ[i] * Hi
+                end
             end
         end
-    elseif (cons_h == true || lag_h == true) && f.cons !== nothing
+    elseif cons_h == true && f.cons !== nothing
         cons_h! = (res, θ) -> f.cons_h(res, θ, p)
+        cons_h_weighted! = nothing
     else
         cons_h! = nothing
+        cons_h_weighted! = nothing
     end
 
     lag_hess_prototype = f.lag_hess_prototype
@@ -203,8 +236,8 @@ function instantiate_function(
 
         function lag_h!(H::AbstractMatrix, θ, σ, λ)
             if σ == zero(eltype(θ))
-                cons_h!(H, θ)
-                H *= λ
+                # When σ=0, use the weighted sum function
+                cons_h_weighted!(H, θ, λ)
             else
                 hessian!(lagrangian, H, lag_prep, soadtype, θ,
                     Constant(σ), Constant(λ), Constant(p))