Fix subgrid_id in output (#1466)

Fixes #895
Deltares · May 16, 2024 · c17d3c4 · c17d3c4
1 parent db1d1ea
commit c17d3c4
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 8 deletions.
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
@@ -630,6 +630,7 @@ end
 
 "Subgrid linearly interpolates basin levels."
 struct Subgrid
+    subgrid_id::Vector{Int32}
     basin_index::Vector{Int32}
     interpolations::Vector{ScalarInterpolation}
     level::Vector{Float64}

diff --git a/core/src/read.jl b/core/src/read.jl
@@ -970,7 +970,8 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
     node_to_basin = Dict(node_id => index for (index, node_id) in enumerate(basin.node_id))
     tables = load_structvector(db, config, BasinSubgridV1)
 
-    basin_ids = Int32[]
+    subgrid_ids = Int32[]
+    basin_index = Int32[]
     interpolations = ScalarInterpolation[]
     has_error = false
     for group in IterTools.groupby(row -> row.subgrid_id, tables)
@@ -987,16 +988,18 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
             pushfirst!(subgrid_level, first(subgrid_level))
             pushfirst!(basin_level, nextfloat(-Inf))
             new_interp = LinearInterpolation(subgrid_level, basin_level; extrapolate = true)
-            push!(basin_ids, node_to_basin[node_id])
+            push!(subgrid_ids, subgrid_id)
+            push!(basin_index, node_to_basin[node_id])
             push!(interpolations, new_interp)
         else
             has_error = true
         end
     end
 
     has_error && error("Invalid Basin / subgrid table.")
+    level = fill(NaN, length(subgrid_ids))
 
-    return Subgrid(basin_ids, interpolations, fill(NaN, length(basin_ids)))
+    return Subgrid(subgrid_ids, basin_index, interpolations, level)
 end
 
 function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation

diff --git a/core/src/write.jl b/core/src/write.jl
@@ -315,12 +315,11 @@ function subgrid_level_table(
     (; t, saveval) = saved.subgrid_level
     subgrid = integrator.p.subgrid
 
-    nelem = length(subgrid.basin_index)
+    nelem = length(subgrid.subgrid_id)
     ntsteps = length(t)
-    unique_elem_id = collect(1:nelem)
 
     time = repeat(datetime_since.(t, config.starttime); inner = nelem)
-    subgrid_id = repeat(unique_elem_id; outer = ntsteps)
+    subgrid_id = repeat(subgrid.subgrid_id; outer = ntsteps)
     subgrid_level = FlatVector(saveval)
     return (; time, subgrid_id, subgrid_level)
 end

diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
@@ -147,8 +147,7 @@
         basin_level = basin.level[1]
         @test length(p.subgrid.level) == 3
         @test diff(p.subgrid.level) ≈ [-1.0, 2.0]
-        # TODO The original subgrid IDs are lost and mapped to 1, 2, 3
-        @test subgrid.subgrid_id[1:3] == [11, 22, 33] broken = true
+        @test subgrid.subgrid_id[1:3] == [11, 22, 33]
         @test subgrid.subgrid_level[1:3] ≈
               [basin_level, basin_level - 1.0, basin_level + 1.0]
         @test subgrid.subgrid_level[(end - 2):end] == p.subgrid.level