Cheap and accurate integrated flux computation in conservation based ODE systems

[SouthEndMusic]

Background

In Ribasim we solve a system of ODEs for the water storage in water bodies of the form

$$ \frac{\text{d}S}{\text{d}t} = BC(t) + \sum q_\text{in} - \sum q_\text{out} $$

i.e. a simple conservation equation per water body with boundary conditions and exchanges with neighboring water bodies. Our first OrdinaryDiffEq.jl implementation of this used the water storages in the water bodies as states. A downside however is that the solver then doesn't tell us the integrated fluxes between the water bodies, which is important for many of our applications. These fluxes can not be reconstructed from the storage changes as this is an underdetermined problem. We can compute the integrated fluxes from the instantaneous fluxes only when using simple algorithms (e.g. Euler forward/backward), but we settled on QNDF because it gives us the best performance (large models have ~50.000 water bodies).

Our next step was to reformulate our ODE in terms of integrated fluxes, i.e the ODEs simply become

$$ \frac{\text{d}Q}{dt} = q, $$

where the storages are formulated from the $Q$ in the right hand side. This way the solver gives us accurate integrated fluxes. This however has a large performance drawback, because there are many more flux terms than there are storage terms.

Current developments

I'm currently working on a PR where I formulate the ODE system in terms of integrated flux states, but solve internally in terms of storage states. The key insight is that the computation from integrated fluxes to storages is linear, so that our ODE system is of the form

$$ \frac{\text{d}\textbf{u}_\text{integrated flux}}{\text{d}t} = f(\textbf{u}_\text{integrated flux}, p, t) = g(\textbf{u}_\text{storage}, p, t) = g( A\textbf{u}_\text{integrated flux} + b(t), p, t), $$

using that

$$ \textbf{u}_\text{storage} = A\textbf{u}_\text{integrated flux} + b(t) $$

where $A$ is a highly sparse matrix with $1$ and $-1$ entries based on the connectivity of the water bodies and $b(t)$ consists of the initial storages and exactly integrated influxes.

There are several things we can do to take advantage of this structure of $f$, using that $\textbf{u}_\text{storage}$ lives in a much lower dimensional space ($\mathbb{R}_{\ge 0}^k$) than $\textbf{u}_\text{integrated flux}$ ($\mathbb{R}_{\ge 0}^n$):

• For Jacobian computation we can make use of the chain rule to only use AD to compute the $(n, k)$ sized Jacobian of $g$ instead of the $(n, n)$ sized Jacobian of $f$;
• For the linear solve we can use the Woodbury matrix identity to solve in terms of $k$ unknowns instead of $n$ unknowns, see here for the precise derivation.

Recommendations and open questions

• It required quite a bit of overloading internals of OrdinaryDiffEq.jl to get the above to work. Since it is quite a general procedure (applicable for every conservation based ODE system where accurate integrated flux computation is important), it might merit its own subpackage
• I would like to also like to compute EEst in the reduced space, but this isn't easy to do because of how this computation is incorporated in perform_step! based on calculate_residuals! and internalnorm in an algorithm specific way
• Currently we have to constantly update our reltol in the integrated flux formulation because of the ever growing state values (even in a dynamic equilibrium situation), it would be nice if this could be handled internally
• Generally speaking, it seems that the OrdinaryDiffEq.jl internals aren't really designed to incorporate the described kind of 'state space reduction', is there a nice way to get around this?
• Maybe there is a way to achieve the above with a state object that contains both S and Q?

[Huite]

I expect having both states S and Q to be more straightforward: you have an ordinary dS/dt function to integrate. Q is basically tacked on and along for the ride. You have a rhs f(u) with u = [S ; Q], for your Newton system you look only at part of the rhs f_s(S) and reconstruct Q cheaply.

[ChrisRackauckas]

where the storages are formulated from the in the right hand side. This way the solver gives us accurate integrated fluxes. This however has a large performance drawback, because there are many more flux terms than there are storage terms.

IntegratingCallback handles that without the major overhead. Wouldn't that be the cheapest way to do this?

[SouthEndMusic]

Scanning the internals of that callback it isn't directly clear to me what happens, but iirc earlier experiments with that callback resulted in integrated fluxes that did not agree accurately enough with the changes in storage; does it use an interpolation of the solution over the latest timestep?

[ChrisRackauckas]

It should be high order. And Gauss Kronrod is error controlled.

does it use an interpolation of the solution over the latest timestep?

Yes, so as long as you don't have a discontinuity it should be same order as the integration.

[SouthEndMusic]

I care more about the integrated fluxes matching with the storage changes than the integrated fluxes matching the ODE system with high accuracy. So if the interpolation + numerical integration introduces different errors than the ODE step that's bad, which is why I prefer not to compute these 2 things separately.

[SouthEndMusic]

Here's a different idea that requires much less code than I have now: say we go back to the formulation in terms of storage states. In an IntegratingCallback we compute an initial guess $Q$ of the integrated fluxes. Then we solve

$$ \text{argmin}_\tilde{Q} \| Q - \tilde{Q} \|_2 \quad \text{subject to} \quad A\tilde{Q} = \Delta S $$

i.e. find the $\tilde{Q}$ closest to $Q$ such that the storage changes agree with the fluxes, which is a simple projection problem with solution $\tilde{Q} = Q + M(\Delta S - AQ)$ where $M = A^T(AA^T)^{-1}$ can be computed once during initialization.