Deltares / Ribasim

Water resources modeling
https://ribasim.org/
MIT License
39 stars 5 forks source link

implicit scheme #40

Closed Huite closed 1 year ago

Huite commented 1 year ago

Here's a rough sketch:

"""
This is the parameters struct used in the DiffEq, which contains all variables except u & du.
"""
struct Parameters
    level_connections::Vector{LevelConnection?}
    outflow_tables::Vector{OutflowTable?}
    precipitation::Vector{Precipitation}
end

struct LevelConnection
   a::Int
   b::Int
   conductance::Float
end

struct Precipitation
    index::Int
    value::Float
end

function formulate!(p::Precipitation, u, du)
    i = p.index
    du[i] += p.value
    return
end

function formulate!(lc::LevelConnection, u, du)
    a = lc.a
    b = lc.b
    diff = u[b] - u[a]
    flow = lc.conductance * diff
    du[a] = -flow
    du[b] = flow
    return
end

function ode_func!(du, u, p, t)
    @unpack level_connections, outflow_tables, precipitation = p

    for p in precipitation
        formulate!(p, u, du)
    end

    for connection in level_connection
       formulate!(level_connection, u, du)
    end

    return
end

u0 = zeros(n_reservoir)
tspan = (0.0, 100.0)
problem = ODEProblem(ode_func!, u0, tspan, parameters)
sol = solve(problem, callback=callbacks)
Huite commented 1 year ago

For an educational example: https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/ As far as I can tell, we should be able to help a little with the Jacobian if we're choosing an implicit solution.

Huite commented 1 year ago

Another thought:

We know the state variables exactly in advance, i.e. the storage volume. So maybe a nice approach is to filter reservoir nodes from the graph, and enumerate them. This then creates a lookup from the "global node number" to the index in the state variable.

To setup the connections is then a matter of looking up the global node numbers.

In general, I believe it's worth to always do the translation step to internal indices, because the ode_func may be called quite a lot for adaptive timestepping.

gijsber commented 1 year ago

extend implementation to implicit scheme

gijsber commented 1 year ago

set to open because of need to implement implicit scheme

gijsber commented 1 year ago

~4d need a function to create derivative per struct for every boundary condition: create_jacobian if the user chooses a implicit solver we provide the jacobian

visr commented 1 year ago

119 switches the default algorithm to an implicit scheme, QNDF, this already works. We just need to disable automatic differentiation, since that is not supported yet. So I turned it off by default. The error if you do turn it on is rather informative, so pasting it here. With this probably this issue becomes less urgent, though points 2 and 3 would still be good to look at at some point.

First call to automatic differentiation for the Jacobian
failed. This means that the user `f` function is not compatible
with automatic differentiation. Methods to fix this include:

1. Turn off automatic differentiation (e.g. Rosenbrock23() becomes
   Rosenbrock23(autodiff=false)). More details can befound at
   https://docs.sciml.ai/DiffEqDocs/stable/features/performance_overloads/
2. Improving the compatibility of `f` with ForwardDiff.jl automatic
   differentiation (using tools like PreallocationTools.jl). More details
   can be found at https://docs.sciml.ai/DiffEqDocs/stable/basics/faq/#Autodifferentiation-and-Dual-Numbers
3. Defining analytical Jacobians. More details can be
   found at https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction

Note: turning off automatic differentiation tends to have a very minimal
performance impact (for this use case, because it's forward mode for a
square Jacobian. This is different from optimization gradient scenarios).
However, one should be careful as some methods are more sensitive to
accurate gradients than others. Specifically, Rodas methods like `Rodas4`
and `Rodas5P` require accurate Jacobians in order to have good convergence,
while many other methods like BDF (`QNDF`, `FBDF`), SDIRK (`KenCarp4`),
and Rosenbrock-W (`Rosenbrock23`) do not. Thus if using an algorithm which
is sensitive to autodiff and solving at a low tolerance, please change the
algorithm as well.
Huite commented 1 year ago

On reflection:

Huite commented 1 year ago

Here's an educational example, also with steady-state:

##
using OrdinaryDiffEq
using Plots

const Float = Float64

##
# Let's set up a simple example three connected linear reservoirs. The
# reservoirs are parametrized by a reservoir coefficient. They are forced by
# precipitation and evaporation.

function plot_solution(solution)
    plot(solution.t, [u[1] for u in solution.u])
    plot!(solution.t, [u[2] for u in solution.u])
    plot!(solution.t, [u[3] for u in solution.u])
end

function sigmoid(x, xcrit)
    x = clamp(x / xcrit, 0.0, 1.0)
    return x^2 * (-2 * x + 3)
end

struct Parameters
    k::Vector{Float}
    A::Vector{Float}
    P::Vector{Float}
    E::Vector{Float}
end

flow(p, u, index) = p.k[index] * u[index]
precipitation(p, index) = p.P[index] * p.A[index]
evaporation(p, u, index) = p.E[index] * p.A[index] * sigmoid(u[index], 10.0)

##
# f = du/dt
# Water flows from the first to the second, from the second to the third, and
# from the third it leaves the model domain.
function f!(du, u, p, t)
    du[1] = (-flow(p, u, 1) + precipitation(p, 1) - evaporation(p, u, 1))
    du[2] = (-flow(p, u, 2) + precipitation(p, 2) - evaporation(p, u, 2) + flow(p, u, 1))
    du[3] = (-flow(p, u, 3) + precipitation(p, 2) - evaporation(p, u, 3) + flow(p, u, 2))
end

##

u = [0.0, 0.0, 500.0]
du = [0.0, 0.0, 0.0]
k = [0.01, 0.01, 0.01]
P = [0.001, 0.001, 0.001]
A = [10_000.0, 10_000.0, 10_000.0]
E = [0.0005, 0.0005, 0.005]
p = Parameters(k, A, P, E)
tspan = (0.0, 100.0)

prob = ODEProblem(f!, u, tspan, p)
solution = solve(prob, Euler(); dt = 10.0)
plot_solution(solution)

prob = ODEProblem(f!, u, tspan, p)
solution = solve(prob, ImplicitEuler(); dt = 10.0)
plot_solution(solution)

prob = ODEProblem(f!, u, tspan, p)
solution = solve(prob, QNDF(); dt = 10.0)
plot_solution(solution)

prob = ODEProblem(f!, u, tspan, p)
solution = solve(prob, Rodas5(); dt = 10.0)
plot_solution(solution)

##
# Now let's set up a steady-state problem instead.

using NonlinearSolve
using SciMLNLSolve

problem = NonlinearProblem(prob)
solution = solve(problem, NLSolveJL())
Huite commented 1 year ago

Euler, dt=10.0: image

ImplicitEuler, dt=10.0: image

QNDF, dt=10.0: image

Rodas5, dt=10.0: image

visr commented 1 year ago

Today I discussed with @Huite how we could configure steady state runs.

In the existing TOML solver section: https://github.com/Deltares/Ribasim/blob/3677e2c283f246a0cdd6db0ea447964aa2453d8c/docs/core/usage.qmd#L52-L62

We could add a transient key that defaults to true. If it is set to false the ODEProblem will be converted to a NonlinearProblem, with an approriate algorithm value. We could use the existing saveat setting to find the timestamp(s) for which we want a steady state solution.

visr commented 1 year ago

This should be split up into separate remaining issues, since the "implicit scheme" works.