Water balance computation

src/Wflow.jl

@@ -139,6 +139,7 @@ struct Model{R <: Routing, L <: AbstractLandModel, T <: AbstractModelType} <:
     network::Network    # connectivity information, directed graph
     routing::R          # routing model (horizontal fluxes), moves along network
     land::L             # land model simulating vertical fluxes, independent of each other
+    # water_balance::WaterBalance{Float64}
     clock::Clock        # to keep track of simulation time
     reader::NCReader    # provides the model with dynamic input
     writer::Writer      # writes model output
@@ -148,6 +149,7 @@ end
 # prevent a large printout of model components and arrays
 Base.show(io::IO, ::AbstractModel{T}) where {T} = print(io, "model of type ", T)
 
+include("water_balance.jl")
 include("forcing.jl")
 include("parameters.jl")
 include("groundwater/connectivity.jl")
@@ -265,6 +267,8 @@ function run_timestep!(model::Model; update_func = update!, write_model_output =
     if write_model_output
         write_output(model)
     end
+    # compute_water_balance!(model, model.water_balance.water_balance_vertical)
+    # compute_water_balance!(model, model.water_balance.water_balance_overland)
     return nothing
 end
 

src/sbm_gwf_model.jl

@@ -433,11 +433,15 @@ function initialize_sbm_gwf_model(config::Config)
 
     routing = Routing(; subsurface_flow, overland_flow, river_flow)
 
+    # n = ?
+    # water_balance = WaterBalance(n, Float64)
+
     model = Model(
         config,
         network,
         routing,
         land_hydrology,
+        # water_balance,
         clock,
         reader,
         writer,

src/sediment_model.jl

@@ -110,8 +110,20 @@ function initialize_sediment_model(config::Config)
 
     routing = Routing(; overland_flow = overland_flow_sediment, river_flow = river_sediment)
 
-    model =
-        Model(config, network, routing, soilloss, clock, reader, writer, SedimentModel())
+    # n = ?
+    # water_balance = WaterBalance(n, Float64)
+
+    model = Model(
+        config,
+        network,
+        routing,
+        soilloss,
+        # water_balance,
+        clock,
+        reader,
+        writer,
+        SedimentModel(),
+    )
 
     set_states!(model)
     @info "Initialized model"

src/water_balance.jl

@@ -0,0 +1,97 @@
+struct WaterBalanceVertical{T <: AbstractFloat} # Terms in [mm]
+    storage_total::Vector{T}
+    storage_total_prev::Vector{T}
+    inflow_total::Vector{T}
+    outflow_total::Vector{T}
+    balance_error::Vector{T}
+    relative_error::Vector{T}
+end
+
+function WaterBalanceVertical(n::Integer, float_type::Type{T}) where {T <: AbstractFloat}
+    WaterBalanceVertical(ntuple(_ -> fill(float_type(Nan), n), 6)...)
+end
+
+function compute_water_balance!(model, water_balance::WaterBalanceVertical)::Nothing
+    (;
+        storage_total,
+        storage_total_prev,
+        inflow_total,
+        outflow_total,
+        balance_error,
+        relative_error,
+    ) = water_balance
+    (; vertical, lateral) = model
+
+    (; soil, runoff) = vertical
+    (; satwaterdepth, ustoredepth, actevap) = soil.variables
+    (; runoff_land, net_runoff_river) = runoff.variables
+
+    (; subsurface) = lateral
+    (; ssf, ssfin) = subsurface.variables
+    (; flow_length, flow_width) = subsurface.parameters
+
+    storage_total .= satwaterdepth
+    storage_total .+= ustoredepth
+
+    inflow_total .= vertical.atmospheric_forcing.precipitation
+    inflow_total .+= runoff.variables.net_runoff_river
+    inflow_total .+= 1000.0 * ssfin ./ (flow_length .* flow_width)
+
+    outflow_total .= runoff_land
+    outflow_total .+= net_runoff_river
+    outflow_total .+= actevap
+    outflow_total .+= 1000.0 * ssf ./ (flow_length .* flow_width)
+
+    @. balance_error = (inflow_total - outflow_total - (storage_total - storage_total_prev))
+    @. relative_error = balance_error / ((inflow_total + outflow_total) / 2)
+
+    copyto!(storage_total_prev, storage_total)
+
+    #TODO : What to do with results?
+
+    return nothing
+end
+
+struct WaterBalanceOverLand{T <: AbstractFloat} # Terms in [?]
+    volume_t_prev::Vector{T}
+    volume_rate::Vector{T}
+    water_in::Vector{T}
+    water_out::Vector{T}
+    balance_error::Vector{T}
+    relative_error::Vector{T}
+end
+
+function WaterBalanceOverLand(n::Integer, float_type::Type{T}) where {T <: AbstractFloat}
+    WaterBalanceOverLand(ntuple(_ -> fill(float_type(Nan), n), 6)...)
+end
+
+# function compute_water_balance!(model, water_balance::WaterBalanceOverLand)::Nothing
+#     (; vertical, lateral) = model
+#     (; volume_t_prev, volume_rate, water_in, water_out, balance_error, relative_error) = water_balance_overland
+
+#     (; land, river) = lateral
+#     (; volume) = land.variables
+#     volume_t = volume
+
+#     water_in .= lateral.land.qin_av
+#     @. water_in += lateral.land.qlat_av * lateral.land.dl
+
+#     water_out .= lateral.land.q_av
+
+#     @. volume_rate = (volume_t - volume_t_previous) / vertical.dt
+
+#     @. balance_error = (water_in - water_out) - volume_rate
+
+#     #TODO : What to do with results?
+
+#     return nothing
+# end
+
+struct WaterBalance{T}
+    water_balance_vertical::WaterBalanceVertical{T}
+    water_balance_over_land::WaterBalanceOverLand{T}
+end
+
+function WaterBalance(n::Integer, float_type::Type{T}) where {T <: AbstractFloat}
+    WaterBalance(WaterBalanceVertical(n, float_type), WaterBalanceOverLand(n, float_type))
+end