From 5c2328a2abd7986de7960c692216671f3a6ae16e Mon Sep 17 00:00:00 2001 From: zero Date: Sat, 4 Jul 2026 19:54:05 -0600 Subject: [PATCH 1/2] Nonlinear diff: skip evaluator rebuild when structure is unchanged _cache_evaluator! rebuilt the MOI.Nonlinear evaluator (AD tapes, Jacobian/Hessian sparsity detection) and the index mappings on every forward_differentiate!/reverse_differentiate! call, although the Cache is a pure function of the Form's structure: parameter values are read live from form.model.parameters by the evaluator, and the primal/dual point lives in model.x/y/s. Return the existing model.cache when there is one, and invalidate it on every structural mutation of the model (add_variable(s), add_constraint, objective function/sense changes; MOI.empty! already reset it). Results are bit-identical to the rebuild path; the per-differentiation setup cost is removed. --- src/NonLinearProgram/NonLinearProgram.jl | 58 ++++++++++ test/nlp_program.jl | 135 +++++++++++++++++++++++ 2 files changed, 193 insertions(+) diff --git a/src/NonLinearProgram/NonLinearProgram.jl b/src/NonLinearProgram/NonLinearProgram.jl index 3df0ecae..216eb56d 100644 --- a/src/NonLinearProgram/NonLinearProgram.jl +++ b/src/NonLinearProgram/NonLinearProgram.jl @@ -352,6 +352,48 @@ function MOI.empty!(model::Model) return end +# The `Cache` built by `_cache_evaluator!` is a pure function of the *structure* of the +# `Form` (variables, constraints, parameters, objective): the evaluator tapes, the +# Jacobian/Hessian sparsity and the index mappings do not depend on parameter values (the +# evaluator reads those live from `form.model.parameters`) nor on the primal/dual point +# (which lives in `model.x`, `model.y`, `model.s`). It therefore stays valid until the +# structure changes. The methods below intercept every structural mutation of the model +# and invalidate the cache, so that `_cache_evaluator!` can return an existing cache +# instead of rebuilding it on every differentiation call. + +function _invalidate_evaluator_cache!(model::Model) + model.cache = nothing + return +end + +function MOI.add_variable(model::Model) + _invalidate_evaluator_cache!(model) + return MOI.add_variable(model.model) +end + +function MOI.add_variables(model::Model, n) + _invalidate_evaluator_cache!(model) + return MOI.add_variables(model.model, n) +end + +function MOI.add_constraint( + model::Model, + func::MOI.AbstractFunction, + set::MOI.AbstractSet, +) + _invalidate_evaluator_cache!(model) + return MOI.add_constraint(model.model, func, set) +end + +function MOI.set( + model::Model, + attr::Union{MOI.ObjectiveSense,MOI.ObjectiveFunction}, + value, +) + _invalidate_evaluator_cache!(model) + return MOI.set(model.model, attr, value) +end + include("nlp_utilities.jl") include("vno_bridge.jl") @@ -449,7 +491,23 @@ _get_num_primal_vars(model::Model) = _get_num_primal_vars(model.model) _get_num_params(form::Form) = length(_all_params(form)) _get_num_params(model::Model) = _get_num_params(model.model) +""" + _cache_evaluator!(model::Model) + +Build the `Cache` for `model` (the `MOI.Nonlinear` evaluator with its AD tapes and +Jacobian/Hessian sparsity, plus the primal/dual/bound index mappings), or return the +existing `model.cache` when there is one. + +The cache only depends on the structure of the `Form`, and every structural mutation +resets `model.cache` to `nothing` (see `_invalidate_evaluator_cache!`), so an existing +cache is always current: rebuilding it would produce an identical evaluator and identical +mappings. Returning it directly avoids paying the tape construction and sparsity +detection on every differentiation call. +""" function _cache_evaluator!(model::Model) + if model.cache !== nothing + return model.cache + end form = model.model # Retrieve and sort primal variables by NLP index params = sort(_all_params(model); by = x -> x.value) diff --git a/test/nlp_program.jl b/test/nlp_program.jl index a9ec9e1f..bd8eddd4 100644 --- a/test/nlp_program.jl +++ b/test/nlp_program.jl @@ -1126,6 +1126,141 @@ function test_reverse_bounds_upper() @test isapprox(dp, 2.88888; atol = 1e-4) end +# Reach the `NonLinearProgram.Model` inside the differentiation backend of a JuMP model +# built with `DiffOpt.nonlinear_diff_model`. Used to check evaluator-cache reuse. +function _inner_nlp_diff_model(model::JuMP.Model) + diffopt = JuMP.unsafe_backend(model) + inner = diffopt.diff + while !(inner isa DiffOpt.NonLinearProgram.Model) + inner = inner.model + end + return inner +end + +function _build_cache_reuse_model(P) + 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] * (1 + 0.1 * sin(p[i])) >= p[i]) + @constraint(model, sum(x) >= 0.4 * 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) + ) + return model, p, x +end + +function test_evaluator_cache_reused_across_differentiations() + P = 4 + model, p, x = _build_cache_reuse_model(P) + optimize!(model) + @assert is_solved_and_feasible(model) + # The first differentiation builds the evaluator cache. + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], cos(1.3 * i)) + end + DiffOpt.reverse_differentiate!(model) + nlp = _inner_nlp_diff_model(model) + cache = nlp.cache + @test cache !== nothing + # Later differentiation calls on the same structure return the same Cache object + # instead of rebuilding tapes and sparsity. + seed = [sin(0.7 * i) for i in 1:P] + DiffOpt.empty_input_sensitivities!(model) + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], seed[i]) + end + DiffOpt.reverse_differentiate!(model) + @test nlp.cache === cache + dp_reused = [ + MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value + for i in 1:P + ] + DiffOpt.empty_input_sensitivities!(model) + DiffOpt.set_forward_parameter(model, p[1], 1.0) + DiffOpt.forward_differentiate!(model) + @test nlp.cache === cache + dx_reused = + [MOI.get(model, DiffOpt.ForwardVariablePrimal(), x[i]) for i in 1:P] + # A fresh model whose only differentiation uses the same seeds must produce + # identical results: the rebuilt cache is the same pure function of the structure, + # so reuse changes no floating-point operation. + model2, p2, x2 = _build_cache_reuse_model(P) + optimize!(model2) + for i in 1:P + DiffOpt.set_reverse_variable(model2, x2[i], seed[i]) + end + DiffOpt.reverse_differentiate!(model2) + dp_fresh = [ + MOI.get(model2, DiffOpt.ReverseConstraintSet(), ParameterRef(p2[i])).value for i in 1:P + ] + @test dp_reused == dp_fresh + DiffOpt.empty_input_sensitivities!(model2) + DiffOpt.set_forward_parameter(model2, p2[1], 1.0) + DiffOpt.forward_differentiate!(model2) + dx_fresh = + [MOI.get(model2, DiffOpt.ForwardVariablePrimal(), x2[i]) for i in 1:P] + @test dx_reused == dx_fresh + return +end + +function test_evaluator_cache_invalidated_on_structural_change() + P = 3 + model, p, x = _build_cache_reuse_model(P) + optimize!(model) + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], 1.0) + end + DiffOpt.reverse_differentiate!(model) + nlp = _inner_nlp_diff_model(model) + @test nlp.cache !== nothing + # Every structural mutation of the inner model must drop the cache. + MOI.set(nlp, MOI.ObjectiveSense(), MOI.MIN_SENSE) + @test nlp.cache === nothing + DiffOpt.reverse_differentiate!(model) + @test nlp.cache !== nothing + MOI.add_variable(nlp) + @test nlp.cache === nothing + return +end + +function test_evaluator_cache_after_public_structural_edit() + # Through the public API, a structural edit resets the whole differentiation model + # (`DiffOpt.Optimizer` sets `diff = nothing`), so differentiation after the edit + # builds a new cache and must match a fresh model built directly in the final state. + P = 3 + model, p, x = _build_cache_reuse_model(P) + optimize!(model) + for i in 1:P + DiffOpt.set_reverse_variable(model, x[i], sin(1.0 + i)) + end + DiffOpt.reverse_differentiate!(model) + nlp_before = _inner_nlp_diff_model(model) + @constraint(model, sum(x) <= 10.0 * P) + optimize!(model) + DiffOpt.reverse_differentiate!(model) + nlp_after = _inner_nlp_diff_model(model) + @test nlp_after !== nlp_before + dp = [ + MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value + for i in 1:P + ] + model2, p2, x2 = _build_cache_reuse_model(P) + @constraint(model2, sum(x2) <= 10.0 * P) + optimize!(model2) + for i in 1:P + DiffOpt.set_reverse_variable(model2, x2[i], sin(1.0 + i)) + end + DiffOpt.reverse_differentiate!(model2) + dp2 = [ + MOI.get(model2, DiffOpt.ReverseConstraintSet(), ParameterRef(p2[i])).value for i in 1:P + ] + @test dp ≈ dp2 rtol = 1e-10 + return +end + end # module TestNLPProgram.runtests() From c1bdb5318b3fc18708340789713a8f8586f4d6f3 Mon Sep 17 00:00:00 2001 From: zero Date: Sun, 5 Jul 2026 15:43:26 -0600 Subject: [PATCH 2/2] Run JuliaFormatter (margin=80) to satisfy format-check --- .../Thermal_Generation_Dispatch_Example.jl | 16 ++++++++++++---- .../Thermal_Generation_Dispatch_Example_new.jl | 16 ++++++++++++---- docs/src/examples/autotuning-ridge.jl | 5 +---- docs/src/examples/autotuning-ridge_new.jl | 5 +---- docs/src/examples/sensitivity-analysis-ridge.jl | 16 ++-------------- .../examples/sensitivity-analysis-ridge_new.jl | 16 ++-------------- 6 files changed, 30 insertions(+), 44 deletions(-) diff --git a/docs/src/examples/Thermal_Generation_Dispatch_Example.jl b/docs/src/examples/Thermal_Generation_Dispatch_Example.jl index 1fe273aa..b5846c0d 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 391f3f2f..0872c2bd 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 17619aef..28cb6661 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 b7fa24a7..87f05ca9 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 a4d3c046..5604d0b6 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 aa41f80a..59d98fec 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.