Use interpolations to update forcing in Basin (#2006)

Fixes another bullet point from #601, namely:

- use interpolation objects instead of `update_basin`, which still
relies on time being sorted first

Now we can also use `parse_static_and_time` for Basin, bringing our node
types more in line with each other and more code reuse. For this I had
to extend it a bit to support other interpolation types, and support IDs
not being in time nor static. We still add time of new data as tstops to
ensure the solver knows about net data as it happens.

Labeling this as breaking since before this we allowed a node ID to be
in both the static and time tables. In `parse_static_and_time` (the
other node types) we specifically check for this with a clear error
message. I had to update `basic_transient`, to remove the static table
there.

Author: visr

core/src/callback.jl

@@ -10,14 +10,7 @@ function create_callbacks(
     u0::ComponentVector,
     saveat,
 )::Tuple{CallbackSet, SavedResults}
-    (;
-        starttime,
-        basin,
-        flow_boundary,
-        level_boundary,
-        user_demand,
-        tabulated_rating_curve,
-    ) = parameters
+    (; starttime, basin, flow_boundary, level_boundary, user_demand) = parameters
     callbacks = SciMLBase.DECallback[]
 
     # Check for negative storage
@@ -42,7 +35,13 @@ function create_callbacks(
     end
 
     # Update Basin forcings
-    tstops = get_tstops(basin.time.time, starttime)
+    # All variables are given at the same time, so just precipitation works
+    times = [itp.t for itp in basin.forcing.precipitation]
+    tstops = Float64[]
+    for t in times
+        append!(tstops, t)
+    end
+    unique!(sort!(tstops))
     basin_cb = PresetTimeCallback(tstops, update_basin!; save_positions = (false, false))
     push!(callbacks, basin_cb)
 
@@ -699,27 +698,45 @@ function save_subgrid_level(u, t, integrator)
     return copy(integrator.p.subgrid.level)
 end
 
-"Load updates from 'Basin / time' into the parameters"
+"Update one current vertical flux from an interpolation at time t."
+function set_flux!(
+    fluxes::AbstractVector{Float64},
+    interpolations::Vector{ScalarConstantInterpolation},
+    i::Int,
+    t,
+)::Nothing
+    val = interpolations[i](t)
+    # keep old value if new value is NaN
+    if !isnan(val)
+        fluxes[i] = val
+    end
+    return nothing
+end
+
+"""
+Update all current vertical fluxes from an interpolation at time t.
+
+This runs in a callback rather than the RHS since that gives issues with the discontinuities
+in the ConstantInterpolations we use, failing the vertical_flux_means test.
+"""
 function update_basin!(integrator)::Nothing
-    (; p) = integrator
+    (; p, t) = integrator
     (; basin) = p
-    (; node_id, time, vertical_flux) = basin
-    t = datetime_since(integrator.t, integrator.p.starttime)
 
-    rows = searchsorted(time.time, t)
-    timeblock = view(time, rows)
-
-    table = (;
-        vertical_flux.precipitation,
-        vertical_flux.potential_evaporation,
-        vertical_flux.drainage,
-        vertical_flux.infiltration,
-    )
+    update_basin!(basin, t)
+    return nothing
+end
 
-    for row in timeblock
-        i = searchsortedfirst(node_id, NodeID(NodeType.Basin, row.node_id, 0))
-        set_table_row!(table, row, i)
+function update_basin!(basin::Basin, t)::Nothing
+    (; vertical_flux, forcing) = basin
+    for id in basin.node_id
+        i = id.idx
+        set_flux!(vertical_flux.precipitation, forcing.precipitation, i, t)
+        set_flux!(vertical_flux.potential_evaporation, forcing.potential_evaporation, i, t)
+        set_flux!(vertical_flux.infiltration, forcing.infiltration, i, t)
+        set_flux!(vertical_flux.drainage, forcing.drainage, i, t)
     end
+
     return nothing
 end
 

core/src/model.jl

@@ -83,6 +83,7 @@ function Model(config::Config)::Model
         # tell the solver to stop when new data comes in
         tstops = Vector{Float64}[]
         for schema_version in [
+            BasinTimeV1,
             FlowBoundaryTimeV1,
             LevelBoundaryTimeV1,
             UserDemandTimeV1,

core/src/parameter.jl

@@ -108,6 +108,10 @@ Base.isless(id_1::NodeID, id_2::NodeID)::Bool = id_1.value < id_2.value
 Base.isless(id_1::Integer, id_2::NodeID)::Bool = id_1 < id_2.value
 Base.isless(id_1::NodeID, id_2::Integer)::Bool = id_1.value < id_2
 
+"ConstantInterpolation from a Float64 to a Float64"
+const ScalarConstantInterpolation =
+    ConstantInterpolation{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64, (1,)}
+
 "LinearInterpolation from a Float64 to a Float64"
 const ScalarInterpolation = LinearInterpolation{
     Vector{Float64},
@@ -373,6 +377,20 @@ end
         Dict{String, ScalarInterpolation}[]
 end
 
+"""
+Data source for Basin parameter updates over time
+
+This is used for both static and dynamic values,
+the length of each Vector is the number of Basins.
+"""
+@kwdef struct BasinForcing
+    precipitation::Vector{ScalarConstantInterpolation} = ScalarConstantInterpolation[]
+    potential_evaporation::Vector{ScalarConstantInterpolation} =
+        ScalarConstantInterpolation[]
+    drainage::Vector{ScalarConstantInterpolation} = ScalarConstantInterpolation[]
+    infiltration::Vector{ScalarConstantInterpolation} = ScalarConstantInterpolation[]
+end
+
 """
 Requirements:
 
@@ -389,7 +407,7 @@ else
     T = Vector{Float64}
 end
 """
-@kwdef struct Basin{V, C, CD, D} <: AbstractParameterNode
+@kwdef struct Basin{V, CD, D} <: AbstractParameterNode
     node_id::Vector{NodeID}
     inflow_ids::Vector{Vector{NodeID}} = [NodeID[]]
     outflow_ids::Vector{Vector{NodeID}} = [NodeID[]]
@@ -420,8 +438,7 @@ end
     # Values for allocation if applicable
     demand::Vector{Float64} = zeros(length(node_id))
     allocated::Vector{Float64} = zeros(length(node_id))
-    # Data source for parameter updates
-    time::StructVector{BasinTimeV1, C, Int}
+    forcing::BasinForcing = BasinForcing()
     # Storage for each Basin at the previous time step
     storage_prev::Vector{Float64} = zeros(length(node_id))
     # Level for each Basin at the previous time step
@@ -929,11 +946,11 @@ const ModelGraph = MetaGraph{
     Float64,
 }
 
-@kwdef mutable struct Parameters{C1, C2, C3, C4, C6, C7, C8, C9, C10, C11}
+@kwdef mutable struct Parameters{C1, C3, C4, C6, C7, C8, C9, C10, C11}
     const starttime::DateTime
     const graph::ModelGraph
     const allocation::Allocation
-    const basin::Basin{C1, C2, C3, C4}
+    const basin::Basin{C1, C3, C4}
     const linear_resistance::LinearResistance
     const manning_resistance::ManningResistance
     const tabulated_rating_curve::TabulatedRatingCurve

core/src/read.jl

@@ -15,6 +15,8 @@ function parse_static_and_time(
     time::Union{StructVector, Nothing} = nothing,
     defaults::NamedTuple = (; active = true),
     time_interpolatables::Vector{Symbol} = Symbol[],
+    interpolation_type::Type{<:AbstractInterpolation} = LinearInterpolation,
+    is_complete::Bool = true,
 )::Tuple{NamedTuple, Bool}
     # E.g. `PumpStatic`
     static_type = eltype(static)
@@ -42,7 +44,16 @@ function parse_static_and_time(
         # If the type is a union, then the associated parameter is optional and
         # the type is of the form Union{Missing,ActualType}
         parameter_type = if parameter_name in time_interpolatables
-            ScalarInterpolation
+            # We need the concrete type to store in the parameters
+            # The interpolation_type is not concrete because they don't have the
+            # constructors we use
+            if interpolation_type == LinearInterpolation
+                ScalarInterpolation
+            elseif interpolation_type == ConstantInterpolation
+                ScalarConstantInterpolation
+            else
+                error("Unknown interpolation type.")
+            end
         elseif isa(parameter_type, Union)
             nonmissingtype(parameter_type)
         else
@@ -120,7 +131,7 @@ function parse_static_and_time(
                         val = defaults[parameter_name]
                     end
                     if parameter_name in time_interpolatables
-                        val = LinearInterpolation(
+                        val = interpolation_type(
                             [val, val],
                             trivial_timespan;
                             cache_parameters = true,
@@ -149,19 +160,16 @@ function parse_static_and_time(
             for parameter_name in parameter_names
                 # If the parameter is interpolatable, create an interpolation object
                 if parameter_name in time_interpolatables
-                    val, is_valid = get_scalar_interpolation(
+                    val = get_scalar_interpolation(
                         config.starttime,
                         t_end,
                         time,
                         node_id,
                         parameter_name;
                         default_value = hasproperty(defaults, parameter_name) ?
                                         defaults[parameter_name] : NaN,
+                        interpolation_type,
                     )
-                    if !is_valid
-                        errors = true
-                        @error "A $parameter_name time series for $node_id has repeated times, this can not be interpolated."
-                    end
                 else
                     # Activity of transient nodes is assumed to be true
                     if parameter_name == :active
@@ -173,6 +181,19 @@ function parse_static_and_time(
                 end
                 getfield(out, parameter_name)[node_id.idx] = val
             end
+        elseif !is_complete
+            # Apply the defaults just like if it was in static but missing
+            for parameter_name in parameter_names
+                val = defaults[parameter_name]
+                if parameter_name in time_interpolatables
+                    val = interpolation_type(
+                        [val, val],
+                        trivial_timespan;
+                        cache_parameters = true,
+                    )
+                end
+                getfield(out, parameter_name)[node_id.idx] = val
+            end
         else
             @error "$node_id data not in any table."
             errors = true
@@ -651,24 +672,20 @@ function ConcentrationData(
         for group in IterTools.groupby(row -> row.substance, data_id)
             first_row = first(group)
             substance = first_row.substance
-            itp, no_duplication = get_scalar_interpolation(
+            itp = get_scalar_interpolation(
                 config.starttime,
                 t_end,
                 StructVector(group),
                 NodeID(:Basin, first_row.node_id, 0),
-                :concentration,
+                :concentration;
+                interpolation_type = LinearInterpolation,
             )
             concentration_external_id["concentration_external.$substance"] = itp
             if any(itp.u .< 0)
                 errors = true
                 @error "Found negative concentration(s) in `Basin / concentration_external`." node_id =
                     id, substance
             end
-            if !no_duplication
-                errors = true
-                @error "There are repeated time values for in `Basin / concentration_external`." node_id =
-                    id substance
-            end
         end
         push!(concentration_external, concentration_external_id)
     end
@@ -691,26 +708,52 @@ function ConcentrationData(
 end
 
 function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
-    node_id = get_node_ids(db, NodeType.Basin)
-    n = length(node_id)
-
     # both static and time are optional, but we need fallback defaults
     static = load_structvector(db, config, BasinStaticV1)
     time = load_structvector(db, config, BasinTimeV1)
     state = load_structvector(db, config, BasinStateV1)
 
-    # Forcing
-    precipitation = zeros(n)
-    potential_evaporation = zeros(n)
-    drainage = zeros(n)
-    infiltration = zeros(n)
-    table = (; precipitation, potential_evaporation, drainage, infiltration)
+    _, _, node_id, valid =
+        static_and_time_node_ids(db, static, time, NodeType.Basin; is_complete = false)
+    if !valid
+        error("Problems encountered when parsing Basin static and time node IDs.")
+    end
+
+    time_interpolatables =
+        [:precipitation, :potential_evaporation, :drainage, :infiltration]
+    parsed_parameters, valid = parse_static_and_time(
+        db,
+        config,
+        Basin;
+        static,
+        time,
+        time_interpolatables,
+        interpolation_type = ConstantInterpolation,
+        defaults = (;
+            precipitation = NaN,
+            potential_evaporation = NaN,
+            drainage = NaN,
+            infiltration = NaN,
+        ),
+        is_complete = false,
+    )
 
-    set_static_value!(table, node_id, static)
-    set_current_value!(table, node_id, time, config.starttime)
-    check_no_nans(table, "Basin")
+    forcing = BasinForcing(;
+        parsed_parameters.precipitation,
+        parsed_parameters.potential_evaporation,
+        parsed_parameters.drainage,
+        parsed_parameters.infiltration,
+    )
 
-    vertical_flux = ComponentVector(; table...)
+    # Current forcing is stored as separate array for BMI access
+    # These are updated from the interpolation objects at runtime
+    n = length(node_id)
+    vertical_flux = ComponentVector(;
+        precipitation = zeros(n),
+        potential_evaporation = zeros(n),
+        drainage = zeros(n),
+        infiltration = zeros(n),
+    )
 
     # Profiles
     area, level = create_storage_tables(db, config)
@@ -736,11 +779,14 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
         vertical_flux,
         storage_to_level,
         level_to_area,
-        time,
+        forcing,
         concentration_data,
         concentration_time,
     )
 
+    # Ensure the initial data is loaded at t0 for BMI
+    update_basin!(basin, 0.0)
+
     storage0 = get_storages_from_levels(basin, state.level)
     @assert length(storage0) == n "Basin / state length differs from number of Basins"
     basin.storage0 .= storage0
@@ -1074,39 +1120,30 @@ function user_demand_time!(
 
         active[user_demand_idx] = true
         demand_from_timeseries[user_demand_idx] = true
-        return_factor_itp, is_valid_return = get_scalar_interpolation(
+        return_factor_itp = get_scalar_interpolation(
             config.starttime,
             t_end,
             StructVector(group),
             NodeID(:UserDemand, first_row.node_id, 0),
             :return_factor;
+            interpolation_type = LinearInterpolation,
         )
-        if is_valid_return
-            return_factor[user_demand_idx] = return_factor_itp
-        else
-            @error "The return_factor(t) relationship for UserDemand $(first_row.node_id) from the time table has repeated timestamps, this can not be interpolated."
-            errors = true
-        end
+        return_factor[user_demand_idx] = return_factor_itp
 
         min_level[user_demand_idx] = first_row.min_level
 
         priority_idx = findsorted(priorities, first_row.priority)
-        demand_p_itp, is_valid_demand = get_scalar_interpolation(
+        demand_p_itp = get_scalar_interpolation(
             config.starttime,
             t_end,
             StructVector(group),
             NodeID(:UserDemand, first_row.node_id, 0),
             :demand;
             default_value = 0.0,
+            interpolation_type = LinearInterpolation,
         )
         demand[user_demand_idx, priority_idx] = demand_p_itp(0.0)
-
-        if is_valid_demand
-            demand_itp[user_demand_idx][priority_idx] = demand_p_itp
-        else
-            @error "The demand(t) relationship for UserDemand $(first_row.node_id) of priority $(first_row.priority_idx) from the time table has repeated timestamps, this can not be interpolated."
-            errors = true
-        end
+        demand_itp[user_demand_idx][priority_idx] = demand_p_itp
     end
     return errors
 end