Storage and flow states with custom Jacobian + linear solve: the formulation to rule them all

[SouthEndMusic]

fixes #3113.

I did little to no vibe development for this branch, to make things as explainable as possible.

File structure changes

• What was solve.jl where the rhs for the ODE lives is renamed to formulate_flows.jl
• A new solve.jl is introduced which contains the custom mass matrix + Jacobian + linear solve implementation

Discontinuity handling

Handling discontinuities of the rhs deserves some attention. Now that I made all positive forcings part of the state, I noticed that the solver isn't able to integrate them exactly, even though they are constant over the timestep. This is because of the memory QNDF uses of previous timesteps, where the forcing is sometimes different. To fix this and related problems, I had to set the integrator.derivative_discontinuity flag when callbacks change parameters which gets most of the way there, and add a step_limiter! to get there exactly. @Huite this of course relates to our discussion about performance and what you care about to see in the result, but I think I remember discussing with @visr a while back that we should care about preserving the forcing volume.

It could be interesting to investigate whether setting integrator.derivative_discontinuity as a significant impact on performace. If it does, it's worth putting in some more effort to catch also the discontinuities not caused by callbacks but by non-differentiable time interpolations. There are also still hard clamps for certain flows, see #1918.

How to interpret the solver config options

The implementations in this PR hook into the solver at a quite low level, so they are less flexible when it comes to certain solver options than we had before. More precisely:

• sparse doesn't mean much for the AD, because the custom implementation represents the Jacobian in a lazy way, and doesn't use sparsity detection and coloring etc. on the rhs. The only sparse arrays still present are J_inner and the Jacobian of the continuous control compound variables with respect to the storages and flows, and don't think it makes sense to keep the sparse option just for those.
• autodiff: The way I implemented the Jacobian computation is locked into the ForwardDiff API. I could refactor it so that it is compatible with DifferentiationInterface again as mentioned here, which requires defining a new CVector type for all the inputs into a single vector.

I think it should be possible to bypass my whole implementation, but we need to think about what we want to support there and how that fits into the config.

Update

I discussed this with @visr and we concluded that:

• It is good to refactor the current implementation so that the AD is implemented in DifferentiationInterface terms. That's a bit more readable and can then make use of the autodiff flag.
• It should indeed be possible to bypass all the custom solve logic. Then the old sparsity detection + sparse matrix coloring + default linear solve will be used. Let's use a new flag optimized_solve = true for this.
• To avoid confusion and because it is not really necessary anymore the sparse flag can be removed.

The state: Carrays $\rightarrow$ CVectors

The use of vector components in this implementation is more complex because they are now nested: (storage, (flows...), pid_integral). This was rather broken in the CArrays implementation, so I reworked it into CVectors which additional constraints like that the ranges must be contiguous. I also added some pretty printing:

julia> model.integrator.u
33-element Ribasim.CVectors.CVector{Float64, Vector{Float64}, @NamedTuple{storage::UnitRange{Int64}, flow::@NamedTuple{pump::UnitRange{Int64}, outlet::UnitRange{Int64}, flow_boundary::UnitRange{Int64}, tabulated_rating_curve::UnitRange{Int64}, linear_resistance::UnitRange{Int64}, manning_resistance::UnitRange{Int64}, user_demand_inflow::UnitRange{Int64}, user_demand_outflow::UnitRange{Int64}, evaporation::UnitRange{Int64}, infiltration::UnitRange{Int64}, drainage::UnitRange{Int64}, surface_runoff::UnitRange{Int64}, precipitation::UnitRange{Int64}}, pid_integral::UnitRange{Int64}}}
  storage: [775.2364363043691, 775.2342825678, 572.6029998376725, 1130.004999999995]
  flow:
    pump: [4375.716386283238]
    outlet: Float64[]
    flow_boundary: [3162.2399999999943, 3162.2399999999943]
    tabulated_rating_curve: [426.97568814891827, 1280.9263593105832, 0.0]
    linear_resistance: [3614.272680166237, -676.5253671870286]
    manning_resistance: [2773.3327467955037]
    user_demand_inflow: Float64[]
    user_demand_outflow: Float64[]
    evaporation: [346.6708169001252, 346.67104958117375, 227.84697318966766, 364.4387061170138]
    infiltration: [0.0, 0.0, 0.0, 0.0]
    drainage: [0.0, 0.0, 0.0, 0.0]
    surface_runoff: [366.0, 366.0, 366.0, 366.0]
    precipitation: [366.0, 366.0, 366.0, 366.0]
  pid_integral: Float64[]

The linear solve construction

What I like about the matrix construction in build_J_inner! is that it very explicitly reveals the structure of the linear solve.

PID control

The PID control values are no longer cached, they are computed on-demand in formulate_pump_and_outlet_flow!.

Continuous control

• AD...
• The code now supports continuous control based on flow, it just has to be allowed in the input and tested

Greedy UserDemand inflow

I think the greedy UserDemand inflow when allocation is not active can be cleaned up a bit: now on main the equal split logic lives both in the flow formulation and in the step limiter, but I think it can be done in a callback where the equal split is assigned to allocated so that other places in the code do not have to know whether it was done by the allocation algorithm or by the simple split.

Latency

As discussed with @visr, precompiling Ribasim takes quite a while because of our custom state vector; a lot of what's precompiled from OrdinaryDiffEq.jl in their precompile workload is only valid for plain vector states. We could get around this by passing a plain vector state to the solver and only wrapping it for indexing convenience where needed.

Allocation

Allocation is outside the scope of this PR, but I noticed some problems with optimize_multi_objective! while debugging:

• Infeasibility is only checked after all objectives have been optimized for. If one of them was infeasible, the subsequent goal programming constraints are uninformative garbage. So infeasibility should be checked after each JuMP.optiomize!, not just after all optimizations are done
• The route priority objective seems to be implemented in an odd way: not only does it not make use of allocation_model.objectives where all objectives should live, it's also optimized for for each other objective, not just once at the end.

DiscreteControl

The only model that is still misbehaving on this branch is continuous_concentration_condition (confusing name btw, as it doesn't use ContinuousControl). This is a typical 'flipper' model because the controlled LinearResistance directly influences the salinity in the Basin being conditioned on. By how DiscreteControl is currently set up (at a per solver timestep basis), how quickly this oscillation happens is fully dependent on the adaptive time stepper, which apparently behaves differently now than with only the flow states. To me the cleanest solution is #2247.

[SouthEndMusic]

@visr I did some profiling with the hws_transient model you sent me:

• This branch (without any work on tolerances):
  • With backtracking: 2 min 42 s
  • No backtracking: 1 min 20 s

water_balance! and callbacks dominate at least at this scale, not at all the AD or linear solve. This indicates that a lot of performance can be gained from threading.

• Main (no backtracking): 1 min 56 s

[visr]

That is promising! I'd leave backtracking and threading out of this PR, and mainly check the performance on the full LHM, which I just sent by email.

Although there main hardly runs, so we'd compare to the version of Lex' branchbthat we ran with.

[SouthEndMusic]

There was backtracking in this branch because Lex added it, I removed it again. I did a quick experiment with running the transient_hws model ignoring the flow residual, but that didn't speed it up. I'll have a look at the full LHM

[SouthEndMusic]

I've been looking at some LHM profiles. The runtime at the start is really inconsistent over simulated time so I'm not sure what a good benchmark is. I did notice that:

• Updating the Jacobian is negligible
• Linear solves take up about 15% of the time, where most of the time is spent on assembling the matrix. There is a simple optimization opportunity here, because most of the assembly time is from terms that only depend on the Jacobian, so if the Jacobian didn't refresh, those terms can be maintained
• water_balance! takes up about 50% of the time, very threadable
• Callbacks take up about 25% of the time, also threadable