diff --git a/docs/src/examples/Thermal_Generation_Dispatch_Example.jl b/docs/src/examples/Thermal_Generation_Dispatch_Example.jl index 1fe273aa3..b5846c0d5 100644 --- a/docs/src/examples/Thermal_Generation_Dispatch_Example.jl +++ b/docs/src/examples/Thermal_Generation_Dispatch_Example.jl @@ -158,7 +158,9 @@ Plots.plot( d, data_results[1, :, 1:(I+1)]; title = "Generation by Demand", - label = ["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"], + label = [ + "Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit" + ], xlabel = "Demand [unit]", ylabel = "Generation [unit]", ) @@ -168,7 +170,9 @@ Plots.plot( d, data_results[1, :, (I+2):(2*(I+1))]; title = "Sensitivity of Generation by Demand", - label = ["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"], + label = [ + "T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity" + ], xlabel = "Demand [unit]", ylabel = "Sensitivity [-]", ) @@ -179,7 +183,9 @@ Plots.plot( d, data_results[2, :, 1:(I+1)]; title = "Generation by Demand", - label = ["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"], + label = [ + "Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit" + ], xlabel = "Demand [unit]", ylabel = "Generation [unit]", ) @@ -189,7 +195,9 @@ Plots.plot( d, data_results[2, :, (I+2):(2*(I+1))]; title = "Sensitivity of Generation by Demand", - label = ["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"], + label = [ + "T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity" + ], xlabel = "Demand [unit]", ylabel = "Sensitivity [-]", ) diff --git a/docs/src/examples/Thermal_Generation_Dispatch_Example_new.jl b/docs/src/examples/Thermal_Generation_Dispatch_Example_new.jl index 391f3f2f3..0872c2bd9 100644 --- a/docs/src/examples/Thermal_Generation_Dispatch_Example_new.jl +++ b/docs/src/examples/Thermal_Generation_Dispatch_Example_new.jl @@ -154,7 +154,9 @@ Plots.plot( d, data_results[1, :, 1:(I+1)]; title = "Generation by Demand", - label = ["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"], + label = [ + "Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit" + ], xlabel = "Demand [unit]", ylabel = "Generation [unit]", ) @@ -164,7 +166,9 @@ Plots.plot( d, data_results[1, :, (I+2):(2*(I+1))]; title = "Sensitivity of Generation by Demand", - label = ["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"], + label = [ + "T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity" + ], xlabel = "Demand [unit]", ylabel = "Sensitivity [-]", ) @@ -175,7 +179,9 @@ Plots.plot( d, data_results[2, :, 1:(I+1)]; title = "Generation by Demand", - label = ["Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit"], + label = [ + "Thermal Generation 1" "Thermal Generation 2" "Thermal Generation 3" "Generation Deficit" + ], xlabel = "Demand [unit]", ylabel = "Generation [unit]", ) @@ -185,7 +191,9 @@ Plots.plot( d, data_results[2, :, (I+2):(2*(I+1))]; title = "Sensitivity of Generation by Demand", - label = ["T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity"], + label = [ + "T. Gen. 1 Sensitivity" "T. Gen. 2 Sensitivity" "T. Gen. 3 Sensitivity" "Gen. Deficit Sensitivity" + ], xlabel = "Demand [unit]", ylabel = "Sensitivity [-]", ) diff --git a/docs/src/examples/autotuning-ridge.jl b/docs/src/examples/autotuning-ridge.jl index 17619aef1..28cb6661f 100644 --- a/docs/src/examples/autotuning-ridge.jl +++ b/docs/src/examples/autotuning-ridge.jl @@ -86,10 +86,7 @@ for α in αs ŷ_test = X_test * ŵ ŷ_train = X_train * ŵ push!(mse_test, LinearAlgebra.norm(ŷ_test - y_test)^2 / (2 * Ntest * D)) - push!( - mse_train, - LinearAlgebra.norm(ŷ_train - y_train)^2 / (2 * Ntrain * D), - ) + push!(mse_train, LinearAlgebra.norm(ŷ_train - y_train)^2 / (2 * Ntrain * D)) end # Visualize the Mean Score Error metric diff --git a/docs/src/examples/autotuning-ridge_new.jl b/docs/src/examples/autotuning-ridge_new.jl index b7fa24a76..87f05ca95 100644 --- a/docs/src/examples/autotuning-ridge_new.jl +++ b/docs/src/examples/autotuning-ridge_new.jl @@ -93,10 +93,7 @@ for α in αs ŷ_test = X_test * ŵ ŷ_train = X_train * ŵ push!(mse_test, LinearAlgebra.norm(ŷ_test - y_test)^2 / (2 * Ntest * D)) - push!( - mse_train, - LinearAlgebra.norm(ŷ_train - y_train)^2 / (2 * Ntrain * D), - ) + push!(mse_train, LinearAlgebra.norm(ŷ_train - y_train)^2 / (2 * Ntrain * D)) end # Visualize the Mean Score Error metric diff --git a/docs/src/examples/sensitivity-analysis-ridge.jl b/docs/src/examples/sensitivity-analysis-ridge.jl index a4d3c0465..5604d0b69 100644 --- a/docs/src/examples/sensitivity-analysis-ridge.jl +++ b/docs/src/examples/sensitivity-analysis-ridge.jl @@ -140,13 +140,7 @@ p = Plots.scatter( label = "", ) mi, ma = minimum(X), maximum(X) -Plots.plot!( - p, - [mi, ma], - [mi * ŵ + b̂, ma * ŵ + b̂]; - color = :blue, - label = "", -) +Plots.plot!(p, [mi, ma], [mi * ŵ + b̂, ma * ŵ + b̂]; color = :blue, label = "") Plots.title!("Regression slope sensitivity with respect to x") # @@ -159,13 +153,7 @@ p = Plots.scatter( label = "", ) mi, ma = minimum(X), maximum(X) -Plots.plot!( - p, - [mi, ma], - [mi * ŵ + b̂, ma * ŵ + b̂]; - color = :blue, - label = "", -) +Plots.plot!(p, [mi, ma], [mi * ŵ + b̂, ma * ŵ + b̂]; color = :blue, label = "") Plots.title!("Regression slope sensitivity with respect to y") # Note the points with less central `x` values induce a greater y sensitivity of the slope. diff --git a/docs/src/examples/sensitivity-analysis-ridge_new.jl b/docs/src/examples/sensitivity-analysis-ridge_new.jl index aa41f80a2..59d98fecb 100644 --- a/docs/src/examples/sensitivity-analysis-ridge_new.jl +++ b/docs/src/examples/sensitivity-analysis-ridge_new.jl @@ -140,13 +140,7 @@ p = Plots.scatter( label = "", ) mi, ma = minimum(X), maximum(X) -Plots.plot!( - p, - [mi, ma], - [mi * ŵ + b̂, ma * ŵ + b̂]; - color = :blue, - label = "", -) +Plots.plot!(p, [mi, ma], [mi * ŵ + b̂, ma * ŵ + b̂]; color = :blue, label = "") Plots.title!("Regression slope sensitivity with respect to x") # @@ -159,13 +153,7 @@ p = Plots.scatter( label = "", ) mi, ma = minimum(X), maximum(X) -Plots.plot!( - p, - [mi, ma], - [mi * ŵ + b̂, ma * ŵ + b̂]; - color = :blue, - label = "", -) +Plots.plot!(p, [mi, ma], [mi * ŵ + b̂, ma * ŵ + b̂]; color = :blue, label = "") Plots.title!("Regression slope sensitivity with respect to y") # Note the points with less central `x` values induce a greater y sensitivity of the slope. diff --git a/src/NonLinearProgram/NonLinearProgram.jl b/src/NonLinearProgram/NonLinearProgram.jl index 3df0ecaec..011f261be 100644 --- a/src/NonLinearProgram/NonLinearProgram.jl +++ b/src/NonLinearProgram/NonLinearProgram.jl @@ -564,13 +564,6 @@ function DiffOpt.reverse_differentiate!(model::Model; tol = 1e-6) cache = _cache_evaluator!(model) form = model.model - # Compute Jacobian - Δs, df_dp = _compute_sensitivity(model; tol = tol) - Δp = if !iszero(model.input_cache.dobj) - df_dp'model.input_cache.dobj - else - zeros(length(cache.params)) - end num_primal = length(cache.primal_vars) # Fetch primal sensitivities Δx = zeros(num_primal) @@ -602,11 +595,31 @@ function DiffOpt.reverse_differentiate!(model::Model; tol = 1e-6) Δdual[num_constraints+num_low+i] = model.input_cache.dy[idx] end end - # Extract Parameter sensitivities - Δw = zeros(size(Δs, 1)) + # Assemble the seed over the KKT unknowns [x; s; λ; ν_L; ν_U]. The dimension is the + # size of the KKT Jacobian `M` (see `_build_sensitivity_matrices`): primal variables, + # one slack per inequality, one dual per constraint, and lower/upper bound duals for + # both primal variables and slacks. + num_leq = length(cache.leq_locations) + num_geq = length(cache.geq_locations) + kkt_dim = + num_primal + + 2 * (num_leq + num_geq) + + num_constraints + + num_low + + num_up + Δw = zeros(kkt_dim) Δw[1:num_primal] = Δx Δw[cache.index_duals] = Δdual - Δp += Δs' * Δw + # A reverse request is a vector-Jacobian product: one adjoint solve against the KKT + # factorization instead of materializing the dense (kkt_dim × P) sensitivity + # `∂s = -K⁻¹N` and contracting it. The objective seed folds into the same + # right-hand side. See `_compute_sensitivity_adjoint`. + Δp = _compute_sensitivity_adjoint( + model, + Δw, + model.input_cache.dobj; + tol = tol, + ) Δp_dict = Dict{MOI.ConstraintIndex,Float64}( form.var2ci[var_idx] => Δp[form.var2param[var_idx].value] diff --git a/src/NonLinearProgram/nlp_utilities.jl b/src/NonLinearProgram/nlp_utilities.jl index 0aa97c099..b536e6b8b 100644 --- a/src/NonLinearProgram/nlp_utilities.jl +++ b/src/NonLinearProgram/nlp_utilities.jl @@ -372,16 +372,19 @@ function _build_sensitivity_matrices( # [V_U 0 0 (X_U - X)] # ] M = [ - W A' I_L I_U; - A spzeros(num_cons, num_cons) spzeros(num_cons, num_low) spzeros(num_cons, num_up); - V_L spzeros(num_low, num_cons) X_lb spzeros(num_low, num_up); - V_U spzeros(num_up, num_cons) spzeros(num_up, num_low) X_ub; + W A' I_L I_U + A spzeros(num_cons, num_cons) spzeros(num_cons, num_low) spzeros( + num_cons, + num_up, + ) + V_L spzeros(num_low, num_cons) X_lb spzeros(num_low, num_up) + V_U spzeros(num_up, num_cons) spzeros(num_up, num_low) X_ub ] # N matrix N = [ - ∇ₓₚL; - ∇ₚC; - spzeros(num_low + num_up, num_parms); + ∇ₓₚL + ∇ₚC + spzeros(num_low + num_up, num_parms) ] return M, N @@ -502,3 +505,85 @@ function _compute_sensitivity(model::Model; tol = 1e-6) df_dp = df_dx'∂s[1:num_vars, :] + df_dp_direct' # ∇ₚfᵒ(x,p) = ∇ₓf(x,p) * ∇ₚxᵒ(p) + ∇ₚf(x,p) * 𝐈ₚ return ∂s, df_dp end + +""" + _compute_sensitivity_adjoint(model::Model, Δw::AbstractVector, dobj::Real; tol=1e-6) + +Reverse-mode (vector-Jacobian product) counterpart of [`_compute_sensitivity`](@ref): return +`Δp = ∂sᵀ Δw + df_dpᵀ dobj` without materializing the dense `∂s = -K⁻¹N` (a `size(K, 1) × P` +solve). Since `∂sᵀ Δw = -Nᵀ (K⁻ᵀ Δw)`, one adjoint solve on the same factorization suffices. + +Two details make this exactly the quantity the dense path contracts: + + * the JuMP-convention sign adjustments `_compute_sensitivity` applies to the dual *rows* of + `∂s` are applied to the corresponding *entries* of the seed instead (algebraically + identical); + * the objective seed folds into the same right-hand side: with + `df_dp = ∇ₓfᵀ ∂s[1:n, :] + ∇ₚfᵀ`, the first term rides along as `∇ₓf ⋅ dobj` added to the + primal entries of the seed, and the direct `∇ₚf ⋅ dobj` term is added after the solve. + +The inertia-correction / singular fallback matches the dense path: a failed factorization +(`K === nothing`) contributes a zero solve component. +""" +function _compute_sensitivity_adjoint( + model::Model, + Δw::AbstractVector, + dobj::Real; + tol = 1e-6, +) + # Solution and bounds — identical to `_compute_sensitivity` + X, + V_L, + X_L, + V_U, + X_U, + leq_locations, + geq_locations, + ineq_locations, + has_up, + has_low, + cons = _compute_solution_and_bounds(model; tol = tol) + M, N = _build_sensitivity_matrices( + model, + cons, + X, + V_L, + X_L, + V_U, + X_U, + leq_locations, + geq_locations, + ineq_locations, + has_up, + has_low, + ) + K = model.input_cache.factorization(M, model) + num_vars = _get_num_primal_vars(model) + num_cons = _get_num_constraints(model) + num_w = num_vars + length(ineq_locations) + num_lower = length(has_low) + _sense_multiplier = _sense_mult(model) + # Seed, with the sign adjustments the dense path applies to the rows of ∂s + w = collect(Float64, Δw) + # Duals + w[(num_w+1):(num_w+num_cons)] .*= -_sense_multiplier + # Dual bounds lower + w[(num_w+num_cons+1):(num_w+num_cons+num_lower)] .*= _sense_multiplier + # Dual bounds upper + w[(num_w+num_cons+num_lower+1):end] .*= -_sense_multiplier + grad = iszero(dobj) ? nothing : _compute_gradient(model) + if grad !== nothing + primal_idx = [i.value for i in model.cache.primal_vars] + w[1:num_vars] .+= grad[primal_idx] .* dobj # the ∂s-dependent part of df_dpᵀ dobj + end + Δp = if K === nothing + zeros(_get_num_params(model)) + else + -(N' * (K' \ w)) # ONE adjoint solve; UMFPACK solves Kᵀz = w on the same factorization + end + if grad !== nothing + params_idx = [i.value for i in model.cache.params] + Δp .+= grad[params_idx] .* dobj # the direct ∇ₚf ⋅ dobj term + end + return Δp +end diff --git a/test/nlp_program.jl b/test/nlp_program.jl index a9ec9e1fa..bed7e1f6a 100644 --- a/test/nlp_program.jl +++ b/test/nlp_program.jl @@ -1126,6 +1126,88 @@ function test_reverse_bounds_upper() @test isapprox(dp, 2.88888; atol = 1e-4) end +function test_reverse_adjoint_matches_forward_jvp() + # `reverse_differentiate!` computes Δp with a single adjoint solve + # (`_compute_sensitivity_adjoint`) instead of materializing the dense ∂s = -K⁻¹N and + # contracting. Cross-check reverse (VJP) against forward (JVP): for every parameter eᵢ, + # ⟨Δw, ∂x/∂pᵢ⟩ from `forward_differentiate!` must equal Δpᵢ from `reverse_differentiate!`, + # across parameter counts and with a simultaneous objective seed. + for P in (3, 11) + model = DiffOpt.nonlinear_diff_model(Ipopt.Optimizer) + set_silent(model) + @variable(model, p[i=1:P] in MOI.Parameter(1.0 + 0.1 * i)) + @variable(model, x[1:P]) + @constraint(model, [i = 1:P], x[i] * (1 + 0.1 * sin(p[i])) >= p[i]) + @constraint(model, sum(x) >= 0.6 * P) + @objective( + model, + Min, + sum((x[i] - p[i])^2 for i in 1:P) + 0.01 * sum(x[i]^4 for i in 1:P) + ) + optimize!(model) + @assert is_solved_and_feasible(model) + Δx_seed = [sin(0.7 * i) for i in 1:P] + Δf_seed = 0.3 + # Reverse: one adjoint solve, primal + objective seeds together + DiffOpt.empty_input_sensitivities!(model) + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], Δx_seed[i]) + end + MOI.set(model, DiffOpt.ReverseObjectiveValue(), Δf_seed) + MOI.set(model, DiffOpt.AllowObjectiveAndSolutionInput(), true) + DiffOpt.reverse_differentiate!(model) + Δp_reverse = [ + MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value for i in 1:P + ] + # Forward, one parameter at a time: Δpᵢ == ⟨Δx_seed, ∂x/∂pᵢ⟩ + Δf_seed ⋅ ∂f/∂pᵢ + for i in 1:P + DiffOpt.empty_input_sensitivities!(model) + DiffOpt.set_forward_parameter(model, p[i], 1.0) + DiffOpt.forward_differentiate!(model) + jvp = sum( + Δx_seed[j] * + MOI.get(model, DiffOpt.ForwardVariablePrimal(), x[j]) for + j in 1:P + ) + jvp += Δf_seed * MOI.get(model, DiffOpt.ForwardObjectiveValue()) + @test isapprox(Δp_reverse[i], jvp; rtol = 1e-7, atol = 1e-9) + end + end + return +end + +function test_reverse_adjoint_singular_factorization_fallback() + # Covers the `K === nothing` singular-factorization fallback in + # `_compute_sensitivity_adjoint` (nlp_utilities.jl) by injecting a factorization + # hook that returns `nothing`. With no objective seed, Δp must be exactly zeros. + P = 3 + model = DiffOpt.nonlinear_diff_model(Ipopt.Optimizer) + set_silent(model) + @variable(model, p[i=1:P] in MOI.Parameter(1.0 + 0.2 * i)) + @variable(model, x[1:P] >= 0) + @constraint(model, [i = 1:P], x[i] >= p[i]) + @constraint(model, sum(x) >= 0.5 * P) + @objective(model, Min, sum((x[i] - p[i])^2 for i in 1:P)) + optimize!(model) + @assert is_solved_and_feasible(model) + MOI.set( + model, + DiffOpt.NonLinearKKTJacobianFactorization(), + (M, m) -> nothing, + ) + DiffOpt.empty_input_sensitivities!(model) + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], 1.0) + end + DiffOpt.reverse_differentiate!(model) + Δp = [ + MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value + for i in 1:P + ] + @test all(iszero, Δp) + return +end + end # module TestNLPProgram.runtests()