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

Issue:
Refactor routing
### Feature type

Changing existing functionality

### Improvement Description

Refactor routing components (including groundwater): split structs into `variables`, `parameters` and `boundary_conditions` (if applicable). See also issue #225.

### Implementation Description

_No response_

### Additional Context

_No response_
Refactor routing
### Feature type

Changing existing functionality

### Improvement Description

Refactor routing components (including groundwater): split structs into `variables`, `parameters` and `boundary_conditions` (if applicable). See also issue #225.

### Implementation Description

_No response_

### Additional Context

_No response_


Discussion / hints:

Fix patch:
diff --git a/docs/changelog.qmd b/docs/changelog.qmd
index c5ad6760f..44d4798f7 100644
--- a/docs/changelog.qmd
+++ b/docs/changelog.qmd
@@ -16,6 +16,11 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Mutating BMI model control functions (`update`, `update_until` and `finalize`) and
   extended mutating BMI functions (`load_state` and `save_state`) should return `nothing`.
 - Added downloading of testdata to Dockerfile, to ensure an image was able to build.
+- The reservoir (`reservoir_index_f`) and lake (`lake_index_f`) indices as part of
+  `network.river` were not correct. These were mapped to their own index in the
+  `SimpleReservoir` and `Lake` struct, and not to the corresponding river index. This
+  resulted in incorrect surface water abstractions from reservoir and lake volumes, and
+  surface water abstractions were set at zero at the wrong river locations.
 
 ### Changed
 - Removed vertical concepts `HBV` and `FLEXTopo`.
@@ -31,6 +36,13 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   model `variables`, `parameters` and `boundary_conditions` (if applicable), including
   associated functions for initializing and updating these model components. The original
   long update function of the `SBM` soil part has been split into separate functions.
+- Refactor the lateral (routing) components: as for the vertical `SBM` concept split the
+  structs into `variables`, `parameters` and `boundary_conditions` (if applicable).
+- Timestepping method parameters for solving the kinematic wave and local inertial
+  approaches for river and overland flow are moved to a `TimeStepping` struct. The
+  timestepping implementation for the kinematic wave is now similar to the local inertial
+  method: a stable timestep is computed for each sub timestep (or a fixed sub timestep is
+  used) as part of a while loop (for each model timestep).
 
 ### Added
 - Support direct output of snow and glacier melt, and add computation of snow water
diff --git a/src/Wflow.jl b/src/Wflow.jl
index 233a96fe0..ff4abfaab 100644
--- a/src/Wflow.jl
+++ b/src/Wflow.jl
@@ -49,7 +49,7 @@ using Parameters: @with_kw
 using Polyester: @batch
 using ProgressLogging: @progress
 using StaticArrays: SVector, pushfirst, setindex
-using Statistics: mean, median, quantile!
+using Statistics: mean, median, quantile!, quantile
 using TerminalLoggers
 using TOML: TOML
 
@@ -141,8 +141,17 @@ Base.show(io::IO, m::Model) = print(io, "model of type ", typeof(m))
 
 include("forcing.jl")
 include("parameters.jl")
-include("flow.jl")
-include("horizontal_process.jl")
+include("groundwater/connectivity.jl")
+include("groundwater/aquifer.jl")
+include("groundwater/boundary_conditions.jl")
+include("routing/timestepping.jl")
+include("routing/subsurface.jl")
+include("routing/reservoir.jl")
+include("routing/lake.jl")
+include("routing/surface_kinwave.jl")
+include("routing/surface_local_inertial.jl")
+include("routing/surface_routing.jl")
+include("routing/routing_process.jl")
 include("vegetation/rainfall_interception.jl")
 include("vegetation/canopy.jl")
 include("snow/snow_process.jl")
@@ -155,12 +164,8 @@ include("soil/soil_process.jl")
 include("sbm.jl")
 include("demand/water_demand.jl")
 include("sediment.jl")
-include("reservoir_lake.jl")
 include("sbm_model.jl")
 include("sediment_model.jl")
-include("groundwater/connectivity.jl")
-include("groundwater/aquifer.jl")
-include("groundwater/boundary_conditions.jl")
 include("sbm_gwf_model.jl")
 include("utils.jl")
 include("bmi.jl")
diff --git a/src/bmi.jl b/src/bmi.jl
index 4063baae7..c97493f92 100644
--- a/src/bmi.jl
+++ b/src/bmi.jl
@@ -153,7 +153,7 @@ function BMI.get_var_grid(model::Model, name::String)
     s = split(name, "[")
     key = symbols(first(s))
     if exchange(param(model, key))
-        type = typeof(param(model, key[1:(end - 1)]))
+        type = typeof(param(model, key[1:2]))
         return if :reservoir in key
             0
         elseif :lake in key
@@ -162,9 +162,9 @@ function BMI.get_var_grid(model::Model, name::String)
             2
         elseif :river in key
             3
-        elseif type <: ShallowWaterLand && occursin("x", s[end])
+        elseif type <: LocalInertialOverlandFlow && occursin("x", s[end])
             4
-        elseif type <: ShallowWaterLand && occursin("y", s[end])
+        elseif type <: LocalInertialOverlandFlow && occursin("y", s[end])
             5
         else
             6
@@ -374,7 +374,7 @@ function BMI.get_grid_edge_nodes(model::Model, grid::Int, edge_nodes::Vector{Int
     m = div(n, 2)
     # inactive nodes (boundary/ghost points) are set at -999
     if grid == 3
-        nodes_at_edge = adjacent_nodes_at_link(network.river.graph)
+        nodes_at_edge = adjacent_nodes_at_edge(network.river.graph)
         nodes_at_edge.dst[nodes_at_edge.dst .== m + 1] .= -999
         edge_nodes[range(1, n; step = 2)] = nodes_at_edge.src
         edge_nodes[range(2, n; step = 2)] = nodes_at_edge.dst
diff --git a/src/demand/water_demand.jl b/src/demand/water_demand.jl
index ed3e7c77c..156d9f895 100644
--- a/src/demand/water_demand.jl
+++ b/src/demand/water_demand.jl
@@ -6,6 +6,7 @@ struct NoIrrigationPaddy{T} <: AbstractIrrigationModel{T} end
 struct NoIrrigationNonPaddy{T} <: AbstractIrrigationModel{T} end
 struct NoNonIrrigationDemand <: AbstractDemandModel end
 struct NoAllocationLand{T} <: AbstractAllocationModel{T} end
+struct NoAllocationRiver{T} <: AbstractAllocationModel{T} end
 
 "Struct to store non-irrigation water demand variables"
 @get_units @grid_loc @with_kw struct NonIrrigationDemandVariables{T}
@@ -30,26 +31,26 @@ get_demand_gross(model::NonIrrigationDemand) = model.demand.demand_gross
 get_demand_gross(model::NoNonIrrigationDemand) = 0.0
 
 "Initialize non-irrigation water demand model for a water use `sector`"
-function NonIrrigationDemand(nc, config, inds, dt, sector)
+function NonIrrigationDemand(dataset, config, indices, dt, sector)
     demand_gross =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.demand.$(sector).demand_gross";
-            sel = inds,
+            sel = indices,
             defaults = 0.0,
             type = Float,
         ) .* (dt / basetimestep)
     demand_net =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.demand.$(sector).demand_net";
-            sel = inds,
+            sel = indices,
             defaults = 0.0,
             type = Float,
         ) .* (dt / basetimestep)
-    n = length(inds)
+    n = length(indices)
     returnflow_f = return_flow_fraction.(demand_gross, demand_net)
 
     demand = PrescibedDemand{Float}(; demand_gross, demand_net)
@@ -82,39 +83,39 @@ end
 end
 
 "Initialize non-paddy irrigation model"
-function NonPaddy(nc, config, inds, dt)
+function NonPaddy(dataset, config, indices, dt)
     efficiency = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.nonpaddy.parameters.irrigation_efficiency";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
     )
     areas = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.nonpaddy.parameters.irrigation_areas";
-        sel = inds,
+        sel = indices,
         defaults = 1,
         optional = false,
         type = Int,
     )
     irrigation_trigger = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.nonpaddy.parameters.irrigation_trigger";
-        sel = inds,
+        sel = indices,
         defaults = 1,
         optional = false,
         type = Bool,
     )
     max_irri_rate =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.demand.nonpaddy.parameters.maximum_irrigation_rate";
-            sel = inds,
+            sel = indices,
             defaults = 25.0,
             type = Float,
         ) .* (dt / basetimestep)
@@ -125,7 +126,7 @@ function NonPaddy(nc, config, inds, dt)
         irrigation_areas = areas,
         irrigation_trigger,
     )
-    vars = NonPaddyVariables{Float}(; demand_gross = fill(mv, length(inds)))
+    vars = NonPaddyVariables{Float}(; demand_gross = fill(mv, length(indices)))
 
     nonpaddy = NonPaddy{Float}(; variables = vars, parameters = params)
 
@@ -137,7 +138,7 @@ get_demand_gross(model::NonPaddy) = model.variables.demand_gross
 get_demand_gross(model::NoIrrigationNonPaddy) = 0.0
 
 """
-    update_demand_gross!(nonpaddy::NonPaddy, soil::SbmSoilModel)
+    update_demand_gross!(model::NonPaddy, soil::SbmSoilModel)
 
 Update gross water demand `demand_gross` of the non-paddy irrigation model for a single
 timestep.
@@ -148,7 +149,7 @@ zone of the SBM soil model. Irrigation brings the root zone back to field capaci
 by the infiltration capacity, taking into account limited irrigation efficiency and limited
 by a maximum irrigation rate.
 """
-function update_demand_gross!(nonpaddy::NonPaddy, soil::SbmSoilModel)
+function update_demand_gross!(model::NonPaddy, soil::SbmSoilModel)
     (; hb, theta_s, theta_r, c, sumlayers, act_thickl, pathfrac, infiltcapsoil) =
         soil.parameters
     (; h3, n_unsatlayers, zi, ustorelayerdepth, f_infiltration_reduction) = soil.variables
@@ -157,7 +158,7 @@ function update_demand_gross!(nonpaddy::NonPaddy, soil::SbmSoilModel)
         irrigation_trigger,
         maximum_irrigation_rate,
         irrigation_efficiency,
-    ) = nonpaddy.parameters
+    ) = model.parameters
     rootingdepth = get_rootingdepth(soil)
 
     for i in eachindex(irrigation_areas)
@@ -178,7 +179,7 @@ function update_demand_gross!(nonpaddy::NonPaddy, soil::SbmSoilModel)
 
                 # check if maximum irrigation rate has been applied at the previous time step.
                 max_irri_rate_applied =
-                    nonpaddy.variables.demand_gross[i] == maximum_irrigation_rate[i]
+                    model.variables.demand_gross[i] == maximum_irrigation_rate[i]
                 if depletion >= raw # start irrigation
                     irri_dem_gross += depletion
                     # add depletion to irrigation gross demand when the maximum irrigation rate has been 
@@ -197,11 +198,11 @@ function update_demand_gross!(nonpaddy::NonPaddy, soil::SbmSoilModel)
         else
             irri_dem_gross = 0.0
         end
-        nonpaddy.variables.demand_gross[i] = irri_dem_gross
+        model.variables.demand_gross[i] = irri_dem_gross
     end
     return nothing
 end
-update_demand_gross!(nonpaddy::NoIrrigationNonPaddy, soil::SbmSoilModel) = nothing
+update_demand_gross!(model::NoIrrigationNonPaddy, soil::SbmSoilModel) = nothing
 
 "Struct to store paddy irrigation model variables"
 @get_units @grid_loc @with_kw struct PaddyVariables{T}
@@ -228,65 +229,65 @@ end
 end
 
 "Initialize paddy irrigation model"
-function Paddy(nc, config, inds, dt)
+function Paddy(dataset, config, indices, dt)
     h_min = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.h_min";
-        sel = inds,
+        sel = indices,
         defaults = 20.0,
         type = Float,
     )
     h_opt = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.h_opt";
-        sel = inds,
+        sel = indices,
         defaults = 50.0,
         type = Float,
     )
     h_max = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.h_max";
-        sel = inds,
+        sel = indices,
         defaults = 80.0,
         type = Float,
     )
     efficiency = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.irrigation_efficiency";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
     )
     areas = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.irrigation_areas";
-        sel = inds,
+        sel = indices,
         optional = false,
         type = Bool,
     )
     irrigation_trigger = ncread(
-        nc,
+        dataset,
         config,
         "vertical.demand.paddy.parameters.irrigation_trigger";
-        sel = inds,
+        sel = indices,
         optional = false,
         type = Bool,
     )
     max_irri_rate =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.demand.paddy.parameters.maximum_irrigation_rate";
-            sel = inds,
+            sel = indices,
             defaults = 25.0,
             type = Float,
         ) .* (dt / basetimestep)
-    n = length(inds)
+    n = length(indices)
     params = PaddyParameters{Float}(;
         irrigation_efficiency = efficiency,
         maximum_irrigation_rate = max_irri_rate,
@@ -395,7 +396,7 @@ function update_demand_gross!(model::Paddy)
     return nothing
 end
 
-update_demand_gross!(paddy::NoIrrigationPaddy) = nothing
+update_demand_gross!(model::NoIrrigationPaddy) = nothing
 
 "Struct to store water demand model variables"
 @get_units @grid_loc @with_kw struct DemandVariables{T}
@@ -436,34 +437,34 @@ end
 end
 
 "Initialize water demand model"
-function Demand(nc, config, inds, dt)
+function Demand(dataset, config, indices, dt)
     domestic = if get(config.model.water_demand, "domestic", false)
-        NonIrrigationDemand(nc, config, inds, dt, "domestic")
+        NonIrrigationDemand(dataset, config, indices, dt, "domestic")
     else
         NoNonIrrigationDemand()
     end
     industry = if get(config.model.water_demand, "industry", false)
-        NonIrrigationDemand(nc, config, inds, dt, "industry")
+        NonIrrigationDemand(dataset, config, indices, dt, "industry")
     else
         NoNonIrrigationDemand()
     end
     livestock = if get(config.model.water_demand, "livestock", false)
-        NonIrrigationDemand(nc, config, inds, dt, "livestock")
+        NonIrrigationDemand(dataset, config, indices, dt, "livestock")
     else
         NoNonIrrigationDemand()
     end
     paddy = if get(config.model.water_demand, "paddy", false)
-        Paddy(nc, config, inds, dt)
+        Paddy(dataset, config, indices, dt)
     else
         NoIrrigationPaddy{Float}()
     end
     nonpaddy = if get(config.model.water_demand, "nonpaddy", false)
-        NonPaddy(nc, config, inds, dt)
+        NonPaddy(dataset, config, indices, dt)
     else
         NoIrrigationNonPaddy{Float}()
     end
 
-    n = length(inds)
+    n = length(indices)
     vars = DemandVariables(Float, n)
     demand = Demand(; domestic, industry, livestock, paddy, nonpaddy, variables = vars)
     return demand
@@ -492,6 +493,9 @@ end
     variables::AllocationRiverVariables{T}
 end
 
+get_nonirrigation_returnflow(model::AllocationRiver) = model.variables.nonirri_returnflow
+get_nonirrigation_returnflow(model::NoAllocationRiver) = 0.0
+
 "Initialize water allocation for the river domain"
 function AllocationRiver(n)
     vars = AllocationRiverVariables(Float, n)
@@ -540,25 +544,25 @@ end
 end
 
 "Initialize water allocation for the land domain"
-function AllocationLand(nc, config, inds)
+function AllocationLand(dataset, config, indices)
     frac_sw_used = ncread(
-        nc,
+        dataset,
         config,
         "vertical.allocation.parameters.frac_sw_used";
-        sel = inds,
+        sel = indices,
         defaults = 1,
         type = Float,
     )
     areas = ncread(
-        nc,
+        dataset,
         config,
         "vertical.allocation.parameters.areas";
-        sel = inds,
+        sel = indices,
         defaults = 1,
         type = Int,
     )
 
-    n = length(inds)
+    n = length(indices)
 
     params = AllocationLandParameters(; areas = areas, frac_sw_used = frac_sw_used)
     vars = AllocationLandVariables(Float, n)
@@ -569,8 +573,13 @@ end
 # wrapper methods
 get_irrigation_allocated(model::AllocationLand) = model.variables.irri_alloc
 get_irrigation_allocated(model::NoAllocationLand) = 0.0
+get_nonirrigation_returnflow(model::AllocationLand) = model.variables.nonirri_returnflow
+get_nonirrigation_returnflow(model::NoAllocationLand) = 0.0
 
-"Return return flow fraction based on gross water demand `demand_gross` and net water demand `demand_net`"
+"""
+Return return flow fraction based on gross water demand `demand_gross` and net water demand
+`demand_net`
+"""
 function return_flow_fraction(demand_gross, demand_net)
     fraction = bounded_divide(demand_net, demand_gross)
     returnflow_fraction = 1.0 - fraction
@@ -588,24 +597,29 @@ end
 # return zero (gross water demand) if non-irrigation water demand sector is not defined
 return_flow_fraction!(model::NoNonIrrigationDemand) = nothing
 
-"Update water allocation for river and land domains based on local surface water (river) availability."
-function surface_water_allocation_local!(land_allocation, demand, river, network, dt)
-    (; surfacewater_alloc) = land_allocation.variables
+"""
+Update water allocation for river and land domains based on local surface water (river)
+availability.
+"""
+function surface_water_allocation_local!(model::AllocationLand, demand, river, network, dt)
+    (; surfacewater_alloc) = model.variables
     (; surfacewater_demand) = demand.variables
     (; act_surfacewater_abst_vol, act_surfacewater_abst, available_surfacewater) =
         river.allocation.variables
+    (; inflow) = river.boundary_conditions
+    (; volume) = river.variables
     # maps from the land domain to the internal river domain (linear index), excluding water bodies
-    index_river = network.land.index_river_wb
+    index_river = network.land.river_inds_excl_waterbody
     for i in eachindex(surfacewater_demand)
         if index_river[i] > 0.0
             # the available volume is limited by a fixed scaling factor of 0.8 to prevent
             # rivers completely drying out. check for abstraction through inflow (external
             # negative inflow) and adjust available volume.
-            if river.inflow[index_river[i]] < 0.0
-                inflow = river.inflow[index_river[i]] * dt
-                available_volume = max(river.volume[index_river[i]] * 0.80 + inflow, 0.0)
+            if inflow[index_river[i]] < 0.0
+                river_inflow = inflow[index_river[i]] * dt
+                available_volume = max(volume[index_river[i]] * 0.80 + river_inflow, 0.0)
             else
-                available_volume = river.volume[index_river[i]] * 0.80
+                available_volume = volume[index_river[i]] * 0.80
             end
             # satisfy surface water demand with available local river volume
             surfacewater_demand_vol = surfacewater_demand[i] * 0.001 * network.land.area[i]
@@ -624,17 +638,21 @@ function surface_water_allocation_local!(land_allocation, demand, river, network
     return nothing
 end
 
-"Update water allocation for river and land domains based on surface water (river) availability for allocation areas."
-function surface_water_allocation_area!(land_allocation, demand, river, network)
-    inds_river = network.river.indices_allocation_areas
-    inds_land = network.land.indices_allocation_areas
-    res_index = network.river.reservoir_index
-    lake_index = network.river.lake_index
+"""
+Update water allocation for river and land domains based on surface water (river)
+availability for allocation areas.
+"""
+function surface_water_allocation_area!(model::AllocationLand, demand, river, network)
+    inds_river = network.river.allocation_area_indices
+    inds_land = network.land.allocation_area_indices
+    inds_reservoir = network.river.reservoir_indices
+    inds_lake = network.river.lake_indices
 
     (; available_surfacewater, act_surfacewater_abst_vol, act_surfacewater_abst) =
         river.allocation.variables
-    (; surfacewater_alloc) = land_allocation.variables
+    (; surfacewater_alloc) = model.variables
     (; surfacewater_demand) = demand.variables
+    (; reservoir, lake) = river.boundary_conditions
 
     # loop over allocation areas
     for i in eachindex(inds_river)
@@ -646,15 +664,15 @@ function surface_water_allocation_area!(land_allocation, demand, river, network)
         # surface water availability (allocation area)
         sw_available = 0.0
         for j in inds_river[i]
-            if res_index[j] > 0
+            if inds_reservoir[j] > 0
                 # for reservoir locations use reservoir volume
-                k = res_index[j]
-                available_surfacewater[j] = river.reservoir.volume[k] * 0.98 # limit available reservoir volume
+                k = inds_reservoir[j]
+                available_surfacewater[j] = reservoir.volume[k] * 0.98 # limit available reservoir volume
                 sw_available += available_surfacewater[j]
-            elseif lake_index[j] > 0
+            elseif inds_lake[j] > 0
                 # for lake locations use lake volume
-                k = lake_index[j]
-                available_surfacewater[j] = river.lake.storage[k] * 0.98 # limit available lake volume
+                k = inds_lake[j]
+                available_surfacewater[j] = lake.storage[k] * 0.98 # limit available lake volume
                 sw_available += available_surfacewater[j]
 
             else
@@ -678,7 +696,7 @@ function surface_water_allocation_area!(land_allocation, demand, river, network)
         for j in inds_river[i]
             act_surfacewater_abst_vol[j] += frac_abstract_sw * available_surfacewater[j]
             act_surfacewater_abst[j] =
-                (act_surfacewater_abst_vol[j] / network.river.area[j]) * 1000.0
+                (act_surfacewater_abst_vol[j] / network.river.cell_area[j]) * 1000.0
         end
 
         # water allocated to each land cell.
@@ -690,21 +708,26 @@ function surface_water_allocation_area!(land_allocation, demand, river, network)
 end
 
 "Update water allocation for land domain based on local groundwater availability."
-function groundwater_allocation_local!(land_allocation, demand, groundwater_volume, network)
+function groundwater_allocation_local!(
+    model::AllocationLand,
+    demand,
+    groundwater_volume,
+    network,
+)
     (;
         surfacewater_alloc,
         act_groundwater_abst_vol,
         available_groundwater,
         act_groundwater_abst,
         groundwater_alloc,
-    ) = land_allocation.variables
+    ) = model.variables
     (; groundwater_demand, total_gross_demand) = demand.variables
 
     for i in eachindex(groundwater_demand)
         # groundwater demand based on allocation from surface water.
         groundwater_demand[i] = max(total_gross_demand[i] - surfacewater_alloc[i], 0.0)
-        # land index excluding water bodies
-        if network.index_wb[i]
+        # excluding water bodies
+        if !network.waterbody[i]
             # satisfy groundwater demand with available local groundwater volume
             groundwater_demand_vol = groundwater_demand[i] * 0.001 * network.area[i]
             available_volume = groundwater_volume[i] * 0.75 # limit available groundwater volume
@@ -722,16 +745,20 @@ function groundwater_allocation_local!(land_allocation, demand, groundwater_volu
     return nothing
 end
 
-"Update water allocation for land domain based on groundwater availability for allocation areas."
-function groundwater_allocation_area!(land_allocation, demand, network)
-    inds_river = network.river.indices_allocation_areas
-    inds_land = network.land.indices_allocation_areas
+"""
+Update water allocation for land domain based on groundwater availability for allocation
+areas.
+
+"""
+function groundwater_allocation_area!(model::AllocationLand, demand, network)
+    inds_river = network.river.allocation_area_indices
+    inds_land = network.land.allocation_area_indices
     (;
         act_groundwater_abst_vol,
         available_groundwater,
         act_groundwater_abst,
         groundwater_alloc,
-    ) = land_allocation.variables
+    ) = model.variables
 
     (; groundwater_demand) = demand.variables
 
@@ -766,27 +793,26 @@ function groundwater_allocation_area!(land_allocation, demand, network)
 end
 
 "Return and update non-irrigation sector (domestic, livestock, industry) return flow"
-function return_flow(non_irri::NonIrrigationDemand, nonirri_demand_gross, nonirri_alloc)
-    for i in eachindex(non_irri.variables.returnflow)
-        frac = bounded_divide(non_irri.demand.demand_gross[i], nonirri_demand_gross[i])
+function return_flow(model::NonIrrigationDemand, nonirri_demand_gross, nonirri_alloc)
+    for i in eachindex(model.variables.returnflow)
+        frac = bounded_divide(model.demand.demand_gross[i], nonirri_demand_gross[i])
         allocate = frac * nonirri_alloc[i]
-        non_irri.variables.returnflow[i] =
-            non_irri.variables.returnflow_fraction[i] * allocate
+        model.variables.returnflow[i] = model.variables.returnflow_fraction[i] * allocate
     end
-    return non_irri.variables.returnflow
+    return model.variables.returnflow
 end
 
 # return zero (return flow) if non-irrigation sector is not defined
-return_flow(non_irri::NoNonIrrigationDemand, nonirri_demand_gross, nonirri_alloc) = 0.0
+return_flow(model::NoNonIrrigationDemand, nonirri_demand_gross, nonirri_alloc) = 0.0
 
 # wrapper methods
-groundwater_volume(model::LateralSSF) = model.volume
-groundwater_volume(model) = model.flow.aquifer.volume
+groundwater_volume(model::LateralSSF) = model.variables.volume
+groundwater_volume(model) = model.flow.aquifer.variables.volume
 
 """
-    update_water_allocation!(land_allocation, demand::Demand, lateral, network, dt)
+    update_water_allocation!((model::AllocationLand, demand, lateral, network, dt)
 
-Update water allocation for the land domain `land_allocation` and water allocation for the
+Update water allocation for the land domain `AllocationLand` and water allocation for the
 river domain (part of `lateral`) based on the water `demand` model for a single timestep.
 First, surface water abstraction is computed to satisfy local water demand (non-irrigation
 and irrigation), and then updated (including lakes and reservoirs) to satisfy the remaining
@@ -794,11 +820,11 @@ water demand for allocation areas. Then groundwater abstraction is computed to s
 remaining local water demand, and then updated to satisfy the remaining water demand for
 allocation areas. Finally, non-irrigation return flows are updated.
 """
-function update_water_allocation!(land_allocation, demand::Demand, lateral, network, dt)
+function update_water_allocation!(model::AllocationLand, demand, lateral, network, dt)
     river = lateral.river
-    index_river = network.land.index_river_wb
-    res_index_f = network.river.reservoir_index_f
-    lake_index_f = network.river.lake_index_f
+    index_river = network.land.river_inds_excl_waterbody
+    inds_reservoir = network.reservoir.river_indices
+    inds_lake = network.lake.river_indices
     (;
         groundwater_alloc,
         surfacewater_alloc,
@@ -808,13 +834,14 @@ function update_water_allocation!(land_allocation, demand::Demand, lateral, netw
         irri_alloc,
         nonirri_alloc,
         nonirri_returnflow,
-    ) = land_allocation.variables
+    ) = model.variables
 
     (; surfacewater_demand, nonirri_demand_gross, irri_demand_gross, total_gross_demand) =
         demand.variables
 
-    (; frac_sw_used) = land_allocation.parameters
+    (; frac_sw_used) = model.parameters
     (; act_surfacewater_abst, act_surfacewater_abst_vol) = river.allocation.variables
+    (; abstraction, reservoir, lake) = river.boundary_conditions
 
     surfacewater_alloc .= 0.0
     act_surfacewater_abst .= 0.0
@@ -824,23 +851,21 @@ function update_water_allocation!(land_allocation, demand::Demand, lateral, netw
         frac_sw_used * nonirri_demand_gross + frac_sw_used * irri_demand_gross
 
     # local surface water demand and allocation (river, excluding reservoirs and lakes)
-    surface_water_allocation_local!(land_allocation, demand, river, network, dt)
+    surface_water_allocation_local!(model, demand, river, network, dt)
     # surface water demand and allocation for areas
-    surface_water_allocation_area!(land_allocation, demand, river, network)
+    surface_water_allocation_area!(model, demand, river, network)
 
-    @. river.abstraction = act_surfacewater_abst_vol / dt
+    @. abstraction = act_surfacewater_abst_vol / dt
 
     # for reservoir and lake locations set river abstraction at zero and abstract volume
     # from reservoir and lake, including an update of lake waterlevel
-    if !isnothing(river.reservoir)
-        @. river.abstraction[res_index_f] = 0.0
-        @. river.reservoir.volume -= act_surfacewater_abst_vol[res_index_f]
-    elseif !isnothing(river.lake)
-        @. river.abstraction[lake_index_f] = 0.0
-        lakes = river.lake
-        @. lakes.storage -= act_surfacewater_abst_vol[lake_index_f]
-        @. lakes.waterlevel =
-            waterlevel(lakes.storfunc, lakes.area, lakes.storage, lakes.sh)
+    if !isnothing(reservoir)
+        @. abstraction[inds_reservoir] = 0.0
+        @. reservoir.volume -= act_surfacewater_abst_vol[inds_reservoir]
+    elseif !isnothing(lake)
+        @. abstraction[inds_lake] = 0.0
+        @. lake.storage -= act_surfacewater_abst_vol[inds_lake]
+        @. lake.waterlevel = waterlevel(lake.storfunc, lake.area, lake.storage, lake.sh)
     end
 
     groundwater_alloc .= 0.0
@@ -848,13 +873,13 @@ function update_water_allocation!(land_allocation, demand::Demand, lateral, netw
     act_groundwater_abst .= 0.0
     # local groundwater demand and allocation
     groundwater_allocation_local!(
-        land_allocation,
+        model,
         demand,
         groundwater_volume(lateral.subsurface),
         network.land,
     )
     # groundwater demand and allocation for areas
-    groundwater_allocation_area!(land_allocation, demand, network)
+    groundwater_allocation_area!(model, demand, network)
 
     # irrigation allocation
     for i in eachindex(total_alloc)
@@ -882,17 +907,17 @@ function update_water_allocation!(land_allocation, demand::Demand, lateral, netw
         end
     end
 end
-update_water_allocation!(allocation, demand::NoDemand, lateral, network, dt) = nothing
+update_water_allocation!(model::NoAllocationLand, demand, lateral, network, dt) = nothing
 
 """
-    update_demand_gross!(demand::Demand)
+    update_demand_gross!(model::Demand)
 
 Update total irrigation gross water demand `irri_demand_gross`, total non-irrigation gross
 water demand `nonirri_demand_gross` and total gross water demand `total_gross_demand`.
 """
-function update_demand_gross!(demand::Demand)
-    (; nonpaddy, paddy, domestic, industry, livestock) = demand
-    (; irri_demand_gross, nonirri_demand_gross, total_gross_demand) = demand.variables
+function update_demand_gross!(model::Demand)
+    (; nonpaddy, paddy, domestic, industry, livestock) = model
+    (; irri_demand_gross, nonirri_demand_gross, total_gross_demand) = model.variables
     # get gross water demands
     industry_dem = get_demand_gross(industry)
     domestic_dem = get_demand_gross(domestic)
@@ -908,18 +933,18 @@ function update_demand_gross!(demand::Demand)
     return nothing
 end
 
-update_demand_gross!(demand::NoDemand) = nothing
+update_demand_gross!(model::NoDemand) = nothing
 
 """
-    update_water_demand!(demand::Demand, soil)
+    update_water_demand!(model::Demand, soil)
 
 Update the return flow fraction `returnflow_fraction` of `industry`, `domestic` and
 `livestock`, gross water demand `demand_gross` of `paddy` and `nonpaddy` models, and the
 total gross water demand, total irrigation gross water demand and total non-irrigation gross
 water demand as part of the water `demand` model.
 """
-function update_water_demand!(demand::Demand, soil)
-    (; nonpaddy, paddy, domestic, industry, livestock) = demand
+function update_water_demand!(model::Demand, soil)
+    (; nonpaddy, paddy, domestic, industry, livestock) = model
 
     return_flow_fraction!(industry)
     return_flow_fraction!(domestic)
@@ -927,8 +952,8 @@ function update_water_demand!(demand::Demand, soil)
 
     update_demand_gross!(nonpaddy, soil)
     update_demand_gross!(paddy)
-    update_demand_gross!(demand)
+    update_demand_gross!(model)
 
     return nothing
 end
-update_water_demand!(demand::NoDemand, soil) = nothing
\ No newline at end of file
+update_water_demand!(model::NoDemand, soil) = nothing
\ No newline at end of file
diff --git a/src/flow.jl b/src/flow.jl
deleted file mode 100644
index dab39db52..000000000
--- a/src/flow.jl
+++ /dev/null
@@ -1,1857 +0,0 @@
-
-abstract type SurfaceFlow end
-
-@get_units @grid_loc @with_kw struct SurfaceFlowRiver{T, R, L, A} <: 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⁻¹]
-    abstraction::Vector{T} | "m3 s-1"       # Abstraction (computed as part of water demand and allocation) [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                                   # 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
-    allocation::A                           # Water allocation
-    kinwave_it::Bool                        # Boolean for iterations kinematic wave
-end
-
-@get_units @grid_loc @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                               # 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)
-    @info "Kinematic wave approach is used for overland flow." iterate
-    if tstep > 0
-        @info "Using a fixed sub-timestep (seconds) $tstep for kinematic wave overland flow."
-    end
-
-    n_land =
-        ncread(nc, config, "lateral.land.n"; sel = inds, defaults = 0.072, type = Float)
-    n = length(inds)
-
-    sf_land = SurfaceFlowLand(;
-        beta = Float(0.6),
-        sl = sl,
-        n = n_land,
-        dl = dl,
-        q = zeros(Float, n),
-        qin = zeros(Float, n),
-        q_av = zeros(Float, n),
-        qlat = zeros(Float, n),
-        inwater = zeros(Float, n),
-        volume = zeros(Float, n),
-        h = zeros(Float, n),
-        h_av = zeros(Float, n),
-        dt = Float(tosecond(dt)),
-        its = tstep > 0 ? Int(cld(tosecond(dt), tstep)) : tstep,
-        width = width,
-        alpha_pow = Float((2.0 / 3.0) * 0.6),
-        alpha_term = fill(mv, n),
-        alpha = fill(mv, n),
-        cel = zeros(Float, n),
-        to_river = zeros(Float, n),
-        kinwave_it = iterate,
-    )
-
-    return sf_land
-end
-
-function initialize_surfaceflow_river(
-    nc,
-    config,
-    inds;
-    dl,
-    width,
-    reservoir_index,
-    reservoir,
-    lake_index,
-    lake,
-    iterate,
-    tstep,
-    dt,
-)
-    @info "Kinematic wave approach is used for river flow." iterate
-    if tstep > 0
-        @info "Using a fixed sub-timestep (seconds) $tstep for kinematic wave river flow."
-    end
-
-    n_river =
-        ncread(nc, config, "lateral.river.n"; sel = inds, defaults = 0.036, type = Float)
-    bankfull_depth = ncread(
-        nc,
-        config,
-        "lateral.river.bankfull_depth";
-        alias = "lateral.river.h_bankfull",
-        sel = inds,
-        defaults = 1.0,
-        type = Float,
-    )
-    if haskey(config.input.lateral.river, "h_bankfull")
-        @warn string(
-            "The `h_bankfull` key in `[input.lateral.river]` is now called ",
-            "`bankfull_depth`. Please update your TOML file.",
-        )
-    end
-    sl = ncread(
-        nc,
-        config,
-        "lateral.river.slope";
-        optional = false,
-        sel = inds,
-        type = Float,
-    )
-    clamp!(sl, 0.00001, Inf)
-
-    do_water_demand = haskey(config.model, "water_demand")
-    n = length(inds)
-
-    sf_river = SurfaceFlowRiver(;
-        beta = Float(0.6),
-        sl = sl,
-        n = n_river,
-        dl = dl,
-        q = zeros(Float, n),
-        qin = zeros(Float, n),
-        q_av = zeros(Float, n),
-        qlat = zeros(Float, n),
-        inwater = zeros(Float, n),
-        inflow = zeros(Float, n),
-        abstraction = zeros(Float, n),
-        inflow_wb = zeros(Float, n),
-        volume = zeros(Float, n),
-        h = zeros(Float, n),
-        h_av = zeros(Float, n),
-        bankfull_depth = bankfull_depth,
-        dt = Float(tosecond(dt)),
-        its = tstep > 0 ? Int(cld(tosecond(dt), tstep)) : tstep,
-        width = width,
-        alpha_pow = Float((2.0 / 3.0) * 0.6),
-        alpha_term = fill(mv, n),
-        alpha = fill(mv, n),
-        cel = zeros(Float, n),
-        reservoir_index = reservoir_index,
-        lake_index = lake_index,
-        reservoir = reservoir,
-        lake = lake,
-        kinwave_it = iterate,
-        allocation = do_water_demand ? AllocationRiver(n) : nothing,
-    )
-
-    return sf_river
-end
-
-function update!(sf::SurfaceFlowLand, network, frac_toriver)
-    (; subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes) = network
-
-    ns = length(subdomain_order)
-
-    @. sf.alpha_term = pow(sf.n / sqrt(sf.sl), sf.beta)
-    # use fixed alpha value based flow width
-    @. sf.alpha = sf.alpha_term * pow(sf.width, sf.alpha_pow)
-    @. sf.qlat = sf.inwater / sf.dl
-
-    sf.q_av .= 0.0
-    sf.h_av .= 0.0
-    sf.to_river .= 0.0
-
-    dt, its = stable_timestep(sf)
-    for _ in 1:its
-        sf.qin .= 0.0
-        for k in 1:ns
-            threaded_foreach(eachindex(subdomain_order[k]); basesize = 1) do i
-                m = subdomain_order[k][i]
-                for (n, v) in zip(indices_subdomain[m], topo_subdomain[m])
-                    # for a river cell without a reservoir or lake part of the upstream
-                    # surface flow goes to the river (frac_toriver) and part goes to the
-                    # surface flow reservoir (1.0 - frac_toriver), upstream nodes with a
-                    # reservoir or lake are excluded
-                    sf.to_river[v] += sum_at(
-                        i -> sf.q[i] * frac_toriver[i],
-                        upstream_nodes[n],
-                        eltype(sf.to_river),
-                    )
-                    if sf.width[v] > 0.0
-                        sf.qin[v] = sum_at(
-                            i -> sf.q[i] * (1.0 - frac_toriver[i]),
-                            upstream_nodes[n],
-                            eltype(sf.q),
-                        )
-                    end
-
-                    sf.q[v] = kinematic_wave(
-                        sf.qin[v],
-                        sf.q[v],
-                        sf.qlat[v],
-                        sf.alpha[v],
-                        sf.beta,
-                        dt,
-                        sf.dl[v],
-                    )
-
-                    # update h, only if surface width > 0.0
-                    if sf.width[v] > 0.0
-                        crossarea = sf.alpha[v] * pow(sf.q[v], sf.beta)
-                        sf.h[v] = crossarea / sf.width[v]
-                    end
-                    sf.q_av[v] += sf.q[v]
-                    sf.h_av[v] += sf.h[v]
-                end
-            end
-        end
-    end
-    sf.q_av ./= its
-    sf.h_av ./= its
-    sf.to_river ./= its
-    sf.volume .= sf.dl .* sf.width .* sf.h
-    return nothing
-end
-
-function update!(sf::SurfaceFlowRiver, network, doy)
-    (; graph, subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes) = network
-
-    ns = length(subdomain_order)
-
-    @. sf.alpha_term = pow(sf.n / sqrt(sf.sl), sf.beta)
-    # use fixed alpha value based on 0.5 * bankfull_depth
-    @. sf.alpha = sf.alpha_term * pow(sf.width + sf.bankfull_depth, sf.alpha_pow)
-
-    @. sf.qlat = sf.inwater / sf.dl
-
-    sf.q_av .= 0.0
-    sf.h_av .= 0.0
-    # because of possible iterations set reservoir and lake inflow and total outflow at
-    # start to zero, the total sum of inflow and outflow at each sub time step is calculated
-    if !isnothing(sf.reservoir)
-        sf.reservoir.inflow .= 0.0
-        sf.reservoir.totaloutflow .= 0.0
-        sf.reservoir.actevap .= 0.0
-    end
-    if !isnothing(sf.lake)
-        sf.lake.inflow .= 0.0
-        sf.lake.totaloutflow .= 0.0
-        sf.lake.actevap .= 0.0
-    end
-
-    dt, its = stable_timestep(sf)
-    for _ in 1:its
-        sf.qin .= 0.0
-        for k in 1:ns
-            threaded_foreach(eachindex(subdomain_order[k]); basesize = 1) do i
-                m = subdomain_order[k][i]
-                for (n, v) in zip(indices_subdomain[m], topo_subdomain[m])
-                    # sf.qin by outflow from upstream reservoir or lake location is added
-                    sf.qin[v] += sum_at(sf.q, upstream_nodes[n])
-                    # Inflow supply/abstraction is added to qlat (divide by flow length)
-                    # If inflow < 0, abstraction is limited
-                    if sf.inflow[v] < 0.0
-                        max_abstract = min(
-                            (sf.inwater[v] + sf.qin[v] + sf.volume[v] / dt) * 0.80,
-                            -sf.inflow[v],
-                        )
-                        inflow = -max_abstract / sf.dl[v]
-                    else
-                        inflow = sf.inflow[v] / sf.dl[v]
-                    end
-                    inflow -= sf.abstraction[v] / sf.dl[v]
-
-                    sf.q[v] = kinematic_wave(
-                        sf.qin[v],
-                        sf.q[v],
-                        sf.qlat[v] + inflow,
-                        sf.alpha[v],
-                        sf.beta,
-                        dt,
-                        sf.dl[v],
-                    )
-
-                    if !isnothing(sf.reservoir) && sf.reservoir_index[v] != 0
-                        # run reservoir model and copy reservoir outflow to inflow (qin) of
-                        # downstream river cell
-                        i = sf.reservoir_index[v]
-                        update!(sf.reservoir, i, sf.q[v] + sf.inflow_wb[v], dt)
-
-                        downstream_nodes = outneighbors(graph, v)
-                        n_downstream = length(downstream_nodes)
-                        if n_downstream == 1
-                            j = only(downstream_nodes)
-                            sf.qin[j] = sf.reservoir.outflow[i]
-                        elseif n_downstream == 0
-                            error(
-                                """A reservoir without a downstream river node is not supported.
-                                Add a downstream river node or move the reservoir to an upstream node (model schematization).
-                                """,
-                            )
-                        else
-                            error("bifurcations not supported")
-                        end
-
-                    elseif !isnothing(sf.lake) && sf.lake_index[v] != 0
-                        # run lake model and copy lake outflow to inflow (qin) of downstream river
-                        # cell
-                        i = sf.lake_index[v]
-                        update!(sf.lake, i, sf.q[v] + sf.inflow_wb[v], doy, dt)
-
-                        downstream_nodes = outneighbors(graph, v)
-                        n_downstream = length(downstream_nodes)
-                        if n_downstream == 1
-                            j = only(downstream_nodes)
-                            sf.qin[j] = sf.lake.outflow[i]
-                        elseif n_downstream == 0
-                            error(
-                                """A lake without a downstream river node is not supported.
-                                Add a downstream river node or move the lake to an upstream node (model schematization).
-                                """,
-                            )
-                        else
-                            error("bifurcations not supported")
-                        end
-                    end
-
-                    # update h
-                    crossarea = sf.alpha[v] * pow(sf.q[v], sf.beta)
-                    sf.h[v] = crossarea / sf.width[v]
-                    sf.volume[v] = sf.dl[v] * sf.width[v] * sf.h[v]
-                    sf.q_av[v] += sf.q[v]
-                    sf.h_av[v] += sf.h[v]
-                end
-            end
-        end
-    end
-    sf.q_av ./= its
-    sf.h_av ./= its
-    sf.volume .= sf.dl .* sf.width .* sf.h
-    return nothing
-end
-
-function stable_timestep(sf::S) where {S <: SurfaceFlow}
-    n = length(sf.q)
-    # two options for iteration, fixed or based on courant number.
-    if sf.kinwave_it
-        if sf.its > 0
-            its = sf.its
-        else
-            # calculate celerity
-            courant = zeros(n)
-            for v in 1:n
-                if sf.q[v] > 0.0
-                    sf.cel[v] =
-                        1.0 / (sf.alpha[v] * sf.beta * pow(sf.q[v], (sf.beta - 1.0)))
-                    courant[v] = (sf.dt / sf.dl[v]) * sf.cel[v]
-                end
-            end
-            filter!(x -> x ≠ 0.0, courant)
-            its = isempty(courant) ? 1 : ceil(Int, (1.25 * quantile!(courant, 0.95)))
-        end
-    else
-        its = 1
-    end
-
-    # sub time step
-    dt = sf.dt / its
-    return dt, its
-end
-
-@get_units @grid_loc struct KhExponential{T}
-    # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
-    kh_0::Vector{T} | "m d-1"
-    # A scaling parameter [m⁻¹] (controls exponential decline of kh_0)
-    f::Vector{T} | "m-1"
-end
-
-@get_units @grid_loc struct KhExponentialConstant{T}
-    # Exponential horizontal hydraulic conductivity profile type
-    exponential::KhExponential
-    # Depth [m] from soil surface for which exponential decline of kv_0 is valid
-    z_exp::Vector{T} | "m"
-end
-
-@get_units @grid_loc struct KhLayered{T}
-    # Horizontal hydraulic conductivity [m d⁻¹]
-    kh::Vector{T} | "m d-1"
-end
-
-abstract type SubsurfaceFlow end
-
-@get_units @grid_loc @with_kw struct LateralSSF{T, Kh} <: SubsurfaceFlow
-    kh_profile::Kh                         # Horizontal hydraulic conductivity profile type [-]  
-    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                                  # 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)
-    exfiltwater::Vector{T} | "m dt-1"      # Exfiltration [m Δt⁻¹] (groundwater above surface level, saturated excess conditions)
-    recharge::Vector{T} | "m2 dt-1"        # Net recharge to saturated store [m² Δt⁻¹]
-    ssf::Vector{T} | "m3 d-1"              # Subsurface flow [m³ d⁻¹]
-    ssfin::Vector{T} | "m3 d-1"            # Inflow from upstream cells [m³ d⁻¹]
-    ssfmax::Vector{T} | "m2 d-1"           # Maximum subsurface flow [m² d⁻¹]
-    to_river::Vector{T} | "m3 d-1"         # Part of subsurface flow [m³ d⁻¹] that flows to the river
-    volume::Vector{T} | "m3"               # Subsurface volume [m³]
-
-    function LateralSSF{T, Kh}(args...) where {T, Kh}
-        equal_size_vectors(args)
-        return new(args...)
-    end
-end
-
-function update!(ssf::LateralSSF, network, frac_toriver)
-    (; subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes, area) = network
-
-    ns = length(subdomain_order)
-    for k in 1:ns
-        threaded_foreach(eachindex(subdomain_order[k]); basesize = 1) do i
-            m = subdomain_order[k][i]
-            for (n, v) in zip(indices_subdomain[m], topo_subdomain[m])
-                # for a river cell without a reservoir or lake part of the upstream
-                # subsurface flow goes to the river (frac_toriver) and part goes to the
-                # subsurface flow reservoir (1.0 - frac_toriver) upstream nodes with a
-                # reservoir or lake are excluded
-                ssf.ssfin[v] = sum_at(
-                    i -> ssf.ssf[i] * (1.0 - frac_toriver[i]),
-                    upstream_nodes[n],
-                    eltype(ssf.ssfin),
-                )
-                ssf.to_river[v] = sum_at(
-                    i -> ssf.ssf[i] * frac_toriver[i],
-                    upstream_nodes[n],
-                    eltype(ssf.to_river),
-                )
-                ssf.ssf[v], ssf.zi[v], ssf.exfiltwater[v] = kinematic_wave_ssf(
-                    ssf.ssfin[v],
-                    ssf.ssf[v],
-                    ssf.zi[v],
-                    ssf.recharge[v],
-                    ssf.slope[v],
-                    ssf.theta_s[v] - ssf.theta_r[v],
-                    ssf.soilthickness[v],
-                    ssf.dt,
-                    ssf.dl[v],
-                    ssf.dw[v],
-                    ssf.ssfmax[v],
-                    ssf.kh_profile,
-                    v,
-                )
-                ssf.volume[v] =
-                    (ssf.theta_s[v] - ssf.theta_r[v]) *
-                    (ssf.soilthickness[v] - ssf.zi[v]) *
-                    area[v]
-            end
-        end
-    end
-    return nothing
-end
-
-@get_units@grid_loc @with_kw struct GroundwaterExchange{T} <: SubsurfaceFlow
-    dt::T                               # 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_water_depth(subsurface::SubsurfaceFlow) = subsurface.zi
-get_exfiltwater(subsurface::SubsurfaceFlow) = subsurface.exfiltwater
-
-@get_units @grid_loc @with_kw struct ShallowWaterRiver{T, R, L, F, A}
-    n::Int                                              # number of cells [-]
-    ne::Int                                             # number of edges/links [-]
-    active_n::Vector{Int} | "-"                         # active nodes [-]
-    active_e::Vector{Int} | "-" | "edge"                # active edges/links [-]
-    g::T                                                # acceleration due to gravity [m s⁻²]
-    alpha::T                                            # stability coefficient (Bates et al., 2010) [-]
-    h_thresh::T                                         # depth threshold for calculating flow [m]
-    dt::T                                               # 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⁻¹]
-    abstraction::Vector{T} | "m3 s-1"                   # abstraction (computed as part of water demand and allocation) [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
-    allocation::A                                       # Water allocation
-end
-
-function initialize_shallowwater_river(
-    nc,
-    config,
-    inds;
-    graph,
-    ldd,
-    dl,
-    width,
-    reservoir_index,
-    reservoir,
-    lake_index,
-    lake,
-    dt,
-    floodplain,
-)
-    # The local inertial approach makes use of a staggered grid (Bates et al. (2010)),
-    # with nodes and links. This information is extracted from the directed graph of the
-    # river. Discharge q is calculated at links between nodes and mapped to the source
-    # nodes for gridded output (index of link is equal to source node index, e.g.:
-    # Edge 1 => 5
-    # Edge 2 => 1
-    # Edge 3 => 2
-    # Edge 4 => 9
-    # ⋮ )
-    alpha = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
-    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at link
-    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number
-    floodplain_1d = floodplain
-
-    @info "Local inertial approach is used for river flow." alpha h_thresh froude_limit floodplain_1d
-    @warn string(
-        "Providing the boundary condition `riverlength_bc` as part of the `[model]` setting ",
-        "in the TOML file has been deprecated as of Wflow v0.8.0.\n The boundary condition should ",
-        "be provided as part of the file `$(config.input.path_static)`.",
-    )
-    # The following boundary conditions can be set at ghost nodes, downstream of river
-    # outlets (pits): river length and river depth
-    index_pit = findall(x -> x == 5, ldd)
-    inds_pit = inds[index_pit]
-    riverlength_bc = ncread(
-        nc,
-        config,
-        "lateral.river.riverlength_bc";
-        sel = inds_pit,
-        defaults = 1.0e04,
-        type = Float,
-    )
-    riverdepth_bc = ncread(
-        nc,
-        config,
-        "lateral.river.riverdepth_bc";
-        sel = inds_pit,
-        defaults = 0.0,
-        type = Float,
-    )
-    bankfull_elevation_2d = ncread(
-        nc,
-        config,
-        "lateral.river.bankfull_elevation";
-        optional = false,
-        type = Float,
-        fill = 0,
-    )
-    bankfull_depth_2d = ncread(
-        nc,
-        config,
-        "lateral.river.bankfull_depth";
-        optional = false,
-        type = Float,
-        fill = 0,
-    )
-    bankfull_depth = bankfull_depth_2d[inds]
-    zb = bankfull_elevation_2d[inds] - bankfull_depth # river bed elevation
-
-    bankfull_volume = bankfull_depth .* width .* dl
-
-    n_river =
-        ncread(nc, config, "lateral.river.n"; sel = inds, defaults = 0.036, type = Float)
-
-    n = length(inds)
-
-    # set river depth h to zero (including reservoir and lake locations)
-    h = fill(0.0, n)
-
-    # set ghost points for boundary condition (downstream river outlet): river width, bed
-    # elevation, manning n is copied from the upstream cell.
-    add_vertex_edge_graph!(graph, index_pit)
-    append!(dl, riverlength_bc)
-    append!(h, riverdepth_bc)
-    append!(zb, zb[index_pit])
-    append!(width, width[index_pit])
-    append!(n_river, n_river[index_pit])
-    append!(bankfull_depth, bankfull_depth[index_pit])
-
-    # for each link the src and dst node is required
-    nodes_at_link = adjacent_nodes_at_link(graph)
-    _ne = ne(graph)
-
-    if floodplain
-        zb_floodplain = zb .+ bankfull_depth
-        floodplain = initialize_floodplain_1d(
-            nc,
-            config,
-            inds,
-            width,
-            dl,
-            zb_floodplain,
-            index_pit,
-            _ne,
-            nodes_at_link,
-        )
-    else
-        floodplain = nothing
-    end
-
-    # determine z, width, length and manning's n at links
-    zb_max = fill(Float(0), _ne)
-    width_at_link = fill(Float(0), _ne)
-    length_at_link = fill(Float(0), _ne)
-    mannings_n_sq = fill(Float(0), _ne)
-    for i in 1:_ne
-        src_node = nodes_at_link.src[i]
-        dst_node = nodes_at_link.dst[i]
-        zb_max[i] = max(zb[src_node], zb[dst_node])
-        width_at_link[i] = min(width[src_node], width[dst_node])
-        length_at_link[i] = 0.5 * (dl[dst_node] + dl[src_node])
-        mannings_n =
-            (n_river[dst_node] * dl[dst_node] + n_river[src_node] * dl[src_node]) /
-            (dl[dst_node] + dl[src_node])
-        mannings_n_sq[i] = mannings_n * mannings_n
-    end
-
-    q_av = zeros(_ne)
-    waterbody = !=(0).(reservoir_index .+ lake_index)
-    active_index = findall(x -> x == 0, waterbody)
-
-    do_water_demand = haskey(config.model, "water_demand")
-    sw_river = ShallowWaterRiver(;
-        n = n,
-        ne = _ne,
-        active_n = active_index,
-        active_e = active_index,
-        g = 9.80665,
-        alpha = alpha,
-        h_thresh = h_thresh,
-        dt = tosecond(dt),
-        q = zeros(_ne),
-        q0 = zeros(_ne),
-        q_av = q_av,
-        q_channel_av = isnothing(floodplain) ? q_av : zeros(_ne),
-        zb_max = zb_max,
-        mannings_n_sq = mannings_n_sq,
-        mannings_n = n_river,
-        h = h,
-        zs_max = zeros(_ne),
-        zs_src = zeros(_ne),
-        zs_dst = zeros(_ne),
-        hf = zeros(_ne),
-        h_av = zeros(n),
-        width = width,
-        width_at_link = width_at_link,
-        a = zeros(_ne),
-        r = zeros(_ne),
-        volume = fill(0.0, n),
-        error = zeros(n),
-        inflow = zeros(n),
-        abstraction = zeros(n),
-        inflow_wb = zeros(n),
-        inwater = zeros(n),
-        dl = dl,
-        dl_at_link = length_at_link,
-        bankfull_volume = bankfull_volume,
-        bankfull_depth = bankfull_depth,
-        zb = zb,
-        froude_limit = froude_limit,
-        reservoir_index = findall(x -> x > 0, reservoir_index),
-        lake_index = findall(x -> x > 0, lake_index),
-        waterbody = waterbody,
-        reservoir = reservoir,
-        lake = lake,
-        floodplain = floodplain,
-        allocation = do_water_demand ? initialize_allocation_river(n) : nothing,
-    )
-    return sw_river, nodes_at_link
-end
-
-"Return the upstream inflow for a waterbody in `ShallowWaterRiver`"
-function get_inflow_waterbody(sw::ShallowWaterRiver, src_edge)
-    q_in = sum_at(sw.q, src_edge)
-    if !isnothing(sw.floodplain)
-        q_in = q_in + sum_at(sw.floodplain.q, src_edge)
-    end
-    return q_in
-end
-
-function shallowwater_river_update!(sw::ShallowWaterRiver, network, dt, doy, update_h)
-    (; nodes_at_link, links_at_node) = network
-
-    sw.q0 .= sw.q
-    if !isnothing(sw.floodplain)
-        sw.floodplain.q0 .= sw.floodplain.q
-    end
-    @tturbo for j in eachindex(sw.active_e)
-        i = sw.active_e[j]
-        i_src = nodes_at_link.src[i]
-        i_dst = nodes_at_link.dst[i]
-        sw.zs_src[i] = sw.zb[i_src] + sw.h[i_src]
-        sw.zs_dst[i] = sw.zb[i_dst] + sw.h[i_dst]
-
-        sw.zs_max[i] = max(sw.zs_src[i], sw.zs_dst[i])
-        sw.hf[i] = (sw.zs_max[i] - sw.zb_max[i])
-
-        sw.a[i] = sw.width_at_link[i] * sw.hf[i] # flow area (rectangular channel)
-        sw.r[i] = sw.a[i] / (sw.width_at_link[i] + 2.0 * sw.hf[i]) # hydraulic radius (rectangular channel)
-
-        sw.q[i] = IfElse.ifelse(
-            sw.hf[i] > sw.h_thresh,
-            local_inertial_flow(
-                sw.q0[i],
-                sw.zs_src[i],
-                sw.zs_dst[i],
-                sw.hf[i],
-                sw.a[i],
-                sw.r[i],
-                sw.dl_at_link[i],
-                sw.mannings_n_sq[i],
-                sw.g,
-                sw.froude_limit,
-                dt,
-            ),
-            0.0,
-        )
-
-        # limit q in case water is not available
-        sw.q[i] = IfElse.ifelse(sw.h[i_src] <= 0.0, min(sw.q[i], 0.0), sw.q[i])
-        sw.q[i] = IfElse.ifelse(sw.h[i_dst] <= 0.0, max(sw.q[i], 0.0), sw.q[i])
-
-        sw.q_av[i] += sw.q[i] * dt
-    end
-    if !isnothing(sw.floodplain)
-        @tturbo @. sw.floodplain.hf = max(sw.zs_max - sw.floodplain.zb_max, 0.0)
-
-        n = 0
-        @inbounds for i in sw.active_e
-            @inbounds if sw.floodplain.hf[i] > sw.h_thresh
-                n += 1
-                sw.floodplain.hf_index[n] = i
-            else
-                sw.floodplain.q[i] = 0.0
-            end
-        end
-
-        @tturbo for j in 1:n
-            i = sw.floodplain.hf_index[j]
-            i_src = nodes_at_link.src[i]
-            i_dst = nodes_at_link.dst[i]
-
-            i0 = 0
-            for k in eachindex(sw.floodplain.profile.depth)
-                i0 += 1 * (sw.floodplain.profile.depth[k] <= sw.floodplain.hf[i])
-            end
-            i1 = max(i0, 1)
-            i2 = IfElse.ifelse(i1 == length(sw.floodplain.profile.depth), i1, i1 + 1)
-
-            a_src = flow_area(
-                sw.floodplain.profile.width[i2, i_src],
-                sw.floodplain.profile.a[i1, i_src],
-                sw.floodplain.profile.depth[i1],
-                sw.floodplain.hf[i],
-            )
-            a_src = max(a_src - (sw.floodplain.hf[i] * sw.width[i_src]), 0.0)
-
-            a_dst = flow_area(
-                sw.floodplain.profile.width[i2, i_dst],
-                sw.floodplain.profile.a[i1, i_dst],
-                sw.floodplain.profile.depth[i1],
-                sw.floodplain.hf[i],
-            )
-            a_dst = max(a_dst - (sw.floodplain.hf[i] * sw.width[i_dst]), 0.0)
-
-            sw.floodplain.a[i] = min(a_src, a_dst)
-
-            sw.floodplain.r[i] = IfElse.ifelse(
-                a_src < a_dst,
-                a_src / wetted_perimeter(
-                    sw.floodplain.profile.p[i1, i_src],
-                    sw.floodplain.profile.depth[i1],
-                    sw.floodplain.hf[i],
-                ),
-                a_dst / wetted_perimeter(
-                    sw.floodplain.profile.p[i1, i_dst],
-                    sw.floodplain.profile.depth[i1],
-                    sw.floodplain.hf[i],
-                ),
-            )
-
-            sw.floodplain.q[i] = IfElse.ifelse(
-                sw.floodplain.a[i] > 1.0e-05,
-                local_inertial_flow(
-                    sw.floodplain.q0[i],
-                    sw.zs_src[i],
-                    sw.zs_dst[i],
-                    sw.floodplain.hf[i],
-                    sw.floodplain.a[i],
-                    sw.floodplain.r[i],
-                    sw.dl_at_link[i],
-                    sw.floodplain.mannings_n_sq[i],
-                    sw.g,
-                    sw.froude_limit,
-                    dt,
-                ),
-                0.0,
-            )
-
-            # limit floodplain q in case water is not available
-            sw.floodplain.q[i] = IfElse.ifelse(
-                sw.floodplain.h[i_src] <= 0.0,
-                min(sw.floodplain.q[i], 0.0),
-                sw.floodplain.q[i],
-            )
-            sw.floodplain.q[i] = IfElse.ifelse(
-                sw.floodplain.h[i_dst] <= 0.0,
-                max(sw.floodplain.q[i], 0.0),
-                sw.floodplain.q[i],
-            )
-
-            sw.floodplain.q[i] =
-                IfElse.ifelse(sw.floodplain.q[i] * sw.q[i] < 0.0, 0.0, sw.floodplain.q[i])
-            sw.floodplain.q_av[i] += sw.floodplain.q[i] * dt
-        end
-    end
-    # For reservoir and lake locations the local inertial solution is replaced by the
-    # reservoir or lake model. These locations are handled as boundary conditions in the
-    # local inertial model (fixed h).
-    for v in eachindex(sw.reservoir_index)
-        i = sw.reservoir_index[v]
-
-        q_in = get_inflow_waterbody(sw, links_at_node.src[i])
-        update!(sw.reservoir, v, q_in + sw.inflow_wb[i], dt)
-        sw.q[i] = sw.reservoir.outflow[v]
-        sw.q_av[i] += sw.q[i] * dt
-    end
-    for v in eachindex(sw.lake_index)
-        i = sw.lake_index[v]
-
-        q_in = get_inflow_waterbody(sw, links_at_node.src[i])
-        update!(sw.lake, v, q_in + sw.inflow_wb[i], doy, dt)
-        sw.q[i] = sw.lake.outflow[v]
-        sw.q_av[i] += sw.q[i] * dt
-    end
-    if update_h
-        @batch per = thread minbatch = 2000 for i in sw.active_n
-            q_src = sum_at(sw.q, links_at_node.src[i])
-            q_dst = sum_at(sw.q, links_at_node.dst[i])
-            sw.volume[i] =
-                sw.volume[i] + (q_src - q_dst + sw.inwater[i] - sw.abstraction[i]) * dt
-
-            if sw.volume[i] < 0.0
-                sw.error[i] = sw.error[i] + abs(sw.volume[i])
-                sw.volume[i] = 0.0 # set volume to zero
-            end
-            sw.volume[i] = max(sw.volume[i] + sw.inflow[i] * dt, 0.0) # add external inflow
-
-            if !isnothing(sw.floodplain)
-                q_src = sum_at(sw.floodplain.q, links_at_node.src[i])
-                q_dst = sum_at(sw.floodplain.q, links_at_node.dst[i])
-                sw.floodplain.volume[i] = sw.floodplain.volume[i] + (q_src - q_dst) * dt
-                # TODO check following approach:
-                # if floodplain volume negative, extract from river volume first
-                if sw.floodplain.volume[i] < 0.0
-                    sw.floodplain.error[i] =
-                        sw.floodplain.error[i] + abs(sw.floodplain.volume[i])
-                    sw.floodplain.volume[i] = 0.0
-                end
-                volume_total = sw.volume[i] + sw.floodplain.volume[i]
-                if volume_total > sw.bankfull_volume[i]
-                    flood_volume = volume_total - sw.bankfull_volume[i]
-                    h = flood_depth(sw.floodplain.profile, flood_volume, sw.dl[i], i)
-                    sw.h[i] = sw.bankfull_depth[i] + h
-                    sw.volume[i] = sw.h[i] * sw.width[i] * sw.dl[i]
-                    sw.floodplain.volume[i] = max(volume_total - sw.volume[i], 0.0)
-                    sw.floodplain.h[i] = sw.floodplain.volume[i] > 0.0 ? h : 0.0
-                else
-                    sw.h[i] = volume_total / (sw.dl[i] * sw.width[i])
-                    sw.volume[i] = volume_total
-                    sw.floodplain.h[i] = 0.0
-                    sw.floodplain.volume[i] = 0.0
-                end
-                sw.floodplain.h_av[i] += sw.floodplain.h[i] * dt
-            else
-                sw.h[i] = sw.volume[i] / (sw.dl[i] * sw.width[i])
-            end
-            sw.h_av[i] += sw.h[i] * dt
-        end
-    end
-    return nothing
-end
-
-function update!(sw::ShallowWaterRiver{T}, network, doy; update_h = true) where {T}
-    if !isnothing(sw.reservoir)
-        sw.reservoir.inflow .= 0.0
-        sw.reservoir.totaloutflow .= 0.0
-        sw.reservoir.actevap .= 0.0
-    end
-    if !isnothing(sw.lake)
-        sw.lake.inflow .= 0.0
-        sw.lake.totaloutflow .= 0.0
-        sw.lake.actevap .= 0.0
-    end
-    if !isnothing(sw.floodplain)
-        sw.floodplain.q_av .= 0.0
-        sw.floodplain.h_av .= 0.0
-    end
-    sw.q_av .= 0.0
-    sw.h_av .= 0.0
-
-    t = T(0.0)
-    while t < sw.dt
-        dt = stable_timestep(sw)
-        if t + dt > sw.dt
-            dt = sw.dt - t
-        end
-        shallowwater_river_update!(sw, network, dt, doy, update_h)
-        t = t + dt
-    end
-    sw.q_av ./= sw.dt
-    sw.h_av ./= sw.dt
-
-    if !isnothing(sw.floodplain)
-        sw.floodplain.q_av ./= sw.dt
-        sw.floodplain.h_av ./= sw.dt
-        sw.q_channel_av .= sw.q_av
-        sw.q_av .= sw.q_channel_av .+ sw.floodplain.q_av
-    end
-
-    return nothing
-end
-
-# Stores links 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)
-
-@get_units @grid_loc @with_kw struct ShallowWaterLand{T}
-    n::Int                                              # number of cells [-]
-    xl::Vector{T} | "m"                                 # cell length x direction [m]
-    yl::Vector{T} | "m"                                 # cell length y direction [m]
-    xwidth::Vector{T} | "m" | "edge"                    # effective flow width x direction (floodplain) [m]
-    ywidth::Vector{T} | "m" | "edge"                    # effective flow width y direction (floodplain) [m]
-    g::T                                                # acceleration due to gravity [m s⁻²]
-    theta::T                                            # weighting factor (de Almeida et al., 2012) [-]
-    alpha::T                                            # stability coefficient (de Almeida et al., 2012) [-]
-    h_thresh::T                                         # depth threshold for calculating flow [m]
-    dt::T                                               # 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                                  # 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(
-    nc,
-    config,
-    inds;
-    modelsize_2d,
-    indices_reverse, # maps from the 2D external domain to the 1D internal domain (Int for linear indexing).
-    xlength,
-    ylength,
-    riverwidth,
-    graph_riv,
-    ldd_riv,
-    inds_riv,
-    river,
-    waterbody,
-    dt,
-)
-    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number
-    alpha = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
-    theta = get(config.model, "inertial_flow_theta", 0.8)::Float64 # weighting factor
-    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at link
-
-    @info "Local inertial approach is used for overlandflow." alpha theta h_thresh froude_limit
-
-    n_land =
-        ncread(nc, config, "lateral.land.n"; sel = inds, defaults = 0.072, type = Float)
-    elevation_2d = ncread(
-        nc,
-        config,
-        "lateral.land.elevation";
-        optional = false,
-        type = Float,
-        fill = 0,
-    )
-    elevation = elevation_2d[inds]
-    n = length(inds)
-
-    # initialize links between cells in x and y direction.
-    indices = Indices(; xu = zeros(n), xd = zeros(n), yu = zeros(n), yd = zeros(n))
-
-    # links without neigbors are handled by an extra index (at n + 1, with n links), which
-    # is set to a value of 0.0 m³ s⁻¹ for qx and qy fields at initialization.
-    # links are defined as follows for the x and y direction, respectively:
-    # node i => node xu (node i + CartesianIndex(1, 0))
-    # node i => node yu (node i + CartesianIndex(0, 1))
-    # where i is the index of inds
-    nrow, ncol = modelsize_2d
-    for (v, i) in enumerate(inds)
-        for (m, neighbor) in enumerate(neighbors)
-            j = i + neighbor
-            dir = dirs[m]
-            if (1 <= j[1] <= nrow) && (1 <= j[2] <= ncol) && (indices_reverse[j] != 0)
-                getfield(indices, dir)[v] = indices_reverse[j]
-            else
-                getfield(indices, dir)[v] = n + 1
-            end
-        end
-    end
-
-    # determine z at links in x and y direction
-    zx_max = fill(Float(0), n)
-    zy_max = fill(Float(0), n)
-    for i in 1:n
-        xu = indices.xu[i]
-        if xu <= n
-            zx_max[i] = max(elevation[i], elevation[xu])
-        end
-        yu = indices.yu[i]
-        if yu <= n
-            zy_max[i] = max(elevation[i], elevation[yu])
-        end
-    end
-
-    # set the effective flow width for river cells in the x and y direction at cell edges.
-    # for waterbody cells (reservoir or lake), h is set to zero (fixed) and not updated, and
-    # overland flow from a downstream cell is not possible (effective flowwidth is zero).
-    we_x = copy(xlength)
-    we_y = copy(ylength)
-    set_effective_flowwidth!(
-        we_x,
-        we_y,
-        indices,
-        graph_riv,
-        riverwidth,
-        ldd_riv,
-        waterbody,
-        indices_reverse[inds_riv],
-    )
-
-    sw_land = ShallowWaterLand{Float}(;
-        n = n,
-        xl = xlength,
-        yl = ylength,
-        xwidth = we_x,
-        ywidth = we_y,
-        g = 9.80665,
-        theta = theta,
-        alpha = alpha,
-        h_thresh = h_thresh,
-        dt = tosecond(dt),
-        qx0 = zeros(n + 1),
-        qy0 = zeros(n + 1),
-        qx = zeros(n + 1),
-        qy = zeros(n + 1),
-        zx_max = zx_max,
-        zy_max = zy_max,
-        mannings_n_sq = n_land .* n_land,
-        volume = zeros(n),
-        error = zeros(n),
-        runoff = zeros(n),
-        inflow_wb = zeros(n),
-        h = zeros(n),
-        h_av = zeros(n),
-        z = elevation,
-        froude_limit = froude_limit,
-        rivercells = river,
-    )
-
-    return sw_land, indices
-end
-
-"""
-    stable_timestep(sw::ShallowWaterRiver)
-    stable_timestep(sw::ShallowWaterLand)
-
-Compute a stable timestep size for the local inertial approach, based on Bates et al. (2010).
-
-dt = alpha * (Δx / sqrt(g max(h))
-"""
-function stable_timestep(sw::ShallowWaterRiver{T})::T where {T}
-    dt_min = T(Inf)
-    @batch per = thread reduction = ((min, dt_min),) for i in 1:(sw.n)
-        @fastmath @inbounds dt = sw.alpha * sw.dl[i] / sqrt(sw.g * sw.h[i])
-        dt_min = min(dt, dt_min)
-    end
-    dt_min = isinf(dt_min) ? T(10.0) : dt_min
-    return dt_min
-end
-
-function stable_timestep(sw::ShallowWaterLand{T})::T where {T}
-    dt_min = T(Inf)
-    @batch per = thread reduction = ((min, dt_min),) for i in 1:(sw.n)
-        @fastmath @inbounds dt = if sw.rivercells[i] == 0
-            sw.alpha * min(sw.xl[i], sw.yl[i]) / sqrt(sw.g * sw.h[i])
-        else
-            T(Inf)
-        end
-        dt_min = min(dt, dt_min)
-    end
-    dt_min = isinf(dt_min) ? T(10.0) : dt_min
-    return dt_min
-end
-
-function update!(
-    sw::ShallowWaterLand{T},
-    swr::ShallowWaterRiver{T},
-    network,
-    doy;
-    update_h = false,
-) where {T}
-    (; nodes_at_link, links_at_node) = network.river
-
-    if !isnothing(swr.reservoir)
-        swr.reservoir.inflow .= 0.0
-        swr.reservoir.totaloutflow .= 0.0
-        swr.reservoir.actevap .= 0.0
-    end
-    if !isnothing(swr.lake)
-        swr.lake.inflow .= 0.0
-        swr.lake.totaloutflow .= 0.0
-        swr.lake.actevap .= 0.0
-    end
-    swr.q_av .= 0.0
-    swr.h_av .= 0.0
-    sw.h_av .= 0.0
-
-    t = T(0.0)
-    while t < swr.dt
-        dt_river = stable_timestep(swr)
-        dt_land = stable_timestep(sw)
-        dt = min(dt_river, dt_land)
-        if t + dt > swr.dt
-            dt = swr.dt - t
-        end
-        shallowwater_river_update!(swr, network.river, dt, doy, update_h)
-        shallowwater_update!(sw, swr, network, dt)
-        t = t + dt
-    end
-    swr.q_av ./= swr.dt
-    swr.h_av ./= swr.dt
-    sw.h_av ./= sw.dt
-
-    return nothing
-end
-
-function shallowwater_update!(
-    sw::ShallowWaterLand{T},
-    swr::ShallowWaterRiver{T},
-    network,
-    dt,
-) where {T}
-    indices = network.land.staggered_indices
-    inds_riv = network.land.index_river
-
-    (; links_at_node) = network.river
-
-    sw.qx0 .= sw.qx
-    sw.qy0 .= sw.qy
-
-    # update qx
-    @batch per = thread minbatch = 6000 for i in 1:(sw.n)
-        yu = indices.yu[i]
-        yd = indices.yd[i]
-        xu = indices.xu[i]
-        xd = indices.xd[i]
-
-        # the effective flow width is zero when the river width exceeds the cell width (dy
-        # for flow in x dir) and floodplain flow is not calculated.
-        if xu <= sw.n && sw.ywidth[i] != T(0.0)
-            zs_x = sw.z[i] + sw.h[i]
-            zs_xu = sw.z[xu] + sw.h[xu]
-            zs_max = max(zs_x, zs_xu)
-            hf = (zs_max - sw.zx_max[i])
-
-            if hf > sw.h_thresh
-                length = T(0.5) * (sw.xl[i] + sw.xl[xu]) # can be precalculated
-                sw.qx[i] = local_inertial_flow(
-                    sw.theta,
-                    sw.qx0[i],
-                    sw.qx0[xd],
-                    sw.qx0[xu],
-                    zs_x,
-                    zs_xu,
-                    hf,
-                    sw.ywidth[i],
-                    length,
-                    sw.mannings_n_sq[i],
-                    sw.g,
-                    sw.froude_limit,
-                    dt,
-                )
-                # limit qx in case water is not available
-                if sw.h[i] <= T(0.0)
-                    sw.qx[i] = min(sw.qx[i], T(0.0))
-                end
-                if sw.h[xu] <= T(0.0)
-                    sw.qx[i] = max(sw.qx[i], T(0.0))
-                end
-            else
-                sw.qx[i] = T(0.0)
-            end
-        end
-
-        # update qy
-
-        # the effective flow width is zero when the river width exceeds the cell width (dx
-        # for flow in y dir) and floodplain flow is not calculated.
-        if yu <= sw.n && sw.xwidth[i] != T(0.0)
-            zs_y = sw.z[i] + sw.h[i]
-            zs_yu = sw.z[yu] + sw.h[yu]
-            zs_max = max(zs_y, zs_yu)
-            hf = (zs_max - sw.zy_max[i])
-
-            if hf > sw.h_thresh
-                length = T(0.5) * (sw.yl[i] + sw.yl[yu]) # can be precalculated
-                sw.qy[i] = local_inertial_flow(
-                    sw.theta,
-                    sw.qy0[i],
-                    sw.qy0[yd],
-                    sw.qy0[yu],
-                    zs_y,
-                    zs_yu,
-                    hf,
-                    sw.xwidth[i],
-                    length,
-                    sw.mannings_n_sq[i],
-                    sw.g,
-                    sw.froude_limit,
-                    dt,
-                )
-                # limit qy in case water is not available
-                if sw.h[i] <= T(0.0)
-                    sw.qy[i] = min(sw.qy[i], T(0.0))
-                end
-                if sw.h[yu] <= T(0.0)
-                    sw.qy[i] = max(sw.qy[i], T(0.0))
-                end
-            else
-                sw.qy[i] = T(0.0)
-            end
-        end
-    end
-
-    # change in volume and water levels based on horizontal fluxes for river and land cells
-    @batch per = thread minbatch = 6000 for i in 1:(sw.n)
-        yd = indices.yd[i]
-        xd = indices.xd[i]
-
-        if sw.rivercells[i]
-            if swr.waterbody[inds_riv[i]]
-                # for reservoir or lake set inflow from land part, these are boundary points
-                # and update of volume and h is not required
-                swr.inflow_wb[inds_riv[i]] =
-                    sw.inflow_wb[i] +
-                    sw.runoff[i] +
-                    (sw.qx[xd] - sw.qx[i] + sw.qy[yd] - sw.qy[i])
-            else
-                sw.volume[i] +=
-                    (
-                        sum_at(swr.q, links_at_node.src[inds_riv[i]]) -
-                        sum_at(swr.q, links_at_node.dst[inds_riv[i]]) + sw.qx[xd] -
-                        sw.qx[i] + sw.qy[yd] - sw.qy[i] +
-                        swr.inflow[inds_riv[i]] +
-                        sw.runoff[i] - swr.abstraction[inds_riv[i]]
-                    ) * dt
-                if sw.volume[i] < T(0.0)
-                    sw.error[i] = sw.error[i] + abs(sw.volume[i])
-                    sw.volume[i] = T(0.0) # set volume to zero
-                end
-                if sw.volume[i] >= swr.bankfull_volume[inds_riv[i]]
-                    swr.h[inds_riv[i]] =
-                        swr.bankfull_depth[inds_riv[i]] +
-                        (sw.volume[i] - swr.bankfull_volume[inds_riv[i]]) /
-                        (sw.xl[i] * sw.yl[i])
-                    sw.h[i] = swr.h[inds_riv[i]] - swr.bankfull_depth[inds_riv[i]]
-                    swr.volume[inds_riv[i]] =
-                        swr.h[inds_riv[i]] * swr.dl[inds_riv[i]] * swr.width[inds_riv[i]]
-                else
-                    swr.h[inds_riv[i]] =
-                        sw.volume[i] / (swr.dl[inds_riv[i]] * swr.width[inds_riv[i]])
-                    sw.h[i] = T(0.0)
-                    swr.volume[inds_riv[i]] = sw.volume[i]
-                end
-                swr.h_av[inds_riv[i]] += swr.h[inds_riv[i]] * dt
-            end
-        else
-            sw.volume[i] +=
-                (sw.qx[xd] - sw.qx[i] + sw.qy[yd] - sw.qy[i] + sw.runoff[i]) * dt
-            if sw.volume[i] < T(0.0)
-                sw.error[i] = sw.error[i] + abs(sw.volume[i])
-                sw.volume[i] = T(0.0) # set volume to zero
-            end
-            sw.h[i] = sw.volume[i] / (sw.xl[i] * sw.yl[i])
-        end
-        sw.h_av[i] += sw.h[i] * dt
-    end
-    return nothing
-end
-
-"""
-    FloodPlainProfile
-
-Floodplain `volume` is a function of `depth` (flood depth intervals). Based on the
-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 @grid_loc @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 @grid_loc @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" | "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
-end
-
-"Determine the initial floodplain volume"
-function initialize_volume!(river, nriv::Int)
-    for i in 1:nriv
-        i1, i2 =
-            interpolation_indices(river.floodplain.h[i], river.floodplain.profile.depth)
-        a = flow_area(
-            river.floodplain.profile.width[i2, i],
-            river.floodplain.profile.a[i1, i],
-            river.floodplain.profile.depth[i1],
-            river.floodplain.h[i],
-        )
-        a = max(a - (river.width[i] * river.floodplain.h[i]), 0.0)
-        river.floodplain.volume[i] = river.dl[i] * a
-    end
-    return nothing
-end
-
-"helper function to get interpolation indices"
-function interpolation_indices(x, v::AbstractVector)
-    i1 = 1
-    for i in eachindex(v)
-        if v[i] <= x
-            i1 = i
-        end
-    end
-    if i1 == length(v)
-        i2 = i1
-    else
-        i2 = i1 + 1
-    end
-    return i1, i2
-end
-
-"""
-    flow_area(width, area, depth, h)
-
-Compute floodplain flow area based on flow depth `h` and floodplain `depth`, `area` and
-`width` of a floodplain profile.
-"""
-function flow_area(width, area, depth, h)
-    dh = h - depth  # depth at i1
-    area = area + (width * dh) # area at i1, width at i2
-    return area
-end
-
-"""
-    function wetted_perimeter(p, depth, h)
-
-Compute floodplain wetted perimeter based on flow depth `h` and floodplain `depth` and
-wetted perimeter `p` of a floodplain profile.
-"""
-function wetted_perimeter(p, depth, h)
-    dh = h - depth # depth at i1
-    p = p + (2.0 * dh) # p at i1
-    return p
-end
-
-"Compute flood depth by interpolating flood volume `flood_volume` using flood depth intervals."
-function flood_depth(profile::FloodPlainProfile{T}, flood_volume, dl, i::Int)::T where {T}
-    i1, i2 = interpolation_indices(flood_volume, @view profile.volume[:, i])
-    ΔA = (flood_volume - profile.volume[i1, i]) / dl
-    dh = ΔA / profile.width[i2, i]
-    flood_depth = profile.depth[i1] + dh
-    return flood_depth
-end
-
-"Initialize floodplain geometry and `FloodPlain` parameters"
-function initialize_floodplain_1d(
-    nc,
-    config,
-    inds,
-    riverwidth,
-    riverlength,
-    zb,
-    index_pit,
-    n_edges,
-    nodes_at_link,
-)
-    n_floodplain = ncread(
-        nc,
-        config,
-        "lateral.river.floodplain.n";
-        sel = inds,
-        defaults = 0.072,
-        type = Float,
-    )
-    volume = ncread(
-        nc,
-        config,
-        "lateral.river.floodplain.volume";
-        sel = inds,
-        type = Float,
-        dimname = :flood_depth,
-    )
-    n = length(inds)
-
-    # for convenience (interpolation) flood depth 0.0 m is added, with associated area (a),
-    # volume, width (river width) and wetted perimeter (p).
-    volume = vcat(fill(Float(0), n)', volume)
-    start_volume = volume
-    flood_depths = Float.(nc["flood_depth"][:])
-    pushfirst!(flood_depths, 0.0)
-    n_depths = length(flood_depths)
-
-    p = zeros(Float, n_depths, n)
-    a = zeros(Float, n_depths, n)
-    segment_volume = zeros(Float, n_depths, n)
-    width = zeros(Float, n_depths, n)
-    width[1, :] = riverwidth[1:n]
-
-    # determine flow area (a), width and wetted perimeter (p)FloodPlain
-    h = diff(flood_depths)
-    incorrect_vol = 0
-    riv_cells = 0
-    error_vol = 0
-    for i in 1:n
-        riv_cell = 0
-        diff_volume = diff(volume[:, i])
-
-        for j in 1:(n_depths - 1)
-            # assume rectangular shape of flood depth segment
-            width[j + 1, i] = diff_volume[j] / (h[j] * riverlength[i])
-            # check provided flood volume (floodplain width should be constant or increasing
-            # as a function of flood depth)
-            if width[j + 1, i] < width[j, i]
-                # raise warning only if difference is larger than rounding error of 0.01 m³
-                if ((width[j, i] - width[j + 1, i]) * h[j] * riverlength[i]) > 0.01
-                    incorrect_vol += 1
-                    riv_cell = 1
-                    error_vol =
-                        error_vol +
-                        ((width[j, i] - width[j + 1, i]) * h[j] * riverlength[i])
-                end
-                width[j + 1, i] = width[j, i]
-            end
-            a[j + 1, i] = width[j + 1, i] * h[j]
-            p[j + 1, i] = (width[j + 1, i] - width[j, i]) + 2.0 * h[j]
-            segment_volume[j + 1, i] = a[j + 1, i] * riverlength[i]
-            if j == 1
-                # for interpolation wetted perimeter at flood depth 0.0 is required
-                p[j, i] = p[j + 1, i] - 2.0 * h[j]
-            end
-        end
-
-        p[2:end, i] = cumsum(p[2:end, i])
-        a[:, i] = cumsum(a[:, i])
-        volume[:, i] = cumsum(segment_volume[:, i])
-
-        riv_cells += riv_cell
-    end
-
-    if incorrect_vol > 0
-        perc_riv_cells = round(100.0 * (riv_cells / n); digits = 2)
-        perc_error_vol = round(100.0 * (error_vol / sum(start_volume[end, :])); digits = 2)
-        @warn string(
-            "The provided volume of $incorrect_vol rectangular floodplain schematization",
-            " segments for $riv_cells river cells ($perc_riv_cells % of total river cells)",
-            " is not correct and has been increased with $perc_error_vol % of provided volume.",
-        )
-    end
-
-    # set floodplain parameters for ghost points
-    volume = hcat(volume, volume[:, index_pit])
-    width = hcat(width, width[:, index_pit])
-    a = hcat(a, a[:, index_pit])
-    p = hcat(p, p[:, index_pit])
-
-    # initialize floodplain profile parameters
-    profile = FloodPlainProfile{Float, n_depths}(;
-        volume = volume,
-        width = width,
-        depth = flood_depths,
-        a = a,
-        p = p,
-    )
-
-    # manning roughness at edges
-    append!(n_floodplain, n_floodplain[index_pit]) # copy to ghost nodes
-    mannings_n_sq = fill(Float(0), n_edges)
-    zb_max = fill(Float(0), n_edges)
-    for i in 1:n_edges
-        src_node = nodes_at_link.src[i]
-        dst_node = nodes_at_link.dst[i]
-        mannings_n =
-            (
-                n_floodplain[dst_node] * riverlength[dst_node] +
-                n_floodplain[src_node] * riverlength[src_node]
-            ) / (riverlength[dst_node] + riverlength[src_node])
-        mannings_n_sq[i] = mannings_n * mannings_n
-        zb_max[i] = max(zb[src_node], zb[dst_node])
-    end
-
-    floodplain = FloodPlain(;
-        profile = profile,
-        mannings_n = n_floodplain,
-        mannings_n_sq = mannings_n_sq,
-        volume = zeros(n),
-        error = zeros(n),
-        h = zeros(n + length(index_pit)),
-        h_av = zeros(n),
-        a = zeros(n_edges),
-        r = zeros(n_edges),
-        hf = zeros(n_edges),
-        zb_max = zb_max,
-        q = zeros(n_edges),
-        q_av = zeros(n_edges),
-        q0 = zeros(n_edges),
-        hf_index = zeros(Int, n_edges),
-    )
-    return floodplain
-end
-
-"""
-    set_river_inwater!(model::Model, ssf_toriver)
-
-Set `inwater` of the lateral river component for a `Model`. `ssf_toriver` is the subsurface
-flow to the river.
-"""
-function set_river_inwater!(model::Model, ssf_toriver)
-    (; lateral, vertical, network, config) = model
-    (; net_runoff_river) = vertical.runoff.variables
-    inds = network.index_river
-    do_water_demand = haskey(config.model, "water_demand")
-    if do_water_demand
-        @. lateral.river.inwater = (
-            ssf_toriver[inds] +
-            lateral.land.to_river[inds] +
-            # net_runoff_river
-            (net_runoff_river[inds] * network.land.area[inds] * 0.001) / vertical.dt +
-            (
-                lateral.river.allocation.variables.nonirri_returnflow *
-                0.001 *
-                network.river.area
-            ) / vertical.dt
-        )
-    else
-        @. lateral.river.inwater = (
-            ssf_toriver[inds] +
-            lateral.land.to_river[inds] +
-            # net_runoff_river
-            (net_runoff_river[inds] * network.land.area[inds] * 0.001) / vertical.dt
-        )
-    end
-    return nothing
-end
-
-"""
-    set_land_inwater!(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
-
-Set `inwater` of the lateral land component for the `SbmGwfModel` type.
-"""
-function set_land_inwater!(
-    model::Model{N, L, V, R, W, T},
-) where {N, L, V, R, W, T <: SbmGwfModel}
-    (; lateral, vertical, network, config) = model
-    (; net_runoff) = vertical.soil.variables
-    do_drains = get(config.model, "drains", false)::Bool
-    drainflux = zeros(length(net_runoff))
-    do_water_demand = haskey(config.model, "water_demand")
-    if do_drains
-        drainflux[lateral.subsurface.drain.index] =
-            -lateral.subsurface.drain.flux ./ tosecond(basetimestep)
-    end
-    if do_water_demand
-        @. lateral.land.inwater =
-            (net_runoff + vertical.allocation.variables.nonirri_returnflow) *
-            network.land.area *
-            0.001 / lateral.land.dt + drainflux
-    else
-        @. lateral.land.inwater =
-            (net_runoff * network.land.area * 0.001) / lateral.land.dt + drainflux
-    end
-    return nothing
-end
-
-"""
-    set_land_inwater!(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
-
-Set `inwater` of the lateral land component for the `SbmModel` type.
-"""
-function set_land_inwater!(
-    model::Model{N, L, V, R, W, T},
-) where {N, L, V, R, W, T <: SbmModel}
-    (; lateral, vertical, network, config) = model
-    (; net_runoff) = vertical.soil.variables
-    do_water_demand = haskey(config.model, "water_demand")
-    if do_water_demand
-        @. lateral.land.inwater =
-            (net_runoff + vertical.allocation.variables.nonirri_returnflow) *
-            network.land.area *
-            0.001 / lateral.land.dt
-    else
-        @. lateral.land.inwater = (net_runoff * network.land.area * 0.001) / lateral.land.dt
-    end
-    return nothing
-end
-
-# Computation of inflow from the lateral components `land` and `subsurface` to water bodies
-# depends on the routing scheme (see different `get_inflow_waterbody` below). For the river
-# kinematic wave, the variables `to_river` can be excluded, because this part is added to
-# the river kinematic wave (kinematic wave is also solved for the water body cell). For
-# local inertial river routing, `to_river` is included, because for the local inertial
-# solution the water body cells are excluded (boundary condition). For `GroundwaterFlow`
-# (Darcian flow in 4 directions), the lateral subsurface flow is excluded (for now) and
-# inflow consists of overland flow.
-"""
-    set_inflow_waterbody!(
-        model::Model{N,L,V,R,W,T},
-    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,SurfaceFlow,SurfaceFlow}},V,R,W,T}
-
-Set inflow from the subsurface and land components to a water body (reservoir or lake)
-`inflow_wb` from a model type that contains the lateral components `SurfaceFlow`.
-"""
-function set_inflow_waterbody!(
-    model::Model{N, L, V, R, W, T},
-) where {N, L <: NamedTuple{<:Any, <:Tuple{Any, SurfaceFlow, SurfaceFlow}}, V, R, W, T}
-    (; lateral, network) = model
-    (; subsurface, land, river) = lateral
-    inds = network.index_river
-
-    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
-        if typeof(subsurface) <: LateralSSF || typeof(subsurface) <: GroundwaterExchange
-            @. river.inflow_wb =
-                subsurface.ssf[inds] / tosecond(basetimestep) + land.q_av[inds]
-        elseif typof(subsurface.flow) <: GroundwaterFlow || isnothing(subsurface)
-            river.inflow_wb .= land.q_av[inds]
-        end
-    end
-    return nothing
-end
-
-"""
-    set_inflow_waterbody!(
-        model::Model{N,L,V,R,W,T},
-    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,SurfaceFlow,ShallowWaterRiver}},V,R,W,T}
-
-Set inflow from the subsurface and land components to a water body (reservoir or lake)
-`inflow_wb` from a model type that contains the lateral components `SurfaceFlow` and
-`ShallowWaterRiver`.
-"""
-function set_inflow_waterbody!(
-    model::Model{N, L, V, R, W, T},
-) where {
-    N,
-    L <: NamedTuple{<:Any, <:Tuple{Any, SurfaceFlow, ShallowWaterRiver}},
-    V,
-    R,
-    W,
-    T,
-}
-    (; lateral, network) = model
-    (; subsurface, land, river) = lateral
-    inds = network.index_river
-
-    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
-        if typeof(subsurface) <: LateralSSF || typeof(subsurface) <: GroundwaterExchange
-            @. river.inflow_wb =
-                (subsurface.ssf[inds] + subsurface.to_river[inds]) /
-                tosecond(basetimestep) +
-                land.q_av[inds] +
-                land.to_river[inds]
-        elseif typeof(subsurface.flow) <: GroundwaterFlow || isnothing(subsurface)
-            @. river.inflow_wb = lateral.land.q_av[inds] + lateral.land.to_river[inds]
-        end
-    end
-    return nothing
-end
-
-"""
-    set_inflow_waterbody!(
-        model::Model{N,L,V,R,W,T},
-    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,ShallowWaterLand,ShallowWaterRiver}},V,R,W,T}
-
-Set inflow from the subsurface and land components to a water body (reservoir or lake)
-`inflow_wb` from a model type that contains the lateral components `ShallowWaterLand` and
-`ShallowWaterRiver`.
-"""
-function set_inflow_waterbody!(
-    model::Model{N, L, V, R, W, T},
-) where {
-    N,
-    L <: NamedTuple{<:Any, <:Tuple{Any, ShallowWaterLand, ShallowWaterRiver}},
-    V,
-    R,
-    W,
-    T,
-}
-    (; lateral, network) = model
-    (; subsurface, land, river) = lateral
-    inds = network.index_river
-
-    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
-        if typeof(subsurface) <: LateralSSF || typeof(subsurface) <: GroundwaterExchange
-            @. land.inflow_wb[inds] =
-                (subsurface.ssf[inds] + subsurface.to_river[inds]) / tosecond(basetimestep)
-        end
-    end
-    return nothing
-end
-
-"""
-    surface_routing!(model; ssf_toriver = 0.0)
-
-Run surface routing (land and river). Kinematic wave for overland flow and kinematic wave or
-local inertial model for river flow.
-"""
-function surface_routing!(model; ssf_toriver = 0.0)
-    (; lateral, network, clock) = model
-
-    # run kinematic wave for overland flow
-    set_land_inwater!(model)
-    update!(lateral.land, network.land, network.frac_toriver)
-
-    # run river flow
-    set_river_inwater!(model, ssf_toriver)
-    set_inflow_waterbody!(model)
-    update!(lateral.river, network.river, julian_day(clock.time - clock.dt))
-    return nothing
-end
-
-"""
-    surface_routing(
-        model::Model{N,L,V,R,W,T};
-        ssf_toriver = 0.0,
-    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,ShallowWaterLand,ShallowWaterRiver}},V,R,W,T}
-
-Run surface routing (land and river) for a model type that contains the lateral components
-`ShallowWaterLand` and `ShallowWaterRiver`.
-"""
-function surface_routing!(
-    model::Model{N, L, V, R, W, T};
-    ssf_toriver = 0.0,
-) where {
-    N,
-    L <: NamedTuple{<:Any, <:Tuple{Any, ShallowWaterLand, ShallowWaterRiver}},
-    V,
-    R,
-    W,
-    T,
-}
-    (; lateral, vertical, network, clock) = model
-    (; net_runoff) = vertical.soil.variables
-    (; net_runoff_river) = vertical.runoff.variables
-
-    @. lateral.land.runoff = (
-        (net_runoff / 1000.0) * (network.land.area) / vertical.dt +
-        ssf_toriver +
-        # net_runoff_river
-        ((net_runoff_river * network.land.area * 0.001) / vertical.dt)
-    )
-    set_inflow_waterbody!(model)
-    update!(lateral.land, lateral.river, network, julian_day(clock.time - clock.dt))
-    return nothing
-end
diff --git a/src/glacier/glacier.jl b/src/glacier/glacier.jl
index ac2c83633..38d0be93a 100644
--- a/src/glacier/glacier.jl
+++ b/src/glacier/glacier.jl
@@ -9,12 +9,12 @@ abstract type AbstractGlacierModel{T} end
 end
 
 "Initialize glacier model variables"
-function GlacierVariables(nc, config, inds)
+function GlacierVariables(dataset, config, indices)
     glacier_store = ncread(
-        nc,
+        dataset,
         config,
         "vertical.glacier.variables.glacier_store";
-        sel = inds,
+        sel = indices,
         defaults = 5500.0,
         type = Float,
         fill = 0.0,
@@ -54,41 +54,41 @@ end
 struct NoGlacierModel{T} <: AbstractGlacierModel{T} end
 
 "Initialize glacier HBV model parameters"
-function GlacierHbvParameters(nc, config, inds, dt)
+function GlacierHbvParameters(dataset, config, indices, dt)
     g_ttm = ncread(
-        nc,
+        dataset,
         config,
         "vertical.glacier.parameters.g_ttm";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
         fill = 0.0,
     )
     g_cfmax =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.glacier.parameters.g_cfmax";
-            sel = inds,
+            sel = indices,
             defaults = 3.0,
             type = Float,
             fill = 0.0,
         ) .* (dt / basetimestep)
     g_sifrac =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.glacier.parameters.g_sifrac";
-            sel = inds,
+            sel = indices,
             defaults = 0.001,
             type = Float,
             fill = 0.0,
         ) .* (dt / basetimestep)
     glacier_frac = ncread(
-        nc,
+        dataset,
         config,
         "vertical.glacier.parameters.glacier_frac";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
         fill = 0.0,
@@ -100,9 +100,9 @@ function GlacierHbvParameters(nc, config, inds, dt)
 end
 
 "Initialize glacier HBV model"
-function GlacierHbvModel(nc, config, inds, dt, bc)
-    params = GlacierHbvParameters(nc, config, inds, dt)
-    vars = GlacierVariables(nc, config, inds)
+function GlacierHbvModel(dataset, config, indices, dt, bc)
+    params = GlacierHbvParameters(dataset, config, indices, dt)
+    vars = GlacierVariables(dataset, config, indices)
     model =
         GlacierHbvModel(; boundary_conditions = bc, parameters = params, variables = vars)
     return model
diff --git a/src/groundwater/aquifer.jl b/src/groundwater/aquifer.jl
index 3a8619bf6..7b6689f4c 100644
--- a/src/groundwater/aquifer.jl
+++ b/src/groundwater/aquifer.jl
@@ -84,18 +84,26 @@ 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 @grid_loc struct ConfinedAquifer{T} <: Aquifer
-    head::Vector{T} | "m"  # hydraulic head [m]
+@get_units @grid_loc @with_kw struct ConfinedAquiferParameters{T}
     k::Vector{T} | "m d-1"  # horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
     bottom::Vector{T} | "m" # bottom of groundwater layer
     area::Vector{T} | "m2" # area of cell
     specific_storage::Vector{T} | "m m-1 m-1" # [m m⁻¹ m⁻¹]
     storativity::Vector{T} | "m m-1" # [m m⁻¹]
+end
+
+@get_units @grid_loc @with_kw struct ConfinedAquiferVariables{T}
+    head::Vector{T} | "m"  # hydraulic head [m]
     conductance::Vector{T} | "m2 d-1" # Confined aquifer conductance is constant
     volume::Vector{T} | "m3" # total volume of water that can be released
 end
 
+@with_kw struct ConfinedAquifer{T} <: Aquifer
+    parameters::ConfinedAquiferParameters{T}
+    variables::ConfinedAquiferVariables{T}
+end
+
 """
     UnconfinedAquifer{T} <: Aquifer
 
@@ -107,22 +115,68 @@ 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 @grid_loc struct UnconfinedAquifer{T} <: Aquifer
-    head::Vector{T} | "m"  # hydraulic head [m]
+@get_units @grid_loc @with_kw struct UnconfinedAquiferParameters{T}
     k::Vector{T} | "m d-1"  # reference horizontal conductivity [m d⁻¹]
     top::Vector{T} | "m" # top of groundwater layer [m]
     bottom::Vector{T} | "m" # bottom of groundwater layer
     area::Vector{T} | "m2"
     specific_yield::Vector{T} | "m m-1" # [m m⁻¹]
-    conductance::Vector{T} | "m2 d-1" #
-    volume::Vector{T} | "m3" # total volume of water that can be released
     f::Vector{T} | "-" # factor controlling the reduction of reference horizontal conductivity
     # Unconfined aquifer conductance is computed with degree of saturation (only when
     # conductivity_profile is set to "exponential")
 end
 
-storativity(A::UnconfinedAquifer) = A.specific_yield
-storativity(A::ConfinedAquifer) = A.storativity
+function UnconfinedAquiferParameters(dataset, config, indices, top, bottom, area)
+    k = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.conductivity";
+        sel = indices,
+        type = Float,
+    )
+    specific_yield = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.specific_yield";
+        sel = indices,
+        type = Float,
+    )
+    f = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.gwf_f";
+        sel = indices,
+        type = Float,
+        defaults = 3.0,
+    )
+
+    parameters =
+        UnconfinedAquiferParameters{Float}(; k, top, bottom, area, specific_yield, f)
+    return parameters
+end
+
+@get_units @grid_loc @with_kw struct UnconfinedAquiferVariables{T}
+    head::Vector{T} | "m"  # hydraulic head [m]
+    conductance::Vector{T} | "m2 d-1" # conductance 
+    volume::Vector{T} | "m3" # total volume of water that can be released m   
+end
+
+@with_kw struct UnconfinedAquifer{T} <: Aquifer
+    parameters::UnconfinedAquiferParameters{T}
+    variables::UnconfinedAquiferVariables{T}
+end
+
+function UnconfinedAquifer(dataset, config, indices, top, bottom, area, conductance, head)
+    parameters = UnconfinedAquiferParameters(dataset, config, indices, top, bottom, area)
+
+    volume = @. (min(top, head) - bottom) * area * parameters.specific_yield
+    variables = UnconfinedAquiferVariables{Float}(head, conductance, volume)
+    aquifer = UnconfinedAquifer{Float}(parameters, variables)
+    return aquifer
+end
+
+storativity(A::UnconfinedAquifer) = A.parameters.specific_yield
+storativity(A::ConfinedAquifer) = A.parameters.storativity
 
 """
     harmonicmean_conductance(kH1, kH2, l1, l2, width)
@@ -147,19 +201,20 @@ function harmonicmean_conductance(kH1, kH2, l1, l2, width)
 end
 
 function saturated_thickness(aquifer::UnconfinedAquifer, index::Int)
-    return min(aquifer.top[index], aquifer.head[index]) - aquifer.bottom[index]
+    return min(aquifer.parameters.top[index], aquifer.variables.head[index]) -
+           aquifer.parameters.bottom[index]
 end
 
 function saturated_thickness(aquifer::ConfinedAquifer, index::Int)
-    return aquifer.top[index] - aquifer.bottom[index]
+    return aquifer.parameters.top[index] - aquifer.parameters.bottom[index]
 end
 
 function saturated_thickness(aquifer::UnconfinedAquifer)
-    @. min(aquifer.top, aquifer.head) - aquifer.bottom
+    @. min(aquifer.parameters.top, aquifer.variables.head) - aquifer.parameters.bottom
 end
 
 function saturated_thickness(aquifer::ConfinedAquifer)
-    @. aquifer.top - aquifer.bottom
+    @. aquifer.parameters.top - aquifer.parameters.bottom
 end
 
 """
@@ -177,10 +232,10 @@ function horizontal_conductance(
     aquifer::A,
     connectivity::Connectivity,
 ) where {A <: Aquifer}
-    k1 = aquifer.k[i]
-    k2 = aquifer.k[j]
-    H1 = aquifer.top[i] - aquifer.bottom[i]
-    H2 = aquifer.top[j] - aquifer.bottom[j]
+    k1 = aquifer.parameters.k[i]
+    k2 = aquifer.parameters.k[j]
+    H1 = aquifer.parameters.top[i] - aquifer.parameters.bottom[i]
+    H2 = aquifer.parameters.top[j] - aquifer.parameters.bottom[j]
     length1 = connectivity.length1[nzi]
     length2 = connectivity.length2[nzi]
     width = connectivity.width[nzi]
@@ -204,7 +259,7 @@ function initialize_conductance!(
         # Loop over connections for cell j
         for nzi in connections(connectivity, i)
             j = connectivity.rowval[nzi]
-            aquifer.conductance[nzi] =
+            aquifer.variables.conductance[nzi] =
                 horizontal_conductance(i, j, nzi, aquifer, connectivity)
         end
     end
@@ -218,7 +273,7 @@ function conductance(
     conductivity_profile::String,
     connectivity::Connectivity,
 )
-    return aquifer.conductance[nzi]
+    return aquifer.variables.conductance[nzi]
 end
 
 """
@@ -257,17 +312,21 @@ function conductance(
 )
     if conductivity_profile == "exponential"
         # Extract required variables
-        zi1 = aquifer.top[i] - aquifer.head[i]
-        zi2 = aquifer.top[j] - aquifer.head[j]
-        thickness1 = aquifer.top[i] - aquifer.bottom[i]
-        thickness2 = aquifer.top[j] - aquifer.bottom[j]
+        zi1 = aquifer.parameters.top[i] - aquifer.variables.head[i]
+        zi2 = aquifer.parameters.top[j] - aquifer.variables.head[j]
+        thickness1 = aquifer.parameters.top[i] - aquifer.parameters.bottom[i]
+        thickness2 = aquifer.parameters.top[j] - aquifer.parameters.bottom[j]
         # calculate conductivity values corrected for depth of water table
         k1 =
-            (aquifer.k[i] / aquifer.f[i]) *
-            (exp(-aquifer.f[i] * zi1) - exp(-aquifer.f[i] * thickness1))
+            (aquifer.parameters.k[i] / aquifer.parameters.f[i]) * (
+                exp(-aquifer.parameters.f[i] * zi1) -
+                exp(-aquifer.parameters.f[i] * thickness1)
+            )
         k2 =
-            (aquifer.k[j] / aquifer.f[j]) *
-            (exp(-aquifer.f[j] * zi2) - exp(-aquifer.f[j] * thickness2))
+            (aquifer.parameters.k[j] / aquifer.parameters.f[j]) * (
+                exp(-aquifer.parameters.f[j] * zi2) -
+                exp(-aquifer.parameters.f[j] * thickness2)
+            )
         return harmonicmean_conductance(
             k1,
             k2,
@@ -276,16 +335,18 @@ function conductance(
             connectivity.width[nzi],
         )
     elseif conductivity_profile == "uniform"
-        head_i = aquifer.head[i]
-        head_j = aquifer.head[j]
+        head_i = aquifer.variables.head[i]
+        head_j = aquifer.variables.head[j]
         if head_i >= head_j
             saturation =
-                saturated_thickness(aquifer, i) / (aquifer.top[i] - aquifer.bottom[i])
+                saturated_thickness(aquifer, i) /
+                (aquifer.parameters.top[i] - aquifer.parameters.bottom[i])
         else
             saturation =
-                saturated_thickness(aquifer, j) / (aquifer.top[j] - aquifer.bottom[j])
+                saturated_thickness(aquifer, j) /
+                (aquifer.parameters.top[j] - aquifer.parameters.bottom[j])
         end
-        return saturation * aquifer.conductance[nzi]
+        return saturation * aquifer.variables.conductance[nzi]
     else
         error(
             """An unknown "conductivity_profile" is specified in the TOML file ($conductivity_profile).
@@ -301,7 +362,7 @@ function flux!(Q, aquifer, connectivity, conductivity_profile)
         for nzi in connections(connectivity, i)
             # connection from i -> j
             j = connectivity.rowval[nzi]
-            delta_head = aquifer.head[i] - aquifer.head[j]
+            delta_head = aquifer.variables.head[i] - aquifer.variables.head[j]
             cond = conductance(aquifer, i, j, nzi, conductivity_profile, connectivity)
             Q[i] -= cond * delta_head
         end
@@ -309,11 +370,32 @@ function flux!(Q, aquifer, connectivity, conductivity_profile)
     return Q
 end
 
-@get_units @grid_loc struct ConstantHead{T}
+@get_units @grid_loc @with_kw struct ConstantHeadVariables{T}
     head::Vector{T} | "m"
+end
+
+@get_units @grid_loc @with_kw struct ConstantHead{T}
+    variables::ConstantHeadVariables{T}
     index::Vector{Int} | "-"
 end
 
+function ConstantHead(dataset, config, indices)
+    constanthead = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.constant_head";
+        sel = indices,
+        type = Float,
+        fill = mv,
+    )
+    n = length(indices)
+    index_constanthead = filter(i -> !isequal(constanthead[i], mv), 1:n)
+    head = constanthead[index_constanthead]
+    variables = ConstantHeadVariables{Float}(head)
+    constant_head = ConstantHead{Float}(; variables, index = index_constanthead)
+    return constant_head
+end
+
 """
     stable_timestep(aquifer)
 
@@ -324,25 +406,28 @@ The following criterion can be found in Chu & Willis (1984)
 """
 function stable_timestep(aquifer, conductivity_profile::String)
     dt_min = Inf
-    for i in eachindex(aquifer.head)
+    for i in eachindex(aquifer.variables.head)
         if conductivity_profile == "exponential"
-            zi = aquifer.top[i] - aquifer.head[i]
-            thickness = aquifer.top[i] - aquifer.bottom[i]
+            zi = aquifer.parameters.top[i] - aquifer.variables.head[i]
+            thickness = aquifer.parameters.top[i] - aquifer.parameters.bottom[i]
             value =
-                (aquifer.k[i] / aquifer.f[i]) *
-                (exp(-aquifer.f[i] * zi) - exp(-aquifer.f[i] * thickness))
+                (aquifer.parameters.k[i] / aquifer.parameters.f[i]) * (
+                    exp(-aquifer.parameters.f[i] * zi) -
+                    exp(-aquifer.parameters.f[i] * thickness)
+                )
         elseif conductivity_profile == "uniform"
-            value = aquifer.k[i] * saturated_thickness(aquifer, i)
+            value = aquifer.parameters.k[i] * saturated_thickness(aquifer, i)
         end
 
-        dt = aquifer.area[i] * storativity(aquifer)[i] / value
+        dt = aquifer.parameters.area[i] * storativity(aquifer)[i] / value
         dt_min = dt < dt_min ? dt : dt_min
     end
     return 0.25 * dt_min
 end
 
-minimum_head(aquifer::ConfinedAquifer) = aquifer.head
-minimum_head(aquifer::UnconfinedAquifer) = max.(aquifer.head, aquifer.bottom)
+minimum_head(aquifer::ConfinedAquifer) = aquifer.variables.head
+minimum_head(aquifer::UnconfinedAquifer) =
+    max.(aquifer.variables.head, aquifer.parameters.bottom)
 
 function update!(gwf, Q, dt, conductivity_profile)
     Q .= 0.0  # TODO: Probably remove this when linking with other components
@@ -350,13 +435,15 @@ function update!(gwf, Q, dt, conductivity_profile)
     for boundary in gwf.boundaries
         flux!(Q, boundary, gwf.aquifer)
     end
-    gwf.aquifer.head .+= (Q ./ gwf.aquifer.area .* dt ./ storativity(gwf.aquifer))
+    gwf.aquifer.variables.head .+=
+        (Q ./ gwf.aquifer.parameters.area .* dt ./ storativity(gwf.aquifer))
     # Set constant head (dirichlet) boundaries
-    gwf.aquifer.head[gwf.constanthead.index] .= gwf.constanthead.head
+    gwf.aquifer.variables.head[gwf.constanthead.index] .= gwf.constanthead.variables.head
     # Make sure no heads ends up below an unconfined aquifer bottom
-    gwf.aquifer.head .= minimum_head(gwf.aquifer)
-    gwf.aquifer.volume .=
-        saturated_thickness(gwf.aquifer) .* gwf.aquifer.area .* storativity(gwf.aquifer)
+    gwf.aquifer.variables.head .= minimum_head(gwf.aquifer)
+    gwf.aquifer.variables.volume .=
+        saturated_thickness(gwf.aquifer) .* gwf.aquifer.parameters.area .*
+        storativity(gwf.aquifer)
     return nothing
 end
 
@@ -381,9 +468,10 @@ end
 function get_water_depth(
     gwf::GroundwaterFlow{A, C, CH, B},
 ) where {A <: UnconfinedAquifer, C, CH, B}
-    gwf.aquifer.head .= min.(gwf.aquifer.head, gwf.aquifer.top)
-    gwf.aquifer.head[gwf.constanthead.index] .= gwf.constanthead.head
-    wtd = gwf.aquifer.top .- gwf.aquifer.head
+    gwf.aquifer.variables.head .=
+        min.(gwf.aquifer.variables.head, gwf.aquifer.parameters.top)
+    gwf.aquifer.variables.head[gwf.constanthead.index] .= gwf.constanthead.variables.head
+    wtd = gwf.aquifer.parameters.top .- gwf.aquifer.variables.head
     return wtd
 end
 
@@ -391,8 +479,19 @@ function get_exfiltwater(
     gwf::GroundwaterFlow{A, C, CH, B},
 ) where {A <: UnconfinedAquifer, C, CH, B}
     exfiltwater =
-        (gwf.aquifer.head .- min.(gwf.aquifer.head, gwf.aquifer.top)) .*
-        storativity(gwf.aquifer)
+        (
+            gwf.aquifer.variables.head .-
+            min.(gwf.aquifer.variables.head, gwf.aquifer.parameters.top)
+        ) .* storativity(gwf.aquifer)
     exfiltwater[gwf.constanthead.index] .= 0
     return exfiltwater
+end
+
+function get_flux_to_river(subsurface)
+    (; flow, river) = subsurface
+    ncell = flow.connectivity.ncell
+    flux = zeros(ncell)
+    index = river.index
+    flux[index] = -river.variables.flux ./ tosecond(basetimestep) # [m³ s⁻¹]
+    return flux
 end
\ No newline at end of file
diff --git a/src/groundwater/boundary_conditions.jl b/src/groundwater/boundary_conditions.jl
index 48e883527..52927733c 100644
--- a/src/groundwater/boundary_conditions.jl
+++ b/src/groundwater/boundary_conditions.jl
@@ -1,6 +1,6 @@
 function check_flux(flux, aquifer::UnconfinedAquifer, index::Int)
     # Check if cell is dry
-    if aquifer.head[index] <= aquifer.bottom[index]
+    if aquifer.variables.head[index] <= aquifer.parameters.bottom[index]
         # If cell is dry, no negative flux is allowed
         return max(0, flux)
     else
@@ -11,91 +11,196 @@ end
 # Do nothing for a confined aquifer: aquifer can always provide flux
 check_flux(flux, aquifer::ConfinedAquifer, index::Int) = flux
 
-@get_units @grid_loc struct River{T} <: AquiferBoundaryCondition
-    stage::Vector{T} | "m"
+@get_units @grid_loc @with_kw struct RiverParameters{T}
     infiltration_conductance::Vector{T} | "m2 d-1"
     exfiltration_conductance::Vector{T} | "m2 d-1"
     bottom::Vector{T} | "m"
+end
+
+@get_units @grid_loc @with_kw struct RiverVariables{T}
+    stage::Vector{T} | "m"
     flux::Vector{T} | "m3 d-1"
+end
+
+function RiverVariables(n)
+    variables = RiverVariables{Float}(; stage = fill(mv, n), flux = fill(mv, n))
+    return variables
+end
+
+@get_units @grid_loc @with_kw struct River{T} <: AquiferBoundaryCondition
+    parameters::RiverParameters{T}
+    variables::RiverVariables{T}
     index::Vector{Int} | "-"
 end
 
+function River(dataset, config, indices, index)
+    infiltration_conductance = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.infiltration_conductance";
+        sel = indices,
+        type = Float,
+    )
+    exfiltration_conductance = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.exfiltration_conductance";
+        sel = indices,
+        type = Float,
+    )
+    bottom = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.river_bottom";
+        sel = indices,
+        type = Float,
+    )
+
+    parameters =
+        RiverParameters{Float}(infiltration_conductance, exfiltration_conductance, bottom)
+    n = length(indices)
+    variables = RiverVariables(n)
+    river = River(parameters, variables, index)
+    return river
+end
+
 function flux!(Q, river::River, aquifer)
     for (i, index) in enumerate(river.index)
-        head = aquifer.head[index]
-        stage = river.stage[i]
+        head = aquifer.variables.head[index]
+        stage = river.variables.stage[i]
         if stage > head
-            cond = river.infiltration_conductance[i]
-            delta_head = min(stage - river.bottom[i], stage - head)
+            cond = river.parameters.infiltration_conductance[i]
+            delta_head = min(stage - river.parameters.bottom[i], stage - head)
         else
-            cond = river.exfiltration_conductance[i]
+            cond = river.parameters.exfiltration_conductance[i]
             delta_head = stage - head
         end
-        river.flux[i] = check_flux(cond * delta_head, aquifer, index)
-        Q[index] += river.flux[i]
+        river.variables.flux[i] = check_flux(cond * delta_head, aquifer, index)
+        Q[index] += river.variables.flux[i]
     end
     return Q
 end
 
-@get_units @grid_loc struct Drainage{T} <: AquiferBoundaryCondition
+@get_units @grid_loc @with_kw struct DrainageParameters{T}
     elevation::Vector{T} | "m"
     conductance::Vector{T} | "m2 d-1"
+end
+
+@get_units @grid_loc @with_kw struct DrainageVariables{T}
     flux::Vector{T} | "m3 d-1"
+end
+
+@get_units @grid_loc @with_kw struct Drainage{T} <: AquiferBoundaryCondition
+    parameters::DrainageParameters{T}
+    variables::DrainageVariables{T}
     index::Vector{Int} | "-"
 end
 
+function Drainage(dataset, config, indices, index)
+    drain_elevation = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.drain_elevation";
+        sel = indices,
+        type = Float,
+        fill = mv,
+    )
+    drain_conductance = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.drain_conductance";
+        sel = indices,
+        type = Float,
+        fill = mv,
+    )
+    elevation = drain_elevation[index]
+    conductance = drain_conductance[index]
+    parameters = DrainageParameters{Float}(; elevation, conductance)
+    variables = DrainageVariables{Float}(; flux = fill(mv, length(index)))
+
+    drains = Drainage{Float}(parameters, variables, index)
+    return drains
+end
+
 function flux!(Q, drainage::Drainage, aquifer)
     for (i, index) in enumerate(drainage.index)
-        cond = drainage.conductance[i]
-        delta_head = min(0, drainage.elevation[i] - aquifer.head[index])
-        drainage.flux[i] = check_flux(cond * delta_head, aquifer, index)
-        Q[index] += drainage.flux[i]
+        cond = drainage.parameters.conductance[i]
+        delta_head =
+            min(0, drainage.parameters.elevation[i] - aquifer.variables.head[index])
+        drainage.variables.flux[i] = check_flux(cond * delta_head, aquifer, index)
+        Q[index] += drainage.variables.flux[i]
     end
     return Q
 end
 
-@get_units @grid_loc struct HeadBoundary{T} <: AquiferBoundaryCondition
-    head::Vector{T} | "m"
+@get_units @grid_loc @with_kw struct HeadBoundaryParameters{T}
     conductance::Vector{T} | "m2 d-1"
+end
+
+@get_units @grid_loc @with_kw struct HeadBoundaryVariables{T}
+    head::Vector{T} | "m"
     flux::Vector{T} | "m3 d-1"
+end
+
+@get_units @grid_loc @with_kw struct HeadBoundary{T} <: AquiferBoundaryCondition
+    parameters::HeadBoundaryParameters{T}
+    variables::HeadBoundaryVariables{T}
     index::Vector{Int} | "-"
 end
 
 function flux!(Q, headboundary::HeadBoundary, aquifer)
     for (i, index) in enumerate(headboundary.index)
-        cond = headboundary.conductance[i]
-        delta_head = headboundary.head[i] - aquifer.head[index]
-        headboundary.flux[i] = check_flux(cond * delta_head, aquifer, index)
-        Q[index] += headboundary.flux[i]
+        cond = headboundary.parameters.conductance[i]
+        delta_head = headboundary.variables.head[i] - aquifer.variables.head[index]
+        headboundary.variables.flux[i] = check_flux(cond * delta_head, aquifer, index)
+        Q[index] += headboundary.variables.flux[i]
     end
     return Q
 end
 
-@get_units @grid_loc struct Recharge{T} <: AquiferBoundaryCondition
+@get_units @grid_loc @with_kw struct RechargeVariables{T}
     rate::Vector{T} | "m d-1"
     flux::Vector{T} | "m3 d-1"
+end
+
+@get_units @grid_loc @with_kw struct Recharge{T} <: AquiferBoundaryCondition
+    variables::RechargeVariables{T}
     index::Vector{Int} | "-"
 end
 
+function Recharge(rate, flux, index)
+    variables = RechargeVariables{Float}(rate, flux)
+    recharge = Recharge{Float}(variables, index)
+    return recharge
+end
+
 function flux!(Q, recharge::Recharge, aquifer)
     for (i, index) in enumerate(recharge.index)
-        recharge.flux[i] =
-            check_flux(recharge.rate[i] * aquifer.area[index], aquifer, index)
-        Q[index] += recharge.flux[i]
+        recharge.variables.flux[i] = check_flux(
+            recharge.variables.rate[i] * aquifer.parameters.area[index],
+            aquifer,
+            index,
+        )
+        Q[index] += recharge.variables.flux[i]
     end
     return Q
 end
 
-@get_units @grid_loc struct Well{T} <: AquiferBoundaryCondition
+@get_units @grid_loc @with_kw struct WellVariables{T}
     volumetric_rate::Vector{T} | "m3 d-1"
     flux::Vector{T} | "m3 d-1"
+end
+
+@get_units @grid_loc @with_kw struct Well{T} <: AquiferBoundaryCondition
+    variables::WellVariables{T}
     index::Vector{Int} | "-"
 end
 
 function flux!(Q, well::Well, aquifer)
     for (i, index) in enumerate(well.index)
-        well.flux[i] = check_flux(well.volumetric_rate[i], aquifer, index)
-        Q[index] += well.flux[i]
+        well.variables.flux[i] =
+            check_flux(well.variables.volumetric_rate[i], aquifer, index)
+        Q[index] += well.variables.flux[i]
     end
     return Q
 end
diff --git a/src/io.jl b/src/io.jl
index 93ec265a0..90d2da595 100644
--- a/src/io.jl
+++ b/src/io.jl
@@ -184,18 +184,18 @@ end
 function get_param_res(model)
     return Dict(
         symbols"vertical.atmospheric_forcing.precipitation" =>
-            model.lateral.river.reservoir.precipitation,
+            model.lateral.river.boundary_conditions.reservoir.boundary_conditions.precipitation,
         symbols"vertical.atmospheric_forcing.potential_evaporation" =>
-            model.lateral.river.reservoir.evaporation,
+            model.lateral.river.boundary_conditions.reservoir.boundary_conditions.evaporation,
     )
 end
 
 function get_param_lake(model)
     return Dict(
         symbols"vertical.atmospheric_forcing.precipitation" =>
-            model.lateral.river.lake.precipitation,
+            model.lateral.river.boundary_conditions.lake.boundary_conditions.precipitation,
         symbols"vertical.atmospheric_forcing.potential_evaporation" =>
-            model.lateral.river.lake.evaporation,
+            model.lateral.river.boundary_conditions.lake.boundary_conditions.evaporation,
     )
 end
 
diff --git a/src/parameters.jl b/src/parameters.jl
index f2d2d582a..91dc3e662 100644
--- a/src/parameters.jl
+++ b/src/parameters.jl
@@ -19,47 +19,47 @@
 end
 
 "Initialize (shared) vegetation parameters"
-function VegetationParameters(nc, config, inds)
-    n = length(inds)
+function VegetationParameters(dataset, config, indices)
+    n = length(indices)
     rootingdepth = ncread(
-        nc,
+        dataset,
         config,
         "vertical.vegetation_parameter_set.rootingdepth";
-        sel = inds,
+        sel = indices,
         defaults = 750.0,
         type = Float,
     )
     kc = ncread(
-        nc,
+        dataset,
         config,
         "vertical.vegetation_parameter_set.kc";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
     )
     if haskey(config.input.vertical.vegetation_parameter_set, "leaf_area_index")
         storage_specific_leaf = ncread(
-            nc,
+            dataset,
             config,
             "vertical.vegetation_parameter_set.storage_specific_leaf";
             optional = false,
-            sel = inds,
+            sel = indices,
             type = Float,
         )
         storage_wood = ncread(
-            nc,
+            dataset,
             config,
             "vertical.vegetation_parameter_set.storage_wood";
             optional = false,
-            sel = inds,
+            sel = indices,
             type = Float,
         )
         kext = ncread(
-            nc,
+            dataset,
             config,
             "vertical.vegetation_parameter_set.kext";
             optional = false,
-            sel = inds,
+            sel = indices,
             type = Float,
         )
         vegetation_parameter_set = VegetationParameters(;
@@ -74,18 +74,18 @@ function VegetationParameters(nc, config, inds)
         )
     else
         canopygapfraction = ncread(
-            nc,
+            dataset,
             config,
             "vertical.vegetation_parameter_set.canopygapfraction";
-            sel = inds,
+            sel = indices,
             defaults = 0.1,
             type = Float,
         )
         cmax = ncread(
-            nc,
+            dataset,
             config,
             "vertical.vegetation_parameter_set.cmax";
-            sel = inds,
+            sel = indices,
             defaults = 1.0,
             type = Float,
         )
diff --git a/src/reservoir_lake.jl b/src/reservoir_lake.jl
deleted file mode 100644
index edc55a91f..000000000
--- a/src/reservoir_lake.jl
+++ /dev/null
@@ -1,601 +0,0 @@
-@get_units @grid_loc @with_kw struct SimpleReservoir{T}
-    dt::T                                               # 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⁻¹]
-    demand::Vector{T} | "m3 s-1"                        # minimum (environmental) flow requirement downstream of the reservoir [m³ s⁻¹]
-    targetminfrac::Vector{T} | "-"                      # target minimum full fraction (of max storage) [-]
-    targetfullfrac::Vector{T} | "-"                     # target fraction full (of max storage) [-]
-    volume::Vector{T} | "m3"                            # reservoir volume [m³]
-    inflow::Vector{T} | "m3"                            # total inflow into reservoir [m³]
-    outflow::Vector{T} | "m3 s-1"                       # outflow from reservoir [m³ s⁻¹]
-    totaloutflow::Vector{T} | "m3"                      # total outflow from reservoir [m³]
-    percfull::Vector{T} | "-"                           # fraction full (of max storage) [-]
-    demandrelease::Vector{T} | "m3 s-1"                 # minimum (environmental) flow released from reservoir [m³ s⁻¹]
-    precipitation::Vector{T}                            # average precipitation for reservoir area [mm Δt⁻¹]
-    evaporation::Vector{T}                              # average potential evaporation for reservoir area [mm Δt⁻¹]
-    actevap::Vector{T}                                  # average actual evaporation for reservoir area [mm Δt⁻¹]
-
-    function SimpleReservoir{T}(args...) where {T}
-        equal_size_vectors(args)
-        return new(args...)
-    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
-    # note that these locations are only the reservoir outlet pixels
-    reslocs = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.locs";
-        optional = false,
-        sel = inds_riv,
-        type = Int,
-        fill = 0,
-    )
-
-    # this holds the same ids as reslocs, but covers the entire reservoir
-    rescoverage_2d = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.areas";
-        optional = false,
-        allow_missing = true,
-    )
-    # for each reservoir, a list of 2D indices, needed for getting the mean precipitation
-    inds_res_cov = Vector{CartesianIndex{2}}[]
-
-    rev_inds_reservoir = zeros(Int, size(rescoverage_2d))
-
-    # construct a map from the rivers to the reservoirs and
-    # a map of the reservoirs to the 2D model grid
-    resindex = fill(0, nriv)
-    inds_res = CartesianIndex{2}[]
-    rescounter = 0
-    for (i, ind) in enumerate(inds_riv)
-        res_id = reslocs[i]
-        if res_id > 0
-            push!(inds_res, ind)
-            rescounter += 1
-            resindex[i] = rescounter
-            rev_inds_reservoir[ind] = rescounter
-
-            # get all indices related to this reservoir outlet
-            # done in this loop to ensure that the order is equal to the order in the
-            # SimpleReservoir struct
-            res_cov = findall(isequal(res_id), rescoverage_2d)
-            push!(inds_res_cov, res_cov)
-        end
-    end
-
-    resdemand = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.demand";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-    resmaxrelease = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.maxrelease";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-    resmaxvolume = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.maxvolume";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-    resarea = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.area";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-    res_targetfullfrac = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.targetfullfrac";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-    res_targetminfrac = ncread(
-        nc,
-        config,
-        "lateral.river.reservoir.targetminfrac";
-        optional = false,
-        sel = inds_res,
-        type = Float,
-        fill = 0,
-    )
-
-    # for surface water routing reservoir locations are considered pits in the flow network
-    # all upstream flow goes to the river and flows into the reservoir
-    pits[inds_res] .= true
-
-    n = length(resarea)
-    @info "Read `$n` reservoir locations."
-    reservoirs = SimpleReservoir{Float}(;
-        dt = dt,
-        demand = resdemand,
-        maxrelease = resmaxrelease,
-        maxvolume = resmaxvolume,
-        area = resarea,
-        targetfullfrac = res_targetfullfrac,
-        targetminfrac = res_targetminfrac,
-        volume = res_targetfullfrac .* resmaxvolume,
-        inflow = fill(mv, n),
-        outflow = fill(mv, n),
-        totaloutflow = fill(mv, n),
-        percfull = fill(mv, n),
-        demandrelease = fill(mv, n),
-        precipitation = fill(mv, n),
-        evaporation = fill(mv, n),
-        actevap = fill(mv, n),
-    )
-
-    return reservoirs,
-    resindex,
-    (
-        indices_outlet = inds_res,
-        indices_coverage = inds_res_cov,
-        reverse_indices = rev_inds_reservoir,
-    ),
-    pits
-end
-
-"""
-Update a single reservoir at position `i`.
-
-This is called from within the kinematic wave loop, therefore updating only for a single
-element rather than all at once.
-"""
-function update!(res::SimpleReservoir, i, inflow, timestepsecs)
-
-    # limit lake evaporation based on total available volume [m³]
-    precipitation = 0.001 * res.precipitation[i] * (timestepsecs / res.dt) * res.area[i]
-    available_volume = res.volume[i] + inflow * timestepsecs + precipitation
-    evap = 0.001 * res.evaporation[i] * (timestepsecs / res.dt) * res.area[i]
-    actevap = min(available_volume, evap) # [m³/timestepsecs]
-
-    vol = res.volume[i] + (inflow * timestepsecs) + precipitation - actevap
-    vol = max(vol, 0.0)
-
-    percfull = vol / res.maxvolume[i]
-    # first determine minimum (environmental) flow using a simple sigmoid curve to scale for target level
-    fac = scurve(percfull, res.targetminfrac[i], Float(1.0), Float(30.0))
-    demandrelease = min(fac * res.demand[i] * timestepsecs, vol)
-    vol = vol - demandrelease
-
-    wantrel = max(0.0, vol - (res.maxvolume[i] * res.targetfullfrac[i]))
-    # Assume extra maximum Q if spilling
-    overflow_q = max((vol - res.maxvolume[i]), 0.0)
-    torelease = min(wantrel, overflow_q + res.maxrelease[i] * timestepsecs - demandrelease)
-    vol = vol - torelease
-    outflow = torelease + demandrelease
-    percfull = vol / res.maxvolume[i]
-
-    # update values in place
-    res.outflow[i] = outflow / timestepsecs
-    res.inflow[i] += inflow * timestepsecs
-    res.totaloutflow[i] += outflow
-    res.demandrelease[i] = demandrelease / timestepsecs
-    res.percfull[i] = percfull
-    res.volume[i] = vol
-    res.actevap[i] += 1000.0 * (actevap / res.area[i])
-
-    return nothing
-end
-
-@get_units @grid_loc @with_kw struct Lake{T}
-    dt::T                                       # 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³]
-    threshold::Vector{T} | "m"                  # water level threshold H₀ [m] below that level outflow is zero
-    storfunc::Vector{Int} | "-"                 # type of lake storage curve, 1: S = AH, 2: S = f(H) from lake data and interpolation
-    outflowfunc::Vector{Int} | "-"              # type of lake rating curve, 1: Q = f(H) from lake data and interpolation, 2: General Q = b(H - H₀)ᵉ, 3: Case of Puls Approach Q = b(H - H₀)²
-    b::Vector{T} | "m3/2 s-1 (if e=3/2)"        # rating curve coefficient
-    e::Vector{T} | "-"                          # rating curve exponent
-    sh::Vector{Union{SH, Missing}}               # data for storage curve
-    hq::Vector{Union{HQ, Missing}}               # data for rating curve
-    waterlevel::Vector{T} | "m"                 # waterlevel H [m] of lake
-    inflow::Vector{T} | "m3"                    # inflow to the lake [m³]
-    storage::Vector{T} | "m3"                   # storage lake [m³]
-    outflow::Vector{T} | "m3 s-1"               # outflow lake [m³ s⁻¹]
-    totaloutflow::Vector{T} | "m3"              # total outflow lake [m³]
-    precipitation::Vector{T}                    # average precipitation for lake area [mm Δt⁻¹]
-    evaporation::Vector{T}                      # average potential evaporation for lake area [mm Δt⁻¹]
-    actevap::Vector{T}                          # average actual evapotranspiration for lake area [mm Δt⁻¹]
-
-    function Lake{T}(args...) where {T}
-        equal_size_vectors(args)
-        return new(args...)
-    end
-end
-
-function initialize_lake(config, nc, inds_riv, nriv, pits, dt)
-    # read only lake data if lakes true
-    # allow lakes only in river cells
-    # note that these locations are only the lake outlet pixels
-    lakelocs_2d = ncread(
-        nc,
-        config,
-        "lateral.river.lake.locs";
-        optional = false,
-        type = Int,
-        fill = 0,
-    )
-    lakelocs = lakelocs_2d[inds_riv]
-
-    # this holds the same ids as lakelocs, but covers the entire lake
-    lakecoverage_2d = ncread(
-        nc,
-        config,
-        "lateral.river.lake.areas";
-        optional = false,
-        allow_missing = true,
-    )
-    # for each lake, a list of 2D indices, needed for getting the mean precipitation
-    inds_lake_cov = Vector{CartesianIndex{2}}[]
-
-    rev_inds_lake = zeros(Int, size(lakecoverage_2d))
-
-    # construct a map from the rivers to the lakes and
-    # a map of the lakes to the 2D model grid
-    lakeindex = fill(0, nriv)
-    inds_lake = CartesianIndex{2}[]
-    lakecounter = 0
-    for (i, ind) in enumerate(inds_riv)
-        lake_id = lakelocs[i]
-        if lake_id > 0
-            push!(inds_lake, ind)
-            lakecounter += 1
-            lakeindex[i] = lakecounter
-            rev_inds_lake[ind] = lakecounter
-
-            # get all indices related to this lake outlet
-            # done in this loop to ensure that the order is equal to the order in the
-            # NaturalLake struct
-            lake_cov = findall(isequal(lake_id), lakecoverage_2d)
-            push!(inds_lake_cov, lake_cov)
-        end
-    end
-
-    lakearea = ncread(
-        nc,
-        config,
-        "lateral.river.lake.area";
-        optional = false,
-        sel = inds_lake,
-        type = Float,
-        fill = 0,
-    )
-    lake_b = ncread(
-        nc,
-        config,
-        "lateral.river.lake.b";
-        optional = false,
-        sel = inds_lake,
-        type = Float,
-        fill = 0,
-    )
-    lake_e = ncread(
-        nc,
-        config,
-        "lateral.river.lake.e";
-        optional = false,
-        sel = inds_lake,
-        type = Float,
-        fill = 0,
-    )
-    lake_threshold = ncread(
-        nc,
-        config,
-        "lateral.river.lake.threshold";
-        optional = false,
-        sel = inds_lake,
-        type = Float,
-        fill = 0,
-    )
-    linked_lakelocs = ncread(
-        nc,
-        config,
-        "lateral.river.lake.linkedlakelocs";
-        sel = inds_lake,
-        defaults = 0,
-        type = Int,
-        fill = 0,
-    )
-    lake_storfunc = ncread(
-        nc,
-        config,
-        "lateral.river.lake.storfunc";
-        optional = false,
-        sel = inds_lake,
-        type = Int,
-        fill = 0,
-    )
-    lake_outflowfunc = ncread(
-        nc,
-        config,
-        "lateral.river.lake.outflowfunc";
-        optional = false,
-        sel = inds_lake,
-        type = Int,
-        fill = 0,
-    )
-    lake_waterlevel = ncread(
-        nc,
-        config,
-        "lateral.river.lake.waterlevel";
-        optional = false,
-        sel = inds_lake,
-        type = Float,
-        fill = 0,
-    )
-
-    # for surface water routing lake locations are considered pits in the flow network
-    # all upstream flow goes to the river and flows into the lake
-    pits[inds_lake] .= true
-
-    # This is currently the same length as all river cells, but will be the
-    # length of all lake cells. To do that we need to introduce a mapping.
-    n_lakes = length(inds_lake)
-    lakelocs = lakelocs_2d[inds_lake]
-
-    @info "Read `$n_lakes` lake locations."
-
-    sh = Vector{Union{SH, Missing}}(missing, n_lakes)
-    hq = Vector{Union{HQ, Missing}}(missing, n_lakes)
-    lowerlake_ind = fill(0, n_lakes)
-    # lake CSV parameter files are expected in the same directory as path_static
-    path = dirname(input_path(config, config.input.path_static))
-
-    for i in 1:n_lakes
-        lakeloc = lakelocs[i]
-        if linked_lakelocs[i] > 0
-            lowerlake_ind[i] = only(findall(x -> x == linked_lakelocs[i], lakelocs))
-        end
-
-        if lake_storfunc[i] == 2
-            csv_path = joinpath(path, "lake_sh_$lakeloc.csv")
-            @info(
-                "Read a storage curve from CSV file `$csv_path`, for lake location `$lakeloc`"
-            )
-            sh[i] = read_sh_csv(csv_path)
-        end
-
-        if lake_outflowfunc[i] == 1
-            csv_path = joinpath(path, "lake_hq_$lakeloc.csv")
-            @info(
-                "Read a rating curve from CSV file `$csv_path`, for lake location `$lakeloc`"
-            )
-            hq[i] = read_hq_csv(csv_path)
-        end
-
-        if lake_outflowfunc[i] == 3 && lake_storfunc[i] != 1
-            @warn(
-                "For the modified pulse approach (LakeOutflowFunc = 3) the LakeStorFunc should be 1"
-            )
-        end
-    end
-    n = length(lakearea)
-    lakes = Lake{Float}(;
-        dt = dt,
-        lowerlake_ind = lowerlake_ind,
-        area = lakearea,
-        maxstorage = maximum_storage(lake_storfunc, lake_outflowfunc, lakearea, sh, hq),
-        threshold = lake_threshold,
-        storfunc = lake_storfunc,
-        outflowfunc = lake_outflowfunc,
-        b = lake_b,
-        e = lake_e,
-        waterlevel = lake_waterlevel,
-        sh = sh,
-        hq = hq,
-        inflow = fill(mv, n),
-        storage = initialize_storage(lake_storfunc, lakearea, lake_waterlevel, sh),
-        outflow = fill(mv, n),
-        totaloutflow = fill(mv, n),
-        precipitation = fill(mv, n),
-        evaporation = fill(mv, n),
-        actevap = fill(mv, n),
-    )
-
-    return lakes,
-    lakeindex,
-    (
-        indices_outlet = inds_lake,
-        indices_coverage = inds_lake_cov,
-        reverse_indices = rev_inds_lake,
-    ),
-    pits
-end
-
-"Determine the initial storage depending on the storage function"
-function initialize_storage(storfunc, area, waterlevel, sh)
-    storage = similar(area)
-    for i in eachindex(storage)
-        if storfunc[i] == 1
-            storage[i] = area[i] * waterlevel[i]
-        else
-            storage[i] = interpolate_linear(waterlevel[i], sh[i].H, sh[i].S)
-        end
-    end
-    return storage
-end
-
-"Determine the water level depending on the storage function"
-function waterlevel(storfunc, area, storage, sh)
-    waterlevel = similar(area)
-    for i in eachindex(storage)
-        if storfunc[i] == 1
-            waterlevel[i] = storage[i] / area[i]
-        else
-            waterlevel[i] = interpolate_linear(storage[i], sh[i].S, sh[i].H)
-        end
-    end
-    return waterlevel
-end
-
-"Determine the maximum storage for lakes with a rating curve of type 1"
-function maximum_storage(storfunc, outflowfunc, area, sh, hq)
-    maxstorage = Vector{Union{Float, Missing}}(missing, length(area))
-    # maximum storage is based on the maximum water level (H) value in the H-Q table
-    for i in eachindex(maxstorage)
-        if outflowfunc[i] == 1
-            if storfunc[i] == 2
-                maxstorage[i] = interpolate_linear(maximum(hq[i].H), sh[i].H, sh[i].S)
-            else
-                maxstorage[i] = area[i] * maximum(hq[i].H)
-            end
-        end
-    end
-    return maxstorage
-end
-
-function interpolate_linear(x, xp, fp)
-    if x <= minimum(xp)
-        return minimum(fp)
-    elseif x >= maximum(xp)
-        return maximum(fp)
-    else
-        idx = findall(i -> i <= x, xp)
-        i1 = last(idx)
-        i2 = i1 + 1
-        return fp[i1] * (1.0 - (x - xp[i1]) / (xp[i2] - xp[i1])) +
-               fp[i2] * (x - xp[i1]) / (xp[i2] - xp[i1])
-    end
-end
-
-"""
-Update a single lake at position `i`.
-
-This is called from within the kinematic wave loop, therefore updating only for a single
-element rather than all at once.
-"""
-function update!(lake::Lake, i, inflow, doy, timestepsecs)
-    lo = lake.lowerlake_ind[i]
-    has_lowerlake = lo != 0
-
-    # limit lake evaporation based on total available volume [m³]
-    precipitation = 0.001 * lake.precipitation[i] * (timestepsecs / lake.dt) * lake.area[i]
-    available_volume = lake.storage[i] + inflow * timestepsecs + precipitation
-    evap = 0.001 * lake.evaporation[i] * (timestepsecs / lake.dt) * lake.area[i]
-    actevap = min(available_volume, evap) # [m³/timestepsecs]
-
-    ### Modified Puls Approach (Burek et al., 2013, LISFLOOD) ###
-    # outflowfunc = 3
-    # Calculate lake factor and SI parameter
-    if lake.outflowfunc[i] == 3
-        lakefactor = lake.area[i] / (timestepsecs * pow(lake.b[i], 0.5))
-        si_factor = (lake.storage[i] + precipitation - actevap) / timestepsecs + inflow
-        # Adjust SIFactor for ResThreshold != 0
-        si_factor_adj = si_factor - lake.area[i] * lake.threshold[i] / timestepsecs
-        # Calculate the new lake outflow/waterlevel/storage
-        if si_factor_adj > 0.0
-            quadratic_sol_term =
-                -lakefactor + pow((pow(lakefactor, 2.0) + 4.0 * si_factor_adj), 0.5)
-            if quadratic_sol_term > 0.0
-                outflow = pow(0.5 * quadratic_sol_term, 2.0)
-            else
-                outflow = 0.0
-            end
-        else
-            outflow = 0.0
-        end
-        outflow = min(outflow, si_factor)
-        storage = (si_factor - outflow) * timestepsecs
-        waterlevel = storage / lake.area[i]
-    end
-
-    ### Linearisation for specific storage/rating curves ###
-    if lake.outflowfunc[i] == 1 || lake.outflowfunc[i] == 2
-        diff_wl = has_lowerlake ? lake.waterlevel[i] - lake.waterlevel[lo] : 0.0
-
-        storage_input = (lake.storage[i] + precipitation - actevap) / timestepsecs + inflow
-        if lake.outflowfunc[i] == 1
-            outflow =
-                interpolate_linear(lake.waterlevel[i], lake.hq[i].H, lake.hq[i].Q[:, doy])
-            outflow = min(outflow, storage_input)
-        else
-            if diff_wl >= 0.0
-                if lake.waterlevel[i] > lake.threshold[i]
-                    dh = lake.waterlevel[i] - lake.threshold[i]
-                    outflow = lake.b[i] * pow(dh, lake.e[i])
-                    maxflow = (dh * lake.area[i]) / timestepsecs
-                    outflow = min(outflow, maxflow)
-                else
-                    outflow = Float(0)
-                end
-            else
-                if lake.waterlevel[lo] > lake.threshold[i]
-                    dh = lake.waterlevel[lo] - lake.threshold[i]
-                    outflow = -1.0 * lake.b[i] * pow(dh, lake.e[i])
-                    maxflow = (dh * lake.area[lo]) / timestepsecs
-                    outflow = max(outflow, -maxflow)
-                else
-                    outflow = Float(0)
-                end
-            end
-        end
-        storage = (storage_input - outflow) * timestepsecs
-
-        # update storage and outflow for lake with rating curve of type 1.
-        if lake.outflowfunc[i] == 1
-            overflow = max(0.0, (storage - lake.maxstorage[i]) / timestepsecs)
-            storage = min(storage, lake.maxstorage[i])
-            outflow = outflow + overflow
-        end
-
-        waterlevel = if lake.storfunc[i] == 1
-            lake.waterlevel[i] + (storage - lake.storage[i]) / lake.area[i]
-        else
-            interpolate_linear(storage, lake.sh[i].S, lake.sh[i].H)
-        end
-
-        # update lower lake (linked lakes) in case flow from lower lake to upper lake occurs
-        if diff_wl < 0.0
-            lowerlake_storage = lake.storage[lo] + outflow * timestepsecs
-
-            lowerlake_waterlevel = if lake.storfunc[lo] == 1
-                lake.waterlevel[lo] + (lowerlake_storage - lake.storage[lo]) / lake.area[lo]
-            else
-                interpolate_linear(lowerlake_storage, lake.sh[lo].S, lake.sh[lo].H)
-            end
-
-            # update values for the lower lake in place
-            lake.outflow[lo] = -outflow
-            lake.totaloutflow[lo] += -outflow * timestepsecs
-            lake.storage[lo] = lowerlake_storage
-            lake.waterlevel[lo] = lowerlake_waterlevel
-        end
-    end
-
-    # update values in place
-    lake.outflow[i] = max(outflow, 0.0) # for a linked lake flow can be negative
-    lake.waterlevel[i] = waterlevel
-    lake.inflow[i] += inflow * timestepsecs
-    lake.totaloutflow[i] += outflow * timestepsecs
-    lake.storage[i] = storage
-    lake.actevap[i] += 1000.0 * (actevap / lake.area[i])
-
-    return nothing
-end
diff --git a/src/routing/lake.jl b/src/routing/lake.jl
new file mode 100644
index 000000000..425915897
--- /dev/null
+++ b/src/routing/lake.jl
@@ -0,0 +1,440 @@
+"Struct for storing lake model parameters"
+@get_units @grid_loc @with_kw struct LakeParameters{T}
+    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³]
+    threshold::Vector{T} | "m"                      # water level threshold H₀ [m] below that level outflow is zero
+    storfunc::Vector{Int} | "-"                     # type of lake storage curve, 1: S = AH, 2: S = f(H) from lake data and interpolation
+    outflowfunc::Vector{Int} | "-"                  # type of lake rating curve, 1: Q = f(H) from lake data and interpolation, 2: General Q = b(H - H₀)ᵉ, 3: Case of Puls Approach Q = b(H - H₀)²
+    b::Vector{T} | "m3/2 s-1 (if e=3/2)"            # rating curve coefficient
+    e::Vector{T} | "-"                              # rating curve exponent
+    sh::Vector{Union{SH, Missing}}                  # data for storage curve
+    hq::Vector{Union{HQ, Missing}}                  # data for rating curve    
+end
+
+"Initialize lake model parameters"
+function LakeParameters(config, dataset, inds_riv, nriv, pits)
+    # read only lake data if lakes true
+    # allow lakes only in river cells
+    # note that these locations are only the lake outlet pixels
+    lakelocs_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.locs";
+        optional = false,
+        type = Int,
+        fill = 0,
+    )
+    lakelocs = lakelocs_2d[inds_riv]
+
+    # this holds the same ids as lakelocs, but covers the entire lake
+    lakecoverage_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.areas";
+        optional = false,
+        allow_missing = true,
+    )
+    # for each lake, a list of 2D indices, needed for getting the mean precipitation
+    inds_lake_cov = Vector{CartesianIndex{2}}[]
+
+    rev_inds_lake = zeros(Int, size(lakecoverage_2d))
+
+    # construct a map from the rivers to the lakes and
+    # a map of the lakes to the 2D model grid
+    inds_lake_map2river = fill(0, nriv)
+    inds_lake = CartesianIndex{2}[]
+    lakecounter = 0
+    for (i, ind) in enumerate(inds_riv)
+        lake_id = lakelocs[i]
+        if lake_id > 0
+            push!(inds_lake, ind)
+            lakecounter += 1
+            inds_lake_map2river[i] = lakecounter
+            rev_inds_lake[ind] = lakecounter
+
+            # get all indices related to this lake outlet
+            # done in this loop to ensure that the order is equal to the order in the
+            # NaturalLake struct
+            lake_cov = findall(isequal(lake_id), lakecoverage_2d)
+            push!(inds_lake_cov, lake_cov)
+        end
+    end
+
+    lakearea = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.area";
+        optional = false,
+        sel = inds_lake,
+        type = Float,
+        fill = 0,
+    )
+    lake_b = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.b";
+        optional = false,
+        sel = inds_lake,
+        type = Float,
+        fill = 0,
+    )
+    lake_e = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.e";
+        optional = false,
+        sel = inds_lake,
+        type = Float,
+        fill = 0,
+    )
+    lake_threshold = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.threshold";
+        optional = false,
+        sel = inds_lake,
+        type = Float,
+        fill = 0,
+    )
+    linked_lakelocs = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.linkedlakelocs";
+        sel = inds_lake,
+        defaults = 0,
+        type = Int,
+        fill = 0,
+    )
+    lake_storfunc = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.storfunc";
+        optional = false,
+        sel = inds_lake,
+        type = Int,
+        fill = 0,
+    )
+    lake_outflowfunc = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.outflowfunc";
+        optional = false,
+        sel = inds_lake,
+        type = Int,
+        fill = 0,
+    )
+    lake_waterlevel = ncread(
+        dataset,
+        config,
+        "lateral.river.lake.waterlevel";
+        optional = false,
+        sel = inds_lake,
+        type = Float,
+        fill = 0,
+    )
+
+    # for surface water routing lake locations are considered pits in the flow network
+    # all upstream flow goes to the river and flows into the lake
+    pits[inds_lake] .= true
+
+    # This is currently the same length as all river cells, but will be the
+    # length of all lake cells. To do that we need to introduce a mapping.
+    n_lakes = length(inds_lake)
+    lakelocs = lakelocs_2d[inds_lake]
+
+    @info "Read `$n_lakes` lake locations."
+
+    sh = Vector{Union{SH, Missing}}(missing, n_lakes)
+    hq = Vector{Union{HQ, Missing}}(missing, n_lakes)
+    lowerlake_ind = fill(0, n_lakes)
+    # lake CSV parameter files are expected in the same directory as path_static
+    path = dirname(input_path(config, config.input.path_static))
+
+    for i in 1:n_lakes
+        lakeloc = lakelocs[i]
+        if linked_lakelocs[i] > 0
+            lowerlake_ind[i] = only(findall(x -> x == linked_lakelocs[i], lakelocs))
+        end
+
+        if lake_storfunc[i] == 2
+            csv_path = joinpath(path, "lake_sh_$lakeloc.csv")
+            @info(
+                "Read a storage curve from CSV file `$csv_path`, for lake location `$lakeloc`"
+            )
+            sh[i] = read_sh_csv(csv_path)
+        end
+
+        if lake_outflowfunc[i] == 1
+            csv_path = joinpath(path, "lake_hq_$lakeloc.csv")
+            @info(
+                "Read a rating curve from CSV file `$csv_path`, for lake location `$lakeloc`"
+            )
+            hq[i] = read_hq_csv(csv_path)
+        end
+
+        if lake_outflowfunc[i] == 3 && lake_storfunc[i] != 1
+            @warn(
+                "For the modified pulse approach (LakeOutflowFunc = 3) the LakeStorFunc should be 1"
+            )
+        end
+    end
+    parameters = LakeParameters{Float}(;
+        lowerlake_ind = lowerlake_ind,
+        area = lakearea,
+        maxstorage = maximum_storage(lake_storfunc, lake_outflowfunc, lakearea, sh, hq),
+        threshold = lake_threshold,
+        storfunc = lake_storfunc,
+        outflowfunc = lake_outflowfunc,
+        b = lake_b,
+        e = lake_e,
+        sh = sh,
+        hq = hq,
+    )
+    lake_network = (
+        indices_outlet = inds_lake,
+        indices_coverage = inds_lake_cov,
+        reverse_indices = rev_inds_lake,
+        river_indices = findall(x -> x ≠ 0, inds_lake_map2river),
+    )
+    return parameters, lake_network, inds_lake_map2river, lake_waterlevel, pits
+end
+
+"Struct for storing Lake model parameters"
+@get_units @grid_loc @with_kw struct LakeVariables{T}
+    waterlevel::Vector{T} | "m"                 # waterlevel H [m] of lake
+    storage::Vector{T} | "m3"                   # storage lake [m³]
+    outflow::Vector{T} | "m3 s-1"               # outflow of lake outlet [m³ s⁻¹]
+    outflow_av::Vector{T} | "m3 s-1"            # average outflow lake [m³ s⁻¹] for model timestep Δt (including flow from lower to upper lake)
+    actevap::Vector{T}                          # average actual evapotranspiration for lake area [mm Δt⁻¹] 
+end
+
+"Initialize lake model variables"
+function LakeVariables(n, lake_waterlevel)
+    variables = LakeVariables{Float}(;
+        waterlevel = lake_waterlevel,
+        inflow = fill(mv, n),
+        storage = initialize_storage(lake_storfunc, lakearea, lake_waterlevel, sh),
+        outflow = fill(mv, n),
+        outflow_av = fill(mv, n),
+        actevap = fill(mv, n),
+    )
+    return variables
+end
+
+"Struct for storing lake model boundary conditions"
+@get_units @grid_loc @with_kw struct LakeBC{T}
+    inflow::Vector{T} | "m3 s-1"                # inflow to the lake [m³ s⁻¹] for model timestep Δt
+    precipitation::Vector{T}                    # average precipitation for lake area [mm Δt⁻¹]
+    evaporation::Vector{T}                      # average potential evaporation for lake area [mm Δt⁻¹]
+end
+
+"Initialize lake model boundary conditions"
+function LakeBC(n)
+    bc = LakeBC{Float}(;
+        inflow = fill(mv, n),
+        precipitation = fill(mv, n),
+        evaporation = fill(mv, n),
+    )
+    return bc
+end
+
+"Lake model"
+@with_kw struct Lake{T}
+    boundary_conditions::LakeBC{T}
+    parameters::LakeParameters{T}
+    variables::LakeVariables{T}
+end
+
+"Initialize lake model"
+function Lake(dataset, config, indices_river, n_river_cells, pits)
+    parameters, lake_network, inds_lake_map2river, lake_waterlevel, pits =
+        LakeParameters(dataset, config, indices_river, n_river_cells, pits)
+
+    n_lakes = length(parameters.area)
+    variables = LakeVariables(n_lakes, lake_waterlevel)
+    boundary_conditions = LakeBC(n_lakes)
+
+    lake = Lake{Float}(; boundary_conditions, parameters, variables)
+
+    return lake, lake_network, inds_lake_map2river, pits
+end
+
+"Determine the initial storage depending on the storage function"
+function initialize_storage(storfunc, area, waterlevel, sh)
+    storage = similar(area)
+    for i in eachindex(storage)
+        if storfunc[i] == 1
+            storage[i] = area[i] * waterlevel[i]
+        else
+            storage[i] = interpolate_linear(waterlevel[i], sh[i].H, sh[i].S)
+        end
+    end
+    return storage
+end
+
+"Determine the water level depending on the storage function"
+function waterlevel(storfunc, area, storage, sh)
+    waterlevel = similar(area)
+    for i in eachindex(storage)
+        if storfunc[i] == 1
+            waterlevel[i] = storage[i] / area[i]
+        else
+            waterlevel[i] = interpolate_linear(storage[i], sh[i].S, sh[i].H)
+        end
+    end
+    return waterlevel
+end
+
+"Determine the maximum storage for lakes with a rating curve of type 1"
+function maximum_storage(storfunc, outflowfunc, area, sh, hq)
+    maxstorage = Vector{Union{Float, Missing}}(missing, length(area))
+    # maximum storage is based on the maximum water level (H) value in the H-Q table
+    for i in eachindex(maxstorage)
+        if outflowfunc[i] == 1
+            if storfunc[i] == 2
+                maxstorage[i] = interpolate_linear(maximum(hq[i].H), sh[i].H, sh[i].S)
+            else
+                maxstorage[i] = area[i] * maximum(hq[i].H)
+            end
+        end
+    end
+    return maxstorage
+end
+
+function interpolate_linear(x, xp, fp)
+    if x <= minimum(xp)
+        return minimum(fp)
+    elseif x >= maximum(xp)
+        return maximum(fp)
+    else
+        idx = findall(i -> i <= x, xp)
+        i1 = last(idx)
+        i2 = i1 + 1
+        return fp[i1] * (1.0 - (x - xp[i1]) / (xp[i2] - xp[i1])) +
+               fp[i2] * (x - xp[i1]) / (xp[i2] - xp[i1])
+    end
+end
+
+"""
+Update a single lake at position `i`.
+
+This is called from within the kinematic wave loop, therefore updating only for a single
+element rather than all at once.
+"""
+function update!(model::Lake, i, inflow, doy, dt, dt_forcing)
+    lake_bc = model.boundary_conditions
+    lake_p = model.parameters
+    lake_v = model.variables
+
+    lo = lake_p.lowerlake_ind[i]
+    has_lowerlake = lo != 0
+
+    # limit lake evaporation based on total available volume [m³]
+    precipitation = 0.001 * lake_bc.precipitation[i] * (dt / dt_forcing) * lake_p.area[i]
+    available_volume = lake_v.storage[i] + inflow * dt + precipitation
+    evap = 0.001 * lake_bc.evaporation[i] * (dt / dt_forcing) * lake_p.area[i]
+    actevap = min(available_volume, evap) # [m³/dt]
+
+    ### Modified Puls Approach (Burek et al., 2013, LISFLOOD) ###
+    # outflowfunc = 3
+    # Calculate lake factor and SI parameter
+    if lake_p.outflowfunc[i] == 3
+        lakefactor = lake_p.area[i] / (dt * pow(lake_p.b[i], 0.5))
+        si_factor = (lake_v.storage[i] + precipitation - actevap) / dt + inflow
+        # Adjust SIFactor for ResThreshold != 0
+        si_factor_adj = si_factor - lake_p.area[i] * lake_p.threshold[i] / dt
+        # Calculate the new lake outflow/waterlevel/storage
+        if si_factor_adj > 0.0
+            quadratic_sol_term =
+                -lakefactor + pow((pow(lakefactor, 2.0) + 4.0 * si_factor_adj), 0.5)
+            if quadratic_sol_term > 0.0
+                outflow = pow(0.5 * quadratic_sol_term, 2.0)
+            else
+                outflow = 0.0
+            end
+        else
+            outflow = 0.0
+        end
+        outflow = min(outflow, si_factor)
+        storage = (si_factor - outflow) * dt
+        waterlevel = storage / lake_p.area[i]
+    end
+
+    ### Linearisation for specific storage/rating curves ###
+    if lake_p.outflowfunc[i] == 1 || lake_p.outflowfunc[i] == 2
+        diff_wl = has_lowerlake ? lake_v.waterlevel[i] - lake_v.waterlevel[lo] : 0.0
+
+        storage_input = (lake_v.storage[i] + precipitation - actevap) / dt + inflow
+        if lake_p.outflowfunc[i] == 1
+            outflow = interpolate_linear(
+                lake_v.waterlevel[i],
+                lake_p.hq[i].H,
+                lake_p.hq[i].Q[:, doy],
+            )
+            outflow = min(outflow, storage_input)
+        else
+            if diff_wl >= 0.0
+                if lake_v.waterlevel[i] > lake_p.threshold[i]
+                    dh = lake_v.waterlevel[i] - lake_p.threshold[i]
+                    outflow = lake_p.b[i] * pow(dh, lake_p.e[i])
+                    maxflow = (dh * lake_p.area[i]) / dt
+                    outflow = min(outflow, maxflow)
+                else
+                    outflow = Float(0)
+                end
+            else
+                if lake_v.waterlevel[lo] > lake_p.threshold[i]
+                    dh = lake_v.waterlevel[lo] - lake_p.threshold[i]
+                    outflow = -1.0 * lake_p.b[i] * pow(dh, lake_p.e[i])
+                    maxflow = (dh * lake_p.area[lo]) / dt
+                    outflow = max(outflow, -maxflow)
+                else
+                    outflow = Float(0)
+                end
+            end
+        end
+        storage = (storage_input - outflow) * dt
+
+        # update storage and outflow for lake with rating curve of type 1.
+        if lake_p.outflowfunc[i] == 1
+            overflow = max(0.0, (storage - lake_p.maxstorage[i]) / dt)
+            storage = min(storage, lake_p.maxstorage[i])
+            outflow = outflow + overflow
+        end
+
+        waterlevel = if lake_p.storfunc[i] == 1
+            lake_v.waterlevel[i] + (storage - lake_v.storage[i]) / lake_p.area[i]
+        else
+            interpolate_linear(storage, lake_p.sh[i].S, lake_p.sh[i].H)
+        end
+
+        # update lower lake (linked lakes) in case flow from lower lake to upper lake occurs
+        if diff_wl < 0.0
+            lowerlake_storage = lake_v.storage[lo] + outflow * dt
+
+            lowerlake_waterlevel = if lake_p.storfunc[lo] == 1
+                lake_v.waterlevel[lo] +
+                (lowerlake_storage - lake_v.storage[lo]) / lake_p.area[lo]
+            else
+                interpolate_linear(lowerlake_storage, lake_p.sh[lo].S, lake_p.sh[lo].H)
+            end
+
+            # update values for the lower lake in place
+            lake_v.outflow[lo] = -outflow
+            lake_v.outflow_av[lo] += -outflow * dt
+            lake_v.storage[lo] = lowerlake_storage
+            lake_v.waterlevel[lo] = lowerlake_waterlevel
+        end
+    end
+
+    # update values in place
+    lake_v.outflow[i] = outflow
+    lake_v.waterlevel[i] = waterlevel
+    lake_bc.inflow[i] += inflow * dt
+    lake_v.outflow_av[i] += outflow * dt
+    lake_v.storage[i] = storage
+    lake_v.actevap[i] += 1000.0 * (actevap / lake_p.area[i])
+
+    return nothing
+end
diff --git a/src/routing/reservoir.jl b/src/routing/reservoir.jl
new file mode 100644
index 000000000..d22a37587
--- /dev/null
+++ b/src/routing/reservoir.jl
@@ -0,0 +1,245 @@
+"Struct for storing reservoir model parameters"
+@get_units @grid_loc @with_kw struct ReservoirParameters{T}
+    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⁻¹]
+    demand::Vector{T} | "m3 s-1"            # minimum (environmental) flow requirement downstream of the reservoir [m³ s⁻¹]
+    targetminfrac::Vector{T} | "-"          # target minimum full fraction (of max storage) [-]
+    targetfullfrac::Vector{T} | "-"         # target fraction full (of max storage
+end
+
+"Initialize reservoir model parameters"
+function ReservoirParameters(dataset, config, indices_river, n_river_cells, pits)
+    # read only reservoir data if reservoirs true
+    # allow reservoirs only in river cells
+    # note that these locations are only the reservoir outlet pixels
+    reslocs = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.locs";
+        optional = false,
+        sel = indices_river,
+        type = Int,
+        fill = 0,
+    )
+
+    # this holds the same ids as reslocs, but covers the entire reservoir
+    rescoverage_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.areas";
+        optional = false,
+        allow_missing = true,
+    )
+    # for each reservoir, a list of 2D indices, needed for getting the mean precipitation
+    inds_res_cov = Vector{CartesianIndex{2}}[]
+
+    rev_inds_reservoir = zeros(Int, size(rescoverage_2d))
+
+    # construct a map from the rivers to the reservoirs and
+    # a map of the reservoirs to the 2D model grid
+    inds_reservoir_map2river = fill(0, n_river_cells)
+    inds_res = CartesianIndex{2}[]
+    rescounter = 0
+    for (i, ind) in enumerate(indices_river)
+        res_id = reslocs[i]
+        if res_id > 0
+            push!(inds_res, ind)
+            rescounter += 1
+            inds_reservoir_map2river[i] = rescounter
+            rev_inds_reservoir[ind] = rescounter
+
+            # get all indices related to this reservoir outlet
+            # done in this loop to ensure that the order is equal to the order in the
+            # SimpleReservoir struct
+            res_cov = findall(isequal(res_id), rescoverage_2d)
+            push!(inds_res_cov, res_cov)
+        end
+    end
+
+    resdemand = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.demand";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+    resmaxrelease = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.maxrelease";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+    resmaxvolume = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.maxvolume";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+    resarea = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.area";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+    res_targetfullfrac = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.targetfullfrac";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+    res_targetminfrac = ncread(
+        dataset,
+        config,
+        "lateral.river.reservoir.targetminfrac";
+        optional = false,
+        sel = inds_res,
+        type = Float,
+        fill = 0,
+    )
+
+    # for surface water routing reservoir locations are considered pits in the flow network
+    # all upstream flow goes to the river and flows into the reservoir
+    pits[inds_res] .= true
+
+    reservoir_network = (
+        indices_outlet = inds_res,
+        indices_coverage = inds_res_cov,
+        reverse_indices = rev_inds_reservoir,
+        river_indices = findall(x -> x ≠ 0, inds_reservoir_map2river),
+    )
+
+    parameters = ReservoirParameters{Float}(;
+        demand = resdemand,
+        maxrelease = resmaxrelease,
+        maxvolume = resmaxvolume,
+        area = resarea,
+        targetfullfrac = res_targetfullfrac,
+        targetminfrac = res_targetminfrac,
+    )
+
+    return parameters, reservoir_network, inds_reservoir_map2river, pits
+end
+
+"Struct for storing reservoir model variables"
+@get_units @grid_loc @with_kw struct ReservoirVariables{T}
+    volume::Vector{T} | "m3"                # reservoir volume [m³]
+    outflow::Vector{T} | "m3 s-1"           # outflow from reservoir [m³ s⁻¹]
+    outflow_av::Vector{T} | "m3 s-1"        # average outflow from reservoir [m³ s⁻¹] for model timestep Δt
+    percfull::Vector{T} | "-"               # fraction full (of max storage) [-]
+    demandrelease::Vector{T} | "m3 s-1"     # minimum (environmental) flow released from reservoir [m³ s⁻¹]
+    actevap::Vector{T}                      # average actual evaporation for reservoir area [mm Δt⁻¹]
+end
+
+"Initialize reservoir model variables"
+function ReservoirVariables(n, parameters)
+    (; targetfullfrac, maxvolume) = parameters
+    variables = ReservoirVariables{Float}(;
+        volume = targetfullfrac .* maxvolume,
+        outflow = fill(mv, n),
+        outflow_av = fill(mv, n),
+        percfull = fill(mv, n),
+        demandrelease = fill(mv, n),
+        actevap = fill(mv, n),
+    )
+    return variables
+end
+
+"Struct for storing reservoir model boundary conditions"
+@get_units @grid_loc @with_kw struct ReservoirBC{T}
+    inflow::Vector{T} | "m3 s-1"    # inflow into reservoir [m³ s⁻¹] for model timestep Δt
+    precipitation::Vector{T}        # average precipitation for reservoir area [mm Δt⁻¹]
+    evaporation::Vector{T}          # average potential evaporation for reservoir area [mm Δt⁻¹]
+end
+
+"Initialize reservoir model boundary conditions"
+function ReservoirBC(n)
+    bc = ReservoirBC{Float}(;
+        inflow = fill(mv, n),
+        precipitation = fill(mv, n),
+        evaporation = fill(mv, n),
+    )
+    return bc
+end
+
+"Reservoir model `SimpleReservoir`"
+@with_kw struct SimpleReservoir{T}
+    boundary_conditions::ReservoirBC{T}
+    parameters::ReservoirParameters{T}
+    variables::ReservoirVariables{T}
+end
+
+"Initialize reservoir model `SimpleReservoir`"
+function SimpleReservoir(dataset, config, indices_river, n_river_cells, pits)
+    parameters, reservoir_network, inds_reservoir_map2river, pits =
+        ReservoirParameters(dataset, config, indices_river, n_river_cells, pits)
+
+    n_reservoirs = length(parameters.area)
+    @info "Read `$n_reservoirs` reservoir locations."
+
+    variables = ReservoirVariables(n_reservoirs, parameters)
+    boundary_conditions = ReservoirBC(n_reservoirs)
+    reservoir = SimpleReservoir{Float}(; boundary_conditions, parameters, variables)
+
+    return reservoir, reservoir_network, inds_reservoir_map2river, pits
+end
+
+"""
+Update a single reservoir at position `i`.
+
+This is called from within the kinematic wave loop, therefore updating only for a single
+element rather than all at once.
+"""
+function update!(model::SimpleReservoir, i, inflow, dt, dt_forcing)
+    res_bc = model.boundary_conditions
+    res_p = model.parameters
+    res_v = model.variables
+
+    # limit lake evaporation based on total available volume [m³]
+    precipitation = 0.001 * res_bc.precipitation[i] * (dt / dt_forcing) * res_p.area[i]
+    available_volume = res_v.volume[i] + inflow * dt + precipitation
+    evap = 0.001 * res_bc.evaporation[i] * (dt / dt_forcing) * res_p.area[i]
+    actevap = min(available_volume, evap) # [m³/dt]
+
+    vol = res_v.volume[i] + (inflow * dt) + precipitation - actevap
+    vol = max(vol, 0.0)
+
+    percfull = vol / res_p.maxvolume[i]
+    # first determine minimum (environmental) flow using a simple sigmoid curve to scale for target level
+    fac = scurve(percfull, res_p.targetminfrac[i], Float(1.0), Float(30.0))
+    demandrelease = min(fac * res_p.demand[i] * dt, vol)
+    vol = vol - demandrelease
+
+    wantrel = max(0.0, vol - (res_p.maxvolume[i] * res_p.targetfullfrac[i]))
+    # Assume extra maximum Q if spilling
+    overflow_q = max((vol - res_p.maxvolume[i]), 0.0)
+    torelease = min(wantrel, overflow_q + res_p.maxrelease[i] * dt - demandrelease)
+    vol = vol - torelease
+    outflow = torelease + demandrelease
+    percfull = vol / res_p.maxvolume[i]
+
+    # update values in place
+    res_v.outflow[i] = outflow / dt
+    res_bc.inflow[i] += inflow * dt
+    res_v.outflow_av[i] += outflow
+    res_v.demandrelease[i] = demandrelease / dt
+    res_v.percfull[i] = percfull
+    res_v.volume[i] = vol
+    res_v.actevap[i] += 1000.0 * (actevap / res_p.area[i])
+
+    return nothing
+end
\ No newline at end of file
diff --git a/src/horizontal_process.jl b/src/routing/routing_process.jl
similarity index 98%
rename from src/horizontal_process.jl
rename to src/routing/routing_process.jl
index fe9c57252..e51c5f9af 100644
--- a/src/horizontal_process.jl
+++ b/src/routing/routing_process.jl
@@ -1,17 +1,17 @@
 "Convert a gridded drainage direction to a directed graph"
-function flowgraph(ldd::AbstractVector, inds::AbstractVector, pcr_dir::AbstractVector)
+function flowgraph(ldd::AbstractVector, indices::AbstractVector, pcr_dir::AbstractVector)
     # prepare a directed graph to be filled
-    n = length(inds)
+    n = length(indices)
     graph = DiGraph(n)
 
     # loop over ldd, adding the edge to the downstream node
-    for (from_node, from_index) in enumerate(inds)
+    for (from_node, from_index) in enumerate(indices)
         ldd_val = ldd[from_node]
         # skip pits to prevent cycles
         ldd_val == 5 && continue
         to_index = from_index + pcr_dir[ldd_val]
         # find the node id of the downstream cell
-        to_node = searchsortedfirst(inds, to_index)
+        to_node = searchsortedfirst(indices, to_index)
         add_edge!(graph, from_node, to_node)
     end
     if is_cyclic(graph)
diff --git a/src/routing/subsurface.jl b/src/routing/subsurface.jl
new file mode 100644
index 000000000..3bf73f642
--- /dev/null
+++ b/src/routing/subsurface.jl
@@ -0,0 +1,253 @@
+abstract type SubsurfaceFlow end
+
+"Exponential depth profile of horizontal hydraulic conductivity at the soil surface"
+@get_units @grid_loc struct KhExponential{T}
+    # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
+    kh_0::Vector{T} | "m d-1"
+    # A scaling parameter [m⁻¹] (controls exponential decline of kh_0)
+    f::Vector{T} | "m-1"
+end
+
+"Exponential constant depth profile of horizontal hydraulic conductivity"
+@get_units @grid_loc struct KhExponentialConstant{T}
+    # Exponential horizontal hydraulic conductivity profile type
+    exponential::KhExponential
+    # Depth [m] from soil surface for which exponential decline of kv_0 is valid
+    z_exp::Vector{T} | "m"
+end
+
+"Layered depth profile of horizontal hydraulic conductivity"
+@get_units @grid_loc struct KhLayered{T}
+    # Horizontal hydraulic conductivity [m d⁻¹]
+    kh::Vector{T} | "m d-1"
+end
+
+"Struct for storing lateral subsurface flow model parameters"
+@get_units @grid_loc @with_kw struct LateralSsfParameters{T, Kh}
+    kh_profile::Kh                         # Horizontal hydraulic conductivity profile type [-]  
+    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 [-]   
+    slope::Vector{T} | "m m-1"             # Slope [m m⁻¹]
+    flow_length::Vector{T} | "m"           # Flow length [m]
+    flow_width::Vector{T} | "m"            # Flow width [m]
+end
+
+"Initialize lateral subsurface flow model parameters"
+function LateralSsfParameters(
+    dataset,
+    config,
+    indices;
+    soil,
+    slope,
+    flow_length,
+    flow_width,
+)
+    khfrac = ncread(
+        dataset,
+        config,
+        "lateral.subsurface.ksathorfrac";
+        sel = indices,
+        defaults = 1.0,
+        type = Float,
+    )
+    n_cells = length(khfrac)
+
+    (; theta_s, theta_r, soilthickness) = soil
+    soilthickness = soilthickness .* 0.001
+
+    kh_profile_type = get(config.input.vertical, "ksat_profile", "exponential")::String
+    dt = Second(config.timestepsecs) / basetimestep
+    if kh_profile_type == "exponential"
+        (; kv_0, f) = soil.kv_profile
+        kh_0 = khfrac .* kv_0 .* 0.001 .* dt
+        kh_profile = KhExponential(kh_0, f .* 1000.0)
+    elseif kh_profile_type == "exponential_constant"
+        (; z_exp) = soil.kv_profile
+        (; kv_0, f) = soil.kv_profile.exponential
+        kh_0 = khfrac .* kv_0 .* 0.001 .* dt
+        exp_profile = KhExponential(kh_0, f .* 1000.0)
+        kh_profile = KhExponentialConstant(exp_profile, z_exp .* 0.001)
+    elseif kh_profile_type == "layered" || kh_profile_type == "layered_exponential"
+        kh_profile = KhLayered(fill(mv, n_cells))
+    end
+    parameters = LateralSsfParameters(
+        kh_profile,
+        khfrac,
+        soilthickness,
+        theta_s,
+        theta_r,
+        slope,
+        flow_length,
+        flow_width,
+    )
+    return parameters
+end
+
+"Struct for storing lateral subsurface flow model variables"
+@get_units @grid_loc @with_kw struct LateralSsfVariables{T}
+    zi::Vector{T} | "m"                    # Pseudo-water table depth [m] (top of the saturated zone)
+    exfiltwater::Vector{T} | "m dt-1"      # Exfiltration [m Δt⁻¹] (groundwater above surface level, saturated excess conditions)
+    recharge::Vector{T} | "m2 dt-1"        # Net recharge to saturated store [m² Δt⁻¹]
+    ssf::Vector{T} | "m3 d-1"              # Subsurface flow [m³ d⁻¹]
+    ssfin::Vector{T} | "m3 d-1"            # Inflow from upstream cells [m³ d⁻¹]
+    ssfmax::Vector{T} | "m2 d-1"           # Maximum subsurface flow [m² d⁻¹]
+    to_river::Vector{T} | "m3 d-1"         # Part of subsurface flow [m³ d⁻¹] that flows to the river
+    volume::Vector{T} | "m3"               # Subsurface volume [m³]
+end
+
+"Initialize lateral subsurface flow model variables"
+function LateralSsfVariables(ssf, zi, xl, yl)
+    n = length(zi)
+    volume = @. (ssf.theta_s - ssf.theta_r) * (ssf.soilthickness - zi) * (xl * yl)
+    variables = LateralSsfVariables(;
+        zi,
+        exfiltwater = fill(mv, n),
+        recharge = fill(mv, n),
+        ssf = fill(mv, n),
+        ssfin = fill(mv, n),
+        ssfmax = fill(mv, n),
+        to_river = zeros(n),
+        volume,
+    )
+    return variables
+end
+
+"Struct for storing lateral subsurface flow model boundary conditions"
+@get_units @grid_loc @with_kw struct LateralSsfBC{T}
+    recharge::Vector{T} | "m2 dt-1"        # Net recharge to saturated store [m² Δt⁻¹]
+end
+
+"Lateral subsurface flow model"
+@with_kw struct LateralSSF{T, Kh} <: SubsurfaceFlow
+    boundary_conditions::LateralSsfBC{T}
+    parameters::LateralSsfParameters{T, Kh}
+    variables::LateralSsfVariables{T}
+end
+
+"Initialize lateral subsurface flow model"
+function LateralSSF(
+    dataset,
+    config,
+    indices;
+    soil,
+    slope,
+    flow_length,
+    flow_width,
+    x_length,
+    y_length,
+)
+    parameters = LateralSsfParameters(
+        dataset,
+        config,
+        indices;
+        soil = soil.parameters,
+        slope,
+        flow_length,
+        flow_width,
+    )
+    zi = 0.001 * soil.variables.zi
+    variables = LateralSsfVariables(parameters, zi, x_length, y_length)
+    boundary_conditions = LateralSsfBC(; recharge = fill(mv, length(zi)))
+    ssf = LateralSSF(; boundary_conditions, parameters, variables)
+    return ssf
+end
+
+"Update lateral subsurface model for a single timestep"
+function update!(model::LateralSSF, network, dt)
+    (;
+        order_of_subdomains,
+        order_subdomain,
+        subdomain_indices,
+        upstream_nodes,
+        area,
+        frac_to_river,
+    ) = network
+
+    (; recharge) = model.boundary_conditions
+    (; ssfin, ssf, ssfmax, to_river, zi, exfiltwater, volume) = model.variables
+    (; slope, theta_s, theta_r, soilthickness, flow_length, flow_width, kh_profile) =
+        model.parameters
+
+    ns = length(order_of_subdomains)
+    for k in 1:ns
+        threaded_foreach(eachindex(order_of_subdomains[k]); basesize = 1) do i
+            m = order_of_subdomains[k][i]
+            for (n, v) in zip(subdomain_indices[m], order_subdomain[m])
+                # for a river cell without a reservoir or lake part of the upstream
+                # subsurface flow goes to the river (frac_to_river) and part goes to the
+                # subsurface flow reservoir (1.0 - frac_to_river) upstream nodes with a
+                # reservoir or lake are excluded
+                ssfin[v] = sum_at(
+                    i -> ssf[i] * (1.0 - frac_to_river[i]),
+                    upstream_nodes[n],
+                    eltype(ssfin),
+                )
+                to_river[v] = sum_at(
+                    i -> ssf[i] * frac_to_river[i],
+                    upstream_nodes[n],
+                    eltype(to_river),
+                )
+                ssf[v], zi[v], exfiltwater[v] = kinematic_wave_ssf(
+                    ssfin[v],
+                    ssf[v],
+                    zi[v],
+                    recharge[v],
+                    slope[v],
+                    theta_s[v] - theta_r[v],
+                    soilthickness[v],
+                    dt,
+                    flow_length[v],
+                    flow_width[v],
+                    ssfmax[v],
+                    kh_profile,
+                    v,
+                )
+                volume[v] = (theta_s[v] - theta_r[v]) * (soilthickness[v] - zi[v]) * area[v]
+            end
+        end
+    end
+    return nothing
+end
+
+"""
+Struct for storing groundwater exchange variables for coupling with an external groundwater
+model.
+"""
+@get_units@grid_loc @with_kw struct GroundwaterExchangeVariables{T}
+    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
+
+"Initialize groundwater exchange variables"
+function GroundwaterExchangeVariables(n)
+    variables = GroundwaterExchangeVariables{Float}(;
+        exfiltwater = fill(mv, n),
+        zi = fill(mv, n),
+        to_river = fill(mv, n),
+        ssf = zeros(n),
+    )
+    return variables
+end
+
+"Groundwater exchange"
+@with_kw struct GroundwaterExchange{T} <: SubsurfaceFlow
+    variables::GroundwaterExchangeVariables{T}
+end
+
+"Initialize groundwater exchange"
+function GroundwaterExchange(n)
+    variables = GroundwaterExchangeVariables(n)
+    ssf = GroundwaterExchange{Float}(; variables)
+    return ssf
+end
+
+# wrapper methods
+get_water_depth(subsurface::SubsurfaceFlow) = subsurface.variables.zi
+get_exfiltwater(subsurface::SubsurfaceFlow) = subsurface.variables.exfiltwater
+
+get_flux_to_river(subsurface::SubsurfaceFlow) =
+    subsurface.variables.to_river ./ tosecond(basetimestep) # [m³ s⁻¹]
\ No newline at end of file
diff --git a/src/routing/surface_kinwave.jl b/src/routing/surface_kinwave.jl
new file mode 100644
index 000000000..7160117aa
--- /dev/null
+++ b/src/routing/surface_kinwave.jl
@@ -0,0 +1,638 @@
+abstract type AbstractRiverFlowModel end
+
+"Struct for storing (shared) variables for river and overland flow models"
+@get_units @grid_loc @with_kw struct FlowVariables{T}
+    q::Vector{T} | "m3 s-1"                 # Discharge [m³ s⁻¹]
+    qlat::Vector{T} | "m2 s-1"              # Lateral inflow per unit length [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⁻¹]
+    volume::Vector{T} | "m3"                # Kinematic wave volume [m³] (based on water depth h)
+    h::Vector{T} | "m"                      # Water depth [m]
+    h_av::Vector{T} | "m"                   # Average water depth [m]
+end
+
+"Initialize timestepping for kinematic wave (river and overland flow models)"
+function init_kinematic_wave_timestepping(config, n; domain, dt_fixed)
+    adaptive = get(config.model, "kin_wave_iteration", false)::Bool
+    @info "Kinematic wave approach is used for $domain flow, adaptive timestepping = $adaptive."
+
+    if adaptive
+        stable_timesteps = zeros(n)
+        timestepping = TimeStepping(; stable_timesteps, adaptive)
+    else
+        dt_fixed = get(config.model, "kw_$(domain)_tstep", dt_fixed)
+        @info "Using a fixed sub-timestep (seconds) $dt_fixed for kinematic wave $domain flow."
+        timestepping = TimeStepping(; dt_fixed, adaptive)
+    end
+    return timestepping
+end
+
+"Initialize variables for river or overland flow models"
+function FlowVariables(n)
+    variables = FlowVariables(;
+        q = zeros(Float, n),
+        qlat = zeros(Float, n),
+        qin = zeros(Float, n),
+        q_av = zeros(Float, n),
+        volume = zeros(Float, n),
+        h = zeros(Float, n),
+        h_av = zeros(Float, n),
+    )
+    return variables
+end
+
+"Struct for storing Manning flow parameters"
+@get_units @grid_loc @with_kw struct ManningFlowParameters{T}
+    beta::T                                 # constant in Manning's equation [-]
+    slope::Vector{T} | "m m-1"              # Slope [m m⁻¹]
+    mannings_n::Vector{T} | "s m-1/3"       # Manning's roughness [s m⁻⅓]
+    flow_length::Vector{T} | "m"            # Flow length [m]
+    flow_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
+end
+
+"Initialize Manning flow parameters"
+function ManningFlowParameters(slope, mannings_n, flow_length, flow_width)
+    n = length(slope)
+    parameters = ManningFlowParameters(;
+        beta = Float(0.6),
+        slope,
+        mannings_n,
+        flow_length,
+        flow_width,
+        alpha_pow = Float((2.0 / 3.0) * 0.6),
+        alpha_term = fill(mv, n),
+        alpha = fill(mv, n),
+    )
+    return parameters
+end
+
+"Struct for storing river flow model parameters"
+@get_units @grid_loc @with_kw struct RiverFlowParameters{T}
+    flow::ManningFlowParameters{T}
+    bankfull_depth::Vector{T} | "m"         # Bankfull water level [m]
+end
+
+"Overload `getproperty` for river flow model parameters"
+function Base.getproperty(v::RiverFlowParameters, s::Symbol)
+    if s === :bankfull_depth
+        getfield(v, s)
+    elseif s === :flow
+        getfield(v, :flow)
+    else
+        getfield(getfield(v, :flow), s)
+    end
+end
+
+"Initialize river flow model parameters"
+function RiverFlowParameters(dataset, config, indices, river_length, river_width)
+    mannings_n = ncread(
+        dataset,
+        config,
+        "lateral.river.mannings_n";
+        sel = indices,
+        defaults = 0.036,
+        type = Float,
+    )
+    bankfull_depth = ncread(
+        dataset,
+        config,
+        "lateral.river.bankfull_depth";
+        alias = "lateral.river.h_bankfull",
+        sel = indices,
+        defaults = 1.0,
+        type = Float,
+    )
+    if haskey(config.input.lateral.river, "h_bankfull")
+        @warn string(
+            "The `h_bankfull` key in `[input.lateral.river]` is now called ",
+            "`bankfull_depth`. Please update your TOML file.",
+        )
+    end
+    slope = ncread(
+        dataset,
+        config,
+        "lateral.river.slope";
+        optional = false,
+        sel = indices,
+        type = Float,
+    )
+    clamp!(slope, 0.00001, Inf)
+
+    flow_parameter_set = ManningFlowParameters(slope, mannings_n, river_length, river_width)
+    parameters =
+        RiverFlowParameters(; flow = flow_parameter_set, bankfull_depth = bankfull_depth)
+    return parameters
+end
+
+"Struct for storing river flow model boundary conditions"
+@get_units @grid_loc @with_kw struct RiverFlowBC{T, R, L}
+    inwater::Vector{T} | "m3 s-1"           # Lateral inflow [m³ s⁻¹]
+    inflow::Vector{T} | "m3 s-1"            # External inflow (abstraction/supply/demand) [m³ s⁻¹]
+    inflow_waterbody::Vector{T} | "m3 s-1"  # inflow waterbody (lake or reservoir model) from land part [m³ s⁻¹]
+    abstraction::Vector{T} | "m3 s-1"       # Abstraction (computed as part of water demand and allocation) [m³ s⁻¹]
+    reservoir::R                            # Reservoir model struct of arrays
+    lake::L                                 # Lake model struct of arrays
+end
+
+"Initialize river flow model boundary conditions"
+function RiverFlowBC(n, reservoir, lake)
+    bc = RiverFlowBC(;
+        inwater = zeros(Float, n),
+        inflow = zeros(Float, n),
+        inflow_waterbody = zeros(Float, n),
+        abstraction = zeros(Float, n),
+        reservoir = reservoir,
+        lake = lake,
+    )
+    return bc
+end
+
+"River flow model using the kinematic wave method and the Manning flow equation"
+@with_kw struct KinWaveRiverFlow{T, R, L, A} <: AbstractRiverFlowModel
+    timestepping::TimeStepping{T}
+    boundary_conditions::RiverFlowBC{T, R, L}
+    parameters::RiverFlowParameters{T}
+    variables::FlowVariables{T}
+    allocation::A   # Water allocation
+end
+
+"Initialize river flow model `KinWaveRiverFlow`"
+function KinWaveRiverFlow(
+    dataset,
+    config,
+    indices;
+    river_length,
+    river_width,
+    reservoir,
+    lake,
+)
+    n = length(indices)
+
+    timestepping =
+        init_kinematic_wave_timestepping(config, n; domain = "river", dt_fixed = 900.0)
+
+    do_water_demand = haskey(config.model, "water_demand")
+    allocation = do_water_demand ? AllocationRiver(n) : NoAllocationRiver{Float}()
+
+    variables = FlowVariables(n)
+    parameters = RiverFlowParameters(dataset, config, indices, river_length, river_width)
+    boundary_conditions = RiverFlowBC(n, reservoir, lake)
+
+    sf_river = KinWaveRiverFlow(;
+        timestepping,
+        boundary_conditions,
+        parameters,
+        variables,
+        allocation,
+    )
+
+    return sf_river
+end
+
+"Struct for storing overland flow model variables"
+@get_units @grid_loc @with_kw struct LandFlowVariables{T}
+    flow::FlowVariables{T}
+    to_river::Vector{T} | "m3 s-1"      # Part of overland flow [m³ s⁻¹] that flows to the river
+end
+
+"Overload `getproperty` for overland flow model variables"
+function Base.getproperty(v::LandFlowVariables, s::Symbol)
+    if s === :to_river
+        getfield(v, s)
+    elseif s === :flow
+        getfield(v, :flow)
+    else
+        getfield(getfield(v, :flow), s)
+    end
+end
+
+"Struct for storing overland flow model boundary conditions"
+@get_units @grid_loc @with_kw struct LandFlowBC{T}
+    inwater::Vector{T} | "m3 s-1"       # Lateral inflow [m³ s⁻¹]
+end
+
+"Overland flow model using the kinematic wave method and the Manning flow equation"
+@with_kw struct KinWaveOverlandFlow{T}
+    timestepping::TimeStepping{T}
+    boundary_conditions::LandFlowBC{T}
+    parameters::ManningFlowParameters{T}
+    variables::LandFlowVariables{T}
+end
+
+"Initialize Overland flow model `KinWaveOverlandFlow`"
+function KinWaveOverlandFlow(dataset, config, indices; slope, flow_length, flow_width)
+    mannings_n = ncread(
+        dataset,
+        config,
+        "lateral.land.mannings_n";
+        sel = indices,
+        defaults = 0.072,
+        type = Float,
+    )
+
+    n = length(indices)
+    timestepping =
+        init_kinematic_wave_timestepping(config, n; domain = "land", dt_fixed = 3600.0)
+
+    variables = LandFlowVariables(; flow = FlowVariables(n), to_river = zeros(Float, n))
+    parameters = ManningFlowParameters(slope, mannings_n, flow_length, flow_width)
+    boundary_conditions = LandFlowBC(; inwater = zeros(Float, n))
+    sf_land =
+        KinWaveOverlandFlow(; timestepping, boundary_conditions, variables, parameters)
+
+    return sf_land
+end
+
+"""
+Helper function to set waterbody variables `inflow`,`outflow_av` and `actevap` to zero. This
+is done at the start of each simulation timestep, during the timestep the total (weighted)
+sum is computed from values at each sub timestep.
+"""
+function set_waterbody_vars!(waterbody::W) where {W <: Union{SimpleReservoir, Lake}}
+    waterbody.boundary_conditions.inflow .= 0.0
+    waterbody.variables.outflow_av .= 0.0
+    waterbody.variables.actevap .= 0.0
+    return nothing
+end
+set_waterbody_vars!(waterbody) = nothing
+
+"""
+Helper function to compute the average of waterbody variables `inflow` and `outflow_av`. This
+is done at the end of each simulation timestep.
+"""
+function average_waterbody_vars!(waterbody::W, dt) where {W <: Union{SimpleReservoir, Lake}}
+    waterbody.variables.outflow_av ./= dt
+    waterbody.boundary_conditions.inflow ./= dt
+end
+average_waterbody_vars!(waterbody, dt) = nothing
+
+"Update overland flow model `KinWaveOverlandFlow` for a single timestep"
+function kinwave_land_update!(model::KinWaveOverlandFlow, network, dt)
+    (;
+        order_of_subdomains,
+        order_subdomain,
+        subdomain_indices,
+        upstream_nodes,
+        frac_to_river,
+    ) = network
+
+    (; beta, alpha, flow_width, flow_length) = model.parameters
+    (; h, h_av, q, q_av, qin, qlat, to_river) = model.variables
+
+    ns = length(order_of_subdomains)
+    qin .= 0.0
+    for k in 1:ns
+        threaded_foreach(eachindex(order_of_subdomains[k]); basesize = 1) do i
+            m = order_of_subdomains[k][i]
+            for (n, v) in zip(subdomain_indices[m], order_subdomain[m])
+                # for a river cell without a reservoir or lake part of the upstream
+                # surface flow goes to the river (frac_to_river) and part goes to the
+                # surface flow reservoir (1.0 - frac_to_river), upstream nodes with a
+                # reservoir or lake are excluded
+                to_river[v] +=
+                    sum_at(
+                        i -> q[i] * frac_to_river[i],
+                        upstream_nodes[n],
+                        eltype(to_river),
+                    ) * dt
+                if flow_width[v] > 0.0
+                    qin[v] = sum_at(
+                        i -> q[i] * (1.0 - frac_to_river[i]),
+                        upstream_nodes[n],
+                        eltype(q),
+                    )
+                end
+
+                q[v] = kinematic_wave(
+                    qin[v],
+                    q[v],
+                    qlat[v],
+                    alpha[v],
+                    beta,
+                    dt,
+                    flow_length[v],
+                )
+
+                # update h, only if flow width > 0.0
+                if flow_width[v] > 0.0
+                    crossarea = alpha[v] * pow(q[v], beta)
+                    h[v] = crossarea / flow_width[v]
+                end
+                q_av[v] += q[v] * dt
+                h_av[v] += h[v] * dt
+            end
+        end
+    end
+end
+
+"""
+Update overland flow model `KinWaveOverlandFlow` for a single timestep `dt`. Timestepping within
+`dt` is either with a fixed timestep `dt_fixed` or adaptive.
+"""
+function update!(model::KinWaveOverlandFlow, network, dt)
+    (; inwater) = model.boundary_conditions
+    (; alpha_term, mannings_n, slope, beta, alpha_pow, alpha, flow_width, flow_length) =
+        model.parameters
+    (; h, h_av, q_av, qlat, volume, to_river) = model.variables
+    (; adaptive) = model.timestepping
+
+    @. alpha_term = pow(mannings_n / sqrt(slope), beta)
+    # use fixed alpha value based flow width
+    @. alpha = alpha_term * pow(flow_width, alpha_pow)
+    @. qlat = inwater / flow_length
+
+    q_av .= 0.0
+    h_av .= 0.0
+    to_river .= 0.0
+
+    t = 0.0
+    while t < dt
+        dt_s = adaptive ? stable_timestep(model, 0.02) : model.timestepping.dt_fixed
+        dt_s = check_timestepsize(dt_s, t, dt)
+        kinwave_land_update!(model, network, dt_s)
+        t = t + dt_s
+    end
+    q_av ./= dt
+    h_av ./= dt
+    to_river ./= dt
+    volume .= flow_length .* flow_width .* h
+    return nothing
+end
+
+"Update river flow model `KinWaveRiverFlow` for a single timestep"
+function kinwave_river_update!(model::KinWaveRiverFlow, network, doy, dt, dt_forcing)
+    (;
+        graph,
+        order_of_subdomains,
+        order_subdomain,
+        subdomain_indices,
+        upstream_nodes,
+        reservoir_indices,
+        lake_indices,
+    ) = network.river
+
+    (; reservoir, lake, inwater, inflow, abstraction, inflow_waterbody) =
+        model.boundary_conditions
+
+    (; beta, alpha, flow_width, flow_length) = model.parameters
+    (; h, h_av, q, q_av, qin, qlat, volume) = model.variables
+
+    ns = length(order_of_subdomains)
+    qin .= 0.0
+    for k in 1:ns
+        threaded_foreach(eachindex(order_of_subdomains[k]); basesize = 1) do i
+            m = order_of_subdomains[k][i]
+            for (n, v) in zip(subdomain_indices[m], order_subdomain[m])
+                # qin by outflow from upstream reservoir or lake location is added
+                qin[v] += sum_at(q, upstream_nodes[n])
+                # Inflow supply/abstraction is added to qlat (divide by flow length)
+                # If inflow < 0, abstraction is limited
+                if inflow[v] < 0.0
+                    max_abstract =
+                        min((inwater[v] + qin[v] + volume[v] / dt) * 0.80, -inflow[v])
+                    _inflow = -max_abstract / flow_length[v]
+                else
+                    _inflow = inflow[v] / flow_length[v]
+                end
+                _inflow -= abstraction[v] / flow_length[v]
+
+                q[v] = kinematic_wave(
+                    qin[v],
+                    q[v],
+                    qlat[v] + _inflow,
+                    alpha[v],
+                    beta,
+                    dt,
+                    flow_length[v],
+                )
+
+                if !isnothing(reservoir) && reservoir_indices[v] != 0
+                    # run reservoir model and copy reservoir outflow to inflow (qin) of
+                    # downstream river cell
+                    i = reservoir_indices[v]
+                    update!(reservoir, i, q[v] + inflow_waterbody[v], dt, dt_forcing)
+
+                    downstream_nodes = outneighbors(graph, v)
+                    n_downstream = length(downstream_nodes)
+                    if n_downstream == 1
+                        j = only(downstream_nodes)
+                        qin[j] = reservoir.variables.outflow[i]
+                    elseif n_downstream == 0
+                        error(
+                            """A reservoir without a downstream river node is not supported.
+                            Add a downstream river node or move the reservoir to an upstream node (model schematization).
+                            """,
+                        )
+                    else
+                        error("bifurcations not supported")
+                    end
+
+                elseif !isnothing(lake) && lake_indices[v] != 0
+                    # run lake model and copy lake outflow to inflow (qin) of downstream river
+                    # cell
+                    i = lake_indices[v]
+                    update!(lake, i, q[v] + inflow_waterbody[v], doy, dt, dt_forcing)
+
+                    downstream_nodes = outneighbors(graph, v)
+                    n_downstream = length(downstream_nodes)
+                    if n_downstream == 1
+                        j = only(downstream_nodes)
+                        qin[j] = max(lake.variables.outflow[i], 0.0)
+                    elseif n_downstream == 0
+                        error(
+                            """A lake without a downstream river node is not supported.
+                            Add a downstream river node or move the lake to an upstream node (model schematization).
+                            """,
+                        )
+                    else
+                        error("bifurcations not supported")
+                    end
+                end
+                # update h
+                crossarea = alpha[v] * pow(q[v], beta)
+                h[v] = crossarea / flow_width[v]
+                volume[v] = flow_length[v] * flow_width[v] * h[v]
+                q_av[v] += q[v] * dt
+                h_av[v] += h[v] * dt
+            end
+        end
+    end
+end
+
+"""
+Update river flow model `KinWaveRiverFlow` for a single timestep `dt`. Timestepping within
+`dt` is either with a fixed timestep `dt_fixed` or adaptive.
+"""
+function update!(model::KinWaveRiverFlow, network, doy, dt)
+    (; reservoir, lake, inwater) = model.boundary_conditions
+
+    (;
+        alpha_term,
+        mannings_n,
+        slope,
+        beta,
+        alpha_pow,
+        alpha,
+        flow_width,
+        flow_length,
+        bankfull_depth,
+    ) = model.parameters
+    (; h, h_av, q_av, qlat, volume) = model.variables
+    (; adaptive) = model.timestepping
+
+    @. alpha_term = pow(mannings_n / sqrt(slope), beta)
+    # use fixed alpha value based on 0.5 * bankfull_depth
+    @. alpha = alpha_term * pow(flow_width + bankfull_depth, alpha_pow)
+    @. qlat = inwater / flow_length
+
+    q_av .= 0.0
+    h_av .= 0.0
+
+    set_waterbody_vars!(reservoir)
+    set_waterbody_vars!(lake)
+
+    t = 0.0
+    while t < dt
+        dt_s = adaptive ? stable_timestep(model, 0.05) : model.timestepping.dt_fixed
+        dt_s = check_timestepsize(dt_s, t, dt)
+        kinwave_river_update!(model, network, doy, dt_s, dt)
+        t = t + dt_s
+    end
+
+    average_waterbody_vars!(reservoir, dt)
+    average_waterbody_vars!(lake, dt)
+
+    q_av ./= dt
+    h_av ./= dt
+    volume .= flow_length .* flow_width .* h
+    return nothing
+end
+
+"""
+Compute a stable timestep size for the kinematice wave method for a river or overland flow
+model using a nonlinear scheme (Chow et al., 1988). 
+
+A stable time step is computed for each vector element based on the Courant timestep size
+criterion. A quantile of the vector is computed based on probability `p` to remove potential
+very low timestep sizes. Li et al. (1975) found that the nonlinear scheme is unconditonally
+stable and that a wide range of dt/dx values can be used without loss of accuracy.
+"""
+function stable_timestep(
+    model::S,
+    p,
+) where {S <: Union{KinWaveOverlandFlow, KinWaveRiverFlow}}
+    (; q) = model.variables
+    (; alpha, beta, flow_length) = model.parameters
+    (; stable_timesteps) = model.timestepping
+
+    n = length(q)
+    stable_timesteps .= Inf
+    k = 0
+    for i in 1:n
+        if q[i] > 0.0
+            k += 1
+            c = 1.0 / (alpha[i] * beta * pow(q[i], (beta - 1.0)))
+            stable_timesteps[k] = (flow_length[i] / c)
+        end
+    end
+
+    dt_min = if k == 1
+        stable_timesteps[k]
+    elseif k > 0
+        quantile!(@view(stable_timesteps[1:k]), p)
+    else
+        600.0
+    end
+
+    return dt_min
+end
+
+"""
+Update boundary condition lateral inflow `inwater` of a river flow model for a single
+timestep.
+"""
+function update_lateral_inflow!(
+    model::AbstractRiverFlowModel,
+    external_models::NamedTuple,
+    river_cell_area,
+    land_area,
+    river_indices,
+    dt,
+)
+    (; allocation, runoff, land, subsurface) = external_models
+    (; inwater) = model.boundary_conditions
+    (; net_runoff_river) = runoff.variables
+
+    inwater .= (
+        get_flux_to_river(subsurface)[river_indices] .+
+        land.variables.to_river[river_indices] .+
+        (net_runoff_river[river_indices] .* land_area[river_indices] .* 0.001) ./ dt .+
+        (get_nonirrigation_returnflow(allocation) .* 0.001 .* river_cell_area) ./ dt
+    )
+    return nothing
+end
+
+"""
+Update boundary condition lateral inflow `inwater` of a kinematic wave overland flow model
+`KinWaveOverlandFlow` for a single timestep.
+"""
+function update_lateral_inflow!(
+    model::KinWaveOverlandFlow,
+    external_models::NamedTuple,
+    area,
+    config,
+    dt,
+)
+    (; soil, subsurface, allocation) = external_models
+    (; net_runoff) = soil.variables
+    (; inwater) = model.boundary_conditions
+
+    do_drains = get(config.model, "drains", false)::Bool
+    if do_drains
+        drainflux = zeros(length(net_runoff))
+        drainflux[subsurface.drain.index] =
+            -subsurface.drain.variables.flux ./ tosecond(basetimestep)
+    else
+        drainflux = 0.0
+    end
+    inwater .=
+        (net_runoff .+ get_nonirrigation_returnflow(allocation)) .* area * 0.001 ./ dt .+
+        drainflux
+
+    return nothing
+end
+
+"""
+Update boundary condition inflow to a waterbody from land `inflow_waterbody` of a model
+`AbstractRiverFlowModel` for a single timestep.
+"""
+function update_inflow_waterbody!(
+    model::AbstractRiverFlowModel,
+    external_models::NamedTuple,
+    river_indices,
+)
+    (; land, subsurface) = external_models
+    (; reservoir, lake, inflow_waterbody) = model.boundary_conditions
+
+    if !isnothing(reservoir) || !isnothing(lake)
+        inflow_land = get_inflow_waterbody(model, land)
+        inflow_subsurface = get_inflow_waterbody(model, subsurface)
+
+        @. inflow_waterbody = inflow_land[river_indices] + inflow_subsurface[river_indices]
+    end
+    return nothing
+end
+
+# For the river kinematic wave, the variable `to_river` can be excluded, because this part
+# is added to the river kinematic wave.
+get_inflow_waterbody(::KinWaveRiverFlow, model::KinWaveOverlandFlow) = model.variables.q_av
+get_inflow_waterbody(::KinWaveRiverFlow, model::LateralSSF) =
+    model.variables.ssf ./ tosecond(basetimestep)
+
+# Exclude subsurface flow for other groundwater components than `LateralSSF`.
+get_inflow_waterbody(::AbstractRiverFlowModel, model::GroundwaterFlow) =
+    model.flow.connectivity.ncell .* 0.0
+get_inflow_waterbody(::KinWaveRiverFlow, model) = model.variables.to_river .* 0.0
\ No newline at end of file
diff --git a/src/routing/surface_local_inertial.jl b/src/routing/surface_local_inertial.jl
new file mode 100644
index 000000000..3050a9a11
--- /dev/null
+++ b/src/routing/surface_local_inertial.jl
@@ -0,0 +1,1409 @@
+"Struct for storing local inertial river flow model parameters"
+@get_units @grid_loc @with_kw struct LocalInertialRiverFlowParameters{T}
+    n::Int                                              # number of cells [-]
+    ne::Int                                             # number of edges [-]
+    active_n::Vector{Int} | "-"                         # active nodes [-]
+    active_e::Vector{Int} | "-" | "edge"                # active edges [-]
+    g::T                                                # acceleration due to gravity [m s⁻²]
+    froude_limit::Bool                                  # if true a check is performed if froude number > 1.0 (algorithm is modified) [-]
+    h_thresh::T                                         # depth threshold for calculating flow [m]
+    zb::Vector{T} | "m"                                 # river bed elevation   
+    zb_max::Vector{T} | "m"                             # maximum channel bed elevation
+    bankfull_volume::Vector{T} | "m3"                   # bankfull volume
+    bankfull_depth::Vector{T} | "m"                     # bankfull depth
+    mannings_n_sq::Vector{T} | "(s m-1/3)2" | "edge"    # Manning's roughness squared at edge
+    mannings_n::Vector{T} | "s m-1/3"                   # Manning's roughness at node
+    flow_length::Vector{T} | "m"                        # flow (river) length
+    flow_length_at_edge::Vector{T} | "m" | "edge"       # flow (river) length at edge
+    flow_width::Vector{T} | "m"                         # flow (river) width
+    flow_width_at_edge::Vector{T} | "m" | "edge"        # flow (river) width at edge
+    waterbody::Vector{Bool} | "-"                       # water body cells (reservoir or lake)
+end
+
+"Initialize local inertial river flow model parameters"
+function LocalInertialRiverFlowParameters(
+    dataset,
+    config,
+    indices;
+    river_length,
+    river_width,
+    waterbody,
+    n_edges,
+    nodes_at_edge,
+    index_pit,
+    inds_pit,
+)
+    cfl = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
+    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at edge
+    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number
+    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
+
+    @info "Local inertial approach is used for river flow." cfl h_thresh froude_limit floodplain_1d
+    @warn string(
+        "Providing the boundary condition `riverlength_bc` as part of the `[model]` setting ",
+        "in the TOML file has been deprecated as of Wflow v0.8.0.\n The boundary condition should ",
+        "be provided as part of the file `$(config.input.path_static)`.",
+    )
+
+    riverlength_bc = ncread(
+        dataset,
+        config,
+        "lateral.river.riverlength_bc";
+        sel = inds_pit,
+        defaults = 1.0e04,
+        type = Float,
+    )
+    bankfull_elevation_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.bankfull_elevation";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    bankfull_depth_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.bankfull_depth";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    bankfull_depth = bankfull_depth_2d[indices]
+    zb = bankfull_elevation_2d[indices] - bankfull_depth # river bed elevation
+
+    bankfull_volume = bankfull_depth .* river_width .* river_length
+    mannings_n = ncread(
+        dataset,
+        config,
+        "lateral.river.mannings_n";
+        sel = indices,
+        defaults = 0.036,
+        type = Float,
+    )
+
+    n = length(indices)
+
+    # set ghost points for boundary condition (downstream river outlet): river width, bed
+    # elevation, manning n is copied from the upstream cell.
+    append!(river_length, riverlength_bc)
+    append!(zb, zb[index_pit])
+    append!(river_width, river_width[index_pit])
+    append!(mannings_n, mannings_n[index_pit])
+    append!(bankfull_depth, bankfull_depth[index_pit])
+
+    # determine z, width, length and manning's n at edges
+    zb_max = fill(Float(0), n_edges)
+    width_at_edge = fill(Float(0), n_edges)
+    length_at_edge = fill(Float(0), n_edges)
+    mannings_n_sq = fill(Float(0), n_edges)
+    for i in 1:n_edges
+        src_node = nodes_at_edge.src[i]
+        dst_node = nodes_at_edge.dst[i]
+        zb_max[i] = max(zb[src_node], zb[dst_node])
+        width_at_edge[i] = min(river_width[src_node], river_width[dst_node])
+        length_at_edge[i] = 0.5 * (river_length[dst_node] + river_length[src_node])
+        mannings_n_i =
+            (
+                mannings_n[dst_node] * river_length[dst_node] +
+                mannings_n[src_node] * river_length[src_node]
+            ) / (river_length[dst_node] + river_length[src_node])
+        mannings_n_sq[i] = mannings_n_i * mannings_n_i
+    end
+    active_index = findall(x -> x == 0, waterbody)
+
+    parameters = LocalInertialRiverFlowParameters(;
+        n,
+        ne = n_edges,
+        active_n = active_index,
+        active_e = active_index,
+        g = 9.80665,
+        froude_limit,
+        h_thresh,
+        zb,
+        zb_max,
+        bankfull_volume,
+        bankfull_depth,
+        mannings_n,
+        mannings_n_sq,
+        flow_length = river_length,
+        flow_length_at_edge = length_at_edge,
+        flow_width = river_width,
+        flow_width_at_edge = width_at_edge,
+        waterbody,
+    )
+    return parameters
+end
+
+"Struct for storing local inertial river flow model variables"
+@get_units @grid_loc @with_kw struct LocalInertialRiverFlowVariables{T}
+    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⁻¹]
+    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
+    h_av::Vector{T} | "m"                               # average water depth
+    a::Vector{T} | "m2" | "edge"                        # flow area at edge
+    r::Vector{T} | "m" | "edge"                         # wetted perimeter at edge
+    volume::Vector{T} | "m3"                            # river volume
+    error::Vector{T} | "m3"                             # error volume    
+end
+
+"Initialize shallow water river flow model variables"
+function LocalInertialRiverFlowVariables(dataset, config, indices, n_edges, inds_pit)
+    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
+    riverdepth_bc = ncread(
+        dataset,
+        config,
+        "lateral.river.riverdepth_bc";
+        sel = inds_pit,
+        defaults = 0.0,
+        type = Float,
+    )
+
+    n = length(indices)
+    # set river depth h to zero (including reservoir and lake locations)
+    h = zeros(n)
+    q_av = zeros(n_edges)
+    # set ghost points for boundary condition (downstream river outlet): river depth `h`
+    append!(h, riverdepth_bc)
+    variables = LocalInertialRiverFlowVariables(;
+        q = zeros(n_edges),
+        q0 = zeros(n_edges),
+        q_av = q_av,
+        q_channel_av = floodplain_1d ? zeros(n_edges) : q_av,
+        h = h,
+        zs_max = zeros(n_edges),
+        zs_src = zeros(n_edges),
+        zs_dst = zeros(n_edges),
+        hf = zeros(n_edges),
+        h_av = zeros(n),
+        a = zeros(n_edges),
+        r = zeros(n_edges),
+        volume = zeros(n),
+        error = zeros(n),
+    )
+    return variables
+end
+
+"Shallow water river flow model using the local inertial method"
+@with_kw struct LocalInertialRiverFlow{T, R, L, F, A} <: AbstractRiverFlowModel
+    timestepping::TimeStepping{T}
+    boundary_conditions::RiverFlowBC{T, R, L}
+    parameters::LocalInertialRiverFlowParameters{T}
+    variables::LocalInertialRiverFlowVariables{T}
+    floodplain::F                                       # Floodplain (1D) schematization
+    allocation::A                                       # Water allocation
+end
+
+"Initialize shallow water river flow model `LocalIntertialRiverFlow`"
+function LocalInertialRiverFlow(
+    dataset,
+    config,
+    indices;
+    graph_river,
+    ldd_river,
+    river_length,
+    river_width,
+    reservoir,
+    lake,
+    waterbody,
+)
+    # The local inertial approach makes use of a staggered grid (Bates et al. (2010)),
+    # with nodes and edges. This information is extracted from the directed graph of the
+    # river. Discharge q is calculated at edges between nodes and mapped to the source
+    # nodes for gridded output (index of edge is equal to source node index, e.g.:
+    # Edge 1 => 5
+    # Edge 2 => 1
+    # Edge 3 => 2
+    # Edge 4 => 9
+    # ⋮ )
+
+    # The following boundary conditions can be set at ghost nodes, downstream of river
+    # outlets (pits): river length and river depth
+    cfl = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
+    timestepping = TimeStepping(; cfl)
+
+    index_pit = findall(x -> x == 5, ldd_river)
+    inds_pit = indices[index_pit]
+
+    add_vertex_edge_graph!(graph_river, index_pit)
+    nodes_at_edge = adjacent_nodes_at_edge(graph_river)
+    n_edges = ne(graph_river)
+
+    parameters = LocalInertialRiverFlowParameters(
+        dataset,
+        config,
+        indices;
+        river_length,
+        river_width,
+        waterbody,
+        n_edges,
+        nodes_at_edge,
+        index_pit,
+        inds_pit,
+    )
+    variables = LocalInertialRiverFlowVariables(dataset, config, indices, n_edges, inds_pit)
+
+    n = length(indices)
+    boundary_conditions = RiverFlowBC(n, reservoir, lake)
+
+    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
+    if floodplain_1d
+        zb_floodplain = parameters.zb .+ parameters.bankfull_depth
+        floodplain = FloodPlain(
+            dataset,
+            config,
+            indices;
+            river_width,
+            river_length,
+            zb_floodplain,
+            index_pit,
+            n_edges,
+            nodes_at_edge,
+        )
+    else
+        floodplain = nothing
+    end
+
+    do_water_demand = haskey(config.model, "water_demand")
+    sw_river = LocalInertialRiverFlow(;
+        timestepping,
+        boundary_conditions,
+        parameters,
+        variables,
+        floodplain,
+        allocation = do_water_demand ? AllocationRiver(n) : NoAllocationRiver{Float}(),
+    )
+    return sw_river, nodes_at_edge
+end
+
+"Return the upstream inflow for a waterbody in `LocalInertialRiverFlow`"
+function get_inflow_waterbody(model::LocalInertialRiverFlow, src_edge)
+    q_in = sum_at(model.variables.q, src_edge)
+    if !isnothing(model.floodplain)
+        q_in = q_in + sum_at(model.floodplain.variables.q, src_edge)
+    end
+    return q_in
+end
+
+# For local inertial river routing, `to_river` is included, as water body cells are excluded
+# (boundary condition).
+get_inflow_waterbody(::LocalInertialRiverFlow, model::KinWaveOverlandFlow) =
+    model.variables.q_av .+ model.variables.to_river
+get_inflow_waterbody(::LocalInertialRiverFlow, model::LateralSSF) =
+    (model.variables.ssf .+ model.variables.to_river) ./ tosecond(basetimestep)
+
+"Update local inertial river flow model `LocalIntertialRiverFlow` for a single timestep"
+function local_inertial_river_update!(
+    model::LocalInertialRiverFlow,
+    network,
+    dt,
+    dt_forcing,
+    doy,
+    update_h,
+)
+    (; nodes_at_edge, edges_at_node) = network.river
+    (; inwater, abstraction, inflow) = model.boundary_conditions
+    river_v = model.variables
+    river_p = model.parameters
+
+    river_v.q0 .= river_v.q
+    if !isnothing(model.floodplain)
+        model.floodplain.variables.q0 .= model.floodplain.variables.q
+    end
+    @tturbo for j in eachindex(river_p.active_e)
+        i = river_p.active_e[j]
+        i_src = nodes_at_edge.src[i]
+        i_dst = nodes_at_edge.dst[i]
+        river_v.zs_src[i] = river_p.zb[i_src] + river_v.h[i_src]
+        river_v.zs_dst[i] = river_p.zb[i_dst] + river_v.h[i_dst]
+
+        river_v.zs_max[i] = max(river_v.zs_src[i], river_v.zs_dst[i])
+        river_v.hf[i] = (river_v.zs_max[i] - river_p.zb_max[i])
+
+        river_v.a[i] = river_p.flow_width_at_edge[i] * river_v.hf[i] # flow area (rectangular channel)
+        river_v.r[i] = river_v.a[i] / (river_p.flow_width_at_edge[i] + 2.0 * river_v.hf[i]) # hydraulic radius (rectangular channel)
+
+        river_v.q[i] = IfElse.ifelse(
+            river_v.hf[i] > river_p.h_thresh,
+            local_inertial_flow(
+                river_v.q0[i],
+                river_v.zs_src[i],
+                river_v.zs_dst[i],
+                river_v.hf[i],
+                river_v.a[i],
+                river_v.r[i],
+                river_p.flow_length_at_edge[i],
+                river_p.mannings_n_sq[i],
+                river_p.g,
+                river_p.froude_limit,
+                dt,
+            ),
+            0.0,
+        )
+
+        # limit q in case water is not available
+        river_v.q[i] =
+            IfElse.ifelse(river_v.h[i_src] <= 0.0, min(river_v.q[i], 0.0), river_v.q[i])
+        river_v.q[i] =
+            IfElse.ifelse(river_v.h[i_dst] <= 0.0, max(river_v.q[i], 0.0), river_v.q[i])
+
+        river_v.q_av[i] += river_v.q[i] * dt
+    end
+    if !isnothing(model.floodplain)
+        floodplain_p = model.floodplain.parameters
+        floodplain_v = model.floodplain.variables
+
+        @tturbo @. floodplain_v.hf = max(river_v.zs_max - floodplain_p.zb_max, 0.0)
+
+        n = 0
+        @inbounds for i in river_p.active_e
+            @inbounds if river_v.hf[i] > river_p.h_thresh
+                n += 1
+                floodplain_v.hf_index[n] = i
+            else
+                floodplain_v.q[i] = 0.0
+            end
+        end
+
+        @tturbo for j in 1:n
+            i = floodplain_v.hf_index[j]
+            i_src = nodes_at_edge.src[i]
+            i_dst = nodes_at_edge.dst[i]
+
+            i0 = 0
+            for k in eachindex(floodplain_p.profile.depth)
+                i0 += 1 * (floodplain_p.profile.depth[k] <= floodplain_v.hf[i])
+            end
+            i1 = max(i0, 1)
+            i2 = IfElse.ifelse(i1 == length(floodplain_p.profile.depth), i1, i1 + 1)
+
+            a_src = flow_area(
+                floodplain_p.profile.width[i2, i_src],
+                floodplain_p.profile.a[i1, i_src],
+                floodplain_p.profile.depth[i1],
+                floodplain_v.hf[i],
+            )
+            a_src = max(a_src - (floodplain_v.hf[i] * river_p.flow_width[i_src]), 0.0)
+
+            a_dst = flow_area(
+                floodplain_p.profile.width[i2, i_dst],
+                floodplain_p.profile.a[i1, i_dst],
+                floodplain_p.profile.depth[i1],
+                floodplain_v.hf[i],
+            )
+            a_dst = max(a_dst - (floodplain_v.hf[i] * river_p.flow_width[i_dst]), 0.0)
+
+            floodplain_v.a[i] = min(a_src, a_dst)
+
+            floodplain_v.r[i] = IfElse.ifelse(
+                a_src < a_dst,
+                a_src / wetted_perimeter(
+                    floodplain_p.profile.p[i1, i_src],
+                    floodplain_p.profile.depth[i1],
+                    floodplain_v.hf[i],
+                ),
+                a_dst / wetted_perimeter(
+                    floodplain_p.profile.p[i1, i_dst],
+                    floodplain_p.profile.depth[i1],
+                    floodplain_v.hf[i],
+                ),
+            )
+
+            floodplain_v.q[i] = IfElse.ifelse(
+                floodplain_v.a[i] > 1.0e-05,
+                local_inertial_flow(
+                    floodplain_v.q0[i],
+                    river_v.zs_src[i],
+                    river_v.zs_dst[i],
+                    floodplain_v.hf[i],
+                    floodplain_v.a[i],
+                    floodplain_v.r[i],
+                    river_p.flow_length_at_edge[i],
+                    floodplain_p.mannings_n_sq[i],
+                    river_p.g,
+                    river_p.froude_limit,
+                    dt,
+                ),
+                0.0,
+            )
+
+            # limit floodplain q in case water is not available
+            floodplain_v.q[i] = IfElse.ifelse(
+                floodplain_v.h[i_src] <= 0.0,
+                min(floodplain_v.q[i], 0.0),
+                floodplain_v.q[i],
+            )
+            floodplain_v.q[i] = IfElse.ifelse(
+                floodplain_v.h[i_dst] <= 0.0,
+                max(floodplain_v.q[i], 0.0),
+                floodplain_v.q[i],
+            )
+
+            floodplain_v.q[i] = IfElse.ifelse(
+                floodplain_v.q[i] * river_v.q[i] < 0.0,
+                0.0,
+                floodplain_v.q[i],
+            )
+            floodplain_v.q_av[i] += floodplain_v.q[i] * dt
+        end
+    end
+    # For reservoir and lake locations the local inertial solution is replaced by the
+    # reservoir or lake model. These locations are handled as boundary conditions in the
+    # local inertial model (fixed h).
+    (; reservoir, inflow_waterbody) = model.boundary_conditions
+    inds_reservoir = network.reservoir.river_indices
+    for v in eachindex(inds_reservoir)
+        i = inds_reservoir[v]
+
+        q_in = get_inflow_waterbody(model, edges_at_node.src[i])
+        update!(reservoir, v, q_in + inflow_waterbody[i], dt, dt_forcing)
+        river_v.q[i] = reservoir.variables.outflow[v]
+        river_v.q_av[i] += river_v.q[i] * dt
+    end
+    (; lake, inflow_waterbody) = model.boundary_conditions
+    inds_lake = network.lake.river_indices
+    for v in eachindex(inds_lake)
+        i = inds_lake[v]
+
+        q_in = get_inflow_waterbody(model, edges_at_node.src[i])
+        update!(lake, v, q_in + inflow_waterbody[i], doy, dt, dt_forcing)
+        river_v.q[i] = max(lake.variables.outflow[v], 0.0)
+        river_v.q_av[i] += river_v.q[i] * dt
+    end
+    if update_h
+        @batch per = thread minbatch = 2000 for i in river_p.active_n
+            q_src = sum_at(river_v.q, edges_at_node.src[i])
+            q_dst = sum_at(river_v.q, edges_at_node.dst[i])
+            river_v.volume[i] =
+                river_v.volume[i] + (q_src - q_dst + inwater[i] - abstraction[i]) * dt
+
+            if river_v.volume[i] < 0.0
+                river_v.error[i] = river_v.error[i] + abs(river_v.volume[i])
+                river_v.volume[i] = 0.0 # set volume to zero
+            end
+            river_v.volume[i] = max(river_v.volume[i] + inflow[i] * dt, 0.0) # add external inflow
+
+            if !isnothing(model.floodplain)
+                floodplain_v = model.floodplain.variables
+                floodplain_p = model.floodplain.parameters
+                q_src = sum_at(floodplain_v.q, edges_at_node.src[i])
+                q_dst = sum_at(floodplain_v.q, edges_at_node.dst[i])
+                floodplain_v.volume[i] = floodplain_v.volume[i] + (q_src - q_dst) * dt
+                # TODO check following approach:
+                # if floodplain volume negative, extract from river volume first
+                if floodplain_v.volume[i] < 0.0
+                    floodplain_v.error[i] =
+                        floodplain_v.error[i] + abs(floodplain_v.volume[i])
+                    floodplain_v.volume[i] = 0.0
+                end
+                volume_total = river_v.volume[i] + floodplain_v.volume[i]
+                if volume_total > river_p.bankfull_volume[i]
+                    flood_volume = volume_total - river_p.bankfull_volume[i]
+                    h = flood_depth(
+                        floodplain_p.profile,
+                        flood_volume,
+                        river_p.flow_length[i],
+                        i,
+                    )
+                    river_v.h[i] = river_p.bankfull_depth[i] + h
+                    river_v.volume[i] =
+                        river_v.h[i] * river_p.flow_width[i] * river_p.flow_length[i]
+                    floodplain_v.volume[i] = max(volume_total - river_v.volume[i], 0.0)
+                    floodplain_v.h[i] = floodplain_v.volume[i] > 0.0 ? h : 0.0
+                else
+                    river_v.h[i] =
+                        volume_total / (river_p.flow_length[i] * river_p.flow_width[i])
+                    river_v.volume[i] = volume_total
+                    floodplain_v.h[i] = 0.0
+                    floodplain_v.volume[i] = 0.0
+                end
+                floodplain_v.h_av[i] += floodplain_v.h[i] * dt
+            else
+                river_v.h[i] =
+                    river_v.volume[i] / (river_p.flow_length[i] * river_p.flow_width[i])
+            end
+            river_v.h_av[i] += river_v.h[i] * dt
+        end
+    end
+    return nothing
+end
+
+"""
+Update local inertial river flow model `LocalInertialRiverFlow` for a single timestep `dt`. An adaptive
+timestepping method is used (computing a sub timestep `dt_s`).
+"""
+function update!(
+    model::LocalInertialRiverFlow{T},
+    network,
+    doy,
+    dt;
+    update_h = true,
+) where {T}
+    (; reservoir, lake) = model.boundary_conditions
+
+    set_waterbody_vars!(reservoir)
+    set_waterbody_vars!(lake)
+
+    if !isnothing(model.floodplain)
+        model.floodplain.variables.q_av .= 0.0
+        model.floodplain.variables.h_av .= 0.0
+    end
+    model.variables.q_av .= 0.0
+    model.variables.h_av .= 0.0
+
+    t = T(0.0)
+    while t < dt
+        dt_s = stable_timestep(model)
+        if t + dt_s > dt
+            dt_s = dt - t
+        end
+        local_inertial_river_update!(model, network, dt_s, dt, doy, update_h)
+        t = t + dt_s
+    end
+    model.variables.q_av ./= dt
+    model.variables.h_av ./= dt
+
+    average_waterbody_vars!(reservoir, dt)
+    average_waterbody_vars!(lake, dt)
+
+    if !isnothing(model.floodplain)
+        model.floodplain.variables.q_av ./= dt
+        model.floodplain.variables.h_av ./= dt
+        model.variables.q_channel_av .= model.variables.q_av
+        model.variables.q_av .=
+            model.variables.q_channel_av .+ model.floodplain.variables.q_av
+    end
+
+    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)
+
+"Struct to store local inertial overland flow model variables"
+@get_units @grid_loc @with_kw struct LocalInertialOverlandFlowVariables{T}
+    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
+    volume::Vector{T} | "m3"                # total volume of cell (including river volume for river cells)
+    error::Vector{T} | "m3"                 # error volume
+    h::Vector{T} | "m"                      # water depth of cell (for river cells the reference is the river bed elevation `zb`)
+    h_av::Vector{T} | "m"                   # average water depth (for river cells the reference is the river bed elevation `zb`)
+end
+
+"Initialize local inertial overland flow model variables"
+function LocalInertialOverlandFlowVariables(n)
+    variables = LocalInertialOverlandFlowVariables(;
+        qx0 = zeros(n + 1),
+        qy0 = zeros(n + 1),
+        qx = zeros(n + 1),
+        qy = zeros(n + 1),
+        volume = zeros(n),
+        error = zeros(n),
+        h = zeros(n),
+        h_av = zeros(n),
+    )
+    return variables
+end
+
+"Struct to store local inertial overland flow model parameters"
+@get_units @grid_loc @with_kw struct LocalInertialOverlandFlowParameters{T}
+    n::Int                                              # number of cells [-]
+    x_length::Vector{T} | "m"                           # cell length x direction [m]
+    y_length::Vector{T} | "m"                           # cell length y direction [m]
+    xwidth::Vector{T} | "m" | "edge"                    # effective flow width x direction (floodplain) [m]
+    ywidth::Vector{T} | "m" | "edge"                    # effective flow width y direction (floodplain) [m]
+    g::T                                                # acceleration due to gravity [m s⁻²]
+    theta::T                                            # weighting factor (de Almeida et al., 2012) [-]
+    h_thresh::T                                         # depth threshold for calculating flow [m]
+    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
+    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
+end
+
+"Initialize shallow water overland flow model parameters"
+function LocalInertialOverlandFlowParameters(
+    dataset,
+    config,
+    indices;
+    modelsize_2d,
+    reverse_indices, # maps from the 2D external domain to the 1D internal domain (Int for linear indexing).
+    x_length,
+    y_length,
+    river_width,
+    graph_river,
+    ldd_river,
+    inds_river,
+    river_location,
+    waterbody,
+)
+    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number
+    cfl = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
+    theta = get(config.model, "inertial_flow_theta", 0.8)::Float64 # weighting factor
+    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at edge
+
+    @info "Local inertial approach is used for overlandflow." cfl theta h_thresh froude_limit
+
+    mannings_n = ncread(
+        dataset,
+        config,
+        "lateral.land.mannings_n";
+        sel = indices,
+        defaults = 0.072,
+        type = Float,
+    )
+    elevation_2d = ncread(
+        dataset,
+        config,
+        "lateral.land.elevation";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    elevation = elevation_2d[indices]
+    n = length(indices)
+
+    # initialize edges between cells in x and y direction.
+    staggered_indices =
+        Indices(; xu = zeros(n), xd = zeros(n), yu = zeros(n), yd = zeros(n))
+
+    # edges without neigbors are handled by an extra index (at n + 1, with n edges), which
+    # is set to a value of 0.0 m³ s⁻¹ for qx and qy fields at initialization.
+    # edges are defined as follows for the x and y direction, respectively:
+    # node i => node xu (node i + CartesianIndex(1, 0))
+    # node i => node yu (node i + CartesianIndex(0, 1))
+    # where i is the index of indices
+    nrow, ncol = modelsize_2d
+    for (v, i) in enumerate(indices)
+        for (m, neighbor) in enumerate(neighbors)
+            j = i + neighbor
+            dir = dirs[m]
+            if (1 <= j[1] <= nrow) && (1 <= j[2] <= ncol) && (reverse_indices[j] != 0)
+                getfield(staggered_indices, dir)[v] = reverse_indices[j]
+            else
+                getfield(staggered_indices, dir)[v] = n + 1
+            end
+        end
+    end
+
+    # determine z at edges in x and y direction
+    zx_max = fill(Float(0), n)
+    zy_max = fill(Float(0), n)
+    for i in 1:n
+        xu = staggered_indices.xu[i]
+        if xu <= n
+            zx_max[i] = max(elevation[i], elevation[xu])
+        end
+        yu = staggered_indices.yu[i]
+        if yu <= n
+            zy_max[i] = max(elevation[i], elevation[yu])
+        end
+    end
+
+    # set the effective flow width for river cells in the x and y direction at cell edges.
+    # for waterbody cells (reservoir or lake), h is set to zero (fixed) and not updated, and
+    # overland flow from a downstream cell is not possible (effective flowwidth is zero).
+    we_x = copy(x_length)
+    we_y = copy(y_length)
+    set_effective_flowwidth!(
+        we_x,
+        we_y,
+        staggered_indices,
+        graph_river,
+        river_width,
+        ldd_river,
+        waterbody,
+        reverse_indices[inds_river],
+    )
+    parameters = LocalInertialOverlandFlowParameters(;
+        n,
+        x_length,
+        y_length,
+        xwidth = we_x,
+        ywidth = we_y,
+        g = 9.80665,
+        theta,
+        h_thresh,
+        zx_max,
+        zy_max,
+        mannings_n_sq = mannings_n .* mannings_n,
+        z = elevation,
+        froude_limit,
+        rivercells = river_location,
+    )
+    return parameters, staggered_indices
+end
+
+"Struct to store local inertial overland flow model boundary conditions"
+@get_units @grid_loc @with_kw struct LocalInertialOverlandFlowBC{T}
+    runoff::Vector{T} | "m3 s-1"               # runoff from hydrological model
+    inflow_waterbody::Vector{T} | "m3 s-1"     # inflow to water body from hydrological model
+end
+
+"Struct to store shallow water overland flow model boundary conditions"
+function LocalInertialOverlandFlowBC(n)
+    bc = LocalInertialOverlandFlowBC(; runoff = zeros(n), inflow_waterbody = zeros(n))
+    return bc
+end
+
+"Local inertial overland flow model using the local inertial method"
+@with_kw struct LocalInertialOverlandFlow{T}
+    timestepping::TimeStepping{T}
+    boundary_conditions::LocalInertialOverlandFlowBC{T}
+    parameters::LocalInertialOverlandFlowParameters{T}
+    variables::LocalInertialOverlandFlowVariables{T}
+end
+
+"Initialize local inertial overland flow model"
+function LocalInertialOverlandFlow(
+    dataset,
+    config,
+    indices;
+    modelsize_2d,
+    reverse_indices, # maps from the 2D external domain to the 1D internal domain (Int for linear indexing).
+    x_length,
+    y_length,
+    river_width,
+    graph_river,
+    ldd_river,
+    inds_river,
+    river_location,
+    waterbody,
+)
+    cfl = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
+    timestepping = TimeStepping(; cfl)
+
+    n = length(indices)
+    boundary_conditions = LocalInertialOverlandFlowBC(n)
+    parameters, staggered_indices = LocalInertialOverlandFlowParameters(
+        dataset,
+        config,
+        indices;
+        modelsize_2d,
+        reverse_indices, # maps from the 2D external domain to the 1D internal domain (Int for linear indexing).
+        x_length,
+        y_length,
+        river_width,
+        graph_river,
+        ldd_river,
+        inds_river,
+        river_location,
+        waterbody,
+    )
+    variables = LocalInertialOverlandFlowVariables(n)
+
+    sw_land = LocalInertialOverlandFlow{Float}(;
+        timestepping,
+        boundary_conditions,
+        parameters,
+        variables,
+    )
+
+    return sw_land, staggered_indices
+end
+
+"""
+    stable_timestep(model::LocalInertialRiverFlow)
+    stable_timestep(model::LocalInertialOverlandFlow)
+
+Compute a stable timestep size for the local inertial approach, based on Bates et al. (2010).
+
+dt = cfl * (Δx / sqrt(g max(h))
+"""
+function stable_timestep(model::LocalInertialRiverFlow{T})::T where {T}
+    dt_min = T(Inf)
+    (; cfl) = model.timestepping
+    (; n, flow_length, g) = model.parameters
+    (; h) = model.variables
+    @batch per = thread reduction = ((min, dt_min),) for i in 1:(n)
+        @fastmath @inbounds dt = cfl * flow_length[i] / sqrt(g * h[i])
+        dt_min = min(dt, dt_min)
+    end
+    dt_min = isinf(dt_min) ? T(60.0) : dt_min
+    return dt_min
+end
+
+function stable_timestep(model::LocalInertialOverlandFlow{T})::T where {T}
+    dt_min = T(Inf)
+    (; cfl) = model.timestepping
+    (; n, g, x_length, y_length, rivercells) = model.parameters
+    (; h) = model.variables
+    @batch per = thread reduction = ((min, dt_min),) for i in 1:(n)
+        @fastmath @inbounds dt = if rivercells[i] == 0
+            cfl * min(x_length[i], y_length[i]) / sqrt(g * h[i])
+        else
+            T(Inf)
+        end
+        dt_min = min(dt, dt_min)
+    end
+    dt_min = isinf(dt_min) ? T(60.0) : dt_min
+    return dt_min
+end
+
+"""
+Update boundary conditions `runoff` and inflow to a waterbody from land `inflow_waterbody` for
+overland flow model `LocalInertialOverlandFlow` for a single timestep.
+"""
+function update_boundary_conditions!(
+    model::LocalInertialOverlandFlow,
+    external_models::NamedTuple,
+    network,
+    dt,
+)
+    (; river, soil, subsurface, runoff) = external_models
+    (; inflow_waterbody) = model.boundary_conditions
+    (; reservoir, lake) = river.boundary_conditions
+    (; net_runoff) = soil.variables
+    (; net_runoff_river) = runoff.variables
+
+    model.boundary_conditions.runoff .=
+        net_runoff ./ 1000.0 .* network.land.area ./ dt .+ get_flux_to_river(subsurface) .+
+        net_runoff_river .* network.land.area .* 0.001 ./ dt
+
+    if !isnothing(reservoir) || !isnothing(lake)
+        inflow_land = get_inflow_waterbody(river, model)
+        inflow_subsurface = get_inflow_waterbody(river, subsurface)
+
+        @. inflow_waterbody[network.river_indices] =
+            inflow_land[network.river_indices] + inflow_subsurface[network.river_indices]
+    end
+    return nothing
+end
+
+"""
+Update combined river `LocalInertialRiverFlow` and overland flow `LocalInertialOverlandFlow` models for a
+single timestep `dt`. An adaptive timestepping method is used (computing a sub timestep
+`dt_s`).
+"""
+function update!(
+    land::LocalInertialOverlandFlow{T},
+    river::LocalInertialRiverFlow{T},
+    network,
+    doy,
+    dt;
+    update_h = false,
+) where {T}
+    (; reservoir, lake) = river.boundary_conditions
+
+    if !isnothing(reservoir)
+        reservoir.boundary_conditions.inflow .= 0.0
+        reservoir.variables.totaloutflow .= 0.0
+        reservoir.variables.actevap .= 0.0
+    end
+    if !isnothing(lake)
+        lake.boundary_conditions.inflow .= 0.0
+        lake.variables.totaloutflow .= 0.0
+        lake.variables.actevap .= 0.0
+    end
+    river.variables.q_av .= 0.0
+    river.variables.h_av .= 0.0
+    land.variables.h_av .= 0.0
+
+    t = T(0.0)
+    while t < dt
+        dt_river = stable_timestep(river)
+        dt_land = stable_timestep(land)
+        dt_s = min(dt_river, dt_land)
+        if t + dt_s > dt
+            dt_s = dt - t
+        end
+        local_inertial_river_update!(river, network, dt_s, dt, doy, update_h)
+        local_inertial_update!(land, river, network, dt_s)
+        t = t + dt_s
+    end
+    river.variables.q_av ./= dt
+    river.variables.h_av ./= dt
+    land.variables.h_av ./= dt
+
+    return nothing
+end
+
+"""
+Update combined river `LocalInertialRiverFlow`and overland flow `LocalInertialOverlandFlow` models for a
+single timestep `dt`.
+"""
+function local_inertial_update!(
+    land::LocalInertialOverlandFlow{T},
+    river::LocalInertialRiverFlow{T},
+    network,
+    dt,
+) where {T}
+    indices = network.land.staggered_indices
+    inds_river = network.land.river_indices
+
+    (; edges_at_node) = network.river
+
+    river_bc = river.boundary_conditions
+    river_v = river.variables
+    river_p = river.parameters
+    land_bc = land.boundary_conditions
+    land_v = land.variables
+    land_p = land.parameters
+
+    land_v.qx0 .= land_v.qx
+    land_v.qy0 .= land_v.qy
+
+    # update qx
+    @batch per = thread minbatch = 6000 for i in 1:(land_p.n)
+        yu = indices.yu[i]
+        yd = indices.yd[i]
+        xu = indices.xu[i]
+        xd = indices.xd[i]
+
+        # the effective flow width is zero when the river width exceeds the cell width (dy
+        # for flow in x dir) and floodplain flow is not calculated.
+        if xu <= land_p.n && land_p.ywidth[i] != T(0.0)
+            zs_x = land_p.z[i] + land_v.h[i]
+            zs_xu = land_p.z[xu] + land_v.h[xu]
+            zs_max = max(zs_x, zs_xu)
+            hf = (zs_max - land_p.zx_max[i])
+
+            if hf > land_p.h_thresh
+                length = T(0.5) * (land_p.x_length[i] + land_p.x_length[xu]) # can be precalculated
+                land_v.qx[i] = local_inertial_flow(
+                    land_p.theta,
+                    land_v.qx0[i],
+                    land_v.qx0[xd],
+                    land_v.qx0[xu],
+                    zs_x,
+                    zs_xu,
+                    hf,
+                    land_p.ywidth[i],
+                    length,
+                    land_p.mannings_n_sq[i],
+                    land_p.g,
+                    land_p.froude_limit,
+                    dt,
+                )
+                # limit qx in case water is not available
+                if land_v.h[i] <= T(0.0)
+                    land_v.qx[i] = min(land_v.qx[i], T(0.0))
+                end
+                if land_v.h[xu] <= T(0.0)
+                    land_v.qx[i] = max(land_v.qx[i], T(0.0))
+                end
+            else
+                land_v.qx[i] = T(0.0)
+            end
+        end
+
+        # update qy
+
+        # the effective flow width is zero when the river width exceeds the cell width (dx
+        # for flow in y dir) and floodplain flow is not calculated.
+        if yu <= land_p.n && land_p.xwidth[i] != T(0.0)
+            zs_y = land_p.z[i] + land_v.h[i]
+            zs_yu = land_p.z[yu] + land_v.h[yu]
+            zs_max = max(zs_y, zs_yu)
+            hf = (zs_max - land_p.zy_max[i])
+
+            if hf > land_p.h_thresh
+                length = T(0.5) * (land_p.y_length[i] + land_p.y_length[yu]) # can be precalculated
+                land_v.qy[i] = local_inertial_flow(
+                    land_p.theta,
+                    land_v.qy0[i],
+                    land_v.qy0[yd],
+                    land_v.qy0[yu],
+                    zs_y,
+                    zs_yu,
+                    hf,
+                    land_p.xwidth[i],
+                    length,
+                    land_p.mannings_n_sq[i],
+                    land_p.g,
+                    land_p.froude_limit,
+                    dt,
+                )
+                # limit qy in case water is not available
+                if land_v.h[i] <= T(0.0)
+                    land_v.qy[i] = min(land_v.qy[i], T(0.0))
+                end
+                if land_v.h[yu] <= T(0.0)
+                    land_v.qy[i] = max(land_v.qy[i], T(0.0))
+                end
+            else
+                land_v.qy[i] = T(0.0)
+            end
+        end
+    end
+
+    # change in volume and water levels based on horizontal fluxes for river and land cells
+    @batch per = thread minbatch = 6000 for i in 1:(land_p.n)
+        yd = indices.yd[i]
+        xd = indices.xd[i]
+
+        if land_p.rivercells[i]
+            if river_p.waterbody[inds_river[i]]
+                # for reservoir or lake set inflow from land part, these are boundary points
+                # and update of volume and h is not required
+                river_bc.inflow_waterbody[inds_river[i]] =
+                    land_bc.inflow_waterbody[i] +
+                    land_bc.runoff[i] +
+                    (land_v.qx[xd] - land_v.qx[i] + land_v.qy[yd] - land_v.qy[i])
+            else
+                land_v.volume[i] +=
+                    (
+                        sum_at(river_v.q, edges_at_node.src[inds_river[i]]) -
+                        sum_at(river_v.q, edges_at_node.dst[inds_river[i]]) +
+                        land_v.qx[xd] - land_v.qx[i] + land_v.qy[yd] - land_v.qy[i] +
+                        river_bc.inflow[inds_river[i]] +
+                        land_bc.runoff[i] - river_bc.abstraction[inds_river[i]]
+                    ) * dt
+                if land_v.volume[i] < T(0.0)
+                    land_v.error[i] = land_v.error[i] + abs(land_v.volume[i])
+                    land_v.volume[i] = T(0.0) # set volume to zero
+                end
+                if land_v.volume[i] >= river_p.bankfull_volume[inds_river[i]]
+                    river_v.h[inds_river[i]] =
+                        river_p.bankfull_depth[inds_river[i]] +
+                        (land_v.volume[i] - river_p.bankfull_volume[inds_river[i]]) /
+                        (land_p.x_length[i] * land_p.y_length[i])
+                    land_v.h[i] =
+                        river_v.h[inds_river[i]] - river_p.bankfull_depth[inds_river[i]]
+                    river_v.volume[inds_river[i]] =
+                        river_v.h[inds_river[i]] *
+                        river_p.flow_length[inds_river[i]] *
+                        river_p.flow_width[inds_river[i]]
+                else
+                    river_v.h[inds_river[i]] =
+                        land_v.volume[i] / (
+                            river_p.flow_length[inds_river[i]] *
+                            river_p.flow_width[inds_river[i]]
+                        )
+                    land_v.h[i] = T(0.0)
+                    river_v.volume[inds_river[i]] = land_v.volume[i]
+                end
+                river_v.h_av[inds_river[i]] += river_v.h[inds_river[i]] * dt
+            end
+        else
+            land_v.volume[i] +=
+                (
+                    land_v.qx[xd] - land_v.qx[i] + land_v.qy[yd] - land_v.qy[i] +
+                    land_bc.runoff[i]
+                ) * dt
+            if land_v.volume[i] < T(0.0)
+                land_v.error[i] = land_v.error[i] + abs(land_v.volume[i])
+                land_v.volume[i] = T(0.0) # set volume to zero
+            end
+            land_v.h[i] = land_v.volume[i] / (land_p.x_length[i] * land_p.y_length[i])
+        end
+        land_v.h_av[i] += land_v.h[i] * dt
+    end
+    return nothing
+end
+
+"""
+    FloodPlainProfile
+
+Floodplain `volume` is a function of `depth` (flood depth intervals). Based on the
+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 @grid_loc @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
+
+"Initialize floodplain profile `FloodPlainProfile`"
+function FloodPlainProfile(dataset, config, indices; river_width, river_length, index_pit)
+    volume = ncread(
+        dataset,
+        config,
+        "lateral.river.floodplain.volume";
+        sel = indices,
+        type = Float,
+        dimname = :flood_depth,
+    )
+    n = length(indices)
+
+    # for convenience (interpolation) flood depth 0.0 m is added, with associated area (a),
+    # volume, width (river width) and wetted perimeter (p).
+    volume = vcat(fill(Float(0), n)', volume)
+    start_volume = volume
+    flood_depths = Float.(dataset["flood_depth"][:])
+    pushfirst!(flood_depths, 0.0)
+    n_depths = length(flood_depths)
+
+    p = zeros(Float, n_depths, n)
+    a = zeros(Float, n_depths, n)
+    segment_volume = zeros(Float, n_depths, n)
+    width = zeros(Float, n_depths, n)
+    width[1, :] = river_width[1:n]
+
+    # determine flow area (a), width and wetted perimeter (p) FloodPlain
+    h = diff(flood_depths)
+    incorrect_vol = 0
+    riv_cells = 0
+    error_vol = 0
+    for i in 1:n
+        riv_cell = 0
+        diff_volume = diff(volume[:, i])
+
+        for j in 1:(n_depths - 1)
+            # assume rectangular shape of flood depth segment
+            width[j + 1, i] = diff_volume[j] / (h[j] * river_length[i])
+            # check provided flood volume (floodplain width should be constant or increasing
+            # as a function of flood depth)
+            if width[j + 1, i] < width[j, i]
+                # raise warning only if difference is larger than rounding error of 0.01 m³
+                if ((width[j, i] - width[j + 1, i]) * h[j] * river_length[i]) > 0.01
+                    incorrect_vol += 1
+                    riv_cell = 1
+                    error_vol =
+                        error_vol +
+                        ((width[j, i] - width[j + 1, i]) * h[j] * river_length[i])
+                end
+                width[j + 1, i] = width[j, i]
+            end
+            a[j + 1, i] = width[j + 1, i] * h[j]
+            p[j + 1, i] = (width[j + 1, i] - width[j, i]) + 2.0 * h[j]
+            segment_volume[j + 1, i] = a[j + 1, i] * river_length[i]
+            if j == 1
+                # for interpolation wetted perimeter at flood depth 0.0 is required
+                p[j, i] = p[j + 1, i] - 2.0 * h[j]
+            end
+        end
+
+        p[2:end, i] = cumsum(p[2:end, i])
+        a[:, i] = cumsum(a[:, i])
+        volume[:, i] = cumsum(segment_volume[:, i])
+
+        riv_cells += riv_cell
+    end
+
+    if incorrect_vol > 0
+        perc_riv_cells = round(100.0 * (riv_cells / n); digits = 2)
+        perc_error_vol = round(100.0 * (error_vol / sum(start_volume[end, :])); digits = 2)
+        @warn string(
+            "The provided volume of $incorrect_vol rectangular floodplain schematization",
+            " segments for $riv_cells river cells ($perc_riv_cells % of total river cells)",
+            " is not correct and has been increased with $perc_error_vol % of provided volume.",
+        )
+    end
+
+    # set floodplain parameters for ghost points
+    volume = hcat(volume, volume[:, index_pit])
+    width = hcat(width, width[:, index_pit])
+    a = hcat(a, a[:, index_pit])
+    p = hcat(p, p[:, index_pit])
+
+    # initialize floodplain profile parameters
+    profile =
+        FloodPlainProfile{Float, n_depths}(; volume, width, depth = flood_depths, a, p)
+    return profile
+end
+
+"Struct to store floodplain flow model parameters"
+@get_units @grid_loc @with_kw struct FloodPlainParameters{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" | "edge"    # manning's roughness squared
+    zb_max::Vector{T} | "m" | "edge"                    # maximum bankfull elevation (edge)
+end
+
+"Initialize floodplain flow model parameters"
+function FloodPlainParameters(
+    dataset,
+    config,
+    indices;
+    river_width,
+    river_length,
+    zb_floodplain,
+    nodes_at_edge,
+    n_edges,
+    index_pit,
+)
+    profile =
+        FloodPlainProfile(dataset, config, indices; river_width, river_length, index_pit)
+
+    mannings_n = ncread(
+        dataset,
+        config,
+        "lateral.river.floodplain.mannings_n";
+        sel = indices,
+        defaults = 0.072,
+        type = Float,
+    )
+    # manning roughness at edges
+    append!(mannings_n, mannings_n[index_pit]) # copy to ghost nodes
+    mannings_n_sq = fill(Float(0), n_edges)
+    zb_max = fill(Float(0), n_edges)
+    for i in 1:n_edges
+        src_node = nodes_at_edge.src[i]
+        dst_node = nodes_at_edge.dst[i]
+        mannings_n_i =
+            (
+                mannings_n[dst_node] * river_length[dst_node] +
+                mannings_n[src_node] * river_length[src_node]
+            ) / (river_length[dst_node] + river_length[src_node])
+        mannings_n_sq[i] = mannings_n_i * mannings_n_i
+        zb_max[i] = max(zb_floodplain[src_node], zb_floodplain[dst_node])
+    end
+    parameters = FloodPlainParameters(profile, mannings_n, mannings_n_sq, zb_max)
+    return parameters
+end
+
+"Struct to store floodplain flow model variables"
+@get_units @grid_loc @with_kw struct FloodPlainVariables{T}
+    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
+    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
+end
+
+"Initialize floodplain flow model variables"
+function FloodPlainVariables(n, n_edges, index_pit)
+    variables = FloodPlainVariables(;
+        volume = zeros(n),
+        error = zeros(n),
+        h = zeros(n + length(index_pit)),
+        h_av = zeros(n),
+        a = zeros(n_edges),
+        r = zeros(n_edges),
+        hf = zeros(n_edges),
+        q = zeros(n_edges),
+        q_av = zeros(n_edges),
+        q0 = zeros(n_edges),
+        hf_index = zeros(Int, n_edges),
+    )
+    return variables
+end
+
+"Floodplain flow model"
+@with_kw struct FloodPlain{T, P}
+    parameters::FloodPlainParameters{T, P}
+    variables::FloodPlainVariables{T}
+end
+
+"Determine the initial floodplain volume"
+function initialize_volume!(river, nriv::Int)
+    (; flow_width, flow_length) = river.parameters
+    (; floodplain) = river
+    profile = floodplain.parameters
+    river = for i in 1:nriv
+        i1, i2 = interpolation_indices(floodplain.variables.h[i], profile.depth)
+        a = flow_area(
+            profile.width[i2, i],
+            profile.a[i1, i],
+            profile.depth[i1],
+            floodplain.variables.h[i],
+        )
+        a = max(a - (flow_width[i] * floodplain.h[i]), 0.0)
+        floodplain.variables.volume[i] = flow_length[i] * a
+    end
+    return nothing
+end
+
+"helper function to get interpolation indices"
+function interpolation_indices(x, v::AbstractVector)
+    i1 = 1
+    for i in eachindex(v)
+        if v[i] <= x
+            i1 = i
+        end
+    end
+    if i1 == length(v)
+        i2 = i1
+    else
+        i2 = i1 + 1
+    end
+    return i1, i2
+end
+
+"""
+    flow_area(width, area, depth, h)
+
+Compute floodplain flow area based on flow depth `h` and floodplain `depth`, `area` and
+`width` of a floodplain profile.
+"""
+function flow_area(width, area, depth, h)
+    dh = h - depth  # depth at i1
+    area = area + (width * dh) # area at i1, width at i2
+    return area
+end
+
+"""
+    function wetted_perimeter(p, depth, h)
+
+Compute floodplain wetted perimeter based on flow depth `h` and floodplain `depth` and
+wetted perimeter `p` of a floodplain profile.
+"""
+function wetted_perimeter(p, depth, h)
+    dh = h - depth # depth at i1
+    p = p + (2.0 * dh) # p at i1
+    return p
+end
+
+"Compute flood depth by interpolating flood volume `flood_volume` using flood depth intervals."
+function flood_depth(
+    profile::FloodPlainProfile{T},
+    flood_volume,
+    flow_length,
+    i::Int,
+)::T where {T}
+    i1, i2 = interpolation_indices(flood_volume, @view profile.volume[:, i])
+    ΔA = (flood_volume - profile.volume[i1, i]) / flow_length
+    dh = ΔA / profile.width[i2, i]
+    flood_depth = profile.depth[i1] + dh
+    return flood_depth
+end
+
+"Initialize floodplain geometry and `FloodPlain` variables and parameters"
+function FloodPlain(
+    dataset,
+    config,
+    indices;
+    river_width,
+    river_length,
+    zb_floodplain,
+    index_pit,
+    n_edges,
+    nodes_at_edge,
+)
+    n = length(indices)
+    parameters = FloodPlainParameters(
+        dataset,
+        config,
+        indices;
+        river_width,
+        river_length,
+        zb_floodplain,
+        nodes_at_edge,
+        n_edges,
+        index_pit,
+    )
+    variables = FloodPlainVariables(n, n_edges, index_pit)
+
+    floodplain = FloodPlain(; parameters, variables)
+    return floodplain
+end
\ No newline at end of file
diff --git a/src/routing/surface_routing.jl b/src/routing/surface_routing.jl
new file mode 100644
index 000000000..8f0547c72
--- /dev/null
+++ b/src/routing/surface_routing.jl
@@ -0,0 +1,66 @@
+"""
+    surface_routing!(model)
+
+Run surface routing (land and river) for a single timestep. Kinematic wave for overland flow
+and kinematic wave or local inertial model for river flow.
+"""
+function surface_routing!(model)
+    (; vertical, lateral, network, config, clock) = model
+    (; soil, runoff, allocation) = vertical
+    (; land, river, subsurface) = lateral
+
+    dt = tosecond(clock.dt)
+    # update lateral inflow for kinematic wave overland flow
+    update_lateral_inflow!(
+        land,
+        (; soil, allocation, subsurface),
+        network.land.area,
+        config,
+        dt,
+    )
+    # run kinematic wave overland flow
+    update!(land, network.land, dt)
+
+    # update lateral inflow river flow
+    update_lateral_inflow!(
+        river,
+        (; allocation = river.allocation, runoff, land, subsurface),
+        network.river.cell_area,
+        network.land.area,
+        network.river.land_indices,
+        dt,
+    )
+    update_inflow_waterbody!(river, (; land, subsurface), network.river.land_indices)
+    update!(river, network, julian_day(clock.time - clock.dt), dt)
+    return nothing
+end
+
+"""
+    surface_routing!(
+        model::Model{N,L,V,R,W,T}
+    ) where {N,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, V, R, W, T},
+) where {
+    N,
+    L <: NamedTuple{<:Any, <:Tuple{Any, LocalInertialOverlandFlow, LocalInertialRiverFlow}},
+    V,
+    R,
+    W,
+    T,
+}
+    (; lateral, vertical, network, clock) = model
+    (; land, river, subsurface) = lateral
+    (; soil, runoff) = vertical
+
+    dt = tosecond(clock.dt)
+    update_boundary_conditions!(land, (; river, subsurface, soil, runoff), network, dt)
+
+    update!(land, river, network, julian_day(clock.time - clock.dt), dt)
+
+    return nothing
+end
\ No newline at end of file
diff --git a/src/routing/timestepping.jl b/src/routing/timestepping.jl
new file mode 100644
index 000000000..6e9b59c22
--- /dev/null
+++ b/src/routing/timestepping.jl
@@ -0,0 +1,16 @@
+
+"Timestepping for solving kinematic wave and local inertial river and overland flow routing."
+@with_kw struct TimeStepping{T}
+    stable_timesteps::Vector{T} = Float[]
+    dt_fixed::T = 0.0
+    adaptive::Bool = true
+    cfl::T = 0.70
+end
+
+"Check timestep size"
+function check_timestepsize(timestepsize, currenttime, endtime)
+    if currenttime + timestepsize > endtime
+        timestepsize = endtime - currenttime
+    end
+    return timestepsize
+end
\ No newline at end of file
diff --git a/src/sbm.jl b/src/sbm.jl
index 7293ec082..0e9f286c7 100644
--- a/src/sbm.jl
+++ b/src/sbm.jl
@@ -9,33 +9,32 @@
     soil::SbmSoilModel
     demand::D
     allocation::A
-    dt::T
 end
 
 "Initialize land hydrology model with SBM soil model"
-function LandHydrologySBM(nc, config, riverfrac, inds)
+function LandHydrologySBM(dataset, config, riverfrac, indices)
     dt = Second(config.timestepsecs)
-    n = length(inds)
+    n = length(indices)
 
     atmospheric_forcing = AtmosphericForcing(n)
-    vegetation_parameter_set = VegetationParameters(nc, config, inds)
+    vegetation_parameter_set = VegetationParameters(dataset, config, indices)
     if dt >= Hour(23)
         interception_model =
-            GashInterceptionModel(nc, config, inds, vegetation_parameter_set)
+            GashInterceptionModel(dataset, config, indices, vegetation_parameter_set)
     else
         interception_model = RutterInterceptionModel(vegetation_parameter_set, n)
     end
 
     modelsnow = get(config.model, "snow", false)::Bool
     if modelsnow
-        snow_model = SnowHbvModel(nc, config, inds, dt)
+        snow_model = SnowHbvModel(dataset, config, indices, dt)
     else
         snow_model = NoSnowModel{Float}()
     end
     modelglacier = get(config.model, "glacier", false)::Bool
     if modelsnow && modelglacier
         glacier_bc = SnowStateBC{Float}(; snow_storage = snow_model.variables.snow_storage)
-        glacier_model = GlacierHbvModel(nc, config, inds, dt, glacier_bc)
+        glacier_model = GlacierHbvModel(dataset, config, indices, dt, glacier_bc)
     elseif modelsnow == false && modelglacier == true
         @warn string(
             "Glacier processes can be modelled when snow modelling is enabled. To include ",
@@ -45,9 +44,9 @@ function LandHydrologySBM(nc, config, riverfrac, inds)
     else
         glacier_model = NoGlacierModel{Float}()
     end
-    runoff_model = OpenWaterRunoff(nc, config, inds, riverfrac)
+    runoff_model = OpenWaterRunoff(dataset, config, indices, riverfrac)
 
-    soil_model = SbmSoilModel(nc, config, vegetation_parameter_set, inds, dt)
+    soil_model = SbmSoilModel(dataset, config, vegetation_parameter_set, indices, dt)
     @. vegetation_parameter_set.rootingdepth = min(
         soil_model.parameters.soilthickness * 0.99,
         vegetation_parameter_set.rootingdepth,
@@ -55,8 +54,9 @@ function LandHydrologySBM(nc, config, riverfrac, inds)
 
     do_water_demand = haskey(config.model, "water_demand")
     allocation =
-        do_water_demand ? AllocationLand(nc, config, inds) : NoAllocationLand{Float}()
-    demand = do_water_demand ? Demand(nc, config, inds, dt) : NoDemand{Float}()
+        do_water_demand ? AllocationLand(dataset, config, indices) :
+        NoAllocationLand{Float}()
+    demand = do_water_demand ? Demand(dataset, config, indices, dt) : NoDemand{Float}()
 
     args = (demand, allocation)
     land_hydrology_model = LandHydrologySBM{Float, typeof.(args)...}(;
@@ -69,25 +69,15 @@ function LandHydrologySBM(nc, config, riverfrac, inds)
         soil = soil_model,
         demand = demand,
         allocation = allocation,
-        dt = tosecond(dt),
     )
     return land_hydrology_model
 end
 
 "Update land hydrology model with SBM soil model for a single timestep"
-function update!(model::LandHydrologySBM, lateral, network, config)
+function update!(model::LandHydrologySBM, lateral, network, config, dt)
     do_water_demand = haskey(config.model, "water_demand")::Bool
-    (;
-        glacier,
-        snow,
-        interception,
-        runoff,
-        soil,
-        demand,
-        allocation,
-        atmospheric_forcing,
-        dt,
-    ) = model
+    (; glacier, snow, interception, runoff, soil, demand, allocation, atmospheric_forcing) =
+        model
 
     update!(interception, atmospheric_forcing)
 
@@ -162,8 +152,11 @@ function update_total_water_storage!(
     # Burn the river routing values
     for (i, index_river) in enumerate(river_network)
         total_storage[index_river] = (
-            (river_routing.h_av[i] * river_routing.width[i] * river_routing.dl[i]) /
-            (area[index_river]) * 1000 # Convert to mm
+            (
+                river_routing.variables.h_av[i] *
+                river_routing.parameters.flow_width[i] *
+                river_routing.parameters.flow_length[i]
+            ) / (area[index_river]) * 1000 # Convert to mm
         )
     end
 
@@ -178,7 +171,7 @@ function update_total_water_storage!(
     threaded_foreach(1:n; basesize = 1000) do i
         sub_surface = ustoredepth[i] + satwaterdepth[i]
         lateral = (
-            land_routing.h_av[i] * (1 - riverfrac[i]) * 1000 # convert to mm
+            land_routing.variables.h_av[i] * (1 - riverfrac[i]) * 1000 # convert to mm
         )
 
         # Add everything to the total water storage
diff --git a/src/sbm_gwf_model.jl b/src/sbm_gwf_model.jl
index c8b52595d..31d037e5f 100644
--- a/src/sbm_gwf_model.jl
+++ b/src/sbm_gwf_model.jl
@@ -28,11 +28,7 @@ function initialize_sbm_gwf_model(config::Config)
     do_drains = get(config.model, "drains", false)::Bool
     do_constanthead = get(config.model, "constanthead", false)::Bool
 
-    kw_river_tstep = get(config.model, "kw_river_tstep", 0)
-    kw_land_tstep = get(config.model, "kw_land_tstep", 0)
-    kinwave_it = get(config.model, "kin_wave_iteration", false)::Bool
     routing_options = ("kinematic-wave", "local-inertial")
-    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
     river_routing = get_options(
         config.model,
         "river_routing",
@@ -43,162 +39,181 @@ function initialize_sbm_gwf_model(config::Config)
         get_options(config.model, "land_routing", routing_options, "kinematic-wave")::String
     do_water_demand = haskey(config.model, "water_demand")
 
-    nc = NCDataset(static_path)
+    dataset = NCDataset(static_path)
 
-    subcatch_2d = ncread(nc, config, "subcatchment"; optional = false, allow_missing = true)
+    subcatch_2d =
+        ncread(dataset, config, "subcatchment"; optional = false, allow_missing = true)
     # indices based on catchment
-    inds, rev_inds = active_indices(subcatch_2d, missing)
-    n = length(inds)
+    indices, reverse_indices = active_indices(subcatch_2d, missing)
+    n_land_cells = length(indices)
     modelsize_2d = size(subcatch_2d)
 
-    river_2d =
-        ncread(nc, config, "river_location"; optional = false, type = Bool, fill = false)
-    river = river_2d[inds]
-    riverwidth_2d =
-        ncread(nc, config, "lateral.river.width"; optional = false, type = Float, fill = 0)
-    riverwidth = riverwidth_2d[inds]
-    riverlength_2d =
-        ncread(nc, config, "lateral.river.length"; optional = false, type = Float, fill = 0)
-    riverlength = riverlength_2d[inds]
-
-    altitude = ncread(nc, config, "altitude"; optional = false, sel = inds, type = Float)
+    river_location_2d = ncread(
+        dataset,
+        config,
+        "river_location";
+        optional = false,
+        type = Bool,
+        fill = false,
+    )
+    river_location = river_location_2d[indices]
+    river_width_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.width";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    river_width = river_width_2d[indices]
+    river_length_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.length";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    river_length = river_length_2d[indices]
+
+    altitude =
+        ncread(dataset, config, "altitude"; optional = false, sel = indices, type = Float)
+
     # read x, y coordinates and calculate cell length [m]
-    y_nc = read_y_axis(nc)
-    x_nc = read_x_axis(nc)
-    y = permutedims(repeat(y_nc; outer = (1, length(x_nc))))[inds]
-    cellength = abs(mean(diff(x_nc)))
+    y_coords = read_y_axis(dataset)
+    x_coords = read_x_axis(dataset)
+    y = permutedims(repeat(y_coords; outer = (1, length(x_coords))))[indices]
+    cell_length = abs(mean(diff(x_coords)))
 
-    sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
-    xl, yl = cell_lengths(y, cellength, sizeinmetres)
-    riverfrac = river_fraction(river, riverlength, riverwidth, xl, yl)
+    size_in_metres = get(config.model, "sizeinmetres", false)::Bool
+    x_length, y_length = cell_lengths(y, cell_length, size_in_metres)
+    river_fraction =
+        get_river_fraction(river_location, river_length, river_width, x_length, y_length)
 
-    inds_riv, rev_inds_riv = active_indices(river_2d, 0)
-    nriv = length(inds_riv)
+    inds_river, reverse_inds_river = active_indices(river_location_2d, 0)
+    n_river_cells = length(inds_river)
 
     # initialize vertical SBM concept
-    lhm = LandHydrologySBM(nc, config, riverfrac, inds)
+    land_hydrology = LandHydrologySBM(dataset, config, river_fraction, indices)
 
     # reservoirs
     pits = zeros(Bool, modelsize_2d)
     if do_reservoirs
-        reservoirs, resindex, reservoir, pits =
-            initialize_simple_reservoir(config, nc, inds_riv, nriv, pits, tosecond(dt))
+        reservoir, reservoir_network, inds_reservoir_map2river, pits =
+            SimpleReservoir(dataset, config, inds_river, n_river_cells, pits)
     else
-        reservoir = ()
-        reservoirs = nothing
-        resindex = fill(0, nriv)
+        reservoir_network = (river_indices = [],)
+        inds_reservoir_map2river = fill(0, n_river_cells)
+        reservoir = nothing
     end
 
     # lakes
     if do_lakes
-        lakes, lakeindex, lake, pits =
-            initialize_lake(config, nc, inds_riv, nriv, pits, tosecond(dt))
+        lake, lake_network, inds_lake_map2river, pits =
+            Lake(dataset, config, inds_river, n_river_cells, pits)
     else
-        lake = ()
-        lakes = nothing
-        lakeindex = fill(0, nriv)
+        lake_network = (river_indices = [],)
+        inds_lake_map2river = fill(0, n_river_cells)
+        lake = nothing
     end
 
     # overland flow (kinematic wave)
-    landslope =
-        ncread(nc, config, "lateral.land.slope"; optional = false, sel = inds, type = Float)
-    clamp!(landslope, 0.00001, Inf)
-    ldd_2d = ncread(nc, config, "ldd"; optional = false, allow_missing = true)
+    land_slope = ncread(
+        dataset,
+        config,
+        "lateral.land.slope";
+        optional = false,
+        sel = indices,
+        type = Float,
+    )
+    clamp!(land_slope, 0.00001, Inf)
+    ldd_2d = ncread(dataset, config, "ldd"; optional = false, allow_missing = true)
 
-    ldd = ldd_2d[inds]
+    ldd = ldd_2d[indices]
 
-    dl = map(detdrainlength, ldd, xl, yl)
-    dw = (xl .* yl) ./ dl
-    sw = map(det_surfacewidth, dw, riverwidth, river)
+    flow_length = map(get_flow_length, ldd, x_length, y_length)
+    flow_width = (x_length .* y_length) ./ flow_length
+    surface_flow_width = map(det_surfacewidth, flow_width, river_width, river_location)
 
-    graph = flowgraph(ldd, inds, pcr_dir)
-    ldd_riv = ldd_2d[inds_riv]
-    graph_riv = flowgraph(ldd_riv, inds_riv, pcr_dir)
+    graph = flowgraph(ldd, indices, pcr_dir)
+    ldd_river = ldd_2d[inds_river]
+    graph_river = flowgraph(ldd_river, inds_river, pcr_dir)
 
-    # the indices of the river cells in the land(+river) cell vector
-    index_river = filter(i -> !isequal(river[i], 0), 1:n)
-    frac_toriver = fraction_runoff_toriver(graph, ldd, index_river, landslope, n)
+    # land indices where river is located
+    inds_land_map2river = filter(i -> !isequal(river_location[i], 0), 1:n_land_cells)
+    frac_to_river = fraction_runoff_to_river(graph, ldd, inds_land_map2river, land_slope)
 
-    inds_allocation_areas = Vector{Int}[]
-    inds_riv_allocation_areas = Vector{Int}[]
+    allocation_area_inds = Vector{Int}[]
+    river_allocation_area_inds = Vector{Int}[]
     if do_water_demand
-        areas = unique(lhm.allocation.parameters.areas)
+        areas = unique(land_hydrology.allocation.parameters.areas)
         for a in areas
-            area_index = findall(==(a), lhm.allocation.parameters.areas)
-            push!(inds_allocation_areas, area_index)
-            area_riv_index = findall(==(a), lhm.allocation.parameters.areas[index_river])
-            push!(inds_riv_allocation_areas, area_riv_index)
+            area_index = findall(x -> x == a, land_hydrology.allocation.parameters.areas)
+            push!(allocation_area_inds, area_index)
+            area_riv_index = findall(
+                x -> x == a,
+                land_hydrology.allocation.parameters.areas[inds_land_map2river],
+            )
+            push!(river_allocation_area_inds, area_riv_index)
         end
     end
 
     if land_routing == "kinematic-wave"
-        olf = initialize_surfaceflow_land(
-            nc,
+        overland_flow = KinWaveOverlandFlow(
+            dataset,
             config,
-            inds;
-            sl = landslope,
-            dl,
-            width = map(det_surfacewidth, dw, riverwidth, river),
-            iterate = kinwave_it,
-            tstep = kw_land_tstep,
-            dt,
+            indices;
+            slope = land_slope,
+            flow_length,
+            flow_width = surface_flow_width,
         )
     elseif land_routing == "local-inertial"
-        index_river_nf = rev_inds_riv[inds] # not filtered (with zeros)
-        olf, indices = initialize_shallowwater_land(
-            nc,
+        inds_river_map2land = reverse_inds_river[indices] # not filtered (with zeros)
+        overland_flow, staggered_indices = LocalInertialOverlandFlow(
+            dataset,
             config,
-            inds;
+            indices;
             modelsize_2d,
-            indices_reverse = rev_inds,
-            xlength = xl,
-            ylength = yl,
-            riverwidth = riverwidth_2d[inds_riv],
-            graph_riv,
-            ldd_riv,
-            inds_riv,
-            river,
-            waterbody = !=(0).(resindex + lakeindex),
-            dt,
+            reverse_indices,
+            x_length,
+            y_length,
+            river_width = river_width_2d[inds_river],
+            graph_river,
+            ldd_river,
+            inds_river,
+            river_location,
+            waterbody = !=(0).(inds_reservoir_map2river + inds_lake_map2river),
         )
     end
 
     # river flow (kinematic wave)
-    riverlength = riverlength_2d[inds_riv]
-    riverwidth = riverwidth_2d[inds_riv]
-    minimum(riverlength) > 0 || error("river length must be positive on river cells")
-    minimum(riverwidth) > 0 || error("river width must be positive on river cells")
+    river_length = river_length_2d[inds_river]
+    river_width = river_width_2d[inds_river]
+    minimum(river_length) > 0 || error("river length must be positive on river cells")
+    minimum(river_width) > 0 || error("river width must be positive on river cells")
 
     if river_routing == "kinematic-wave"
-        rf = initialize_surfaceflow_river(
-            nc,
+        river_flow = KinWaveRiverFlow(
+            dataset,
             config,
-            inds_riv;
-            dl = riverlength,
-            width = riverwidth,
-            reservoir_index = resindex,
-            reservoir = reservoirs,
-            lake_index = lakeindex,
-            lake = lakes,
-            iterate = kinwave_it,
-            tstep = kw_river_tstep,
-            dt = dt,
+            inds_river;
+            river_length,
+            river_width,
+            reservoir = reservoir,
+            lake = lake,
         )
     elseif river_routing == "local-inertial"
-        rf, nodes_at_link = initialize_shallowwater_river(
-            nc,
+        river_flow, nodes_at_edge = LocalInertialRiverFlow(
+            dataset,
             config,
-            inds_riv;
-            graph = graph_riv,
-            ldd = ldd_riv,
-            dl = riverlength,
-            width = riverwidth,
-            reservoir_index = resindex,
-            reservoir = reservoirs,
-            lake_index = lakeindex,
-            lake = lakes,
-            dt = dt,
-            floodplain = floodplain_1d,
+            inds_river;
+            graph_river,
+            ldd_river,
+            river_length,
+            river_width,
+            reservoir,
+            lake,
+            waterbody = !=(0).(inds_reservoir_map2river + inds_lake_map2river),
         )
     else
         error(
@@ -210,141 +225,76 @@ function initialize_sbm_gwf_model(config::Config)
 
     # unconfined aquifer
     if do_constanthead
-        constanthead = ncread(
-            nc,
-            config,
-            "lateral.subsurface.constant_head";
-            sel = inds,
-            type = Float,
-            fill = mv,
-        )
-        index_constanthead = filter(i -> !isequal(constanthead[i], mv), 1:n)
-        constant_head = ConstantHead(constanthead[index_constanthead], index_constanthead)
+        constant_head = ConstantHead(dataset, config, indices)
     else
-        constant_head = ConstantHead{Float}(Float[], Int64[])
+        variables = ConstantHeadVariables{Float}(; head = Float[])
+        constant_head = ConstantHead{Float}(; variables, index = Int64[])
     end
 
-    conductivity =
-        ncread(nc, config, "lateral.subsurface.conductivity"; sel = inds, type = Float)
-    specific_yield =
-        ncread(nc, config, "lateral.subsurface.specific_yield"; sel = inds, type = Float)
-    gwf_f = ncread(
-        nc,
-        config,
-        "lateral.subsurface.gwf_f";
-        sel = inds,
-        type = Float,
-        defaults = 3.0,
-    )
-    conductivity_profile =
-        get(config.input.lateral.subsurface, "conductivity_profile", "uniform")::String
-
-    connectivity = Connectivity(inds, rev_inds, xl, yl)
-    initial_head = altitude .- lhm.soil.variables.zi / 1000.0 # cold state for groundwater head based on SBM zi
-    initial_head[index_river] = altitude[index_river]
+    connectivity = Connectivity(indices, reverse_indices, x_length, y_length)
 
+    initial_head = altitude .- land_hydrology.soil.variables.zi / 1000.0 # cold state for groundwater head based on SBM zi
+    initial_head[inds_land_map2river] = altitude[inds_land_map2river]
     if do_constanthead
-        initial_head[constant_head.index] = constant_head.head
+        initial_head[constant_head.index] = constant_head.variables.head
     end
 
-    bottom = altitude .- lhm.soil.parameters.soilthickness ./ Float(1000.0)
-    area = xl .* yl
-    volume = @. (min(altitude, initial_head) - bottom) * area * specific_yield # total volume than can be released
-
+    bottom = altitude .- land_hydrology.soil.parameters.soilthickness ./ Float(1000.0)
+    area = x_length .* y_length
+    conductance = zeros(Float, connectivity.nconnection)
     aquifer = UnconfinedAquifer(
-        initial_head,
-        conductivity,
+        dataset,
+        config,
+        indices,
         altitude,
         bottom,
         area,
-        specific_yield,
-        zeros(Float, connectivity.nconnection),  # conductance
-        volume,
-        gwf_f,
+        conductance,
+        initial_head,
     )
 
     # river boundary of unconfined aquifer
-    infiltration_conductance = ncread(
-        nc,
-        config,
-        "lateral.subsurface.infiltration_conductance";
-        sel = inds_riv,
-        type = Float,
-    )
-    exfiltration_conductance = ncread(
-        nc,
-        config,
-        "lateral.subsurface.exfiltration_conductance";
-        sel = inds_riv,
-        type = Float,
-    )
-    river_bottom =
-        ncread(nc, config, "lateral.subsurface.river_bottom"; sel = inds_riv, type = Float)
-
-    river_flux = fill(mv, nriv)
-    river_stage = fill(mv, nriv)
-    river = River(
-        river_stage,
-        infiltration_conductance,
-        exfiltration_conductance,
-        river_bottom,
-        river_flux,
-        index_river,
-    )
+    river = River(dataset, config, inds_river, inds_land_map2river)
 
     # recharge boundary of unconfined aquifer
-    r = fill(mv, n)
-    recharge = Recharge(r, zeros(Float, n), collect(1:n))
+    recharge = Recharge(
+        fill(mv, n_land_cells),
+        zeros(Float, n_land_cells),
+        collect(1:n_land_cells),
+    )
 
     # drain boundary of unconfined aquifer (optional)
     if do_drains
-        drain_2d = ncread(nc, config, "lateral.subsurface.drain"; type = Bool, fill = false)
-
-        drain = drain_2d[inds]
-        # check if drain occurs where overland flow is not possible (sw = 0.0)
-        # and correct if this is the case
-        false_drain = filter(i -> !isequal(drain[i], 0) && sw[i] == Float(0), 1:n)
+        drain_2d =
+            ncread(dataset, config, "lateral.subsurface.drain"; type = Bool, fill = false)
+
+        drain = drain_2d[indices]
+        # check if drain occurs where overland flow is not possible (surface_flow_width =
+        # 0.0) and correct if this is the case
+        false_drain = filter(
+            i -> !isequal(drain[i], 0) && surface_flow_width[i] == Float(0),
+            1:n_land_cells,
+        )
         n_false_drain = length(false_drain)
         if n_false_drain > 0
-            drain_2d[inds[false_drain]] .= 0
+            drain_2d[indices[false_drain]] .= 0
             drain[false_drain] .= 0
             @info "$n_false_drain drain locations are removed that occur where overland flow
              is not possible (overland flow width is zero)"
         end
-        inds_drain, rev_inds_drain = active_indices(drain_2d, 0)
 
-        drain_elevation = ncread(
-            nc,
-            config,
-            "lateral.subsurface.drain_elevation";
-            sel = inds,
-            type = Float,
-            fill = mv,
-        )
-        drain_conductance = ncread(
-            nc,
-            config,
-            "lateral.subsurface.drain_conductance";
-            sel = inds,
-            type = Float,
-            fill = mv,
-        )
-        index_drain = filter(i -> !isequal(drain[i], 0), 1:n)
-        drain_flux = fill(mv, length(index_drain))
-        drains = Drainage(
-            drain_elevation[index_drain],
-            drain_conductance[index_drain],
-            drain_flux,
-            index_drain,
-        )
-        drain = (indices = inds_drain, reverse_indices = rev_inds_drain)
+        indices_drain, reverse_inds_drain = active_indices(drain_2d, 0)
+        inds_land_map2drain = filter(i -> !isequal(drain[i], 0), 1:n_land_cells)
+
+        drains = Drainage(dataset, config, indices, inds_land_map2drain)
+        drain = (indices = indices_drain, reverse_indices = reverse_inds_drain)
         aquifer_boundaries = AquiferBoundaryCondition[recharge, river, drains]
     else
         aquifer_boundaries = AquiferBoundaryCondition[recharge, river]
         drain = ()
     end
 
-    gwf = GroundwaterFlow{Float}(;
+    groundwater_flow = GroundwaterFlow{Float}(;
         aquifer,
         connectivity,
         constanthead = constant_head,
@@ -354,14 +304,17 @@ function initialize_sbm_gwf_model(config::Config)
     # map GroundwaterFlow and its boundaries
     if do_drains
         subsurface_map = (
-            flow = gwf,
-            recharge = gwf.boundaries[1],
-            river = gwf.boundaries[2],
-            drain = gwf.boundaries[3],
+            flow = groundwater_flow,
+            recharge = groundwater_flow.boundaries[1],
+            river = groundwater_flow.boundaries[2],
+            drain = groundwater_flow.boundaries[3],
         )
     else
-        subsurface_map =
-            (flow = gwf, recharge = gwf.boundaries[1], river = gwf.boundaries[2])
+        subsurface_map = (
+            flow = groundwater_flow,
+            recharge = groundwater_flow.boundaries[1],
+            river = groundwater_flow.boundaries[2],
+        )
     end
 
     # setup subdomains for the land and river kinematic wave domain, if nthreads = 1
@@ -372,48 +325,54 @@ function initialize_sbm_gwf_model(config::Config)
     end
     if land_routing == "kinematic-wave"
         toposort = topological_sort_by_dfs(graph)
-        index_pit_land = findall(x -> x == 5, ldd)
+        land_pit_inds = findall(x -> x == 5, ldd)
         min_streamorder_land = get(config.model, "min_streamorder_land", 5)
-        subbas_order, indices_subbas, topo_subbas = kinwave_set_subdomains(
+        order_of_subdomains, subdomain_inds, toposort_subdomain = kinwave_set_subdomains(
             graph,
             toposort,
-            index_pit_land,
+            land_pit_inds,
             streamorder,
             min_streamorder_land,
         )
     end
     if river_routing == "kinematic-wave"
         min_streamorder_river = get(config.model, "min_streamorder_river", 6)
-        toposort_riv = topological_sort_by_dfs(graph_riv)
-        index_pit_river = findall(x -> x == 5, ldd_riv)
-        subriv_order, indices_subriv, topo_subriv = kinwave_set_subdomains(
-            graph_riv,
-            toposort_riv,
-            index_pit_river,
-            streamorder[index_river],
-            min_streamorder_river,
-        )
+        toposort_river = topological_sort_by_dfs(graph_river)
+        river_pit_inds = findall(x -> x == 5, ldd_river)
+        order_of_river_subdomains, river_subdomain_inds, toposort_river_subdomain =
+            kinwave_set_subdomains(
+                graph_river,
+                toposort_river,
+                river_pit_inds,
+                streamorder[inds_land_map2river],
+                min_streamorder_river,
+            )
     end
 
-    modelmap =
-        (vertical = lhm, lateral = (subsurface = subsurface_map, land = olf, river = rf))
+    modelmap = (
+        vertical = land_hydrology,
+        lateral = (subsurface = subsurface_map, land = overland_flow, river = river_flow),
+    )
     indices_reverse = (
-        land = rev_inds,
-        river = rev_inds_riv,
-        reservoir = isempty(reservoir) ? nothing : reservoir.reverse_indices,
-        lake = isempty(lake) ? nothing : lake.reverse_indices,
-        drain = isempty(drain) ? nothing : rev_inds_drain,
+        land = reverse_indices,
+        river = reverse_inds_river,
+        reservoir = isnothing(reservoir) ? nothing : reservoir.reverse_indices,
+        lake = isnothing(lake) ? nothing : lake.reverse_indices,
+        drain = isempty(drain) ? nothing : reverse_inds_drain,
     )
     writer = prepare_writer(
         config,
         modelmap,
         indices_reverse,
-        x_nc,
-        y_nc,
-        nc;
-        extra_dim = (name = "layer", value = Float64.(1:(lhm.soil.parameters.maxlayers))),
+        x_coords,
+        y_coords,
+        dataset;
+        extra_dim = (
+            name = "layer",
+            value = Float64.(1:(land_hydrology.soil.parameters.maxlayers)),
+        ),
     )
-    close(nc)
+    close(dataset)
 
     # for each domain save:
     # - the directed acyclic graph (graph),
@@ -430,93 +389,98 @@ function initialize_sbm_gwf_model(config::Config)
     # functions
     if land_routing == "kinematic-wave"
         land = (
-            graph = graph,
-            upstream_nodes = filter_upsteam_nodes(graph, pits[inds]),
-            subdomain_order = subbas_order,
-            topo_subdomain = topo_subbas,
-            indices_subdomain = indices_subbas,
+            graph,
+            upstream_nodes = filter_upsteam_nodes(graph, pits[indices]),
+            order_of_subdomains,
+            order_subdomain = toposort_subdomain,
+            subdomain_indices = subdomain_inds,
             order = toposort,
-            indices = inds,
-            reverse_indices = rev_inds,
-            area = xl .* yl,
-            slope = landslope,
-            altitude = altitude,
-            indices_allocation_areas = inds_allocation_areas,
+            indices,
+            reverse_indices,
+            area = x_length .* y_length,
+            slope = land_slope,
+            frac_to_river,
+            altitude,
+            allocation_area_indices = allocation_area_inds,
         )
     elseif land_routing == "local-inertial"
         land = (
-            graph = graph,
+            graph,
             order = toposort,
-            indices = inds,
-            reverse_indices = rev_inds,
-            area = xl .* yl,
-            slope = landslope,
-            altitude = altitude,
-            index_river = index_river_nf,
-            staggered_indices = indices,
-            indices_allocation_areas = inds_allocation_areas,
+            indices,
+            reverse_indices,
+            area = x_length .* y_length,
+            slope = land_slope,
+            frac_to_river,
+            altitude,
+            river_indices = inds_river_map2land,
+            staggered_indices,
+            allocation_area_indices = allocation_area_inds,
         )
     end
     if do_water_demand
         # exclude waterbodies for local surface and ground water abstraction
-        inds_riv_2d = copy(rev_inds_riv)
-        inds_2d = ones(Bool, modelsize_2d)
-        if !isempty(reservoir)
-            inds_cov = collect(Iterators.flatten(reservoir.indices_coverage))
+        inds_riv_2d = copy(reverse_inds_river)
+        inds_2d = zeros(Bool, modelsize_2d)
+        if !isnothing(reservoir)
+            inds_cov = collect(Iterators.flatten(reservoir_network.indices_coverage))
             inds_riv_2d[inds_cov] .= 0
-            inds_2d[inds_cov] .= 0
+            inds_2d[inds_cov] .= 1
         end
-        if !isempty(lake)
-            inds_cov = collect(Iterators.flatten(lake.indices_coverage))
+        if !isnothing(lake)
+            inds_cov = collect(Iterators.flatten(lake_network.indices_coverage))
             inds_riv_2d[inds_cov] .= 0
-            inds_2d[inds_cov] .= 0
+            inds_2d[inds_cov] .= 1
         end
-        land = merge(land, (index_river_wb = inds_riv_2d[inds], index_wb = inds_2d[inds]))
+        land = merge(
+            land,
+            (
+                river_inds_excl_waterbody = inds_riv_2d[indices],
+                waterbody = inds_2d[indices],
+            ),
+        )
     end
     if river_routing == "kinematic-wave"
         river = (
-            graph = graph_riv,
-            indices = inds_riv,
-            reverse_indices = rev_inds_riv,
+            graph = graph_river,
+            indices = inds_river,
+            reverse_indices = reverse_inds_river,
             # reservoir and lake index
-            reservoir_index = resindex,
-            lake_index = lakeindex,
-            reservoir_index_f = filter(x -> x ≠ 0, resindex),
-            lake_index_f = filter(x -> x ≠ 0, lakeindex),
+            reservoir_indices = inds_reservoir_map2river,
+            lake_indices = inds_lake_map2river,
+            land_indices = inds_land_map2river,
             # specific for kinematic_wave
-            upstream_nodes = filter_upsteam_nodes(graph_riv, pits[inds_riv]),
-            subdomain_order = subriv_order,
-            topo_subdomain = topo_subriv,
-            indices_subdomain = indices_subriv,
-            order = toposort_riv,
+            upstream_nodes = filter_upsteam_nodes(graph_river, pits[inds_river]),
+            order_of_subdomains = order_of_river_subdomains,
+            order_subdomain = toposort_river_subdomain,
+            subdomain_indices = river_subdomain_inds,
+            order = toposort_river,
             # water allocation areas
-            indices_allocation_areas = inds_riv_allocation_areas,
-            area = xl[index_river] .* yl[index_river],
+            allocation_area_indices = river_allocation_area_inds,
+            cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     elseif river_routing == "local-inertial"
         river = (
-            graph = graph_riv,
-            indices = inds_riv,
-            reverse_indices = rev_inds_riv,
-            # reservoir and lake index
-            reservoir_index = resindex,
-            lake_index = lakeindex,
-            reservoir_index_f = filter(x -> x ≠ 0, resindex),
-            lake_index_f = filter(x -> x ≠ 0, lakeindex),
+            graph = graph_river,
+            indices = inds_river,
+            reverse_indices = reverse_inds_river,
+            reservoir_indices = inds_reservoir_map2river,
+            lake_indices = inds_lake_map2river,
+            land_indices = inds_land_map2river,
             # specific for local-inertial
-            nodes_at_link = nodes_at_link,
-            links_at_node = adjacent_links_at_node(graph_riv, nodes_at_link),
+            nodes_at_edge = nodes_at_edge,
+            edges_at_node = adjacent_edges_at_node(graph_river, nodes_at_edge),
             # water allocation areas
-            indices_allocation_areas = inds_riv_allocation_areas,
-            area = xl[index_river] .* yl[index_river],
+            allocation_area_indices = river_allocation_area_inds,
+            cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     end
 
     model = Model(
         config,
-        (; land, river, reservoir, lake, drain, index_river, frac_toriver),
-        (subsurface = subsurface_map, land = olf, river = rf),
-        lhm,
+        (; land, river, reservoir = reservoir_network, lake = lake_network, drain),
+        (subsurface = subsurface_map, land = overland_flow, river = river_flow),
+        land_hydrology,
         clock,
         reader,
         writer,
@@ -534,19 +498,20 @@ function update!(model::Model{N, L, V, R, W, T}) where {N, L, V, R, W, T <: SbmG
     (; soil, runoff, demand) = vertical
 
     do_water_demand = haskey(config.model, "water_demand")
-    inds_riv = network.index_river
     aquifer = lateral.subsurface.flow.aquifer
+    dt = tosecond(clock.dt)
 
-    update!(vertical, lateral, network, config)
+    update!(vertical, lateral, network, config, dt)
 
     # set river stage (groundwater) to average h from kinematic wave
-    lateral.subsurface.river.stage .= lateral.river.h_av .+ lateral.subsurface.river.bottom
+    lateral.subsurface.river.variables.stage .=
+        lateral.river.variables.h_av .+ lateral.subsurface.river.parameters.bottom
 
     # determine stable time step for groundwater flow
     conductivity_profile =
         get(config.input.lateral.subsurface, "conductivity_profile", "uniform")
     dt_gw = stable_timestep(aquifer, conductivity_profile) # time step in day (Float64)
-    dt_sbm = (vertical.dt / tosecond(basetimestep)) # vertical.dt is in seconds (Float64)
+    dt_sbm = (dt / tosecond(basetimestep)) # dt is in seconds (Float64)
     if dt_gw < dt_sbm
         @warn(
             "stable time step dt $dt_gw for groundwater flow is smaller than `LandHydrologySBM` model dt $dt_sbm"
@@ -556,9 +521,10 @@ function update!(model::Model{N, L, V, R, W, T}) where {N, L, V, R, W, T <: SbmG
     Q = zeros(lateral.subsurface.flow.connectivity.ncell)
     # exchange of recharge between SBM soil model and groundwater flow domain
     # recharge rate groundwater is required in units [m d⁻¹]
-    @. lateral.subsurface.recharge.rate = soil.variables.recharge / 1000.0 * (1.0 / dt_sbm)
+    @. lateral.subsurface.recharge.variables.rate =
+        soil.variables.recharge / 1000.0 * (1.0 / dt_sbm)
     if do_water_demand
-        @. lateral.subsurface.recharge.rate -=
+        @. lateral.subsurface.recharge.variables.rate -=
             vertical.allocation.variables.act_groundwater_abst / 1000.0 * (1.0 / dt_sbm)
     end
     # update groundwater domain
@@ -567,9 +533,7 @@ function update!(model::Model{N, L, V, R, W, T}) where {N, L, V, R, W, T <: SbmG
     # update SBM soil model (runoff, ustorelayerdepth and satwaterdepth)
     update!(soil, (; runoff, demand, subsurface = lateral.subsurface.flow))
 
-    ssf_toriver = zeros(length(soil.variables.zi))
-    ssf_toriver[inds_riv] = -lateral.subsurface.river.flux ./ lateral.river.dt
-    surface_routing!(model; ssf_toriver = ssf_toriver)
+    surface_routing!(model)
 
     return nothing
 end
diff --git a/src/sbm_model.jl b/src/sbm_model.jl
index cd14aa7dc..e03d48cbf 100644
--- a/src/sbm_model.jl
+++ b/src/sbm_model.jl
@@ -19,11 +19,7 @@ function initialize_sbm_model(config::Config)
     do_lakes = get(config.model, "lakes", false)::Bool
     do_pits = get(config.model, "pits", false)::Bool
 
-    kw_river_tstep = get(config.model, "kw_river_tstep", 0)
-    kw_land_tstep = get(config.model, "kw_land_tstep", 0)
-    kinwave_it = get(config.model, "kin_wave_iteration", false)::Bool
     routing_options = ("kinematic-wave", "local-inertial")
-    floodplain_1d = get(config.model, "floodplain_1d", false)::Bool
     river_routing = get_options(
         config.model,
         "river_routing",
@@ -41,236 +37,214 @@ function initialize_sbm_model(config::Config)
     masswasting = get(config.model, "masswasting", false)::Bool
     @info "General model settings" reservoirs lakes snow masswasting glacier
 
-    nc = NCDataset(static_path)
+    dataset = NCDataset(static_path)
 
-    subcatch_2d = ncread(nc, config, "subcatchment"; optional = false, allow_missing = true)
-    # indices based on catchment
-    inds, rev_inds = active_indices(subcatch_2d, missing)
-    n = length(inds)
+    subcatch_2d =
+        ncread(dataset, config, "subcatchment"; optional = false, allow_missing = true)
+    # indices based on sub-catchments
+    indices, reverse_indices = active_indices(subcatch_2d, missing)
+    n_land_cells = length(indices)
     modelsize_2d = size(subcatch_2d)
 
-    river_2d =
-        ncread(nc, config, "river_location"; optional = false, type = Bool, fill = false)
-    river = river_2d[inds]
-    riverwidth_2d =
-        ncread(nc, config, "lateral.river.width"; optional = false, type = Float, fill = 0)
-    riverwidth = riverwidth_2d[inds]
-    riverlength_2d =
-        ncread(nc, config, "lateral.river.length"; optional = false, type = Float, fill = 0)
-    riverlength = riverlength_2d[inds]
+    river_location_2d = ncread(
+        dataset,
+        config,
+        "river_location";
+        optional = false,
+        type = Bool,
+        fill = false,
+    )
+    river_location = river_location_2d[indices]
+    river_width_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.width";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    river_width = river_width_2d[indices]
+    river_length_2d = ncread(
+        dataset,
+        config,
+        "lateral.river.length";
+        optional = false,
+        type = Float,
+        fill = 0,
+    )
+    river_length = river_length_2d[indices]
 
     # read x, y coordinates and calculate cell length [m]
-    y_nc = read_y_axis(nc)
-    x_nc = read_x_axis(nc)
-    y = permutedims(repeat(y_nc; outer = (1, length(x_nc))))[inds]
-    cellength = abs(mean(diff(x_nc)))
+    y_coords = read_y_axis(dataset)
+    x_coords = read_x_axis(dataset)
+    y = permutedims(repeat(y_coords; outer = (1, length(x_coords))))[indices]
+    cell_length = abs(mean(diff(x_coords)))
 
-    sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
-    xl, yl = cell_lengths(y, cellength, sizeinmetres)
-    riverfrac = river_fraction(river, riverlength, riverwidth, xl, yl)
+    size_in_metres = get(config.model, "sizeinmetres", false)::Bool
+    x_length, y_length = cell_lengths(y, cell_length, size_in_metres)
+    river_fraction =
+        get_river_fraction(river_location, river_length, river_width, x_length, y_length)
 
-    lhm = LandHydrologySBM(nc, config, riverfrac, inds)
+    land_hydrology = LandHydrologySBM(dataset, config, river_fraction, indices)
 
-    inds_riv, rev_inds_riv = active_indices(river_2d, 0)
-    nriv = length(inds_riv)
+    inds_river, reverse_inds_river = active_indices(river_location_2d, 0)
+    n_river_cells = length(inds_river)
 
     # reservoirs
     pits = zeros(Bool, modelsize_2d)
     if do_reservoirs
-        reservoirs, resindex, reservoir, pits =
-            initialize_simple_reservoir(config, nc, inds_riv, nriv, pits, tosecond(dt))
+        reservoir, reservoir_network, inds_reservoir_map2river, pits =
+            SimpleReservoir(dataset, config, inds_river, n_river_cells, pits)
     else
-        reservoir = ()
-        reservoirs = nothing
-        resindex = fill(0, nriv)
+        reservoir_network = (river_indices = [],)
+        inds_reservoir_map2river = fill(0, n_river_cells)
+        reservoir = nothing
     end
 
     # lakes
     if do_lakes
-        lakes, lakeindex, lake, pits =
-            initialize_lake(config, nc, inds_riv, nriv, pits, tosecond(dt))
+        lake, lake_network, inds_lake_map2river, pits =
+            Lake(dataset, config, inds_river, n_river_cells, pits)
     else
-        lake = ()
-        lakes = nothing
-        lakeindex = fill(0, nriv)
+        lake_network = (river_indices = [],)
+        inds_lake_map2river = fill(0, n_river_cells)
+        lake = nothing
     end
 
-    ldd_2d = ncread(nc, config, "ldd"; optional = false, allow_missing = true)
-    ldd = ldd_2d[inds]
+    ldd_2d = ncread(dataset, config, "ldd"; optional = false, allow_missing = true)
+    ldd = ldd_2d[indices]
     if do_pits
-        pits_2d = ncread(nc, config, "pits"; optional = false, type = Bool, fill = false)
-        ldd = set_pit_ldd(pits_2d, ldd, inds)
+        pits_2d =
+            ncread(dataset, config, "pits"; optional = false, type = Bool, fill = false)
+        ldd = set_pit_ldd(pits_2d, ldd, indices)
     end
 
-    landslope =
-        ncread(nc, config, "lateral.land.slope"; optional = false, sel = inds, type = Float)
-    clamp!(landslope, 0.00001, Inf)
-
-    dl = map(detdrainlength, ldd, xl, yl)
-    dw = (xl .* yl) ./ dl
+    land_slope = ncread(
+        dataset,
+        config,
+        "lateral.land.slope";
+        optional = false,
+        sel = indices,
+        type = Float,
+    )
+    clamp!(land_slope, 0.00001, Inf)
+    flow_length = map(get_flow_length, ldd, x_length, y_length)
+    flow_width = (x_length .* y_length) ./ flow_length
 
     # check if lateral subsurface flow component is defined for the SBM model, when coupled
     # to another groundwater model, this component is not defined in the TOML file.
-    subsurface_flow = haskey(config.input.lateral, "subsurface")
-    if subsurface_flow
-        khfrac = ncread(
-            nc,
+    do_lateral_ssf = haskey(config.input.lateral, "subsurface")
+    if do_lateral_ssf
+        subsurface_flow = LateralSSF(
+            dataset,
             config,
-            "lateral.subsurface.ksathorfrac";
-            sel = inds,
-            defaults = 1.0,
-            type = Float,
-        )
-
-        (; theta_s, theta_r, soilthickness) = lhm.soil.parameters
-        (; zi) = lhm.soil.variables
-        ssf_soilthickness = soilthickness .* 0.001
-        ssf_zi = zi .* 0.001
-
-        kh_profile_type = get(config.input.vertical, "ksat_profile", "exponential")::String
-        if kh_profile_type == "exponential"
-            (; kv_0, f) = lhm.soil.parameters.kv_profile
-            kh_0 = khfrac .* kv_0 .* 0.001 .* (basetimestep / dt)
-            kh_profile = KhExponential(kh_0, f .* 1000.0)
-        elseif kh_profile_type == "exponential_constant"
-            (; z_exp) = lhm.soil.parameters.kv_profile
-            (; kv_0, f) = lhm.soil.parameters.kv_profile.exponential
-            kh_0 = khfrac .* kv_0 .* 0.001 .* (basetimestep / dt)
-            exp_profile = KhExponential(kh_0, f .* 1000.0)
-            kh_profile = KhExponentialConstant(exp_profile, z_exp .* 0.001)
-        elseif kh_profile_type == "layered" || kh_profile_type == "layered_exponential"
-            kh_profile = KhLayered(fill(mv, n))
-        end
-
-        # unit for lateral subsurface flow component is [m³ d⁻¹], kv_0 [mm Δt⁻¹]
-        ssf = LateralSSF{Float, typeof(kh_profile)}(;
-            kh_profile = kh_profile,
-            khfrac = khfrac,
-            zi = ssf_zi,
-            soilthickness = ssf_soilthickness,
-            theta_s,
-            theta_r,
-            dt = dt / basetimestep,
-            slope = landslope,
-            dl = dl,
-            dw = dw,
-            exfiltwater = fill(mv, n),
-            recharge = fill(mv, n),
-            ssf = fill(mv, n),
-            ssfin = fill(mv, n),
-            ssfmax = fill(mv, n),
-            to_river = zeros(n),
-            volume = (theta_s .- theta_r) .* (ssf_soilthickness .- ssf_zi) .* (xl .* yl),
+            indices;
+            soil = land_hydrology.soil,
+            slope = land_slope,
+            flow_length,
+            flow_width,
+            x_length,
+            y_length,
         )
         # update variables `ssf`, `ssfmax` and `kh` (layered profile) based on ksat_profile
+        kh_profile_type = get(config.input.vertical, "ksat_profile", "exponential")::String
         if kh_profile_type == "exponential" || kh_profile_type == "exponential_constant"
-            initialize_lateralssf!(ssf, kh_profile)
+            initialize_lateralssf!(subsurface_flow, subsurface_flow.parameters.kh_profile)
         elseif kh_profile_type == "layered" || kh_profile_type == "layered_exponential"
-            (; kv_profile) = lhm.soil.parameters
-            initialize_lateralssf!(ssf, lhm.soil, kv_profile, tosecond(dt))
+            (; kv_profile) = land_hydrology.soil.parameters
+            initialize_lateralssf!(
+                subsurface_flow,
+                land_hydrology.soil,
+                kv_profile,
+                tosecond(dt),
+            )
         end
     else
         # when the SBM model is coupled (BMI) to a groundwater model, the following
         # variables are expected to be exchanged from the groundwater model.
-        ssf = GroundwaterExchange{Float}(;
-            dt = dt / basetimestep,
-            exfiltwater = fill(mv, n),
-            zi = fill(mv, n),
-            to_river = fill(mv, n),
-            ssf = zeros(n),
-        )
+        subsurface_flow = GroundwaterExchange(n_land_cells)
     end
 
-    graph = flowgraph(ldd, inds, pcr_dir)
-    ldd_riv = ldd_2d[inds_riv]
+    graph = flowgraph(ldd, indices, pcr_dir)
+    ldd_river = ldd_2d[inds_river]
     if do_pits
-        ldd_riv = set_pit_ldd(pits_2d, ldd_riv, inds_riv)
+        ldd_river = set_pit_ldd(pits_2d, ldd_river, inds_river)
     end
-    graph_riv = flowgraph(ldd_riv, inds_riv, pcr_dir)
+    graph_river = flowgraph(ldd_river, inds_river, pcr_dir)
 
-    # the indices of the river cells in the land(+river) cell vector
-    index_river = filter(i -> !isequal(river[i], 0), 1:n)
-    frac_toriver = fraction_runoff_toriver(graph, ldd, index_river, landslope, n)
+    # land indices where river is located
+    inds_land_map2river = filter(i -> !isequal(river_location[i], 0), 1:n_land_cells)
+    frac_to_river = fraction_runoff_to_river(graph, ldd, inds_land_map2river, land_slope)
 
-    inds_allocation_areas = Vector{Int}[]
-    inds_riv_allocation_areas = Vector{Int}[]
+    allocation_area_inds = Vector{Int}[]
+    river_allocation_area_inds = Vector{Int}[]
     if do_water_demand
-        areas = unique(lhm.allocation.parameters.areas)
+        areas = unique(land_hydrology.allocation.parameters.areas)
         for a in areas
-            area_index = findall(x -> x == a, lhm.allocation.parameters.areas)
-            push!(inds_allocation_areas, area_index)
-            area_riv_index =
-                findall(x -> x == a, lhm.allocation.parameters.areas[index_river])
-            push!(inds_riv_allocation_areas, area_riv_index)
+            area_index = findall(x -> x == a, land_hydrology.allocation.parameters.areas)
+            push!(allocation_area_inds, area_index)
+            area_riv_index = findall(
+                x -> x == a,
+                land_hydrology.allocation.parameters.areas[inds_land_map2river],
+            )
+            push!(river_allocation_area_inds, area_riv_index)
         end
     end
 
     if land_routing == "kinematic-wave"
-        olf = initialize_surfaceflow_land(
-            nc,
+        overland_flow = KinWaveOverlandFlow(
+            dataset,
             config,
-            inds;
-            sl = landslope,
-            dl,
-            width = map(det_surfacewidth, dw, riverwidth, river),
-            iterate = kinwave_it,
-            tstep = kw_land_tstep,
-            dt,
+            indices;
+            slope = land_slope,
+            flow_length,
+            flow_width = map(det_surfacewidth, flow_width, river_width, river_location),
         )
     elseif land_routing == "local-inertial"
-        index_river_nf = rev_inds_riv[inds] # not filtered (with zeros)
-        olf, indices = initialize_shallowwater_land(
-            nc,
+        inds_river_map2land = reverse_inds_river[indices] # not filtered (with zeros)
+        overland_flow, staggered_indices = LocalInertialOverlandFlow(
+            dataset,
             config,
-            inds;
+            indices;
             modelsize_2d,
-            indices_reverse = rev_inds,
-            xlength = xl,
-            ylength = yl,
-            riverwidth = riverwidth_2d[inds_riv],
-            graph_riv,
-            ldd_riv,
-            inds_riv,
-            river,
-            waterbody = !=(0).(resindex + lakeindex),
-            dt,
+            reverse_indices,
+            x_length,
+            y_length,
+            river_width = river_width_2d[inds_river],
+            graph_river,
+            ldd_river,
+            inds_river,
+            river_location,
+            waterbody = !=(0).(inds_reservoir_map2river + inds_lake_map2river),
         )
     end
 
-    riverlength = riverlength_2d[inds_riv]
-    riverwidth = riverwidth_2d[inds_riv]
-    minimum(riverlength) > 0 || error("river length must be positive on river cells")
-    minimum(riverwidth) > 0 || error("river width must be positive on river cells")
+    river_length = river_length_2d[inds_river]
+    river_width = river_width_2d[inds_river]
+    minimum(river_length) > 0 || error("river length must be positive on river cells")
+    minimum(river_width) > 0 || error("river width must be positive on river cells")
     if river_routing == "kinematic-wave"
-        rf = initialize_surfaceflow_river(
-            nc,
+        river_flow = KinWaveRiverFlow(
+            dataset,
             config,
-            inds_riv;
-            dl = riverlength,
-            width = riverwidth,
-            reservoir_index = resindex,
-            reservoir = reservoirs,
-            lake_index = lakeindex,
-            lake = lakes,
-            iterate = kinwave_it,
-            tstep = kw_river_tstep,
-            dt = dt,
+            inds_river;
+            river_length,
+            river_width,
+            reservoir = reservoir,
+            lake = lake,
         )
     elseif river_routing == "local-inertial"
-        rf, nodes_at_link = initialize_shallowwater_river(
-            nc,
+        river_flow, nodes_at_edge = LocalInertialRiverFlow(
+            dataset,
             config,
-            inds_riv;
-            graph = graph_riv,
-            ldd = ldd_riv,
-            dl = riverlength,
-            width = riverwidth,
-            reservoir_index = resindex,
-            reservoir = reservoirs,
-            lake_index = lakeindex,
-            lake = lakes,
-            dt = dt,
-            floodplain = floodplain_1d,
+            inds_river;
+            graph_river,
+            ldd_river,
+            river_length,
+            river_width,
+            reservoir,
+            lake,
+            waterbody = !=(0).(inds_reservoir_map2river + inds_lake_map2river),
         )
     else
         error(
@@ -285,60 +259,64 @@ function initialize_sbm_model(config::Config)
     toposort = topological_sort_by_dfs(graph)
     if land_routing == "kinematic-wave" ||
        river_routing == "kinematic-wave" ||
-       subsurface_flow
+       do_lateral_ssf
         streamorder = stream_order(graph, toposort)
     end
-    if land_routing == "kinematic-wave" || subsurface_flow
+    if land_routing == "kinematic-wave" || do_lateral_ssf
         toposort = topological_sort_by_dfs(graph)
-        index_pit_land = findall(x -> x == 5, ldd)
+        land_pit_inds = findall(x -> x == 5, ldd)
         min_streamorder_land = get(config.model, "min_streamorder_land", 5)
-        subbas_order, indices_subbas, topo_subbas = kinwave_set_subdomains(
+        order_of_subdomains, subdomain_inds, toposort_subdomain = kinwave_set_subdomains(
             graph,
             toposort,
-            index_pit_land,
+            land_pit_inds,
             streamorder,
             min_streamorder_land,
         )
     end
     if river_routing == "kinematic-wave"
         min_streamorder_river = get(config.model, "min_streamorder_river", 6)
-        toposort_riv = topological_sort_by_dfs(graph_riv)
-        index_pit_river = findall(x -> x == 5, ldd_riv)
-        subriv_order, indices_subriv, topo_subriv = kinwave_set_subdomains(
-            graph_riv,
-            toposort_riv,
-            index_pit_river,
-            streamorder[index_river],
-            min_streamorder_river,
-        )
+        toposort_river = topological_sort_by_dfs(graph_river)
+        river_pit_inds = findall(x -> x == 5, ldd_river)
+        order_of_river_subdomains, river_subdomain_inds, toposort_river_subdomain =
+            kinwave_set_subdomains(
+                graph_river,
+                toposort_river,
+                river_pit_inds,
+                streamorder[inds_land_map2river],
+                min_streamorder_river,
+            )
     end
 
     if nthreads() > 1
         if river_routing == "kinematic-wave"
             @info "Parallel execution of kinematic wave" min_streamorder_land min_streamorder_river
-        elseif land_routing == "kinematic-wave" || subsurface_flow
+        elseif land_routing == "kinematic-wave" || do_lateral_ssf
             @info "Parallel execution of kinematic wave" min_streamorder_land
         end
     end
 
-    modelmap = (vertical = lhm, lateral = (subsurface = ssf, land = olf, river = rf))
+    modelmap = (
+        vertical = land_hydrology,
+        lateral = (subsurface = subsurface_flow, land = overland_flow, river = river_flow),
+    )
     indices_reverse = (
-        land = rev_inds,
-        river = rev_inds_riv,
-        reservoir = isempty(reservoir) ? nothing : reservoir.reverse_indices,
-        lake = isempty(lake) ? nothing : lake.reverse_indices,
+        land = reverse_indices,
+        river = reverse_inds_river,
+        reservoir = isnothing(reservoir) ? nothing : reservoir_network.reverse_indices,
+        lake = isnothing(lake) ? nothing : lake_network.reverse_indices,
     )
-    (; maxlayers) = lhm.soil.parameters
+    (; maxlayers) = land_hydrology.soil.parameters
     writer = prepare_writer(
         config,
         modelmap,
         indices_reverse,
-        x_nc,
-        y_nc,
-        nc;
+        x_coords,
+        y_coords,
+        dataset;
         extra_dim = (name = "layer", value = Float64.(1:(maxlayers))),
     )
-    close(nc)
+    close(dataset)
 
     # for each domain save:
     # - the directed acyclic graph (graph),
@@ -355,80 +333,83 @@ function initialize_sbm_model(config::Config)
     # functions
     land = (
         graph = graph,
-        upstream_nodes = filter_upsteam_nodes(graph, pits[inds]),
-        subdomain_order = subbas_order,
-        topo_subdomain = topo_subbas,
-        indices_subdomain = indices_subbas,
+        upstream_nodes = filter_upsteam_nodes(graph, pits[indices]),
+        order_of_subdomains,
+        order_subdomain = toposort_subdomain,
+        subdomain_indices = subdomain_inds,
         order = toposort,
-        indices = inds,
-        reverse_indices = rev_inds,
-        area = xl .* yl,
-        slope = landslope,
-        indices_allocation_areas = inds_allocation_areas,
+        indices,
+        reverse_indices,
+        area = x_length .* y_length,
+        slope = land_slope,
+        frac_to_river,
+        allocation_area_indices = allocation_area_inds,
     )
     if land_routing == "local-inertial"
-        land = merge(land, (index_river = index_river_nf, staggered_indices = indices))
+        land = merge(land, (river_indices = inds_river_map2land, staggered_indices))
     end
     if do_water_demand
         # exclude waterbodies for local surface and ground water abstraction
-        inds_riv_2d = copy(rev_inds_riv)
-        inds_2d = ones(Bool, modelsize_2d)
-        if !isempty(reservoir)
-            inds_cov = collect(Iterators.flatten(reservoir.indices_coverage))
+        inds_riv_2d = copy(reverse_inds_river)
+        inds_2d = zeros(Bool, modelsize_2d)
+        if !isnothing(reservoir)
+            inds_cov = collect(Iterators.flatten(reservoir_network.indices_coverage))
             inds_riv_2d[inds_cov] .= 0
-            inds_2d[inds_cov] .= 0
+            inds_2d[inds_cov] .= 1
         end
-        if !isempty(lake)
-            inds_cov = collect(Iterators.flatten(lake.indices_coverage))
+        if !isnothing(lake)
+            inds_cov = collect(Iterators.flatten(lake_network.indices_coverage))
             inds_riv_2d[inds_cov] .= 0
-            inds_2d[inds_cov] .= 0
+            inds_2d[inds_cov] .= 1
         end
-        land = merge(land, (index_river_wb = inds_riv_2d[inds], index_wb = inds_2d[inds]))
+        land = merge(
+            land,
+            (
+                river_inds_excl_waterbody = inds_riv_2d[indices],
+                waterbody = inds_2d[indices],
+            ),
+        )
     end
     if river_routing == "kinematic-wave"
         river = (
-            graph = graph_riv,
-            indices = inds_riv,
-            reverse_indices = rev_inds_riv,
-            # reservoir and lake index
-            reservoir_index = resindex,
-            lake_index = lakeindex,
-            reservoir_index_f = filter(x -> x ≠ 0, resindex),
-            lake_index_f = filter(x -> x ≠ 0, lakeindex),
+            graph = graph_river,
+            indices = inds_river,
+            reverse_indices = reverse_inds_river,
+            reservoir_indices = inds_reservoir_map2river,
+            lake_indices = inds_lake_map2river,
+            land_indices = inds_land_map2river,
             # specific for kinematic_wave
-            upstream_nodes = filter_upsteam_nodes(graph_riv, pits[inds_riv]),
-            subdomain_order = subriv_order,
-            topo_subdomain = topo_subriv,
-            indices_subdomain = indices_subriv,
-            order = toposort_riv,
+            upstream_nodes = filter_upsteam_nodes(graph_river, pits[inds_river]),
+            order_of_subdomains = order_of_river_subdomains,
+            order_subdomain = toposort_river_subdomain,
+            subdomain_indices = river_subdomain_inds,
+            order = toposort_river,
             # water allocation areas
-            indices_allocation_areas = inds_riv_allocation_areas,
-            area = xl[index_river] .* yl[index_river],
+            allocation_area_indices = river_allocation_area_inds,
+            cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     elseif river_routing == "local-inertial"
         river = (
-            graph = graph_riv,
-            indices = inds_riv,
-            reverse_indices = rev_inds_riv,
-            # reservoir and lake index
-            reservoir_index = resindex,
-            lake_index = lakeindex,
-            reservoir_index_f = filter(x -> x ≠ 0, resindex),
-            lake_index_f = filter(x -> x ≠ 0, lakeindex),
+            graph = graph_river,
+            indices = inds_river,
+            reverse_indices = reverse_inds_river,
+            reservoir_indices = inds_reservoir_map2river,
+            lake_indices = inds_lake_map2river,
+            land_indices = inds_land_map2river,
             # specific for local-inertial
-            nodes_at_link = nodes_at_link,
-            links_at_node = adjacent_links_at_node(graph_riv, nodes_at_link),
+            nodes_at_edge = nodes_at_edge,
+            edges_at_node = adjacent_edges_at_node(graph_river, nodes_at_edge),
             # water allocation areas
-            indices_allocation_areas = inds_riv_allocation_areas,
-            area = xl[index_river] .* yl[index_river],
+            allocation_area_indices = river_allocation_area_inds,
+            cell_area = x_length[inds_land_map2river] .* y_length[inds_land_map2river],
         )
     end
 
     model = Model(
         config,
-        (; land, river, reservoir, lake, index_river, frac_toriver),
-        (subsurface = ssf, land = olf, river = rf),
-        lhm,
+        (; land, river, reservoir = reservoir_network, lake = lake_network),
+        (subsurface = subsurface_flow, land = overland_flow, river = river_flow),
+        land_hydrology,
         clock,
         reader,
         writer,
@@ -443,22 +424,25 @@ end
 
 "update SBM model for a single timestep"
 function update!(model::Model{N, L, V, R, W, T}) where {N, L, V, R, W, T <: SbmModel}
-    (; lateral, vertical, network, config) = model
+    (; lateral, vertical, network, clock, config) = model
+    dt = tosecond(clock.dt)
     do_water_demand = haskey(config.model, "water_demand")
     (; kv_profile) = vertical.soil.parameters
 
     update_until_recharge!(model)
     # exchange of recharge between SBM soil model and subsurface flow domain
-    lateral.subsurface.recharge .= vertical.soil.variables.recharge ./ 1000.0
+    lateral.subsurface.boundary_conditions.recharge .=
+        vertical.soil.variables.recharge ./ 1000.0
     if do_water_demand
-        @. lateral.subsurface.recharge -=
+        @. lateral.subsurface.boundary_conditions.recharge -=
             vertical.allocation.variables.act_groundwater_abst / 1000.0
     end
-    lateral.subsurface.recharge .*= lateral.subsurface.dw
-    lateral.subsurface.zi .= vertical.soil.variables.zi ./ 1000.0
+    lateral.subsurface.boundary_conditions.recharge .*=
+        lateral.subsurface.parameters.flow_width
+    lateral.subsurface.variables.zi .= vertical.soil.variables.zi ./ 1000.0
     # update lateral subsurface flow domain (kinematic wave)
-    kh_layered_profile!(vertical.soil, lateral.subsurface, kv_profile, vertical.dt)
-    update!(lateral.subsurface, network.land, network.frac_toriver)
+    kh_layered_profile!(vertical.soil, lateral.subsurface, kv_profile, dt)
+    update!(lateral.subsurface, network.land, clock.dt / basetimestep)
     update_after_subsurfaceflow!(model)
     update_total_water_storage!(model)
     return nothing
@@ -473,8 +457,9 @@ through BMI, to couple the SBM model to an external groundwater model.
 function update_until_recharge!(
     model::Model{N, L, V, R, W, T},
 ) where {N, L, V, R, W, T <: SbmModel}
-    (; lateral, vertical, network, config) = model
-    update!(vertical, lateral, network, config)
+    (; lateral, vertical, network, clock, config) = model
+    dt = tosecond(clock.dt)
+    update!(vertical, lateral, network, config, dt)
     return nothing
 end
 
@@ -494,8 +479,7 @@ function update_after_subsurfaceflow!(
     # update SBM soil model (runoff, ustorelayerdepth and satwaterdepth)
     update!(soil, (; runoff, demand, subsurface))
 
-    ssf_toriver = lateral.subsurface.to_river ./ tosecond(basetimestep)
-    surface_routing!(model; ssf_toriver = ssf_toriver)
+    surface_routing!(model)
 
     return nothing
 end
@@ -514,7 +498,7 @@ function update_total_water_storage!(
     # TODO Maybe look at routing in the near future
     update_total_water_storage!(
         vertical,
-        network.index_river,
+        network.river.land_indices,
         network.land.area,
         lateral.river,
         lateral.land,
@@ -526,6 +510,10 @@ function set_states!(
     model::Model{N, L, V, R, W, T},
 ) where {N, L, V, R, W, T <: Union{SbmModel, SbmGwfModel}}
     (; lateral, vertical, network, config) = model
+    land_v = lateral.land.variables
+    land_p = lateral.land.parameters
+    river_v = lateral.river.variables
+    river_p = lateral.river.parameters
 
     reinit = get(config.model, "reinit", true)::Bool
     routing_options = ("kinematic-wave", "local-inertial")
@@ -551,26 +539,24 @@ function set_states!(
         vertical.soil.variables.zi .= zi
         if land_routing == "kinematic-wave"
             # make sure land cells with zero flow width are set to zero q and h
-            for i in eachindex(lateral.land.width)
-                if lateral.land.width[i] <= 0.0
-                    lateral.land.q[i] = 0.0
-                    lateral.land.h[i] = 0.0
+            for i in eachindex(land_p.flow_width)
+                if land_p.flow_width[i] <= 0.0
+                    land_v.q[i] = 0.0
+                    land_v.h[i] = 0.0
                 end
             end
-            lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
+            land_v.volume .= land_v.h .* land_p.flow_width .* land_p.flow_length
         elseif land_routing == "local-inertial"
             for i in eachindex(lateral.land.volume)
-                if lateral.land.rivercells[i]
+                if land_p.rivercells[i]
                     j = network.land.index_river[i]
-                    if lateral.land.h[i] > 0.0
-                        lateral.land.volume[i] =
-                            lateral.land.h[i] * lateral.land.xl[i] * lateral.land.yl[i] +
-                            lateral.river.bankfull_volume[j]
+                    if land_v.h[i] > 0.0
+                        land_v.volume[i] =
+                            land_v.h[i] * land_p.xl[i] * land_p.yl[i] +
+                            land_p.bankfull_volume[j]
                     else
-                        lateral.land.volume[i] =
-                            lateral.river.h[j] *
-                            lateral.river.width[j] *
-                            lateral.river.dl[j]
+                        land_v.volume[i] =
+                            river_v.h[j] * river_p.flow_width[j] * river_p.flow_length[j]
                     end
                 else
                     lateral.land.volume[i] =
@@ -579,9 +565,8 @@ function set_states!(
             end
         end
         # only set active cells for river (ignore boundary conditions/ghost points)
-        lateral.river.volume[1:nriv] .=
-            lateral.river.h[1:nriv] .* lateral.river.width[1:nriv] .*
-            lateral.river.dl[1:nriv]
+        river_v.volume[1:nriv] .=
+            river_v.h[1:nriv] .* river_p.flow_width[1:nriv] .* river_p.flow_length[1:nriv]
 
         if floodplain_1d
             initialize_volume!(lateral.river, nriv)
diff --git a/src/sediment.jl b/src/sediment.jl
index ec5c8f8d1..fbe08e418 100644
--- a/src/sediment.jl
+++ b/src/sediment.jl
@@ -316,8 +316,8 @@ function initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
         eroslagg = fill(mv, n),
         ### Transport capacity part ###
         # Parameters
-        dl = map(detdrainlength, ldd, xl, yl),
-        dw = map(detdrainwidth, ldd, xl, yl),
+        dl = map(get_flow_length, ldd, xl, yl),
+        dw = map(get_flow_width, ldd, xl, yl),
         cGovers = cGovers,
         D50 = D50,
         dmclay = dmclay,
diff --git a/src/sediment_model.jl b/src/sediment_model.jl
index cf2b78eaf..886640c19 100644
--- a/src/sediment_model.jl
+++ b/src/sediment_model.jl
@@ -51,7 +51,7 @@ function initialize_sediment_model(config::Config)
 
     sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
     xl, yl = cell_lengths(y, cellength, sizeinmetres)
-    riverfrac = river_fraction(river, riverlength, riverwidth, xl, yl)
+    riverfrac = get_river_fraction(river, riverlength, riverwidth, xl, yl)
 
     eros = initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
 
@@ -107,7 +107,7 @@ function initialize_sediment_model(config::Config)
     graph_riv = flowgraph(ldd_riv, inds_riv, pcr_dir)
 
     index_river = filter(i -> !isequal(river[i], 0), 1:n)
-    frac_toriver = fraction_runoff_toriver(graph, ldd, index_river, landslope, n)
+    frac_toriver = fraction_runoff_to_river(graph, ldd, index_river, landslope)
 
     rs = initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
 
diff --git a/src/snow/snow.jl b/src/snow/snow.jl
index a60521b13..c67c54ea5 100644
--- a/src/snow/snow.jl
+++ b/src/snow/snow.jl
@@ -68,45 +68,45 @@ end
 struct NoSnowModel{T} <: AbstractSnowModel{T} end
 
 "Initialize snow HBV model parameters"
-function SnowHbvParameters(nc, config, inds, dt)
+function SnowHbvParameters(dataset, config, indices, dt)
     cfmax =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.snow.parameters.cfmax";
-            sel = inds,
+            sel = indices,
             defaults = 3.75653,
             type = Float,
         ) .* (dt / basetimestep)
     tt = ncread(
-        nc,
+        dataset,
         config,
         "vertical.snow.parameters.tt";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
     )
     tti = ncread(
-        nc,
+        dataset,
         config,
         "vertical.snow.parameters.tti";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
     )
     ttm = ncread(
-        nc,
+        dataset,
         config,
         "vertical.snow.parameters.ttm";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
     )
     whc = ncread(
-        nc,
+        dataset,
         config,
         "vertical.snow.parameters.whc";
-        sel = inds,
+        sel = indices,
         defaults = 0.1,
         type = Float,
     )
@@ -116,9 +116,9 @@ function SnowHbvParameters(nc, config, inds, dt)
 end
 
 "Initialize snow HBV model"
-function SnowHbvModel(nc, config, inds, dt)
-    n = length(inds)
-    params = SnowHbvParameters(nc, config, inds, dt)
+function SnowHbvModel(dataset, config, indices, dt)
+    n = length(indices)
+    params = SnowHbvParameters(dataset, config, indices, dt)
     vars = SnowVariables(Float, n)
     bc = SnowBC(Float, n)
     model = SnowHbvModel(; boundary_conditions = bc, parameters = params, variables = vars)
diff --git a/src/soil/soil.jl b/src/soil/soil.jl
index 157869c2c..bdd649e9d 100644
--- a/src/soil/soil.jl
+++ b/src/soil/soil.jl
@@ -200,18 +200,28 @@ end
 end
 
 "Initialize SBM soil model hydraulic conductivity depth profile"
-function sbm_kv_profiles(nc, config, inds, kv_0, f, maxlayers, nlayers, sumlayers, dt)
+function sbm_kv_profiles(
+    dataset,
+    config,
+    indices,
+    kv_0,
+    f,
+    maxlayers,
+    nlayers,
+    sumlayers,
+    dt,
+)
     kv_profile_type = get(config.input.vertical, "ksat_profile", "exponential")::String
-    n = length(inds)
+    n = length(indices)
     if kv_profile_type == "exponential"
         kv_profile = KvExponential(kv_0, f)
     elseif kv_profile_type == "exponential_constant"
         z_exp = ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.z_exp";
             optional = false,
-            sel = inds,
+            sel = indices,
             type = Float,
         )
         exp_profile = KvExponential(kv_0, f)
@@ -219,10 +229,10 @@ function sbm_kv_profiles(nc, config, inds, kv_0, f, maxlayers, nlayers, sumlayer
     elseif kv_profile_type == "layered" || kv_profile_type == "layered_exponential"
         kv =
             ncread(
-                nc,
+                dataset,
                 config,
                 "vertical.soil.parameters.kv";
-                sel = inds,
+                sel = indices,
                 defaults = 1000.0,
                 type = Float,
                 dimname = :layer,
@@ -236,11 +246,11 @@ function sbm_kv_profiles(nc, config, inds, kv_0, f, maxlayers, nlayers, sumlayer
             kv_profile = KvLayered(svectorscopy(kv, Val{maxlayers}()))
         else
             z_layered = ncread(
-                nc,
+                dataset,
                 config,
                 "vertical.soil.parameters.z_layered";
                 optional = false,
-                sel = inds,
+                sel = indices,
                 type = Float,
             )
             nlayers_kv = fill(0, n)
@@ -331,7 +341,7 @@ end
 end
 
 "Initialize SBM soil model parameters"
-function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
+function SbmSoilParameters(dataset, config, vegetation_parameter_set, indices, dt)
     config_thicknesslayers = get(config.model, "thicknesslayers", Float[])
     if length(config_thicknesslayers) > 0
         thicknesslayers = SVector(Tuple(push!(Float.(config_thicknesslayers), mv)))
@@ -344,152 +354,152 @@ function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
     end
     w_soil =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.w_soil";
-            sel = inds,
+            sel = indices,
             defaults = 0.1125,
             type = Float,
         ) .* (dt / basetimestep)
     cf_soil = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.cf_soil";
-        sel = inds,
+        sel = indices,
         defaults = 0.038,
         type = Float,
     )
     # soil parameters
     theta_s = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.theta_s";
-        sel = inds,
+        sel = indices,
         defaults = 0.6,
         type = Float,
     )
     theta_r = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.theta_r";
-        sel = inds,
+        sel = indices,
         defaults = 0.01,
         type = Float,
     )
     kv_0 =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.kv_0";
-            sel = inds,
+            sel = indices,
             defaults = 3000.0,
             type = Float,
         ) .* (dt / basetimestep)
     f = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.f";
-        sel = inds,
+        sel = indices,
         defaults = 0.001,
         type = Float,
     )
     hb = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.hb";
-        sel = inds,
+        sel = indices,
         defaults = -10.0,
         type = Float,
     )
     h1 = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.h1";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
     )
     h2 = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.h2";
-        sel = inds,
+        sel = indices,
         defaults = -100.0,
         type = Float,
     )
     h3_high = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.h3_high";
-        sel = inds,
+        sel = indices,
         defaults = -400.0,
         type = Float,
     )
     h3_low = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.h3_low";
-        sel = inds,
+        sel = indices,
         defaults = -1000.0,
         type = Float,
     )
     h4 = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.h4";
-        sel = inds,
+        sel = indices,
         defaults = -15849.0,
         type = Float,
     )
     alpha_h1 = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.alpha_h1";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
     )
     soilthickness = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.soilthickness";
-        sel = inds,
+        sel = indices,
         defaults = 2000.0,
         type = Float,
     )
     infiltcappath =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.infiltcappath";
-            sel = inds,
+            sel = indices,
             defaults = 10.0,
             type = Float,
         ) .* (dt / basetimestep)
     infiltcapsoil =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.infiltcapsoil";
-            sel = inds,
+            sel = indices,
             defaults = 100.0,
             type = Float,
         ) .* (dt / basetimestep)
     maxleakage =
         ncread(
-            nc,
+            dataset,
             config,
             "vertical.soil.parameters.maxleakage";
-            sel = inds,
+            sel = indices,
             defaults = 0.0,
             type = Float,
         ) .* (dt / basetimestep)
 
     c = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.c";
-        sel = inds,
+        sel = indices,
         defaults = 10.0,
         type = Float,
         dimname = :layer,
@@ -500,10 +510,10 @@ function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
         error("$parname needs a layer dimension of size $maxlayers, but is $size1")
     end
     kvfrac = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.kvfrac";
-        sel = inds,
+        sel = indices,
         defaults = 1.0,
         type = Float,
         dimname = :layer,
@@ -515,36 +525,36 @@ function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
     end
     # fraction compacted area
     pathfrac = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.pathfrac";
-        sel = inds,
+        sel = indices,
         defaults = 0.01,
         type = Float,
     )
 
     # vegetation parameters
     rootdistpar = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.rootdistpar";
-        sel = inds,
+        sel = indices,
         defaults = -500.0,
         type = Float,
     )
     cap_hmax = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.cap_hmax";
-        sel = inds,
+        sel = indices,
         defaults = 2000.0,
         type = Float,
     )
     cap_n = ncread(
-        nc,
+        dataset,
         config,
         "vertical.soil.parameters.cap_n";
-        sel = inds,
+        sel = indices,
         defaults = 2.0,
         type = Float,
     )
@@ -554,20 +564,20 @@ function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
     nlayers = number_of_active_layers.(act_thickl)
 
     if length(config_thicknesslayers) > 0
-        # root fraction read from nc file, in case of multiple soil layers and TOML file
+        # root fraction read from dataset file, in case of multiple soil layers and TOML file
         # includes "vertical.rootfraction"
         if haskey(config.input.vertical.soil.parameters, "rootfraction")
             rootfraction = ncread(
-                nc,
+                dataset,
                 config,
                 "vertical.soil.parameters.rootfraction";
-                sel = inds,
+                sel = indices,
                 optional = false,
                 type = Float,
                 dimname = :layer,
             )
         else
-            n = length(inds)
+            n = length(indices)
             (; rootingdepth) = vegetation_parameter_set
             # default root fraction in case of multiple soil layers
             rootfraction = zeros(Float, maxlayers, n)
@@ -590,12 +600,21 @@ function SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
         rootfraction = ones(Float, maxlayers, n)
     end
 
-    kv_profile =
-        sbm_kv_profiles(nc, config, inds, kv_0, f, maxlayers, nlayers, sumlayers, dt)
+    kv_profile = sbm_kv_profiles(
+        dataset,
+        config,
+        indices,
+        kv_0,
+        f,
+        maxlayers,
+        nlayers,
+        sumlayers,
+        dt,
+    )
 
     soilwatercapacity = @. soilthickness * (theta_s - theta_r)
 
-    n = length(inds)
+    n = length(indices)
     sbm_params = SbmSoilParameters(;
         maxlayers,
         nlayers,
@@ -639,9 +658,9 @@ end
 end
 
 "Initialize SBM soil model"
-function SbmSoilModel(nc, config, vegetation_parameter_set, inds, dt)
-    n = length(inds)
-    params = SbmSoilParameters(nc, config, vegetation_parameter_set, inds, dt)
+function SbmSoilModel(dataset, config, vegetation_parameter_set, indices, dt)
+    n = length(indices)
+    params = SbmSoilParameters(dataset, config, vegetation_parameter_set, indices, dt)
     vars = SbmSoilVariables(n, params)
     bc = SbmSoilBC(n)
     model = SbmSoilModel(; boundary_conditions = bc, parameters = params, variables = vars)
diff --git a/src/states.jl b/src/states.jl
index 4cd991654..b0515c3a6 100644
--- a/src/states.jl
+++ b/src/states.jl
@@ -217,30 +217,37 @@ function extract_required_states(config::Config)
         add_to_required_states(required_states, (:vertical, :soil, :variables), soil_states)
     # Add subsurface states to dict
     if model_type == "sbm_gwf"
-        key_entry = (:lateral, :subsurface, :flow, :aquifer)
+        key_entry = (:lateral, :subsurface, :flow, :aquifer, :variables)
     else
-        key_entry = (:lateral, :subsurface)
+        key_entry = (:lateral, :subsurface, :variables)
     end
     required_states = add_to_required_states(required_states, key_entry, ssf_states)
     # Add land states to dict
     required_states =
-        add_to_required_states(required_states, (:lateral, :land), land_states)
+        add_to_required_states(required_states, (:lateral, :land, :variables), land_states)
     # Add river states to dict
-    required_states =
-        add_to_required_states(required_states, (:lateral, :river), river_states)
+    if model_type == "sediment"
+        key_entry = (:lateral, :river)
+    else
+        key_entry = (:lateral, :river, :variables)
+    end
+    required_states = add_to_required_states(required_states, key_entry, river_states)
     # Add floodplain states to dict
     required_states = add_to_required_states(
         required_states,
-        (:lateral, :river, :floodplain),
+        (:lateral, :river, :floodplain, :variables),
         floodplain_states,
     )
     # Add lake states to dict
-    required_states =
-        add_to_required_states(required_states, (:lateral, :river, :lake), lake_states)
+    required_states = add_to_required_states(
+        required_states,
+        (:lateral, :river, :boundary_conditions, :lake, :variables),
+        lake_states,
+    )
     # Add reservoir states to dict
     required_states = add_to_required_states(
         required_states,
-        (:lateral, :river, :reservoir),
+        (:lateral, :river, :boundary_conditions, :reservoir, :variables),
         reservoir_states,
     )
     # Add paddy states to dict
diff --git a/src/surfacewater/runoff.jl b/src/surfacewater/runoff.jl
index 6203b4020..b53e55e5b 100644
--- a/src/surfacewater/runoff.jl
+++ b/src/surfacewater/runoff.jl
@@ -34,13 +34,13 @@ end
 end
 
 "Initialize open water runoff parameters"
-function OpenWaterRunoffParameters(nc, config, inds, riverfrac)
+function OpenWaterRunoffParameters(dataset, config, indices, riverfrac)
     # fraction open water
     waterfrac = ncread(
-        nc,
+        dataset,
         config,
         "vertical.runoff.parameters.waterfrac";
-        sel = inds,
+        sel = indices,
         defaults = 0.0,
         type = Float,
     )
@@ -73,11 +73,11 @@ end
 end
 
 "Initialize open water runoff model"
-function OpenWaterRunoff(nc, config, inds, riverfrac)
+function OpenWaterRunoff(dataset, config, indices, riverfrac)
     n = length(riverfrac)
     vars = OpenWaterRunoffVariables(Float, n)
     bc = OpenWaterRunoffBC(Float, n)
-    params = OpenWaterRunoffParameters(nc, config, inds, riverfrac)
+    params = OpenWaterRunoffParameters(dataset, config, indices, riverfrac)
     model =
         OpenWaterRunoff(; boundary_conditions = bc, parameters = params, variables = vars)
     return model
@@ -115,15 +115,15 @@ function update_boundary_conditions!(
     network,
 )
     (; water_flux_surface, waterlevel_river, waterlevel_land) = model.boundary_conditions
-    inds_riv = network.index_river
+    (; land_indices) = network.river
     (; snow, glacier, interception) = external_models
 
     get_water_flux_surface!(water_flux_surface, snow, glacier, interception)
 
     # extract water levels h_av [m] from the land and river domains this is used to limit
     # open water evaporation
-    waterlevel_land .= lateral.land.h_av .* 1000.0
-    waterlevel_river[inds_riv] .= lateral.river.h_av .* 1000.0
+    waterlevel_land .= lateral.land.variables.h_av .* 1000.0
+    waterlevel_river[land_indices] .= lateral.river.variables.h_av .* 1000.0
     return nothing
 end
 
diff --git a/src/utils.jl b/src/utils.jl
index 40c514ec1..2599b3199 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -124,23 +124,23 @@ function cell_lengths(y::AbstractVector, cellength::Real, sizeinmetres::Bool)
     return xl, yl
 end
 
-function river_fraction(
+function get_river_fraction(
     river::AbstractVector,
     riverlength::AbstractVector,
     riverwidth::AbstractVector,
-    xl::AbstractVector,
-    yl::AbstractVector,
+    x_length::AbstractVector,
+    y_length::AbstractVector,
 )
     n = length(river)
-    riverfrac = fill(mv, n)
+    river_fraction = fill(mv, n)
     for i in 1:n
-        riverfrac[i] = if river[i]
-            min((riverlength[i] * riverwidth[i]) / (xl[i] * yl[i]), 1.0)
+        river_fraction[i] = if river[i]
+            min((riverlength[i] * riverwidth[i]) / (x_length[i] * y_length[i]), 1.0)
         else
             0.0
         end
     end
-    return riverfrac
+    return river_fraction
 end
 
 """
@@ -382,42 +382,42 @@ function number_of_active_layers(thickness::SVector)
 end
 
 """
-    detdrainlength(ldd, xl, yl)
+    get_flow_length(ldd, x_length, y_length)
 
-Determines the drainaige length for a non square grid. Input `ldd` (drainage network), `xl` (length of cells in x direction),
-`yl` (length of cells in y direction). Output is drainage length.
+Return the flow length for a non square grid. Input `ldd` (drainage network), `x_length`
+(length of cells in x direction), `y_length` (length of cells in y direction). Output is
+flow length.
 """
-function detdrainlength(ldd, xl, yl)
+function get_flow_length(ldd, x_length, y_length)
     # take into account non-square cells
-    # if ldd is 8 or 2 use ylength
-    # if ldd is 4 or 6 use xlength
+    # if ldd is 8 or 2 use y_length
+    # if ldd is 4 or 6 use x_length
     if ldd == 2 || ldd == 8
-        yl
+        y_length
     elseif ldd == 4 || ldd == 6
-        xl
+        x_length
     else
-        hypot(xl, yl)
+        hypot(x_length, y_length)
     end
 end
 
 """
-    detdrainwidth(ldd, xl, yl)
+    get_flow_width(ldd, x_length, y_length)
 
-Determines the drainaige width for a non square grid. Input `ldd` (drainage network), `xl`
-(length of cells in x direction), `yl` (length of cells in y direction). Output is drainage
-width.
+Return the Flow width for a non square grid. Input `ldd` (drainage network), `x_length`
+(length of cells in x direction), `y_length` (length of cells in y direction). Output is
+flow width.
 """
-function detdrainwidth(ldd, xl, yl)
+function get_flow_width(ldd, x_length, y_length)
     # take into account non-square cells
-    # if ldd is 8 or 2 use xlength
-    # if ldd is 4 or 6 use ylength
-    slantwidth = (xl + yl) * 0.5
+    # if ldd is 8 or 2 use x_length
+    # if ldd is 4 or 6 use y_length
     if ldd == 2 || ldd == 8
-        xl
+        x_length
     elseif ldd == 4 || ldd == 6
-        yl
+        y_length
     else
-        slantwidth
+        (x_length + y_length) * 0.5
     end
 end
 
@@ -462,22 +462,23 @@ function svectorscopy(x::Matrix{T}, ::Val{N}) where {T, N}
 end
 
 """
-    fraction_runoff_toriver(graph, ldd, index_river, slope, n)
+    fraction_runoff_to_river(graph, ldd, index_river, slope)
 
-Determine ratio `frac` between `slope` river cell `index_river` and `slope` of each upstream
-neighbor (based on directed acyclic graph `graph`).
+Return ratio `fraction` between `slope` river cell `inds_river` and `slope` of each
+upstream neighbor (based on directed acyclic graph `graph`).
 """
-function fraction_runoff_toriver(graph, ldd, index_river, slope, n)
-    frac = zeros(n)
-    for i in index_river
+function fraction_runoff_to_river(graph, ldd, inds_river, slope)
+    n = length(slope)
+    fraction = zeros(n)
+    for i in inds_river
         nbs = inneighbors(graph, i)
         for j in nbs
             if ldd[j] != ldd[i]
-                frac[j] = slope[j] / (slope[i] + slope[j])
+                fraction[j] = slope[j] / (slope[i] + slope[j])
             end
         end
     end
-    return frac
+    return fraction
 end
 
 """
@@ -521,29 +522,29 @@ tosecond(x::T) where {T <: DatePeriod} = Float64(Dates.value(Second(x)))
 tosecond(x::T) where {T <: TimePeriod} = x / convert(T, Second(1))
 
 """
-    adjacent_nodes_at_link(graph)
+    adjacent_nodes_at_edge(graph)
 
-Return the source node `src` and destination node `dst` of each link of a directed `graph`.
+Return the source node `src` and destination node `dst` of each edge of a directed `graph`.
 """
-function adjacent_nodes_at_link(graph)
-    links = collect(edges(graph))
-    return (src = src.(links), dst = dst.(links))
+function adjacent_nodes_at_edge(graph)
+    _edges = collect(edges(graph))
+    return (src = src.(_edges), dst = dst.(_edges))
 end
 
 """
-    adjacent_links_at_node(graph, nodes_at_link)
+    adjacent_edges_at_node(graph, nodes_at_edge)
 
-Return the source link `src` and destination link `dst` of each node of a directed `graph`.
+Return the source edge `src` and destination edge `dst` of each node of a directed `graph`.
 """
-function adjacent_links_at_node(graph, nodes_at_link)
+function adjacent_edges_at_node(graph, nodes_at_edge)
     nodes = vertices(graph)
-    src_link = Vector{Int}[]
-    dst_link = copy(src_link)
+    src_edge = Vector{Int}[]
+    dst_edge = copy(src_edge)
     for i in 1:nv(graph)
-        push!(src_link, findall(isequal(nodes[i]), nodes_at_link.dst))
-        push!(dst_link, findall(isequal(nodes[i]), nodes_at_link.src))
+        push!(src_edge, findall(isequal(nodes[i]), nodes_at_edge.dst))
+        push!(dst_edge, findall(isequal(nodes[i]), nodes_at_edge.src))
     end
-    return (src = src_link, dst = dst_link)
+    return (src = src_edge, dst = dst_edge)
 end
 
 "Add `vertex` and `edge` to `pits` of a directed `graph`"
@@ -557,34 +558,34 @@ function add_vertex_edge_graph!(graph, pits)
 end
 
 """
-    set_effective_flowwidth!(we_x, we_y, indices, graph_riv, riverwidth, ldd_riv, inds_rev_riv)
+    set_effective_flowwidth!(we_x, we_y, indices, graph_river, river_width, ldd_river, rev_inds_river)
 
 For river cells (D8 flow direction) in a staggered grid the effective flow width at cell
 edges (floodplain) `we_x` in the x-direction and `we_y` in the y-direction is corrected by
-subtracting the river width `riverwidth` from the cell edges. For diagonal directions, the
-`riverwidth` is split between the two adjacent cell edges. A cell edge at linear index `idx`
-is defined as the edge between node `idx` and the adjacent node (+ CartesianIndex(1, 0)) for
-x and (+ CartesianIndex(0, 1)) for y. For cells that contain a `waterbody` (reservoir or
-lake), the effective flow width is set to zero.
+subtracting the river width `river_width` from the cell edges. For diagonal directions, the
+`river_width` is split between the two adjacent cell edges. A cell edge at linear index
+`idx` is defined as the edge between node `idx` and the adjacent node (+ CartesianIndex(1,
+0)) for x and (+ CartesianIndex(0, 1)) for y. For cells that contain a `waterbody`
+(reservoir or lake), the effective flow width is set to zero.
 """
 function set_effective_flowwidth!(
     we_x,
     we_y,
     indices,
-    graph_riv,
-    riverwidth,
-    ldd_riv,
+    graph_river,
+    river_width,
+    ldd_river,
     waterbody,
-    inds_rev_riv,
+    rev_inds_river,
 )
-    toposort = topological_sort_by_dfs(graph_riv)
+    toposort = topological_sort_by_dfs(graph_river)
     n = length(we_x)
     for v in toposort
-        dst = outneighbors(graph_riv, v)
+        dst = outneighbors(graph_river, v)
         isempty(dst) && continue
-        w = min(riverwidth[v], riverwidth[only(dst)])
-        dir = pcr_dir[ldd_riv[v]]
-        idx = inds_rev_riv[v]
+        w = min(river_width[v], river_width[only(dst)])
+        dir = pcr_dir[ldd_river[v]]
+        idx = rev_inds_river[v]
         # loop over river D8 directions
         if dir == CartesianIndex(1, 1)
             we_x[idx] = waterbody[v] ? 0.0 : max(we_x[idx] - 0.5 * w, 0.0)
@@ -727,8 +728,8 @@ function kh_layered_profile!(
 )
     (; nlayers, sumlayers, act_thickl, soilthickness) = soil.parameters
     (; n_unsatlayers, zi) = soil.variables
-    (; kh) = subsurface.kh_profile
-    (; khfrac) = subsurface
+    (; kh) = subsurface.parameters.kh_profile
+    (; khfrac) = subsurface.parameters
 
     t_factor = (tosecond(basetimestep) / dt)
     for i in eachindex(kh)
@@ -764,8 +765,8 @@ function kh_layered_profile!(
     (; nlayers, sumlayers, act_thickl, soilthickness) = soil.parameters
     (; nlayers_kv, z_layered, kv, f) = kv_profile
     (; n_unsatlayers, zi) = soil.variables
-    (; kh) = subsurface.kh_profile
-    (; khfrac) = subsurface
+    (; kh) = subsurface.parameters.kh_profile
+    (; khfrac) = subsurface.parameters
     t_factor = (tosecond(basetimestep) / dt)
 
     for i in eachindex(kh)
@@ -834,17 +835,19 @@ conductivity profile `kh_profile`.
 """
 function initialize_lateralssf!(model::LateralSSF, kh_profile::KhExponential)
     (; kh_0, f) = kh_profile
-    (; ssf, ssfmax, zi, slope, soilthickness, dw) = model
+    (; ssf, ssfmax, zi) = model.variables
+    (; slope, soilthickness, flow_width) = model.parameters
 
     @. ssfmax = ((kh_0 * slope) / f) * (1.0 - exp(-f * soilthickness))
-    @. ssf = ((kh_0 * slope) / f) * (exp(-f * zi) - exp(-f * soilthickness)) * dw
+    @. ssf = ((kh_0 * slope) / f) * (exp(-f * zi) - exp(-f * soilthickness)) * flow_width
     return nothing
 end
 
 function initialize_lateralssf!(model::LateralSSF, kh_profile::KhExponentialConstant)
     (; kh_0, f) = kh_profile.exponential
     (; z_exp) = kh_profile
-    (; ssf, ssfmax, zi, slope, soilthickness, dw) = model
+    (; ssf, ssfmax, zi) = model.variables
+    (; slope, soilthickness, flow_width) = model.parameters
     ssf_constant = @. kh_0 * exp(-f * z_exp) * slope * (soilthickness - z_exp)
     for i in eachindex(ssf)
         ssfmax[i] =
@@ -854,10 +857,14 @@ function initialize_lateralssf!(model::LateralSSF, kh_profile::KhExponentialCons
                 (
                     ((kh_0[i] * slope[i]) / f[i]) *
                     (exp(-f[i] * zi[i]) - exp(-f[i] * z_exp[i])) + ssf_constant[i]
-                ) * dw[i]
+                ) * flow_width[i]
         else
             ssf[i] =
-                kh_0[i] * exp(-f[i] * zi[i]) * slope[i] * (soilthickness[i] - zi[i]) * dw[i]
+                kh_0[i] *
+                exp(-f[i] * zi[i]) *
+                slope[i] *
+                (soilthickness[i] - zi[i]) *
+                flow_width[i]
         end
     end
     return nothing
@@ -876,12 +883,14 @@ function initialize_lateralssf!(
     kv_profile::KvLayered,
     dt,
 )
-    (; kh) = subsurface.kh_profile
+    (; kh) = subsurface.parameters.kh_profile
     (; nlayers, act_thickl) = soil.parameters
-    (; ssf, ssfmax, zi, khfrac, soilthickness, slope, dw) = subsurface
+    (; ssf, ssfmax, zi) = subsurface.variables
+    (; khfrac, soilthickness, slope, flow_width) = subsurface.parameters
+
     kh_layered_profile!(soil, subsurface, kv_profile, dt)
     for i in eachindex(ssf)
-        ssf[i] = kh[i] * (soilthickness[i] - zi[i]) * slope[i] * dw[i]
+        ssf[i] = kh[i] * (soilthickness[i] - zi[i]) * slope[i] * flow_width[i]
         kh_max = 0.0
         for j in 1:nlayers[i]
             kh_max += kv_profile.kv[i][j] * act_thickl[i][j]
@@ -898,14 +907,15 @@ function initialize_lateralssf!(
     kv_profile::KvLayeredExponential,
     dt,
 )
-    (; ssf, ssfmax, zi, khfrac, soilthickness, slope, dw) = subsurface
+    (; ssf, ssfmax, zi) = subsurface.variables
+    (; khfrac, soilthickness, slope, flow_width) = subsurface.parameters
     (; nlayers, act_thickl) = soil.parameters
-    (; kh) = subsurface.kh_profile
+    (; kh) = subsurface.parameters.kh_profile
     (; kv, f, nlayers_kv, z_layered) = kv_profile
 
     kh_layered_profile!(soil, subsurface, kv_profile, dt)
     for i in eachindex(ssf)
-        ssf[i] = kh[i] * (soilthickness[i] - zi[i]) * slope[i] * dw[i]
+        ssf[i] = kh[i] * (soilthickness[i] - zi[i]) * slope[i] * flow_width[i]
         kh_max = 0.0
         for j in 1:nlayers[i]
             if j <= nlayers_kv[i]
diff --git a/src/vegetation/canopy.jl b/src/vegetation/canopy.jl
index 98762ea50..2ab44e7ce 100644
--- a/src/vegetation/canopy.jl
+++ b/src/vegetation/canopy.jl
@@ -39,16 +39,16 @@ end
 end
 
 "Initialize Gash interception model"
-function GashInterceptionModel(nc, config, inds, vegetation_parameter_set)
+function GashInterceptionModel(dataset, config, indices, vegetation_parameter_set)
     e_r = ncread(
-        nc,
+        dataset,
         config,
         "vertical.interception.parameters.e_r";
-        sel = inds,
+        sel = indices,
         defaults = 0.1,
         type = Float,
     )
-    n = length(inds)
+    n = length(indices)
     params =
         GashParameters(; e_r = e_r, vegetation_parameter_set = vegetation_parameter_set)
     vars = InterceptionVariables(Float, n)


Test patch:
diff --git a/server/test/client.jl b/server/test/client.jl
index 1688af26f..58f039923 100644
--- a/server/test/client.jl
+++ b/server/test/client.jl
@@ -32,13 +32,13 @@ end
 
 @testset "model information functions" begin
     @test request((fn = "get_component_name",)) == Dict("component_name" => "sbm")
-    @test request((fn = "get_input_item_count",)) == Dict("input_item_count" => 203)
-    @test request((fn = "get_output_item_count",)) == Dict("output_item_count" => 203)
+    @test request((fn = "get_input_item_count",)) == Dict("input_item_count" => 202)
+    @test request((fn = "get_output_item_count",)) == Dict("output_item_count" => 202)
     to_check = [
         "vertical.soil.parameters.nlayers",
         "vertical.soil.parameters.theta_r",
-        "lateral.river.q",
-        "lateral.river.reservoir.outflow",
+        "lateral.river.variables.q",
+        "lateral.river.boundary_conditions.reservoir.variables.outflow",
     ]
     retrieved_vars = request((fn = "get_input_var_names",))["input_var_names"]
     @test all(x -> x in retrieved_vars, to_check)
@@ -49,12 +49,12 @@ end
 zi_size = 0
 vwc_1_size = 0
 @testset "variable information and get and set functions" begin
-    @test request((fn = "get_var_itemsize", name = "lateral.subsurface.ssf")) ==
+    @test request((fn = "get_var_itemsize", name = "lateral.subsurface.variables.ssf")) ==
           Dict("var_itemsize" => sizeof(Wflow.Float))
     @test request((fn = "get_var_type", name = "vertical.n"))["status"] == "ERROR"
     @test request((fn = "get_var_units", name = "vertical.soil.parameters.theta_s")) ==
           Dict("var_units" => "-")
-    @test request((fn = "get_var_location", name = "lateral.river.q")) ==
+    @test request((fn = "get_var_location", name = "lateral.river.variables.q")) ==
           Dict("var_location" => "node")
     zi_nbytes =
         request((fn = "get_var_nbytes", name = "vertical.soil.variables.zi"))["var_nbytes"]
@@ -68,19 +68,20 @@ vwc_1_size = 0
     vwc_1_itemsize =
         request((fn = "get_var_itemsize", name = "vertical.soil.variables.vwc[1]"))["var_itemsize"]
     vwc_1_size = Int(vwc_1_nbytes / vwc_1_itemsize)
-    @test request((fn = "get_var_grid", name = "lateral.river.h")) == Dict("var_grid" => 3)
+    @test request((fn = "get_var_grid", name = "lateral.river.variables.h")) ==
+          Dict("var_grid" => 3)
     msg = (fn = "get_value", name = "vertical.soil.variables.zi", dest = fill(0.0, zi_size))
     @test mean(request(msg)["value"]) ≈ 277.3620724821974
     msg = (fn = "get_value_ptr", name = "vertical.soil.parameters.theta_s")
     @test mean(request(msg)["value_ptr"]) ≈ 0.4409211971535584
     msg = (
         fn = "get_value_at_indices",
-        name = "lateral.river.q",
+        name = "lateral.river.variables.q",
         dest = [0.0, 0.0, 0.0],
         inds = [1, 5, 10],
     )
     @test request(msg)["value_at_indices"] ≈
-          [2.198747900215207f0, 2.6880427720508515f0, 3.4848783702629564f0]
+          [2.1007361866518766, 2.5702292750107687, 3.2904803551115727]
     msg =
         (fn = "set_value", name = "vertical.soil.variables.zi", src = fill(300.0, zi_size))
     @test request(msg) == Dict("status" => "OK")
diff --git a/server/test/sbm_config.toml b/server/test/sbm_config.toml
index 39898a488..9e805a61f 100644
--- a/server/test/sbm_config.toml
+++ b/server/test/sbm_config.toml
@@ -31,18 +31,18 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 h = "h_river"
 h_av = "h_av_river"
 q = "q_river"
 
-[state.lateral.river.reservoir]
+[state.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 h = "h_land"
 h_av = "h_av_land"
 q = "q_land"
@@ -112,7 +112,7 @@ offset = 0.0
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 bankfull_elevation = "RiverZ"
@@ -132,7 +132,7 @@ targetminfrac = "ResTargetMinFrac"
 ksathorfrac = "KsatHorFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [model]
@@ -161,17 +161,17 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 h = "h_river"
 q = "q_river"
 
-[output.lateral.river.reservoir]
+[output.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[output.lateral.subsurface]
+[output.lateral.subsurface.variables]
 ssf = "ssf"
 
-[output.lateral.land]
+[output.lateral.land.variables]
 h = "h_land"
 q = "q_land"
 
@@ -181,7 +181,7 @@ path = "output_scalar_moselle.nc"
 [[netcdf.variable]]
 name = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[netcdf.variable]]
 coordinate.x = 6.255
@@ -202,13 +202,13 @@ path = "output_moselle.csv"
 
 [[csv.column]]
 header = "Q"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 reducer = "maximum"
 
 [[csv.column]]
 header = "volume"
 index = 1
-parameter = "lateral.river.reservoir.volume"
+parameter = "lateral.river.boundary_conditions.reservoir.variables.volume"
 
 [[csv.column]]
 coordinate.x = 6.255
@@ -232,7 +232,7 @@ parameter = "vertical.atmospheric_forcing.temperature"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[csv.column]]
 header = "recharge"
@@ -255,8 +255,19 @@ components = [
   "vertical.snow.boundary_conditions",
   "vertical.snow.variables",
   "vertical.snow.parameters",
-  "lateral.subsurface",
-  "lateral.land",
-  "lateral.river",
-  "lateral.river.reservoir",
+  "lateral.subsurface.boundary_conditions",
+  "lateral.subsurface.variables",
+  "lateral.subsurface.parameters",
+  "lateral.subsurface.parameters.kh_profile",
+  "lateral.land.boundary_conditions",
+  "lateral.land.variables",
+  "lateral.land.variables.flow",
+  "lateral.land.parameters",
+  "lateral.river.variables",
+  "lateral.river.parameters",
+  "lateral.river.parameters.flow",
+  "lateral.river.boundary_conditions",
+  "lateral.river.boundary_conditions.reservoir.boundary_conditions",
+  "lateral.river.boundary_conditions.reservoir.parameters",
+  "lateral.river.boundary_conditions.reservoir.variables",
 ]
diff --git a/test/bmi.jl b/test/bmi.jl
index 510487f22..df56fcc42 100644
--- a/test/bmi.jl
+++ b/test/bmi.jl
@@ -17,13 +17,13 @@ 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) == 203
-            @test BMI.get_output_item_count(model) == 203
+            @test BMI.get_input_item_count(model) == 202
+            @test BMI.get_output_item_count(model) == 202
             to_check = [
                 "vertical.soil.parameters.nlayers",
                 "vertical.soil.parameters.theta_r",
-                "lateral.river.q",
-                "lateral.river.reservoir.outflow",
+                "lateral.river.variables.q",
+                "lateral.river.boundary_conditions.reservoir.variables.outflow",
             ]
             retrieved_vars = BMI.get_input_var_names(model)
             @test all(x -> x in retrieved_vars, to_check)
@@ -33,18 +33,28 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 
         @testset "variable information functions" begin
             @test BMI.get_var_grid(model, "vertical.soil.parameters.theta_s") == 6
-            @test BMI.get_var_grid(model, "lateral.river.h") == 3
-            @test BMI.get_var_grid(model, "lateral.river.reservoir.inflow") == 0
-            @test_throws ErrorException BMI.get_var_grid(model, "lateral.river.lake.volume")
-            @test BMI.get_var_type(model, "lateral.river.reservoir.inflow") == "$Float"
+            @test BMI.get_var_grid(model, "lateral.river.variables.h") == 3
+            @test BMI.get_var_grid(
+                model,
+                "lateral.river.boundary_conditions.reservoir.boundary_conditions.inflow",
+            ) == 0
+            @test_throws ErrorException BMI.get_var_grid(
+                model,
+                "lateral.river.boundary_conditions.lake.volume",
+            )
+            @test BMI.get_var_type(
+                model,
+                "lateral.river.boundary_conditions.reservoir.boundary_conditions.inflow",
+            ) == "$Float"
             @test BMI.get_var_units(model, "vertical.soil.parameters.theta_s") == "-"
-            @test BMI.get_var_itemsize(model, "lateral.subsurface.ssf") == sizeof(Float)
-            @test BMI.get_var_nbytes(model, "lateral.river.q") ==
-                  length(model.lateral.river.q) * sizeof(Float)
-            @test BMI.get_var_location(model, "lateral.river.q") == "node"
+            @test BMI.get_var_itemsize(model, "lateral.subsurface.variables.ssf") ==
+                  sizeof(Float)
+            @test BMI.get_var_nbytes(model, "lateral.river.variables.q") ==
+                  length(model.lateral.river.variables.q) * sizeof(Float)
+            @test BMI.get_var_location(model, "lateral.river.variables.q") == "node"
             @test_throws ErrorException(
-                "lateral.land.alpha_pow not listed as variable for BMI exchange",
-            ) BMI.get_var_itemsize(model, "lateral.land.alpha_pow")
+                "lateral.land.parameters.alpha_pow not listed as variable for BMI exchange",
+            ) BMI.get_var_itemsize(model, "lateral.land.parameters.alpha_pow")
         end
 
         BMI.update(model)
@@ -69,10 +79,10 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             ) ≈ getindex.(model.vertical.soil.variables.vwc, 2)[1:3]
             @test BMI.get_value_at_indices(
                 model,
-                "lateral.river.q",
+                "lateral.river.variables.q",
                 zeros(Float, 3),
                 [1, 100, 5617],
-            ) ≈ [0.623325399343309, 5.227139951657074, 0.027942874327781947]
+            ) ≈ [0.6525631197206111, 7.493760826794606, 0.02319714614721354]
             BMI.set_value(
                 model,
                 "vertical.soil.variables.zi",
@@ -141,8 +151,8 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
     @testset "BMI grid edges" begin
         tomlpath = joinpath(@__DIR__, "sbm_swf_config.toml")
         model = BMI.initialize(Wflow.Model, tomlpath)
-        @test BMI.get_var_grid(model, "lateral.land.qx") == 4
-        @test BMI.get_var_grid(model, "lateral.land.qy") == 5
+        @test BMI.get_var_grid(model, "lateral.land.variables.qx") == 4
+        @test BMI.get_var_grid(model, "lateral.land.variables.qy") == 5
         @test BMI.get_grid_edge_count(model, 4) == 50063
         @test BMI.get_grid_edge_count(model, 5) == 50063
         @test BMI.get_grid_edge_nodes(model, 4, fill(0, 2 * 50063))[1:4] == [1, -999, 2, 3]
@@ -174,12 +184,12 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
         # set zi and exfiltwater from external source (e.g. a groundwater model)
         BMI.set_value(
             model,
-            "lateral.subsurface.zi",
+            "lateral.subsurface.variables.zi",
             fill(0.25, BMI.get_grid_node_count(model, 6)),
         )
         BMI.set_value(
             model,
-            "lateral.subsurface.exfiltwater",
+            "lateral.subsurface.variables.exfiltwater",
             fill(1.0e-5, BMI.get_grid_node_count(model, 6)),
         )
         # update SBM after subsurface flow
@@ -193,9 +203,9 @@ tomlpath = joinpath(@__DIR__, "sbm_config.toml")
             @test sbm.snow.variables.snow_storage[1] ≈ 3.4847899611762876f0
             @test sbm.soil.variables.recharge[5] ≈ 0.0f0
             @test sbm.soil.variables.zi[5] ≈ 250.0f0
-            @test sub.zi[5] ≈ 0.25f0
-            @test sub.exfiltwater[1] ≈ 1.0f-5
-            @test sub.ssf[1] ≈ 0.0f0
+            @test sub.variables.zi[5] ≈ 0.25f0
+            @test sub.variables.exfiltwater[1] ≈ 1.0f-5
+            @test sub.variables.ssf[1] ≈ 0.0f0
         end
 
         BMI.finalize(model)
diff --git a/test/groundwater.jl b/test/groundwater.jl
index def5ce7e5..427c4c96b 100644
--- a/test/groundwater.jl
+++ b/test/groundwater.jl
@@ -35,28 +35,35 @@ function homogenous_aquifer(nrow, ncol)
     connectivity = Wflow.Connectivity(indices, reverse_indices, dx, dy)
     ncell = connectivity.ncell
 
-    conf_aqf = Wflow.ConfinedAquifer(
-        [0.0, 7.5, 20.0],  # head
-        fill(10.0, ncell),  # k
-        fill(10.0, ncell),  # top
-        fill(0.0, ncell),  # bottom
-        fill(100.0, ncell),  # area
-        fill(0.1, ncell), # specific storage
-        fill(1.0, ncell),  # storativity
-        fill(0.0, connectivity.nconnection),  # conductance
-        fill(0.0, ncell), # total volume that can be released
+    parameters = Wflow.ConfinedAquiferParameters(;
+        k = fill(10.0, ncell),
+        top = fill(10.0, ncell),
+        bottom = fill(0.0, ncell),
+        area = fill(100.0, ncell),
+        specific_storage = fill(0.1, ncell),
+        storativity = fill(1.0, ncell),
     )
-    unconf_aqf = Wflow.UnconfinedAquifer(
-        [0.0, 7.5, 20.0],  # head
-        fill(10.0, ncell),   # k
-        fill(10.0, ncell),  # top
-        fill(0.0, ncell),  # bottom
-        fill(100.0, ncell),  # area
-        fill(0.15, ncell),  # specific yield
-        fill(0.0, connectivity.nconnection),  # conductance
-        fill(0.0, ncell), # total volume that can be released
-        fill(3.0, ncell), # conductance reduction factor
+    variables = Wflow.ConfinedAquiferVariables(;
+        head = [0.0, 7.5, 20.0],
+        conductance = fill(0.0, connectivity.nconnection),
+        volume = fill(0.0, ncell),
     )
+    conf_aqf = Wflow.ConfinedAquifer(; parameters, variables)
+
+    parameters = Wflow.UnconfinedAquiferParameters(;
+        k = fill(10.0, ncell),
+        top = fill(10.0, ncell),
+        bottom = fill(0.0, ncell),
+        area = fill(100.0, ncell),
+        specific_yield = fill(0.15, ncell),
+        f = fill(3.0, ncell),
+    )
+    variables = Wflow.UnconfinedAquiferVariables(;
+        head = [0.0, 7.5, 20.0],
+        conductance = fill(0.0, connectivity.nconnection),
+        volume = fill(0.0, ncell),
+    )
+    unconf_aqf = Wflow.UnconfinedAquifer(; parameters, variables)
     return (connectivity, conf_aqf, unconf_aqf)
 end
 
@@ -244,19 +251,19 @@ end
         end
 
         @testset "minimum_head-confined" begin
-            original_head = copy(conf_aqf.head)
-            conf_aqf.head[1] = -10.0
+            original_head = copy(conf_aqf.variables.head)
+            conf_aqf.variables.head[1] = -10.0
             @test Wflow.check_flux(-1.0, conf_aqf, 1) == -1.0
             @test Wflow.minimum_head(conf_aqf)[1] == -10.0
-            conf_aqf.head .= original_head
+            conf_aqf.variables.head .= original_head
         end
 
         @testset "minimum_head-unconfined" begin
-            original_head = copy(unconf_aqf.head)
-            unconf_aqf.head[1] = -10.0
+            original_head = copy(unconf_aqf.variables.head)
+            unconf_aqf.variables.head[1] = -10.0
             @test Wflow.check_flux(-1.0, unconf_aqf, 1) == 0.0
             @test Wflow.minimum_head(conf_aqf)[1] == 0.0
-            unconf_aqf.head .= original_head
+            unconf_aqf.variables.head .= original_head
         end
 
         @testset "stable_timestep" begin
@@ -288,14 +295,13 @@ end
         end
 
         @testset "river" begin
-            river = Wflow.River(
-                [2.0, 2.0],
-                [100.0, 100.0],
-                [200.0, 200.0],
-                [1.0, 1.0],
-                [0.0, 0.0],
-                [1, 3],
+            parameters = Wflow.RiverParameters(;
+                infiltration_conductance = [100.0, 100.0],
+                exfiltration_conductance = [200.0, 200.0],
+                bottom = [1.0, 1.0],
             )
+            variables = Wflow.RiverVariables(; stage = [2.0, 2.0], flux = [0.0, 0.0])
+            river = Wflow.River(; parameters, variables, index = [1, 3])
             Q = zeros(3)
             Wflow.flux!(Q, river, conf_aqf)
             # infiltration, below bottom, flux is (stage - bottom) * inf_cond
@@ -305,7 +311,12 @@ end
         end
 
         @testset "drainage" begin
-            drainage = Wflow.Drainage([2.0, 2.0], [100.0, 100.0], [0.0, 0.0], [1, 2])
+            parameters = Wflow.DrainageParameters(;
+                elevation = [2.0, 2.0],
+                conductance = [100.0, 100.0],
+            )
+            variables = Wflow.DrainageVariables(; flux = [0.0, 0.0])
+            drainage = Wflow.Drainage(; parameters, variables, index = [1, 2])
             Q = zeros(3)
             Wflow.flux!(Q, drainage, conf_aqf)
             @test Q[1] == 0.0
@@ -313,8 +324,10 @@ end
         end
 
         @testset "headboundary" begin
-            headboundary =
-                Wflow.HeadBoundary([2.0, 2.0], [100.0, 100.0], [0.0, 0.0], [1, 2])
+            parameters = Wflow.HeadBoundaryParameters(; conductance = [100.0, 100.0])
+            variables = Wflow.HeadBoundaryVariables(; head = [2.0, 2.0], flux = [0.0, 0.0])
+
+            headboundary = Wflow.HeadBoundary(; parameters, variables, index = [1, 2])
             Q = zeros(3)
             Wflow.flux!(Q, headboundary, conf_aqf)
             @test Q[1] == 100.0 * (2.0 - 0.0)
@@ -322,14 +335,19 @@ end
         end
 
         @testset "recharge" begin
-            recharge = Wflow.Recharge([1.0e-3, 1.0e-3, 1.0e-3], [0.0, 0.0, 0.0], [1, 2, 3])
+            variables = Wflow.RechargeVariables(;
+                rate = [1.0e-3, 1.0e-3, 1.0e-3],
+                flux = [0.0, 0.0, 0.0],
+            )
+            recharge = Wflow.Recharge(; variables, index = [1, 2, 3])
             Q = zeros(3)
             Wflow.flux!(Q, recharge, conf_aqf)
             @test all(Q .== 1.0e-3 * 100.0)
         end
 
         @testset "well" begin
-            well = Wflow.Well([-1000.0], [0.0], [1])
+            variables = Wflow.WellVariables(; volumetric_rate = [-1000.0], flux = [0.0])
+            well = Wflow.Well(; variables, index = [1])
             Q = zeros(3)
             Wflow.flux!(Q, well, conf_aqf)
             @test Q[1] == -1000.0
@@ -338,7 +356,8 @@ end
 
     @testset "integration: steady 1D" begin
         connectivity, aquifer, _ = homogenous_aquifer(3, 1)
-        constanthead = Wflow.ConstantHead([2.0, 4.0], [1, 3])
+        variables = Wflow.ConstantHeadVariables(; head = [2.0, 4.0])
+        constanthead = Wflow.ConstantHead(; variables, index = [1, 3])
         conductivity_profile = "uniform"
         gwf = Wflow.GroundwaterFlow{Wflow.Float}(;
             aquifer = aquifer,
@@ -347,7 +366,8 @@ end
             boundaries = Wflow.AquiferBoundaryCondition[],
         )
         # Set constant head (dirichlet) boundaries
-        gwf.aquifer.head[gwf.constanthead.index] .= gwf.constanthead.head
+        gwf.aquifer.variables.head[gwf.constanthead.index] .=
+            gwf.constanthead.variables.head
 
         Q = zeros(3)
         dt = 0.25 # days
@@ -355,12 +375,13 @@ end
             Wflow.update!(gwf, Q, dt, conductivity_profile)
         end
 
-        @test gwf.aquifer.head ≈ [2.0, 3.0, 4.0]
+        @test gwf.aquifer.variables.head ≈ [2.0, 3.0, 4.0]
     end
 
     @testset "integration: steady 1D, exponential conductivity" begin
         connectivity, aquifer, _ = homogenous_aquifer(3, 1)
-        constanthead = Wflow.ConstantHead([2.0, 4.0], [1, 3])
+        variables = Wflow.ConstantHeadVariables(; head = [2.0, 4.0])
+        constanthead = Wflow.ConstantHead(; variables, index = [1, 3])
         conductivity_profile = "exponential"
         gwf = Wflow.GroundwaterFlow{Wflow.Float}(;
             aquifer = aquifer,
@@ -369,7 +390,8 @@ end
             boundaries = Wflow.AquiferBoundaryCondition[],
         )
         # Set constant head (dirichlet) boundaries
-        gwf.aquifer.head[gwf.constanthead.index] .= gwf.constanthead.head
+        gwf.aquifer.variables.head[gwf.constanthead.index] .=
+            gwf.constanthead.variables.head
 
         Q = zeros(3)
         dt = 0.25 # days
@@ -377,7 +399,7 @@ end
             Wflow.update!(gwf, Q, dt, conductivity_profile)
         end
 
-        @test gwf.aquifer.head ≈ [2.0, 3.0, 4.0]
+        @test gwf.aquifer.variables.head ≈ [2.0, 3.0, 4.0]
     end
 
     @testset "integration: unconfined transient 1D" begin
@@ -402,19 +424,25 @@ end
         connectivity = Wflow.Connectivity(indices, reverse_indices, dx, dy)
         ncell = connectivity.ncell
         xc = collect(range(0.0; stop = aquifer_length - cellsize, step = cellsize))
-        aquifer = Wflow.UnconfinedAquifer(
-            initial_head.(xc),
-            fill(conductivity, ncell),
-            fill(top, ncell),
-            fill(bottom, ncell),
-            fill(cellsize * cellsize, ncell),
-            fill(specific_yield, ncell),
-            fill(0.0, connectivity.nconnection),
-            fill(0.0, ncell),
-            fill(gwf_f, ncell),
+
+        variables = Wflow.UnconfinedAquiferVariables(;
+            head = initial_head.(xc),
+            conductance = fill(0.0, connectivity.nconnection),
+            volume = fill(0.0, ncell),
+        )
+        parameters = Wflow.UnconfinedAquiferParameters(;
+            k = fill(conductivity, ncell),
+            top = fill(top, ncell),
+            bottom = fill(bottom, ncell),
+            area = fill(cellsize * cellsize, ncell),
+            specific_yield = fill(specific_yield, ncell),
+            f = fill(gwf_f, ncell),
         )
+
+        aquifer = Wflow.UnconfinedAquifer(; parameters, variables)
         # constant head on left boundary, 0 at 0
-        constanthead = Wflow.ConstantHead([0.0], [1])
+        variables = Wflow.ConstantHeadVariables(; head = [0.0])
+        constanthead = Wflow.ConstantHead(; variables, index = [1])
         gwf = Wflow.GroundwaterFlow{Wflow.Float}(;
             aquifer = aquifer,
             connectivity = connectivity,
@@ -431,7 +459,7 @@ end
         for i in 1:nstep
             Wflow.update!(gwf, Q, dt, conductivity_profile)
             # Gradient dh/dx is positive, all flow to the left
-            @test all(diff(gwf.aquifer.head) .> 0.0)
+            @test all(diff(gwf.aquifer.variables.head) .> 0.0)
         end
 
         head_analytical = [
@@ -444,7 +472,7 @@ end
                 beta,
             ) for x in xc
         ]
-        difference = gwf.aquifer.head .- head_analytical
+        difference = gwf.aquifer.variables.head .- head_analytical
         # @test all(difference .< ?)  #TODO
     end
 
@@ -470,19 +498,25 @@ end
         connectivity = Wflow.Connectivity(indices, reverse_indices, dx, dy)
         ncell = connectivity.ncell
         xc = collect(range(0.0; stop = aquifer_length - cellsize, step = cellsize))
-        aquifer = Wflow.UnconfinedAquifer(
-            initial_head.(xc),
-            fill(conductivity, ncell),
-            fill(top, ncell),
-            fill(bottom, ncell),
-            fill(cellsize * cellsize, ncell),
-            fill(specific_yield, ncell),
-            fill(0.0, connectivity.nconnection),
-            fill(0.0, ncell),
-            fill(gwf_f, ncell),
+
+        variables = Wflow.UnconfinedAquiferVariables(;
+            head = initial_head.(xc),
+            conductance = fill(0.0, connectivity.nconnection),
+            volume = fill(0.0, ncell),
+        )
+        parameters = Wflow.UnconfinedAquiferParameters(;
+            k = fill(conductivity, ncell),
+            top = fill(top, ncell),
+            bottom = fill(bottom, ncell),
+            area = fill(cellsize * cellsize, ncell),
+            specific_yield = fill(specific_yield, ncell),
+            f = fill(gwf_f, ncell),
         )
+
+        aquifer = Wflow.UnconfinedAquifer(; parameters, variables)
         # constant head on left boundary, 0 at 0
-        constanthead = Wflow.ConstantHead([0.0], [1])
+        variables = Wflow.ConstantHeadVariables(; head = [0.0])
+        constanthead = Wflow.ConstantHead(; variables, index = [1])
         gwf = Wflow.GroundwaterFlow{Wflow.Float}(;
             aquifer = aquifer,
             connectivity = connectivity,
@@ -499,7 +533,7 @@ end
         for i in 1:nstep
             Wflow.update!(gwf, Q, dt, conductivity_profile)
             # Gradient dh/dx is positive, all flow to the left
-            @test all(diff(gwf.aquifer.head) .> 0.0)
+            @test all(diff(gwf.aquifer.variables.head) .> 0.0)
         end
 
         head_analytical = [
@@ -512,7 +546,7 @@ end
                 beta,
             ) for x in xc
         ]
-        difference = gwf.aquifer.head .- head_analytical
+        difference = gwf.aquifer.variables.head .- head_analytical
         # @test all(difference .< ?)  #TODO
     end
 
@@ -541,23 +575,29 @@ end
         indices, reverse_indices = Wflow.active_indices(domain, false)
         connectivity = Wflow.Connectivity(indices, reverse_indices, dx, dy)
         ncell = connectivity.ncell
-        aquifer = Wflow.ConfinedAquifer(
-            fill(startinghead, ncell),
-            fill(conductivity, ncell),
-            fill(top, ncell),
-            fill(bottom, ncell),
-            fill(cellsize * cellsize, ncell),
-            fill(specific_storage, ncell),
-            fill(storativity, ncell),
-            fill(0.0, connectivity.nconnection), # conductance, to be set
-            fill(0.0, ncell), # total volume that can be released, to be set
+
+        parameters = Wflow.ConfinedAquiferParameters(;
+            k = fill(conductivity, ncell),
+            top = fill(top, ncell),
+            bottom = fill(bottom, ncell),
+            area = fill(cellsize * cellsize, ncell),
+            specific_storage = fill(specific_storage, ncell),
+            storativity = fill(storativity, ncell),
+        )
+        variables = Wflow.ConfinedAquiferVariables(;
+            head = fill(startinghead, ncell),
+            conductance = fill(0.0, connectivity.nconnection),
+            volume = fill(0.0, ncell),
         )
+        aquifer = Wflow.ConfinedAquifer(; parameters, variables)
 
         cell_index = reshape(collect(range(1, ncell; step = 1)), shape)
         indices = vcat(cell_index[1, :], cell_index[end, :])# , cell_index[:, 1], cell_index[:, end],)
-        constanthead = Wflow.ConstantHead(fill(10.0, size(indices)), indices)
+        variables = Wflow.ConstantHeadVariables(; head = fill(10.0, size(indices)))
+        constanthead = Wflow.ConstantHead(; variables, index = indices)
         # Place a well in the middle of the domain
-        well = Wflow.Well([discharge], [0.0], [reverse_indices[wellrow, wellrow]])
+        variables = Wflow.WellVariables(; volumetric_rate = [discharge], flux = [0.0])
+        well = Wflow.Well(; variables, index = [reverse_indices[wellrow, wellrow]])
         gwf = Wflow.GroundwaterFlow{Wflow.Float}(;
             aquifer = aquifer,
             connectivity = connectivity,
@@ -576,7 +616,7 @@ end
         end
 
         # test for symmetry on x and y axes
-        head = reshape(gwf.aquifer.head, shape)
+        head = reshape(gwf.aquifer.variables.head, shape)
         @test head[1:halfnrow, :] ≈ head[end:-1:(halfnrow + 2), :]
         @test head[:, 1:halfnrow] ≈ head[:, end:-1:(halfnrow + 2)]
 
diff --git a/test/io.jl b/test/io.jl
index a5700b10a..744d952f2 100644
--- a/test/io.jl
+++ b/test/io.jl
@@ -213,7 +213,10 @@ end
 @test Wflow.param(model, "lateral.doesnt_exist", -1) == -1
 
 @testset "warm states" begin
-    @test Wflow.param(model, "lateral.river.reservoir.volume")[1] ≈ 3.2807224993363418e7
+    @test Wflow.param(
+        model,
+        "lateral.river.boundary_conditions.reservoir.variables.volume",
+    )[1] ≈ 3.2807224993363418e7
     @test Wflow.param(model, "vertical.soil.variables.satwaterdepth")[9115] ≈
           477.13548089422125
     @test Wflow.param(model, "vertical.snow.variables.snow_storage")[5] ≈ 11.019233179897599
@@ -223,13 +226,13 @@ end
     @test Wflow.param(model, "vertical.snow.variables.snow_water")[5] ≈ 0.0
     @test Wflow.param(model, "vertical.interception.variables.canopy_storage")[50063] ≈ 0.0
     @test Wflow.param(model, "vertical.soil.variables.zi")[50063] ≈ 296.8028609104624
-    @test Wflow.param(model, "lateral.subsurface.ssf")[10606] ≈ 39.972334552895816
-    @test Wflow.param(model, "lateral.river.q")[149] ≈ 53.48673634956338
-    @test Wflow.param(model, "lateral.river.h")[149] ≈ 1.167635369628945
-    @test Wflow.param(model, "lateral.river.volume")[149] ≈ 63854.60119358985
-    @test Wflow.param(model, "lateral.land.q")[2075] ≈ 3.285909284322251
-    @test Wflow.param(model, "lateral.land.h")[2075] ≈ 0.052076262033771775
-    @test Wflow.param(model, "lateral.land.volume")[2075] ≈ 29920.754983235012
+    @test Wflow.param(model, "lateral.subsurface.variables.ssf")[10606] ≈ 39.972334552895816
+    @test Wflow.param(model, "lateral.river.variables.q")[149] ≈ 53.48673634956338
+    @test Wflow.param(model, "lateral.river.variables.h")[149] ≈ 1.167635369628945
+    @test Wflow.param(model, "lateral.river.variables.volume")[149] ≈ 63854.60119358985
+    @test Wflow.param(model, "lateral.land.variables.q")[2075] ≈ 3.285909284322251
+    @test Wflow.param(model, "lateral.land.variables.h")[2075] ≈ 0.052076262033771775
+    @test Wflow.param(model, "lateral.land.variables.volume")[2075] ≈ 29920.754983235012
 end
 
 @testset "reducer" begin
@@ -453,11 +456,14 @@ end
     @test (:vertical, :soil, :variables, :satwaterdepth) in required_states
     @test (:vertical, :soil, :variables, :ustorelayerdepth) in required_states
     @test (:vertical, :interception, :variables, :canopy_storage) in required_states
-    @test (:lateral, :subsurface, :ssf) in required_states
-    @test (:lateral, :river, :q) in required_states
-    @test (:lateral, :river, :h_av) in required_states
-    @test (:lateral, :land, :h_av) in required_states
-    @test !((:lateral, :river, :lake, :waterlevel) in required_states)
+    @test (:lateral, :subsurface, :variables, :ssf) in required_states
+    @test (:lateral, :river, :variables, :q) in required_states
+    @test (:lateral, :river, :variables, :h_av) in required_states
+    @test (:lateral, :land, :variables, :h_av) in required_states
+    @test !(
+        (:lateral, :river, :boundary_conditions, :lake, :variables, :waterlevel) in
+        required_states
+    )
 
     # Adding an unused state the see if the right warning message is thrown
     config.state.vertical.soil.variables.additional_state = "additional_state"
@@ -481,8 +487,8 @@ end
     @test (:vertical, :soil, :variables, :satwaterdepth) in required_states
     @test (:vertical, :soil, :variables, :ustorelayerdepth) in required_states
     @test (:vertical, :interception, :variables, :canopy_storage) in required_states
-    @test (:lateral, :subsurface, :flow, :aquifer, :head) in required_states
-    @test (:lateral, :river, :q) in required_states
-    @test (:lateral, :river, :h_av) in required_states
-    @test (:lateral, :land, :h_av) in required_states
+    @test (:lateral, :subsurface, :flow, :aquifer, :variables, :head) in required_states
+    @test (:lateral, :river, :variables, :q) in required_states
+    @test (:lateral, :river, :variables, :h_av) in required_states
+    @test (:lateral, :land, :variables, :h_av) in required_states
 end
diff --git a/test/reservoir_lake.jl b/test/reservoir_lake.jl
index 261021927..0dcd881cd 100644
--- a/test/reservoir_lake.jl
+++ b/test/reservoir_lake.jl
@@ -1,65 +1,82 @@
-res = Wflow.SimpleReservoir{Float64}(;
-    dt = 86400.0,
+
+res_bc =
+    Wflow.ReservoirBC{Float}(; inflow = [0.0], precipitation = [4.2], evaporation = [1.5])
+res_params = Wflow.ReservoirParameters{Float64}(;
     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],
+)
+res_vars = Wflow.ReservoirVariables{Float64}(;
+    volume = [1.925e7],
+    outflow_av = [0.0],
     actevap = [0.0],
     outflow = [NaN],
     percfull = [NaN],
     demandrelease = [NaN],
 )
+
+res = Wflow.SimpleReservoir{Float64}(;
+    boundary_conditions = res_bc,
+    parameters = res_params,
+    variables = res_vars,
+)
 @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
-    @test res.volume[1] ≈ 2.0e7
-    @test res.percfull[1] ≈ 0.80
-    @test res.demandrelease[1] ≈ 52.5229994727611
-    @test res.precipitation[1] ≈ 4.2
-    @test res.evaporation[1] ≈ 1.5
-    @test res.actevap[1] ≈ 1.5
+    Wflow.update!(res, 1, 100.0, 86400.0, 86400.0)
+    res.variables.outflow_av ./= 86400.0
+    @test res.variables.outflow[1] ≈ 91.3783714867453
+    @test res.variables.outflow_av[1] == res.variables.outflow[1]
+    @test res.variables.volume[1] ≈ 2.0e7
+    @test res.variables.percfull[1] ≈ 0.80
+    @test res.variables.demandrelease[1] ≈ 52.5229994727611
+    @test res.boundary_conditions.precipitation[1] ≈ 4.2
+    @test res.boundary_conditions.evaporation[1] ≈ 1.5
+    @test res.variables.actevap[1] ≈ 1.5
 end
 
-lake = Wflow.Lake{Float64}(;
-    dt = 86400.0,
+lake_bc = Wflow.LakeBC{Float}(; inflow = [0.0], precipitation = [20.0], evaporation = [3.2])
+lake_params = Wflow.LakeParameters{Float}(;
     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],
+)
+lake_vars = Wflow.LakeVariables{Float}(;
+    outflow_av = [0.0],
     storage = Wflow.initialize_storage([1], [180510409.0], [18.5], [missing]),
     waterlevel = [18.5],
-    precipitation = [20.0],
-    evaporation = [3.2],
     actevap = [0.0],
     outflow = [NaN],
 )
+
+lake = Wflow.Lake{Float64}(;
+    boundary_conditions = lake_bc,
+    parameters = lake_params,
+    variables = lake_vars,
+)
 @testset "Update lake" begin
-    Wflow.update!(lake, 1, 2500.0, 181, 86400.0)
-    @test Wflow.waterlevel(lake.storfunc, lake.area, lake.storage, lake.sh)[1] ≈
+    lake_p = lake.parameters
+    lake_v = lake.variables
+    lake_bc = lake.boundary_conditions
+    Wflow.update!(lake, 1, 2500.0, 181, 86400.0, 86400.0)
+    lake_v.outflow_av ./= 86400.0
+    @test Wflow.waterlevel(lake_p.storfunc, lake_p.area, lake_v.storage, lake_p.sh)[1] ≈
           19.672653848925634
-    @test lake.outflow[1] ≈ 85.14292808113598
-    @test lake.totaloutflow[1] ≈ 7.356348986210149e6
-    @test lake.storage[1] ≈ 3.55111879238499e9
-    @test lake.waterlevel[1] ≈ 19.672653848925634
-    @test lake.precipitation[1] ≈ 20.0
-    @test lake.evaporation[1] ≈ 3.2
-    @test lake.actevap[1] ≈ 3.2
+    @test lake_v.outflow[1] ≈ 85.14292808113598
+    @test lake_v.outflow_av[1] ≈ lake_v.outflow[1]
+    @test lake_v.storage[1] ≈ 3.55111879238499e9
+    @test lake_v.waterlevel[1] ≈ 19.672653848925634
+    @test lake_bc.precipitation[1] ≈ 20.0
+    @test lake_bc.evaporation[1] ≈ 3.2
+    @test lake_v.actevap[1] ≈ 3.2
 end
 
 datadir = joinpath(@__DIR__, "data")
@@ -71,8 +88,7 @@ sh = [
     @test keys(sh[1]) == (:H, :S)
     @test typeof(values(sh[1])) == Tuple{Vector{Float}, Vector{Float}}
 
-    lake = Wflow.Lake{Float}(;
-        dt = 86400.0,
+    lake_params = Wflow.LakeParameters{Float}(;
         lowerlake_ind = [2, 0],
         area = [472461536.0, 60851088.0],
         maxstorage = Wflow.maximum_storage(
@@ -84,16 +100,15 @@ sh = [
         ),
         threshold = [393.7, 0.0],
         storfunc = [2, 2],
-        inflow = [0.0, 0.0],
-        totaloutflow = [0.0, 0.0],
         outflowfunc = [2, 1],
         b = [140.0, 0.0],
         e = [1.5, 1.5],
         sh = sh,
         hq = [missing, Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))],
+    )
+    lake_vars = Wflow.LakeVariables{Float}(;
+        outflow_av = [0.0, 0.0],
         waterlevel = [395.03027, 394.87833],
-        precipitation = [10.0, 10.0],
-        evaporation = [2.0, 2.0],
         actevap = [0.0, 0.0],
         outflow = [NaN, NaN],
         storage = Wflow.initialize_storage(
@@ -103,28 +118,43 @@ sh = [
             sh,
         ),
     )
+    lake_bc = Wflow.LakeBC{Float}(;
+        inflow = [0.0, 0.0],
+        precipitation = [10.0, 10.0],
+        evaporation = [2.0, 2.0],
+    )
 
-    Wflow.update!(lake, 1, 500.0, 15, 86400.0)
-    Wflow.update!(lake, 2, 500.0, 15, 86400.0)
-    @test lake.outflow ≈ [214.80170846121263, 236.83281600000214] atol = 1e-2
-    @test lake.totaloutflow ≈ [1.855886761104877e7, 2.0462355302400187e7] atol = 1e3
-    @test lake.storage ≈ [1.2737435094769483e9, 2.6019755340159863e8] atol = 1e4
-    @test lake.waterlevel ≈ [395.0912274997361, 395.2101079057371] atol = 1e-2
-    lake.actevap .= 0.0
-    lake.totaloutflow .= 0.0
-    lake.inflow .= 0.0
-    Wflow.update!(lake, 1, 500.0, 15, 86400.0)
-    Wflow.update!(lake, 2, 500.0, 15, 86400.0)
-    @test lake.outflow ≈ [0.0, 239.66710359986183] atol = 1e-2
-    @test lake.totaloutflow ≈ [-2.2446764487487033e7, 4.3154002238515094e7] atol = 1e3
-    @test lake.storage ≈ [1.3431699662524352e9, 2.6073035986708355e8] atol = 1e4
-    @test lake.waterlevel ≈ [395.239782021054, 395.21771942667266] atol = 1e-2
-    @test lake.actevap ≈ [2.0, 2.0]
+    lake = Wflow.Lake{Float}(;
+        boundary_conditions = lake_bc,
+        parameters = lake_params,
+        variables = lake_vars,
+    )
+
+    Wflow.update!(lake, 1, 500.0, 15, 86400.0, 86400.0)
+    Wflow.update!(lake, 2, 500.0, 15, 86400.0, 86400.0)
+    lake.variables.outflow_av ./= 86400.0
+    lake_v = lake.variables
+    lake_bc = lake.boundary_conditions
+    @test lake_v.outflow ≈ [214.80170846121263, 236.83281600000214]
+    @test lake_v.outflow_av ≈ lake_v.outflow
+    @test lake_v.storage ≈ [1.2737435094769483e9, 2.6019755340159863e8]
+    lake_v.actevap .= 0.0
+    lake_v.outflow_av .= 0.0
+    lake_bc.inflow .= 0.0
+    Wflow.update!(lake, 1, 500.0, 15, 86400.0, 86400.0)
+    Wflow.update!(lake, 2, 500.0, 15, 86400.0, 86400.0)
+    lake.variables.outflow_av ./= 86400.0
+    @test lake_v.outflow ≈ [-259.8005149014703, 239.66710359986183]
+    @test lake_v.outflow_av ≈ [-259.8005149014703, 499.4676185013321]
+    @test lake_v.storage ≈ [1.3431699662524352e9, 2.6073035986708355e8]
+    @test lake_v.waterlevel ≈ [395.239782021054, 395.21771942667266]
+    @test lake_v.actevap ≈ [2.0, 2.0]
 end
 
 @testset "overflowing lake with sh and hq" begin
-    lake = Wflow.Lake{Float}(;
-        dt = 86400.0,
+    lake_bc =
+        Wflow.LakeBC{Float}(; inflow = [0.00], precipitation = [10.0], evaporation = [2.0])
+    lake_params = Wflow.LakeParameters{Float}(;
         lowerlake_ind = [0],
         area = [200_000_000],
         maxstorage = Wflow.maximum_storage(
@@ -136,26 +166,33 @@ end
         ),
         threshold = [0.0],
         storfunc = [2],
-        inflow = [0.00],
-        totaloutflow = [0.0],
         outflowfunc = [1],
         b = [0.0],
         e = [0.0],
         sh = [Wflow.read_sh_csv(joinpath(datadir, "input", "lake_sh_2.csv"))],
         hq = [Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))],
+    )
+    lake_vars = Wflow.LakeVariables{Float}(;
+        outflow_av = [0.0],
         waterlevel = [397.75],
-        precipitation = [10.0],
-        evaporation = [2.0],
         actevap = [0.0],
         outflow = [NaN],
         storage = [410_760_000],
     )
+    lake = Wflow.Lake{Float}(;
+        boundary_conditions = lake_bc,
+        parameters = lake_params,
+        variables = lake_vars,
+    )
 
-    Wflow.update!(lake, 1, 1500.0, 15, 86400.0)
-    @test Wflow.waterlevel(lake.storfunc, lake.area, lake.storage, lake.sh) ≈ [398.0] atol =
-        1e-2
-    @test lake.outflow ≈ [1303.67476852] atol = 1e-2
-    @test lake.totaloutflow ≈ [11.26375000e7] atol = 1e3
-    @test lake.storage ≈ [4.293225e8] atol = 1e4
-    @test lake.waterlevel ≈ [398.000000] atol = 1e-2
+    Wflow.update!(lake, 1, 1500.0, 15, 86400.0, 86400.0)
+    lake.variables.outflow_av ./= 86400.0
+    lake_p = lake.parameters
+    lake_v = lake.variables
+    @test Wflow.waterlevel(lake_p.storfunc, lake_p.area, lake_v.storage, lake_p.sh) ≈
+          [398.0] atol = 1e-2
+    @test lake_v.outflow ≈ [1303.67476852]
+    @test lake_v.outflow_av ≈ lake_v.outflow
+    @test lake_v.storage ≈ [4.293225e8]
+    @test lake_v.waterlevel ≈ [398.000000]
 end
diff --git a/test/horizontal_process.jl b/test/routing_process.jl
similarity index 77%
rename from test/horizontal_process.jl
rename to test/routing_process.jl
index 7cb8bc780..931060fbe 100644
--- a/test/horizontal_process.jl
+++ b/test/routing_process.jl
@@ -135,7 +135,7 @@ end
     x = [dx:dx:L;]
     zb = first.([quadgk(s, xi, L; rtol = 1e-12) for xi in x])
 
-    # initialize ShallowWaterRiver
+    # initialize local inertial river flow model
     graph = DiGraph(n)
     for i in 1:n
         add_edge!(graph, i, i + 1)
@@ -145,82 +145,95 @@ end
     width = fill(10.0, n)
     n_river = fill(0.03, n)
 
-    # for each link the src and dst node is required
-    nodes_at_link = Wflow.adjacent_nodes_at_link(graph)
+    # for each edge the src and dst node is required
+    nodes_at_edge = Wflow.adjacent_nodes_at_edge(graph)
     _ne = ne(graph)
 
-    # determine z, width, length and manning's n at links
+    # determine z, width, length and manning's n at edges
     zb_max = fill(0.0, _ne)
-    width_at_link = fill(0.0, _ne)
-    length_at_link = fill(0.0, _ne)
+    width_at_edge = fill(0.0, _ne)
+    length_at_edge = fill(0.0, _ne)
     mannings_n_sq = fill(0.0, _ne)
     for i in 1:_ne
-        zb_max[i] = max(zb[nodes_at_link.src[i]], zb[nodes_at_link.dst[i]])
-        width_at_link[i] = min(width[nodes_at_link.dst[i]], width[nodes_at_link.src[i]])
-        length_at_link[i] = 0.5 * (dl[nodes_at_link.dst[i]] + dl[nodes_at_link.src[i]])
+        zb_max[i] = max(zb[nodes_at_edge.src[i]], zb[nodes_at_edge.dst[i]])
+        width_at_edge[i] = min(width[nodes_at_edge.dst[i]], width[nodes_at_edge.src[i]])
+        length_at_edge[i] = 0.5 * (dl[nodes_at_edge.dst[i]] + dl[nodes_at_edge.src[i]])
         mannings_n =
             (
-                n_river[nodes_at_link.dst[i]] * dl[nodes_at_link.dst[i]] +
-                n_river[nodes_at_link.src[i]] * dl[nodes_at_link.src[i]]
-            ) / (dl[nodes_at_link.dst[i]] + dl[nodes_at_link.src[i]])
+                n_river[nodes_at_edge.dst[i]] * dl[nodes_at_edge.dst[i]] +
+                n_river[nodes_at_edge.src[i]] * dl[nodes_at_edge.src[i]]
+            ) / (dl[nodes_at_edge.dst[i]] + dl[nodes_at_edge.src[i]])
         mannings_n_sq[i] = mannings_n * mannings_n
     end
 
+    river_network = (
+        nodes_at_edge = nodes_at_edge,
+        edges_at_node = Wflow.adjacent_edges_at_node(graph, nodes_at_edge),
+    )
     network = (
-        nodes_at_link = nodes_at_link,
-        links_at_node = Wflow.adjacent_links_at_node(graph, nodes_at_link),
+        river = river_network,
+        reservoir = (river_indices = [],),
+        lake = (river_indices = [],),
     )
 
-    alpha = 0.7
-    dt = 1.0
     h_thresh = 1.0e-03
     froude_limit = true
     h_init = zeros(n - 1)
     push!(h_init, h_a[n])
 
-    sw_river = Wflow.ShallowWaterRiver(;
+    timestepping = Wflow.TimeStepping(; cfl = 0.7)
+    parameters = Wflow.LocalInertialRiverFlowParameters(;
         n = n,
         ne = _ne,
         active_n = collect(1:(n - 1)),
         active_e = collect(1:_ne),
         g = 9.80665,
-        alpha = alpha,
         h_thresh = h_thresh,
-        dt = dt,
+        zb_max = zb_max,
+        mannings_n_sq = mannings_n_sq,
+        mannings_n = n_river,
+        flow_width = width,
+        flow_width_at_edge = width_at_edge,
+        flow_length = dl,
+        flow_length_at_edge = length_at_edge,
+        bankfull_volume = fill(Wflow.mv, n),
+        bankfull_depth = fill(Wflow.mv, n),
+        zb = zb,
+        froude_limit = froude_limit,
+        waterbody = zeros(n),
+    )
+
+    variables = Wflow.LocalInertialRiverFlowVariables(;
         q0 = zeros(_ne),
         q = zeros(_ne),
         q_av = zeros(_ne),
         q_channel_av = zeros(_ne),
-        zb_max = zb_max,
-        mannings_n_sq = mannings_n_sq,
-        mannings_n = n_river,
         h = h_init,
         zs_max = zeros(_ne),
         zs_src = zeros(_ne),
         zs_dst = zeros(_ne),
         hf = zeros(_ne),
         h_av = zeros(n),
-        width = width,
-        width_at_link = width_at_link,
         a = zeros(_ne),
         r = zeros(_ne),
         volume = fill(0.0, n),
         error = zeros(n),
+    )
+
+    boundary_conditions = Wflow.RiverFlowBC(;
         inflow = zeros(n),
         abstraction = zeros(n),
-        inflow_wb = zeros(n),
+        inflow_waterbody = zeros(n),
         inwater = zeros(n),
-        dl = dl,
-        dl_at_link = length_at_link,
-        bankfull_volume = fill(Wflow.mv, n),
-        bankfull_depth = fill(Wflow.mv, n),
-        zb = zb,
-        froude_limit = froude_limit,
-        reservoir_index = Int[],
-        lake_index = Int[],
-        waterbody = zeros(n),
         reservoir = nothing,
         lake = nothing,
+    )
+
+    sw_river = Wflow.LocalInertialRiverFlow(;
+        timestepping,
+        boundary_conditions,
+        parameters,
+        variables,
         floodplain = nothing,
         allocation = nothing,
     )
@@ -228,16 +241,16 @@ end
     # run until steady state is reached
     epsilon = 1.0e-12
     while true
-        sw_river.inwater[1] = 20.0
-        h0 = mean(sw_river.h)
+        sw_river.boundary_conditions.inwater[1] = 20.0
+        h0 = mean(sw_river.variables.h)
         dt = Wflow.stable_timestep(sw_river)
-        Wflow.shallowwater_river_update!(sw_river, network, dt, 0.0, true)
-        d = abs(h0 - mean(sw_river.h))
+        Wflow.local_inertial_river_update!(sw_river, network, dt, 86400.0, 0.0, true)
+        d = abs(h0 - mean(sw_river.variables.h))
         if d <= epsilon
             break
         end
     end
 
     # test for mean absolute error [cm]
-    @test mean(abs.(sw_river.h .- h_a)) * 100.0 ≈ 1.873574206931199
+    @test mean(abs.(sw_river.variables.h .- h_a)) * 100.0 ≈ 1.873574206931199
 end
diff --git a/test/run_sbm.jl b/test/run_sbm.jl
index 2d03bed96..282c5bc2d 100644
--- a/test/run_sbm.jl
+++ b/test/run_sbm.jl
@@ -14,26 +14,26 @@ flush(model.writer.csv_io)  # ensure the buffer is written fully to disk
     row = csv_first_row(model.writer.csv_path)
 
     @test row.time == DateTime("2000-01-02T00:00:00")
-    @test row.Q ≈ 8.15299947254324f0
+    @test row.Q ≈ 14.283603959245331f0
     @test row.volume ≈ 2.7535003939625636f7
     @test row.temp_bycoord ≈ 2.390000104904175f0
     @test row.vwc_layer2_bycoord ≈ 0.25938809638672006f0
     @test row.temp_byindex ≈ 2.390000104904175f0
-    @test row.Q_6336050 ≈ 0.006583064321841488f0
-    @test row.Q_6336510 ≈ 0.029864230092642642f0
-    @test row.Q_6836100 ≈ 0.19995488963854305f0
-    @test row.Q_6336500 ≈ 0.006277726622788425f0
-    @test row.Q_6836190 ≈ 0.0031262850749354237f0
-    @test row.Q_6336800 ≈ 0.008278375560053742f0
-    @test row.Q_6336900 ≈ 0.0066141980189014385f0
-    @test row.Q_6336930 ≈ 0.09141703511009937f0
-    @test row.Q_6336910 ≈ 0.007475453481320056f0
-    @test row.Q_6136500 ≈ 0.001834989281902289f0
-    @test row.Q_6136520 ≈ 0.0022266031120691397f0
-    @test row.Q_6136150 ≈ 0.006310361139139334f0
-    @test row.Q_6136151 ≈ 0.007946301730645885f0
-    @test row.Q_6136160 ≈ 3.927719795530719f0
-    @test row.Q_6136202 ≈ 1.4162246003743886f0
+    @test row.Q_6336050 ≈ 0.0077488610359143706f0
+    @test row.Q_6336510 ≈ 0.029475452452069627f0
+    @test row.Q_6836100 ≈ 0.018908801563313697f0
+    @test row.Q_6336500 ≈ 0.006568230884299706f0
+    @test row.Q_6836190 ≈ 0.004495680971248625f0
+    @test row.Q_6336800 ≈ 0.00916756751611816f0
+    @test row.Q_6336900 ≈ 0.008200812406926841f0
+    @test row.Q_6336930 ≈ 0.026477954532813916f0
+    @test row.Q_6336910 ≈ 0.009127079213306247f0
+    @test row.Q_6136500 ≈ 0.001154334630885042f0
+    @test row.Q_6136520 ≈ 0.0014432161225527867f0
+    @test row.Q_6136150 ≈ 0.0073146444422822225f0
+    @test row.Q_6136151 ≈ 0.006614628762328246f0
+    @test row.Q_6136160 ≈ 6.186486795085275f0
+    @test row.Q_6136202 ≈ 1.8472398588556238f0
     @test row.recharge_1 ≈ -0.0020800523945940217f0
 end
 
@@ -41,26 +41,26 @@ end
     ds = model.writer.dataset_scalar
     @test ds["time"][1] == DateTime("2000-01-02T00:00:00")
     @test ds["Q"][:][1:20] ≈ [
-        0.7425387f0,
-        1.4162246f0,
-        1.4425076f0,
-        1.4044669f0,
-        5.738109f0,
-        2.7616737f0,
-        2.1128905f0,
-        4.105428f0,
-        0.008651769f0,
-        3.9277198f0,
-        4.069447f0,
-        0.006356805f0,
-        0.007946302f0,
-        0.008135906f0,
-        0.0037393502f0,
-        0.70888275f0,
-        0.0024000728f0,
-        1.3347782f0,
-        3.8374817f0,
-        1.676597f0,
+        0.79544294f0,
+        1.8472399f0,
+        1.8459954f0,
+        1.5499605f0,
+        13.334657f0,
+        4.865496f0,
+        0.09143141f0,
+        4.437221f0,
+        0.010515818f0,
+        6.1864867f0,
+        6.1764274f0,
+        0.0060072034f0,
+        0.0066146287f0,
+        0.006407934f0,
+        0.00416659f0,
+        1.2095966f0,
+        0.002213421f0,
+        1.832764f0,
+        8.601765f0,
+        2.7730224f0,
     ]
     @test ds["Q_gauges"].attrib["cf_role"] == "timeseries_id"
     @test ds["temp_index"][:] ≈ [2.39f0]
@@ -84,7 +84,7 @@ end
     @test snow.variables.snow_storage[5] ≈ 3.768513390588815f0
     @test mean(snow.variables.snow_storage) ≈ 0.038019723676094325f0
     @test sbm.variables.total_storage[50063] ≈ 559.9035608052374f0
-    @test sbm.variables.total_storage[429] ≈ 597.4578475404879f0 # river cell
+    @test sbm.variables.total_storage[429] ≈ 597.1603591432968f0 # river cell
 end
 
 # run the second timestep
@@ -103,11 +103,11 @@ Wflow.run_timestep!(model)
     @test snow.variables.snow_storage[5] ≈ 3.843412524052313f0
     @test mean(snow.variables.snow_storage) ≈ 0.03461317061870949f0
     @test sbm.variables.total_storage[50063] ≈ 560.0152135062889f0
-    @test sbm.variables.total_storage[429] ≈ 617.2238533241972f0 # river cell
+    @test sbm.variables.total_storage[429] ≈ 616.8798940139915f0 # river cell
 end
 
 @testset "subsurface flow" begin
-    ssf = model.lateral.subsurface.ssf
+    ssf = model.lateral.subsurface.variables.ssf
     @test sum(ssf) ≈ 6.3761585406186976f7
     @test ssf[network.land.order[1]] ≈ 718.2802566393531f0
     @test ssf[network.land.order[end - 100]] ≈ 2337.771227118579f0
@@ -115,35 +115,36 @@ end
 end
 
 @testset "overland flow" begin
-    q = model.lateral.land.q_av
-    @test sum(q) ≈ 291.4923871784623f0
+    q = model.lateral.land.variables.q_av
+    @test sum(q) ≈ 285.59889671823953f0
     @test q[26625] ≈ 0.0
     @test q[39308] ≈ 0.0
     @test q[network.land.order[end]] ≈ 1.0f-30
 end
 
 @testset "river flow" begin
-    q = model.lateral.river.q_av
-    @test sum(q) ≈ 3625.0013368279815f0
-    @test q[1622] ≈ 0.0006503254947860838f0
-    @test q[43] ≈ 12.06416878694095f0
-    @test q[network.river.order[end]] ≈ 0.039200124520463835f0
+    q = model.lateral.river.variables.q_av
+    @test sum(q) ≈ 3848.832553197247f0
+    @test q[1622] ≈ 0.0007520477205381629f0
+    @test q[43] ≈ 11.936893577401598f0
+    @test q[network.river.order[end]] ≈ 0.04418878169426842f0
 end
 
 @testset "reservoir simple" begin
-    res = model.lateral.river.reservoir
-    @test res.outflow[1] ≈ 0.21750000119148086f0
-    @test res.inflow[1] ≈ 43.18479982574888f0
-    @test res.volume[1] ≈ 2.751299001489657f7
-    @test res.precipitation[1] ≈ 0.17999997735023499f0
-    @test res.evaporation[1] ≈ 0.5400000810623169f0
+    res = model.lateral.river.boundary_conditions.reservoir
+    @test res.variables.outflow[1] ≈ 0.21750000119148086f0
+    @test res.boundary_conditions.inflow[1] ≈ 0.00051287944327482
+    @test res.variables.volume[1] ≈ 2.751299001489657f7
+    @test res.boundary_conditions.precipitation[1] ≈ 0.17999997735023499f0
+    @test res.boundary_conditions.evaporation[1] ≈ 0.5400000810623169f0
 end
 
 # set these variables for comparison in "changed dynamic parameters"
 precip = copy(model.vertical.atmospheric_forcing.precipitation)
 evap = copy(model.vertical.atmospheric_forcing.potential_evaporation)
 lai = copy(model.vertical.vegetation_parameter_set.leaf_area_index)
-res_evap = copy(model.lateral.river.reservoir.evaporation)
+res_evap =
+    copy(model.lateral.river.boundary_conditions.reservoir.boundary_conditions.evaporation)
 
 Wflow.close_files(model; delete_output = false)
 
@@ -166,11 +167,11 @@ model = Wflow.run(config)
 end
 
 @testset "river flow at basin outlets and downstream of one pit" begin
-    q = model.lateral.river.q_av
-    @test q[4009] ≈ 8.543145028037452f0 # pit/ outlet, CartesianIndex(141, 228)
+    q = model.lateral.river.variables.q_av
+    @test q[4009] ≈ 8.556311783589425f0 # pit/ outlet, CartesianIndex(141, 228)
     @test q[4020] ≈ 0.006779014715290862f0 # downstream of pit 4009, CartesianIndex(141, 229)
-    @test q[2508] ≈ 150.5640617045796f0 # pit/ outlet
-    @test q[5808] ≈ 0.12438899196040518f0 # pit/ outlet
+    @test q[2508] ≈ 150.5479054707659f0 # pit/ outlet
+    @test q[5808] ≈ 0.12337347858092176f0 # pit/ outlet
 end
 
 # test changing forcing and cyclic LAI parameter
@@ -192,30 +193,33 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "changed dynamic parameters" begin
-    res = model.lateral.river.reservoir
+    res = model.lateral.river.boundary_conditions.reservoir
     vertical = model.vertical
     @test vertical.atmospheric_forcing.precipitation[2] / precip[2] ≈ 2.0f0
     @test (vertical.atmospheric_forcing.potential_evaporation[100] - 1.50) / evap[100] ≈
           3.0f0
     @test vertical.vegetation_parameter_set.leaf_area_index[100] / lai[100] ≈ 1.6f0
-    @test (res.evaporation[2] - 1.50) / res_evap[2] ≈ 3.0000012203408635f0
+    @test (res.boundary_conditions.evaporation[2] - 1.50) / res_evap[2] ≈
+          3.0000012203408635f0
 end
 
 # test cyclic river inflow
 tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 config = Wflow.Config(tomlpath)
 
-config.input.cyclic =
-    ["vertical.vegetation_parameter_set.leaf_area_index", "lateral.river.inflow"]
-config.input.lateral.river.inflow = "inflow"
+config.input.cyclic = [
+    "vertical.vegetation_parameter_set.leaf_area_index",
+    "lateral.river.boundary_conditions.inflow",
+]
+Dict(config.input.lateral.river)["boundary_conditions"] = Dict("inflow" => "inflow")
 
 model = Wflow.initialize_sbm_model(config)
 Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river inflow (cyclic)" begin
-    @test model.lateral.river.inflow[44] ≈ 0.75
-    @test model.lateral.river.q_av[44] ≈ 10.723729440690567f0
+    @test model.lateral.river.boundary_conditions.inflow[44] ≈ 0.75
+    @test model.lateral.river.variables.q_av[44] ≈ 10.554267976107754f0
 end
 
 # test fixed forcing (precipitation = 2.5)
@@ -227,7 +231,12 @@ Wflow.load_fixed_forcing!(model)
 @testset "fixed precipitation forcing (initialize)" begin
     @test maximum(model.vertical.atmospheric_forcing.precipitation) ≈ 2.5
     @test minimum(model.vertical.atmospheric_forcing.precipitation) ≈ 0.0
-    @test all(isapprox.(model.lateral.river.reservoir.precipitation, 2.5))
+    @test all(
+        isapprox.(
+            model.lateral.river.boundary_conditions.reservoir.boundary_conditions.precipitation,
+            2.5,
+        ),
+    )
 end
 
 Wflow.run_timestep!(model)
@@ -235,7 +244,12 @@ Wflow.run_timestep!(model)
 @testset "fixed precipitation forcing (first timestep)" begin
     @test maximum(model.vertical.atmospheric_forcing.precipitation) ≈ 2.5
     @test minimum(model.vertical.atmospheric_forcing.precipitation) ≈ 0.0
-    @test all(isapprox.(model.lateral.river.reservoir.precipitation, 2.5))
+    @test all(
+        isapprox.(
+            model.lateral.river.boundary_conditions.reservoir.boundary_conditions.precipitation,
+            2.5,
+        ),
+    )
 end
 
 Wflow.close_files(model; delete_output = false)
@@ -250,16 +264,16 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river flow and depth (local inertial)" begin
-    q = model.lateral.river.q_av
-    @test sum(q) ≈ 3922.0644366362544f0
-    @test q[1622] ≈ 7.315676375562105f-5
-    @test q[43] ≈ 11.92787156357907f0
-    @test q[501] ≈ 3.57855182713785f0
-    h = model.lateral.river.h_av
-    @test h[1622] ≈ 0.001987887644883981f0
-    @test h[43] ≈ 0.4366415244811759f0
-    @test h[501] ≈ 0.057317706869865745f0
-    q_channel = model.lateral.river.q_channel_av
+    q = model.lateral.river.variables.q_av
+    @test sum(q) ≈ 3854.717278465037f0
+    @test q[1622] ≈ 7.296767063082754f-5
+    @test q[43] ≈ 11.716766734364437f0
+    @test q[501] ≈ 3.4819773071884716f0
+    h = model.lateral.river.variables.h_av
+    @test h[1622] ≈ 0.001986669483044286f0
+    @test h[43] ≈ 0.43311924038778815f0
+    @test h[501] ≈ 0.05635210581824346f0
+    q_channel = model.lateral.river.variables.q_channel_av
     @test q ≈ q_channel
 end
 Wflow.close_files(model; delete_output = false)
@@ -273,20 +287,20 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river and overland flow and depth (local inertial)" begin
-    q = model.lateral.river.q_av
+    q = model.lateral.river.variables.q_av
     @test sum(q) ≈ 2380.64389229669f0
-    @test q[1622] ≈ 7.328535246760549f-5
+    @test q[1622] ≈ 7.30561606758937f-5
     @test q[43] ≈ 5.3566292152594155f0
-    @test q[501] ≈ 1.6042388573126602f0
-    h = model.lateral.river.h_av
-    @test h[1622] ≈ 0.0019891342000364796f0
+    @test q[501] ≈ 1.602564408503896f0
+    h = model.lateral.river.variables.h_av
+    @test h[1622] ≈ 0.001987528017923597f0
     @test h[43] ≈ 0.30026439683630496f0
-    @test h[501] ≈ 0.03195324587192846f0
-    qx = model.lateral.land.qx
-    qy = model.lateral.land.qy
-    @test qx[[26, 35, 631]] ≈ [0.1939736998417174f0, 0.026579954465883678f0, 0.0f0]
-    @test qy[[26, 35, 631]] ≈ [0.12906530420401777f0, 1.7225115950614904f0, 0.0f0]
-    h = model.lateral.land.h
+    @test h[501] ≈ 0.031933708617123746f0
+    qx = model.lateral.land.variables.qx
+    qy = model.lateral.land.variables.qy
+    @test qx[[26, 35, 631]] ≈ [0.18343478752498582f0, 0.000553471702071059f0, 0.0f0]
+    @test qy[[26, 35, 631]] ≈ [0.12607229901243375f0, 0.019605967561619194f0, 0.0f0]
+    h = model.lateral.land.variables.h
     @test h[[26, 35, 631]] ≈
           [0.07367301172613304f0, 0.009139882310161706f0, 0.0007482998926237368f0]
 end
@@ -301,12 +315,12 @@ config.model.floodplain_1d = true
 config.model.river_routing = "local-inertial"
 config.model.land_routing = "kinematic-wave"
 Dict(config.input.lateral.river)["floodplain"] = Dict("volume" => "floodplain_volume")
-Dict(config.state.lateral.river)["floodplain"] =
+Dict(config.state.lateral.river)["floodplain.variables"] =
     Dict("q" => "q_floodplain", "h" => "h_floodplain")
 
 model = Wflow.initialize_sbm_model(config)
 
-fp = model.lateral.river.floodplain.profile
+fp = model.lateral.river.floodplain.parameters.profile
 river = model.lateral.river
 dh = diff(fp.depth)
 Δv = diff(fp.volume[:, 3])
@@ -339,36 +353,36 @@ dh = diff(fp.depth)
         297.8700179533214f0,
         463.35655296229805f0,
     ]
-    @test dh .* fp.width[2:end, 3] * river.dl[3] ≈ Δv
-    @test fp.a[:, 3] * river.dl[3] ≈ fp.volume[:, 3]
+    @test dh .* fp.width[2:end, 3] * river.parameters.flow_length[3] ≈ Δv
+    @test fp.a[:, 3] * river.parameters.flow_length[3] ≈ fp.volume[:, 3]
     # flood depth from flood volume (8000.0)
     flood_vol = 8000.0f0
-    river.volume[3] = flood_vol + river.bankfull_volume[3]
+    river.variables.volume[3] = flood_vol + river.parameters.bankfull_volume[3]
     i1, i2 = Wflow.interpolation_indices(flood_vol, fp.volume[:, 3])
     @test (i1, i2) == (1, 2)
-    flood_depth = Wflow.flood_depth(fp, flood_vol, river.dl[3], 3)
+    flood_depth = Wflow.flood_depth(fp, flood_vol, river.parameters.flow_length[3], 3)
     @test flood_depth ≈ 0.46290938548779076f0
-    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.dl[3] + fp.volume[i1, 3] ≈
-          flood_vol
+    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.parameters.flow_length[3] +
+          fp.volume[i1, 3] ≈ flood_vol
     # flood depth from flood volume (12000.0)
     flood_vol = 12000.0f0
-    river.volume[3] = flood_vol + river.bankfull_volume[3]
+    river.variables.volume[3] = flood_vol + river.parameters.bankfull_volume[3]
     i1, i2 = Wflow.interpolation_indices(flood_vol, fp.volume[:, 3])
     @test (i1, i2) == (2, 3)
-    flood_depth = Wflow.flood_depth(fp, flood_vol, river.dl[3], 3)
+    flood_depth = Wflow.flood_depth(fp, flood_vol, river.parameters.flow_length[3], 3)
     @test flood_depth ≈ 0.6619575699132112f0
-    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.dl[3] + fp.volume[i1, 3] ≈
-          flood_vol
+    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.parameters.flow_length[3] +
+          fp.volume[i1, 3] ≈ flood_vol
     # test extrapolation of segment
     flood_vol = 95000.0f0
-    river.volume[3] = flood_vol + river.bankfull_volume[3]
+    river.variables.volume[3] = flood_vol + river.parameters.bankfull_volume[3]
     i1, i2 = Wflow.interpolation_indices(flood_vol, fp.volume[:, 3])
     @test (i1, i2) == (6, 6)
-    flood_depth = Wflow.flood_depth(fp, flood_vol, river.dl[3], 3)
+    flood_depth = Wflow.flood_depth(fp, flood_vol, river.parameters.flow_length[3], 3)
     @test flood_depth ≈ 2.749036625585836f0
-    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.dl[3] + fp.volume[i1, 3] ≈
-          flood_vol
-    river.volume[3] = 0.0 # reset volume
+    @test (flood_depth - fp.depth[i1]) * fp.width[i2, 3] * river.parameters.flow_length[3] +
+          fp.volume[i1, 3] ≈ flood_vol
+    river.variables.volume[3] = 0.0 # reset volume
     # flow area and wetted perimeter based on hf
     h = 0.5
     i1, i2 = Wflow.interpolation_indices(h, fp.depth)
@@ -404,17 +418,17 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river flow (local inertial) with floodplain schematization simulation" begin
-    q = model.lateral.river.q_av
-    @test sum(q) ≈ 3910.4728717811836f0
-    @test q[1622] ≈ 7.315676384849305f-5
-    @test q[43] ≈ 11.92787156357908f0
-    @test q[501] ≈ 3.510668846752431f0
-    @test q[5808] ≈ 0.002223993845806248f0
-    h = model.lateral.river.h_av
-    @test h[1622] ≈ 0.001987887580593841f0
-    @test h[43] ≈ 0.436641524481545f0
-    @test h[501] ≈ 0.05670770509802258f0
-    @test h[5808] ≈ 0.005929945681367346f0
+    q = model.lateral.river.variables.q_av
+    @test sum(q) ≈ 3843.944494991296f0
+    @test q[1622] ≈ 7.296767071929629f-5
+    @test q[43] ≈ 11.716766734364418f0
+    @test q[501] ≈ 3.424364314225289f0
+    @test q[5808] ≈ 0.002228981516146531f0
+    h = model.lateral.river.variables.h_av
+    @test h[1622] ≈ 0.0019866694251020806f0
+    @test h[43] ≈ 0.433119240388070f0
+    @test h[501] ≈ 0.055832770820860404f0
+    @test h[5808] ≈ 0.005935591961908253f0
 end
 
 # set boundary condition local inertial routing from netCDF file
@@ -425,16 +439,16 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "change boundary condition for local inertial routing (including floodplain)" begin
-    q = model.lateral.river.q_av
-    @test sum(q) ≈ 3910.683449719468f0
-    @test q[1622] ≈ 7.315757521099307f-5
-    @test q[43] ≈ 11.927871563591228f0
-    @test q[501] ≈ 3.5106678593721496f0
-    @test q[5808] ≈ 0.060518234525259465f0
-    h = model.lateral.river.h_av
-    @test h[1622] ≈ 0.0019878952928530183f0
-    @test h[43] ≈ 0.4366415249636809f0
-    @test h[501] ≈ 0.056707564314724804f0
+    q = model.lateral.river.variables.q_av
+    @test sum(q) ≈ 3844.1544889903134f0
+    @test q[1622] ≈ 7.296767071929629f-5
+    @test q[43] ≈ 11.716766734416717f0
+    @test q[501] ≈ 3.424329413571391f0
+    @test q[5808] ≈ 0.055269620065756024f0
+    h = model.lateral.river.variables.h_av
+    @test h[1622] ≈ 0.0019866694251020806f0
+    @test h[43] ≈ 0.4331192403230577f0
+    @test h[501] ≈ 0.0558281185092927f0
     @test h[5808] ≈ 2.0000006940603936f0
 end
 Wflow.close_files(model; delete_output = false)
@@ -457,8 +471,8 @@ Wflow.close_files(model; delete_output = false)
         kvfrac = soil.parameters.kvfrac
         kv_z = Wflow.hydraulic_conductivity_at_depth(kv_profile, kvfrac, z, i, 2)
         @test kv_z ≈ kvfrac[i][2] * kv_profile.kv_0[i] * exp(-kv_profile.f[i] * z)
-        @test subsurface.ssfmax[i] ≈ 28.32720603576582f0
-        @test subsurface.ssf[i] ≈ 11683.330684556406f0
+        @test subsurface.variables.ssfmax[i] ≈ 28.32720603576582f0
+        @test subsurface.variables.ssf[i] ≈ 11683.330684556406f0
     end
 
     @testset "exponential constant profile" begin
@@ -478,8 +492,8 @@ Wflow.close_files(model; delete_output = false)
         kv_1000 = Wflow.hydraulic_conductivity_at_depth(kv_profile, kvfrac, 1000.0, i, 3)
         @test kv_400 ≈ kv_1000
         @test all(kv_profile.z_exp .== 400.0)
-        @test subsurface.ssfmax[i] ≈ 49.38558575188426f0
-        @test subsurface.ssf[i] ≈ 24810.460986497365f0
+        @test subsurface.variables.ssfmax[i] ≈ 49.38558575188426f0
+        @test subsurface.variables.ssf[i] ≈ 24810.460986497365f0
     end
 
     @testset "layered profile" begin
@@ -493,9 +507,9 @@ Wflow.close_files(model; delete_output = false)
         @test Wflow.hydraulic_conductivity_at_depth(kv_profile, kvfrac, z, i, 2) ≈
               kv_profile.kv[i][2]
         Wflow.kh_layered_profile!(soil, subsurface, kv_profile, 86400.0)
-        @test subsurface.kh_profile.kh[i] ≈ 47.508932674632355f0
-        @test subsurface.ssfmax[i] ≈ 30.237094380100316f0
-        @test subsurface.ssf[i] ≈ 14546.518932613191f0
+        @test subsurface.parameters.kh_profile.kh[i] ≈ 47.508932674632355f0
+        @test subsurface.variables.ssfmax[i] ≈ 30.237094380100316f0
+        @test subsurface.variables.ssf[i] ≈ 14546.518932613191f0
     end
 
     @testset "layered exponential profile" begin
@@ -510,20 +524,20 @@ Wflow.close_files(model; delete_output = false)
               kv_profile.kv[i][2]
         @test kv_profile.nlayers_kv[i] == 2
         Wflow.kh_layered_profile!(soil, subsurface, kv_profile, 86400.0)
-        @test subsurface.kh_profile.kh[i] ≈ 33.76026208801769f0
+        @test subsurface.parameters.kh_profile.kh[i] ≈ 33.76026208801769f0
         @test all(kv_profile.z_layered[1:10] .== 400.0)
-        @test subsurface.ssfmax[i] ≈ 23.4840490395906f0
-        @test subsurface.ssf[i] ≈ 10336.88327617503f0
+        @test subsurface.variables.ssfmax[i] ≈ 23.4840490395906f0
+        @test subsurface.variables.ssf[i] ≈ 10336.88327617503f0
     end
 
     @testset "river flow layered exponential profile" begin
         model = Wflow.initialize_sbm_model(config)
         Wflow.run_timestep!(model)
         Wflow.run_timestep!(model)
-        q = model.lateral.river.q_av
-        @test sum(q) ≈ 3159.38300016008f0
-        @test q[1622] ≈ 0.0005972577112819149f0
-        @test q[43] ≈ 10.017642376280731f0
+        q = model.lateral.river.variables.q_av
+        @test sum(q) ≈ 3302.1390089922525f0
+        @test q[1622] ≈ 0.0006990391393246408f0
+        @test q[43] ≈ 9.673592182882691f0
     end
 
     Wflow.close_files(model; delete_output = false)
diff --git a/test/run_sbm_gwf.jl b/test/run_sbm_gwf.jl
index 60b77a8de..81038ed5a 100644
--- a/test/run_sbm_gwf.jl
+++ b/test/run_sbm_gwf.jl
@@ -13,7 +13,7 @@ flush(model.writer.csv_io)  # ensure the buffer is written fully to disk
     row = csv_first_row(model.writer.csv_path)
 
     @test row.time == DateTime("2000-06-01T00:00:00")
-    @test row.Q_av ≈ 0.01620324716944374f0
+    @test row.Q_av ≈ 0.01619703129434486f0
     @test row.head ≈ 1.6471323360175287f0
 end
 
@@ -39,34 +39,34 @@ Wflow.run_timestep!(model)
 end
 
 @testset "overland flow (kinematic wave)" begin
-    q = model.lateral.land.q_av
-    @test sum(q) ≈ 2.229860508650628f-7
+    q = model.lateral.land.variables.q_av
+    @test sum(q) ≈ 2.2321111203610908f-7
 end
 
 @testset "river domain (kinematic wave)" begin
-    q = model.lateral.river.q_av
+    q = model.lateral.river.variables.q_av
     river = model.lateral.river
-    @test sum(q) ≈ 0.035443370536496675f0
-    @test q[6] ≈ 0.008031554512314907f0
-    @test river.volume[6] ≈ 4.532124903256408f0
-    @test river.inwater[6] ≈ 0.0004073892212290558f0
+    @test sum(q) ≈ 0.035468154534622556f0
+    @test q[6] ≈ 0.00803825101724232f0
+    @test river.variables.volume[6] ≈ 4.532124903256408f0
+    @test river.boundary_conditions.inwater[6] ≈ 0.00040826084140456616f0
     @test q[13] ≈ 0.0006017024138583771f0
-    @test q[network.river.order[end]] ≈ 0.008559590281509943f0
+    @test q[network.river.order[end]] ≈ 0.00856866488665273f0
 end
 
 @testset "groundwater" begin
     gw = model.lateral.subsurface
-    @test gw.river.stage[1] ≈ 1.2123636929067039f0
-    @test gw.flow.aquifer.head[17:21] ≈ [
+    @test gw.river.variables.stage[1] ≈ 1.2123636929067039f0
+    @test gw.flow.aquifer.variables.head[17:21] ≈ [
         1.2866380350225155f0,
         1.3477853512604643f0,
         1.7999999523162842f0,
         1.6225103807809076f0,
         1.4053590307668113f0,
     ]
-    @test gw.river.flux[1] ≈ -51.34674583702381f0
-    @test gw.drain.flux[1] ≈ 0.0
-    @test gw.recharge.rate[19] ≈ -0.0014241196552847502f0
+    @test gw.river.variables.flux[1] ≈ -51.34674583702381f0
+    @test gw.drain.variables.flux[1] ≈ 0.0
+    @test gw.recharge.variables.rate[19] ≈ -0.0014241196552847502f0
 end
 
 @testset "no drains" begin
@@ -91,12 +91,12 @@ Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river domain (local inertial)" begin
-    q = model.lateral.river.q_av
+    q = model.lateral.river.variables.q_av
     river = model.lateral.river
     @test sum(q) ≈ 0.02727911500112358f0
     @test q[6] ≈ 0.006111263175002127f0
-    @test river.volume[6] ≈ 7.6120096530771075f0
-    @test river.inwater[6] ≈ 0.00022087679662860144f0
+    @test river.variables.volume[6] ≈ 7.6120096530771075f0
+    @test river.boundary_conditions.inwater[6] ≈ 0.0002210785332342944f0
     @test q[13] ≈ 0.0004638698607639214f0
     @test q[5] ≈ 0.0064668491697542786f0
 end
@@ -112,27 +112,27 @@ config.input.lateral.river.bankfull_elevation = "bankfull_elevation"
 config.input.lateral.river.bankfull_depth = "bankfull_depth"
 config.input.lateral.land.elevation = "wflow_dem"
 
-pop!(Dict(config.state.lateral.land), "q")
-config.state.lateral.land.h_av = "h_av_land"
-config.state.lateral.land.qx = "qx_land"
-config.state.lateral.land.qy = "qy_land"
+pop!(Dict(config.state.lateral.land.variables), "q")
+config.state.lateral.land.variables.h_av = "h_av_land"
+config.state.lateral.land.variables.qx = "qx_land"
+config.state.lateral.land.variables.qy = "qy_land"
 
 model = Wflow.initialize_sbm_gwf_model(config)
 Wflow.run_timestep!(model)
 Wflow.run_timestep!(model)
 
 @testset "river and land domain (local inertial)" begin
-    q = model.lateral.river.q_av
+    q = model.lateral.river.variables.q_av
     @test sum(q) ≈ 0.027286431923384962f0
     @test q[6] ≈ 0.00611309161099138f0
     @test q[13] ≈ 0.0004639786629631376f0
     @test q[5] ≈ 0.006468859889145798f0
-    h = model.lateral.river.h_av
+    h = model.lateral.river.variables.h_av
     @test h[6] ≈ 0.08120137914886108f0
     @test h[5] ≈ 0.07854966203902745f0
     @test h[13] ≈ 0.08323543174453409f0
-    qx = model.lateral.land.qx
-    qy = model.lateral.land.qy
+    qx = model.lateral.land.variables.qx
+    qy = model.lateral.land.variables.qy
     @test all(qx .== 0.0f0)
     @test all(qy .== 0.0f0)
 end
@@ -157,50 +157,34 @@ Wflow.run_timestep!(model)
 end
 
 @testset "overland flow warm start (kinematic wave)" begin
-    q = model.lateral.land.q_av
-    @test sum(q) ≈ 1.4589771292158736f-5
+    q = model.lateral.land.variables.q_av
+    @test sum(q) ≈ 1.4224503548471601f-5
 end
 
 @testset "river domain warm start (kinematic wave)" begin
-    q = model.lateral.river.q_av
+    q = model.lateral.river.variables.q_av
     river = model.lateral.river
     @test sum(q) ≈ 0.01191742350356312f0
     @test q[6] ≈ 0.0024353072305122064f0
-    @test river.volume[6] ≈ 2.2277585577366357f0
-    @test river.inwater[6] ≈ -1.3019072795599315f-5
+    @test river.variables.volume[6] ≈ 2.2277585577366357f0
+    @test river.boundary_conditions.inwater[6] ≈ -1.3042629584651168f-5
     @test q[13] ≈ 7.332742814063803f-5
     @test q[network.river.order[end]] ≈ 0.002472526149620472f0
 end
 
 @testset "groundwater warm start" begin
     gw = model.lateral.subsurface
-    @test gw.river.stage[1] ≈ 1.2031171676781156f0
-    @test gw.flow.aquifer.head[17:21] ≈ [
+    @test gw.river.variables.stage[1] ≈ 1.2031171676781156f0
+    @test gw.flow.aquifer.variables.head[17:21] ≈ [
         1.2277456867225283f0,
         1.286902494792006f0,
         1.7999999523162842f0,
         1.5901747932190804f0,
         1.2094238817776854f0,
     ]
-    @test gw.river.flux[1] ≈ -6.692884222603261f0
-    @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) == true
-    @test Wflow.exchange(aquifer.k) == true
-    @test Wflow.grid_loc(aquifer, :head) == "node"
-    @test Wflow.grid_loc(aquifer, :k) == "node"
-    recharge = model.lateral.subsurface.recharge
-    @test Wflow.exchange(recharge.rate) == true
-    @test Wflow.exchange(recharge.flux) == true
-    @test Wflow.grid_loc(recharge, :rate) == "node"
-    @test Wflow.grid_loc(recharge, :flux) == "node"
-    constanthead = model.lateral.subsurface.flow.constanthead
-    @test Wflow.exchange(constanthead) == false
-    @test Wflow.grid_loc(constanthead, :head) == "node"
+    @test gw.river.variables.flux[1] ≈ -6.692884222603261f0
+    @test gw.drain.variables.flux[1] ≈ 0.0
+    @test gw.recharge.variables.rate[19] ≈ -0.0014241196552847502f0
 end
 
 Wflow.close_files(model; delete_output = false)
diff --git a/test/run_sbm_piave.jl b/test/run_sbm_piave.jl
index a409cad08..ac666bd0a 100644
--- a/test/run_sbm_piave.jl
+++ b/test/run_sbm_piave.jl
@@ -5,9 +5,9 @@ function run_piave(model, steps)
     riv_vol = zeros(steps)
     for i in 1:steps
         Wflow.run_timestep!(model)
-        ssf_vol[i] = mean(model.lateral.subsurface.volume)
-        riv_vol[i] = mean(model.lateral.river.volume)
-        q[i] = model.lateral.river.q_av[1]
+        ssf_vol[i] = mean(model.lateral.subsurface.variables.volume)
+        riv_vol[i] = mean(model.lateral.river.variables.volume)
+        q[i] = model.lateral.river.variables.q_av[1]
     end
     return q, riv_vol, ssf_vol
 end
@@ -27,52 +27,52 @@ Wflow.close_files(model; delete_output = false)
 @testset "piave with and without water demand" begin
     idx = 1:3:28
     @test q_demand[idx] ≈ [
-        218.52770747627918f0,
-        193.02890386240665f0,
-        271.68778616960685f0,
-        161.80734173386602f0,
-        164.81279958487485f0,
-        150.4149716788231f0,
-        121.02659706677429f0,
-        108.15426854851842f0,
-        174.63022487569546f0,
-        108.20883498755789f0,
+        218.52013823809472f0,
+        193.0134951603773f0,
+        272.4111837647947f0,
+        161.88264628787172f0,
+        164.8199089671644f0,
+        150.2681168314876f0,
+        121.20070337007452f0,
+        108.10106381132412f0,
+        175.13799714754256f0,
+        108.26190463186364f0,
     ]
     @test q_[idx] ≈ [
-        219.9042687883228f0,
-        197.73238933696658f0,
-        276.99163278840734f0,
-        165.74244080863346f0,
-        172.99856395134805f0,
-        153.91963517651158f0,
-        128.52653903788593f0,
-        112.02923578000491f0,
-        178.19599851038708f0,
-        109.99414262238396f0,
+        219.87655632704903f0,
+        197.7038754807009f0,
+        277.7110869134211f0,
+        165.79913971520423f0,
+        173.04466296857905f0,
+        153.84187794694486f0,
+        128.71609293239374f0,
+        112.02394903669563f0,
+        178.8207608992179f0,
+        110.0540286256144f0,
     ]
     @test riv_vol_demand[idx] ≈ [
-        60125.790043761706f0,
-        55476.580177102805f0,
-        62195.90597660078f0,
-        48963.00368970478f0,
-        51871.46897847274f0,
-        43794.737044826295f0,
-        41710.964839545784f0,
-        39623.67826731185f0,
-        44719.74807290461f0,
-        39494.768128540185f0,
+        60108.85346812383f0,
+        55483.137044195726f0,
+        62098.68355844664f0,
+        48989.58100040463f0,
+        51879.53087453071f0,
+        43784.59321641879f0,
+        41700.35166194379f0,
+        39630.55231098758f0,
+        44664.51932077705f0,
+        39490.08957716613f0,
     ]
     @test riv_vol[idx] ≈ [
-        60620.42943610538f0,
-        55927.61042067993f0,
-        62448.21099253601f0,
-        49283.04780863544f0,
-        52293.53301649927f0,
-        44029.56815119591f0,
-        41996.885591977836f0,
-        39751.39389784692f0,
-        44619.74874948371f0,
-        39285.722489743566f0,
+        60612.8152223446f0,
+        55934.74021061718f0,
+        62385.538865370385f0,
+        49305.35863834837f0,
+        52307.51338053202f0,
+        44027.55259408914f0,
+        41998.69770760942f0,
+        39768.94289967264f0,
+        44586.63915460806f0,
+        39296.07358095679f0,
     ]
     @test ssf_vol_demand[idx] ≈ [
         244149.81771426898f0,
diff --git a/test/runtests.jl b/test/runtests.jl
index 6981af08d..f768e4610 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -78,7 +78,7 @@ include("testing_utils.jl")
 with_logger(NullLogger()) do
     ## run all tests
     @testset "Wflow.jl" begin
-        include("horizontal_process.jl")
+        include("routing_process.jl")
         include("io.jl")
         include("vertical_process.jl")
         include("reservoir_lake.jl")
diff --git a/test/sbm_config.toml b/test/sbm_config.toml
index 39898a488..9e805a61f 100644
--- a/test/sbm_config.toml
+++ b/test/sbm_config.toml
@@ -31,18 +31,18 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 h = "h_river"
 h_av = "h_av_river"
 q = "q_river"
 
-[state.lateral.river.reservoir]
+[state.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 h = "h_land"
 h_av = "h_av_land"
 q = "q_land"
@@ -112,7 +112,7 @@ offset = 0.0
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 bankfull_elevation = "RiverZ"
@@ -132,7 +132,7 @@ targetminfrac = "ResTargetMinFrac"
 ksathorfrac = "KsatHorFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [model]
@@ -161,17 +161,17 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 h = "h_river"
 q = "q_river"
 
-[output.lateral.river.reservoir]
+[output.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[output.lateral.subsurface]
+[output.lateral.subsurface.variables]
 ssf = "ssf"
 
-[output.lateral.land]
+[output.lateral.land.variables]
 h = "h_land"
 q = "q_land"
 
@@ -181,7 +181,7 @@ path = "output_scalar_moselle.nc"
 [[netcdf.variable]]
 name = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[netcdf.variable]]
 coordinate.x = 6.255
@@ -202,13 +202,13 @@ path = "output_moselle.csv"
 
 [[csv.column]]
 header = "Q"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 reducer = "maximum"
 
 [[csv.column]]
 header = "volume"
 index = 1
-parameter = "lateral.river.reservoir.volume"
+parameter = "lateral.river.boundary_conditions.reservoir.variables.volume"
 
 [[csv.column]]
 coordinate.x = 6.255
@@ -232,7 +232,7 @@ parameter = "vertical.atmospheric_forcing.temperature"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[csv.column]]
 header = "recharge"
@@ -255,8 +255,19 @@ components = [
   "vertical.snow.boundary_conditions",
   "vertical.snow.variables",
   "vertical.snow.parameters",
-  "lateral.subsurface",
-  "lateral.land",
-  "lateral.river",
-  "lateral.river.reservoir",
+  "lateral.subsurface.boundary_conditions",
+  "lateral.subsurface.variables",
+  "lateral.subsurface.parameters",
+  "lateral.subsurface.parameters.kh_profile",
+  "lateral.land.boundary_conditions",
+  "lateral.land.variables",
+  "lateral.land.variables.flow",
+  "lateral.land.parameters",
+  "lateral.river.variables",
+  "lateral.river.parameters",
+  "lateral.river.parameters.flow",
+  "lateral.river.boundary_conditions",
+  "lateral.river.boundary_conditions.reservoir.boundary_conditions",
+  "lateral.river.boundary_conditions.reservoir.parameters",
+  "lateral.river.boundary_conditions.reservoir.variables",
 ]
diff --git a/test/sbm_gw.toml b/test/sbm_gw.toml
index a18965254..fe5a00757 100644
--- a/test/sbm_gw.toml
+++ b/test/sbm_gw.toml
@@ -30,15 +30,15 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 h = "h_river"
 h_av = "h_av_river"
 q = "q_river"
 
-[state.lateral.river.reservoir]
+[state.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 h = "h_land"
 h_av = "h_av_land"
 q = "q_land"
@@ -102,7 +102,7 @@ ttm = "TTM"
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 
@@ -117,7 +117,7 @@ targetfullfrac = "ResTargetFullFrac"
 targetminfrac = "ResTargetMinFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [model]
@@ -144,14 +144,14 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 h = "h_river"
 q = "q_river"
 
-[output.lateral.river.reservoir]
+[output.lateral.river.boundary_conditions.reservoir.variables]
 volume = "volume_reservoir"
 
-[output.lateral.land]
+[output.lateral.land.variables]
 h = "h_land"
 q = "q_land"
 
@@ -161,7 +161,7 @@ path = "output_scalar_moselle.nc"
 [[netcdf.variable]]
 name = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[netcdf.variable]]
 coordinate.x = 6.255
@@ -182,13 +182,13 @@ path = "output_moselle.csv"
 
 [[csv.column]]
 header = "Q"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 reducer = "maximum"
 
 [[csv.column]]
 header = "volume"
 index = 1
-parameter = "lateral.river.reservoir.volume"
+parameter = "lateral.river.boundary_conditions.reservoir.variables.volume"
 
 [[csv.column]]
 coordinate.x = 6.255
@@ -205,7 +205,7 @@ parameter = "vertical.atmospheric_forcing.temperature"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[csv.column]]
 header = "recharge"
@@ -215,9 +215,15 @@ reducer = "mean"
 
 [API]
 components = [
-  "vertical",
-  "lateral.subsurface",
-  "lateral.land",
-  "lateral.river",
-  "lateral.river.reservoir",
+  "vertical.variables",
+  "vertical.parameters",
+  "lateral.subsurface.variables",
+  "lateral.subsurface.parameters",
+  "lateral.land.variables",
+  "lateral.land.parameters",
+  "lateral.river.variables",
+  "lateral.river.parameters",
+  "lateral.river.boundary_conditions.reservoir.boundary_conditions",
+  "lateral.river.boundary_conditions.reservoir.parameters",
+  "lateral.river.boundary_conditions.reservoir.variables",
 ]
diff --git a/test/sbm_gwf_config.toml b/test/sbm_gwf_config.toml
index f08431236..c917b409f 100644
--- a/test/sbm_gwf_config.toml
+++ b/test/sbm_gwf_config.toml
@@ -25,17 +25,17 @@ canopy_storage = "canopystorage"
 satwaterdepth = "satwaterdepth"
 ustorelayerdepth = "ustorelayerdepth"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 h = "h_river"
 h_av = "h_av_river"
 q = "q_river"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 h = "h_land"
 h_av = "h_av_land"
 q = "q_land"
 
-[state.lateral.subsurface.flow.aquifer]
+[state.lateral.subsurface.flow.aquifer.variables]
 head = "head"
 
 [input]
@@ -77,12 +77,12 @@ theta_s = "thetaS"
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_river"
+mannings_n = "N_river"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [input.lateral.subsurface]
@@ -119,19 +119,19 @@ ustorelayerdepth = "ustorelayerdepth"
 [output.vertical.soil.parameters]
 soilthickness = "soilthickness"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 q = "q"
 
-[output.lateral.subsurface.flow.aquifer]
+[output.lateral.subsurface.flow.aquifer.variables]
 head = "head"
 
-[output.lateral.subsurface.recharge]
+[output.lateral.subsurface.recharge.variables]
 rate = "rate"
 
-[output.lateral.subsurface.drain]
+[output.lateral.subsurface.drain.variables]
 flux = "drain_flux"
 
-[output.lateral.subsurface.river]
+[output.lateral.subsurface.river.variables]
 flux = "flux"
 
 [csv]
@@ -140,9 +140,9 @@ path = "output_example-sbm-gwf.csv"
 [[csv.column]]
 header = "Q_av"
 index = 5
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"
 
 [[csv.column]]
 header = "head"
 index = 5
-parameter = "lateral.subsurface.flow.aquifer.head"
+parameter = "lateral.subsurface.flow.aquifer.variables.head"
diff --git a/test/sbm_gwf_piave_demand_config.toml b/test/sbm_gwf_piave_demand_config.toml
index cf61cc705..a2a0688ec 100644
--- a/test/sbm_gwf_piave_demand_config.toml
+++ b/test/sbm_gwf_piave_demand_config.toml
@@ -71,7 +71,7 @@ glacier_store = "glacierstore"
 [state.vertical.demand.paddy.variables]
 h = "h_paddy"
 
-[state.lateral.subsurface.flow.aquifer]
+[state.lateral.subsurface.flow.aquifer.variables]
 head = "head"
 
 [input.vertical.glacier.parameters]
@@ -136,15 +136,15 @@ livestock = true
 paddy = true
 nonpaddy = true
 
-[state.lateral.river]
+[state.lateral.river.variables]
 q = "q_river"
 h = "h_river"
 h_av = "h_av_river"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 q = "q_land"
 h = "h_land"
 h_av = "h_av_land"
@@ -191,19 +191,19 @@ specific_yield = "specific_yield"
 gwf_f = "gwf_f"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [output]
 path = "output-piave-gwf.nc"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 q_av = "q_river"
 
 [output.vertical.soil.variables]
 zi = "zi"
 
-[output.lateral.subsurface.flow.aquifer]
+[output.lateral.subsurface.flow.aquifer.variables]
 head = "head"
 
 [csv]
@@ -212,9 +212,9 @@ path = "output.csv"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"
 
 [[csv.column]]
 header = "Q"
 map = "gauges_grdc"
-parameter = "lateral.river.q_av"
\ No newline at end of file
+parameter = "lateral.river.variables.q_av"
\ No newline at end of file
diff --git a/test/sbm_piave_config.toml b/test/sbm_piave_config.toml
index 8f385c4ff..026ab9b7a 100644
--- a/test/sbm_piave_config.toml
+++ b/test/sbm_piave_config.toml
@@ -111,22 +111,22 @@ tti = "TTI"
 ttm = "TTM"
 cfmax = "Cfmax"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 q = "q_river"
 h = "h_river"
 h_av = "h_av_river"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 q = "q_land"
 h = "h_land"
 h_av = "h_av_land"
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 bankfull_depth = "RiverDepth"
@@ -135,13 +135,13 @@ bankfull_depth = "RiverDepth"
 ksathorfrac = "KsatHorFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [output]
 path = "output-piave.nc"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 q_av = "q_river"
 
 [csv]
@@ -150,9 +150,9 @@ path = "output-piave.csv"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"
 
 [[csv.column]]
 header = "Q"
 map = "gauges_grdc"
-parameter = "lateral.river.q_av"
\ No newline at end of file
+parameter = "lateral.river.variables.q_av"
\ No newline at end of file
diff --git a/test/sbm_piave_demand_config.toml b/test/sbm_piave_demand_config.toml
index 17f5828d6..ffc012efb 100644
--- a/test/sbm_piave_demand_config.toml
+++ b/test/sbm_piave_demand_config.toml
@@ -131,15 +131,15 @@ livestock = true
 paddy = true
 nonpaddy = true
 
-[state.lateral.river]
+[state.lateral.river.variables]
 q = "q_river"
 h = "h_river"
 h_av = "h_av_river"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 q = "q_land"
 h = "h_land"
 h_av = "h_av_land"
@@ -170,7 +170,7 @@ irrigation_trigger = "irrigation_trigger"
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 bankfull_depth = "RiverDepth"
@@ -179,13 +179,13 @@ bankfull_depth = "RiverDepth"
 ksathorfrac = "KsatHorFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 
 [output]
 path = "output-piave-demand.nc"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 q_av = "q_river"
 
 [output.vertical.soil.variables]
@@ -197,12 +197,12 @@ path = "output-piave-demand.csv"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"
 
 [[csv.column]]
 header = "Q"
 map = "gauges_grdc"
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"
 
 [[csv.column]]
 coordinate.x = 12.7243
diff --git a/test/sbm_simple.toml b/test/sbm_simple.toml
index b6e16d6d1..d2a6f10cd 100644
--- a/test/sbm_simple.toml
+++ b/test/sbm_simple.toml
@@ -99,7 +99,7 @@ path = "output_moselle_simple.csv"
 coordinate.x = 7.378
 coordinate.y = 50.204
 header = "Q"
-parameter = "lateral.river.q"
+parameter = "lateral.river.variables.q"
 
 [[csv.column]]
 header = "recharge"
diff --git a/test/sbm_swf_config.toml b/test/sbm_swf_config.toml
index 79ec88b3c..0ebbbf872 100644
--- a/test/sbm_swf_config.toml
+++ b/test/sbm_swf_config.toml
@@ -26,7 +26,7 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[state.lateral.river]
+[state.lateral.river.variables]
 h = "h_river"
 h_av = "h_av_river"
 q = "q_river"
@@ -34,10 +34,10 @@ q = "q_river"
 [state.lateral.river.reservoir]
 volume = "volume_reservoir"
 
-[state.lateral.subsurface]
+[state.lateral.subsurface.variables]
 ssf = "ssf"
 
-[state.lateral.land]
+[state.lateral.land.variables]
 h = "h_land"
 h_av = "h_av_land"
 qx = "qx_land"
@@ -103,7 +103,7 @@ cfmax = "Cfmax"
 
 [input.lateral.river]
 length = "wflow_riverlength"
-n = "N_River"
+mannings_n = "N_River"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
 bankfull_elevation = "RiverZ"
@@ -118,11 +118,12 @@ maxrelease = "ResMaxRelease"
 maxvolume = "ResMaxVolume"
 targetfullfrac = "ResTargetFullFrac"
 targetminfrac = "ResTargetMinFrac"
+
 [input.lateral.subsurface]
 ksathorfrac = "KsatHorFrac"
 
 [input.lateral.land]
-n = "N"
+mannings_n = "N"
 slope = "Slope"
 elevation = "FloodplainZ"
 
@@ -152,16 +153,16 @@ ustorelayerdepth = "ustorelayerdepth"
 snow_storage = "snow"
 snow_water = "snowwater"
 
-[output.lateral.river]
+[output.lateral.river.variables]
 h = "h_river"
 h_av = "hav_river"
 q = "q_river"
 q_av = "qav_river"
 
-[output.lateral.subsurface]
+[output.lateral.subsurface.variables]
 ssf = "ssf"
 
-[output.lateral.land]
+[output.lateral.land.variables]
 h = "h_land"
 qx = "qx_land"
 qy = "qy_land"
@@ -172,5 +173,5 @@ path = "output_moselle_swf.csv"
 [[csv.column]]
 header = "Q"
 map = "gauges"
-parameter = "lateral.river.q_av"
+parameter = "lateral.river.variables.q_av"