Total Water Storage

docs/src/model_docs/params_vertical.md

@@ -116,6 +116,7 @@ specific_leaf = "Sl"
 | **`leaf_area_index`** | leaf area index | m``^2`` m``{-2}`` | - |
 | `waterlevel_land` | water level land | mm | - |
 | `waterlevel_river` | water level river | mm | - |
+| `total_storage` | total water storage (excluding floodplains) | mm | - |
 
 
 ## [HBV](@id params_hbv)

src/sbm.jl

@@ -194,6 +194,8 @@
     waterlevel_land::Vector{T} | "mm"
     # Water level river [mm]
     waterlevel_river::Vector{T} | "mm"
+    # Total water storage (excluding floodplain volume) [mm]
+    total_storage::Vector{T} | "mm"
 
     function SBM{T,N,M}(args...) where {T,N,M}
         equal_size_vectors(args)
@@ -597,6 +599,7 @@ function initialize_sbm(nc, config, riverfrac, inds)
         leaf_area_index = fill(mv, n),
         waterlevel_land = fill(mv, n),
         waterlevel_river = zeros(Float, n), #set to zero to account for cells outside river domain
+        total_storage = zeros(Float, n) # Set the total water storage from initialized values
     )
 
     return sbm
@@ -1049,3 +1052,35 @@ function update_after_subsurfaceflow(sbm::SBM, zi, exfiltsatwater)
         sbm.zi[i] = zi[i]
     end
 end
+
+function update_total_water_storage(sbm::SBM, rnet, rrouting, lrouting)
+
+    # Set the total storage to zero
+    fill!(sbm.total_storage, 0)
+
+    # Burn the river routing values
+    sbm.total_storage[rnet] = (
+        rrouting.h_av .* sbm.riverfrac[rnet] * 0.001 # convert to mm
+    )
+
+    # Chunk the data for parallel computing
+    threaded_foreach(1:sbm.n, basesize=1000) do i
+
+        # Cumulate per vertical type
+        # Maybe re-categorize in the future
+        # Unique are those that are not always present
+        unique = sbm.glacierstore[i] * sbm.glacierfrac[i]
+        surface = (
+            sbm.snow[i] + sbm.snowwater[i] + sbm.canopystorage[i]
+        )
+        sub_surfuce = sbm.ustoredepth[i] + sbm.satwaterdepth[i]
+        lateral = (
+            lrouting.h_av[i] * (1 - sbm.riverfrac[i]) * 0.001 # convert to mm
+        )
+
+        # Add everything to the total water storage
+        sbm.total_storage[i] += (
+            unique + surface + sub_surfuce + lateral
+        )
+    end
+end

src/sbm_model.jl

@@ -376,6 +376,7 @@ function update(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
     # update lateral subsurface flow domain (kinematic wave)
     update(lateral.subsurface, network.land, network.frac_toriver)
     model = update_after_subsurfaceflow(model)
+    model = update_total_water_storage(model)
 end
 
 """
@@ -438,6 +439,22 @@ function update_after_subsurfaceflow(
     return model
 end
 
+function update_total_water_storage(
+    model::Model{N,L,V,R,W,T},
+) where {N,L,V,R,W,T<:SbmModel}
+    @unpack lateral, vertical, network, clock, config = model
+
+    # Update the total water storage based on vertical states
+    # TODO Maybe look at routing in the near future
+    update_total_water_storage(
+        vertical,
+        network.index_river,
+        lateral.river,
+        lateral.land,
+    )
+    return model
+end
+
 function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
     @unpack lateral, vertical, network, config = model
 

test/bmi.jl

@@ -20,9 +20,9 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @testset "model information functions" begin
             @test BMI.get_component_name(model) == "sbm"
-            @test BMI.get_input_item_count(model) == 183
-            @test BMI.get_output_item_count(model) == 183
-            @test BMI.get_input_var_names(model)[[1, 5, 150, 174]] == [
+            @test BMI.get_input_item_count(model) == 184
+            @test BMI.get_output_item_count(model) == 184
+            @test BMI.get_input_var_names(model)[[1, 5, 151, 175]] == [
                 "vertical.nlayers",
                 "vertical.θᵣ",
                 "lateral.river.sl",

test/run_sbm.jl

@@ -80,6 +80,8 @@ end
     @test sbm.runoff[50063] == 0.0
     @test sbm.soilevap[50063] == 0.0
     @test sbm.snow[5] ≈ 3.592840840467347f0
+    @test sbm.total_storage[50063] ≈ 559.70849973344f0
+    @test sbm.total_storage[429] ≈ 583.9800667536956f0 # river cell
 end
 
 # run the second timestep
@@ -92,6 +94,8 @@ model = Wflow.run_timestep(model)
     @test sbm.runoff[50063] == 0.0
     @test sbm.soilevap[50063] ≈ 0.006358004660566856f0
     @test sbm.snow[5] ≈ 3.667748983774868f0
+    @test sbm.total_storage[50063] ≈ 559.7935411649405f0
+    @test sbm.total_storage[429] ≈ 589.7026699582976f0 # river cell
 end
 
 @testset "subsurface flow" begin