WIP: Sandu projection

[SKopecz]

This implements Sandu's projection method, see #124.

[JoshuaLampert]

Would it make sense and be feasible to put JuMP.jl in an extension as it is not very lightweight and not necessary for every user, such that we save the additional loading time if no Sandu projection is used?

[SKopecz]

Is 25a8127 what you had in mind?

[JoshuaLampert]

Yes, looks good. Thanks!

[codecov]

Codecov Report

❌ Patch coverage is 96.38554% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/SanduProjectionExt.jl	93.75%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[coveralls]

Pull Request Test Coverage Report for Build 17466558065

Details

• 80 of 83 (96.39%) changed or added relevant lines in 2 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage decreased (-0.06%) to 97.713%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
ext/SanduProjectionExt.jl	45	48	93.75%

Totals
Change from base Build 16001455410:	-0.06%
Covered Lines:	2051
Relevant Lines:	2099

---

💛 - Coveralls

[SKopecz]

This is a scaled version of the stratospheric reaction problem and can be solved using SanduProjection and Clarabel with default settings.

Unfortunately, scaling destroys the "conservation structure" of a PDS, i.e. terms that canceled before scaling, do not afterwards. I was not able to find a PDS representation of the scaled statospheric problem, which could be solved as good as the original one by MPRK schemes. That's why there is only an ODE version of this problem.

[SKopecz]

Since we only have an ODE version of the scaled stratospheric reaction problem (see comment above). I introduced this auxiliary function to get corresponding invariance matrix.

[SKopecz]

I would have expected that using StaticArrays is not necessary in the extension, as it is used in the parent package.

[SKopecz]

The code

using PositiveIntegrators
using OrdinaryDiffEqRosenbrock: ROS2
using JuMP: Model
using Clarabel
using Profile

AT = linear_invariants_stratreac_scaled()
b = AT * prob_ode_stratreac_scaled.u0
cb = SanduProjection(Model(Clarabel.Optimizer), AT, b, 1.0; save = false)

function doit(args...; n = 1)
    for _ in 1:n
        solve(prob_ode_stratreac_scaled, ROS2();
              save_everystep = false, callback = cb)
    end
end

@profview doit(1)
@profview doit(10^6)

generates the following flame graph.

There seems to be quite some runtime dispatch. But most of it does not seem to be related to SanduProjection. Zooming into the discrete callback section, we see:

As I see it, the runtime dispatch is a problem of Clarabel.

What do you think?

[ranocha]

If there are some instabilities within Clarabel.jl, I would not worry about them too much here unless they really hurt the performance (in which case we might have a look at their source code).

I am more worried about the red parts of the residual computation. It looks like we lose quite a lot of performance there. Would you mind trying to setup a minimal working example? Does it happen for every simple ODE using SVector? Does it change anything if you pass the global variables as arguments to doit, e.g.,

function doit(ode, cb; n = 1)
    for _ in 1:n
        solve(ode, ROS2();
              save_everystep = false, callback = cb)
    end
end

and @profview doit(prob_ode_stratreac_scaled, cb, 10^6)?

[SKopecz]

Here's a MWE. It's not only an issue for StaticArrays, but also arrays in general. But it's definitely a bigger concern for StaticArrays.

using OrdinaryDiffEqRosenbrock
using StaticArrays
using Profile

function doit(ode, n = 1)
    for _ in 1:n
        solve(ode, ROS2())
    end
end

function f(u, p, t)
    if 1.0 < t
        sigma = t
    else
        sigma = 0.0
    end

    return @SVector [sigma * u[1]]
end
u0 = @SVector ones(1)
probf = ODEProblem(f, u0, (0.0, 1.0))

@profview doit(probf)
@profview doit(probf, 10^5)

function g(u, p, t)
    if 1.0 < t
        sigma = t
    else
        sigma = 0.0
    end

    return [sigma * u[1]]
end
u0 = ones(1)
probg = ODEProblem(g, u0, (0.0, 1.0))

@profview doit(probg)
@profview doit(probg, 10^5)

function h(u, p, t)
    if 1.0 < t
        sigma = t
    else
        sigma = 0.0
    end

    return sigma *  u
end
probh = ODEProblem(h, 1.0, (0.0, 1.0))

@profview doit(probh)
@profview doit(probh, 10^5)

The flame graph for f is:

The flame graph for g is:

The flame graph for h is:

[SKopecz]

Using the following function fixes the problem! I'll modify the stratospheric reaction problems accordingly.

function f(u, p, t)
    if one(t) < t
        sigma = t
    else
        sigma = zero(t)
    end

    return @SVector [sigma * u[1]]
end

[SKopecz]

Now that f_stratreac_scaled is type stable, the situation is as follows.

using PositiveIntegrators
using OrdinaryDiffEqRosenbrock: ROS2

using JuMP
using Clarabel

using Test
using Plots

AT = linear_invariants_stratreac_scaled()
b = AT * prob_ode_stratreac_scaled.u0
cb = SanduProjection(Model(Clarabel.Optimizer), AT, b)

function doit(ode, cb, n = 1)
    for _ in 1:n
        solve(ode, ROS2();
              save_everystep = false, callback = cb)
    end
end

@profview doit(prob_ode_stratreac_scaled, cb, 1)
@profview doit(prob_ode_stratreac_scaled, cb, 10^2)

[SKopecz]

I will revise the other stratospheric reaction functions in another PR, since this should also affect the corresponding benchmark results.

[ranocha]

Now that f_stratreac_scaled is type stable, the situation is as follows.

using PositiveIntegrators
using OrdinaryDiffEqRosenbrock: ROS2

using JuMP
using Clarabel

using Test
using Plots

AT = linear_invariants_stratreac_scaled()
b = AT * prob_ode_stratreac_scaled.u0
cb = SanduProjection(Model(Clarabel.Optimizer), AT, b)

function doit(ode, cb, n = 1)
    for _ in 1:n
        solve(ode, ROS2();
              save_everystep = false, callback = cb)
    end
end

@profview doit(prob_ode_stratreac_scaled, cb, 1)
@profview doit(prob_ode_stratreac_scaled, cb, 10^2)

Nice fix 👍 This flamegraph looks quite bad - the optimizer takes a huge amount of time. Did you expect this?

[SKopecz]

I compared the runtime of all QP packages supported by JuMP, which could solve the statrospheric reaction problem with default settings. With the choice abstol = 1e-7, reltol = 1e-6 all solvers generate approximations that look good, i.e. no peaks are missing.

julia> using JuMP

julia> using Clarabel

julia> @benchmark solve(prob_ode_stratreac_scaled, ROS2(); abstol = 1e-7, reltol = 1e-6,
                   save_everystep = false, callback = SanduProjection(Model(Clarabel.Optimizer), AT, b))
BenchmarkTools.Trial: 18 samples with 1 evaluation per sample.
 Range (min … max):  253.762 ms … 333.667 ms  ┊ GC (min … max): 0.00% … 3.45%
 Time  (median):     279.162 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   290.035 ms ±  27.314 ms  ┊ GC (mean ± σ):  3.14% ± 4.25%

  ▁  ▁▁  ▁   █ ▁   █   ▁                 █▁▁         ▁     ▁▁ ▁  
  █▁▁██▁▁█▁▁▁█▁█▁▁▁█▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁███▁▁▁▁▁▁▁▁▁█▁▁▁▁▁██▁█ ▁
  254 ms           Histogram: frequency by time          334 ms <

 Memory estimate: 72.78 MiB, allocs estimate: 1463271.

julia> using COSMO

julia> @benchmark solve(prob_ode_stratreac_scaled, ROS2(); abstol = 1e-7, reltol = 1e-6,
                   save_everystep = false, callback = SanduProjection(Model(COSMO.Optimizer), AT, b)) 
BenchmarkTools.Trial: 11 samples with 1 evaluation per sample.
 Range (min … max):  466.016 ms … 559.975 ms  ┊ GC (min … max): 3.37% … 8.98%
 Time  (median):     489.071 ms               ┊ GC (median):    6.41%
 Time  (mean ± σ):   495.325 ms ±  32.541 ms  ┊ GC (mean ± σ):  5.08% ± 3.55%

  █▁ ▁   ▁      ▁   ▁    ▁▁                             ▁     ▁  
  ██▁█▁▁▁█▁▁▁▁▁▁█▁▁▁█▁▁▁▁██▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁█ ▁
  466 ms           Histogram: frequency by time          560 ms <

 Memory estimate: 195.80 MiB, allocs estimate: 3746332.

julia> using DAQP

julia> @benchmark solve(prob_ode_stratreac_scaled, ROS2(); abstol = 1e-7, reltol = 1e-6,
                   save_everystep = false, callback = SanduProjection(Model(DAQP.Optimizer), AT, b))
BenchmarkTools.Trial: 22 samples with 1 evaluation per sample.
 Range (min … max):  201.188 ms … 338.161 ms  ┊ GC (min … max): 0.00% … 22.10%
 Time  (median):     213.586 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   230.957 ms ±  41.720 ms  ┊ GC (mean ± σ):  5.28% ±  8.39%

    ▄ █▁▁▁                                                       
  ▆▆█▆████▁▆▁▆▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▆▁▆▁▁▁▁▆▁▁▁▁▁▁▁▁▆ ▁
  201 ms           Histogram: frequency by time          338 ms <

 Memory estimate: 70.27 MiB, allocs estimate: 1582602.

julia> using MadNLP

julia> @benchmark solve(prob_ode_stratreac_scaled, ROS2(); abstol = 1e-7, reltol = 1e-6,
                   save_everystep = false, callback = SanduProjection(Model(MadNLP.Optimizer), AT, b))    
BenchmarkTools.Trial: 16 samples with 1 evaluation per sample.
 Range (min … max):  290.044 ms … 361.920 ms  ┊ GC (min … max): 0.00% … 6.26%
 Time  (median):     329.323 ms               ┊ GC (median):    4.70%
 Time  (mean ± σ):   325.449 ms ±  19.264 ms  ┊ GC (mean ± σ):  4.20% ± 2.68%

  █          ▁    ▁       ▁  ▁   █   █  █ ▁▁    ▁             ▁  
  █▁▁▁▁▁▁▁▁▁▁█▁▁▁▁█▁▁▁▁▁▁▁█▁▁█▁▁▁█▁▁▁█▁▁█▁██▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
  290 ms           Histogram: frequency by time          362 ms <

 Memory estimate: 129.43 MiB, allocs estimate: 1539207.

julia> using SCS

julia> @benchmark solve(prob_ode_stratreac_scaled, ROS2(); abstol = 1e-7, reltol = 1e-6,
                   save_everystep = false, callback = SanduProjection(Model(SCS.Optimizer), AT, b))             
BenchmarkTools.Trial: 38 samples with 1 evaluation per sample.
 Range (min … max):  115.246 ms … 156.251 ms  ┊ GC (min … max): 0.00% … 18.74%
 Time  (median):     132.393 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   133.455 ms ±   9.553 ms  ┊ GC (mean ± σ):  2.00% ±  4.84%

                   █▁ ▁    ▁▄   ▄                                
  ▆▁▁▆▁▆▁▁▁▁▁▁▁▆▁▆▆██▆█▆▆▆▁██▆▁▆█▁▆▁▆▁▁▆▆▁▁▁▁▆▁▁▁▁▆▆▁▁▁▁▆▁▁▁▁▆▆ ▁
  115 ms           Histogram: frequency by time          156 ms <

 Memory estimate: 26.02 MiB, allocs estimate: 337759.

SCS is cleary the fastest. The corresponding flame graph looks like this:

The projection takes about 30% of the runtime. But also with SCS there's quite some run-time dispatch, so it's hard to say what the performance would be without this. The ODE solver performs 33607 steps, calling the projection 189 times, i.e. 0.6%. I have no experience with optimization, but 30% computing time for this small number of projection calls seems a bit excessive to me. Moreover, Sandu's paper says: "With Ros2 the optimization routine is called only a few times and the overhead is small (8%)."

The big question for me is if the instabilities come from our side or not?

So how should we proceed?

[ranocha]

Well, it's better to have a useful feature that might not be perfect than to not have the option to use it. Did Sandu give any hints about solving the optimization problem?

[SKopecz]

Sandu uses the Goldfarb and Idnani algorithm and states that additional adaptions to the specific minimization problem were made.

[SKopecz]

I did some experiments solving a simple minimization problem without any ODE stuff using different optimization packages within JuMP as well as using Clarabel or COSMO directly. All flame graphs showed significant run-time dispatch.

So I guess the problem is not on our side and would leave this as it currently is.