Make Network data structure explicit in structs

src/Wflow.jl

@@ -25,22 +25,23 @@ using DelimitedFiles: readdlm
 using FieldMetadata: @metadata
 using Glob: glob
 using Graphs:
-    Graphs,
-    Graph,
-    DiGraph,
     add_edge!,
-    is_cyclic,
+    add_vertex!,
+    DiGraph,
+    dst,
+    edges,
+    Graph,
+    Graphs,
+    induced_subgraph,
     inneighbors,
+    is_cyclic,
+    ne,
+    nv,
     outneighbors,
-    edges,
-    topological_sort_by_dfs,
+    SimpleDiGraph,
     src,
-    dst,
-    vertices,
-    nv,
-    ne,
-    induced_subgraph,
-    add_vertex!
+    topological_sort_by_dfs,
+    vertices
 using IfElse: IfElse
 using LoggingExtras
 using LoopVectorization: @tturbo
@@ -113,18 +114,19 @@ function Clock(config, reader)
 end
 
 include("io.jl")
+include("connectivity.jl")
 
 abstract type AbstractModel{T} end
 
 """
-    Model{N,L,V,R,W}
+    Model{L, V, R, W, T}
 
 Composite type that represents all different aspects of a Wflow Model, such as the
 network, parameters, clock, configuration and input and output.
 """
-struct Model{N, L, V, R, W, T} <: AbstractModel{T}
+struct Model{L, V, R, W, T} <: AbstractModel{T}
     config::Config  # all configuration options
-    network::N  # connectivity information, directed graph
+    network::Network  # connectivity information, directed graph
     lateral::L  # lateral model that holds lateral state, moves along network
     vertical::V  # vertical model that holds vertical state, independent of each other
     clock::Clock  # to keep track of simulation time

src/bmi.jl

@@ -43,7 +43,7 @@ Update the model for a single timestep.
 # Arguments
 - `run = nothing`: to update a model partially.
 """
-function BMI.update(model::Model; run = nothing)
+function BMI.update(model::Model; run::Union{Nothing, String} = nothing)
     if isnothing(run)
         run_timestep!(model)
     elseif run == "sbm_until_recharge"

src/connectivity.jl

@@ -0,0 +1,82 @@
+@kwdef struct NetworkDrain
+    indices::Vector{CartesianIndex{2}} = CartesianIndex{2}[]
+    reverse_indices::Matrix{Int64} = zeros(Int, 0, 0)
+end
+
+@kwdef struct NetworkLake
+    river_indices::Vector{Int} = Int[]
+end
+
+# Stores edges in x and y direction between cells of a Vector with CartesianIndex(x, y), for
+# staggered grid calculations.
+@with_kw struct Indices
+    xu::Vector{Int} = Int[]     # index of neighbor cell in the (+1, 0) direction
+    xd::Vector{Int} = Int[]     # index of neighbor cell in the (-1, 0) direction
+    yu::Vector{Int} = Int[]     # index of neighbor cell in the (0, +1) direction
+    yd::Vector{Int} = Int[]     # index of neighbor cell in the (0, -1) direction
+end
+
+@kwdef struct NetworkLand
+    allocation_area_indices::Vector{Vector{Int64}} = Vector{Int}[]
+    altitude::Vector{Float64} = Float64[]
+    area::Vector{Float64} = Float64[]
+    frac_to_river::Vector{Float64} = Float64[]
+    graph::SimpleDiGraph{Int} = DiGraph(0)
+    indices::Vector{CartesianIndex{2}} = CartesianIndex{2}[]
+    order::Vector{Int} = Int[]
+    order_of_subdomains::Vector{Vector{Int}} = Vector{Int}[]
+    order_subdomain::Vector{Vector{Int}} = Vector{Int}[]
+    reverse_indices::Matrix{Int} = zeros(Int, 0, 0)
+    river_indices::Vector{Int} = Int[]
+    river_inds_excl_waterbody::Vector{Int} = Int[]
+    slope::Vector{Float64} = Float64[]
+    staggered_indices::Indices = Indices()
+    subdomain_indices::Vector{Vector{Int}} = Vector{Int}[]
+    upstream_nodes::Vector{Vector{Int}} = Vector{Int}[]
+    waterbody::Vector{Bool} = Bool[]
+end
+
+@kwdef struct NetworkReservoir
+    indices_coverage::Vector{Vector{CartesianIndex{2}}} = Vector{CartesianIndex{2}}[]
+    indices_outlet::Vector{CartesianIndex{2}} = CartesianIndex{2}[]
+    reverse_indices::Matrix{Int} = zeros(Int, 0, 0)
+    river_indices::Vector{Int} = Int[]
+end
+
+@kwdef struct NodesAtEdge
+    src::Vector{Int} = Int[]
+    dst::Vector{Int} = Int[]
+end
+
+@kwdef struct EdgesAtNode
+    src::Vector{Vector{Int}} = Vector{Int}[]
+    dst::Vector{Vector{Int}} = Vector{Int}[]
+end
+
+@kwdef struct NetworkRiver
+    allocation_area_indices::Vector{Vector{Int64}} = Vector{Int}[]
+    cell_area::Vector{Float64} = Float64[]
+    edges_at_node::EdgesAtNode = EdgesAtNode()
+    graph::SimpleDiGraph{Int} = DiGraph(0)
+    indices::Vector{CartesianIndex{2}} = CartesianIndex{2}[]
+    lake_indices::Vector{Int} = Int[]
+    land_indices::Vector{Int} = Int[]
+    nodes_at_edge::NodesAtEdge = NodesAtEdge()
+    order::Vector{Int} = Int[]
+    order_of_subdomains::Vector{Vector{Int}} = Vector{Int}[]
+    order_subdomain::Vector{Vector{Int}} = Vector{Int}[]
+    reservoir_indices::Vector{Int} = Int[]
+    reverse_indices::Matrix{Int} = zeros(Int, 0, 0)
+    subdomain_indices::Vector{Vector{Int}} = Vector{Int}[]
+    upstream_nodes::Vector{Vector{Int}} = Vector{Int}[]
+end
+
+@kwdef struct Network
+    drain::NetworkDrain = NetworkDrain()
+    lake::NetworkLake = NetworkLake()
+    land::NetworkLand = NetworkLand()
+    reservoir::NetworkReservoir = NetworkReservoir()
+    river::NetworkRiver = NetworkRiver()
+    frac_to_river::Vector{Float64} = Float64[]
+    index_river::Vector{Int} = Int[]
+end

src/routing/surface_local_inertial.jl

@@ -583,15 +583,6 @@ function update!(
     return nothing
 end
 
-# Stores edges in x and y direction between cells of a Vector with CartesianIndex(x, y), for
-# staggered grid calculations.
-@with_kw struct Indices
-    xu::Vector{Int}     # index of neighbor cell in the (+1, 0) direction
-    xd::Vector{Int}     # index of neighbor cell in the (-1, 0) direction
-    yu::Vector{Int}     # index of neighbor cell in the (0, +1) direction
-    yd::Vector{Int}     # index of neighbor cell in the (0, -1) direction
-end
-
 # maps the fields of struct Indices to the defined Wflow cartesian indices of const
 # neigbors.
 const dirs = (:yd, :xd, :xu, :yu)

src/routing/surface_routing.jl

@@ -37,16 +37,15 @@ end
 
 """
     surface_routing!(
-        model::Model{N,L,V,R,W,T}
-    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,LocalInertialOverlandFlow,LocalInertialRiverFlow}},V,R,W,T}
+        model::Model{L}
+    ) where {L<:NamedTuple{<:Any,<:Tuple{Any,LocalInertialOverlandFlow,LocalInertialRiverFlow}},V,R,W,T}
 
 Run surface routing (land and river) for a model type that contains the lateral components
 `LocalInertialOverlandFlow` and `LocalInertialRiverFlow` for a single timestep.
 """
 function surface_routing!(
-    model::Model{N, L},
+    model::Model{L},
 ) where {
-    N,
     L <: NamedTuple{<:Any, <:Tuple{Any, LocalInertialOverlandFlow, LocalInertialRiverFlow}},
 }
     (; lateral, vertical, network, clock) = model

src/sbm_gwf_model.jl

@@ -468,17 +468,27 @@ function initialize_sbm_gwf_model(config::Config)
             lake_indices = inds_lake_map2river,
             land_indices = inds_land_map2river,
             # specific for local-inertial
-            nodes_at_edge = nodes_at_edge,
-            edges_at_node = adjacent_edges_at_node(graph_river, nodes_at_edge),
+            nodes_at_edge = NodesAtEdge(nodes_at_edge...),
+            edges_at_node = EdgesAtNode(
+                adjacent_edges_at_node(graph_river, nodes_at_edge)...,
+            ),
             # water allocation areas
             allocation_area_indices = river_allocation_area_inds,
             cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     end
 
+    network = Network(;
+        land = NetworkLand(; land...),
+        river = NetworkRiver(; river...),
+        reservoir = NetworkReservoir(; reservoir_network...),
+        lake = NetworkLake(; lake_network...),
+        drain = NetworkDrain(; drain...),
+    )
+
     model = Model(
         config,
-        (; land, river, reservoir = reservoir_network, lake = lake_network, drain),
+        network,
         (subsurface = subsurface_map, land = overland_flow, river = river_flow),
         land_hydrology,
         clock,

src/sbm_model.jl

@@ -397,17 +397,26 @@ function initialize_sbm_model(config::Config)
             lake_indices = inds_lake_map2river,
             land_indices = inds_land_map2river,
             # specific for local-inertial
-            nodes_at_edge = nodes_at_edge,
-            edges_at_node = adjacent_edges_at_node(graph_river, nodes_at_edge),
+            nodes_at_edge = NodesAtEdge(; nodes_at_edge...),
+            edges_at_node = EdgesAtNode(;
+                adjacent_edges_at_node(graph_river, nodes_at_edge)...,
+            ),
             # water allocation areas
             allocation_area_indices = river_allocation_area_inds,
             cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     end
 
+    network = Network(;
+        land = NetworkLand(; land...),
+        river = NetworkRiver(; river...),
+        reservoir = NetworkReservoir(; reservoir_network...),
+        lake = NetworkLake(; lake_network...),
+    )
+
     model = Model(
         config,
-        (; land, river, reservoir = reservoir_network, lake = lake_network),
+        network,
         (subsurface = subsurface_flow, land = overland_flow, river = river_flow),
         land_hydrology,
         clock,
@@ -449,7 +458,7 @@ function update!(model::AbstractModel{<:SbmModel})
 end
 
 """
-    update_until_recharge!(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
+    update_until_recharge!model::AbstractModel{<:SbmModel})
 
 Update SBM model until recharge for a single timestep. This function is also accessible
 through BMI, to couple the SBM model to an external groundwater model.
@@ -462,7 +471,7 @@ function update_until_recharge!(model::AbstractModel{<:SbmModel})
 end
 
 """
-    update_after_subsurfaceflow!(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
+    update_after_subsurfaceflow!(model::AbstractModel{<:SbmModel})
 
 Update SBM model after subsurface flow for a single timestep. This function is also
 accessible through BMI, to couple the SBM model to an external groundwater model.

src/sediment_model.jl

@@ -98,7 +98,7 @@ function initialize_sediment_model(config::Config)
     clamp!(landslope, 0.00001, Inf)
 
     index_river = filter(i -> !isequal(river[i], 0), 1:n)
-    frac_toriver = fraction_runoff_to_river(graph, ldd, index_river, landslope)
+    frac_to_river = fraction_runoff_to_river(graph, ldd, index_river, landslope)
 
     river_sediment = RiverSediment(dataset, config, indices_riv, waterbodies)
 
@@ -133,9 +133,18 @@ function initialize_sediment_model(config::Config)
         reverse_indices = rev_indices_riv,
     )
 
+    network = Network(;
+        land = NetworkLand(; land...),
+        river = NetworkRiver(; river...),
+        reservoir = NetworkReservoir(; reservoir...),
+        lake = NetworkLake(; lake...),
+        index_river,
+        frac_to_river,
+    )
+
     model = Model(
         config,
-        (; land, river, reservoir, lake, index_river, frac_toriver),
+        network,
         (land = overland_flow_sediment, river = river_sediment),
         soilloss,
         clock,

src/utils.jl

@@ -77,7 +77,7 @@ function active_indices(subcatch_2d::AbstractMatrix, nodata)
     return indices, reverse_indices
 end
 
-function active_indices(network::NamedTuple, key::Tuple)
+function active_indices(network::Network, key::Tuple)
     if :reservoir in key
         return network.reservoir.indices_outlet
     elseif :lake in key