🚀 Add NonLinearProgram Support to DiffOpt.jl #260

andrewrosemberg · 2024-12-06T15:07:48Z

This PR introduces a new module, NonLinearProgram, to extend DiffOpt.jl's functionality for differentiating nonlinear optimization problems (NLPs). The implementation integrates with JuMP-based nonlinear models and supports advanced derivative computation through custom evaluator and differentiation logic.

🆕 Features

Nonlinear Model Differentiation:
- Forward and reverse differentiation for user-specified primal variables (focus_vars) and dual variables (focus_duals) with respect to a given set of parameters.
- Extensible API for handling NLP-specific sensitivities using custom utilities.
Core Structures:
- Cache: Stores primal variables, parameters, evaluator, and constraints for efficient reuse.
- ForwCache: Holds results of forward differentiation, including sensitivities for specified variables.
- ReverseCache: Holds results of reverse differentiation (implemented in this PR).
Integration with DiffOpt API:
- Implements DiffOpt.AbstractModel for seamless compatibility with DiffOpt's API.
- Provides forward_differentiate! and reverse_differentiate! functions for NLPs.

🔧 How It Works

Custom Sensitivity Calculations:
- Utilizes implicit function differentiation (similar to sIPOPT) to compute derivatives of the optimization problem with respect to selected parameters.
- Allows differentiation of only specific variables and duals, offering a fine-grained approach compared to existing DiffOpt modules.
Forward Differentiation:
- Computes Jacobian products for selected primal and dual variables using the provided parameters.
- Results are cached for efficient access through the DiffOpt API.
Reverse Differentiation:
- Computes transpose Jacobian products for selected primal and dual variables with respect to the parameters.
- Results are cached in ReverseCache.

📜 Implementation Highlights

Forward Differentiation:
```
DiffOpt.forward_differentiate!(model)
```
- Calculates sensitivity of variables and duals w.r.t.params`.
- Results are cached in ForwCache.
Reverse Differentiation:
```
DiffOpt.reverse_differentiate!(model)
```
- Computes transpose Jacobian products for parameters with respect to the primal and dual variables.
- Results are cached in ReverseCache.
Custom Utilities:
- Leverages create_evaluator, compute_sensitivity, and other utilities from nlp_utilities.jl for efficient derivative computation.

📋 Example Usage

Forward Differentiation

using DiffOpt
using JuMP
using Ipopt

# Define a JuMP model
model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
@variable(model, p ∈ MOI.Parameter(0.1))
@variable(model, x >= p)
@variable(model, y >= 0)
@objective(model, Min, x^2 + y^2)
@constraint(model, con, x + y >= 1)

# Solve
JuMP.optimize!(model)

# set parameter pertubations
MOI.set(model, DiffOpt.ForwardParameter(), params[1], 0.2)

# forward differentiate
DiffOpt.forward_differentiate!(model)

# Retrieve sensitivities
dx = MOI.get(model, DiffOpt.ForwardVariablePrimal(), x)
dy = MOI.get(model, DiffOpt.ForwardVariablePrimal(), y)

Reverse Differentiation

# Set Primal Pertubations
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, 1.0)

# Reverse differentiation
DiffOpt.reverse_differentiate!(model)

# Retrieve reverse sensitivities (example usage)
dp= MOI.get(model, DiffOpt.ReverseParameter(), p)

🚧 TODO

Forward Differentiation: Implemented and validated for custom primal and dual variables.
Reverse Differentiation: Implemented to compute transpose Jacobian products.
Tests: Add comprehensive test cases for forward and reverse differentiation.
Documentation: Extend user documentation with more examples and edge case scenarios.

🛠 Future Work

Allow different methods to factorize the "M" matrix and solve the linear system of equations needed.

frapac

good job @andrewrosemberg ! I appreciate the numerous unit-tests you have shipped with this PR.

I will try to run your code locally on small NLP instances, to assess how fast is the current implementation (my main concern is the total number of allocations). I will try also to test your code on parameterized OPF instances, to assess how far we can get in term of size.

test/data/nlp_problems.jl

src/NonLinearProgram/nlp_utilities.jl

frapac · 2025-01-08T18:44:17Z

@andrewrosemberg Playing with your branch right now, I must say I like DiffOpt's interface!

I did a dummy mistake when solving the problem HS15 and retrieving the sensitivity. A MWE is:

model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
@variable(model, p[1:2] ∈ MOI.Parameter.([100.0, 1.0]))
@variable(model, x[1:2])
set_upper_bound.(x[1], 0.5)
@objective(model, Min, p[1] * (x[2] - x[1]^2)^2 +  (p[2] - x[1])^2)
@constraint(model, x[1] * x[2] >= 1.0)
@constraint(model, x[1] + x[2]^2 >= 0.0)

optimize!(model)
# set parameter pertubations
MOI.set(model, DiffOpt.ForwardParameter(), p[1], 1.0)

# # forward differentiate
DiffOpt.forward_differentiate!(model)

Δx = [
    MOI.get(model, DiffOpt.ForwardVariablePrimal(), var) for
    var in x
]

I forgot to specify the sensitivity for p[2], resulting in an error. I am wondering what would be the expected behavior here:

by default, should we set the sensitivity to 0.0 ?
or should we raise a proper error message if the user has forgotten to specify some sensitivities?

frapac · 2025-01-08T20:01:25Z

Also, if I query the solution after calling forward_differentiate I get an error. Is this expected?

The MWE is:

model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
@variable(model, p[1:2] ∈ MOI.Parameter.([100.0, 1.0]))
@variable(model, x[1:2])
set_upper_bound.(x[1], 0.5)
@objective(model, Min, p[1] * (x[2] - x[1]^2)^2 +  (p[2] - x[1])^2)
@constraint(model, x[1] * x[2] >= 1.0)
@constraint(model, x[1] + x[2]^2 >= 0.0)
optimize!(model)
# set parameter pertubations
MOI.set(model, DiffOpt.ForwardParameter(), p[1], 1.0)
MOI.set(model, DiffOpt.ForwardParameter(), p[2], 1.0)

# # forward differentiate
DiffOpt.forward_differentiate!(model)

JuMP.value.(x)

Output:

┌ Warning: The model has been modified since the last call to `optimize!` (or `optimize!` has not been called yet). If you are iteratively querying solution information and modifying a model, query all the results first, then modify the model.
└ @ JuMP ~/.julia/packages/JuMP/CU7H5/src/optimizer_interface.jl:1085
ERROR: LoadError: OptimizeNotCalled()

andrewrosemberg · 2025-01-08T22:17:04Z

I forgot to specify the sensitivity for p[2], resulting in an error. I am wondering what would be the expected behavior here:

by default, should we set the sensitivity to 0.0 ?

or should we raise a proper error message if the user has forgotten to specify some sensitivities?

@frapac good question! I am fine either way. @joaquimg what would be the consistent approach given the rest of DIffOpt?

andrewrosemberg · 2025-01-08T22:23:14Z

Also, if I query the solution after calling forward_differentiate I get an error. Is this expected?

I don't think this is the desired outcome haha Not sure how to ask MOI to ignore the forward attributes, but I will look into it. @joaquimg do you have an idea on how to avoid this?

frapac · 2025-01-09T07:15:27Z

@andrewrosemberg Another issue I noted: if we are using non-standard indexing in JuMP (DenseAxisArray instead of Vector{VariableRef}) I got an error when retrieving the sensitivity

ERROR: LoadError: MethodError: no method matching similar(::Type{Vector{Float64}}, ::Tuple{UnitRange{Int64}})
The function `similar` exists, but no method is defined for this combination of argument types.

A MWE is:

model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
@variable(model, p[1:2] ∈ MOI.Parameter.([100.0, 1.0]))
# N.B: use non-standard indexing
@variable(model, x[0:2])
set_upper_bound.(x[1], 0.5)
@objective(model, Min, p[1] * (x[2] - x[1]^2)^2 +  (p[2] - x[1])^2)
@constraint(model, x[1] * x[2] >= 1.0)
@constraint(model, x[1] + x[2]^2 >= 0.0)
optimize!(model)
# set parameter pertubations
MOI.set(model, DiffOpt.ForwardParameter(), p[1], 1.0)
MOI.set(model, DiffOpt.ForwardParameter(), p[2], 1.0)

# # forward differentiate
DiffOpt.forward_differentiate!(model)

Δx = [
    MOI.get(model, DiffOpt.ForwardVariablePrimal(), var) for
    var in x
]

frapac · 2025-01-09T07:24:13Z

Pushing it one step further, I have tried to differentiate an ACOPF instance using DiffOpt. I re-used the code in rosetta-opf, and ported it to DiffOpt. You can find a gist here.

Two observations:

on case9, DiffOpt is using inertia correction. This is not expected, as I am sure we satisfy all the good constraints qualifications in case9. I am not sure to understand why UMFPACK fails to factorize the KKT system.
on case1354pegase, I get an out-of-memory error. This is caused by this line, which allocates a dense matrix. I would suggest replacing diagm by SparseArrays.spdiagm

codecov · 2025-02-20T14:27:27Z

Codecov Report

Attention: Patch coverage is 92.49012% with 38 lines in your changes missing coverage. Please review.

Project coverage is 89.07%. Comparing base (72b50c5) to head (5b26b88).
Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
src/NonLinearProgram/NonLinearProgram.jl	91.02%	22 Missing ⚠️
src/moi_wrapper.jl	68.18%	14 Missing ⚠️
src/NonLinearProgram/nlp_utilities.jl	99.47%	1 Missing ⚠️
src/diff_opt.jl	92.85%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #260      +/-   ##
==========================================
+ Coverage   86.67%   89.07%   +2.40%     
==========================================
  Files          13       15       +2     
  Lines        1478     1968     +490     
==========================================
+ Hits         1281     1753     +472     
- Misses        197      215      +18

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

joaquimg · 2025-02-20T19:24:34Z

src/NonLinearProgram/NonLinearProgram.jl

+objective_sense(form::Form) = form.sense
+objective_sense(model::Model) = objective_sense(model.model)


missing underscore

src/NonLinearProgram/nlp_utilities.jl

src/diff_opt.jl

src/jump_moi_overloads.jl

andrewrosemberg added 17 commits December 6, 2024 09:58

NonLinearProgram

0dc8099

index by MOI index

4b16e13

only cache gradient

db0c0ce

update API

d8b38f4

start reverse mode

eaa4c69

add overloads

e170339

update MOI wrapper

b006b57

update code for DiffOpt API

5ade750

working code

ff1052f

usage example

24fb230

add reverse diff

56a4d1e

update code

b9781ba

update tests

1e94996

update tests

48e7b19

add forward_differentiate! tests

501fe83

add reverse_differentiate! tests

28ec54f

update docs

d3563dd

andrewrosemberg changed the title ~~[WIP] 🚀 Add NonLinearProgram Support to DiffOpt.jl~~ 🚀 Add NonLinearProgram Support to DiffOpt.jl Dec 23, 2024

andrewrosemberg added 2 commits December 23, 2024 15:29

format

7867ece

update API reference

4868748

joaquimg mentioned this pull request Dec 27, 2024

Integrate with POI to improve UX #262

Merged

4 tasks

frapac reviewed Jan 4, 2025

View reviewed changes

andrewrosemberg added 2 commits January 6, 2025 12:25

fix typos

1d5dd4a

update reference

89d34ea

andrewrosemberg added 7 commits February 18, 2025 10:46

add comment

a0b2f3d

rm rm kwargs

095a5d9

use correct underscore signature for private funcs

db91862

format

65631b3

change github actions to v3

e2cb3be

reverse checkout version

d508789

add reference sipopt paper

0797d22

joaquimg closed this Feb 18, 2025

joaquimg reopened this Feb 18, 2025

andrewrosemberg added 2 commits February 18, 2025 12:46

update factorization routine API

09fec6b

format

e82060e

andrewrosemberg mentioned this pull request Feb 18, 2025

Test Code Cov NlpProgram #275

Closed

Update ci.yml

1bad199

andrewrosemberg and others added 12 commits February 20, 2025 11:01

improve coverage

660b513

add test inertia correction

afa0182

add test ReverseConstraintDual

6cec90b

format

5105bc7

rm useless checks

1da5f71

add test get ReverseConstraintDual

200299d

format

9633691

rm unecessary funcs

9b4cef5

rm kwargs

1847392

format

9e78f2e

rename factorization attributte

ef100b2

add supports

83deef4

joaquimg reviewed Feb 21, 2025

View reviewed changes

joaquimg mentioned this pull request Feb 21, 2025

Diff information of dual variables #94

Open

Apply suggestions from code review

5b26b88

joaquimg merged commit 275cebc into jump-dev:master Feb 21, 2025
5 of 6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

🚀 Add NonLinearProgram Support to DiffOpt.jl #260

🚀 Add NonLinearProgram Support to DiffOpt.jl #260

andrewrosemberg commented Dec 6, 2024 •

edited

Loading

frapac left a comment

frapac commented Jan 8, 2025

frapac commented Jan 8, 2025

andrewrosemberg commented Jan 8, 2025

andrewrosemberg commented Jan 8, 2025

frapac commented Jan 9, 2025

frapac commented Jan 9, 2025 •

edited

Loading

codecov bot commented Feb 20, 2025 •

edited

Loading

joaquimg Feb 20, 2025

		objective_sense(form::Form) = form.sense
		objective_sense(model::Model) = objective_sense(model.model)

🚀 Add NonLinearProgram Support to DiffOpt.jl #260

🚀 Add NonLinearProgram Support to DiffOpt.jl #260

Conversation

andrewrosemberg commented Dec 6, 2024 • edited Loading

🆕 Features

🔧 How It Works

📜 Implementation Highlights

📋 Example Usage

Forward Differentiation

Reverse Differentiation

🚧 TODO

🛠 Future Work

frapac left a comment

Choose a reason for hiding this comment

frapac commented Jan 8, 2025

frapac commented Jan 8, 2025

andrewrosemberg commented Jan 8, 2025

andrewrosemberg commented Jan 8, 2025

frapac commented Jan 9, 2025

frapac commented Jan 9, 2025 • edited Loading

codecov bot commented Feb 20, 2025 • edited Loading

Codecov Report

joaquimg Feb 20, 2025

Choose a reason for hiding this comment

andrewrosemberg commented Dec 6, 2024 •

edited

Loading

frapac commented Jan 9, 2025 •

edited

Loading

codecov bot commented Feb 20, 2025 •

edited

Loading