From bc1834617f85cf19624ef79e1f55f7e15c62cf64 Mon Sep 17 00:00:00 2001
From: Brendan <brencodeert@outlook.com>
Date: Tue, 16 Jan 2024 11:46:35 +0100
Subject: [PATCH] Added total water storage

---
 src/sbm.jl       | 35 +++++++++++++++++++++++++++++++++++
 src/sbm_model.jl | 17 +++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/src/sbm.jl b/src/sbm.jl
index a85254bf7..e9ae822e3 100644
--- a/src/sbm.jl
+++ b/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
diff --git a/src/sbm_model.jl b/src/sbm_model.jl
index 907bbd9bc..ecc07ca3a 100644
--- a/src/sbm_model.jl
+++ b/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