From 5afb5bafd4e78c75c125d364bdbf3244d88077d1 Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Tue, 14 May 2024 15:06:17 +0200
Subject: [PATCH] Fix subgrid_id in output

Fixes #895
---
 core/src/parameter.jl        | 1 +
 core/src/read.jl             | 9 ++++++---
 core/src/write.jl            | 5 ++---
 core/test/run_models_test.jl | 3 +--
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index df3aa43c8..a4ee97098 100644
--- 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
index 2c7b2df43..785290f14 100644
--- 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,7 +988,8 @@ 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
@@ -995,8 +997,9 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
     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
index 1f2ed07d4..da89faade 100644
--- 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
index 5c9037cd2..45fa53f73 100644
--- 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