Merge main.

Project.toml

@@ -62,3 +62,6 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+
+[sources]
+Ribasim = {path = "core"}

core/Project.toml

@@ -65,7 +65,7 @@ ComponentArrays = "0.13, 0.14, 0.15"
 Configurations = "0.17"
 DBInterface = "2.4"
 DataFrames = "1.4"
-DataInterpolations = "6"
+DataInterpolations = "6.5"
 DataStructures = "0.18"
 Dates = "1"
 DiffEqBase = "6.155"
@@ -96,7 +96,7 @@ PreallocationTools = "0.4"
 SQLite = "1.5.1"
 SciMLBase = "2.36"
 SparseArrays = "1"
-SparseConnectivityTracer = "0.6.3"
+SparseConnectivityTracer = "0.6.8"
 StructArrays = "0.6.13"
 TOML = "1"
 Tables = "1"

core/src/config.jl

@@ -134,6 +134,19 @@ end
     use_allocation::Bool = false
 end
 
+@option struct Experimental <: TableOption
+    concentration::Bool = false
+end
+# For logging enabled experimental features
+function Base.iterate(exp::Experimental, state = 0)
+    state >= nfields(exp) && return
+    return Base.getfield(exp, state + 1), state + 1
+end
+function Base.show(io::IO, exp::Experimental)
+    fields = (field for field in fieldnames(typeof(exp)) if getfield(exp, field))
+    print(io, join(fields, " "))
+end
+
 @option @addnodetypes struct Toml <: TableOption
     starttime::DateTime
     endtime::DateTime
@@ -145,6 +158,7 @@ end
     solver::Solver = Solver()
     logging::Logging = Logging()
     results::Results = Results()
+    experimental::Experimental = Experimental()
 end
 
 struct Config
@@ -250,6 +264,18 @@ const algorithms = Dict{String, Type}(
     "Euler" => Euler,
 )
 
+"""
+Check whether the given function has a method that accepts the given kwarg.
+Note that it is possible that methods exist that accept :a and :b individually,
+but not both.
+"""
+function function_accepts_kwarg(f, kwarg)::Bool
+    for method in methods(f)
+        kwarg in Base.kwarg_decl(method) && return true
+    end
+    return false
+end
+
 "Create an OrdinaryDiffEqAlgorithm from solver config"
 function algorithm(solver::Solver; u0 = [])::OrdinaryDiffEqAlgorithm
     algotype = get(algorithms, solver.algorithm, nothing)
@@ -259,19 +285,22 @@ function algorithm(solver::Solver; u0 = [])::OrdinaryDiffEqAlgorithm
             Available options are: ($(options)).")
     end
     kwargs = Dict{Symbol, Any}()
+
     if algotype <: OrdinaryDiffEqNewtonAdaptiveAlgorithm
         kwargs[:nlsolve] = NLNewton(;
             relax = Ribasim.MonitoredBackTracking(; z_tmp = copy(u0), dz_tmp = copy(u0)),
         )
     end
-    # not all algorithms support this keyword
-    kwargs[:autodiff] = solver.autodiff
-    try
-        algotype(; kwargs...)
-    catch
-        pop!(kwargs, :autodiff)
-        algotype(; kwargs...)
+
+    if function_accepts_kwarg(algotype, :step_limiter!)
+        kwargs[:step_limiter!] = Ribasim.limit_flow!
     end
+
+    if function_accepts_kwarg(algotype, :autodiff)
+        kwargs[:autodiff] = solver.autodiff
+    end
+
+    algotype(; kwargs...)
 end
 
 "Convert the saveat Float64 from our Config to SciML's saveat"

core/src/main.jl

@@ -40,6 +40,9 @@ function main(toml_path::AbstractString)::Cint
                     @warn "The Ribasim version in the TOML config file does not match the used Ribasim CLI version." config.ribasim_version cli.ribasim_version
                 end
                 @info "Starting a Ribasim simulation." toml_path cli.ribasim_version starttime endtime
+                if any(config.experimental)
+                    @warn "The following *experimental* features are enabled: $(config.experimental)"
+                end
 
                 try
                     model = run(config)

core/src/parameter.jl

@@ -104,6 +104,7 @@ const ScalarInterpolation = LinearInterpolation{
     Vector{Float64},
     Vector{Float64},
     Float64,
+    (1,),
 }
 
 set_zero!(v) = v .= zero(eltype(v))
@@ -343,6 +344,7 @@ end
             Vector{Float64},
             ScalarInterpolation,
             Float64,
+            (1,),
         },
     }
     level_to_area::Vector{ScalarInterpolation}

core/src/solve.jl

@@ -320,29 +320,14 @@ function formulate_flow!(
 )::Nothing
     (; allocation) = p
     all_nodes_active = p.all_nodes_active[]
-    for (
-        id,
-        inflow_edge,
-        outflow_edge,
-        active,
-        demand_itp,
-        demand,
-        allocated,
-        return_factor,
-        min_level,
-        demand_from_timeseries,
-    ) in zip(
+    for (id, inflow_edge, outflow_edge, active, allocated, return_factor, min_level) in zip(
         user_demand.node_id,
         user_demand.inflow_edge,
         user_demand.outflow_edge,
         user_demand.active,
-        user_demand.demand_itp,
-        # TODO permute these so the nodes are the last dimension, for performance
-        eachrow(user_demand.demand),
         eachrow(user_demand.allocated),
         user_demand.return_factor,
         user_demand.min_level,
-        user_demand.demand_from_timeseries,
     )
         if !(active || all_nodes_active)
             continue
@@ -356,11 +341,7 @@ function formulate_flow!(
         # effectively allocated = demand.
         for priority_idx in eachindex(allocation.priorities)
             alloc_prio = allocated[priority_idx]
-            demand_prio = if demand_from_timeseries
-                demand_itp[priority_idx](t)
-            else
-                demand[priority_idx]
-            end
+            demand_prio = get_demand(user_demand, id, priority_idx, t)
             alloc = min(alloc_prio, demand_prio)
             q += alloc
         end
@@ -718,3 +699,115 @@ function formulate_flows!(
         formulate_flow!(du, user_demand, p, t, current_storage, current_level)
     end
 end
+
+"""
+Clamp the cumulative flow states within the minimum and maximum
+flow rates for the last time step if these flow rate bounds are known.
+"""
+function limit_flow!(
+    u::ComponentVector,
+    integrator::DEIntegrator,
+    p::Parameters,
+    t::Number,
+)::Nothing
+    (; uprev, dt) = integrator
+    (;
+        pump,
+        outlet,
+        linear_resistance,
+        user_demand,
+        tabulated_rating_curve,
+        basin,
+        allocation,
+    ) = p
+
+    # TabulatedRatingCurve flow is in [0, ∞) and can be inactive
+    for (id, active) in zip(tabulated_rating_curve.node_id, tabulated_rating_curve.active)
+        limit_flow!(
+            u.tabulated_rating_curve,
+            uprev.tabulated_rating_curve,
+            id,
+            0.0,
+            Inf,
+            active,
+            dt,
+        )
+    end
+
+    # Pump flow is in [min_flow_rate, max_flow_rate] and can be inactive
+    for (id, min_flow_rate, max_flow_rate, active) in
+        zip(pump.node_id, pump.min_flow_rate, pump.max_flow_rate, pump.active)
+        limit_flow!(u.pump, uprev.pump, id, min_flow_rate, max_flow_rate, active, dt)
+    end
+
+    # Outlet flow is in [min_flow_rate, max_flow_rate] and can be inactive
+    for (id, min_flow_rate, max_flow_rate, active) in
+        zip(outlet.node_id, outlet.min_flow_rate, outlet.max_flow_rate, outlet.active)
+        limit_flow!(u.outlet, uprev.outlet, id, min_flow_rate, max_flow_rate, active, dt)
+    end
+
+    # LinearResistance flow is in [-max_flow_rate, max_flow_rate] and can be inactive
+    for (id, max_flow_rate, active) in zip(
+        linear_resistance.node_id,
+        linear_resistance.max_flow_rate,
+        linear_resistance.active,
+    )
+        limit_flow!(
+            u.linear_resistance,
+            uprev.linear_resistance,
+            id,
+            -max_flow_rate,
+            max_flow_rate,
+            active,
+            dt,
+        )
+    end
+
+    # UserDemand inflow is in [0, total_demand] and can be inactive
+    for (id, active) in zip(user_demand.node_id, user_demand.active)
+        total_demand = sum(
+            get_demand(user_demand, id, priority_idx, t) for
+            priority_idx in eachindex(allocation.priorities)
+        )
+        limit_flow!(
+            u.user_demand_inflow,
+            uprev.user_demand_inflow,
+            id,
+            0.0,
+            total_demand,
+            active,
+            dt,
+        )
+    end
+
+    # Evaporation is in [0, ∞) (a stricter upper bound would require also estimating the area)
+    # Infiltration is in [0, infiltration]
+    for (id, infiltration) in zip(basin.node_id, basin.vertical_flux.infiltration)
+        limit_flow!(u.evaporation, uprev.evaporation, id, 0.0, Inf, true, dt)
+        limit_flow!(u.infiltration, uprev.infiltration, id, 0.0, infiltration, true, dt)
+    end
+
+    return nothing
+end
+
+function limit_flow!(
+    u_component,
+    uprev_component,
+    id::NodeID,
+    min_flow_rate::Number,
+    max_flow_rate::Number,
+    active::Bool,
+    dt::Number,
+)::Nothing
+    u_prev = uprev_component[id.idx]
+    if active
+        u_component[id.idx] = clamp(
+            u_component[id.idx],
+            u_prev + min_flow_rate * dt,
+            u_prev + max_flow_rate * dt,
+        )
+    else
+        u_component[id.idx] = uprev_component[id.idx]
+    end
+    return nothing
+end

core/src/util.jl

@@ -1134,3 +1134,12 @@ function isoutofdomain(u, p, t)
     formulate_storages!(current_storage, u, u, p, t)
     any(<(0), current_storage)
 end
+
+function get_demand(user_demand, id, priority_idx, t)::Float64
+    (; demand_from_timeseries, demand_itp, demand) = user_demand
+    if demand_from_timeseries[id.idx]
+        demand_itp[id.idx][priority_idx](t)
+    else
+        demand[id.idx, priority_idx]
+    end
+end

core/src/write.jl

@@ -5,7 +5,7 @@ Write all results to the Arrow files as specified in the model configuration.
 """
 function write_results(model::Model)::Model
     (; config) = model
-    (; results) = model.config
+    (; results, experimental) = model.config
     compress = get_compressor(results)
     remove_empty_table = model.integrator.t != 0
 
@@ -25,9 +25,11 @@ function write_results(model::Model)::Model
     write_arrow(path, table, compress; remove_empty_table)
 
     # concentrations
-    table = concentration_table(model)
-    path = results_path(config, RESULTS_FILENAME.concentration)
-    write_arrow(path, table, compress; remove_empty_table)
+    if experimental.concentration
+        table = concentration_table(model)
+        path = results_path(config, RESULTS_FILENAME.concentration)
+        write_arrow(path, table, compress; remove_empty_table)
+    end
 
     # discrete control
     table = discrete_control_table(model)

core/test/docs.toml

@@ -48,3 +48,7 @@ verbosity = "info" # optional, default "info", can otherwise be "debug", "warn"
 # These results files are always written
 compression = true  # optional, default true, using zstd compression
 compression_level = 6 # optional, default 6
+
+[experimental]
+# Experimental features, disabled by default
+concentration = false # tracer calculations

core/test/run_models_test.jl

@@ -184,6 +184,7 @@ end
     using Logging: Debug, with_logger
     using LoggingExtras
     using SciMLBase: successful_retcode
+    using OrdinaryDiffEqBDF: QNDF
     import Tables
     using Dates
 
@@ -197,11 +198,17 @@ end
     end
 
     @test model isa Ribasim.Model
-    p = model.integrator.p
+
+    (; integrator) = model
+    (; p, alg) = integrator
+
     @test p isa Ribasim.Parameters
     @test isconcretetype(typeof(p))
     @test all(isconcretetype, fieldtypes(typeof(p)))
 
+    @test alg isa QNDF
+    @test alg.step_limiter! == Ribasim.limit_flow!
+
     @test successful_retcode(model)
     @test length(model.integrator.sol) == 2 # start and end
     @test model.integrator.p.basin.current_storage[Float64[]] ≈
@@ -242,7 +249,7 @@ end
     precipitation = model.integrator.p.basin.vertical_flux.precipitation
     @test length(precipitation) == 4
     @test model.integrator.p.basin.current_storage[parent(du)] ≈
-          Float32[720.23611, 694.8785, 415.22371, 1334.3859] atol = 2.0 skip = Sys.isapple()
+          Float32[721.17656, 695.8066, 416.66188, 1334.4879] atol = 2.0 skip = Sys.isapple()
 end
 
 @testitem "Allocation example model" begin

docs/_quarto.yml

@@ -32,6 +32,7 @@ website:
         - install.qmd
         - changelog.qmd
         - known_issues.qmd
+        - contact.qmd
 
     - title: "Tutorials"
       contents: