AugLag.jl

AugLag(; inner, kwargs...)

Augmented Lagrangian optimizer that solves constrained Optimization.jl problems by repeatedly calling an inner optimizer on the augmented objective.

Keyword Arguments

• inner: Inner optimizer used for the augmented subproblems.
• τ: Constraint violation contraction threshold.
• γ: Penalty growth factor.
• ϵ, ϵ_primal, ϵ_dual: Convergence tolerances.
• ρ_init: Initial penalty parameter. If nothing, a problem-scaled value is chosen.
• inner_kwargs: Keyword arguments forwarded to the inner optimizer.

classify_constraints(lcons, ucons)

Partition constraint indices given the SciMLBase contract lcons[i] ≤ c_i(θ) ≤ ucons[i]:

• eq_inds: rows where lcons[i] == ucons[i] (equality).
• ineq_upper_inds: rows with lcons[i] != ucons[i] and finite ucons[i], contributing the penalty for c_i ≤ ucons[i].
• ineq_lower_inds: rows with lcons[i] != ucons[i] and finite lcons[i], contributing the penalty for lcons[i] ≤ c_i.

A two-sided inequality (both bounds finite, lcons[i] < ucons[i]) appears in bothineq_upper_inds and ineq_lower_inds and gets a separate Lagrange multiplier per side.

generate_auglag(cache, eq_inds, ineq_upper_inds, ineq_lower_inds,
                λ, μ_upper, μ_lower, ρ_ref)

Build the augmented-Lagrangian subproblem function as an OptimizationFunction with analytical value, gradient, and fg! derived from the user's loss and constraints in cache.f.

The augmented Lagrangian is

L(θ; λ, μ_u, μ_l, ρ) = f(θ)
    + Σᵢ∈eq    [ λᵢ (cᵢ - lᵢ) + (ρ/2)(cᵢ - lᵢ)² ]
    + (1/(2ρ)) Σᵢ∈up max(0, μ_uᵢ + ρ (cᵢ - uᵢ))²
    + (1/(2ρ)) Σᵢ∈lo max(0, μ_lᵢ + ρ (lᵢ - cᵢ))²

so its gradient (used closed-form, not by AD'ing through L) is

∇L = ∇f
    + Σᵢ∈eq    (λᵢ + ρ (cᵢ - lᵢ)) ∇cᵢ
    + Σᵢ∈up    max(0, μ_uᵢ + ρ (cᵢ - uᵢ)) ∇cᵢ
    - Σᵢ∈lo    max(0, μ_lᵢ + ρ (lᵢ - cᵢ)) ∇cᵢ.

λ, μ_upper, μ_lower are mutable vectors and ρ_ref is a Ref{<:Real} (or any zero-arg-getindex container). The closures dereference them at call time, so the outer AugLag loop can update multipliers and the penalty in place between inner solves without rebuilding the function.

cons_tmp and J are preallocated once with element type eltype(cache.u0). This is safe because the analytical gradient does not AD through this function — cache.f.grad and cache.f.cons_j (which the inner solver ultimately calls) handle their own AD internally.

Constraints and the data-iterator p

Constraints are treated as non-stochastic / batch-independent: the user's cons!(res, θ, p) is always invoked with the fullcache.p, regardless of which inner-solve batch is currently being processed. For the typical non-DataLoader case cache.p is just the user's p. For a DataLoaderp, cons! receives the iterator itself, and the user's constraint body may pull the underlying full data from it (e.g. via p.data for an MLUtils.DataLoader) — but the constraint is still a deterministic function of θ and the full data, never of a single batch.

Consistent with this, the constraint Jacobian cache.f.cons_j is invoked without p and uses the p that was closed over at AD-preparation time (the first batch for a data iterator). Since the constraint is by contract batch-independent, that closed-over p is irrelevant to the Jacobian's value.