Repository: Deltares/Wflow.jl (language: python)

Issue:
Metadata without macros
### Feature type

Changing existing functionality

### Improvement Description

The metadata of struct fields is generated with macros. While fancy, it is not very readable/explicit (code), and it is only used for model coupling through BMI and thus not exposed to a wflow user (running a wflow model simulation). Implementing metadata without macros makes the code more readable/explicit.  

### Implementation Description

_No response_

### Additional Context

_No response_
Metadata without macros
### Feature type

Changing existing functionality

### Improvement Description

The metadata of struct fields is generated with macros. While fancy, it is not very readable/explicit (code), and it is only used for model coupling through BMI and thus not exposed to a wflow user (running a wflow model simulation). Implementing metadata without macros makes the code more readable/explicit.  

### Implementation Description

_No response_

### Additional Context

_No response_


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 868c5ee46..60504c78b 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - master
+      - v1
     tags: '*'
 jobs:
   test:
diff --git a/.github/workflows/CIWflowServer.yml b/.github/workflows/CIWflowServer.yml
index a051bcf52..126361940 100644
--- a/.github/workflows/CIWflowServer.yml
+++ b/.github/workflows/CIWflowServer.yml
@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - master
+      - v1
     tags: '*'
 jobs:
   test:
diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 6133e9600..9cc3d318e 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Removed vertical concepts `HBV` and `FLEXTopo`.
+- Removed metadata macros `exchange`, `grid_type` and `grid_location`. The macro `grid_type`
+  is not required because this functionality is already part of `BMI`. The macros `exchange`
+  and `grid_location` are replaced by functions resulting in cleaner code.
 
 ### Added
 
diff --git a/src/Wflow.jl b/src/Wflow.jl
index 199d2d0a1..651df3e21 100644
--- a/src/Wflow.jl
+++ b/src/Wflow.jl
@@ -23,10 +23,6 @@ using LoopVectorization
 using IfElse
 
 @metadata get_units "mm dt-1" String
-# metadata for BMI grid
-@metadata exchange 1 Integer
-@metadata grid_type "unstructured" String
-@metadata grid_location "node" String
 
 const BMI = BasicModelInterface
 const Float = Float64
diff --git a/src/bmi.jl b/src/bmi.jl
index 040a474b0..2ea9f0395 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -116,22 +116,23 @@ function BMI.get_input_var_names(model::Model)
         var_names = Vector{String}()
         for c in config.API.components
             type = typeof(param(model, c))
-            inds = findall(x -> x != 0, exchange(type))
-            field_names = fieldnames(type)[inds]
+            field_names = fieldnames(type)
 
             for name in field_names
-                var = string(c, ".", name)
-                model_var = param(model, var)
-                if eltype(model_var) <: SVector
-                    for i = 1:length(first(model_var))
-                        push!(var_names, string(var, "[", i, "]"))
+                if exchange(param(model, c), name) == 1
+                    var = string(c, ".", name)
+                    model_var = param(model, var)
+                    if eltype(model_var) <: SVector
+                        for i = 1:length(first(model_var))
+                            push!(var_names, string(var, "[", i, "]"))
+                        end
+                    elseif ndims(model_var) > 1
+                        for i = 1:length(first(model_var))
+                            push!(var_names, string(var, "[", i, "]"))
+                        end
+                    else
+                        push!(var_names, var)
                     end
-                elseif ndims(model_var) > 1
-                    for i = 1:length(first(model_var))
-                        push!(var_names, string(var, "[", i, "]"))
-                    end
-                else
-                    push!(var_names, var)
                 end
             end
         end
@@ -152,7 +153,6 @@ function BMI.get_var_grid(model::Model, name::String)
     s = split(name, "[")
     key = symbols(first(s))
     if exchange(param(model, key[1:end-1]), key[end]) == 1
-        gridtype = grid_type(param(model, key))
         type = typeof(param(model, key[1:end-1]))
         return if :reservoir in key
             0
@@ -199,8 +199,9 @@ end
 
 function BMI.get_var_location(model::Model, name::String)
     key = symbols(first(split(name, "[")))
-    if exchange(param(model, key[1:end-1]), key[end]) == 1
-        return grid_location(model, key)
+    type = param(model, key[1:end-1])
+    if exchange(type, key[end]) == 1
+        return grid_location(type, key[end])
     else
         error("$name not listed as variable for BMI exchange")
     end
@@ -409,3 +410,125 @@ end
 function get_start_unix_time(model::Model)
     datetime2unix(DateTime(model.config.starttime))
 end
+
+# Exchange and grid location functions.
+exchange(::Nothing, var) = 0
+grid_location(::Nothing, var) = "none"
+
+function exchange(::SurfaceFlow, var)
+    if var in (:dt, :beta, :its, :alpha_pow, :reservoir, :lake, :kinwave_it)
+        0
+    else
+        1
+    end
+end
+
+function grid_location(::SurfaceFlow, var)
+    if var in (:dt, :its, :kinwave_it)
+        "none"
+    else
+        "node"
+    end
+end
+
+exchange(::Union{LateralSSF,GroundwaterExchange}, var) = var == :dt ? 0 : 1
+grid_location(::Union{LateralSSF,GroundwaterExchange}, var) = var == :dt ? "none" : "node"
+
+function exchange(::ShallowWaterRiver, var)
+    if var in (
+        :dt,
+        :n,
+        :ne,
+        :g,
+        :alpha,
+        :h_thresh,
+        :froude_limit,
+        :reservoir_index,
+        :lake_index,
+        :reservoir,
+        :lake,
+        :floodplain,
+    )
+        0
+    else
+        1
+    end
+end
+
+function grid_location(::ShallowWaterRiver, var)
+    if var in (:n, :ne, :dt, :froude_limit)
+        "none"
+    elseif var in (
+        :active_e,
+        :q,
+        :q0,
+        :q_av,
+        :mannings_n_sq,
+        :zs_max,
+        :hf,
+        :dl_at_link,
+        :width_at_link,
+        :a,
+        :r,
+    )
+        "edge"
+    else
+        "node"
+    end
+end
+
+function exchange(::ShallowWaterLand, var)
+    if var in (:n, :g, :theta, :alpha, :h_thresh, :dt, :froude_limit)
+        0
+    else
+        1
+    end
+end
+
+function grid_location(::ShallowWaterLand, var)
+    if var in (:n, :dt, :froude_limit)
+        "none"
+    elseif var in (:xwidth, :ywidth, :qy0, :qx0, :qx, :qy, :zx_max, :zy_max, :mannings_n_sq)
+        "edge"
+    else
+        "node"
+    end
+end
+
+exchange(::FloodPlainProfile, var) = 1
+grid_location(::FloodPlainProfile, var) = "node"
+
+exchange(::FloodPlain, var) = var == :profile ? 0 : 1
+function grid_location(::FloodPlain, var)
+    if var in (:mannings_n_sq, :a, :r, :hf, :zb_max, :q0, :q, :q_av, :hf_index)
+        "edge"
+    else
+        "node"
+    end
+
+end
+
+exchange(::SimpleReservoir, var) = var == :dt ? 0 : 1
+grid_location(::SimpleReservoir, var) = var == :dt ? "none" : "node"
+
+exchange(::Lake, var) = var == :dt ? 0 : 1
+grid_location(::Lake, var) = var == :dt ? "none" : "node"
+
+exchange(::SBM, var) = var in (:n, :dt, :maxlayers) ? 0 : 1
+grid_location(::SBM, var) = var in (:n, :dt, :maxlayers) ? "none" : "node"
+
+exchange(::Union{LandSediment,OverlandFlowSediment}, var) = var == :n ? 0 : 1
+grid_location(::Union{LandSediment,OverlandFlowSediment}, var) = var == :n ? "none" : "node"
+
+exchange(::RiverSediment, var) = var in (:n, :dt) ? 0 : 1
+grid_location(::RiverSediment, var) = var in (:n, :dt) ? "none" : "node"
+
+
+exchange(::Aquifer, var) = 1
+grid_location(::Aquifer, var) = "node"
+
+exchange(::ConstantHead, var) = 0
+grid_location(::ConstantHead, var) = "node"
+
+exchange(::AquiferBoundaryCondition, var) = 1
+grid_location(::AquiferBoundaryCondition, var) = "node"
diff --git a/src/flow.jl b/src/flow.jl
index 77b058bf0..00e3475be 100644
--- a/src/flow.jl
+++ b/src/flow.jl
@@ -1,35 +1,34 @@
 
 abstract type SurfaceFlow end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowRiver{T,R,L} <:
-                                                               SurfaceFlow
-    beta::T | "-" | 0 | "scalar"                 # constant in Manning's equation
-    sl::Vector{T} | "m m-1"                      # Slope [m m⁻¹]
-    n::Vector{T} | "s m-1/3"                     # Manning's roughness [s m⁻⅓]
-    dl::Vector{T} | "m"                          # Drain length [m]
-    q::Vector{T} | "m3 s-1"                      # Discharge [m³ s⁻¹]
-    qin::Vector{T} | "m3 s-1"                    # Inflow from upstream cells [m³ s⁻¹]
-    q_av::Vector{T} | "m3 s-1"                   # Average discharge [m³ s⁻¹]
-    qlat::Vector{T} | "m2 s-1"                   # Lateral inflow per unit length [m² s⁻¹]
-    inwater::Vector{T} | "m3 s-1"                # Lateral inflow [m³ s⁻¹]
-    inflow::Vector{T} | "m3 s-1"                 # External inflow (abstraction/supply/demand) [m³ s⁻¹]
-    inflow_wb::Vector{T} | "m3 s-1"              # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
-    volume::Vector{T} | "m3"                     # Kinematic wave volume [m³] (based on water level h)
-    h::Vector{T} | "m"                           # Water level [m]
-    h_av::Vector{T} | "m"                        # Average water level [m]
-    bankfull_depth::Vector{T} | "m"              # Bankfull water level [m]
-    dt::T | "s" | 0 | "none" | "none"            # Model time step [s]
-    its::Int | "-" | 0 | "none" | "none"         # Number of fixed iterations
-    width::Vector{T} | "m"                       # Flow width [m]
-    alpha_pow::T | "-" | 0 | "scalar"            # Used in the power part of alpha
-    alpha_term::Vector{T} | "-"                  # Term used in computation of alpha
-    alpha::Vector{T} | "s3/5 m1/5"               # Constant in momentum equation A = alpha*Q^beta, based on Manning's equation
-    cel::Vector{T} | "m s-1"                     # Celerity of the kinematic wave
-    reservoir_index::Vector{Int} | "-"           # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
-    lake_index::Vector{Int} | "-"                # map cell to 0 (no lake) or i (pick lake i in lake field)
-    reservoir::R | "-" | 0                       # Reservoir model struct of arrays
-    lake::L | "-" | 0                            # Lake model struct of arrays
-    kinwave_it::Bool | "-" | 0 | "none" | "none" # Boolean for iterations kinematic wave
+@get_units @with_kw struct SurfaceFlowRiver{T,R,L} <: SurfaceFlow
+    beta::T | "-"                           # constant in Manning's equation
+    sl::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
+    n::Vector{T} | "s m-1/3"                # Manning's roughness [s m⁻⅓]
+    dl::Vector{T} | "m"                     # Drain length [m]
+    q::Vector{T} | "m3 s-1"                 # Discharge [m³ s⁻¹]
+    qin::Vector{T} | "m3 s-1"               # Inflow from upstream cells [m³ s⁻¹]
+    q_av::Vector{T} | "m3 s-1"              # Average discharge [m³ s⁻¹]
+    qlat::Vector{T} | "m2 s-1"              # Lateral inflow per unit length [m² s⁻¹]
+    inwater::Vector{T} | "m3 s-1"           # Lateral inflow [m³ s⁻¹]
+    inflow::Vector{T} | "m3 s-1"            # External inflow (abstraction/supply/demand) [m³ s⁻¹]
+    inflow_wb::Vector{T} | "m3 s-1"         # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
+    volume::Vector{T} | "m3"                # Kinematic wave volume [m³] (based on water level h)
+    h::Vector{T} | "m"                      # Water level [m]
+    h_av::Vector{T} | "m"                   # Average water level [m]
+    bankfull_depth::Vector{T} | "m"         # Bankfull water level [m]
+    dt::T | "s"                             # Model time step [s]
+    its::Int | "-"                          # Number of fixed iterations
+    width::Vector{T} | "m"                  # Flow width [m]
+    alpha_pow::T | "-"                      # Used in the power part of alpha
+    alpha_term::Vector{T} | "-"             # Term used in computation of alpha
+    alpha::Vector{T} | "s3/5 m1/5"          # Constant in momentum equation A = alpha*Q^beta, based on Manning's equation
+    cel::Vector{T} | "m s-1"                # Celerity of the kinematic wave
+    reservoir_index::Vector{Int} | "-"      # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
+    lake_index::Vector{Int} | "-"           # map cell to 0 (no lake) or i (pick lake i in lake field)
+    reservoir::R | "-"                      # Reservoir model struct of arrays
+    lake::L | "-"                           # Lake model struct of arrays
+    kinwave_it::Bool | "-"                  # Boolean for iterations kinematic wave
 
     # TODO unclear why this causes a MethodError
     # function SurfaceFlow{T,R,L}(args...) where {T,R,L}
@@ -38,29 +37,28 @@ abstract type SurfaceFlow end
     # end
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowLand{T} <:
-                                                               SurfaceFlow
-    beta::T | "-" | 0 | "scalar"                       # constant in Manning's equation
-    sl::Vector{T} | "m m-1"                         # Slope [m m⁻¹]
-    n::Vector{T} | "s m-1/3"                        # Manning's roughness [s m⁻⅓]
-    dl::Vector{T} | "m"                             # Drain length [m]
-    q::Vector{T} | "m3 s-1"                         # Discharge [m³ s⁻¹]
-    qin::Vector{T} | "m3 s-1"                       # Inflow from upstream cells [m³ s⁻¹]
-    q_av::Vector{T} | "m3 s-1"                      # Average discharge [m³ s⁻¹]
-    qlat::Vector{T} | "m2 s-1"                      # Lateral inflow per unit length [m² s⁻¹]
-    inwater::Vector{T} | "m3 s-1"                   # Lateral inflow [m³ s⁻¹]
-    volume::Vector{T} | "m3"                        # Kinematic wave volume [m³] (based on water level h)
-    h::Vector{T} | "m"                              # Water level [m]
-    h_av::Vector{T} | "m"                           # Average water level [m]
-    dt::T | "s" | 0 | "none" | "none"               # Model time step [s]
-    its::Int | "-" | 0 | "none" | "none"            # Number of fixed iterations
-    width::Vector{T} | "m"                          # Flow width [m]
-    alpha_pow::T | "-" | 0 | "scalar"               # Used in the power part of alpha
-    alpha_term::Vector{T} | "-"                     # Term used in computation of alpha
-    alpha::Vector{T} | "s3/5 m1/5"                  # Constant in momentum equation A = alpha * Q^beta, based on Manning's equation
-    cel::Vector{T} | "m s-1"                        # Celerity of the kinematic wave
-    to_river::Vector{T} | "m3 s-1"                  # Part of overland flow [m³ s⁻¹] that flows to the river
-    kinwave_it::Bool | "-" | 0 | "none" | "none"    # Boolean for iterations kinematic wave
+@get_units @with_kw struct SurfaceFlowLand{T} <: SurfaceFlow
+    beta::T | "-"                       # constant in Manning's equation
+    sl::Vector{T} | "m m-1"             # Slope [m m⁻¹]
+    n::Vector{T} | "s m-1/3"            # Manning's roughness [s m⁻⅓]
+    dl::Vector{T} | "m"                 # Drain length [m]
+    q::Vector{T} | "m3 s-1"             # Discharge [m³ s⁻¹]
+    qin::Vector{T} | "m3 s-1"           # Inflow from upstream cells [m³ s⁻¹]
+    q_av::Vector{T} | "m3 s-1"          # Average discharge [m³ s⁻¹]
+    qlat::Vector{T} | "m2 s-1"          # Lateral inflow per unit length [m² s⁻¹]
+    inwater::Vector{T} | "m3 s-1"       # Lateral inflow [m³ s⁻¹]
+    volume::Vector{T} | "m3"            # Kinematic wave volume [m³] (based on water level h)
+    h::Vector{T} | "m"                  # Water level [m]
+    h_av::Vector{T} | "m"               # Average water level [m]
+    dt::T | "s"                         # Model time step [s]
+    its::Int | "-"                      # Number of fixed iterations
+    width::Vector{T} | "m"              # Flow width [m]
+    alpha_pow::T | "-"                  # Used in the power part of alpha
+    alpha_term::Vector{T} | "-"         # Term used in computation of alpha
+    alpha::Vector{T} | "s3/5 m1/5"      # Constant in momentum equation A = alpha * Q^beta, based on Manning's equation
+    cel::Vector{T} | "m s-1"            # Celerity of the kinematic wave
+    to_river::Vector{T} | "m3 s-1"      # Part of overland flow [m³ s⁻¹] that flows to the river
+    kinwave_it::Bool | "-"              # Boolean for iterations kinematic wave
 end
 
 function initialize_surfaceflow_land(nc, config, inds; sl, dl, width, iterate, tstep, dt)
@@ -389,16 +387,16 @@ function stable_timestep(sf::S) where {S<:SurfaceFlow}
     return dt, its
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct LateralSSF{T}
-    kh_0::Vector{T} | "m d-1"               # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
+@get_units @with_kw struct LateralSSF{T}
+    kh_0::Vector{T} | "m d-1"              # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
     f::Vector{T} | "m-1"                   # A scaling parameter [m⁻¹] (controls exponential decline of kh_0)
     kh::Vector{T} | "m d-1"                # Horizontal hydraulic conductivity [m d⁻¹]
     khfrac::Vector{T} | "-"                # A muliplication factor applied to vertical hydraulic conductivity `kv` [-]
     soilthickness::Vector{T} | "m"         # Soil thickness [m]
     theta_s::Vector{T} | "-"               # Saturated water content (porosity) [-]
     theta_r::Vector{T} | "-"               # Residual water content [-]
-    dt::T | "d" | 0 | "none" | "none"      # model time step [d]
-    slope::Vector{T} | "m m-1"            # Slope [m m⁻¹]
+    dt::T | "d"                            # model time step [d]
+    slope::Vector{T} | "m m-1"             # Slope [m m⁻¹]
     dl::Vector{T} | "m"                    # Drain length [m]
     dw::Vector{T} | "m"                    # Flow width [m]
     zi::Vector{T} | "m"                    # Pseudo-water table depth [m] (top of the saturated zone)
@@ -480,57 +478,57 @@ function update(ssf::LateralSSF, network, frac_toriver, ksat_profile)
     end
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct GroundwaterExchange{T}
-    dt::T | "d" | 0 | "none" | "none"   # model time step [d]
+@get_units @with_kw struct GroundwaterExchange{T}
+    dt::T | "d"                         # model time step [d]
     exfiltwater::Vector{T} | "m dt-1"   # Exfiltration [m Δt⁻¹]  (groundwater above surface level, saturated excess conditions)
     zi::Vector{T} | "m"                 # Pseudo-water table depth [m] (top of the saturated zone)
     to_river::Vector{T} | "m3 d-1"      # Part of subsurface flow [m³ d⁻¹] that flows to the river
     ssf::Vector{T} | "m3 d-1"           # Subsurface flow [m³ d⁻¹]
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct ShallowWaterRiver{T,R,L,F}
-    n::Int | "-" | 0 | "none" | "none"                      # number of cells
-    ne::Int | "-" | 0 | "none" | "none"                     # number of edges/links
-    active_n::Vector{Int} | "-"                             # active nodes
-    active_e::Vector{Int} | "-" | _ | "edge"                # active edges/links
-    g::T | "m s-2" | 0 | "scalar"                           # acceleration due to gravity
-    alpha::T | "-" | 0 | "scalar"                           # stability coefficient (Bates et al., 2010)
-    h_thresh::T | "m" | 0 | "scalar"                        # depth threshold for calculating flow
-    dt::T | "s" | 0 | "none" | "none"                       # model time step [s]
-    q::Vector{T} | "m3 s-1" | _ | "edge"                    # river discharge (subgrid channel)
-    q0::Vector{T} | "m3 s-1" | _ | "edge"                   # river discharge (subgrid channel) at previous time step
-    q_av::Vector{T} | "m3 s-1" | _ | "edge"                 # average river channel (+ floodplain) discharge [m³ s⁻¹]
-    q_channel_av::Vector{T} | "m3 s-1"                      # average river channel discharge [m³ s⁻¹]
-    zb_max::Vector{T} | "m"                                 # maximum channel bed elevation
-    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # Manning's roughness squared at edge/link
-    mannings_n::Vector{T} | "s m-1/3"                       # Manning's roughness at node
-    h::Vector{T} | "m"                                      # water depth
-    zs_max::Vector{T} | "m" | _ | "edge"                    # maximum water elevation at edge
-    zs_src::Vector{T} | "m"                                 # water elevation of source node of edge
-    zs_dst::Vector{T} | "m"                                 # water elevation of downstream node of edge
-    hf::Vector{T} | "m" | _ | "edge"                        # water depth at edge/link
-    h_av::Vector{T} | "m"                                   # average water depth
-    dl::Vector{T} | "m"                                     # river length
-    dl_at_link::Vector{T} | "m" | _ | "edge"                # river length at edge/link
-    width::Vector{T} | "m"                                  # river width
-    width_at_link::Vector{T} | "m" | _ | "edge"             # river width at edge/link
-    a::Vector{T} | "m2" | _ | "edge"                        # flow area at edge/link
-    r::Vector{T} | "m" | _ | "edge"                         # wetted perimeter at edge/link
-    volume::Vector{T} | "m3"                                # river volume
-    error::Vector{T} | "m3"                                 # error volume
-    inwater::Vector{T} | "m3 s-1"                           # lateral inflow [m³ s⁻¹]
-    inflow::Vector{T} | "m3 s-1"                            # external inflow (abstraction/supply/demand) [m³ s⁻¹]
-    inflow_wb::Vector{T} | "m3 s-1"                         # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
-    bankfull_volume::Vector{T} | "m3"                       # bankfull volume
-    bankfull_depth::Vector{T} | "m"                         # bankfull depth
-    zb::Vector{T} | "m"                                     # river bed elevation
-    froude_limit::Bool | "-" | 0 | "none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
-    reservoir_index::Vector{Int} | "-" | 0                  # river cell index with a reservoir (each index of reservoir_index maps to reservoir i in reservoir field)
-    lake_index::Vector{Int} | "-" | 0                       # river cell index with a lake (each index of lake_index maps to lake i in lake field)
-    waterbody::Vector{Bool} | "-"                           # water body cells (reservoir or lake)
-    reservoir::R | "-" | 0                                  # Reservoir model struct of arrays
-    lake::L | "-" | 0                                       # Lake model struct of arrays
-    floodplain::F | "-" | 0                                 # Floodplain (1D) schematization
+@get_units @with_kw struct ShallowWaterRiver{T,R,L,F}
+    n::Int | "-"                                # number of cells
+    ne::Int | "-"                               # number of edges/links
+    active_n::Vector{Int} | "-"                 # active nodes
+    active_e::Vector{Int} | "-"                 # active edges/links
+    g::T | "m s-2"                              # acceleration due to gravity
+    alpha::T | "-"                              # stability coefficient (Bates et al., 2010)
+    h_thresh::T | "m"                           # depth threshold for calculating flow
+    dt::T | "s"                                 # model time step [s]
+    q::Vector{T} | "m3 s-1"                     # river discharge (subgrid channel)
+    q0::Vector{T} | "m3 s-1"                    # river discharge (subgrid channel) at previous time step
+    q_av::Vector{T} | "m3 s-1"                  # average river channel (+ floodplain) discharge [m³ s⁻¹]
+    q_channel_av::Vector{T} | "m3 s-1"          # average river channel discharge [m³ s⁻¹]
+    zb_max::Vector{T} | "m"                     # maximum channel bed elevation
+    mannings_n_sq::Vector{T} | "(s m-1/3)2"     # Manning's roughness squared at edge/link
+    mannings_n::Vector{T} | "s m-1/3"           # Manning's roughness at node
+    h::Vector{T} | "m"                          # water depth
+    zs_max::Vector{T} | "m"                     # maximum water elevation at edge
+    zs_src::Vector{T} | "m"                     # water elevation of source node of edge
+    zs_dst::Vector{T} | "m"                     # water elevation of downstream node of edge
+    hf::Vector{T} | "m"                         # water depth at edge/link
+    h_av::Vector{T} | "m"                       # average water depth
+    dl::Vector{T} | "m"                         # river length
+    dl_at_link::Vector{T} | "m"                 # river length at edge/link
+    width::Vector{T} | "m"                      # river width
+    width_at_link::Vector{T} | "m"              # river width at edge/link
+    a::Vector{T} | "m2"                         # flow area at edge/link
+    r::Vector{T} | "m"                          # wetted perimeter at edge/link
+    volume::Vector{T} | "m3"                    # river volume
+    error::Vector{T} | "m3"                     # error volume
+    inwater::Vector{T} | "m3 s-1"               # lateral inflow [m³ s⁻¹]
+    inflow::Vector{T} | "m3 s-1"                # external inflow (abstraction/supply/demand) [m³ s⁻¹]
+    inflow_wb::Vector{T} | "m3 s-1"             # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
+    bankfull_volume::Vector{T} | "m3"           # bankfull volume
+    bankfull_depth::Vector{T} | "m"             # bankfull depth
+    zb::Vector{T} | "m"                         # river bed elevation
+    froude_limit::Bool | "-"                    # if true a check is performed if froude number > 1.0 (algorithm is modified)
+    reservoir_index::Vector{Int} | "-"          # river cell index with a reservoir (each index of reservoir_index maps to reservoir i in reservoir field)
+    lake_index::Vector{Int} | "-"               # river cell index with a lake (each index of lake_index maps to lake i in lake field)
+    waterbody::Vector{Bool} | "-"               # water body cells (reservoir or lake)
+    reservoir::R | "-"                          # Reservoir model struct of arrays
+    lake::L | "-"                               # Lake model struct of arrays
+    floodplain::F | "-"                         # Floodplain (1D) schematization
 end
 
 function initialize_shallowwater_river(
@@ -983,33 +981,33 @@ end
 # neigbors.
 const dirs = (:yd, :xd, :xu, :yu)
 
-@get_units @exchange @grid_type @grid_location @with_kw struct ShallowWaterLand{T}
-    n::Int | "-" | 0 | "none" | "none"                      # number of cells
-    xl::Vector{T} | "m"                                     # cell length x direction
-    yl::Vector{T} | "m"                                     # cell length y direction
-    xwidth::Vector{T} | "m" | _ | "edge"                    # effective flow width x direction (floodplain)
-    ywidth::Vector{T} | "m" | _ | "edge"                    # effective flow width y direction (floodplain)
-    g::T | "m2 s-1" | 0 | "scalar"                          # acceleration due to gravity
-    theta::T | "-" | 0 | "scalar"                           # weighting factor (de Almeida et al., 2012)
-    alpha::T | "-" | 0 | "scalar"                           # stability coefficient (de Almeida et al., 2012)
-    h_thresh::T | "m" | 0 | "scalar"                        # depth threshold for calculating flow
-    dt::T | "s" | 0 | "none" | "none"                       # model time step [s]
-    qy0::Vector{T} | "m3 s-1" | _ | "edge"                  # flow in y direction at previous time step
-    qx0::Vector{T} | "m3 s-1" | _ | "edge"                  # flow in x direction at previous time step
-    qx::Vector{T} | "m3 s-1" | _ | "edge"                   # flow in x direction
-    qy::Vector{T} | "m3 s-1" | _ | "edge"                   # flow in y direction
-    zx_max::Vector{T} | "m" | _ | "edge"                    # maximum cell elevation (x direction)
-    zy_max::Vector{T} | "m" | _ | "edge"                    # maximum cell elevation (y direction)
-    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # Manning's roughness squared
-    volume::Vector{T} | "m3"                                # total volume of cell (including river volume for river cells)
-    error::Vector{T} | "m3"                                 # error volume
-    runoff::Vector{T} | "m3 s-1"                            # runoff from hydrological model
-    inflow_wb::Vector{T} | "m3 s-1"                         # inflow to water body from hydrological model
-    h::Vector{T} | "m"                                      # water depth of cell (for river cells the reference is the river bed elevation `zb`)
-    z::Vector{T} | "m"                                      # elevation of cell
-    froude_limit::Bool | "-" | 0 | "none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
-    rivercells::Vector{Bool} | "-"                          # river cells
-    h_av::Vector{T} | "m"                                   # average water depth (for river cells the reference is the river bed elevation `zb`)
+@get_units @with_kw struct ShallowWaterLand{T}
+    n::Int | "-"                                # number of cells
+    xl::Vector{T} | "m"                         # cell length x direction
+    yl::Vector{T} | "m"                         # cell length y direction
+    xwidth::Vector{T} | "m"                     # effective flow width x direction (floodplain)
+    ywidth::Vector{T} | "m"                     # effective flow width y direction (floodplain)
+    g::T | "m2 s-1"                             # acceleration due to gravity
+    theta::T | "-"                              # weighting factor (de Almeida et al., 2012)
+    alpha::T | "-"                              # stability coefficient (de Almeida et al., 2012)
+    h_thresh::T | "m"                           # depth threshold for calculating flow
+    dt::T | "s"                                 # model time step [s]
+    qy0::Vector{T} | "m3 s-1"                   # flow in y direction at previous time step
+    qx0::Vector{T} | "m3 s-1"                   # flow in x direction at previous time step
+    qx::Vector{T} | "m3 s-1"                    # flow in x direction
+    qy::Vector{T} | "m3 s-1"                    # flow in y direction
+    zx_max::Vector{T} | "m"                     # maximum cell elevation (x direction)
+    zy_max::Vector{T} | "m"                     # maximum cell elevation (y direction)
+    mannings_n_sq::Vector{T} | "(s m-1/3)2"     # Manning's roughness squared
+    volume::Vector{T} | "m3"                    # total volume of cell (including river volume for river cells)
+    error::Vector{T} | "m3"                     # error volume
+    runoff::Vector{T} | "m3 s-1"                # runoff from hydrological model
+    inflow_wb::Vector{T} | "m3 s-1"             # inflow to water body from hydrological model
+    h::Vector{T} | "m"                          # water depth of cell (for river cells the reference is the river bed elevation `zb`)
+    z::Vector{T} | "m"                          # elevation of cell
+    froude_limit::Bool | "-"                    # if true a check is performed if froude number > 1.0 (algorithm is modified)
+    rivercells::Vector{Bool} | "-"              # river cells
+    h_av::Vector{T} | "m"                       # average water depth (for river cells the reference is the river bed elevation `zb`)
 end
 
 function initialize_shallowwater_land(
@@ -1366,30 +1364,30 @@ Floodplain `volume` is a function of `depth` (flood depth intervals). Based on t
 cumulative floodplain `volume` a floodplain profile as a function of `flood_depth` is
 derived with floodplain area `a` (cumulative) and wetted perimeter radius `p` (cumulative).
 """
-@get_units @exchange @grid_type @grid_location @with_kw struct FloodPlainProfile{T,N}
-    depth::Vector{T} | "m" | 0                # Flood depth
+@get_units @with_kw struct FloodPlainProfile{T,N}
+    depth::Vector{T} | "m"                    # Flood depth
     volume::Array{T,2} | "m3"                 # Flood volume (cumulative)
     width::Array{T,2} | "m"                   # Flood width
     a::Array{T,2} | "m2"                      # Flow area (cumulative)
     p::Array{T,2} | "m"                       # Wetted perimeter (cumulative)
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct FloodPlain{T,P}
-    profile::P | "-" | 0                                    # floodplain profile
-    mannings_n::Vector{T} | "s m-1/3"                       # manning's roughness
-    mannings_n_sq::Vector{T} | "(s m-1/3)2" | _ | "edge"    # manning's roughness squared
-    volume::Vector{T} | "m3"                                # volume
-    h::Vector{T} | "m"                                      # water depth
-    h_av::Vector{T} | "m"                                   # average water depth
-    error::Vector{T} | "m3"                                 # error volume
-    a::Vector{T} | "m2" | _ | "edge"                        # flow area
-    r::Vector{T} | "m" | _ | "edge"                         # hydraulic radius
-    hf::Vector{T} | "m" | _ | "edge"                        # water depth at edge/link
-    zb_max::Vector{T} | "m" | _ | "edge"                    # maximum bankfull elevation (edge/link)
-    q0::Vector{T} | "m3 s-1" | _ | "edge"                   # discharge at previous time step
-    q::Vector{T} | "m3 s-1" | _ | "edge"                    # discharge
-    q_av::Vector{T} | "m" | _ | "edge"                      # average river discharge
-    hf_index::Vector{Int} | "-" | _ | "edge"                # index with `hf` above depth threshold
+@get_units @with_kw struct FloodPlain{T,P}
+    profile::P | "-"                           # floodplain profile
+    mannings_n::Vector{T} | "s m-1/3"          # manning's roughness
+    mannings_n_sq::Vector{T} | "(s m-1/3)2"    # manning's roughness squared
+    volume::Vector{T} | "m3"                   # volume
+    h::Vector{T} | "m"                         # water depth
+    h_av::Vector{T} | "m"                      # average water depth
+    error::Vector{T} | "m3"                    # error volume
+    a::Vector{T} | "m2"                        # flow area
+    r::Vector{T} | "m"                         # hydraulic radius
+    hf::Vector{T} | "m"                        # water depth at edge/link
+    zb_max::Vector{T} | "m"                    # maximum bankfull elevation (edge/link)
+    q0::Vector{T} | "m3 s-1"                   # discharge at previous time step
+    q::Vector{T} | "m3 s-1"                    # discharge
+    q_av::Vector{T} | "m"                      # average river discharge
+    hf_index::Vector{Int} | "-"                # index with `hf` above depth threshold
 end
 
 "Determine the initial floodplain volume"
diff --git a/src/groundwater/aquifer.jl b/src/groundwater/aquifer.jl
index d650bf799..34a59439c 100644
--- a/src/groundwater/aquifer.jl
+++ b/src/groundwater/aquifer.jl
@@ -86,7 +86,7 @@ NOTA BENE: **specific** storage is per m of aquifer (conf. specific weight).
 **Storativity** or (**storage coefficient**) is for the entire aquifer (conf.
 transmissivity).
 """
-@get_units @exchange @grid_type @grid_location struct ConfinedAquifer{T} <: Aquifer
+@get_units struct ConfinedAquifer{T} <: Aquifer
     head::Vector{T} | "m"  # hydraulic head [m]
     k::Vector{T} | "m d-1"  # horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
@@ -109,7 +109,7 @@ aquifer will yield when all water drains and the pore volume is filled by air
 instead. Specific yield will vary roughly between 0.05 (clay) and 0.45 (peat)
 (Johnson, 1967).
 """
-@get_units @exchange @grid_type @grid_location struct UnconfinedAquifer{T} <: Aquifer
+@get_units struct UnconfinedAquifer{T} <: Aquifer
     head::Vector{T} | "m"  # hydraulic head [m]
     k::Vector{T} | "m d-1"  # reference horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
@@ -122,7 +122,6 @@ instead. Specific yield will vary roughly between 0.05 (clay) and 0.45 (peat)
     # conductivity_profile is set to "exponential")
 end
 
-
 storativity(A::UnconfinedAquifer) = A.specific_yield
 storativity(A::ConfinedAquifer) = A.storativity
 
@@ -306,13 +305,11 @@ function flux!(Q, aquifer, connectivity, conductivity_profile)
     return Q
 end
 
-
-@get_units @exchange @grid_type @grid_location struct ConstantHead{T}
+@get_units struct ConstantHead{T}
     head::Vector{T} | "m"
     index::Vector{Int} | "-"
 end
 
-
 """
     stable_timestep(aquifer)
 
diff --git a/src/groundwater/boundary_conditions.jl b/src/groundwater/boundary_conditions.jl
index 06a04897e..39de1226f 100644
--- a/src/groundwater/boundary_conditions.jl
+++ b/src/groundwater/boundary_conditions.jl
@@ -13,7 +13,7 @@ end
 check_flux(flux, aquifer::ConfinedAquifer, index::Int) = flux
 
 
-@get_units @exchange @grid_type @grid_location struct River{T} <: AquiferBoundaryCondition
+@get_units struct River{T} <: AquiferBoundaryCondition
     stage::Vector{T} | "m"
     infiltration_conductance::Vector{T} | "m2 d-1"
     exfiltration_conductance::Vector{T} | "m2 d-1"
@@ -41,8 +41,7 @@ function flux!(Q, river::River, aquifer)
 end
 
 
-@get_units @exchange @grid_type @grid_location struct Drainage{T} <:
-                                                      AquiferBoundaryCondition
+@get_units struct Drainage{T} <: AquiferBoundaryCondition
     elevation::Vector{T} | "m"
     conductance::Vector{T} | "m2 d-1"
     flux::Vector{T} | "m3 d-1"
@@ -80,8 +79,7 @@ function flux!(Q, headboundary::HeadBoundary, aquifer)
 end
 
 
-@get_units @exchange @grid_type @grid_location struct Recharge{T} <:
-                                                      AquiferBoundaryCondition
+@get_units struct Recharge{T} <: AquiferBoundaryCondition
     rate::Vector{T} | "m d-1"
     flux::Vector{T} | "m3 d-1"
     index::Vector{Int} | "-"
diff --git a/src/reservoir_lake.jl b/src/reservoir_lake.jl
index 0b6acf768..b0ba2611c 100644
--- a/src/reservoir_lake.jl
+++ b/src/reservoir_lake.jl
@@ -1,5 +1,5 @@
-@get_units @exchange @grid_type @grid_location @with_kw struct SimpleReservoir{T}
-    dt::T | "s" | 0 | "none" | "none"                   # Model time step [s]
+@get_units @with_kw struct SimpleReservoir{T}
+    dt::T | "s"                                         # Model time step [s]
     maxvolume::Vector{T} | "m3"                         # maximum storage (above which water is spilled) [m³]
     area::Vector{T} | "m2"                              # reservoir area [m²]
     maxrelease::Vector{T} | "m3 s-1"                    # maximum amount that can be released if below spillway [m³ s⁻¹]
@@ -22,7 +22,6 @@
     end
 end
 
-
 function initialize_simple_reservoir(config, nc, inds_riv, nriv, pits, dt)
     # read only reservoir data if reservoirs true
     # allow reservoirs only in river cells
@@ -204,8 +203,8 @@ function update(res::SimpleReservoir, i, inflow, timestepsecs)
     return res
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct Lake{T}
-    dt::T | "s" | 0 | "none" | "none"           # Model time step [s]
+@get_units @with_kw struct Lake{T}
+    dt::T | "s"                                 # Model time step [s]
     lowerlake_ind::Vector{Int} | "-"            # Index of lower lake (linked lakes)
     area::Vector{T} | "m2"                      # lake area [m²]
     maxstorage::Vector{Union{T,Missing}} | "m3" # lake maximum storage from rating curve 1 [m³]
diff --git a/src/sbm.jl b/src/sbm.jl
index 352ab5c08..7e2676a16 100644
--- a/src/sbm.jl
+++ b/src/sbm.jl
@@ -1,10 +1,10 @@
-@get_units @exchange @grid_type @grid_location @with_kw struct SBM{T,N,M}
+@get_units @with_kw struct SBM{T,N,M}
     # Model time step [s]
-    dt::T | "s" | 0 | "none" | "none"
+    dt::T | "s"
     # Maximum number of soil layers
-    maxlayers::Int | "-" | 0 | "none" | "none"
+    maxlayers::Int | "-"
     # number of cells
-    n::Int | "-" | 0 | "none" | "none"
+    n::Int | "-"
     # Number of soil layers
     nlayers::Vector{Int} | "-"
     # Number of unsaturated soil layers
@@ -213,7 +213,6 @@
     end
 end
 
-
 function initialize_canopy(nc, config, inds)
     n = length(inds)
     # if leaf area index climatology provided use sl, swood and kext to calculate cmax, e_r and canopygapfraction
@@ -348,31 +347,13 @@ function initialize_sbm(nc, config, riverfrac, inds)
         fill = 0.0,
     )
     # soil parameters
-    theta_s = ncread(
-        nc,
-        config,
-        "vertical.theta_s";
-        sel = inds,
-        defaults = 0.6,
-        type = Float,
-    )
-    theta_r = ncread(
-        nc,
-        config,
-        "vertical.theta_r";
-        sel = inds,
-        defaults = 0.01,
-        type = Float,
-    )
+    theta_s =
+        ncread(nc, config, "vertical.theta_s"; sel = inds, defaults = 0.6, type = Float)
+    theta_r =
+        ncread(nc, config, "vertical.theta_r"; sel = inds, defaults = 0.01, type = Float)
     kv_0 =
-        ncread(
-            nc,
-            config,
-            "vertical.kv_0";
-            sel = inds,
-            defaults = 3000.0,
-            type = Float,
-        ) .* (dt / basetimestep)
+        ncread(nc, config, "vertical.kv_0"; sel = inds, defaults = 3000.0, type = Float) .*
+        (dt / basetimestep)
     f = ncread(nc, config, "vertical.f"; sel = inds, defaults = 0.001, type = Float)
     hb = ncread(nc, config, "vertical.hb"; sel = inds, defaults = 10.0, type = Float)
     soilthickness = ncread(
@@ -1002,8 +983,10 @@ function update_until_recharge(sbm::SBM, config)
 
             netcapflux = capflux
             for k = n_usl:-1:1
-                toadd =
-                    min(netcapflux, max(usl[k] * (sbm.theta_s[i] - sbm.theta_r[i]) - usld[k], 0.0))
+                toadd = min(
+                    netcapflux,
+                    max(usl[k] * (sbm.theta_s[i] - sbm.theta_r[i]) - usld[k], 0.0),
+                )
                 usld = setindex(usld, usld[k] + toadd, k)
                 netcapflux = netcapflux - toadd
                 actcapflux = actcapflux + toadd
@@ -1097,7 +1080,10 @@ function update_after_subsurfaceflow(sbm::SBM, zi, exfiltsatwater)
             if k <= n_usl
                 vwc = setindex(
                     vwc,
-                    (usld[k] + (sbm.act_thickl[i][k] - usl[k]) * (sbm.theta_s[i] - sbm.theta_r[i])) / sbm.act_thickl[i][k] + sbm.theta_r[i],
+                    (
+                        usld[k] +
+                        (sbm.act_thickl[i][k] - usl[k]) * (sbm.theta_s[i] - sbm.theta_r[i])
+                    ) / sbm.act_thickl[i][k] + sbm.theta_r[i],
                     k,
                 )
             else
@@ -1114,7 +1100,8 @@ function update_after_subsurfaceflow(sbm::SBM, zi, exfiltsatwater)
                 usld[k]
         end
 
-        rootstore_sat = max(0.0, sbm.rootingdepth[i] - zi[i]) * (sbm.theta_s[i] - sbm.theta_r[i])
+        rootstore_sat =
+            max(0.0, sbm.rootingdepth[i] - zi[i]) * (sbm.theta_s[i] - sbm.theta_r[i])
         rootstore = rootstore_sat + rootstore_unsat
         vwc_root = rootstore / sbm.rootingdepth[i] + sbm.theta_r[i]
         vwc_percroot = (vwc_root / sbm.theta_s[i]) * 100.0
diff --git a/src/sediment.jl b/src/sediment.jl
index dd708a32a..4d6edd153 100644
--- a/src/sediment.jl
+++ b/src/sediment.jl
@@ -1,7 +1,7 @@
 ### Soil erosion ###
-@get_units @exchange @grid_type @grid_location @with_kw struct LandSediment{T}
+@get_units @with_kw struct LandSediment{T}
     # number of cells
-    n::Int | "-" | 0 | "none" | "none"
+    n::Int | "-"
     ### Soil erosion part ###
     # length of cells in y direction [m]
     yl::Vector{T} | "m"
@@ -591,9 +591,9 @@ function tc_yalinpart(ols::LandSediment, i::Int, sinslope::Float, ts::Float)
 end
 
 ### Sediment transport in overland flow ###
-@get_units @exchange @grid_type @grid_location @with_kw struct OverlandFlowSediment{T}
+@get_units @with_kw struct OverlandFlowSediment{T}
     # number of cells
-    n::Int | "-" | 0 | "none" | "none"
+    n::Int | "-"
     # Filter with river cells
     rivcell::Vector{T} | "-"
     # Total eroded soil [ton Δt⁻¹]
@@ -634,7 +634,6 @@ end
     end
 end
 
-
 function partial_update!(inland, rivcell, eroded)
     no_erosion = zero(eltype(eroded))
     for i in eachindex(inland)
@@ -678,11 +677,11 @@ function update(ols::OverlandFlowSediment, network, config)
 end
 
 ### River transport and processes ###
-@get_units @exchange @grid_type @grid_location @with_kw struct RiverSediment{T}
+@get_units @with_kw struct RiverSediment{T}
     # number of cells
-    n::Int | "-" | 0 | "none" | "none"
+    n::Int | "-"
     # Timestep [s]
-    dt::T | "s" | 0 | "none" | "none"
+    dt::T | "s"
     # River geometry (slope [-], length [m], width [m])
     sl::Vector{T} | "m"
     dl::Vector{T} | "m"
@@ -782,7 +781,6 @@ end
     # end
 end
 
-
 function initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
     # Initialize river parameters
     nriv = length(inds_riv)


Test patch:
diff --git a/test/reservoir_lake.jl b/test/reservoir_lake.jl
index bc2446f92..51201e337 100644
--- a/test/reservoir_lake.jl
+++ b/test/reservoir_lake.jl
@@ -1,23 +1,22 @@
-@testset "reservoir simple" begin
-    res = Wflow.SimpleReservoir{Float64}(
-        dt = 86400.0,
-        demand = [52.523],
-        maxrelease = [420.184],
-        maxvolume = [25_000_000.0],
-        volume = [1.925e7],
-        totaloutflow = [0.0],
-        inflow = [0.0],
-        area = [1885665.353626924],
-        targetfullfrac = [0.8],
-        targetminfrac = [0.2425554726620697],
-        precipitation = [4.2],
-        evaporation = [1.5],
-        actevap = [0.0],
-        outflow = [NaN],
-        percfull = [NaN],
-        demandrelease = [NaN],
-    )
-
+res = Wflow.SimpleReservoir{Float64}(
+    dt = 86400.0,
+    demand = [52.523],
+    maxrelease = [420.184],
+    maxvolume = [25_000_000.0],
+    volume = [1.925e7],
+    totaloutflow = [0.0],
+    inflow = [0.0],
+    area = [1885665.353626924],
+    targetfullfrac = [0.8],
+    targetminfrac = [0.2425554726620697],
+    precipitation = [4.2],
+    evaporation = [1.5],
+    actevap = [0.0],
+    outflow = [NaN],
+    percfull = [NaN],
+    demandrelease = [NaN],
+)
+@testset "Update reservoir simple" begin
     Wflow.update(res, 1, 100.0, 86400.0)
     @test res.outflow[1] ≈ 91.3783714867453
     @test res.totaloutflow[1] ≈ 7.895091296454794e6
@@ -29,29 +28,37 @@
     @test res.actevap[1] ≈ 1.5
 end
 
-@testset "lake" begin
-    lake = Wflow.Lake{Float64}(
-        dt = 86400.0,
-        lowerlake_ind = [0],
-        area = [180510409.0],
-        maxstorage = Wflow.maximum_storage([1], [3], [180510409.0], [missing], [missing]),
-        threshold = [0.0],
-        storfunc = [1],
-        outflowfunc = [3],
-        totaloutflow = [0.0],
-        inflow = [0.0],
-        b = [0.22],
-        e = [2.0],
-        sh = [missing],
-        hq = [missing],
-        storage = Wflow.initialize_storage([1], [180510409.0], [18.5], [missing]),
-        waterlevel = [18.5],
-        precipitation = [20.0],
-        evaporation = [3.2],
-        actevap = [0.0],
-        outflow = [NaN],
-    )
+@testset "Exchange and grid location reservoir" begin
+    @test Wflow.exchange(res, :dt) == 0
+    @test Wflow.exchange(res, :volume) == 1
+    @test Wflow.exchange(res, :outflow) == 1
+    @test Wflow.grid_location(res, :dt) == "none"
+    @test Wflow.grid_location(res, :volume) == "node"
+    @test Wflow.grid_location(res, :outflow) == "node"
+end
 
+lake = Wflow.Lake{Float64}(
+    dt = 86400.0,
+    lowerlake_ind = [0],
+    area = [180510409.0],
+    maxstorage = Wflow.maximum_storage([1], [3], [180510409.0], [missing], [missing]),
+    threshold = [0.0],
+    storfunc = [1],
+    outflowfunc = [3],
+    totaloutflow = [0.0],
+    inflow = [0.0],
+    b = [0.22],
+    e = [2.0],
+    sh = [missing],
+    hq = [missing],
+    storage = Wflow.initialize_storage([1], [180510409.0], [18.5], [missing]),
+    waterlevel = [18.5],
+    precipitation = [20.0],
+    evaporation = [3.2],
+    actevap = [0.0],
+    outflow = [NaN],
+)
+@testset "Update lake" begin
     Wflow.update(lake, 1, 2500.0, 181, 86400.0)
     @test lake.outflow[1] ≈ 85.14292808113598
     @test lake.totaloutflow[1] ≈ 7.356348986210149e6
@@ -62,6 +69,15 @@ end
     @test lake.actevap[1] ≈ 3.2
 end
 
+@testset "Exchange and grid location lake" begin
+    @test Wflow.exchange(lake, :dt) == 0
+    @test Wflow.exchange(lake, :storage) == 1
+    @test Wflow.exchange(lake, :outflow) == 1
+    @test Wflow.grid_location(lake, :dt) == "none"
+    @test Wflow.grid_location(lake, :storage) == "node"
+    @test Wflow.grid_location(lake, :outflow) == "node"
+end
+
 datadir = joinpath(@__DIR__, "data")
 sh = [
     Wflow.read_sh_csv(joinpath(datadir, "input", "lake_sh_1.csv")),
diff --git a/test/run_sbm.jl b/test/run_sbm.jl
index 62fdfba99..76db24547 100644
--- a/test/run_sbm.jl
+++ b/test/run_sbm.jl
@@ -138,6 +138,42 @@ end
     @test res.evaporation[1] ≈ 0.5400000810623169f0
 end
 
+@testset "Exchange and grid location SBM" begin
+    sbm = model.vertical
+    @test Wflow.exchange(sbm, :n) == 0
+    @test Wflow.exchange(sbm, :rootingdepth) == 1
+    @test Wflow.grid_location(sbm, :dt) == "none"
+    @test Wflow.grid_location(sbm, :rootingdepth) == "node"
+    @test Wflow.grid_location(sbm, :runoff) == "node"
+end
+
+@testset "Exchange and grid location subsurface flow" begin
+    ssf = model.lateral.subsurface
+    @test Wflow.exchange(ssf, :dt) == 0
+    @test Wflow.exchange(ssf, :ssf) == 1
+    @test Wflow.exchange(ssf, :slope) == 1
+    @test Wflow.grid_location(ssf, :dt) == "none"
+    @test Wflow.grid_location(ssf, :ssf) == "node"
+    @test Wflow.grid_location(ssf, :slope) == "node"
+end
+
+@testset "Exchange and grid location kinematic wave" begin
+    land = model.lateral.land
+    @test Wflow.exchange(land, :dt) == 0
+    @test Wflow.exchange(land, :q_av) == 1
+    @test Wflow.exchange(land, :inwater) == 1
+    @test Wflow.grid_location(land, :dt) == "none"
+    @test Wflow.grid_location(land, :q_av) == "node"
+    @test Wflow.grid_location(land, :inwater) == "node"
+    river = model.lateral.river
+    @test Wflow.exchange(river, :dt) == 0
+    @test Wflow.exchange(river, :q_av) == 1
+    @test Wflow.exchange(river, :inwater) == 1
+    @test Wflow.grid_location(river, :dt) == "none"
+    @test Wflow.grid_location(river, :q_av) == "node"
+    @test Wflow.grid_location(river, :inwater) == "node"
+end
+
 # set these variables for comparison in "changed dynamic parameters"
 precip = copy(model.vertical.precipitation)
 evap = copy(model.vertical.potential_evaporation)
@@ -287,6 +323,19 @@ model = Wflow.run_timestep(model)
     @test h[[26, 35, 631]] ≈
           [0.07361854999908582f0, 0.009155393111676267f0, 0.0007258741013439351f0]
 end
+
+@testset "Exchange and grid location local inertial overland flow" begin
+    land = model.lateral.land
+    @test Wflow.exchange(land, :dt) == 0
+    @test Wflow.exchange(land, :g) == 0
+    @test Wflow.exchange(land, :h) == 1
+    @test Wflow.exchange(land, :qx) == 1
+    @test Wflow.grid_location(land, :dt) == "none"
+    @test Wflow.grid_location(land, :g) == "node"
+    @test Wflow.grid_location(land, :h) == "node"
+    @test Wflow.grid_location(land, :qx) == "edge"
+end
+
 Wflow.close_files(model, delete_output = false)
 
 # test local-inertial option for river flow including 1D floodplain schematization
@@ -410,7 +459,33 @@ model = Wflow.run_timestep(model)
     @test h[1622] ≈ 0.001987887580593841f0
     @test h[43] ≈ 0.436641524481545f0
     @test h[501] ≈ 0.05665942153713204f0
-	@test h[5808] ≈ 0.005933025055544241f0
+    @test h[5808] ≈ 0.005933025055544241f0
+end
+
+@testset "Exchange and grid location local inertial river flow" begin
+    river = model.lateral.river
+    @test Wflow.exchange(river, :dt) == 0
+    @test Wflow.exchange(river, :g) == 0
+    @test Wflow.exchange(river, :h) == 1
+    @test Wflow.exchange(river, :q_av) == 1
+    @test Wflow.grid_location(river, :dt) == "none"
+    @test Wflow.grid_location(river, :g) == "node"
+    @test Wflow.grid_location(river, :h) == "node"
+    @test Wflow.grid_location(river, :q_av) == "edge"
+end
+
+@testset "Exchange and grid location floodplain" begin
+    floodplain = model.lateral.river.floodplain
+    @test Wflow.exchange(floodplain.profile, :depth) == 1
+    @test Wflow.exchange(floodplain.profile, :a) == 1
+    @test Wflow.grid_location(floodplain.profile, :depth) == "node"
+    @test Wflow.grid_location(floodplain.profile, :a) == "node"
+    @test Wflow.exchange(floodplain, :profile) == 0
+    @test Wflow.exchange(floodplain, :h) == 1
+    @test Wflow.exchange(floodplain, :q) == 1
+    @test Wflow.grid_location(floodplain, :profile) == "node"
+    @test Wflow.grid_location(floodplain, :h) == "node"
+    @test Wflow.grid_location(floodplain, :q) == "edge"
 end
 
 # set boundary condition local inertial routing from netCDF file
diff --git a/test/run_sbm_gwf.jl b/test/run_sbm_gwf.jl
index c2d422f03..53c9d22fc 100644
--- a/test/run_sbm_gwf.jl
+++ b/test/run_sbm_gwf.jl
@@ -186,4 +186,21 @@ end
     @test gw.drain.flux[1] ≈ 0.0
     @test gw.recharge.rate[19] ≈ -0.0014241196552847502f0
 end
+
+@testset "Exchange and grid location aquifer, recharge and constant head" begin
+    aquifer = model.lateral.subsurface.flow.aquifer
+    @test Wflow.exchange(aquifer, :head) == 1
+    @test Wflow.exchange(aquifer, :k) == 1
+    @test Wflow.grid_location(aquifer, :head) == "node"
+    @test Wflow.grid_location(aquifer, :k) == "node"
+    recharge = model.lateral.subsurface.recharge
+    @test Wflow.exchange(recharge, :rate) == 1
+    @test Wflow.exchange(recharge, :flux) == 1
+    @test Wflow.grid_location(recharge, :rate) == "node"
+    @test Wflow.grid_location(recharge, :flux) == "node"
+    constanthead = model.lateral.subsurface.flow.constanthead
+    @test Wflow.exchange(constanthead, :head) == 0
+    @test Wflow.grid_location(constanthead, :head) == "node"
+end
+
 Wflow.close_files(model, delete_output = false)
diff --git a/test/run_sediment.jl b/test/run_sediment.jl
index 7d009bddc..1d44f79c5 100644
--- a/test/run_sediment.jl
+++ b/test/run_sediment.jl
@@ -51,4 +51,20 @@ end
     @test lat.river.outclay[5649] ≈ 2.359031898208781f-9
 end
 
+@testset "Exchange and grid location sediment" begin
+    @test Wflow.exchange(model.vertical, :n) == 0
+    @test Wflow.exchange(model.vertical, :erosk) == 1
+    @test Wflow.exchange(model.vertical, :leaf_area_index) == 1
+    @test Wflow.grid_location(model.vertical, :n) == "none"
+    @test Wflow.grid_location(model.vertical, :erosk) == "node"
+    @test Wflow.grid_location(model.vertical, :leaf_area_index) == "node"
+    land = model.lateral.land
+    @test Wflow.exchange(land, :n) == 0
+    @test Wflow.exchange(land, :soilloss) == 1
+    @test Wflow.exchange(land, :inlandsed) == 1
+    @test Wflow.grid_location(land, :n) == "none"
+    @test Wflow.grid_location(land, :soilloss) == "node"
+    @test Wflow.grid_location(land, :inlandsed) == "node"
+end
+
 Wflow.close_files(model)