Fix FractionalFlow sum to 1 validation with control states

The initial fractions come from whatever happens to be the last control state per node, and therefore can trigger false positive validation errors.

Now we validate this per control state. To make that work I had to add the missing control state to the `control_mapping`. For simplicity I coalesce a missing `control_state` to the empty string. I wonder if that is the right approach though, since for e.g. name we default to the empty string rather than support missing. We should probably settle on either for all string columns.

core/src/bmi.jl

@@ -55,7 +55,7 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
         if !valid_fractional_flow(
             connectivity.graph_flow,
             fractional_flow.node_id,
-            fractional_flow.fraction,
+            fractional_flow.control_mapping,
         )
             error("Invalid fractional flow node combinations found.")
         end

core/src/create.jl

@@ -115,22 +115,17 @@ function parse_static_and_time(
                     if parameter_name in time_interpolatables
                         val = LinearInterpolation([val, val], trivial_timespan)
                     end
-                    # If this row defines a control state, collect the parameter values in
-                    # the parameter_values vector
-                    if !ismissing(control_state)
-                        push!(parameter_values, val)
-                    end
+                    # Collect the parameter values in the parameter_values vector
+                    push!(parameter_values, val)
                     # The initial parameter value is overwritten here each time until the last row,
                     # but in the case of control the proper initial parameter values are set later on
                     # in the code
                     getfield(out, parameter_name)[node_idx] = val
                 end
-                # If a control state is associated with this row, add the parameter values to the
-                # control mapping
-                if !ismissing(control_state)
-                    control_mapping[(node_id, control_state)] =
-                        NamedTuple{Tuple(parameter_names)}(Tuple(parameter_values))
-                end
+                # Add the parameter values to the control mapping
+                control_state_key = coalesce(control_state, "")
+                control_mapping[(node_id, control_state_key)] =
+                    NamedTuple{Tuple(parameter_names)}(Tuple(parameter_values))
             end
         elseif node_id in time_node_ids
             # TODO replace (time, node_id) order by (node_id, time)

core/src/validation.jl

@@ -549,18 +549,19 @@ outneighbor, that the fractions leaving a node add up to ≈1 and that the fract
 function valid_fractional_flow(
     graph_flow::DiGraph{Int},
     node_id::Vector{Int},
-    fraction::Vector{Float64},
+    control_mapping::Dict{Tuple{Int64, String}, NamedTuple},
 )::Bool
     errors = false
 
-    # Node ids that have fractional flow outneighbors
+    # Node IDs that have fractional flow outneighbors
     src_ids = Set{Int}()
 
     for id in node_id
         union!(src_ids, inneighbors(graph_flow, id))
     end
 
-    node_id_set = Set(node_id)
+    node_id_set = Set{Int}(node_id)
+    control_states = Set{String}([key[2] for key in keys(control_mapping)])
 
     for src_id in src_ids
         src_outneighbor_ids = Set(outneighbors(graph_flow, src_id))
@@ -571,27 +572,35 @@ function valid_fractional_flow(
             )
         end
 
-        fraction_sum = 0.0
-
-        for ff_id in intersect(src_outneighbor_ids, node_id_set)
-            ff_idx = findsorted(node_id, ff_id)
-            frac = fraction[ff_idx]
-            fraction_sum += frac
+        # Each control state (including missing) must sum to 1
+        for control_state in control_states
+            fraction_sum = 0.0
+
+            for ff_id in intersect(src_outneighbor_ids, node_id_set)
+                parameter_values = get(control_mapping, (ff_id, control_state), nothing)
+                if parameter_values === nothing
+                    continue
+                else
+                    (; fraction) = parameter_values
+                end
+
+                fraction_sum += fraction
+
+                if fraction <= 0
+                    errors = true
+                    @error(
+                        "Fractional flow nodes must have non-negative fractions, got $fraction for #$ff_id."
+                    )
+                end
+            end
 
-            if frac <= 0
+            if fraction_sum ≉ 1
                 errors = true
                 @error(
-                    "Fractional flow nodes must have non-negative fractions, got $frac for #$ff_id."
+                    "The sum of fractional flow fractions leaving a node must be ≈1, got $fraction_sum for #$src_id."
                 )
             end
         end
-
-        if fraction_sum ≉ 1
-            errors = true
-            @error(
-                "The sum of fractional flow fractions leaving a node must be ≈1, got $fraction_sum for #$src_id."
-            )
-        end
     end
     return !errors
 end

core/test/validation.jl

@@ -168,7 +168,7 @@ if !Sys.islinux()
             @test !Ribasim.valid_fractional_flow(
                 connectivity.graph_flow,
                 fractional_flow.node_id,
-                fractional_flow.fraction,
+                fractional_flow.control_mapping,
             )
         end