Add support for SciMLStructures

[SebastianM-C]

Fix: #346

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

This PR together with SciML/ModelingToolkit.jl#4154 allows us to target the optimal control interface and the parameter fitting interface that is developed here from the MTK side by working with parameters as SciMLStructures (MTKParameters in the case of MTK). The MTK PR adds support for the codegen side of the BVPFunction and with that the following snippet runs

using ModelingToolkit
using ModelingToolkit: D_nounits as D, t_nounits as t
using OrdinaryDiffEqTsit5
import InfiniteOpt
using Ipopt
using DiffEqDevTools

@parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0
@variables x(t) y(t)

eqs = [D(x) ~ α * x - β * x * y,
    D(y) ~ -γ * y + δ * x * y]

@mtkcompile sys0 = System(eqs, t)
tspan = (0.0, 1.0)
u0map = [x => 4.0, y => 2.0]
parammap = [α => 1.8, β => 1.0, γ => 6.5, δ => 1.0]

oprob = ODEProblem(sys0, [u0map; parammap], tspan)
osol = solve(oprob, Tsit5())
ts = range(tspan..., length=51)
data = osol(ts, idxs=x).u

costs = [EvalAt(t)(x)-data[i] for (i, t) in enumerate(ts)]
consolidate(u, sub) = sum(abs2.(u))

@mtkcompile sys = System(eqs, t; costs, consolidate)

sys′ = subset_tunables(sys, [γ, α])
jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/50, tune_parameters=true)
jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()), verbose=true)

# plot(jsol.sol)
jsol.sol.ps
# Parameter Indexing Proxy
# ==================================================
# Parameter                | Value
# --------------------------------------------------
# α                        | 1.8000274842583655
# γ                        | 6.499863856791081
# δ                        | 1.0
# β                        | 1.0

#### BVP version

using OptimizationIpopt # Ipopt has some issue with the constraint jacobian, I'll have to see what's up with that
using OptimizationMadNLP, MadNLP
using BoundaryValueDiffEqMIRK
using BoundaryValueDiffEqFIRK

bprob = BVProblem(sys′, u0map, tspan; cse=false, fit_parameters=true)
bsol = solve(bprob, MIRK4(; optimize = MadNLPOptimizer(hessian_approximation = MadNLP.CompactLBFGS)), dt=1/50)
bsol.ps

# Parameter Indexing Proxy
# ==================================================
# Parameter                | Value
# --------------------------------------------------
# α                        | 1.8702268592290692
# γ                        | 6.429985754524716
# δ                        | 1.0
# β                        | 1.0

bsol = solve(bprob, RadauIIa5(; optimize = MadNLPOptimizer(hessian_approximation = MadNLP.CompactLBFGS)), dt=1/50)
bsol.ps

# Parameter Indexing Proxy
# ==================================================
# Parameter                | Value
# --------------------------------------------------
# α                        | 1.800027483065039
# γ                        | 6.499863861325147
# δ                        | 1.0
# β                        | 1.0

I'm not sure where should we add the tests for this, but at least this seems to work with MTK codegen on a simple case.
As a small note, there is a bug in OptimizationIpopt's constraint jabcobian handling, so I used MadNLP instead.

Should I add more testing here or would that go on the MTK side? or both? Since this is a PR to a PR, it would be harder to run tests from the MTK side.

@ChrisRackauckas Let me know what you think and how should I proceed here.

[ChrisRackauckas]

redirect to master

[github-actions]

Benchmark Results

Click to check benchmark results

	master	d5f757c...	master / d5f757c...
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK2()	0.575 ± 0.023 s	0.588 ± 0.021 s	0.978 ± 0.053
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK3()	13.7 ± 0.66 ms	13.7 ± 0.71 ms	0.997 ± 0.071
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK4()	3.14 ± 0.17 ms	3.14 ± 0.23 ms	1 ± 0.092
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK5()	9.29 ± 1.3 ms	9.27 ± 1.2 ms	1 ± 0.19
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK6()	1.62 ± 0.29 ms	1.62 ± 0.29 ms	0.996 ± 0.25
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = false)	1.84 ± 0.65 ms	1.83 ± 0.65 ms	1 ± 0.5
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = true)	3.07 ± 0.99 ms	3.2 ± 1 ms	0.96 ± 0.43
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = false)	0.0643 ± 0.008 s	0.0675 ± 0.0065 s	0.953 ± 0.15
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = true)	0.0787 ± 0.014 s	0.0812 ± 0.015 s	0.969 ± 0.25
Simple Pendulum/IIP/Shooting(Tsit5())	0.24 ± 0.072 ms	0.257 ± 0.078 ms	0.935 ± 0.4
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK2()	0.748 ± 0.005 s	0.764 ± 0.02 s	0.979 ± 0.027
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK3()	16.7 ± 6 ms	16.4 ± 5.1 ms	1.02 ± 0.49
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK4()	3.68 ± 0.21 ms	3.57 ± 0.19 ms	1.03 ± 0.081
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK5()	11 ± 1.9 ms	10.8 ± 1.1 ms	1.02 ± 0.21
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK6()	1.91 ± 0.22 ms	1.87 ± 0.17 ms	1.02 ± 0.15
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = false)	3.66 ± 3 ms	3.7 ± 3 ms	0.989 ± 1.1
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = true)	6.36 ± 5.3 ms	6.13 ± 5.2 ms	1.04 ± 1.2
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = false)	0.128 ± 0.0055 s	0.115 ± 0.0071 s	1.12 ± 0.085
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = true)	0.144 ± 0.012 s	0.138 ± 0.0039 s	1.05 ± 0.095
Simple Pendulum/OOP/Shooting(Tsit5())	0.651 ± 0.065 ms	0.663 ± 0.042 ms	0.982 ± 0.12
time_to_load	5.82 ± 0.048 s	5.62 ± 0.0037 s	1.03 ± 0.0086

### Benchmark Plots A plot of the benchmark results has been uploaded as an artifact to the workflow run for this PR. Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

[SebastianM-C]

@ErikQQY @ChrisRackauckas Can you take a look at this? I think it should be good to go and after a new release of the package(s) it should also fix the tests in SciML/ModelingToolkit.jl#4154
As I understand, the MTK tests should go inside MTK, so I'll follow up with more tests there.

[ErikQQY]

The original unflatten operation is really allocating a lot, maybe we can change to use iterators to construct immediate solution from one-dimensional state vectors, for example reshape state vectors and use eachcol iterator?

[SebastianM-C]

Changed to that in d5f757c.

[ErikQQY]

This generally looks pretty good.

As a small note, there is a bug in OptimizationIpopt's constraint jabcobian handling, so I used MadNLP instead.

Since all of the optimal control functionality is using OptimiationIpopt, just courious if this bug affect the current usage of OptimizaitonIpopt in BoundaryValueDiffEq.jl?

[SebastianM-C]

Since all of the optimal control functionality is using OptimiationIpopt, just courious if this bug affect the current usage of OptimizaitonIpopt in BoundaryValueDiffEq.jl?

I tested now again and I can't reproduce the error anymore, it was happening with the above examples a few weeks ago, but I didn't dig in at the time. I'll look more into OptimizationIpopt, but currently it seems to be working.