Generic loading of user concentration data.

core/src/Ribasim.jl

@@ -140,6 +140,9 @@ using Tables: Tables, AbstractRow, columntable
 # Wrapper around a vector of structs to easily retrieve the same field from all elements.
 using StructArrays: StructVector
 
+# OrderedSet is used to store the order of the substances in the network.
+using DataStructures: OrderedSet
+
 export libribasim
 
 include("schema.jl")

core/src/parameter.jl

@@ -347,6 +347,7 @@ end
     mass::Matrix{Float64}  # basin, substance
     concentration_external::Vector{Dict{String, ScalarInterpolation}} =
         Dict{String, ScalarInterpolation}[]
+    substances::OrderedSet{Symbol}
 end
 
 """

core/src/read.jl

@@ -416,7 +416,7 @@ end
 function LevelBoundary(db::DB, config::Config)::LevelBoundary
     static = load_structvector(db, config, LevelBoundaryStaticV1)
     time = load_structvector(db, config, LevelBoundaryTimeV1)
-    concentration = load_structvector(db, config, LevelBoundaryConcentrationV1)
+    concentration_data = load_structvector(db, config, LevelBoundaryConcentrationV1)
 
     _, _, node_ids, valid = static_and_time_node_ids(db, static, time, "LevelBoundary")
 
@@ -428,9 +428,11 @@ function LevelBoundary(db::DB, config::Config)::LevelBoundary
     parsed_parameters, valid =
         parse_static_and_time(db, config, LevelBoundary; static, time, time_interpolatables)
 
-    concentration = zeros(length(node_ids), 7)
+    substances = get_substances(db, config)
+    concentration = zeros(length(node_ids), length(substances))
     concentration[:, 1] .= 1.0  # Continuity
     concentration[:, 3] .= 1.0  # UserDemand
+    set_concentrations!(concentration, concentration_data, substances, Int32.(node_ids))
 
     if !valid
         error("Errors occurred when parsing LevelBoundary data.")
@@ -447,6 +449,7 @@ end
 function FlowBoundary(db::DB, config::Config, graph::MetaGraph)::FlowBoundary
     static = load_structvector(db, config, FlowBoundaryStaticV1)
     time = load_structvector(db, config, FlowBoundaryTimeV1)
+    concentration_data = load_structvector(db, config, FlowBoundaryConcentrationV1)
 
     _, _, node_ids, valid = static_and_time_node_ids(db, static, time, "FlowBoundary")
 
@@ -467,9 +470,11 @@ function FlowBoundary(db::DB, config::Config, graph::MetaGraph)::FlowBoundary
         end
     end
 
-    concentration = zeros(length(node_ids), 7)
+    substances = get_substances(db, config)
+    concentration = zeros(length(node_ids), length(substances))
     concentration[:, 1] .= 1.0  # Continuity
     concentration[:, 4] .= 1.0  # UserDemand
+    set_concentrations!(concentration, concentration_data, substances, Int32.(node_ids))
 
     if !valid
         error("Errors occurred when parsing FlowBoundary data.")
@@ -583,36 +588,37 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
     static = load_structvector(db, config, BasinStaticV1)
     time = load_structvector(db, config, BasinTimeV1)
     state = load_structvector(db, config, BasinStateV1)
-    concentration_state = load_structvector(db, config, BasinConcentrationStateV1)
-    concentration_boundary = load_structvector(db, config, BasinConcentrationV1)
+    concentration_state_data = load_structvector(db, config, BasinConcentrationStateV1)
+    concentration_boundary_data = load_structvector(db, config, BasinConcentrationV1)
 
     # TODO Also include all other unique substances (in boundary nodes)
     # TODO Move into a function
-    substances = Set(Symbol.(concentration_state.substance))
-    substances = Set{Symbol}()
-    # Add fractions
-    union!(
-        substances,
-        (
-            :Continuity,
-            :Initial,
-            :LevelBoundary,
-            :FlowBoundary,
-            :UserDemand,
-            :Drainage,
-            :Precipitation,
-        ),
-    )
+    substances = get_substances(db, config)
     concentration_state = zeros(n, length(substances))
     concentration_state[:, 1] .= 1.0  # Continuity
     concentration_state[:, 2] .= 1.0  # Initial
+    set_concentrations!(concentration_state, concentration_state_data, substances, node_id)
     mass = copy(concentration_state)
 
     concentration = zeros(2, n, length(substances))
     concentration[1, :, 1] .= 1.0  # Drainage / Continuity
     concentration[1, :, 6] .= 1.0  # Drainage / Drainage
     concentration[2, :, 1] .= 1.0  # Precipitation / Continuity
     concentration[2, :, 7] .= 1.0  # Precipitation / Precipitation
+    set_concentrations!(
+        view(concentration, 1, :, :),
+        concentration_boundary_data,
+        substances,
+        node_id;
+        concentration_column = :drainage,
+    )
+    set_concentrations!(
+        view(concentration, 1, :, :),
+        concentration_boundary_data,
+        substances,
+        node_id;
+        concentration_column = :precipitation,
+    )
 
     set_static_value!(table, node_id, static)
     set_current_value!(table, node_id, time, config.starttime)
@@ -695,6 +701,7 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
         concentration,
         mass,
         concentration_external,
+        substances,
     )
 
     storage0 = get_storages_from_levels(basin, state.level)
@@ -1068,7 +1075,7 @@ end
 function UserDemand(db::DB, config::Config, graph::MetaGraph)::UserDemand
     static = load_structvector(db, config, UserDemandStaticV1)
     time = load_structvector(db, config, UserDemandTimeV1)
-    concentration = load_structvector(db, config, UserDemandConcentrationV1)
+    concentration_data = load_structvector(db, config, UserDemandConcentrationV1)
     ids = get_ids(db, "UserDemand")
 
     _, _, node_ids, valid = static_and_time_node_ids(db, static, time, "UserDemand")
@@ -1126,9 +1133,11 @@ function UserDemand(db::DB, config::Config, graph::MetaGraph)::UserDemand
         config,
     )
 
-    concentration = zeros(length(node_ids), 7)
+    substances = get_substances(db, config)
+    concentration = zeros(length(node_ids), length(substances))
     concentration[:, 1] .= 1.0  # Continuity
     concentration[:, 5] .= 1.0  # UserDemand
+    set_concentrations!(concentration, concentration_data, substances, ids)
 
     if errors || !valid_demand(node_ids, demand_itp, priorities)
         error("Errors occurred when parsing UserDemand data.")
@@ -1496,3 +1505,50 @@ function create_storage_tables(
     end
     return area, level
 end
+
+"Determine all substances present in the input over multiple tables"
+function get_substances(db::DB, config::Config)::OrderedSet{Symbol}
+    # Hardcoded tracers
+    substances = OrderedSet{Symbol}([
+        :Continuity,
+        :Initial,
+        :LevelBoundary,
+        :FlowBoundary,
+        :UserDemand,
+        :Drainage,
+        :Precipitation,
+    ])
+    for table in [
+        BasinConcentrationStateV1,
+        BasinConcentrationV1,
+        FlowBoundaryConcentrationV1,
+        LevelBoundaryConcentrationV1,
+        UserDemandConcentrationV1,
+    ]
+        data = load_structvector(db, config, table)
+        for row in data
+            push!(substances, Symbol(row.substance))
+        end
+    end
+    return substances
+end
+
+"Set values in wide concentration matrix from a long input table."
+function set_concentrations!(
+    concentration,
+    concentration_data,
+    substances,
+    node_ids;
+    concentration_column = :concentration,
+)
+    for substance in unique(concentration_data.substance)
+        data_sub = filter(row -> row.substance == substance, concentration_data)
+        sub_idx = findfirst(==(Symbol(substance)), substances)
+        for group in IterTools.groupby(row -> row.node_id, data_sub)
+            first_row = first(group)
+            node_idx = findfirst(==(first_row.node_id), node_ids)
+            # TODO Now only sets the first static value, should be changed to interpolation
+            concentration[node_idx, sub_idx] = getproperty(first_row, concentration_column)
+        end
+    end
+end

core/src/write.jl

@@ -280,24 +280,19 @@ function concentration_table(
     substance::Vector{String},
     concentration::Vector{Float64},
 }
-    (; saved) = model
+    (; saved, integrator) = model
+    (; p) = integrator
+    (; basin) = p
+
     # The last timestep is not included; there is no period over which to compute flows.
     data = get_storages_and_levels(model)
 
     ntsteps = length(data.time) - 1
     nbasin = length(data.node_id)
-    nsubstance = 7
+    nsubstance = length(basin.substances)
     nrows = ntsteps * nbasin * nsubstance
 
-    substances = [
-        "Continuity",
-        "Initial",
-        "LevelBoundary",
-        "FlowBoundary",
-        "UserDemand",
-        "Drainage",
-        "Precipitation",
-    ]
+    substances = String.(basin.substances)
     concentration = zeros(nrows)
 
     idx_row = 0