jump-dev / jump-dev/DiffOpt.jl
Possible performance improvement?
- Dominant language
- Julia
- Stars
- 145
- Forks
- 21
- PR merge metrics
- No merged PRs in 30d
Description
following an initially-unrelated discussion with @andrewrosemberg
I think there might be room for (substantial 🤔🤞) performance improvements in the `NonLinearProgram` code.
Note: what I write below might apply to other classes of problems, I only checked the nonlinear code so far.
TLDR: I think we can replace {a lot of linear system solves} with {a single linear system solve} during forward/reverse diff
Medium-long explanation:
Forward/reverse diff compute Jacobian-vector products of the form `J*v` or `J'*v`, and the Jacobian `J` is of the form `K\N`, where `K, N` are matrices. Currently, to obtain `w=J*v`, we solve `KJ = N` (matrix RHS), then compute `J*v`. I believe we can instead compute `w = Jv` by solving `Kw = (Nv)`, which requires a single linear system solve (with vector RHS).
**Why it matters** for the ACOPF example I work with, on a 300-bus case with loads as parameters, `K` has size 13k x 13k, and `N` is 13k x 200. Replacing `K\N` with `K\(Nv)` **potentially yields a 200x reduction in memory and time.**
___
Complete thought process and stack traversing.
I'm focusing on forward diff for the sake of example, same remarks apply to reverse diff
* Forward diff returns primal/dual sensitivities in a `ForwCache` structure here:
https://github.com/jump-dev/DiffOpt.jl/blob/df4a7dcee059ee048ff620590ff5478d8b2c79cc/src/NonLinearProgram/NonLinearProgram.jl#L534-L537
* These are obtained from a left multiplication with Jacobian here
https://github.com/jump-dev/DiffOpt.jl/blob/df4a7dcee059ee048ff620590ff5478d8b2c79cc/src/NonLinearProgram/NonLinearProgram.jl#L530-L532
* Said Jacobian `Δs` is computed by the `_compute_sensitivity` function. Stepping into that call, said jacobian is obtained from `_compute_derivatives_no_relax`
https://github.com/jump-dev/DiffOpt.jl/blob/df4a7dcee059ee048ff620590ff5478d8b2c79cc/src/NonLinearProgram/nlp_utilities.jl#L457
* which then points us to a linear system solve
https://github.com/jump-dev/DiffOpt.jl/blob/df4a7dcee059ee048ff620590ff5478d8b2c79cc/src/NonLinearProgram/nlp_utilities.jl#L425-L427
The alternative implementation I propose would cache the factorization of `K` (which seems to already be the case), and lazily compute Jacobian-vector products `w=J*v` by solving `w = K\(Nv)`. The math would be similar (I think) for Jacobian-transpose-vector products
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.