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

Issue:
Refactor modeltype `sediment`
### Feature type

Refactor the sediment model, i.e., re-organize the model's code without adding new or removing existing features.

### Improvement Description

The initialization of the sediment model takes place in one function (`initialize_sediment_model`) of over 140 lines of code.
This is poorly maintainable and should be split up into smaller individual logical steps.
In addition, the parameters used for initializing the sediment model are, currently, stored alongside the parameters of all other models / concepts in wflow.
Sediment model parameters should be extracted into its own parameter `Struct`.
The extraction of the sediment model parameters should include the re-coupling of sediment model parameter to other models that require access to sediment model parameters.

see also main issue https://github.com/Deltares/Wflow.jl/issues/351

### Implementation Description

_No response_

### Additional Context

_No response_

### Tasks

- [ ] #425 
- [ ] #451  
- [ ] #453  
- [ ] Store sediment model parameters in a separate `Struct`
- [ ] Refactor the sediment model to use its own parameter `Struct` instead of the `Model` struct
- [ ] Remove sediment model parameters from the `Model` struct
- [ ] Fix all sediment model parameter access invalidations


Discussion / hints:

Fix patch:
diff --git a/docs/changelog.qmd b/docs/changelog.qmd
index 44d4798f7..6a138394e 100644
--- a/docs/changelog.qmd
+++ b/docs/changelog.qmd
@@ -38,6 +38,11 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   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).
+- Refactor the `sediment` concept: similar to the refactor of `SBM` concept, the sediment
+  model has been split into three main structures. `SoilLoss` was split into rainfall_erosion,
+  overland_flow_erosion and (total) soil_erosion concepts. `OverlandFlowSediment` has been
+  split into transport_capacity, sediment_flux and to_river concepts. `RiverSediment` has been
+  split into transport_capacity, potential_erosion, sediment_flux and concentrations concepts.
 - 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
diff --git a/src/Wflow.jl b/src/Wflow.jl
index ff4abfaab..2f19e7585 100644
--- a/src/Wflow.jl
+++ b/src/Wflow.jl
@@ -163,8 +163,20 @@ include("soil/soil.jl")
 include("soil/soil_process.jl")
 include("sbm.jl")
 include("demand/water_demand.jl")
-include("sediment.jl")
 include("sbm_model.jl")
+include("sediment/erosion/erosion_process.jl")
+include("sediment/erosion/rainfall_erosion.jl")
+include("sediment/erosion/overland_flow_erosion.jl")
+include("sediment/erosion/soil_erosion.jl")
+include("sediment/erosion/river_erosion.jl")
+include("sediment/sediment_transport/deposition.jl")
+include("sediment/sediment_transport/transport_capacity_process.jl")
+include("sediment/sediment_transport/transport_capacity.jl")
+include("sediment/sediment_transport/overland_flow_transport.jl")
+include("sediment/sediment_transport/land_to_river.jl")
+include("sediment/sediment_transport/river_transport.jl")
+include("erosion.jl")
+include("sediment_flux.jl")
 include("sediment_model.jl")
 include("sbm_gwf_model.jl")
 include("utils.jl")
diff --git a/src/erosion.jl b/src/erosion.jl
new file mode 100644
index 000000000..187e91ac4
--- /dev/null
+++ b/src/erosion.jl
@@ -0,0 +1,81 @@
+"Soil loss model"
+@with_kw struct SoilLoss{RE, OFE, SE, T}
+    atmospheric_forcing::AtmosphericForcing{T}
+    hydrological_forcing::HydrologicalForcing{T}
+    geometry::LandGeometry{T}
+    rainfall_erosion::RE
+    overland_flow_erosion::OFE
+    soil_erosion::SE
+end
+
+"Initialize soil loss model"
+function SoilLoss(dataset, config, indices)
+    n = length(indices)
+
+    atmospheric_forcing = AtmosphericForcing(n)
+    hydrological_forcing = HydrologicalForcing(n)
+    geometry = LandGeometry(dataset, config, indices)
+
+    # Rainfall erosion
+    rainfallerosionmodel = get(config.model, "rainfall_erosion", "answers")::String
+    if rainfallerosionmodel == "answers"
+        rainfall_erosion_model = RainfallErosionAnswersModel(dataset, config, indices)
+    elseif rainfallerosionmodel == "eurosem"
+        rainfall_erosion_model = RainfallErosionEurosemModel(dataset, config, indices)
+    else
+        error("Unknown rainfall erosion model: $rainfallerosionmodel")
+    end
+
+    # Overland flow erosion
+    overlandflowerosionmodel = get(config.model, "overland_flow_erosion", "answers")::String
+    if overlandflowerosionmodel == "answers"
+        overland_flow_erosion_model =
+            OverlandFlowErosionAnswersModel(dataset, config, indices)
+    else
+        error("Unknown overland flow erosion model: $overlandflowerosionmodel")
+        # overland_flow_erosion_model = NoOverlandFlowErosionModel()
+    end
+
+    # Total soil erosion and particle differentiation
+    soil_erosion_model = SoilErosionModel(dataset, config, indices)
+
+    soil_loss = SoilLoss{
+        typeof(rainfall_erosion_model),
+        typeof(overland_flow_erosion_model),
+        typeof(soil_erosion_model),
+        Float,
+    }(;
+        atmospheric_forcing = atmospheric_forcing,
+        hydrological_forcing = hydrological_forcing,
+        geometry = geometry,
+        rainfall_erosion = rainfall_erosion_model,
+        overland_flow_erosion = overland_flow_erosion_model,
+        soil_erosion = soil_erosion_model,
+    )
+    return soil_loss
+end
+
+"Update soil loss model for a single timestep"
+function update!(model::SoilLoss, dt)
+    (;
+        atmospheric_forcing,
+        hydrological_forcing,
+        geometry,
+        rainfall_erosion,
+        overland_flow_erosion,
+        soil_erosion,
+    ) = model
+
+    #TODO add interception/canopygapfraction calculation here for eurosem
+    #need SBM refactor
+
+    # Rainfall erosion
+    update_boundary_conditions!(rainfall_erosion, atmospheric_forcing, hydrological_forcing)
+    update!(rainfall_erosion, geometry, dt)
+    # Overland flow erosion
+    update_boundary_conditions!(overland_flow_erosion, hydrological_forcing)
+    update!(overland_flow_erosion, geometry, dt)
+    # Total soil erosion and particle differentiation
+    update_boundary_conditions!(soil_erosion, rainfall_erosion, overland_flow_erosion)
+    update!(soil_erosion)
+end
diff --git a/src/forcing.jl b/src/forcing.jl
index 294aa2908..68f94bb20 100644
--- a/src/forcing.jl
+++ b/src/forcing.jl
@@ -1,5 +1,5 @@
 "Struct to store atmospheric forcing variables"
-@get_units @with_kw struct AtmosphericForcing{T}
+@get_units @grid_loc @with_kw struct AtmosphericForcing{T}
     # Precipitation [mm Δt⁻¹]
     precipitation::Vector{T}
     # Potential reference evapotranspiration [mm Δt⁻¹]
@@ -15,9 +15,37 @@ function AtmosphericForcing(
     potential_evaporation::Vector{T} = fill(mv, n),
     temperature::Vector{T} = fill(mv, n),
 ) where {T}
-    return AtmosphericForcing{T}(;
-        precipitation,
-        potential_evaporation,
-        temperature,
+    return AtmosphericForcing{T}(; precipitation, potential_evaporation, temperature)
+end
+
+"Struct to store hydrological forcing variables"
+@get_units @grid_loc @with_kw struct HydrologicalForcing{T}
+    # Rainfall interception by the vegetation [mm]
+    interception::Vector{T} | "mm"
+    # Overland flow depth [m]
+    waterlevel_land::Vector{T} | "m"
+    # Overland flow discharge [m3 s-1]
+    q_land::Vector{T} | "m3 s-1"
+    # River depth [m]
+    waterlevel_river::Vector{T} | "m"
+    # River discharge [m3 s-1]
+    q_river::Vector{T} | "m3 s-1"
+end
+
+"Initialize hydrological forcing"
+function HydrologicalForcing(
+    n;
+    interception::Vector{T} = fill(mv, n),
+    waterlevel_land::Vector{T} = fill(mv, n),
+    q_land::Vector{T} = fill(mv, n),
+    waterlevel_river::Vector{T} = fill(mv, n),
+    q_river::Vector{T} = fill(mv, n),
+) where {T}
+    return HydrologicalForcing{T}(;
+        interception,
+        waterlevel_land,
+        q_land,
+        waterlevel_river,
+        q_river,
     )
 end
\ No newline at end of file
diff --git a/src/parameters.jl b/src/parameters.jl
index 91dc3e662..0200feb2c 100644
--- a/src/parameters.jl
+++ b/src/parameters.jl
@@ -101,4 +101,87 @@ function VegetationParameters(dataset, config, indices)
         )
     end
     return vegetation_parameter_set
+end
+
+"Struct to store land geometry parameters"
+@get_units @grid_loc @with_kw struct LandGeometry{T}
+    # cell area [m^2]
+    area::Vector{T} | "m^2"
+    # drain width [m]
+    width::Vector{T} | "m"
+    # drain slope
+    slope::Vector{T} | "-"
+end
+
+"Initialize land geometry parameters"
+function LandGeometry(nc, config, inds)
+    # 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)))
+
+    sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
+    xl, yl = cell_lengths(y, cellength, sizeinmetres)
+    area = xl .* yl
+    ldd = ncread(nc, config, "ldd"; optional = false, sel = inds, allow_missing = true)
+    drain_width = map(get_flow_width, ldd, xl, yl)
+    landslope = ncread(
+        nc,
+        config,
+        "vertical.land_parameter_set.slope";
+        optional = false,
+        sel = inds,
+        type = Float,
+    )
+    clamp!(landslope, 0.00001, Inf)
+
+    land_parameter_set =
+        LandGeometry{Float}(; area = area, width = drain_width, slope = landslope)
+    return land_parameter_set
+end
+
+"Struct to store river geometry parameters"
+@get_units @grid_loc @with_kw struct RiverGeometry{T}
+    # river width [m]
+    width::Vector{T} | "m"
+    # river length
+    length::Vector{T} | "m"
+    # slope
+    slope::Vector{T} | "-"
+end
+
+"Initialize river geometry parameters"
+function RiverGeometry(nc, config, inds)
+    riverwidth = ncread(
+        nc,
+        config,
+        "lateral.river_parameter_set.width";
+        optional = false,
+        sel = inds,
+        type = Float,
+    )
+    riverlength = ncread(
+        nc,
+        config,
+        "lateral.river_parameter_set.length";
+        optional = false,
+        sel = inds,
+        type = Float,
+    )
+    riverslope = ncread(
+        nc,
+        config,
+        "lateral.river_parameter_set.slope";
+        optional = false,
+        sel = inds,
+        type = Float,
+    )
+    minimum(riverlength) > 0 || error("river length must be positive on river cells")
+    minimum(riverwidth) > 0 || error("river width must be positive on river cells")
+    clamp!(riverslope, 0.00001, Inf)
+
+    river_parameter_set =
+        RiverGeometry{Float}(; width = riverwidth, length = riverlength, slope = riverslope)
+    return river_parameter_set
 end
\ No newline at end of file
diff --git a/src/sediment.jl b/src/sediment.jl
deleted file mode 100644
index fbe08e418..000000000
--- a/src/sediment.jl
+++ /dev/null
@@ -1,1620 +0,0 @@
-### Soil erosion ###
-@get_units @grid_loc @with_kw struct LandSediment{T}
-    # number of cells
-    n::Int | "-" | "none"
-    ### Soil erosion part ###
-    # length of cells in y direction [m]
-    yl::Vector{T} | "m"
-    # length of cells in x direction [m]
-    xl::Vector{T} | "m"
-    # Fraction of river [-]
-    riverfrac::Vector{T} | "-"
-    # Waterbodies areas [-]
-    wbcover::Vector{T} | "-"
-    # Depth of overland flow [m]
-    h_land::Vector{T} | "m"
-    # Canopy interception [mm Δt⁻¹]
-    interception::Vector{T} | "mm dt-1"
-    # Precipitation [mm Δt⁻¹]
-    precipitation::Vector{T} | "mm dt-1"
-    # Overland flow [m3/s]
-    q_land::Vector{T} | "m3 s-1"
-    # Canopy height [m]
-    canopyheight::Vector{T} | "m"
-    # Canopy gap fraction [-]
-    canopygapfraction::Vector{T} | "-"
-    # Coefficient for EUROSEM rainfall erosion [-]
-    erosk::Vector{T} | "-"
-    # Exponent for EUROSEM rainfall erosion [-]
-    erosspl::Vector{T} | "-"
-    # Coefficient for ANSWERS overland flow erosion [-]
-    erosov::Vector{T} | "-"
-    # Fraction of impervious area per grid cell [-]
-    pathfrac::Vector{T} | "-"
-    # Land slope [-]
-    slope::Vector{T} | "-"
-    # USLE crop management factor [-]
-    usleC::Vector{T} | "-"
-    # USLE soil erodibility factor [-]
-    usleK::Vector{T} | "-"
-    # Sediment eroded by rainfall [ton Δt⁻¹]
-    sedspl::Vector{T} | "t dt-1"
-    # Sediment eroded by overland flow [ton Δt⁻¹]
-    sedov::Vector{T} | "t dt-1"
-    # Total eroded soil [ton Δt⁻¹]
-    soilloss::Vector{T} | "t dt-1"
-    # Eroded soil per particle class [ton Δt⁻¹]
-    erosclay::Vector{T} | "t dt-1"  # clay
-    erossilt::Vector{T} | "t dt-1"  # silt
-    erossand::Vector{T} | "t dt-1"  # sand
-    erossagg::Vector{T} | "t dt-1"  # small aggregates
-    eroslagg::Vector{T} | "t dt-1"  # large aggregates
-    ## Interception related to leaf_area_index climatology ###
-    # Specific leaf storage [mm]
-    sl::Vector{T} | "mm"
-    # Storage woody part of vegetation [mm]
-    swood::Vector{T} | "mm"
-    # Extinction coefficient [-] (to calculate canopy gap fraction)
-    kext::Vector{T} | "-"
-    # Leaf area index [m² m⁻²]
-    leaf_area_index::Vector{T} | "m2 m-2"
-    ### Transport capacity part ###
-    # Drain length [m]
-    dl::Vector{T} | "m"
-    # Flow width [m]
-    dw::Vector{T} | "m"
-    # Govers transport capacity coefficients [-]
-    cGovers::Vector{T} | "-"
-    nGovers::Vector{T} | "-"
-    # Median particle diameter of the topsoil [mm]
-    D50::Vector{T} | "mm"
-    # Median diameter per particle class [um]
-    dmclay::Vector{T} | "µm"
-    dmsilt::Vector{T} | "µm"
-    dmsand::Vector{T} | "µm"
-    dmsagg::Vector{T} | "µm"
-    dmlagg::Vector{T} | "µm"
-    # Fraction of the different particle class [-]
-    fclay::Vector{T} | "-"
-    fsilt::Vector{T} | "-"
-    fsand::Vector{T} | "-"
-    fsagg::Vector{T} | "-"
-    flagg::Vector{T} | "-"
-    # Density of sediment [kg/m3]
-    rhos::Vector{T} | "kg m-3"
-    # Filter with river cells
-    rivcell::Vector{T} | "-"
-    # Total transport capacity of overland flow [ton Δt⁻¹]
-    TCsed::Vector{T} | "t dt-1"
-    # Transport capacity of overland flow per particle class [ton Δt⁻¹]
-    TCclay::Vector{T} | "t dt-1"
-    TCsilt::Vector{T} | "t dt-1"
-    TCsand::Vector{T} | "t dt-1"
-    TCsagg::Vector{T} | "t dt-1"
-    TClagg::Vector{T} | "t dt-1"
-
-    function LandSediment{T}(args...) where {T}
-        equal_size_vectors(args)
-        return new(args...)
-    end
-end
-
-# This function is not used by SBM (was part of sbm.jl)
-function initialize_canopy(nc, config, inds)
-    n = length(inds)
-    # if leaf area index climatology provided use sl, swood and kext to calculate cmax, e_r and canopygapfraction
-    if haskey(config.input.vertical, "leaf_area_index")
-        # TODO confirm if leaf area index climatology is present in the netCDF
-        sl = ncread(
-            nc,
-            config,
-            "vertical.specific_leaf";
-            optional = false,
-            sel = inds,
-            type = Float,
-        )
-        swood = ncread(
-            nc,
-            config,
-            "vertical.storage_wood";
-            optional = false,
-            sel = inds,
-            type = Float,
-        )
-        kext =
-            ncread(nc, config, "vertical.kext"; optional = false, sel = inds, type = Float)
-        cmax = fill(mv, n)
-        e_r = fill(mv, n)
-        canopygapfraction = fill(mv, n)
-    else
-        sl = fill(mv, n)
-        swood = fill(mv, n)
-        kext = fill(mv, n)
-        # cmax, e_r, canopygapfraction only required when leaf area index climatology not provided
-        cmax = ncread(nc, config, "vertical.cmax"; sel = inds, defaults = 1.0, type = Float)
-        e_r =
-            ncread(nc, config, "vertical.eoverr"; sel = inds, defaults = 0.1, type = Float)
-        canopygapfraction = ncread(
-            nc,
-            config,
-            "vertical.canopygapfraction";
-            sel = inds,
-            defaults = 0.1,
-            type = Float,
-        )
-    end
-    return cmax, e_r, canopygapfraction, sl, swood, kext
-end
-
-function initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
-    # Initialize parameters for the soil loss part
-    n = length(inds)
-    do_river = get(config.model, "runrivermodel", false)::Bool
-    # Reservoir / lake
-    do_reservoirs = get(config.model, "doreservoir", false)::Bool
-    do_lakes = get(config.model, "dolake", false)::Bool
-    # Rainfall erosion equation: ["answers", "eurosem"]
-    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
-    # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
-    landtransportmethod = get(config.model, "landtransportmethod", "yalinpart")::String
-
-    altitude =
-        ncread(nc, config, "vertical.altitude"; optional = false, sel = inds, type = Float)
-    canopyheight = ncread(
-        nc,
-        config,
-        "vertical.canopyheight";
-        sel = inds,
-        defaults = 3.0,
-        type = Float,
-    )
-    erosk = ncread(nc, config, "vertical.erosk"; sel = inds, defaults = 0.6, type = Float)
-    erosspl =
-        ncread(nc, config, "vertical.erosspl"; sel = inds, defaults = 2.0, type = Float)
-    erosov = ncread(nc, config, "vertical.erosov"; sel = inds, defaults = 0.9, type = Float)
-    pathfrac =
-        ncread(nc, config, "vertical.pathfrac"; sel = inds, defaults = 0.01, type = Float)
-    slope = ncread(nc, config, "vertical.slope"; sel = inds, defaults = 0.01, type = Float)
-    usleC = ncread(nc, config, "vertical.usleC"; sel = inds, defaults = 0.01, type = Float)
-    usleK = ncread(nc, config, "vertical.usleK"; sel = inds, defaults = 0.1, type = Float)
-
-    cmax, _, canopygapfraction, sl, swood, kext = initialize_canopy(nc, config, inds)
-
-    # Initialise parameters for the transport capacity part
-    clamp!(slope, 0.00001, Inf)
-    ldd_2d = ncread(nc, config, "ldd"; optional = false, allow_missing = true)
-    ldd = ldd_2d[inds]
-
-    dmclay = ncread(nc, config, "vertical.dmclay"; sel = inds, defaults = 2.0, type = Float)
-    dmsilt =
-        ncread(nc, config, "vertical.dmsilt"; sel = inds, defaults = 10.0, type = Float)
-    dmsand =
-        ncread(nc, config, "vertical.dmsand"; sel = inds, defaults = 200.0, type = Float)
-    dmsagg =
-        ncread(nc, config, "vertical.dmsagg"; sel = inds, defaults = 30.0, type = Float)
-    dmlagg =
-        ncread(nc, config, "vertical.dmlagg"; sel = inds, defaults = 500.0, type = Float)
-    pclay = ncread(nc, config, "vertical.pclay"; sel = inds, defaults = 0.1, type = Float)
-    psilt = ncread(nc, config, "vertical.psilt"; sel = inds, defaults = 0.1, type = Float)
-    rhos =
-        ncread(nc, config, "vertical.rhosed"; sel = inds, defaults = 2650.0, type = Float)
-
-    ### Initialize transport capacity variables ###
-    rivcell = float(river)
-    # Percent Sand
-    psand = 100 .- pclay .- psilt
-    # Govers coefficient for transport capacity
-    if landtransportmethod != "yalinpart"
-        # Calculation of D50 and fraction of fine and very fine sand (fvfs) from Fooladmand et al, 2006
-        psand999 = psand .* ((999 - 25) / (1000 - 25))
-        vd50 = log.((1 ./ (0.01 .* (pclay .+ psilt)) .- 1) ./ (1 ./ (0.01 .* pclay) .- 1))
-        wd50 =
-            log.(
-                (1 ./ (0.01 .* (pclay .+ psilt .+ psand999)) .- 1) ./
-                (1 ./ (0.01 .* pclay) .- 1),
-            )
-        ad50 = 1 / log((25 - 1) / (999 - 1))
-        bd50 = ad50 ./ log((25 - 1) / 1)
-        cd50 = ad50 .* log.(vd50 ./ wd50)
-        ud50 = (.-vd50) .^ (1 .- bd50) ./ ((.-wd50) .^ (.-bd50))
-        D50 = 1 .+ (-1 ./ ud50 .* log.(1 ./ (1 ./ (0.01 .* pclay) .- 1))) .^ (1 ./ cd50) #[um]
-        D50 = D50 ./ 1000 # [mm]
-    else
-        D50 = fill(mv, n)
-    end
-    if landtransportmethod == "govers"
-        cGovers = ((D50 .* 1000 .+ 5) ./ 0.32) .^ (-0.6)
-        nGovers = ((D50 .* 1000 .+ 5) ./ 300) .^ (0.25)
-    else
-        cGovers = fill(mv, n)
-        nGovers = fill(mv, n)
-    end
-    if do_river || landtransportmethod == "yalinpart"
-        # Determine sediment size distribution, estimated from primary particle size distribution (Foster et al., 1980)
-        fclay = 0.20 .* pclay ./ 100
-        fsilt = 0.13 .* psilt ./ 100
-        fsand = 0.01 .* psand .* (1 .- 0.01 .* pclay) .^ (2.4)
-        fsagg = 0.28 .* (0.01 .* pclay .- 0.25) .+ 0.5
-        for i in 1:n
-            if pclay[i] > 50.0
-                fsagg[i] = 0.57
-            elseif pclay[i] < 25
-                fsagg[i] = 2.0 * 0.01 * pclay[i]
-            end
-        end
-        flagg = 1.0 .- fclay .- fsilt .- fsand .- fsagg
-    else
-        fclay = fill(mv, n)
-        fsilt = fill(mv, n)
-        fsand = fill(mv, n)
-        fsagg = fill(mv, n)
-        flagg = fill(mv, n)
-    end
-
-    # Reservoir and lakes
-    wbcover = zeros(Float, n)
-    if do_reservoirs
-        rescoverage_2d = ncread(
-            nc,
-            config,
-            "vertical.resareas";
-            optional = false,
-            sel = inds,
-            type = Float,
-            fill = 0.0,
-        )
-        wbcover = wbcover .+ rescoverage_2d
-    end
-    if do_lakes
-        lakecoverage_2d = ncread(
-            nc,
-            config,
-            "vertical.lakeareas";
-            optional = false,
-            sel = inds,
-            type = Float,
-            fill = 0.0,
-        )
-        wbcover = wbcover .+ lakecoverage_2d
-    end
-
-    eros = LandSediment{Float}(;
-        n = n,
-        yl = yl,
-        xl = xl,
-        riverfrac = riverfrac,
-        wbcover = wbcover,
-        ### Soil erosion part ###
-        # Forcing
-        interception = fill(mv, n),
-        h_land = fill(mv, n),
-        precipitation = fill(mv, n),
-        q_land = fill(mv, n),
-        # Parameters
-        canopyheight = canopyheight,
-        canopygapfraction = canopygapfraction,
-        erosk = erosk,
-        erosspl = erosspl,
-        erosov = erosov,
-        pathfrac = pathfrac,
-        slope = slope,
-        usleC = usleC,
-        usleK = usleK,
-        # Interception related to climatology (leaf_area_index)
-        sl = sl,
-        swood = swood,
-        kext = kext,
-        leaf_area_index = fill(mv, n),
-        # Outputs
-        sedspl = fill(mv, n),
-        sedov = fill(mv, n),
-        soilloss = fill(mv, n),
-        erosclay = fill(mv, n),
-        erossilt = fill(mv, n),
-        erossand = fill(mv, n),
-        erossagg = fill(mv, n),
-        eroslagg = fill(mv, n),
-        ### Transport capacity part ###
-        # Parameters
-        dl = map(get_flow_length, ldd, xl, yl),
-        dw = map(get_flow_width, ldd, xl, yl),
-        cGovers = cGovers,
-        D50 = D50,
-        dmclay = dmclay,
-        dmsilt = dmsilt,
-        dmsand = dmsand,
-        dmsagg = dmsagg,
-        dmlagg = dmlagg,
-        fclay = fclay,
-        fsilt = fsilt,
-        fsand = fsand,
-        fsagg = fsagg,
-        flagg = flagg,
-        nGovers = nGovers,
-        rhos = rhos,
-        rivcell = rivcell,
-        # Outputs
-        TCsed = fill(mv, n),
-        TCclay = fill(mv, n),
-        TCsilt = fill(mv, n),
-        TCsand = fill(mv, n),
-        TCsagg = fill(mv, n),
-        TClagg = fill(mv, n),
-    )
-
-    return eros
-end
-
-# Soil erosion
-function update_until_ols!(eros::LandSediment, config)
-    # Options from config
-    do_lai = haskey(config.input.vertical, "leaf_area_index")
-    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
-    dt = Second(config.timestepsecs)
-    ts = Float(dt.value)
-
-    for i in 1:(eros.n)
-
-        ### Splash / Rainfall erosion ###
-        # ANSWERS method
-        if rainerosmethod == "answers"
-            # calculate rainfall intensity [mm/min]
-            rintnsty = eros.precipitation[i] / (ts / 60)
-            # splash erosion [kg/min]
-            sedspl =
-                0.108 * eros.usleC[i] * eros.usleK[i] * eros.xl[i] * eros.yl[i] * rintnsty^2
-            # [ton/timestep]
-            sedspl = sedspl * (ts / 60) * 1e-3
-        end
-        # TODO check eurosem method output values (too high!)
-        if rainerosmethod == "eurosem"
-            if do_lai
-                cmax = eros.sl[i] * eros.leaf_area_index[i] + eros.swood[i]
-                canopygapfraction = exp(-eros.kext[i] * eros.leaf_area_index[i])
-            end
-            # calculate rainfall intensity [mm/h]
-            rintnsty = eros.precipitation[i] / (ts / 3600)
-            # Kinetic energy of direct throughfall [J/m2/mm]
-            # kedir = max(11.87 + 8.73 * log10(max(0.0001, rintnsty)),0.0) #basis used in USLE
-            kedir = max(8.95 + 8.44 * log10(max(0.0001, rintnsty)), 0.0) #variant used in most distributed mdoels
-            # Kinetic energy of leaf drainage [J/m2/mm]
-            pheff = 0.5 * eros.canopyheight[i]
-            keleaf = max((15.8 * pheff^0.5) - 5.87, 0.0)
-
-            #Depths of rainfall (total, leaf drianage, direct) [mm]
-            rdtot = eros.precipitation[i]
-            rdleaf = rdtot * 0.1 * canopygapfraction #stemflow
-            rddir = max(rdtot - rdleaf - eros.interception[i], 0.0) #throughfall
-
-            #Total kinetic energy by rainfall [J/m2]
-            ketot = (rddir * kedir + rdleaf * keleaf) * 0.001
-            # Rainfall / splash erosion [g/m2]
-            sedspl = eros.erosk[i] * ketot * exp(-eros.erosspl[i] * eros.h_land[i])
-            sedspl = sedspl * eros.xl[i] * eros.yl[i] * 1e-6 # ton/cell
-        end
-
-        # Remove the impervious area
-        sedspl = sedspl * (1.0 - eros.pathfrac[i])
-
-        ### Overland flow erosion ###
-        # ANWERS method
-        # Overland flow rate [m2/min]
-        qr_land = eros.q_land[i] * 60 / ((eros.xl[i] + eros.yl[i]) / 2)
-        # Sine of the slope
-        sinslope = sin(atan(eros.slope[i]))
-
-        # Overland flow erosion [kg/min]
-        # For a wide range of slope, it is better to use the sine of slope rather than tangeant
-        sedov =
-            eros.erosov[i] *
-            eros.usleC[i] *
-            eros.usleK[i] *
-            eros.xl[i] *
-            eros.yl[i] *
-            sinslope *
-            qr_land
-        # [ton/timestep]
-        sedov = sedov * (ts / 60) * 1e-3
-        # Remove the impervious area
-        sedov = sedov * (1.0 - eros.pathfrac[i])
-
-        # Assume no erosion in reservoir/lake areas
-        if eros.wbcover[i] > 0
-            sedspl = 0.0
-            sedov = 0.0
-        end
-
-        # Total soil loss [ton/cell/timestep]
-        soilloss = sedspl + sedov
-
-        # update the outputs and states
-        eros.sedspl[i] = sedspl
-        eros.sedov[i] = sedov
-        eros.soilloss[i] = soilloss
-        # Eroded amount per particle class
-        eros.erosclay[i] = soilloss * eros.fclay[i]
-        eros.erossilt[i] = soilloss * eros.fsilt[i]
-        eros.erossand[i] = soilloss * eros.fsand[i]
-        eros.erossagg[i] = soilloss * eros.fsagg[i]
-        eros.eroslagg[i] = soilloss * eros.flagg[i]
-    end
-    return nothing
-end
-
-### Sediment transport capacity in overland flow ###
-function update_until_oltransport!(ols::LandSediment, config::Config)
-    do_river = get(config.model, "runrivermodel", false)::Bool
-    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
-    dt = Second(config.timestepsecs)
-    ts = Float(dt.value)
-
-    for i in 1:(ols.n)
-        sinslope = sin(atan(ols.slope[i]))
-
-        if !do_river
-            # Total transport capacity without particle differentiation
-            if tcmethod == "govers"
-                TC = tc_govers(ols, i, sinslope, ts)
-            end
-
-            if tcmethod == "yalin"
-                TC = tc_yalin(ols, i, sinslope, ts)
-            end
-
-            # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
-            if ols.rivcell[i] == 1.0
-                TC = 0.0
-            end
-
-            # Set particle TC to 0
-            TCclay = 0.0
-            TCsilt = 0.0
-            TCsand = 0.0
-            TCsagg = 0.0
-            TClagg = 0.0
-        end
-
-        if do_river || tcmethod == "yalinpart"
-            TC, TCclay, TCsilt, TCsand, TCsagg, TClagg = tc_yalinpart(ols, i, sinslope, ts)
-        end
-
-        # update the outputs and states
-        ols.TCsed[i] = TC
-        ols.TCclay[i] = TCclay
-        ols.TCsilt[i] = TCsilt
-        ols.TCsand[i] = TCsand
-        ols.TCsagg[i] = TCsagg
-        ols.TClagg[i] = TClagg
-    end
-    return nothing
-end
-
-function tc_govers(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
-    # Transport capacity from govers 1990
-    # Unit stream power
-    if ols.h_land[i] > 0.0
-        velocity = ols.q_land[i] / (ols.dw[i] * ols.h_land[i])
-    else
-        velocity = 0.0
-    end
-    omega = 10 * sinslope * 100 * velocity #cm/s
-    if omega > 0.4
-        TCf = ols.cGovers[i] * (omega - 0.4)^(ols.nGovers[i]) * 2650 #kg/m3
-    else
-        TCf = 0.0
-    end
-    TC = TCf * ols.q_land[i] * ts * 1e-3 #[ton/cell]
-
-    return TC
-end
-
-function tc_yalin(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
-    # Transport capacity from Yalin without particle differentiation
-    delta = max(
-        (
-            ols.h_land[i] * sinslope / (ols.D50[i] * 0.001 * (ols.rhos[i] / 1000 - 1)) /
-            0.06 - 1
-        ),
-        0.0,
-    )
-    alphay = delta * 2.45 / (0.001 * ols.rhos[i])^0.4 * 0.06^(0.5)
-    if ols.q_land[i] > 0.0 && alphay != 0.0
-        TC = (
-            ols.dw[i] / ols.q_land[i] *
-            (ols.rhos[i] - 1000) *
-            ols.D50[i] *
-            0.001 *
-            (9.81 * ols.h_land[i] * sinslope) *
-            0.635 *
-            delta *
-            (1 - log(1 + alphay) / (alphay))
-        ) # [kg/m3]
-        TC = TC * ols.q_land[i] * ts * 1e-3 #[ton]
-    else
-        TC = 0.0
-    end
-    return TC
-end
-
-function tc_yalinpart(ols::LandSediment, i::Int, sinslope::Float, ts::Float)
-    # Transport capacity from Yalin with particle differentiation
-    # Delta parameter of Yalin for each particle class
-    delta = ols.h_land[i] * sinslope / (1e-6 * (ols.rhos[i] / 1000 - 1)) / 0.06
-    dclay = max(1 / ols.dmclay[i] * delta - 1, 0.0)
-    dsilt = max(1 / ols.dmsilt[i] * delta - 1, 0.0)
-    dsand = max(1 / ols.dmsand[i] * delta - 1, 0.0)
-    dsagg = max(1 / ols.dmsagg[i] * delta - 1, 0.0)
-    dlagg = max(1 / ols.dmlagg[i] * delta - 1, 0.0)
-    # Total transportability
-    dtot = dclay + dsilt + dsand + dsagg + dlagg
-
-    # Yalin transport capacity of overland flow for each particle class
-    if ols.q_land[i] > 0.0
-        TCa =
-            ols.dw[i] / ols.q_land[i] *
-            (ols.rhos[i] - 1000) *
-            1e-6 *
-            (9.81 * ols.h_land[i] * sinslope)
-    else
-        TCa = 0.0
-    end
-    TCb = 2.45 / (ols.rhos[i] / 1000)^0.4 * 0.06^0.5
-    if dtot != 0.0 && dclay != 0.0
-        TCclay =
-            TCa * ols.dmclay[i] * dclay / dtot *
-            0.635 *
-            dclay *
-            (1 - log(1 + dclay * TCb) / dclay * TCb) # [kg/m3]
-        TCclay = TCclay * ols.q_land[i] * ts * 1e-3 # [ton]
-    else
-        TCclay = 0.0
-    end
-    if dtot != 0.0 && dsilt != 0.0
-        TCsilt =
-            TCa * ols.dmsilt[i] * dsilt / dtot *
-            0.635 *
-            dsilt *
-            (1 - log(1 + dsilt * TCb) / dsilt * TCb) # [kg/m3]
-        TCsilt = TCsilt * ols.q_land[i] * ts * 1e-3 # [ton]
-    else
-        TCsilt = 0.0
-    end
-    if dtot != 0.0 && dsand != 0.0
-        TCsand =
-            TCa * ols.dmsand[i] * dsand / dtot *
-            0.635 *
-            dsand *
-            (1 - log(1 + dsand * TCb) / dsand * TCb) # [kg/m3]
-        TCsand = TCsand * ols.q_land[i] * ts * 1e-3 # [ton]
-    else
-        TCsand = 0.0
-    end
-    if dtot != 0.0 && dsagg != 0.0
-        TCsagg =
-            TCa * ols.dmsagg[i] * dsagg / dtot *
-            0.635 *
-            dsagg *
-            (1 - log(1 + dsagg * TCb) / dsagg * TCb) # [kg/m3]
-        TCsagg = TCsagg * ols.q_land[i] * ts * 1e-3 # [ton]
-    else
-        TCsagg = 0.0
-    end
-    if dtot != 0.0 && dlagg != 0.0
-        TClagg =
-            TCa * ols.dmlagg[i] * dlagg / dtot *
-            0.635 *
-            dlagg *
-            (1 - log(1 + dlagg * TCb) / dlagg * TCb) # [kg/m3]
-        TClagg = TClagg * ols.q_land[i] * ts * 1e-3 # [ton]
-    else
-        TClagg = 0.0
-    end
-
-    # Assume that ols all reach the river in reservoir/lake areas (very high TC)
-    if ols.wbcover[i] > 0
-        TC = 1e9
-        TCclay = 1e9
-        TCsilt = 1e9
-        TCsand = 1e9
-        TCsagg = 1e9
-        TClagg = 1e9
-    end
-
-    # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
-    if ols.rivcell[i] == 1.0
-        TCclay = 0.0
-        TCsilt = 0.0
-        TCsand = 0.0
-        TCsagg = 0.0
-        TClagg = 0.0
-    end
-
-    # Set total TC to 0
-    TC = 0.0
-
-    return TC, TCclay, TCsilt, TCsand, TCsagg, TClagg
-end
-
-### Sediment transport in overland flow ###
-@get_units @grid_loc @with_kw struct OverlandFlowSediment{T}
-    # number of cells
-    n::Int | "-" | "none"
-    # Filter with river cells
-    rivcell::Vector{T} | "-"
-    # Total eroded soil [ton Δt⁻¹]
-    soilloss::Vector{T} | "t dt-1"
-    # Eroded soil per particle class [ton Δt⁻¹]
-    erosclay::Vector{T} | "t dt-1"
-    erossilt::Vector{T} | "t dt-1"
-    erossand::Vector{T} | "t dt-1"
-    erossagg::Vector{T} | "t dt-1"
-    eroslagg::Vector{T} | "t dt-1"
-    # Total transport capacity of overland flow [ton Δt⁻¹]
-    TCsed::Vector{T} | "t dt-1"
-    # Transport capacity of overland flow per particle class [ton Δt⁻¹]
-    TCclay::Vector{T} | "t dt-1"
-    TCsilt::Vector{T} | "t dt-1"
-    TCsand::Vector{T} | "t dt-1"
-    TCsagg::Vector{T} | "t dt-1"
-    TClagg::Vector{T} | "t dt-1"
-    # Sediment flux in overland flow [ton dt-1]
-    olsed::Vector{T} | "t dt-1"
-    olclay::Vector{T} | "t dt-1"
-    olsilt::Vector{T} | "t dt-1"
-    olsand::Vector{T} | "t dt-1"
-    olsagg::Vector{T} | "t dt-1"
-    ollagg::Vector{T} | "t dt-1"
-    # Sediment reaching the river with overland flow [ton dt-1]
-    inlandsed::Vector{T} | "t dt-1"
-    inlandclay::Vector{T} | "t dt-1"
-    inlandsilt::Vector{T} | "t dt-1"
-    inlandsand::Vector{T} | "t dt-1"
-    inlandsagg::Vector{T} | "t dt-1"
-    inlandlagg::Vector{T} | "t dt-1"
-
-    function OverlandFlowSediment{T}(args...) where {T}
-        equal_size_vectors(args)
-        return new(args...)
-    end
-end
-
-function partial_update!(inland, rivcell, eroded)
-    no_erosion = zero(eltype(eroded))
-    for i in eachindex(inland)
-        inland[i] = rivcell[i] == 1 ? eroded[i] : no_erosion
-    end
-    return inland
-end
-
-function update!(ols::OverlandFlowSediment, network, config)
-    do_river = get(config.model, "runrivermodel", false)::Bool
-    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
-    zeroarr = fill(0.0, ols.n)
-
-    # transport sediment down the network
-    if do_river || tcmethod == "yalinpart"
-        # clay
-        accucapacityflux!(ols.olclay, ols.erosclay, network, ols.TCclay)
-        partial_update!(ols.inlandclay, ols.rivcell, ols.erosclay)
-        # silt
-        accucapacityflux!(ols.olsilt, ols.erossilt, network, ols.TCsilt)
-        partial_update!(ols.inlandsilt, ols.rivcell, ols.erossilt)
-        # sand
-        accucapacityflux!(ols.olsand, ols.erossand, network, ols.TCsand)
-        partial_update!(ols.inlandsand, ols.rivcell, ols.erossand)
-        # small aggregates
-        accucapacityflux!(ols.olsagg, ols.erossagg, network, ols.TCsagg)
-        partial_update!(ols.inlandsagg, ols.rivcell, ols.erossagg)
-        # large aggregates
-        accucapacityflux!(ols.ollagg, ols.eroslagg, network, ols.TClagg)
-        partial_update!(ols.inlandlagg, ols.rivcell, ols.eroslagg)
-
-        # total sediment, all particle classes
-        ols.olsed .= ols.olclay .+ ols.olsilt .+ ols.olsand .+ ols.olsagg .+ ols.ollagg
-        ols.inlandsed .=
-            ols.inlandclay .+ ols.inlandsilt .+ ols.inlandsand .+ ols.inlandsagg .+
-            ols.inlandlagg
-    else
-        accucapacityflux!(ols.olsed, ols.soilloss, network, ols.TCsed)
-        ols.inlandsed .= ifelse.(ols.rivcell .== 1, ols.soilloss, zeroarr)
-    end
-    return nothing
-end
-
-### River transport and processes ###
-@get_units @grid_loc @with_kw struct RiverSediment{T}
-    # number of cells
-    n::Int | "-" | "none"
-    # Timestep [s]
-    dt::T | "s"
-    # River geometry (slope [-], length [m], width [m])
-    sl::Vector{T} | "m"
-    dl::Vector{T} | "m"
-    width::Vector{T} | "m"
-    # Sediment mean diameter in the river bed [mm]
-    d50::Vector{T} | "mm"
-    # Particle mean diameter [mm]
-    dmclay::Vector{T} | "mm"
-    dmsilt::Vector{T} | "mm"
-    dmsand::Vector{T} | "mm"
-    dmsagg::Vector{T} | "mm"
-    dmlagg::Vector{T} | "mm"
-    dmgrav::Vector{T} | "mm"
-    # River bed and bank particle fraction composition [-]
-    fclayriv::Vector{T} | "-"
-    fsiltriv::Vector{T} | "-"
-    fsandriv::Vector{T} | "-"
-    fgravriv::Vector{T} | "-"
-    # Sediment mean diameter for Engelund and Hansen transport equation [mm]
-    d50engelund::Vector{T} | "mm"
-    # Parameters for Bagnold transport equation
-    cbagnold::Vector{T} | "-"
-    ebagnold::Vector{T} | "-"
-    # Parameters for Kodatie transport equation
-    ak::Vector{T} | "-"
-    bk::Vector{T} | "-"
-    ck::Vector{T} | "-"
-    dk::Vector{T} | "-"
-    # Critical bed and bank shear stress [N/m2]
-    TCrbank::Vector{T} | "N m-2"
-    TCrbed::Vector{T} | "N m-2"
-    # Bed and bank erodibilities [m3/N.s]
-    kdbank::Vector{T} | "m3 N-1 s-1"
-    kdbed::Vector{T} | "m3 N-1 s-1"
-    # Sediment density [kg/m3]
-    rhos::Vector{T} | "kg m-3"
-    # River water level [m]
-    h_riv::Vector{T} | "m"
-    # River discharge [m3/s]
-    q_riv::Vector{T} | "m3 s-1"
-    # Sediment input from land erosion [ton Δt⁻¹]
-    inlandclay::Vector{T} | "t dt-1"
-    inlandsilt::Vector{T} | "t dt-1"
-    inlandsand::Vector{T} | "t dt-1"
-    inlandsagg::Vector{T} | "t dt-1"
-    inlandlagg::Vector{T} | "t dt-1"
-    inlandsed::Vector{T} | "t dt-1"
-    # Sediment / particle left in the cell [ton]
-    sedload::Vector{T} | "t"
-    clayload::Vector{T} | "t"
-    siltload::Vector{T} | "t"
-    sandload::Vector{T} | "t"
-    saggload::Vector{T} | "t"
-    laggload::Vector{T} | "t"
-    gravload::Vector{T} | "t"
-    # Sediment / particle stored on the river bed after deposition [ton Δt⁻¹]
-    sedstore::Vector{T} | "t dt-1"
-    claystore::Vector{T} | "t dt-1"
-    siltstore::Vector{T} | "t dt-1"
-    sandstore::Vector{T} | "t dt-1"
-    saggstore::Vector{T} | "t dt-1"
-    laggstore::Vector{T} | "t dt-1"
-    gravstore::Vector{T} | "t dt-1"
-    # Sediment / particle flux [ton Δt⁻¹]
-    outsed::Vector{T} | "t dt-1"
-    outclay::Vector{T} | "t dt-1"
-    outsilt::Vector{T} | "t dt-1"
-    outsand::Vector{T} | "t dt-1"
-    outsagg::Vector{T} | "t dt-1"
-    outlagg::Vector{T} | "t dt-1"
-    outgrav::Vector{T} | "t dt-1"
-    # Total sediment concentrations (SSconc + Bedconc) [g/m3]
-    Sedconc::Vector{T} | "g m-3"
-    # Suspended load concentration [g/m3]
-    SSconc::Vector{T} | "g m-3"
-    # Bed load concentration [g/m3]
-    Bedconc::Vector{T} | "g m-3"
-    # River transport capacity
-    maxsed::Vector{T} | "t dt-1"
-    # Eroded sediment (total, bank and bed)
-    erodsed::Vector{T} | "t dt-1"
-    erodsedbank::Vector{T} | "t dt-1"
-    erodsedbed::Vector{T} | "t dt-1"
-    # Deposited sediment
-    depsed::Vector{T} | "t dt-1"
-    # Sediment in
-    insed::Vector{T} | "t dt-1"
-    # Reservoir and lakes
-    wbcover::Vector{T} | "-"
-    wblocs::Vector{T} | "-"
-    wbarea::Vector{T} | "m2"
-    wbtrap::Vector{T} | "-"
-
-    # function RiverSediment{T}(args...) where {T}
-    #     equal_size_vectors(args)
-    #     return new(args...)
-    # end
-end
-
-function initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
-    # Initialize river parameters
-    nriv = length(inds_riv)
-    # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
-    tcmethodriv = get(config.model, "rivtransportmethod", "bagnold")::String
-    dt = Second(config.timestepsecs)
-    # Reservoir / lakes
-    do_reservoirs = get(config.model, "doreservoir", false)::Bool
-    do_lakes = get(config.model, "dolake", false)::Bool
-    wbcover = zeros(Float, nriv)
-    wblocs = zeros(Float, nriv)
-    wbarea = zeros(Float, nriv)
-    wbtrap = zeros(Float, nriv)
-
-    if do_reservoirs
-        reslocs = ncread(
-            nc,
-            config,
-            "lateral.river.reslocs";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0,
-        )
-        rescoverage_2d = ncread(
-            nc,
-            config,
-            "lateral.river.resareas";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0,
-        )
-        resarea = ncread(
-            nc,
-            config,
-            "lateral.river.resarea";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0.0,
-        )
-        restrapefficiency = ncread(
-            nc,
-            config,
-            "lateral.river.restrapeff";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            defaults = 1.0,
-            fill = 0.0,
-        )
-
-        wbcover = wbcover .+ rescoverage_2d
-        wblocs = wblocs .+ reslocs
-        wbarea = wbarea .+ resarea
-        wbtrap = wbtrap .+ restrapefficiency
-    end
-
-    if do_lakes
-        lakelocs = ncread(
-            nc,
-            config,
-            "lateral.river.lakelocs";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0,
-        )
-        lakecoverage_2d = ncread(
-            nc,
-            config,
-            "lateral.river.lakeareas";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0,
-        )
-        lakearea = ncread(
-            nc,
-            config,
-            "lateral.river.lakearea";
-            optional = false,
-            sel = inds_riv,
-            type = Float,
-            fill = 0.0,
-        )
-
-        wbcover = wbcover .+ lakecoverage_2d
-        wblocs = wblocs .+ lakelocs
-        wbarea = wbarea .+ lakearea
-    end
-
-    riverslope = ncread(
-        nc,
-        config,
-        "lateral.river.slope";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    clamp!(riverslope, 0.00001, Inf)
-    rhos = ncread(
-        nc,
-        config,
-        "lateral.river.rhosed";
-        sel = inds_riv,
-        defaults = 2650.0,
-        type = Float,
-    )
-    dmclay = ncread(
-        nc,
-        config,
-        "lateral.river.dmclay";
-        sel = inds_riv,
-        defaults = 2.0,
-        type = Float,
-    )
-    dmsilt = ncread(
-        nc,
-        config,
-        "lateral.river.dmsilt";
-        sel = inds_riv,
-        defaults = 10.0,
-        type = Float,
-    )
-    dmsand = ncread(
-        nc,
-        config,
-        "lateral.river.dmsand";
-        sel = inds_riv,
-        defaults = 200.0,
-        type = Float,
-    )
-    dmsagg = ncread(
-        nc,
-        config,
-        "lateral.river.dmsagg";
-        sel = inds_riv,
-        defaults = 30.0,
-        type = Float,
-    )
-    dmlagg = ncread(
-        nc,
-        config,
-        "lateral.river.dmlagg";
-        sel = inds_riv,
-        defaults = 500.0,
-        type = Float,
-    )
-    dmgrav = ncread(
-        nc,
-        config,
-        "lateral.river.dmgrav";
-        sel = inds_riv,
-        defaults = 2000.0,
-        type = Float,
-    )
-    fclayriv = ncread(
-        nc,
-        config,
-        "lateral.river.fclayriv";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    fsiltriv = ncread(
-        nc,
-        config,
-        "lateral.river.fsiltriv";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    fsandriv = ncread(
-        nc,
-        config,
-        "lateral.river.fsandriv";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    fgravriv = ncread(
-        nc,
-        config,
-        "lateral.river.fgravriv";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    d50riv = ncread(
-        nc,
-        config,
-        "lateral.river.d50";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    d50engelund = ncread(
-        nc,
-        config,
-        "lateral.river.d50engelund";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    cbagnold = ncread(
-        nc,
-        config,
-        "lateral.river.cbagnold";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-    ebagnold = ncread(
-        nc,
-        config,
-        "lateral.river.ebagnold";
-        optional = false,
-        sel = inds_riv,
-        type = Float,
-    )
-
-    # Initialisation of parameters for Kodatie transport capacity
-    ak = zeros(Float, nriv)
-    bk = zeros(Float, nriv)
-    ck = zeros(Float, nriv)
-    dk = zeros(Float, nriv)
-    if tcmethodriv == "kodatie"
-        for i in 1:nriv
-            if d50riv[i] <= 0.05
-                ak[i] = 281.4
-                bk[i] = 2.622
-                ck[i] = 0.182
-                dk[i] = 0.0
-            elseif d50riv[i] <= 0.25
-                ak[i] = 2829.6
-                bk[i] = 3.646
-                ck[i] = 0.406
-                dk[i] = 0.412
-            elseif d50riv[i] <= 2.0
-                ak[i] = 2123.4
-                bk[i] = 3.3
-                ck[i] = 0.468
-                dk[i] = 0.613
-            else
-                ak[i] = 431884.8
-                bk[i] = 1.0
-                ck[i] = 1.0
-                dk[i] = 2.0
-            end
-        end
-    end
-    # Initialisation of parameters for river erosion
-    # Bed and Bank from Shields diagram, Da Silva & Yalin (2017)
-    E_ = (2.65 - 1) * 9.81
-    E = (E_ .* (d50riv .* 1e-3) .^ 3 ./ 1e-12) .^ 0.33
-    TCrbed = @. Float(
-        E_ *
-        d50riv *
-        (0.13 * E^(-0.392) * exp(-0.015 * E^2) + 0.045 * (1 - exp(-0.068 * E))),
-    )
-    TCrbank = TCrbed
-    # kd from Hanson & Simon 2001
-    kdbank = @. Float(0.2 * TCrbank^(-0.5) * 1e-6)
-    kdbed = @. Float(0.2 * TCrbed^(-0.5) * 1e-6)
-
-    rs = RiverSediment(;
-        n = nriv,
-        dt = Float(dt.value),
-        # Parameters
-        sl = riverslope,
-        dl = riverlength,
-        width = riverwidth,
-        dmclay = dmclay,
-        dmsilt = dmsilt,
-        dmsand = dmsand,
-        dmsagg = dmsagg,
-        dmlagg = dmlagg,
-        dmgrav = dmgrav,
-        fclayriv = fclayriv,
-        fsiltriv = fsiltriv,
-        fsandriv = fsandriv,
-        fgravriv = fgravriv,
-        d50 = d50riv,
-        d50engelund = d50engelund,
-        cbagnold = cbagnold,
-        ebagnold = ebagnold,
-        ak = ak,
-        bk = bk,
-        ck = ck,
-        dk = dk,
-        kdbank = kdbank,
-        kdbed = kdbed,
-        TCrbank = TCrbank,
-        TCrbed = TCrbed,
-        rhos = rhos,
-        # Forcing
-        h_riv = fill(mv, nriv),
-        q_riv = fill(mv, nriv),
-        # Input from land
-        inlandclay = zeros(Float, nriv),
-        inlandsilt = zeros(Float, nriv),
-        inlandsand = zeros(Float, nriv),
-        inlandsagg = zeros(Float, nriv),
-        inlandlagg = zeros(Float, nriv),
-        inlandsed = zeros(Float, nriv),
-        # States
-        sedload = zeros(Float, nriv),
-        clayload = zeros(Float, nriv),
-        siltload = zeros(Float, nriv),
-        sandload = zeros(Float, nriv),
-        saggload = zeros(Float, nriv),
-        laggload = zeros(Float, nriv),
-        gravload = zeros(Float, nriv),
-        sedstore = zeros(Float, nriv),
-        claystore = zeros(Float, nriv),
-        siltstore = zeros(Float, nriv),
-        sandstore = zeros(Float, nriv),
-        saggstore = zeros(Float, nriv),
-        laggstore = zeros(Float, nriv),
-        gravstore = zeros(Float, nriv),
-        outsed = zeros(Float, nriv),
-        outclay = zeros(Float, nriv),
-        outsilt = zeros(Float, nriv),
-        outsand = zeros(Float, nriv),
-        outsagg = zeros(Float, nriv),
-        outlagg = zeros(Float, nriv),
-        outgrav = zeros(Float, nriv),
-        # Outputs
-        Sedconc = zeros(Float, nriv),
-        SSconc = zeros(Float, nriv),
-        Bedconc = zeros(Float, nriv),
-        maxsed = zeros(Float, nriv),
-        erodsed = zeros(Float, nriv),
-        erodsedbank = zeros(Float, nriv),
-        erodsedbed = zeros(Float, nriv),
-        depsed = zeros(Float, nriv),
-        insed = zeros(Float, nriv),
-        # Reservoir / lake
-        wbcover = wbcover,
-        wblocs = wblocs,
-        wbarea = wbarea,
-        wbtrap = wbtrap,
-    )
-
-    return rs
-end
-
-function update!(rs::RiverSediment, network, config)
-    (; graph, order) = network
-    tcmethod = get(config.model, "rivtransportmethod", "bagnold")::String
-
-    # River sediment loads are separated into different particle class.
-    # Clay, silt and sand can both come from land, resuspension or river channel erosion.
-    # Small and large aggregates only come from land erosion or resuspension.
-    # Gravel only comes from resuspension or river channel erosion.
-
-    for v in order
-        ### Sediment input in the cell (left from previous timestep + from land + from upstream outflux) ###
-        upstream_nodes = inneighbors(graph, v)
-
-        inrivclay = 0.0
-        inrivsilt = 0.0
-        inrivsand = 0.0
-        inrivsagg = 0.0
-        inrivlagg = 0.0
-        inrivgrav = 0.0
-        if !isempty(upstream_nodes)
-            for i in upstream_nodes
-                if rs.outclay[i] >= 0.0 # avoid NaN from upstream non-river cells
-                    inrivclay += rs.outclay[i]
-                    inrivsilt += rs.outsilt[i]
-                    inrivsand += rs.outsand[i]
-                    inrivsagg += rs.outsagg[i]
-                    inrivlagg += rs.outlagg[i]
-                    inrivgrav += rs.outgrav[i]
-                end
-            end
-        end
-
-        inclay = rs.clayload[v] + rs.inlandclay[v] + inrivclay
-        insilt = rs.siltload[v] + rs.inlandsilt[v] + inrivsilt
-        insand = rs.sandload[v] + rs.inlandsand[v] + inrivsand
-        insagg = rs.saggload[v] + rs.inlandsagg[v] + inrivsagg
-        inlagg = rs.laggload[v] + rs.inlandlagg[v] + inrivlagg
-        ingrav = rs.gravload[v] + inrivgrav
-
-        insed = inclay + insilt + insand + insagg + inlagg + ingrav
-        rs.insed[v] = insed
-        rs.inlandsed[v] =
-            rs.inlandclay[v] +
-            rs.inlandsilt[v] +
-            rs.inlandsand[v] +
-            rs.inlandsagg[v] +
-            rs.inlandlagg[v]
-
-        ### Transport capacity of the flow ###
-        # Hydraulic radius of the river [m] (rectangular channel)
-        hydrad = rs.h_riv[v] * rs.width[v] / (rs.width[v] + 2 * rs.h_riv[v])
-
-        # Engelund and Hansen transport formula
-        if tcmethod == "engelund"
-            vmean =
-                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
-            vshear = (9.81 * hydrad * rs.sl[v])^0.5
-            # Concentration by weight
-            cw = ifelse(
-                hydrad > 0.0,
-                (
-                    rs.rhos[v] / 1000 * 0.05 * vmean * vshear^3 /
-                    ((rs.rhos[v] / 1000 - 1)^2 * 9.81^2 * rs.d50engelund[v] * hydrad)
-                ),
-                0.0,
-            )
-            cw = min(1.0, cw)
-            # Transport capacity [tons/m3]
-            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
-        elseif tcmethod == "bagnold"
-            maxsed =
-                rs.cbagnold[v] * (rs.q_riv[v] / (rs.h_riv[v] * rs.width[v]))^rs.ebagnold[v]
-        elseif tcmethod == "kodatie"
-            vmean =
-                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
-            maxsed = rs.ak[v] * vmean^rs.bk[v] * rs.h_riv[v]^rs.ck[v] * rs.sl[v]^rs.dk[v]
-            # Transport capacity [tons/m3]
-            maxsed =
-                ifelse(rs.q_riv[v] > 0.0, maxsed * rs.width[v] / (rs.q_riv[v] * rs.dt), 0.0)
-        elseif tcmethod == "yang"
-            ws = 411 * rs.d50[v]^2 / 3600
-            vshear = (9.81 * hydrad * rs.sl[v])^0.5
-            var1 = vshear * rs.d50[v] / 1000 / (1.16 * 1e-6)
-            var2 = ws * rs.d50[v] / 1000 / (1.16 * 1e-6)
-            vcr = min(
-                0.0,
-                ifelse(var1 >= 70.0, 2.05 * ws, ws * (2.5 / (log10(var1) - 0.06) + 0.66)),
-            )
-            # Sand equation
-            if (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
-                logcppm = (
-                    5.435 - 0.286 * log10(var2) - 0.457 * log10(vshear / ws) + 1.799 -
-                    0.409 * log10(var2) -
-                    0.314 *
-                    log10(vshear / ws) *
-                    log10((rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws)
-                )
-                # Gravel equation
-            elseif (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
-                logcppm = (
-                    6.681 - 0.633 * log10(var2) - 4.816 * log10(vshear / ws) + 2.784 -
-                    0.305 * log10(var2) -
-                    0.282 *
-                    log10(vshear / ws) *
-                    log10((rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws)
-                )
-            else
-                logcppm = 0.0
-            end
-            # Sediment concentration by weight
-            cw = 10^logcppm * 1e-6
-            # Transport capacity [ton/m3]
-            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
-        elseif tcmethod == "molinas"
-            ws = 411 * rs.d50[v]^2 / 3600
-            vmean =
-                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
-            if rs.h_riv[v] > 0.0
-                psi = (
-                    vmean^3 / (
-                        (rs.rhos[v] / 1000 - 1) *
-                        9.81 *
-                        rs.h_riv[v] *
-                        ws *
-                        log10(1000 * rs.h_riv[v] / rs.d50[v])^2
-                    )
-                )
-            else
-                psi = 0.0
-            end
-            # Concentration by weight
-            cw = 1430 * (0.86 + psi^0.5) * psi^1.5 / (0.016 + psi) * 1e-6
-            # Transport capacity [ton/m3]
-            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
-        end
-
-        # 1285 g/L: boundary between streamflow and debris flow (Costa, 1988)
-        maxsed = min(maxsed, 1.285)
-        # Transport capacity [ton]
-        maxsed = maxsed * (rs.h_riv[v] * rs.width[v] * rs.dl[v] + rs.q_riv[v] * rs.dt)
-        rs.maxsed[v] = maxsed
-
-        ### River erosion ###
-        # Erosion only if the load is below the transport capacity of the flow.
-        sedex = max(maxsed - insed, 0.0)
-        # No erosion in lake and reservoir cells
-        if rs.wbcover[v] > 0.0
-            sedex = 0.0
-        end
-        # Bed and bank are eroded only if the previously deposited material is not enough
-        rs.sedstore[v] =
-            rs.claystore[v] +
-            rs.siltstore[v] +
-            rs.sandstore[v] +
-            rs.saggstore[v] +
-            rs.laggstore[v] +
-            rs.gravstore[v]
-        if sedex > 0.0 && sedex > rs.sedstore[v]
-            # Effective sediment needed fom river bed and bank erosion [ton]
-            effsedex = sedex - rs.sedstore[v]
-
-            # Repartition of the effective shear stress between the bank and the bed from Knight et al. 1984 [%]
-            SFbank = ifelse(
-                rs.h_riv[v] > 0.0,
-                exp(-3.23 * log10(rs.width[v] / rs.h_riv[v] + 3) + 6.146),
-                0.0,
-            )
-            # Effective shear stress on river bed and banks [N/m2]
-            TEffbank = ifelse(
-                rs.h_riv[v] > 0.0,
-                1000 * 9.81 * hydrad * rs.sl[v] * SFbank / 100 *
-                (1 + rs.width[v] / (2 * rs.h_riv[v])),
-                0.0,
-            )
-            TEffbed =
-                1000 *
-                9.81 *
-                hydrad *
-                rs.sl[v] *
-                (1 - SFbank / 100) *
-                (1 + 2 * rs.h_riv[v] / rs.width[v])
-            # Potential erosion rates of the bed and bank [t/cell/timestep]
-            #(assuming only one bank is eroding)
-            Tex = max(TEffbank - rs.TCrbank[v], 0.0)
-            # 1.4 is bank default bulk density
-            ERbank = max(0.0, rs.kdbank[v] * Tex * rs.dl[v] * rs.h_riv[v] * 1.4 * rs.dt)
-            # 1.5 is bed default bulk density
-            ERbed = max(
-                0.0,
-                rs.kdbed[v] *
-                (TEffbed - rs.TCrbed[v]) *
-                rs.dl[v] *
-                rs.width[v] *
-                1.5 *
-                rs.dt,
-            )
-            # Relative potential erosion rates of the bed and the bank [-]
-            RTEbank = ifelse(ERbank + ERbed > 0.0, ERbank / (ERbank + ERbed), 0.0)
-            RTEbed = 1.0 - RTEbank
-
-            # Bank erosion (difference between effective and potential erosion) [ton]
-            sedbank = ifelse(effsedex * RTEbank <= ERbank, effsedex * RTEbank, ERbank)
-            claybank = rs.fclayriv[v] * sedbank
-            siltbank = rs.fsiltriv[v] * sedbank
-            sandbank = rs.fsandriv[v] * sedbank
-            gravbank = rs.fgravriv[v] * sedbank
-
-            # Bed erosion [ton]
-            sedbed = ifelse(effsedex * RTEbed <= ERbed, effsedex * RTEbed, ERbed)
-            claybed = rs.fclayriv[v] * sedbed
-            siltbed = rs.fsiltriv[v] * sedbed
-            sandbed = rs.fsandriv[v] * sedbed
-            gravbed = rs.fgravriv[v] * sedbed
-        else
-            sedbank = 0.0
-            claybank = 0.0
-            siltbank = 0.0
-            sandbank = 0.0
-            gravbank = 0.0
-
-            sedbed = 0.0
-            claybed = 0.0
-            siltbed = 0.0
-            sandbed = 0.0
-            gravbed = 0.0
-        end
-
-        # Erosion/degradation of the previously deposited sediment (from clay to gravel) [ton]
-        if sedex > 0.0
-            degstoreclay = ifelse(rs.claystore[v] >= sedex, sedex, rs.claystore[v])
-            rs.claystore[v] = rs.claystore[v] - degstoreclay
-            #Update amount of sediment that need to be degraded
-            sedex = sedex - degstoreclay
-            degstoresilt = ifelse(rs.siltstore[v] >= sedex, sedex, rs.siltstore[v])
-            rs.siltstore[v] = rs.siltstore[v] - degstoresilt
-            sedex = max(0.0, sedex - degstoresilt)
-            degstoresagg = ifelse(rs.saggstore[v] >= sedex, sedex, rs.saggstore[v])
-            rs.saggstore[v] = rs.saggstore[v] - degstoresagg
-            sedex = max(0.0, sedex - degstoresagg)
-            degstoresand = ifelse(rs.sandstore[v] >= sedex, sedex, rs.sandstore[v])
-            rs.sandstore[v] = rs.sandstore[v] - degstoresand
-            sedex = max(0.0, sedex - degstoresand)
-            degstorelagg = ifelse(rs.laggstore[v] >= sedex, sedex, rs.laggstore[v])
-            rs.laggstore[v] = rs.laggstore[v] - degstorelagg
-            sedex = max(0.0, sedex - degstorelagg)
-            degstoregrav = ifelse(rs.gravstore[v] >= sedex, sedex, rs.gravstore[v])
-            rs.gravstore[v] = rs.gravstore[v] - degstoregrav
-            sedex = max(0.0, sedex - degstoregrav)
-            degstoresed =
-                degstoreclay +
-                degstoresilt +
-                degstoresagg +
-                degstoresand +
-                degstorelagg +
-                degstoregrav
-        else
-            degstoreclay = 0.0
-            degstoresilt = 0.0
-            degstoresagg = 0.0
-            degstoresand = 0.0
-            degstorelagg = 0.0
-            degstoregrav = 0.0
-            degstoresed = 0.0
-        end
-
-        # Sum all erosion sources per particle class
-        erodsed = sedbank + sedbed + degstoresed
-        erodclay = claybank + claybed + degstoreclay
-        erodsilt = siltbank + siltbed + degstoresilt
-        erodsand = sandbank + sandbed + degstoresand
-        erodsagg = degstoresagg
-        erodlagg = degstorelagg
-        erodgrav = gravbank + gravbed + degstoregrav
-
-        rs.erodsed[v] = erodsed
-        rs.erodsedbank[v] = sedbank
-        rs.erodsedbed[v] = sedbed
-
-        ### Deposition / settling ###
-        # Fractions of deposited particles in river cells from the Einstein formula [-]
-        # Particle fall velocity [m/s] from Stokes
-        xs = ifelse(rs.q_riv[v] > 0.0, 1.055 * rs.dl[v] / (rs.q_riv[v] / rs.width[v]), 0.0)
-        xclay = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmclay[v] / 1000)^2 / 3600)))
-        xsilt = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsilt[v] / 1000)^2 / 3600)))
-        xsand = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsand[v] / 1000)^2 / 3600)))
-        xsagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsagg[v] / 1000)^2 / 3600)))
-        xlagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmlagg[v] / 1000)^2 / 3600)))
-        xgrav = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmgrav[v] / 1000)^2 / 3600)))
-
-        # Sediment deposited in the channel [ton]
-        # From natural settling with Einstein formula (density controlled).
-        settclay = xclay * (inclay + erodclay)
-        settsilt = xsilt * (insilt + erodsilt)
-        settsand = xsand * (insand + erodsand)
-        settsagg = xsagg * (insagg + erodsagg)
-        settlagg = xlagg * (inlagg + erodlagg)
-        settgrav = xgrav * (ingrav + erodgrav)
-
-        # Sediment deposited in the channel (from gravel to clay) [ton]
-        # From transport capacity exceedance (insed > maxsed)
-        insedex = max(insed - maxsed, 0.0)
-        if insedex > 0.0
-            depgrav = ifelse(ingrav >= insedex, insedex, ingrav)
-            insedex = max(insedex - depgrav, 0.0)
-            deplagg = ifelse(inlagg >= insedex, insedex, inlagg)
-            insedex = max(insedex - deplagg, 0.0)
-            depsand = ifelse(insand >= insedex, insedex, insand)
-            insedex = max(insedex - depsand, 0.0)
-            depsagg = ifelse(insagg >= insedex, insedex, insagg)
-            insedex = max(insedex - depsagg, 0.0)
-            depsilt = ifelse(insilt >= insedex, insedex, insilt)
-            insedex = max(insedex - depsilt, 0.0)
-            depclay = ifelse(inclay >= insedex, insedex, inclay)
-            insedex = max(insedex - depclay, 0.0)
-        else
-            depclay = settclay
-            depsilt = settsilt
-            depsand = settsand
-            depsagg = settsagg
-            deplagg = settlagg
-            depgrav = settgrav
-        end
-
-        # No deposition in regular lake and reservoir cells, only at the outlet
-        # Deposition in lake/reservoir from Camp 1945
-        # Extra trapping of large particles for dams
-        if rs.wbcover[v] > 0.0 && rs.wblocs[v] > 0.0
-            # Compute deposition
-            vcres = rs.q_riv[v] / rs.wbarea[v]
-            DCres = 411 / 3600 / vcres
-            depclay = (inclay + erodclay) * min(1.0, (DCres * (rs.dmclay[v] / 1000)^2))
-            depsilt = (insilt + erodsilt) * min(1.0, (DCres * (rs.dmsilt[v] / 1000)^2))
-            depsand = (insand + erodsand) * min(1.0, (DCres * (rs.dmsand[v] / 1000)^2))
-            depsagg = (insagg + erodsagg) * min(1.0, (DCres * (rs.dmsagg[v] / 1000)^2))
-            deplagg = (inlagg + erodlagg) * min(1.0, (DCres * (rs.dmlagg[v] / 1000)^2))
-            depgrav = (ingrav + erodgrav) * min(1.0, (DCres * (rs.dmgrav[v] / 1000)^2))
-            # Trapping of large particles
-            # Use the rouse number for sagg (suspension or bedload)
-            depsand = max(depsand, rs.wbtrap[v] * (insand + erodsand))
-            deplagg = max(deplagg, rs.wbtrap[v] * (inlagg + erodlagg))
-            depgrav = max(depgrav, rs.wbtrap[v] * (ingrav + erodgrav))
-            # threshold diameter between suspended load and mixed load using Rouse number
-            dsuspf =
-                1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
-            depsagg = ifelse(
-                rs.dmsagg[v] > dsuspf,
-                depsagg,
-                max(depsagg, rs.wbtrap[v] * (insagg + erodsagg)),
-            )
-        elseif rs.wbcover[v] > 0.0
-            depsed = 0.0
-            depclay = 0.0
-            depsilt = 0.0
-            depsand = 0.0
-            depsagg = 0.0
-            deplagg = 0.0
-            depgrav = 0.0
-        end
-
-        depsed = depclay + depsilt + depsand + depsagg + deplagg + depgrav
-        rs.depsed[v] = depsed
-
-        # Update the river deposited sediment storage
-        rs.sedstore[v] = rs.sedstore[v] + depsed
-        rs.claystore[v] = rs.claystore[v] + depclay
-        rs.siltstore[v] = rs.siltstore[v] + depsilt
-        rs.sandstore[v] = rs.sandstore[v] + depsand
-        rs.saggstore[v] = rs.saggstore[v] + depsagg
-        rs.laggstore[v] = rs.laggstore[v] + deplagg
-        rs.gravstore[v] = rs.gravstore[v] + depgrav
-
-        ### Ouput loads ###
-        # Sediment transported out of the cell during the timestep [ton]
-        # 0 in case all sediment are deposited in the cell
-        # Reduce the fraction so that there is still some sediment staying in the river cell
-        fwaterout = min(rs.q_riv[v] * rs.dt / (rs.h_riv[v] * rs.width[v] * rs.dl[v]), 1.0)
-        rs.outsed[v] = fwaterout * (insed + erodsed - depsed)
-        rs.outclay[v] = fwaterout * (inclay + erodclay - depclay)
-        rs.outsilt[v] = fwaterout * (insilt + erodsilt - depsilt)
-        rs.outsand[v] = fwaterout * (insand + erodsand - depsand)
-        rs.outsagg[v] = fwaterout * (insagg + erodsagg - depsagg)
-        rs.outlagg[v] = fwaterout * (inlagg + erodlagg - deplagg)
-        rs.outgrav[v] = fwaterout * (ingrav + erodgrav - depgrav)
-
-        ### Mass balance ###
-        # Sediment left in the cell [ton]
-        rs.sedload[v] = insed + erodsed - depsed - rs.outsed[v]
-        rs.clayload[v] = inclay + erodclay - depclay - rs.outclay[v]
-        rs.siltload[v] = insilt + erodsilt - depsilt - rs.outsilt[v]
-        rs.sandload[v] = insand + erodsand - depsand - rs.outsand[v]
-        rs.saggload[v] = insagg + erodsagg - depsagg - rs.outsagg[v]
-        rs.gravload[v] = ingrav + erodgrav - depgrav - rs.outgrav[v]
-
-        ### Concentrations and suspended sediments ###
-        # Conversion from load [ton] to concentration for rivers [mg/L]
-        toconc = ifelse(rs.q_riv[v] > 0.0, 1e6 / (rs.q_riv[v] * rs.dt), 0.0)
-        rs.Sedconc[v] = rs.outsed[v] * toconc
-
-        # Differentiation of bed and suspended load using Rouse number for suspension
-        # threshold diameter between bed load and mixed load using Rouse number
-        dbedf = 1e3 * (2.5 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
-        # threshold diameter between suspended load and mixed load using Rouse number
-        dsuspf = 1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
-        # Rouse with diameter
-        SSclay = ifelse(
-            rs.dmclay[v] <= dsuspf,
-            rs.outclay[v],
-            ifelse(rs.dmclay[v] <= dbedf, rs.outclay[v] / 2, 0.0),
-        )
-        SSsilt = ifelse(
-            rs.dmsilt[v] <= dsuspf,
-            rs.outsilt[v],
-            ifelse(rs.dmsilt[v] <= dbedf, rs.outsilt[v] / 2, 0.0),
-        )
-        SSsagg = ifelse(
-            rs.dmsagg[v] <= dsuspf,
-            rs.outsagg[v],
-            ifelse(rs.dmsagg[v] <= dbedf, rs.outsagg[v] / 2, 0.0),
-        )
-        SSsand = ifelse(
-            rs.dmsand[v] <= dsuspf,
-            rs.outsand[v],
-            ifelse(rs.dmsand[v] <= dbedf, rs.outsand[v] / 2, 0.0),
-        )
-        SSlagg = ifelse(
-            rs.dmlagg[v] <= dsuspf,
-            rs.outlagg[v],
-            ifelse(rs.dmlagg[v] <= dbedf, rs.outlagg[v] / 2, 0.0),
-        )
-        SSgrav = ifelse(
-            rs.dmgrav[v] <= dsuspf,
-            rs.outgrav[v],
-            ifelse(rs.dmgrav[v] <= dbedf, rs.outgrav[v] / 2, 0.0),
-        )
-
-        SS = SSclay + SSsilt + SSsagg + SSsand + SSlagg + SSgrav
-        Bed = rs.outsed[v] - SS
-
-        rs.SSconc[v] = SS * toconc
-        rs.Bedconc[v] = Bed * toconc
-    end
-    return nothing
-end
diff --git a/src/sediment/erosion/erosion_process.jl b/src/sediment/erosion/erosion_process.jl
new file mode 100644
index 000000000..b50f55b25
--- /dev/null
+++ b/src/sediment/erosion/erosion_process.jl
@@ -0,0 +1,300 @@
+"""
+    rainfall_erosion_eurosem(
+        precip,
+        interception,
+        waterlevel,
+        soil_detachability,
+        eurosem_exponent,
+        canopyheight,
+        canopygapfraction,
+        soilcover_fraction,
+        area,
+        dt,
+    )
+
+Rainfall erosion model based on EUROSEM.
+
+# Arguments
+- `precip` (precipitation [mm Δt⁻¹])
+- `interception` (interception [mm Δt⁻¹])
+- `waterlevel` (water level [m])
+- `soil_detachability` (soil detachability [-])
+- `eurosem_exponent` (EUROSEM exponent [-])
+- `canopyheight` (canopy height [m])
+- `canopygapfraction` (canopy gap fraction [-])
+- `soilcover_fraction` (soil cover fraction [-])
+- `area` (area [m2])
+- `dt` (timestep [seconds])
+
+# Output
+- `rainfall_erosion` (soil loss [t Δt⁻¹])
+"""
+function rainfall_erosion_eurosem(
+    precip,
+    interception,
+    waterlevel,
+    soil_detachability,
+    eurosem_exponent,
+    canopyheight,
+    canopygapfraction,
+    soilcover_fraction,
+    area,
+    dt,
+)
+    # calculate rainfall intensity [mm/h]
+    rintnsty = precip / (dt / 3600)
+    # Kinetic energy of direct throughfall [J/m2/mm]
+    # kedir = max(11.87 + 8.73 * log10(max(0.0001, rintnsty)),0.0) #basis used in USLE
+    kedir = max(8.95 + 8.44 * log10(max(0.0001, rintnsty)), 0.0) #variant used in most distributed mdoels
+    # Kinetic energy of leaf drainage [J/m2/mm]
+    pheff = 0.5 * canopyheight
+    keleaf = max((15.8 * pheff^0.5) - 5.87, 0.0)
+
+    #Depths of rainfall (total, leaf drianage, direct) [mm]
+    rdtot = precip
+    rdleaf = rdtot * 0.1 * canopygapfraction #stemflow
+    rddir = max(rdtot - rdleaf - interception, 0.0) #throughfall
+
+    #Total kinetic energy by rainfall [J/m2]
+    ketot = (rddir * kedir + rdleaf * keleaf) * 0.001
+    # Rainfall / splash erosion [g/m2]
+    rainfall_erosion = soil_detachability * ketot * exp(-eurosem_exponent * waterlevel)
+    rainfall_erosion = rainfall_erosion * area * 1e-6 # ton/cell
+
+    # Remove the impervious area
+    rainfall_erosion = rainfall_erosion * (1.0 - soilcover_fraction)
+    return rainfall_erosion
+end
+
+"""
+    rainfall_erosion_answers(
+        precip,
+        usle_k,
+        usle_c,
+        area,
+        dt,
+    )
+
+Rainfall erosion model based on ANSWERS.
+
+# Arguments
+- `precip` (precipitation [mm Δt⁻¹])
+- `usle_k` (USLE soil erodibility [t ha-1 mm-1])
+- `usle_c` (USLE cover and management factor [-])
+- `soilcover_fraction` (soil cover fraction [-])
+- `area` (area [m2])
+- `dt` (timestep [seconds])
+
+# Output
+- `rainfall_erosion` (soil loss [t Δt⁻¹])
+"""
+function rainfall_erosion_answers(precip, usle_k, usle_c, area, dt)
+    # calculate rainfall intensity [mm/min]
+    rintnsty = precip / (dt / 60)
+    # splash erosion [kg/min]
+    rainfall_erosion = 0.108 * usle_c * usle_k * area * rintnsty^2
+    # [ton/timestep]
+    rainfall_erosion = rainfall_erosion * (dt / 60) * 1e-3
+    return rainfall_erosion
+end
+
+"""
+    overland_flow_erosion_answers(
+        overland_flow,
+        waterlevel,
+        usle_k,
+        usle_c,
+        answers_k,
+        slope,
+        soilcover_fraction,
+        area,
+        dt,
+    )
+
+Overland flow erosion model based on ANSWERS.
+
+# Arguments
+- `overland_flow` (overland flow [m3 s-1])
+- `waterlevel` (water level [m])
+- `usle_k` (USLE soil erodibility [t ha-1 mm-1])
+- `usle_c` (USLE cover and management factor [-])
+- `answers_k` (ANSWERS overland flow factor [-])
+- `slope` (slope [-])
+- `area` (area [m2])
+- `dt` (timestep [seconds])
+
+# Output
+- `overland_flow_erosion` (soil loss [t Δt⁻¹])
+"""
+function overland_flow_erosion_answers(
+    overland_flow,
+    usle_k,
+    usle_c,
+    answers_k,
+    slope,
+    area,
+    dt,
+)
+    # Overland flow rate [m2/min]
+    qr_land = overland_flow * 60 / (area .^ 0.5)
+    # Sine of the slope
+    sinslope = sin(atan(slope))
+
+    # Overland flow erosion [kg/min]
+    # For a wide range of slope, it is better to use the sine of slope rather than tangeant
+    erosion = answers_k * usle_c * usle_k * area * sinslope * qr_land
+    # [ton/timestep]
+    erosion = erosion * (dt / 60) * 1e-3
+    return erosion
+end
+
+"""
+    total_soil_erosion(
+        rainfall_erosion,
+        overland_flow_erosion,
+        clay_fraction,
+        silt_fraction,
+        sand_fraction,
+        sagg_fraction,
+        lagg_fraction,
+    )
+
+Calculate total soil erosion and particle differentiation.
+
+# Arguments
+- `rainfall_erosion` (soil loss from rainfall erosion [t Δt⁻¹])
+- `overland_flow_erosion` (soil loss from overland flow erosion [t Δt⁻¹])
+- `clay_fraction` (clay fraction [-])
+- `silt_fraction` (silt fraction [-])
+- `sand_fraction` (sand fraction [-])
+- `sagg_fraction` (small aggregates fraction [-])
+- `lagg_fraction` (large aggregates fraction [-])
+
+# Output
+- `soil_erosion` (total soil loss [t Δt⁻¹])
+- `clay_erosion` (clay loss [t Δt⁻¹])
+- `silt_erosion` (silt loss [t Δt⁻¹])
+- `sand_erosion` (sand loss [t Δt⁻¹])
+- `sagg_erosion` (small aggregates loss [t Δt⁻¹])
+- `lagg_erosion` (large aggregates loss [t Δt⁻¹])
+"""
+function total_soil_erosion(
+    rainfall_erosion,
+    overland_flow_erosion,
+    clay_fraction,
+    silt_fraction,
+    sand_fraction,
+    sagg_fraction,
+    lagg_fraction,
+)
+    # Total soil erosion
+    soil_erosion = rainfall_erosion + overland_flow_erosion
+    # Particle differentiation
+    clay_erosion = soil_erosion * clay_fraction
+    silt_erosion = soil_erosion * silt_fraction
+    sand_erosion = soil_erosion * sand_fraction
+    sagg_erosion = soil_erosion * sagg_fraction
+    lagg_erosion = soil_erosion * lagg_fraction
+    return soil_erosion,
+    clay_erosion,
+    silt_erosion,
+    sand_erosion,
+    sagg_erosion,
+    lagg_erosion
+end
+
+"""
+    function river_erosion_julian_torres(
+        waterlevel,
+        d50,
+        width,
+        length,
+        slope,
+        dt,
+    )
+
+River erosion model based on Julian Torres.
+Repartition of the effective shear stress between the bank and the bed from Knight et al. 1984 [%]
+
+# Arguments
+- `waterlevel` (water level [m])
+- `d50` (median grain size [m])
+- `width` (width [m])
+- `length` (length [m])
+- `slope` (slope [-])
+- `dt` (timestep [seconds])
+
+# Output
+- `bed` (potential river erosion [t Δt⁻¹])
+- `bank` (potential bank erosion [t Δt⁻¹])
+"""
+function river_erosion_julian_torres(waterlevel, d50, width, length, slope, dt)
+    if waterlevel > 0.0
+        # Bed and Bank from Shields diagram, Da Silva & Yalin (2017)
+        E_ = (2.65 - 1) * 9.81
+        E = (E_ * (d50 * 1e-3)^3 / 1e-12)^0.33
+        TCrbed =
+            E_ *
+            d50 *
+            (0.13 * E^(-0.392) * exp(-0.015 * E^2) + 0.045 * (1 - exp(-0.068 * E)))
+        TCrbank = TCrbed
+        # kd from Hanson & Simon 2001
+        kdbank = 0.2 * TCrbank^(-0.5) * 1e-6
+        kdbed = 0.2 * TCrbed^(-0.5) * 1e-6
+
+        # Hydraulic radius of the river [m] (rectangular channel)
+        hydrad = waterlevel * width / (width + 2 * waterlevel)
+
+        # Repartition of the effective shear stress between the bank and the Bed
+        SFbank = exp(-3.23 * log10(width / waterlevel + 3) + 6.146)
+        # Effective shear stress on river bed and banks [N/m2]
+        TEffbank =
+            1000 * 9.81 * hydrad * slope * SFbank / 100 * (1 + width / (2 * waterlevel))
+        TEffbed =
+            1000 * 9.81 * hydrad * slope * (1 - SFbank / 100) * (1 + 2 * waterlevel / width)
+
+        # Potential erosion rates of the bed and bank [t/cell/timestep]
+        #(assuming only one bank is eroding)
+        Tex = max(TEffbank - TCrbank, 0.0)
+        # 1.4 is bank default bulk density
+        ERbank = kdbank * Tex * length * waterlevel * 1.4 * dt
+        # 1.5 is bed default bulk density
+        ERbed = kdbed * (TEffbed - TCrbed) * length * width * 1.5 * dt
+
+        # Potential maximum bed/bank erosion
+        bed = max(ERbed, 0.0)
+        bank = max(ERbank, 0.0)
+
+    else
+        bed = 0.0
+        bank = 0.0
+    end
+
+    return bed, bank
+end
+
+"""
+    function river_erosion_store(
+        excess_sediment,
+        store,
+    )
+
+River erosion of the previously deposited sediment.
+
+# Arguments
+- `excess_sediment` (excess sediment [t Δt⁻¹])
+- `store` (sediment store [t])
+
+# Output
+- `erosion` (river erosion [t Δt⁻¹])
+- `excess_sediment` (updated excess sediment [t Δt⁻¹])
+- `store` (updated sediment store [t])
+"""
+function river_erosion_store(excess_sediment, store)
+    # River erosion of the previously deposited sediment
+    erosion = min(store, excess_sediment)
+    # Update the excess sediment and the sediment store
+    excess_sediment -= erosion
+    store -= erosion
+    return erosion, excess_sediment, store
+end
\ No newline at end of file
diff --git a/src/sediment/erosion/overland_flow_erosion.jl b/src/sediment/erosion/overland_flow_erosion.jl
new file mode 100644
index 000000000..879af50fd
--- /dev/null
+++ b/src/sediment/erosion/overland_flow_erosion.jl
@@ -0,0 +1,118 @@
+abstract type AbstractOverlandFlowErosionModel{T} end
+
+"Struct for storing overland flow erosion model variables"
+@get_units @grid_loc @with_kw struct OverlandFlowErosionVariables{T}
+    # Total soil erosion from overland flow
+    amount::Vector{T} | "t dt-1"
+end
+
+"Initialize overland flow erosion model variables"
+function OverlandFlowErosionVariables(n; amount::Vector{T} = fill(mv, n)) where {T}
+    return OverlandFlowErosionVariables{T}(; amount = amount)
+end
+
+"Struct for storing overland flow erosion model boundary conditions"
+@get_units @grid_loc @with_kw struct OverlandFlowErosionBC{T}
+    # Overland flow [m3 s-1]
+    q::Vector{T}
+end
+
+"Initialize overland flow erosion model boundary conditions"
+function OverlandFlowErosionBC(n; q::Vector{T} = fill(mv, n)) where {T}
+    return OverlandFlowErosionBC{T}(; q = q)
+end
+
+"Struct for storing ANSWERS overland flow erosion model parameters"
+@get_units @grid_loc @with_kw struct OverlandFlowErosionAnswersParameters{T}
+    # Soil erodibility factor
+    usle_k::Vector{T} | "-"
+    # Crop management factor
+    usle_c::Vector{T} | "-"
+    # Answers overland flow factor
+    answers_k::Vector{T} | "-"
+end
+
+"Initialize ANSWERS overland flow erosion model parameters"
+function OverlandFlowErosionAnswersParameters(dataset, config, indices)
+    usle_k = ncread(
+        dataset,
+        config,
+        "vertical.overland_flow_erosion.parameters.usle_k";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    usle_c = ncread(
+        dataset,
+        config,
+        "vertical.overland_flow_erosion.parameters.usle_c";
+        sel = indices,
+        defaults = 0.01,
+        type = Float,
+    )
+    answers_k = ncread(
+        dataset,
+        config,
+        "vertical.overland_flow_erosion.parameters.answers_k";
+        sel = indices,
+        defaults = 0.9,
+        type = Float,
+    )
+    answers_parameters = OverlandFlowErosionAnswersParameters(;
+        usle_k = usle_k,
+        usle_c = usle_c,
+        answers_k = answers_k,
+    )
+    return answers_parameters
+end
+
+"ANSWERS overland flow erosion model"
+@with_kw struct OverlandFlowErosionAnswersModel{T} <: AbstractOverlandFlowErosionModel{T}
+    boundary_conditions::OverlandFlowErosionBC{T}
+    parameters::OverlandFlowErosionAnswersParameters{T}
+    variables::OverlandFlowErosionVariables{T}
+end
+
+"Initialize ANSWERS overland flow erosion model"
+function OverlandFlowErosionAnswersModel(dataset, config, indices)
+    n = length(indices)
+    vars = OverlandFlowErosionVariables(n)
+    params = OverlandFlowErosionAnswersParameters(dataset, config, indices)
+    bc = OverlandFlowErosionBC(n)
+    model = OverlandFlowErosionAnswersModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update boundary conditions for ANSWERS overland flow erosion model"
+function update_boundary_conditions!(
+    model::OverlandFlowErosionAnswersModel,
+    hydrological_forcing::HydrologicalForcing,
+)
+    (; q) = model.boundary_conditions
+    (; q_land) = hydrological_forcing
+    @. q = q_land
+end
+
+"Update ANSWERS overland flow erosion model for a single timestep"
+function update!(model::OverlandFlowErosionAnswersModel, geometry::LandGeometry, dt)
+    (; q) = model.boundary_conditions
+    (; usle_k, usle_c, answers_k) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = overland_flow_erosion_answers(
+            q[i],
+            usle_k[i],
+            usle_c[i],
+            answers_k[i],
+            geometry.slope[i],
+            geometry.area[i],
+            dt,
+        )
+    end
+end
\ No newline at end of file
diff --git a/src/sediment/erosion/rainfall_erosion.jl b/src/sediment/erosion/rainfall_erosion.jl
new file mode 100644
index 000000000..73016f874
--- /dev/null
+++ b/src/sediment/erosion/rainfall_erosion.jl
@@ -0,0 +1,255 @@
+abstract type AbstractRainfallErosionModel{T} end
+
+"Struct for storing rainfall erosion model variables"
+@get_units @grid_loc @with_kw struct RainfallErosionModelVariables{T}
+    # Total soil erosion from rainfall (splash)
+    amount::Vector{T} | "t dt-1"
+end
+
+"Initialize rainfall erosion model variables"
+function RainfallErosionModelVariables(n; amount::Vector{T} = fill(mv, n)) where {T}
+    return RainfallErosionModelVariables{T}(; amount = amount)
+end
+
+"Struct for storing EUROSEM rainfall erosion model boundary conditions"
+@get_units @grid_loc @with_kw struct RainfallErosionEurosemBC{T}
+    # precipitation
+    precipitation::Vector{T} | "mm dt-1"
+    # Interception
+    interception::Vector{T} | "mm dt-1"
+    # Waterlevel on land
+    waterlevel::Vector{T} | "m"
+end
+
+"Initialize EUROSEM rainfall erosion model boundary conditions"
+function RainfallErosionEurosemBC(
+    n;
+    precipitation::Vector{T} = fill(mv, n),
+    interception::Vector{T} = fill(mv, n),
+    waterlevel::Vector{T} = fill(mv, n),
+) where {T}
+    return RainfallErosionEurosemBC{T}(;
+        precipitation = precipitation,
+        interception = interception,
+        waterlevel = waterlevel,
+    )
+end
+
+"Struct for storing EUROSEM rainfall erosion model parameters"
+@get_units @grid_loc @with_kw struct RainfallErosionEurosemParameters{T}
+    # Soil detachability factor
+    soil_detachability::Vector{T} | "g J-1"
+    # Exponent EUROSEM
+    eurosem_exponent::Vector{T} | "-"
+    # Canopy height
+    canopyheight::Vector{T} | "m"
+    # Canopy gap fraction
+    canopygapfraction::Vector{T} | "-"
+    # Fraction of the soil that is covered (eg paved, snow, etc)
+    soilcover_fraction::Vector{T} | "-"
+end
+
+"Initialize EUROSEM rainfall erosion model parameters"
+function RainfallErosionEurosemParameters(dataset, config, indices)
+    soil_detachability = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.soil_detachability";
+        sel = indices,
+        defaults = 0.6,
+        type = Float,
+    )
+    eurosem_exponent = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.eurosem_exponent";
+        sel = indices,
+        defaults = 2.0,
+        type = Float,
+    )
+    canopyheight = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.canopyheight";
+        sel = indices,
+        defaults = 0.5,
+        type = Float,
+    )
+    canopygapfraction = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.canopygapfraction";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    soilcover_fraction = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.pathfrac";
+        sel = indices,
+        defaults = 0.01,
+        type = Float,
+    )
+    eurosem_parameters = RainfallErosionEurosemParameters(;
+        soil_detachability = soil_detachability,
+        eurosem_exponent = eurosem_exponent,
+        canopyheight = canopyheight,
+        canopygapfraction = canopygapfraction,
+        soilcover_fraction = soilcover_fraction,
+    )
+    return eurosem_parameters
+end
+
+"EUROSEM rainfall erosion model"
+@with_kw struct RainfallErosionEurosemModel{T} <: AbstractRainfallErosionModel{T}
+    boundary_conditions::RainfallErosionEurosemBC{T}
+    parameters::RainfallErosionEurosemParameters{T}
+    variables::RainfallErosionModelVariables{T}
+end
+
+"Initialize EUROSEM rainfall erosion model"
+function RainfallErosionEurosemModel(dataset, config, indices)
+    n = length(indices)
+    vars = RainfallErosionModelVariables(n)
+    params = RainfallErosionEurosemParameters(dataset, config, indices)
+    bc = RainfallErosionEurosemBC(n)
+    model = RainfallErosionEurosemModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update EUROSEM rainfall erosion model boundary conditions for a single timestep"
+function update_boundary_conditions!(
+    model::RainfallErosionEurosemModel,
+    atmospheric_forcing::AtmosphericForcing,
+    hydrological_forcing::HydrologicalForcing,
+)
+    (; precipitation, interception, waterlevel) = model.boundary_conditions
+    @. precipitation = atmospheric_forcing.precipitation
+    @. waterlevel = hydrological_forcing.waterlevel_land
+    @. interception = hydrological_forcing.interception
+end
+
+"Update EUROSEM rainfall erosion model for a single timestep"
+function update!(model::RainfallErosionEurosemModel, geometry::LandGeometry, dt)
+    (; precipitation, interception, waterlevel) = model.boundary_conditions
+    (;
+        soil_detachability,
+        eurosem_exponent,
+        canopyheight,
+        canopygapfraction,
+        soilcover_fraction,
+    ) = model.parameters
+    (; amount) = model.variables
+
+    n = length(precipitation)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = rainfall_erosion_eurosem(
+            precipitation[i],
+            interception[i],
+            waterlevel[i],
+            soil_detachability[i],
+            eurosem_exponent[i],
+            canopyheight[i],
+            canopygapfraction[i],
+            soilcover_fraction[i],
+            geometry.area[i],
+            dt,
+        )
+    end
+end
+
+"Struct for storing ANSWERS rainfall erosion model boundary conditions"
+@get_units @grid_loc @with_kw struct RainfallErosionAnswersBC{T}
+    # precipitation
+    precipitation::Vector{T} | "mm dt-1"
+end
+
+"Initialize ANSWERS rainfall erosion model boundary conditions"
+function RainfallErosionAnswersBC(n; precipitation::Vector{T} = fill(mv, n)) where {T}
+    return RainfallErosionAnswersBC{T}(; precipitation = precipitation)
+end
+
+"Struct for storing ANSWERS rainfall erosion model parameters"
+@get_units @grid_loc @with_kw struct RainfallErosionAnswersParameters{T}
+    # Soil erodibility factor
+    usle_k::Vector{T} | "-"
+    # Crop management factor
+    usle_c::Vector{T} | "-"
+end
+
+"Initialize ANSWERS rainfall erosion model parameters"
+function RainfallErosionAnswersParameters(dataset, config, indices)
+    usle_k = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.usle_k";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    usle_c = ncread(
+        dataset,
+        config,
+        "vertical.rainfall_erosion.parameters.usle_c";
+        sel = indices,
+        defaults = 0.01,
+        type = Float,
+    )
+    answers_parameters =
+        RainfallErosionAnswersParameters(; usle_k = usle_k, usle_c = usle_c)
+    return answers_parameters
+end
+
+"ANSWERS rainfall erosion model"
+@with_kw struct RainfallErosionAnswersModel{T} <: AbstractRainfallErosionModel{T}
+    boundary_conditions::RainfallErosionAnswersBC{T}
+    parameters::RainfallErosionAnswersParameters{T}
+    variables::RainfallErosionModelVariables{T}
+end
+
+"Initialize ANSWERS rainfall erosion model"
+function RainfallErosionAnswersModel(dataset, config, indices)
+    n = length(indices)
+    bc = RainfallErosionAnswersBC(n)
+    vars = RainfallErosionModelVariables(n)
+    params = RainfallErosionAnswersParameters(dataset, config, indices)
+    model = RainfallErosionAnswersModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update ANSWERS rainfall erosion model boundary conditions for a single timestep"
+function update_boundary_conditions!(
+    model::RainfallErosionAnswersModel,
+    atmospheric_forcing::AtmosphericForcing,
+    hydrological_forcing::HydrologicalForcing,
+)
+    (; precipitation) = model.boundary_conditions
+    @. precipitation = atmospheric_forcing.precipitation
+end
+
+"Update ANSWERS rainfall erosion model for a single timestep"
+function update!(model::RainfallErosionAnswersModel, geometry::LandGeometry, dt)
+    (; precipitation) = model.boundary_conditions
+    (; usle_k, usle_c) = model.parameters
+    (; amount) = model.variables
+
+    n = length(precipitation)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = rainfall_erosion_answers(
+            precipitation[i],
+            usle_k[i],
+            usle_c[i],
+            geometry.area[i],
+            dt,
+        )
+    end
+end
\ No newline at end of file
diff --git a/src/sediment/erosion/river_erosion.jl b/src/sediment/erosion/river_erosion.jl
new file mode 100644
index 000000000..9e6e470e5
--- /dev/null
+++ b/src/sediment/erosion/river_erosion.jl
@@ -0,0 +1,100 @@
+abstract type AbstractRiverErosionModel{T} end
+
+"Struct for storing river bed and bank erosion model variables"
+@get_units @grid_loc @with_kw struct RiverErosionModelVariables{T}
+    # Potential river bed erosion
+    bed::Vector{T} | "t dt-1"
+    # Potential river bank erosion
+    bank::Vector{T} | "t dt-1"
+end
+
+"Initialize river bed and bank erosion model variables"
+function RiverErosionModelVariables(
+    n;
+    bed::Vector{T} = fill(mv, n),
+    bank::Vector{T} = fill(mv, n),
+) where {T}
+    return RiverErosionModelVariables{T}(; bed = bed, bank = bank)
+end
+
+"Struct for storing river erosion model boundary conditions"
+@get_units @grid_loc @with_kw struct RiverErosionBC{T}
+    # Waterlevel
+    waterlevel::Vector{T} | "t dt-1"
+end
+
+"Initialize river erosion model boundary conditions"
+function RiverErosionBC(n; waterlevel::Vector{T} = fill(mv, n)) where {T}
+    return RiverErosionBC{T}(; waterlevel = waterlevel)
+end
+
+"Struct for storing river erosion model parameters"
+@get_units @grid_loc @with_kw struct RiverErosionParameters{T}
+    # Mean diameter in the river bed/bank
+    d50::Vector{T} | "mm"
+end
+
+"Julian and Torres river erosion model"
+@with_kw struct RiverErosionJulianTorresModel{T} <: AbstractRiverErosionModel{T}
+    boundary_conditions::RiverErosionBC{T}
+    parameters::RiverErosionParameters{T}
+    variables::RiverErosionModelVariables{T}
+end
+
+"Initialize Julian and Torres river erosion parameters"
+function RiverErosionParameters(dataset, config, indices)
+    d50 = ncread(
+        dataset,
+        config,
+        "lateral.river.potential_erosion.parameters.d50";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    river_parameters = RiverErosionParameters(; d50 = d50)
+
+    return river_parameters
+end
+
+"Initialize Julian and Torres river erosion model"
+function RiverErosionJulianTorresModel(dataset, config, indices)
+    n = length(indices)
+    vars = RiverErosionModelVariables(n)
+    params = RiverErosionParameters(dataset, config, indices)
+    bc = RiverErosionBC(n)
+    model = RiverErosionJulianTorresModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update river erosion model boundary conditions"
+function update_boundary_conditions!(
+    model::RiverErosionJulianTorresModel,
+    hydrological_forcing::HydrologicalForcing,
+)
+    (; waterlevel) = model.boundary_conditions
+    (; waterlevel_river) = hydrological_forcing
+    @. waterlevel = waterlevel_river
+end
+
+"Update Julian and Torres river erosion model for a single timestep"
+function update!(model::RiverErosionJulianTorresModel, geometry::RiverGeometry, dt)
+    (; waterlevel) = model.boundary_conditions
+    (; d50) = model.parameters
+    (; bed, bank) = model.variables
+
+    n = length(waterlevel)
+    threaded_foreach(1:n; basesize = 1000) do i
+        bed[i], bank[i] = river_erosion_julian_torres(
+            waterlevel[i],
+            d50[i],
+            geometry.width[i],
+            geometry.length[i],
+            geometry.slope[i],
+            dt,
+        )
+    end
+end
\ No newline at end of file
diff --git a/src/sediment/erosion/soil_erosion.jl b/src/sediment/erosion/soil_erosion.jl
new file mode 100644
index 000000000..6578a0dbe
--- /dev/null
+++ b/src/sediment/erosion/soil_erosion.jl
@@ -0,0 +1,182 @@
+abstract type AbstractSoilErosionModel{T} end
+
+"Struct for storing total soil erosion with differentiation model variables"
+@get_units @grid_loc @with_kw struct SoilErosionModelVariables{T}
+    # Total soil erosion
+    amount::Vector{T} | "t dt-1"
+    # Total clay erosion
+    clay::Vector{T} | "t dt-1"
+    # Total silt erosion
+    silt::Vector{T} | "t dt-1"
+    # Total sand erosion
+    sand::Vector{T} | "t dt-1"
+    # Total small aggregates erosion
+    sagg::Vector{T} | "t dt-1"
+    # Total large aggregates erosion
+    lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize soil erosion model variables"
+function SoilErosionModelVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(mv, n),
+    silt::Vector{T} = fill(mv, n),
+    sand::Vector{T} = fill(mv, n),
+    sagg::Vector{T} = fill(mv, n),
+    lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return SoilErosionModelVariables{T}(;
+        amount = amount,
+        clay = clay,
+        silt = silt,
+        sand = sand,
+        sagg = sagg,
+        lagg = lagg,
+    )
+end
+
+"Struct for storing soil erosion model boundary conditions"
+@get_units @grid_loc @with_kw struct SoilErosionBC{T}
+    # Rainfall erosion
+    rainfall_erosion::Vector{T} | "t dt-1"
+    # Overland flow erosion
+    overland_flow_erosion::Vector{T} | "m dt-1"
+end
+
+"Initialize soil erosion model boundary conditions"
+function SoilErosionBC(
+    n;
+    rainfall_erosion::Vector{T} = fill(mv, n),
+    overland_flow_erosion::Vector{T} = fill(mv, n),
+) where {T}
+    return SoilErosionBC{T}(;
+        rainfall_erosion = rainfall_erosion,
+        overland_flow_erosion = overland_flow_erosion,
+    )
+end
+
+"Struct for storing soil erosion model parameters"
+@get_units @grid_loc @with_kw struct SoilErosionParameters{T}
+    # Soil content clay
+    clay_fraction::Vector{T} | "-"
+    # Soil content silt
+    silt_fraction::Vector{T} | "-"
+    # Soil content sand
+    sand_fraction::Vector{T} | "-"
+    # Soil content small aggregates
+    sagg_fraction::Vector{T} | "-"
+    # Soil content large aggregates
+    lagg_fraction::Vector{T} | "-"
+end
+
+"Initialize soil erosion model parameters"
+function SoilErosionParameters(dataset, config, indices)
+    clay_fraction = ncread(
+        dataset,
+        config,
+        "vertical.soil_erosion.parameters.clay_fraction";
+        sel = indices,
+        defaults = 0.4,
+        type = Float,
+    )
+    silt_fraction = ncread(
+        dataset,
+        config,
+        "vertical.soil_erosion.parameters.silt_fraction";
+        sel = indices,
+        defaults = 0.3,
+        type = Float,
+    )
+    sand_fraction = ncread(
+        dataset,
+        config,
+        "vertical.soil_erosion.parameters.sand_fraction";
+        sel = indices,
+        defaults = 0.3,
+        type = Float,
+    )
+    sagg_fraction = ncread(
+        dataset,
+        config,
+        "vertical.soil_erosion.parameters.sagg_fraction";
+        sel = indices,
+        defaults = 0.0,
+        type = Float,
+    )
+    lagg_fraction = ncread(
+        dataset,
+        config,
+        "vertical.soil_erosion.parameters.lagg_fraction";
+        sel = indices,
+        defaults = 0.0,
+        type = Float,
+    )
+    # Check that soil fractions sum to 1
+    soil_fractions =
+        clay_fraction + silt_fraction + sand_fraction + sagg_fraction + lagg_fraction
+    if any(abs.(soil_fractions .- 1.0) .> 1e-3)
+        error("Particle fractions in the soil must sum to 1")
+    end
+    soil_parameters = SoilErosionParameters(;
+        clay_fraction = clay_fraction,
+        silt_fraction = silt_fraction,
+        sand_fraction = sand_fraction,
+        sagg_fraction = sagg_fraction,
+        lagg_fraction = lagg_fraction,
+    )
+
+    return soil_parameters
+end
+
+"Total soil erosion with differentiation model"
+@with_kw struct SoilErosionModel{T} <: AbstractSoilErosionModel{T}
+    boundary_conditions::SoilErosionBC{T}
+    parameters::SoilErosionParameters{T}
+    variables::SoilErosionModelVariables{T}
+end
+
+"Initialize soil erosion model"
+function SoilErosionModel(dataset, config, indices)
+    n = length(indices)
+    vars = SoilErosionModelVariables(n)
+    params = SoilErosionParameters(dataset, config, indices)
+    bc = SoilErosionBC(n)
+    model =
+        SoilErosionModel(; boundary_conditions = bc, parameters = params, variables = vars)
+    return model
+end
+
+"Update boundary conditions for soil erosion model"
+function update_boundary_conditions!(
+    model::SoilErosionModel,
+    rainfall_erosion::AbstractRainfallErosionModel,
+    overland_flow_erosion::AbstractOverlandFlowErosionModel,
+)
+    re = rainfall_erosion.variables.amount
+    ole = overland_flow_erosion.variables.amount
+    (; rainfall_erosion, overland_flow_erosion) = model.boundary_conditions
+    @. rainfall_erosion = re
+    @. overland_flow_erosion = ole
+end
+
+"Update soil erosion model for a single timestep"
+function update!(model::SoilErosionModel)
+    (; rainfall_erosion, overland_flow_erosion) = model.boundary_conditions
+    (; clay_fraction, silt_fraction, sand_fraction, sagg_fraction, lagg_fraction) =
+        model.parameters
+    (; amount, clay, silt, sand, sagg, lagg) = model.variables
+
+    n = length(rainfall_erosion)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i], clay[i], silt[i], sand[i], sagg[i], lagg[i] = total_soil_erosion(
+            rainfall_erosion[i],
+            overland_flow_erosion[i],
+            clay_fraction[i],
+            silt_fraction[i],
+            sand_fraction[i],
+            sagg_fraction[i],
+            lagg_fraction[i],
+        )
+    end
+end
\ No newline at end of file
diff --git a/src/sediment/sediment_transport/deposition.jl b/src/sediment/sediment_transport/deposition.jl
new file mode 100644
index 000000000..a533fd357
--- /dev/null
+++ b/src/sediment/sediment_transport/deposition.jl
@@ -0,0 +1,49 @@
+"""
+    reservoir_deposition_camp(
+        input,
+        q,
+        waterlevel,
+        wb_area,
+        wb_trapping_efficiency,
+        dm,
+        slope,
+    )
+
+Deposition of sediment in waterbodies from Camp 1945.
+
+# Arguments
+- `input` (sediment input [t Δt⁻¹])
+- `q` (discharge [m³ Δt⁻¹])
+- `waterlevel` (water level [m])
+- `wb_area` (waterbody area [m²])
+- `wb_trapping_efficiency` (waterbody trapping efficiency [-])
+- `dm` (mean diameter [m])
+- `slope` (slope [-])
+
+# Output
+- `deposition` (deposition [t Δt⁻¹])
+"""
+function reservoir_deposition_camp(
+    input,
+    q,
+    waterlevel,
+    wb_area,
+    wb_trapping_efficiency,
+    dm,
+    slope,
+)
+    # Compute critical velocity
+    vcres = q / wb_area
+    DCres = 411 / 3600 / vcres
+    # Natural deposition
+    deposition = input * min(1.0, (DCres * (dm / 1000)^2))
+
+    # Check if particles are travelling in suspension or bed load using Rouse number
+    dsuspf = 1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * waterlevel * slope)^0.5)^0.5
+    # If bed load, we have extra deposition depending on the reservoir type 
+    if dm > dsuspf
+        deposition = max(deposition, wb_trapping_efficiency * input)
+    end
+
+    return deposition
+end
\ No newline at end of file
diff --git a/src/sediment/sediment_transport/land_to_river.jl b/src/sediment/sediment_transport/land_to_river.jl
new file mode 100644
index 000000000..c4c34971b
--- /dev/null
+++ b/src/sediment/sediment_transport/land_to_river.jl
@@ -0,0 +1,180 @@
+abstract type AbstractSedimentToRiverModel{T} end
+
+"Struct to store total sediment reaching the river model variables"
+@get_units @grid_loc @with_kw struct SedimentToRiverVariables{T}
+    # Total sediment reaching the river
+    amount::Vector{T} | "t dt-1"
+end
+
+"Initialize total sediment reaching the river model variables"
+function SedimentToRiverVariables(n; amount::Vector{T} = fill(mv, n)) where {T}
+    return SedimentToRiverVariables{T}(; amount = amount)
+end
+
+"Struct to store total sediment reaching the river model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentToRiverBC{T}
+    # Deposited material
+    deposition::Vector{T} | "t dt-1"
+end
+
+"Initialize total sediment reaching the river model boundary conditions"
+function SedimentToRiverBC(n; deposition::Vector{T} = fill(mv, n)) where {T}
+    return SedimentToRiverBC{T}(; deposition = deposition)
+end
+
+"Struct to store total sediment reaching the river model"
+@with_kw struct SedimentToRiverModel{T} <: AbstractSedimentToRiverModel{T}
+    boundary_conditions::SedimentToRiverBC{T}
+    variables::SedimentToRiverVariables{T}
+end
+
+"Initialize total sediment reaching the river model"
+function SedimentToRiverModel(indices)
+    n = length(indices)
+    vars = SedimentToRiverVariables(n)
+    bc = SedimentToRiverBC(n)
+    model = SedimentToRiverModel(; boundary_conditions = bc, variables = vars)
+    return model
+end
+
+"Update total sediment reaching the river model boundary conditions"
+function update_boundary_conditions!(
+    model::SedimentToRiverModel,
+    transport_model::SedimentLandTransportModel,
+)
+    (; deposition) = model.boundary_conditions
+    @. deposition = transport_model.variables.deposition
+end
+
+"Update total sediment reaching the river model for a single timestep"
+function update!(model::SedimentToRiverModel, rivers)
+    (; deposition) = model.boundary_conditions
+    (; amount) = model.variables
+
+    zeros = fill(0.0, length(amount))
+    amount .= ifelse.(rivers, deposition, zeros)
+end
+
+"Struct to store differentiated sediment reaching the river model variables"
+@get_units @grid_loc @with_kw struct SedimentToRiverDifferentiationVariables{T}
+    # Total sediment flux
+    amount::Vector{T} | "t dt-1"
+    # Clay flux
+    clay::Vector{T} | "t dt-1"
+    # Silt
+    silt::Vector{T} | "t dt-1"
+    # Sand flux
+    sand::Vector{T} | "t dt-1"
+    # Small aggregates flux
+    sagg::Vector{T} | "t dt-1"
+    # Large aggregates flux
+    lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize differentiated sediment reaching the river model variables"
+function SedimentToRiverDifferentiationVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(mv, n),
+    silt::Vector{T} = fill(mv, n),
+    sand::Vector{T} = fill(mv, n),
+    sagg::Vector{T} = fill(mv, n),
+    lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentToRiverDifferentiationVariables{T}(;
+        amount = amount,
+        clay = clay,
+        silt = silt,
+        sand = sand,
+        sagg = sagg,
+        lagg = lagg,
+    )
+end
+
+"Struct to store differentiated sediment reaching the river model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentToRiverDifferentiationBC{T}
+    # Deposited clay
+    deposition_clay::Vector{T} | "t dt-1"
+    # Deposited silt
+    deposition_silt::Vector{T} | "t dt-1"
+    # Deposited sand
+    deposition_sand::Vector{T} | "t dt-1"
+    # Deposited small aggregates
+    deposition_sagg::Vector{T} | "t dt-1"
+    # Deposited large aggregates
+    deposition_lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize differentiated sediment reaching the river model boundary conditions"
+function SedimentToRiverDifferentiationBC(
+    n;
+    deposition_clay::Vector{T} = fill(mv, n),
+    deposition_silt::Vector{T} = fill(mv, n),
+    deposition_sand::Vector{T} = fill(mv, n),
+    deposition_sagg::Vector{T} = fill(mv, n),
+    deposition_lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentToRiverDifferentiationBC{T}(;
+        deposition_clay = deposition_clay,
+        deposition_silt = deposition_silt,
+        deposition_sand = deposition_sand,
+        deposition_sagg = deposition_sagg,
+        deposition_lagg = deposition_lagg,
+    )
+end
+
+"Struct to store differentiated sediment reaching the river model"
+@with_kw struct SedimentToRiverDifferentiationModel{T} <: AbstractSedimentToRiverModel{T}
+    boundary_conditions::SedimentToRiverDifferentiationBC{T}
+    variables::SedimentToRiverDifferentiationVariables{T}
+end
+
+"Initialize differentiated sediment reaching the river model"
+function SedimentToRiverDifferentiationModel(indices)
+    n = length(indices)
+    vars = SedimentToRiverDifferentiationVariables(n)
+    bc = SedimentToRiverDifferentiationBC(n)
+    model =
+        SedimentToRiverDifferentiationModel(; boundary_conditions = bc, variables = vars)
+    return model
+end
+
+"Update differentiated sediment reaching the river model boundary conditions"
+function update_boundary_conditions!(
+    model::SedimentToRiverDifferentiationModel,
+    transport_model::SedimentLandTransportDifferentiationModel,
+)
+    (;
+        deposition_clay,
+        deposition_silt,
+        deposition_sand,
+        deposition_sagg,
+        deposition_lagg,
+    ) = model.boundary_conditions
+    @. deposition_clay = transport_model.variables.deposition_clay
+    @. deposition_silt = transport_model.variables.deposition_silt
+    @. deposition_sand = transport_model.variables.deposition_sand
+    @. deposition_sagg = transport_model.variables.deposition_sagg
+    @. deposition_lagg = transport_model.variables.deposition_lagg
+end
+
+"Update differentiated sediment reaching the river model for a single timestep"
+function update!(model::SedimentToRiverDifferentiationModel, rivers)
+    (;
+        deposition_clay,
+        deposition_silt,
+        deposition_sand,
+        deposition_sagg,
+        deposition_lagg,
+    ) = model.boundary_conditions
+    (; amount, clay, silt, sand, sagg, lagg) = model.variables
+
+    zeros = fill(0.0, length(amount))
+    clay .= ifelse.(rivers .> 0, deposition_clay, zeros)
+    silt .= ifelse.(rivers .> 0, deposition_silt, zeros)
+    sand .= ifelse.(rivers .> 0, deposition_sand, zeros)
+    sagg .= ifelse.(rivers .> 0, deposition_sagg, zeros)
+    lagg .= ifelse.(rivers .> 0, deposition_lagg, zeros)
+
+    amount .= clay .+ silt .+ sand .+ sagg .+ lagg
+end
diff --git a/src/sediment/sediment_transport/overland_flow_transport.jl b/src/sediment/sediment_transport/overland_flow_transport.jl
new file mode 100644
index 000000000..b6ef17ff2
--- /dev/null
+++ b/src/sediment/sediment_transport/overland_flow_transport.jl
@@ -0,0 +1,284 @@
+abstract type AbstractSedimentLandTransportModel{T} end
+
+"Struct to store total sediment flux in overland flow model variables"
+@get_units @grid_loc @with_kw struct SedimentLandTransportVariables{T}
+    # Total sediment flux
+    amount::Vector{T} | "t dt-1"
+    deposition::Vector{T} | "t dt-1"
+end
+
+"Initialize total sediment flux in overland flow model variables"
+function SedimentLandTransportVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    deposition::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentLandTransportVariables{T}(; amount = amount, deposition = deposition)
+end
+
+"Struct to store total sediment flux in overland flow model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentLandTransportBC{T}
+    # Eroded material
+    erosion::Vector{T} | "t dt-1"
+    # Transport capacity
+    transport_capacity::Vector{T} | "t dt-1"
+end
+
+"Initialize total sediment flux in overland flow model boundary conditions"
+function SedimentLandTransportBC(
+    n;
+    erosion::Vector{T} = fill(mv, n),
+    transport_capacity::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentLandTransportBC{T}(;
+        erosion = erosion,
+        transport_capacity = transport_capacity,
+    )
+end
+
+"Struct to store total sediment flux in overland flow model"
+@with_kw struct SedimentLandTransportModel{T} <: AbstractSedimentLandTransportModel{T}
+    boundary_conditions::SedimentLandTransportBC{T}
+    variables::SedimentLandTransportVariables{T}
+end
+
+"Initialize total sediment flux in overland flow model"
+function SedimentLandTransportModel(indices)
+    n = length(indices)
+    vars = SedimentLandTransportVariables(n)
+    bc = SedimentLandTransportBC(n)
+    model = SedimentLandTransportModel(; boundary_conditions = bc, variables = vars)
+    return model
+end
+
+"Update total sediment flux in overland flow model boundary conditions"
+function update_boundary_conditions!(
+    model::SedimentLandTransportModel,
+    erosion_model::SoilErosionModel,
+    transport_capacity_model::AbstractTransportCapacityModel,
+)
+    (; erosion, transport_capacity) = model.boundary_conditions
+    (; amount) = erosion_model.variables
+    @. erosion = amount
+
+    (; amount) = transport_capacity_model.variables
+    @. transport_capacity = amount
+end
+
+"Update total sediment flux in overland flow model for a single timestep"
+function update!(model::SedimentLandTransportModel, network)
+    (; erosion, transport_capacity) = model.boundary_conditions
+    (; amount, deposition) = model.variables
+
+    accucapacityflux!(amount, erosion, network, transport_capacity)
+    deposition .= erosion
+end
+
+"Struct to store differentiated sediment flux in overland flow model variables"
+@get_units @grid_loc @with_kw struct SedimentLandTransportDifferentiationVariables{T}
+    # Total sediment flux
+    amount::Vector{T} | "t dt-1"
+    # Deposition
+    deposition::Vector{T} | "t dt-1"
+    # Clay flux
+    clay::Vector{T} | "t dt-1"
+    # Deposition clay
+    deposition_clay::Vector{T} | "t dt-1"
+    # Silt
+    silt::Vector{T} | "t dt-1"
+    # Deposition silt
+    deposition_silt::Vector{T} | "t dt-1"
+    # Sand flux
+    sand::Vector{T} | "t dt-1"
+    # Deposition sand
+    deposition_sand::Vector{T} | "t dt-1"
+    # Small aggregates flux
+    sagg::Vector{T} | "t dt-1"
+    # Deposition small aggregates
+    deposition_sagg::Vector{T} | "t dt-1"
+    # Large aggregates flux
+    lagg::Vector{T} | "t dt-1"
+    # Deposition large aggregates
+    deposition_lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize differentiated sediment flux in overland flow model variables"
+function SedimentLandTransportDifferentiationVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    deposition::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(mv, n),
+    deposition_clay::Vector{T} = fill(mv, n),
+    silt::Vector{T} = fill(mv, n),
+    deposition_silt::Vector{T} = fill(mv, n),
+    sand::Vector{T} = fill(mv, n),
+    deposition_sand::Vector{T} = fill(mv, n),
+    sagg::Vector{T} = fill(mv, n),
+    deposition_sagg::Vector{T} = fill(mv, n),
+    lagg::Vector{T} = fill(mv, n),
+    deposition_lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentLandTransportDifferentiationVariables{T}(;
+        amount = amount,
+        deposition = deposition,
+        clay = clay,
+        deposition_clay = deposition_clay,
+        silt = silt,
+        deposition_silt = deposition_silt,
+        sand = sand,
+        deposition_sand = deposition_sand,
+        sagg = sagg,
+        deposition_sagg = deposition_sagg,
+        lagg = lagg,
+        deposition_lagg = deposition_lagg,
+    )
+end
+
+"Struct to store differentiated sediment flux in overland flow model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentLandTransportDifferentiationBC{T}
+    # Eroded clay
+    erosion_clay::Vector{T} | "t dt-1"
+    # Eroded silt
+    erosion_silt::Vector{T} | "t dt-1"
+    # Eroded sand
+    erosion_sand::Vector{T} | "t dt-1"
+    # Eroded small aggregates
+    erosion_sagg::Vector{T} | "t dt-1"
+    # Eroded large aggregates
+    erosion_lagg::Vector{T} | "t dt-1"
+    # Transport capacity clay
+    transport_capacity_clay::Vector{T} | "t dt-1"
+    # Transport capacity silt
+    transport_capacity_silt::Vector{T} | "t dt-1"
+    # Transport capacity sand
+    transport_capacity_sand::Vector{T} | "t dt-1"
+    # Transport capacity small aggregates
+    transport_capacity_sagg::Vector{T} | "t dt-1"
+    # Transport capacity large aggregates
+    transport_capacity_lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize differentiated sediment flux in overland flow model boundary conditions"
+function SedimentLandTransportDifferentiationBC(
+    n;
+    erosion_clay::Vector{T} = fill(mv, n),
+    erosion_silt::Vector{T} = fill(mv, n),
+    erosion_sand::Vector{T} = fill(mv, n),
+    erosion_sagg::Vector{T} = fill(mv, n),
+    erosion_lagg::Vector{T} = fill(mv, n),
+    transport_capacity_clay::Vector{T} = fill(mv, n),
+    transport_capacity_silt::Vector{T} = fill(mv, n),
+    transport_capacity_sand::Vector{T} = fill(mv, n),
+    transport_capacity_sagg::Vector{T} = fill(mv, n),
+    transport_capacity_lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentLandTransportDifferentiationBC{T}(;
+        erosion_clay = erosion_clay,
+        erosion_silt = erosion_silt,
+        erosion_sand = erosion_sand,
+        erosion_sagg = erosion_sagg,
+        erosion_lagg = erosion_lagg,
+        transport_capacity_clay = transport_capacity_clay,
+        transport_capacity_silt = transport_capacity_silt,
+        transport_capacity_sand = transport_capacity_sand,
+        transport_capacity_sagg = transport_capacity_sagg,
+        transport_capacity_lagg = transport_capacity_lagg,
+    )
+end
+
+"Struct to store differentiated sediment flux in overland flow model"
+@with_kw struct SedimentLandTransportDifferentiationModel{T} <:
+                AbstractSedimentLandTransportModel{T}
+    boundary_conditions::SedimentLandTransportDifferentiationBC{T}
+    variables::SedimentLandTransportDifferentiationVariables{T}
+end
+
+"Initialize differentiated sediment flux in overland flow model"
+function SedimentLandTransportDifferentiationModel(indices)
+    n = length(indices)
+    vars = SedimentLandTransportDifferentiationVariables(n)
+    bc = SedimentLandTransportDifferentiationBC(n)
+    model = SedimentLandTransportDifferentiationModel(;
+        boundary_conditions = bc,
+        variables = vars,
+    )
+    return model
+end
+
+"Update differentiated sediment flux in overland flow model boundary conditions"
+function update_boundary_conditions!(
+    model::SedimentLandTransportDifferentiationModel,
+    erosion_model::SoilErosionModel,
+    transport_capacity_model::TransportCapacityYalinDifferentiationModel,
+)
+    (;
+        erosion_clay,
+        erosion_silt,
+        erosion_sand,
+        erosion_sagg,
+        erosion_lagg,
+        transport_capacity_clay,
+        transport_capacity_silt,
+        transport_capacity_sand,
+        transport_capacity_sagg,
+        transport_capacity_lagg,
+    ) = model.boundary_conditions
+    (; clay, silt, sand, sagg, lagg) = erosion_model.variables
+    @. erosion_clay = clay
+    @. erosion_silt = silt
+    @. erosion_sand = sand
+    @. erosion_sagg = sagg
+    @. erosion_lagg = lagg
+
+    (; clay, silt, sand, sagg, lagg) = transport_capacity_model.variables
+    @. transport_capacity_clay = clay
+    @. transport_capacity_silt = silt
+    @. transport_capacity_sand = sand
+    @. transport_capacity_sagg = sagg
+    @. transport_capacity_lagg = lagg
+end
+
+"Update differentiated sediment flux in overland flow model for a single timestep"
+function update!(model::SedimentLandTransportDifferentiationModel, network)
+    (;
+        erosion_clay,
+        erosion_silt,
+        erosion_sand,
+        erosion_sagg,
+        erosion_lagg,
+        transport_capacity_clay,
+        transport_capacity_silt,
+        transport_capacity_sand,
+        transport_capacity_sagg,
+        transport_capacity_lagg,
+    ) = model.boundary_conditions
+    (;
+        amount,
+        deposition,
+        clay,
+        deposition_clay,
+        silt,
+        deposition_silt,
+        sand,
+        deposition_sand,
+        sagg,
+        deposition_sagg,
+        lagg,
+        deposition_lagg,
+    ) = model.variables
+
+    accucapacityflux!(clay, erosion_clay, network, transport_capacity_clay)
+    deposition_clay .= erosion_clay
+    accucapacityflux!(silt, erosion_silt, network, transport_capacity_silt)
+    deposition_silt .= erosion_silt
+    accucapacityflux!(sand, erosion_sand, network, transport_capacity_sand)
+    deposition_sand .= erosion_sand
+    accucapacityflux!(sagg, erosion_sagg, network, transport_capacity_sagg)
+    deposition_sagg .= erosion_sagg
+    accucapacityflux!(lagg, erosion_lagg, network, transport_capacity_lagg)
+    deposition_lagg .= erosion_lagg
+    amount .= clay .+ silt .+ sand .+ sagg .+ lagg
+    deposition .=
+        deposition_clay .+ deposition_silt .+ deposition_sand .+ deposition_sagg .+
+        deposition_lagg
+end
\ No newline at end of file
diff --git a/src/sediment/sediment_transport/river_transport.jl b/src/sediment/sediment_transport/river_transport.jl
new file mode 100644
index 000000000..9d5addde5
--- /dev/null
+++ b/src/sediment/sediment_transport/river_transport.jl
@@ -0,0 +1,1009 @@
+abstract type AbstractSedimentRiverTransportModel{T} end
+
+"Struct to store river sediment transport model variables"
+@get_units @grid_loc @with_kw struct SedimentRiverTransportVariables{T}
+    # Sediment flux [ton]
+    amount::Vector{T} | "t dt-1"
+    clay::Vector{T} | "t dt-1"
+    silt::Vector{T} | "t dt-1"
+    sand::Vector{T} | "t dt-1"
+    sagg::Vector{T} | "t dt-1"
+    lagg::Vector{T} | "t dt-1"
+    gravel::Vector{T} | "t dt-1"
+    # Total Sediment deposition [ton]
+    deposition::Vector{T} | "t dt-1"
+    # Total sediment erosion (from store + direct river bed/bank) [ton]
+    erosion::Vector{T} | "t dt-1"
+    # Sediment / particle left in the cell [ton] - states
+    leftover_clay::Vector{T} | "t dt-1"
+    leftover_silt::Vector{T} | "t dt-1"
+    leftover_sand::Vector{T} | "t dt-1"
+    leftover_sagg::Vector{T} | "t dt-1"
+    leftover_lagg::Vector{T} | "t dt-1"
+    leftover_gravel::Vector{T} | "t dt-1"
+    # Sediment / particle stored on the river bed after deposition [ton] -states
+    store_clay::Vector{T} | "t dt-1"
+    store_silt::Vector{T} | "t dt-1"
+    store_sand::Vector{T} | "t dt-1"
+    store_sagg::Vector{T} | "t dt-1"
+    store_lagg::Vector{T} | "t dt-1"
+    store_gravel::Vector{T} | "t dt-1"
+end
+
+"Initialize river sediment transport model variables"
+function SedimentRiverTransportVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(0.0, n),
+    silt::Vector{T} = fill(0.0, n),
+    sand::Vector{T} = fill(0.0, n),
+    sagg::Vector{T} = fill(0.0, n),
+    lagg::Vector{T} = fill(0.0, n),
+    gravel::Vector{T} = fill(0.0, n),
+    deposition::Vector{T} = fill(mv, n),
+    erosion::Vector{T} = fill(mv, n),
+    leftover_clay::Vector{T} = fill(0.0, n),
+    leftover_silt::Vector{T} = fill(0.0, n),
+    leftover_sand::Vector{T} = fill(0.0, n),
+    leftover_sagg::Vector{T} = fill(0.0, n),
+    leftover_lagg::Vector{T} = fill(0.0, n),
+    leftover_gravel::Vector{T} = fill(0.0, n),
+    store_clay::Vector{T} = fill(0.0, n),
+    store_silt::Vector{T} = fill(0.0, n),
+    store_sand::Vector{T} = fill(0.0, n),
+    store_sagg::Vector{T} = fill(0.0, n),
+    store_lagg::Vector{T} = fill(0.0, n),
+    store_gravel::Vector{T} = fill(0.0, n),
+) where {T}
+    return SedimentRiverTransportVariables{T}(;
+        amount = amount,
+        clay = clay,
+        silt = silt,
+        sand = sand,
+        sagg = sagg,
+        lagg = lagg,
+        gravel = gravel,
+        deposition = deposition,
+        erosion = erosion,
+        leftover_clay = leftover_clay,
+        leftover_silt = leftover_silt,
+        leftover_sand = leftover_sand,
+        leftover_sagg = leftover_sagg,
+        leftover_lagg = leftover_lagg,
+        leftover_gravel = leftover_gravel,
+        store_clay = store_clay,
+        store_silt = store_silt,
+        store_sand = store_sand,
+        store_sagg = store_sagg,
+        store_lagg = store_lagg,
+        store_gravel = store_gravel,
+    )
+end
+
+"Struct to store river sediment transport model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentRiverTransportBC{T}
+    # Waterlevel
+    waterlevel::Vector{T} | "t dt-1"
+    # Discharge
+    q::Vector{T} | "m3 s-1"
+    # Transport capacity of the flow
+    transport_capacity::Vector{T} | "t dt-1"
+    # Sediment input from land erosion
+    erosion_land_clay::Vector{T} | "t dt-1"
+    erosion_land_silt::Vector{T} | "t dt-1"
+    erosion_land_sand::Vector{T} | "t dt-1"
+    erosion_land_sagg::Vector{T} | "t dt-1"
+    erosion_land_lagg::Vector{T} | "t dt-1"
+    # Sediment available from direct river erosion
+    potential_erosion_river_bed::Vector{T} | "t dt-1"
+    potential_erosion_river_bank::Vector{T} | "t dt-1"
+end
+
+"Initialize river sediment transport model boundary conditions"
+function SedimentRiverTransportBC(
+    n;
+    waterlevel::Vector{T} = fill(mv, n),
+    q::Vector{T} = fill(mv, n),
+    transport_capacity::Vector{T} = fill(mv, n),
+    erosion_land_clay::Vector{T} = fill(mv, n),
+    erosion_land_silt::Vector{T} = fill(mv, n),
+    erosion_land_sand::Vector{T} = fill(mv, n),
+    erosion_land_sagg::Vector{T} = fill(mv, n),
+    erosion_land_lagg::Vector{T} = fill(mv, n),
+    potential_erosion_river_bed::Vector{T} = fill(mv, n),
+    potential_erosion_river_bank::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentRiverTransportBC{T}(;
+        waterlevel = waterlevel,
+        q = q,
+        transport_capacity = transport_capacity,
+        erosion_land_clay = erosion_land_clay,
+        erosion_land_silt = erosion_land_silt,
+        erosion_land_sand = erosion_land_sand,
+        erosion_land_sagg = erosion_land_sagg,
+        erosion_land_lagg = erosion_land_lagg,
+        potential_erosion_river_bed = potential_erosion_river_bed,
+        potential_erosion_river_bank = potential_erosion_river_bank,
+    )
+end
+
+"Struct to store river sediment transport model parameters"
+@get_units @grid_loc @with_kw struct SedimentRiverTransportParameters{T}
+    # River bed/bank content clay
+    clay_fraction::Vector{T} | "-"
+    # River bed/bank content silt
+    silt_fraction::Vector{T} | "-"
+    # River bed/bank content sand
+    sand_fraction::Vector{T} | "-"
+    # River bed/bank content gravel
+    gravel_fraction::Vector{T} | "-"
+    # Clay mean diameter
+    dm_clay::Vector{T} | "µm"
+    # Silt mean diameter
+    dm_silt::Vector{T} | "µm"
+    # Sand mean diameter
+    dm_sand::Vector{T} | "µm"
+    # Small aggregates mean diameter
+    dm_sagg::Vector{T} | "µm"
+    # Large aggregates mean diameter
+    dm_lagg::Vector{T} | "µm"
+    # Gravel mean diameter
+    dm_gravel::Vector{T} | "µm"
+    # Waterbodies outlets
+    waterbodies_locs::Vector{Bool} | "-"
+    # Waterbodies area
+    waterbodies_area::Vector{T} | "m2"
+    # Waterbodies trapping efficiency
+    waterbodies_trapping_efficiency::Vector{T} | "-"
+end
+
+"Initialize river sediment transport model parameters"
+function SedimentRiverTransportParameters(dataset, config, indices)
+    n = length(indices)
+    clay_fraction = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.clay_fraction";
+        sel = indices,
+        defaults = 0.15,
+        type = Float,
+    )
+    silt_fraction = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.silt_fraction";
+        sel = indices,
+        defaults = 0.65,
+        type = Float,
+    )
+    sand_fraction = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.sand_fraction";
+        sel = indices,
+        defaults = 0.15,
+        type = Float,
+    )
+    gravel_fraction = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.sagg_fraction";
+        sel = indices,
+        defaults = 0.05,
+        type = Float,
+    )
+    # Check that river fractions sum to 1
+    river_fractions = clay_fraction + silt_fraction + sand_fraction + gravel_fraction
+    if any(abs.(river_fractions .- 1.0) .> 1e-3)
+        error("Particle fractions in the river bed must sum to 1")
+    end
+    dm_clay = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_clay";
+        sel = indices,
+        defaults = 2.0,
+        type = Float,
+    )
+    dm_silt = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_silt";
+        sel = indices,
+        defaults = 10.0,
+        type = Float,
+    )
+    dm_sand = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_sand";
+        sel = indices,
+        defaults = 200.0,
+        type = Float,
+    )
+    dm_sagg = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_sagg";
+        sel = indices,
+        defaults = 30.0,
+        type = Float,
+    )
+    dm_lagg = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_lagg";
+        sel = indices,
+        defaults = 500.0,
+        type = Float,
+    )
+    dm_gravel = ncread(
+        dataset,
+        config,
+        "lateral.river.sediment_flux.parameters.dm_gravel";
+        sel = indices,
+        defaults = 2000.0,
+        type = Float,
+    )
+    # Waterbodies
+    wblocs = zeros(Float, n)
+    wbarea = zeros(Float, n)
+    wbtrap = zeros(Float, n)
+    do_reservoirs = get(config.model, "doreservoir", false)::Bool
+    do_lakes = get(config.model, "dolake", false)::Bool
+
+    if do_reservoirs
+        reslocs = ncread(
+            dataset,
+            config,
+            "lateral.river.sediment_flux.parameters.reslocs";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0,
+        )
+        resarea = ncread(
+            dataset,
+            config,
+            "lateral.river.sediment_flux.parameters.resarea";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0.0,
+        )
+        restrapefficiency = ncread(
+            dataset,
+            config,
+            "lateral.river.sediment_flux.parameters.restrapeff";
+            optional = false,
+            sel = indices,
+            type = Float,
+            defaults = 1.0,
+            fill = 0.0,
+        )
+        wblocs = wblocs .+ reslocs
+        wbarea = wbarea .+ resarea
+        wbtrap = wbtrap .+ restrapefficiency
+    end
+
+    if do_lakes
+        lakelocs = ncread(
+            dataset,
+            config,
+            "lateral.river.sediment_flux.parameters.lakelocs";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0,
+        )
+        lakearea = ncread(
+            dataset,
+            config,
+            "lateral.river.sediment_flux.parameters.lakearea";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0.0,
+        )
+        wblocs = wblocs .+ lakelocs
+        wbarea = wbarea .+ lakearea
+    end
+
+    river_parameters = SedimentRiverTransportParameters(;
+        clay_fraction = clay_fraction,
+        silt_fraction = silt_fraction,
+        sand_fraction = sand_fraction,
+        gravel_fraction = gravel_fraction,
+        dm_clay = dm_clay,
+        dm_silt = dm_silt,
+        dm_sand = dm_sand,
+        dm_sagg = dm_sagg,
+        dm_lagg = dm_lagg,
+        dm_gravel = dm_gravel,
+        waterbodies_locs = wblocs .> 0,
+        waterbodies_area = wbarea,
+        waterbodies_trapping_efficiency = wbtrap,
+    )
+
+    return river_parameters
+end
+
+"Struct to store river sediment transport model"
+@with_kw struct SedimentRiverTransportModel{T} <: AbstractSedimentRiverTransportModel{T}
+    boundary_conditions::SedimentRiverTransportBC{T}
+    parameters::SedimentRiverTransportParameters{T}
+    variables::SedimentRiverTransportVariables{T}
+end
+
+"Initialize river sediment transport model"
+function SedimentRiverTransportModel(dataset, config, indices)
+    n = length(indices)
+    vars = SedimentRiverTransportVariables(n)
+    params = SedimentRiverTransportParameters(dataset, config, indices)
+    bc = SedimentRiverTransportBC(n)
+    model = SedimentRiverTransportModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update boundary conditions for river sediment transport model"
+function update_boundary_conditions!(
+    model::SedimentRiverTransportModel,
+    hydrological_forcing::HydrologicalForcing,
+    transport_capacity_model::AbstractTransportCapacityModel,
+    to_river_model::SedimentToRiverDifferentiationModel,
+    potential_erosion_model::AbstractRiverErosionModel,
+    indices_riv,
+)
+    (;
+        waterlevel,
+        q,
+        transport_capacity,
+        erosion_land_clay,
+        erosion_land_silt,
+        erosion_land_sand,
+        erosion_land_sagg,
+        erosion_land_lagg,
+        potential_erosion_river_bed,
+        potential_erosion_river_bank,
+    ) = model.boundary_conditions
+
+    # Hydrological forcing
+    (; q_river, waterlevel_river) = hydrological_forcing
+    @. q = q_river
+    @. waterlevel = waterlevel_river
+    # Transport capacity
+    @. transport_capacity = transport_capacity_model.variables.amount
+    # Input from soil erosion
+    (; clay, silt, sand, sagg, lagg) = to_river_model.variables
+    @. erosion_land_clay = clay[indices_riv]
+    @. erosion_land_silt = silt[indices_riv]
+    @. erosion_land_sand = sand[indices_riv]
+    @. erosion_land_sagg = sagg[indices_riv]
+    @. erosion_land_lagg = lagg[indices_riv]
+    # Maximum direct river bed/bank erosion
+    @. potential_erosion_river_bed = potential_erosion_model.variables.bed
+    @. potential_erosion_river_bank = potential_erosion_model.variables.bank
+end
+
+"Update river sediment transport model for a single timestep"
+function update!(
+    model::SedimentRiverTransportModel,
+    network,
+    geometry::RiverGeometry,
+    waterbodies,
+    dt,
+)
+    (;
+        waterlevel,
+        q,
+        transport_capacity,
+        erosion_land_clay,
+        erosion_land_silt,
+        erosion_land_sand,
+        erosion_land_sagg,
+        erosion_land_lagg,
+        potential_erosion_river_bed,
+        potential_erosion_river_bank,
+    ) = model.boundary_conditions
+    (;
+        clay_fraction,
+        silt_fraction,
+        sand_fraction,
+        gravel_fraction,
+        dm_clay,
+        dm_silt,
+        dm_sand,
+        dm_sagg,
+        dm_lagg,
+        dm_gravel,
+        waterbodies_locs,
+        waterbodies_area,
+        waterbodies_trapping_efficiency,
+    ) = model.parameters
+    (;
+        amount,
+        clay,
+        silt,
+        sand,
+        sagg,
+        lagg,
+        gravel,
+        deposition,
+        erosion,
+        leftover_clay,
+        leftover_silt,
+        leftover_sand,
+        leftover_sagg,
+        leftover_lagg,
+        leftover_gravel,
+        store_clay,
+        store_silt,
+        store_sand,
+        store_sagg,
+        store_lagg,
+        store_gravel,
+    ) = model.variables
+
+    (; graph, order) = network
+
+    # Sediment transport - water balance in the river
+    for v in order
+        ### Sediment input in the cell (left from previous timestep + from land + from upstream outflux) ###
+        input_clay = leftover_clay[v] + erosion_land_clay[v]
+        input_silt = leftover_silt[v] + erosion_land_silt[v]
+        input_sand = leftover_sand[v] + erosion_land_sand[v]
+        input_sagg = leftover_sagg[v] + erosion_land_sagg[v]
+        input_lagg = leftover_lagg[v] + erosion_land_lagg[v]
+        input_gravel = leftover_gravel[v]
+
+        # Add upstream contribution
+        upstream_nodes = inneighbors(graph, v)
+        if !isempty(upstream_nodes)
+            for i in upstream_nodes
+                if clay[i] >= 0.0 # avoid NaN from upstream non-river cells
+                    input_clay += clay[i]
+                    input_silt += silt[i]
+                    input_sand += sand[i]
+                    input_sagg += sagg[i]
+                    input_lagg += lagg[i]
+                    input_gravel += gravel[i]
+                end
+            end
+        end
+
+        input_sediment =
+            input_clay + input_silt + input_sand + input_sagg + input_lagg + input_gravel
+
+        ### River erosion ###
+        # Erosion only if the load is below the transport capacity of the flow.
+        sediment_need = max(transport_capacity[v] - input_sediment, 0.0)
+        # No erosion in reservoirs
+        if waterbodies[v]
+            sediment_need = 0.0
+        end
+
+        # Available sediment stored from previous deposition
+        store_sediment =
+            store_clay[v] +
+            store_silt[v] +
+            store_sand[v] +
+            store_sagg[v] +
+            store_lagg[v] +
+            store_gravel[v]
+
+        # Direct erosion from the river bed/bank
+        if sediment_need > store_sediment
+            # Effective sediment needed fom river bed and bank erosion [ton]
+            effsediment_need = sediment_need - store_sediment
+            # Relative potential erosion rates of the bed and the bank [-]
+            if (potential_erosion_river_bank[v] + potential_erosion_river_bed[v] > 0.0)
+                RTEbank =
+                    potential_erosion_river_bank[v] /
+                    (potential_erosion_river_bank[v] + potential_erosion_river_bed[v])
+            else
+                RTEbank = 0.0
+            end
+            RTEbed = 1.0 - RTEbank
+
+            # Actual bed and bank erosion
+            erosion_bank = min(RTEbank * effsediment_need, potential_erosion_river_bank[v])
+            erosion_bed = min(RTEbed * effsediment_need, potential_erosion_river_bed[v])
+            erosion_river = erosion_bank + erosion_bed
+            # Per particle
+            erosion_clay = erosion_river * clay_fraction[v]
+            erosion_silt = erosion_river * silt_fraction[v]
+            erosion_sand = erosion_river * sand_fraction[v]
+            erosion_gravel = erosion_river * gravel_fraction[v]
+            # No small and large aggregates in the river bed/bank
+            erosion_sagg = 0.0
+            erosion_lagg = 0.0
+        else
+            erosion_clay = 0.0
+            erosion_silt = 0.0
+            erosion_sand = 0.0
+            erosion_gravel = 0.0
+            erosion_sagg = 0.0
+            erosion_lagg = 0.0
+        end
+
+        # Erosion/degradation of the previously deposited sediment (from clay to gravel) [ton]
+        if sediment_need > 0.0
+            # Erosion in priority of the smaller particles
+            # Clay
+            if store_clay[v] > 0.0
+                erosion_store_clay, sediment_need, store_clay[v] =
+                    river_erosion_store(sediment_need, store_clay[v])
+                # Update the clay erosion
+                erosion_clay += erosion_store_clay
+            end
+            # Silt
+            if store_silt[v] > 0.0
+                erosion_store_silt, sediment_need, store_silt[v] =
+                    river_erosion_store(sediment_need, store_silt[v])
+                # Update the silt erosion
+                erosion_silt += erosion_store_silt
+            end
+            # Small aggregates
+            if store_sagg[v] > 0.0
+                erosion_store_sagg, sediment_need, store_sagg[v] =
+                    river_erosion_store(sediment_need, store_sagg[v])
+                # Update the sagg erosion
+                erosion_sagg += erosion_store_sagg
+            end
+            # Sand
+            if store_sand[v] > 0.0
+                erosion_store_sand, sediment_need, store_sand[v] =
+                    river_erosion_store(sediment_need, store_sand[v])
+                # Update the sand erosion
+                erosion_sand += erosion_store_sand
+            end
+            # Large aggregates
+            if store_lagg[v] > 0.0
+                erosion_store_lagg, sediment_need, store_lagg[v] =
+                    river_erosion_store(sediment_need, store_lagg[v])
+                # Update the lagg erosion
+                erosion_lagg += erosion_store_lagg
+            end
+            # Gravel
+            if store_gravel[v] > 0.0
+                erosion_store_gravel, sediment_need, store_gravel[v] =
+                    river_erosion_store(sediment_need, store_gravel[v])
+                # Update the gravel erosion
+                erosion_gravel += erosion_store_gravel
+            end
+        end
+
+        # Compute total erosion
+        erosion[v] =
+            erosion_clay +
+            erosion_silt +
+            erosion_sand +
+            erosion_sagg +
+            erosion_lagg +
+            erosion_gravel
+
+        ### Deposition / settling ###
+        # Different deposition if waterbody outlet or river
+        if waterbodies_locs[v]
+            # Deposition in waterbodies outlets
+            deposition_clay = reservoir_deposition_camp(
+                (input_clay + erosion_clay),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_clay[v],
+                geometry.slope[v],
+            )
+            deposition_silt = reservoir_deposition_camp(
+                (input_silt + erosion_silt),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_silt[v],
+                geometry.slope[v],
+            )
+            deposition_sand = reservoir_deposition_camp(
+                (input_sand + erosion_sand),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_sand[v],
+                geometry.slope[v],
+            )
+            deposition_sagg = reservoir_deposition_camp(
+                (input_sagg + erosion_sagg),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_sagg[v],
+                geometry.slope[v],
+            )
+            deposition_lagg = reservoir_deposition_camp(
+                (input_lagg + erosion_lagg),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_lagg[v],
+                geometry.slope[v],
+            )
+            deposition_gravel = reservoir_deposition_camp(
+                (input_gravel + erosion_gravel),
+                q[v],
+                waterlevel[v],
+                waterbodies_area[v],
+                waterbodies_trapping_efficiency[v],
+                dm_gravel[v],
+                geometry.slope[v],
+            )
+        elseif waterbodies[v]
+            # No deposition in waterbodies, only at the outlets
+            deposition_clay = 0.0
+            deposition_silt = 0.0
+            deposition_sand = 0.0
+            deposition_sagg = 0.0
+            deposition_lagg = 0.0
+            deposition_gravel = 0.0
+        else
+            # Deposition in the river
+            # From transport capacity exceedance
+            excess_sediment = max(input_sediment - transport_capacity[v], 0.0)
+            if excess_sediment > 0.0
+                # Sediment deposited in the channel (from gravel to clay) [ton]
+                # Gravel
+                deposition_gravel = ifelse(
+                    excess_sediment > (input_gravel + erosion_gravel),
+                    (input_gravel + erosion_gravel),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_gravel, 0.0)
+                # Large aggregates
+                deposition_lagg = ifelse(
+                    excess_sediment > (input_lagg + erosion_lagg),
+                    (input_lagg + erosion_lagg),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_lagg, 0.0)
+                # Sand
+                deposition_sand = ifelse(
+                    excess_sediment > (input_sand + erosion_sand),
+                    (input_sand + erosion_sand),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_sand, 0.0)
+                # Small aggregates
+                deposition_sagg = ifelse(
+                    excess_sediment > (input_sagg + erosion_sagg),
+                    (input_sagg + erosion_sagg),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_sagg, 0.0)
+                # Silt
+                deposition_silt = ifelse(
+                    excess_sediment > (input_silt + erosion_silt),
+                    (input_silt + erosion_silt),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_silt, 0.0)
+                # Clay
+                deposition_clay = ifelse(
+                    excess_sediment > (input_clay + erosion_clay),
+                    (input_clay + erosion_clay),
+                    excess_sediment,
+                )
+                excess_sediment = max(excess_sediment - deposition_clay, 0.0)
+            else
+                # Natural deposition from Einstein's formula (density controlled)
+                # Particle fall velocity [m/s] from Stokes
+                xs = ifelse(
+                    q[v] > 0.0,
+                    1.055 * geometry.length[v] / (q[v] / geometry.width[v]),
+                    0.0,
+                )
+                xclay =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_clay[v] / 1000)^2 / 3600)))
+                deposition_clay = xclay * (input_clay + erosion_clay)
+                xsilt =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_silt[v] / 1000)^2 / 3600)))
+                deposition_silt = xsilt * (input_silt + erosion_silt)
+                xsand =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_sand[v] / 1000)^2 / 3600)))
+                deposition_sand = xsand * (input_sand + erosion_sand)
+                xsagg =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_sagg[v] / 1000)^2 / 3600)))
+                deposition_sagg = xsagg * (input_sagg + erosion_sagg)
+                xlagg =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_lagg[v] / 1000)^2 / 3600)))
+                deposition_lagg = xlagg * (input_lagg + erosion_lagg)
+                xgrav =
+                    min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (dm_gravel[v] / 1000)^2 / 3600)))
+                deposition_gravel = xgrav * (input_gravel + erosion_gravel)
+            end
+        end
+
+        # Update the sediment store
+        store_clay[v] += deposition_clay
+        store_silt[v] += deposition_silt
+        store_sand[v] += deposition_sand
+        store_sagg[v] += deposition_sagg
+        store_lagg[v] += deposition_lagg
+        store_gravel[v] += deposition_gravel
+
+        # Compute total deposition
+        deposition[v] =
+            deposition_clay +
+            deposition_silt +
+            deposition_sand +
+            deposition_sagg +
+            deposition_lagg +
+            deposition_gravel
+
+        ### Output loads ###
+        # Sediment transported out of the cell during the timestep [ton]
+        # 0 in case all sediment are deposited in the cell
+        # Reduce the fraction so that there is still some sediment staying in the river cell
+        if waterlevel[v] > 0.0
+            fwaterout = min(
+                q[v] * dt / (waterlevel[v] * geometry.width[v] * geometry.length[v]),
+                1.0,
+            )
+        else
+            fwaterout = 1.0
+        end
+        clay[v] = fwaterout * (input_clay + erosion_clay - deposition_clay)
+        silt[v] = fwaterout * (input_silt + erosion_silt - deposition_silt)
+        sand[v] = fwaterout * (input_sand + erosion_sand - deposition_sand)
+        sagg[v] = fwaterout * (input_sagg + erosion_sagg - deposition_sagg)
+        lagg[v] = fwaterout * (input_lagg + erosion_lagg - deposition_lagg)
+        gravel[v] = fwaterout * (input_gravel + erosion_gravel - deposition_gravel)
+
+        amount[v] = clay[v] + silt[v] + sand[v] + sagg[v] + lagg[v] + gravel[v]
+
+        ### Leftover / mass balance ###
+        # Sediment left in the cell [ton]
+        leftover_clay[v] = input_clay + erosion_clay - deposition_clay - clay[v]
+        leftover_silt[v] = input_silt + erosion_silt - deposition_silt - silt[v]
+        leftover_sand[v] = input_sand + erosion_sand - deposition_sand - sand[v]
+        leftover_sagg[v] = input_sagg + erosion_sagg - deposition_sagg - sagg[v]
+        leftover_lagg[v] = input_lagg + erosion_lagg - deposition_lagg - lagg[v]
+        leftover_gravel[v] = input_gravel + erosion_gravel - deposition_gravel - gravel[v]
+    end
+end
+
+abstract type AbstractSedimentConcentrationsRiverModel{T} end
+
+"Struct to store river sediment concentrations model variables"
+@get_units @grid_loc @with_kw struct SedimentConcentrationsRiverVariables{T}
+    # Total sediment concentration in the river
+    total::Vector{T} | "g m-3"
+    # suspended sediemnt concentration in the river
+    suspended::Vector{T} | "g m-3"
+    # bed load sediment concentration in the river
+    bed::Vector{T} | "g m-3"
+end
+
+"Initialize river sediment concentrations model variables"
+function SedimentConcentrationsRiverVariables(
+    n;
+    total::Vector{T} = fill(mv, n),
+    suspended::Vector{T} = fill(mv, n),
+    bed::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentConcentrationsRiverVariables{T}(;
+        total = total,
+        suspended = suspended,
+        bed = bed,
+    )
+end
+
+"Struct to store river sediment concentrations model boundary conditions"
+@get_units @grid_loc @with_kw struct SedimentConcentrationsRiverBC{T}
+    # Discharge
+    q::Vector{T} | "m3 s-1"
+    waterlevel::Vector{T} | "m"
+    # Clay load
+    clay::Vector{T} | "g m-3"
+    # Silt load
+    silt::Vector{T} | "g m-3"
+    # Sand load
+    sand::Vector{T} | "g m-3"
+    # Small aggregates load
+    sagg::Vector{T} | "g m-3"
+    # Large aggregates load
+    lagg::Vector{T} | "g m-3"
+    # Gravel load
+    gravel::Vector{T} | "g m-3"
+end
+
+"Initialize river sediment concentrations model boundary conditions"
+function SedimentConcentrationsRiverBC(
+    n;
+    q::Vector{T} = fill(mv, n),
+    waterlevel::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(mv, n),
+    silt::Vector{T} = fill(mv, n),
+    sand::Vector{T} = fill(mv, n),
+    sagg::Vector{T} = fill(mv, n),
+    lagg::Vector{T} = fill(mv, n),
+    gravel::Vector{T} = fill(mv, n),
+) where {T}
+    return SedimentConcentrationsRiverBC{T}(;
+        q = q,
+        waterlevel = waterlevel,
+        clay = clay,
+        silt = silt,
+        sand = sand,
+        sagg = sagg,
+        lagg = lagg,
+        gravel = gravel,
+    )
+end
+
+"Struct to store river sediment concentrations model parameters"
+@get_units @grid_loc @with_kw struct SedimentConcentrationsRiverParameters{T}
+    # Clay mean diameter
+    dm_clay::Vector{T} | "µm"
+    # Silt mean diameter
+    dm_silt::Vector{T} | "µm"
+    # Sand mean diameter
+    dm_sand::Vector{T} | "µm"
+    # Small aggregates mean diameter
+    dm_sagg::Vector{T} | "µm"
+    # Large aggregates mean diameter
+    dm_lagg::Vector{T} | "µm"
+    # Gravel mean diameter
+    dm_gravel::Vector{T} | "µm"
+end
+
+"Initialize river sediment concentrations model parameters"
+function SedimentConcentrationsRiverParameters(dataset, config, indices)
+    dm_clay = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_clay";
+        sel = indices,
+        defaults = 2.0,
+        type = Float,
+    )
+    dm_silt = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_silt";
+        sel = indices,
+        defaults = 10.0,
+        type = Float,
+    )
+    dm_sand = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_sand";
+        sel = indices,
+        defaults = 200.0,
+        type = Float,
+    )
+    dm_sagg = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_sagg";
+        sel = indices,
+        defaults = 30.0,
+        type = Float,
+    )
+    dm_lagg = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_lagg";
+        sel = indices,
+        defaults = 500.0,
+        type = Float,
+    )
+    dm_gravel = ncread(
+        dataset,
+        config,
+        "lateral.river.concentrations.parameters.dm_gravel";
+        sel = indices,
+        defaults = 2000.0,
+        type = Float,
+    )
+    conc_parameters = SedimentConcentrationsRiverParameters(;
+        dm_clay = dm_clay,
+        dm_silt = dm_silt,
+        dm_sand = dm_sand,
+        dm_sagg = dm_sagg,
+        dm_lagg = dm_lagg,
+        dm_gravel = dm_gravel,
+    )
+
+    return conc_parameters
+end
+
+"Struct to store river sediment concentrations model"
+@with_kw struct SedimentConcentrationsRiverModel{T} <:
+                AbstractSedimentConcentrationsRiverModel{T}
+    boundary_conditions::SedimentConcentrationsRiverBC{T}
+    parameters::SedimentConcentrationsRiverParameters{T}
+    variables::SedimentConcentrationsRiverVariables{T}
+end
+
+"Initialize river sediment concentrations model"
+function SedimentConcentrationsRiverModel(dataset, config, indices)
+    n = length(indices)
+    vars = SedimentConcentrationsRiverVariables(n)
+    params = SedimentConcentrationsRiverParameters(dataset, config, indices)
+    bc = SedimentConcentrationsRiverBC(n)
+    model = SedimentConcentrationsRiverModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update boundary conditions for river sediment concentrations model"
+function update_boundary_conditions!(
+    model::SedimentConcentrationsRiverModel,
+    hydrological_forcing::HydrologicalForcing,
+    sediment_flux_model::AbstractSedimentRiverTransportModel,
+)
+    (; q, waterlevel, clay, silt, sand, sagg, lagg, gravel) = model.boundary_conditions
+    # Hydrological forcing
+    (; q_river, waterlevel_river) = hydrological_forcing
+    @. q = q_river
+    @. waterlevel = waterlevel_river
+    # Sediment flux per particle
+    @. clay = sediment_flux_model.variables.clay
+    @. silt = sediment_flux_model.variables.silt
+    @. sand = sediment_flux_model.variables.sand
+    @. sagg = sediment_flux_model.variables.sagg
+    @. lagg = sediment_flux_model.variables.lagg
+    @. gravel = sediment_flux_model.variables.gravel
+end
+
+"Update river sediment concentrations model for a single timestep"
+function update!(model::SedimentConcentrationsRiverModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel, clay, silt, sand, sagg, lagg, gravel) = model.boundary_conditions
+    (; dm_clay, dm_silt, dm_sand, dm_sagg, dm_lagg, dm_gravel) = model.parameters
+    (; total, suspended, bed) = model.variables
+
+    zeros = fill(0.0, length(q))
+    # Conversion from load [ton] to concentration for rivers [mg/L]
+    toconc = ifelse.(q .> 0.0, 1e6 ./ (q .* dt), zeros)
+
+    # Differentiation of bed and suspended load using Rouse number for suspension
+    # threshold diameter between bed load and mixed load using Rouse number
+    dbedf =
+        1e3 .*
+        (2.5 .* 3600 .* 0.41 ./ 411 .* (9.81 .* waterlevel .* geometry.slope) .^ 0.5) .^ 0.5
+    # threshold diameter between suspended load and mixed load using Rouse number
+    dsuspf =
+        1e3 .*
+        (1.2 .* 3600 .* 0.41 ./ 411 .* (9.81 .* waterlevel .* geometry.slope) .^ 0.5) .^ 0.5
+
+    # Rouse with diameter
+    SSclay = ifelse.(dm_clay .<= dsuspf, clay, ifelse.(dm_clay .<= dbedf, clay ./ 2, zeros))
+    SSsilt = ifelse.(dm_silt .<= dsuspf, silt, ifelse.(dm_silt .<= dbedf, silt ./ 2, zeros))
+    SSsand = ifelse.(dm_sand .<= dsuspf, sand, ifelse.(dm_sand .<= dbedf, sand ./ 2, zeros))
+    SSsagg = ifelse.(dm_sagg .<= dsuspf, sagg, ifelse.(dm_sagg .<= dbedf, sagg ./ 2, zeros))
+    SSlagg = ifelse.(dm_lagg .<= dsuspf, lagg, ifelse.(dm_lagg .<= dbedf, lagg ./ 2, zeros))
+    SSgrav =
+        ifelse.(
+            dm_gravel .<= dsuspf,
+            gravel,
+            ifelse.(dm_gravel .<= dbedf, gravel ./ 2, zeros),
+        )
+
+    SS = SSclay .+ SSsilt .+ SSsagg .+ SSsand .+ SSlagg .+ SSgrav
+    Bedload = (clay .+ silt .+ sagg .+ sand .+ lagg .+ gravel) .- SS
+
+    @. suspended = SS .* toconc
+    @. bed = Bedload .* toconc
+    @. total = (clay .+ silt .+ sagg .+ sand .+ lagg .+ gravel) .* toconc
+end
\ No newline at end of file
diff --git a/src/sediment/sediment_transport/transport_capacity.jl b/src/sediment/sediment_transport/transport_capacity.jl
new file mode 100644
index 000000000..9dd1e4b4b
--- /dev/null
+++ b/src/sediment/sediment_transport/transport_capacity.jl
@@ -0,0 +1,780 @@
+abstract type AbstractTransportCapacityModel{T} end
+
+"Struct to store total transport capacity model variables"
+@get_units @grid_loc @with_kw struct TransportCapacityModelVariables{T}
+    # Total sediment transport capacity
+    amount::Vector{T} | "t dt-1"
+end
+
+"Initialize total transport capacity model variables"
+function TransportCapacityModelVariables(n; amount::Vector{T} = fill(mv, n)) where {T}
+    return TransportCapacityModelVariables{T}(; amount = amount)
+end
+
+"Struct to store total transport capacity model boundary conditions"
+@get_units @grid_loc @with_kw struct TransportCapacityBC{T}
+    # Discharge
+    q::Vector{T} | "m3 s-1"
+    # Flow depth
+    waterlevel::Vector{T} | "m"
+end
+
+"Initialize total transport capacity model boundary conditions"
+function TransportCapacityBC(
+    n;
+    q::Vector{T} = fill(mv, n),
+    waterlevel::Vector{T} = fill(mv, n),
+) where {T}
+    return TransportCapacityBC{T}(; q = q, waterlevel = waterlevel)
+end
+
+"Update total transport capacity model boundary conditions"
+function update_boundary_conditions!(
+    model::AbstractTransportCapacityModel,
+    hydrological_forcing::HydrologicalForcing,
+    model_type::Symbol,
+)
+    (; q, waterlevel) = model.boundary_conditions
+    (; q_land, waterlevel_land, q_river, waterlevel_river) = hydrological_forcing
+
+    if model_type == :land
+        @. q = q_land
+        @. waterlevel = waterlevel_land
+    elseif model_type == :river
+        @. q = q_river
+        @. waterlevel = waterlevel_river
+    end
+end
+
+##################### Overland Flow #####################
+
+"Struct to store Govers overland flow transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityGoversParameters{T}
+    # Particle density
+    density::Vector{T} | "kg m-3"
+    # Govers transport capacity coefficient
+    c_govers::Vector{T} | "-"
+    # Govers transport capacity exponent
+    n_govers::Vector{T} | "-"
+end
+
+"Initialize Govers overland flow transport capacity model parameters"
+function TransportCapacityGoversParameters(dataset, config, indices)
+    density = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.density";
+        sel = indices,
+        defaults = 2650.0,
+        type = Float,
+    )
+    c_govers = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.c_govers";
+        sel = indices,
+        defaults = 0.000505,
+        type = Float,
+    )
+    n_govers = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.n_govers";
+        sel = indices,
+        defaults = 4.27,
+        type = Float,
+    )
+    tc_parameters = TransportCapacityGoversParameters(;
+        density = density,
+        c_govers = c_govers,
+        n_govers = n_govers,
+    )
+
+    return tc_parameters
+end
+
+"Govers overland flow transport capacity model"
+@with_kw struct TransportCapacityGoversModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityGoversParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Govers overland flow transport capacity model"
+function TransportCapacityGoversModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityGoversParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityGoversModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Govers overland flow transport capacity model for a single timestep"
+function update!(
+    model::TransportCapacityGoversModel,
+    geometry::LandGeometry,
+    waterbodies,
+    rivers,
+    dt,
+)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, c_govers, n_govers) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_govers(
+            q[i],
+            waterlevel[i],
+            c_govers[i],
+            n_govers[i],
+            density[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dt,
+        )
+    end
+end
+
+"Struct to store Yalin overland flow transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityYalinParameters{T}
+    # Particle density
+    density::Vector{T} | "kg m-3"
+    # Particle mean diameter
+    d50::Vector{T} | "mm"
+end
+
+"Initialize Yalin overland flow transport capacity model parameters"
+function TransportCapacityYalinParameters(dataset, config, indices)
+    density = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.density";
+        sel = indices,
+        defaults = 2650.0,
+        type = Float,
+    )
+    d50 = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.d50";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    tc_parameters = TransportCapacityYalinParameters(; density = density, d50 = d50)
+
+    return tc_parameters
+end
+
+"Yalin overland flow transport capacity model"
+@with_kw struct TransportCapacityYalinModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityYalinParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Yalin overland flow transport capacity model"
+function TransportCapacityYalinModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityYalinParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityYalinModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Yalin overland flow transport capacity model for a single timestep"
+function update!(
+    model::TransportCapacityYalinModel,
+    geometry::LandGeometry,
+    waterbodies,
+    rivers,
+    dt,
+)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, d50) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_yalin(
+            q[i],
+            waterlevel[i],
+            density[i],
+            d50[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dt,
+        )
+    end
+end
+
+"Struct to store Yalin differentiated overland flow transport capacity model variables"
+@get_units @grid_loc @with_kw struct TransportCapacityYalinDifferentiationModelVariables{T}
+    # Total sediment transport capacity
+    amount::Vector{T} | "t dt-1"
+    # Transport capacity clay
+    clay::Vector{T} | "t dt-1"
+    # Transport capacity silt
+    silt::Vector{T} | "t dt-1"
+    # Transport capacity sand
+    sand::Vector{T} | "t dt-1"
+    # Transport capacity small aggregates
+    sagg::Vector{T} | "t dt-1"
+    # Transport capacity large aggregates
+    lagg::Vector{T} | "t dt-1"
+end
+
+"Initialize Yalin differentiated overland flow transport capacity model variables"
+function TransportCapacityYalinDifferentiationModelVariables(
+    n;
+    amount::Vector{T} = fill(mv, n),
+    clay::Vector{T} = fill(mv, n),
+    silt::Vector{T} = fill(mv, n),
+    sand::Vector{T} = fill(mv, n),
+    sagg::Vector{T} = fill(mv, n),
+    lagg::Vector{T} = fill(mv, n),
+) where {T}
+    return TransportCapacityYalinDifferentiationModelVariables{T}(;
+        amount = amount,
+        clay = clay,
+        silt = silt,
+        sand = sand,
+        sagg = sagg,
+        lagg = lagg,
+    )
+end
+
+"Struct to store Yalin differentiated overland flow transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityYalinDifferentiationParameters{T}
+    # Particle density
+    density::Vector{T} | "kg m-3"
+    # Clay mean diameter
+    dm_clay::Vector{T} | "µm"
+    # Silt mean diameter
+    dm_silt::Vector{T} | "µm"
+    # Sand mean diameter
+    dm_sand::Vector{T} | "µm"
+    # Small aggregates mean diameter
+    dm_sagg::Vector{T} | "µm"
+    # Large aggregates mean diameter
+    dm_lagg::Vector{T} | "µm"
+end
+
+"Initialize Yalin differentiated overland flow transport capacity model parameters"
+function TransportCapacityYalinDifferentiationParameters(dataset, config, indices)
+    density = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.density";
+        sel = indices,
+        defaults = 2650.0,
+        type = Float,
+    )
+    dm_clay = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.dm_clay";
+        sel = indices,
+        defaults = 2.0,
+        type = Float,
+    )
+    dm_silt = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.dm_silt";
+        sel = indices,
+        defaults = 10.0,
+        type = Float,
+    )
+    dm_sand = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.dm_sand";
+        sel = indices,
+        defaults = 200.0,
+        type = Float,
+    )
+    dm_sagg = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.dm_sagg";
+        sel = indices,
+        defaults = 30.0,
+        type = Float,
+    )
+    dm_lagg = ncread(
+        dataset,
+        config,
+        "lateral.land.transport_capacity.parameters.dm_lagg";
+        sel = indices,
+        defaults = 500.0,
+        type = Float,
+    )
+    tc_parameters = TransportCapacityYalinDifferentiationParameters(;
+        density = density,
+        dm_clay = dm_clay,
+        dm_silt = dm_silt,
+        dm_sand = dm_sand,
+        dm_sagg = dm_sagg,
+        dm_lagg = dm_lagg,
+    )
+
+    return tc_parameters
+end
+
+"Yalin differentiated overland flow transport capacity model"
+@with_kw struct TransportCapacityYalinDifferentiationModel{T} <:
+                AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityYalinDifferentiationParameters{T}
+    variables::TransportCapacityYalinDifferentiationModelVariables{T}
+end
+
+"Initialize Yalin differentiated overland flow transport capacity model"
+function TransportCapacityYalinDifferentiationModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityYalinDifferentiationModelVariables(n)
+    params = TransportCapacityYalinDifferentiationParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityYalinDifferentiationModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Yalin differentiated overland flow transport capacity model for a single timestep"
+function update!(
+    model::TransportCapacityYalinDifferentiationModel,
+    geometry::LandGeometry,
+    waterbodies,
+    rivers,
+    dt,
+)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, dm_clay, dm_silt, dm_sand, dm_sagg, dm_lagg) = model.parameters
+    (; amount, clay, silt, sand, sagg, lagg) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        dtot = transportability_yalin_differentiation(
+            waterlevel[i],
+            density[i],
+            dm_clay[i],
+            dm_silt[i],
+            dm_sand[i],
+            dm_sagg[i],
+            dm_lagg[i],
+            geometry.slope[i],
+        )
+        clay[i] = transport_capacity_yalin_differentiation(
+            q[i],
+            waterlevel[i],
+            density[i],
+            dm_clay[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dtot,
+            dt,
+        )
+        silt[i] = transport_capacity_yalin_differentiation(
+            q[i],
+            waterlevel[i],
+            density[i],
+            dm_silt[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dtot,
+            dt,
+        )
+        sand[i] = transport_capacity_yalin_differentiation(
+            q[i],
+            waterlevel[i],
+            density[i],
+            dm_sand[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dtot,
+            dt,
+        )
+        sagg[i] = transport_capacity_yalin_differentiation(
+            q[i],
+            waterlevel[i],
+            density[i],
+            dm_sagg[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dtot,
+            dt,
+        )
+        lagg[i] = transport_capacity_yalin_differentiation(
+            q[i],
+            waterlevel[i],
+            density[i],
+            dm_lagg[i],
+            geometry.slope[i],
+            geometry.width[i],
+            waterbodies[i],
+            rivers[i],
+            dtot,
+            dt,
+        )
+        amount[i] = clay[i] + silt[i] + sand[i] + sagg[i] + lagg[i]
+    end
+end
+
+"Struct to store common river transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityRiverParameters{T}
+    # Particle density
+    density::Vector{T} | "kg m-3"
+    # Particle mean diameter
+    d50::Vector{T} | "mm"
+end
+
+"Initialize common river transport capacity model parameters"
+function TransportCapacityRiverParameters(dataset, config, indices)
+    density = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.density";
+        sel = indices,
+        defaults = 2650.0,
+        type = Float,
+    )
+    d50 = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.d50";
+        sel = indices,
+        defaults = 0.1,
+        type = Float,
+    )
+    tc_parameters = TransportCapacityRiverParameters(; density = density, d50 = d50)
+
+    return tc_parameters
+end
+
+"Struct to store Bagnold transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityBagnoldParameters{T}
+    # Bagnold transport capacity coefficient
+    c_bagnold::Vector{T} | "-"
+    # Bagnold transport capacity exponent
+    e_bagnold::Vector{T} | "-"
+end
+
+"Initialize Bagnold transport capacity model parameters"
+function TransportCapacityBagnoldParameters(dataset, config, indices)
+    c_bagnold = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.c_bagnold";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    e_bagnold = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.e_bagnold";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    tc_parameters =
+        TransportCapacityBagnoldParameters(; c_bagnold = c_bagnold, e_bagnold = e_bagnold)
+
+    return tc_parameters
+end
+
+"Bagnold river transport capacity model"
+@with_kw struct TransportCapacityBagnoldModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityBagnoldParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Bagnold river transport capacity model"
+function TransportCapacityBagnoldModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityBagnoldParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityBagnoldModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Bagnold river transport capacity model for a single timestep"
+function update!(model::TransportCapacityBagnoldModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel) = model.boundary_conditions
+    (; c_bagnold, e_bagnold) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    # Note: slope is not used here but this allows for a consistent interface of update! functions
+    # Only Bagnold does not use it
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_bagnold(
+            q[i],
+            waterlevel[i],
+            c_bagnold[i],
+            e_bagnold[i],
+            geometry.width[i],
+            geometry.length[i],
+            dt,
+        )
+    end
+end
+
+"Engelund and Hansen river transport capacity model parameters"
+@with_kw struct TransportCapacityEngelundModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityRiverParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Engelund and Hansen river transport capacity model"
+function TransportCapacityEngelundModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityRiverParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityEngelundModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Engelund and Hansen river transport capacity model for a single timestep"
+function update!(model::TransportCapacityEngelundModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, d50) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_engelund(
+            q[i],
+            waterlevel[i],
+            density[i],
+            d50[i],
+            geometry.width[i],
+            geometry.length[i],
+            geometry.slope[i],
+            dt,
+        )
+    end
+end
+
+"Struct to store Kodatie river transport capacity model parameters"
+@get_units @grid_loc @with_kw struct TransportCapacityKodatieParameters{T}
+    # Kodatie transport capacity coefficient a
+    a_kodatie::Vector{T} | "-"
+    # Kodatie transport capacity coefficient b
+    b_kodatie::Vector{T} | "-"
+    # Kodatie transport capacity coefficient c
+    c_kodatie::Vector{T} | "-"
+    # Kodatie transport capacity coefficient d
+    d_kodatie::Vector{T} | "-"
+end
+
+"Initialize Kodatie river transport capacity model parameters"
+function TransportCapacityKodatieParameters(dataset, config, indices)
+    a_kodatie = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.a_kodatie";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    b_kodatie = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.b_kodatie";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    c_kodatie = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.c_kodatie";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    d_kodatie = ncread(
+        dataset,
+        config,
+        "lateral.river.transport_capacity.parameters.d_kodatie";
+        sel = indices,
+        optional = false,
+        type = Float,
+    )
+    tc_parameters = TransportCapacityKodatieParameters(;
+        a_kodatie = a_kodatie,
+        b_kodatie = b_kodatie,
+        c_kodatie = c_kodatie,
+        d_kodatie = d_kodatie,
+    )
+
+    return tc_parameters
+end
+
+"Kodatie river transport capacity model"
+@with_kw struct TransportCapacityKodatieModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityKodatieParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Kodatie river transport capacity model"
+function TransportCapacityKodatieModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityKodatieParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityKodatieModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Kodatie river transport capacity model for a single timestep"
+function update!(model::TransportCapacityKodatieModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel) = model.boundary_conditions
+    (; a_kodatie, b_kodatie, c_kodatie, d_kodatie) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_kodatie(
+            q[i],
+            waterlevel[i],
+            a_kodatie[i],
+            b_kodatie[i],
+            c_kodatie[i],
+            d_kodatie[i],
+            geometry.width[i],
+            geometry.length[i],
+            geometry.slope[i],
+            dt,
+        )
+    end
+end
+
+"Yang river transport capacity model"
+@with_kw struct TransportCapacityYangModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityRiverParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Yang river transport capacity model"
+function TransportCapacityYangModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityRiverParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityYangModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Yang river transport capacity model for a single timestep"
+function update!(model::TransportCapacityYangModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, d50) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_yang(
+            q[i],
+            waterlevel[i],
+            density[i],
+            d50[i],
+            geometry.width[i],
+            geometry.length[i],
+            geometry.slope[i],
+            dt,
+        )
+    end
+end
+
+"Molinas and Wu river transport capacity model"
+@with_kw struct TransportCapacityMolinasModel{T} <: AbstractTransportCapacityModel{T}
+    boundary_conditions::TransportCapacityBC{T}
+    parameters::TransportCapacityRiverParameters{T}
+    variables::TransportCapacityModelVariables{T}
+end
+
+"Initialize Molinas and Wu river transport capacity model"
+function TransportCapacityMolinasModel(dataset, config, indices)
+    n = length(indices)
+    vars = TransportCapacityModelVariables(n)
+    params = TransportCapacityRiverParameters(dataset, config, indices)
+    bc = TransportCapacityBC(n)
+    model = TransportCapacityMolinasModel(;
+        boundary_conditions = bc,
+        parameters = params,
+        variables = vars,
+    )
+    return model
+end
+
+"Update Molinas and Wu river transport capacity model for a single timestep"
+function update!(model::TransportCapacityMolinasModel, geometry::RiverGeometry, dt)
+    (; q, waterlevel) = model.boundary_conditions
+    (; density, d50) = model.parameters
+    (; amount) = model.variables
+
+    n = length(q)
+    threaded_foreach(1:n; basesize = 1000) do i
+        amount[i] = transport_capacity_molinas(
+            q[i],
+            waterlevel[i],
+            density[i],
+            d50[i],
+            geometry.width[i],
+            geometry.length[i],
+            geometry.slope[i],
+            dt,
+        )
+    end
+end
\ No newline at end of file
diff --git a/src/sediment/sediment_transport/transport_capacity_process.jl b/src/sediment/sediment_transport/transport_capacity_process.jl
new file mode 100644
index 000000000..2e8e2338e
--- /dev/null
+++ b/src/sediment/sediment_transport/transport_capacity_process.jl
@@ -0,0 +1,648 @@
+"""
+    mask_transport_capacity(
+        transport_capacity,
+        waterbodies,
+        rivers,
+    )
+
+Mask transport capacity values for waterbodies and rivers.
+
+# Arguments
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+- `waterbodies` (waterbodies mask [-])
+- `rivers` (rivers mask [-])
+
+# Output
+- `transport_capacity` (masked total sediment transport capacity [t dt-1])
+"""
+function mask_transport_capacity(transport_capacity, waterbodies, rivers)
+    # Full deposition in rivers to switch to river concept
+    if rivers
+        tc = 0.0
+        # Sediment flux in waterbodies will all reach the river
+    elseif waterbodies
+        tc = 1e9
+    else
+        tc = transport_capacity
+    end
+
+    return tc
+end
+
+"""
+    limit_and_convert_transport_capacity(
+        transport_capacity,
+        q,
+        waterlevel,
+        width,
+        length,
+        dt,
+    )
+
+Limit to stremaflow and not debris flow and convert transport in ton/m3 to ton.
+
+# Arguments
+- `transport_capacity` (total sediment transport capacity [t m-3])
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `width` (drain width [m])
+- `length` (drain length [m])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function limit_and_convert_transport_capacity(
+    transport_capacity,
+    q,
+    waterlevel,
+    width,
+    length,
+    dt,
+)
+    # 1285 g/L: boundary between streamflow and debris flow (Costa, 1988)
+    transport_capacity = min(transport_capacity, 1.285)
+    # Transport capacity [ton]
+    transport_capacity = transport_capacity * (waterlevel * width * length + q * dt)
+
+    return transport_capacity
+end
+
+"""
+    transport_capacity_govers(
+        q,
+        waterlevel,
+        c_govers,
+        n_govers,
+        density,
+        slope,
+        width,
+        waterbodies,
+        rivers,
+        dt,
+    )
+
+Total sediment transport capacity based on Govers.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `c_govers` (Govers transport capacity coefficient [-])
+- `n_govers` (Govers transport capacity exponent [-])
+- `density` (sediment density [kg m-3])
+- `slope` (slope [-])
+- `width` (drain width [m])
+- `waterbodies` (waterbodies mask [-])
+- `rivers` (rivers mask [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_govers(
+    q,
+    waterlevel,
+    c_govers,
+    n_govers,
+    density,
+    slope,
+    width,
+    waterbodies,
+    rivers,
+    dt,
+)
+    # Transport capacity from govers 1990
+    sinslope = sin(atan(slope)) #slope in radians
+    # Unit stream power
+    if waterlevel > 0.0
+        velocity = q / (width * waterlevel) #m/s
+    else
+        velocity = 0.0
+    end
+    omega = 10 * sinslope * 100 * velocity #cm/s
+    if omega > 0.4
+        TCf = c_govers * (omega - 0.4)^(n_govers) * density #kg/m3
+    else
+        TCf = 0.0
+    end
+    transport_capacity = TCf * q * dt * 1e-3 #[ton/cell]
+
+    # Mask transport capacity values for waterbodies and rivers
+    transport_capacity = mask_transport_capacity(transport_capacity, waterbodies, rivers)
+
+    return transport_capacity
+end
+
+"""
+    transport_capacity_yalin(
+        q,
+        waterlevel,
+        density,
+        d50,
+        slope,
+        width,
+        waterbodies,
+        rivers,
+        dt,
+    )
+
+Total sediment transport capacity based on Yalin.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `d50` (median grain size [m])
+- `slope` (slope [-])
+- `width` (drain width [m])
+- `waterbodies` (waterbodies mask [-])
+- `rivers` (rivers mask [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_yalin(
+    q,
+    waterlevel,
+    density,
+    d50,
+    slope,
+    width,
+    waterbodies,
+    rivers,
+    dt,
+)
+    sinslope = sin(atan(slope)) #slope in radians
+    # Transport capacity from Yalin without particle differentiation
+    delta =
+        max((waterlevel * sinslope / (d50 * 0.001 * (density / 1000 - 1)) / 0.06 - 1), 0.0)
+    alphay = delta * 2.45 / (0.001 * density)^0.4 * 0.06^(0.5)
+    if q > 0.0 && alphay != 0.0
+        TC = (
+            width / q *
+            (density - 1000) *
+            d50 *
+            0.001 *
+            (9.81 * waterlevel * sinslope) *
+            0.635 *
+            delta *
+            (1 - log(1 + alphay) / (alphay))
+        ) # [kg/m3]
+        transport_capacity = TC * q * dt * 1e-3 #[ton]
+    else
+        transport_capacity = 0.0
+    end
+
+    # Mask transport capacity values for waterbodies and rivers
+    transport_capacity = mask_transport_capacity(transport_capacity, waterbodies, rivers)
+
+    return transport_capacity
+end
+
+"""
+    transportability_yalin_differentiation(
+        waterlevel,
+        density,
+        dm_clay,
+        dm_silt,
+        dm_sand,
+        dm_sagg,
+        dm_lagg,
+        slope,
+    )
+
+Total flow transportability based on Yalin with particle differentiation.
+
+# Arguments
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `dm_clay` (clay median grain size [m])
+- `dm_silt` (silt median grain size [m])
+- `dm_sand` (sand median grain size [m])
+- `dm_sagg` (small aggregates median grain size [m])
+- `dm_lagg` (large aggregates median grain size [m])
+- `slope` (slope [-])
+
+# Output
+- `dtot` (total transportability of the flow [-])
+"""
+function transportability_yalin_differentiation(
+    waterlevel,
+    density,
+    dm_clay,
+    dm_silt,
+    dm_sand,
+    dm_sagg,
+    dm_lagg,
+    slope,
+)
+    sinslope = sin(atan(slope)) #slope in radians    
+    # Delta parameter of Yalin for each particle class
+    delta = waterlevel * sinslope / (1e-6 * (density / 1000 - 1)) / 0.06
+    dclay = max(1 / dm_clay * delta - 1, 0.0)
+    dsilt = max(1 / dm_silt * delta - 1, 0.0)
+    dsand = max(1 / dm_sand * delta - 1, 0.0)
+    dsagg = max(1 / dm_sagg * delta - 1, 0.0)
+    dlagg = max(1 / dm_lagg * delta - 1, 0.0)
+    # Total transportability
+    dtot = dclay + dsilt + dsand + dsagg + dlagg
+
+    return dtot
+end
+
+"""
+    transport_capacity_yalin_differentiation(
+        q,
+        waterlevel,
+        density,
+        dm,
+        slope,
+        width,
+        waterbodies,
+        rivers,
+        dtot,
+        dt,
+    )
+
+Transport capacity for a specific grain size based on Yalin with particle differentiation.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `dm` (median grain size [m])
+- `slope` (slope [-])
+- `width` (drain width [m])
+- `waterbodies` (waterbodies mask [-])
+- `rivers` (rivers mask [-])
+- `dtot` (total flow transportability [t dt-1])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_yalin_differentiation(
+    q,
+    waterlevel,
+    density,
+    dm,
+    slope,
+    width,
+    waterbodies,
+    rivers,
+    dtot,
+    dt,
+)
+    sinslope = sin(atan(slope)) #slope in radians
+    # Transport capacity from Yalin with particle differentiation
+    # Delta parameter of Yalin for the specific particle class
+    delta = waterlevel * sinslope / (1e-6 * (density / 1000 - 1)) / 0.06
+    d_part = max(1 / dm * delta - 1, 0.0)
+
+    if q > 0.0
+        TCa = width / q * (density - 1000) * 1e-6 * (9.81 * waterlevel * sinslope)
+    else
+        TCa = 0.0
+    end
+
+    TCb = 2.45 / (0.001 * density)^0.4 * 0.06^0.5
+
+    if dtot != 0.0 && d_part != 0.0
+        TC =
+            TCa * dm * d_part / dtot *
+            0.635 *
+            d_part *
+            (1 - log(1 + d_part * TCb) / d_part * TCb) # [kg/m3]
+        transport_capacity = TC * q * dt * 1e-3 #[ton]
+    else
+        transport_capacity = 0.0
+    end
+
+    # Mask transport capacity values for waterbodies and rivers
+    transport_capacity = mask_transport_capacity(transport_capacity, waterbodies, rivers)
+
+    return transport_capacity
+end
+
+"""
+    function trasnport_capacity_bagnold(
+        q,
+        waterlevel,
+        c_bagnold,
+        e_bagnold,
+        width,
+        length,
+        dt,
+    )
+
+Total sediment transport capacity based on Bagnold.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `c_bagnold` (Bagnold transport capacity coefficient [-])
+- `e_bagnold` (Bagnold transport capacity exponent [-])
+- `width` (drain width [m])
+- `length` (drain length [m])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_bagnold(q, waterlevel, c_bagnold, e_bagnold, width, length, dt)
+    # Transport capacity from Bagnold
+    if waterlevel > 0.0
+        # Transport capacity [tons/m3]
+        transport_capacity = c_bagnold * (q / (waterlevel * width))^e_bagnold
+        transport_capacity = limit_and_convert_transport_capacity(
+            transport_capacity,
+            q,
+            waterlevel,
+            width,
+            length,
+            dt,
+        )
+    else
+        transport_capacity = 0.0
+    end
+
+    return transport_capacity
+end
+
+"""
+    function trasnport_capacity_engelund(
+        q,
+        waterlevel,
+        density,
+        d50,
+        width,
+        length,
+        slope,
+        dt,
+    )
+
+Total sediment transport capacity based on Engelund and Hansen.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `d50` (median grain size [m])
+- `width` (drain width [m])
+- `length` (drain length [m])
+- `slope` (slope [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_engelund(q, waterlevel, density, d50, width, length, slope, dt)
+    # Transport capacity from Engelund and Hansen
+    if waterlevel > 0.0
+        # Hydraulic radius of the river [m] (rectangular channel)
+        hydrad = waterlevel * width / (width + 2 * waterlevel)
+        vshear = sqrt(9.81 * hydrad * slope)
+
+        # Flow velocity [m/s]
+        velocity = (q / (waterlevel * width))
+
+        # Concentration by weight
+        cw =
+            density / 1000 * 0.05 * velocity * vshear^3 /
+            ((density / 1000 - 1)^2 * 9.81^2 * d50 * hydrad)
+        cw = min(1.0, cw)
+
+        # Transport capacity [tons/m3]
+        transport_capacity = cw / (cw + (1 - cw) * density / 1000) * density / 1000
+        transport_capacity = max(transport_capacity, 0.0)
+        # Transport capacity [tons]
+        transport_capacity = limit_and_convert_transport_capacity(
+            transport_capacity,
+            q,
+            waterlevel,
+            width,
+            length,
+            dt,
+        )
+
+    else
+        transport_capacity = 0.0
+    end
+
+    return transport_capacity
+end
+
+"""
+    function trasnport_capacity_kodatie(
+        q,
+        waterlevel,
+        a_kodatie,
+        b_kodatie,
+        c_kodatie,
+        d_kodatie,
+        width,
+        length,
+        slope,
+        dt,
+    )
+
+Total sediment transport capacity based on Kodatie.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `a_kodatie` (Kodatie transport capacity coefficient [-])
+- `b_kodatie` (Kodatie transport capacity coefficient [-])
+- `c_kodatie` (Kodatie transport capacity coefficient [-])
+- `d_kodatie` (Kodatie transport capacity coefficient [-])
+- `width` (drain width [m])
+- `slope` (slope [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_kodatie(
+    q,
+    waterlevel,
+    a_kodatie,
+    b_kodatie,
+    c_kodatie,
+    d_kodatie,
+    width,
+    length,
+    slope,
+    dt,
+)
+    # Transport capacity from Kodatie
+    if waterlevel > 0.0
+        # Flow velocity [m/s]
+        velocity = (q / (waterlevel * width))
+
+        # Concentration
+        transport_capacity =
+            a_kodatie * velocity^b_kodatie * waterlevel^c_kodatie * slope^d_kodatie
+
+        # Transport capacity [tons/m3]
+        transport_capacity = transport_capacity * width / (q * dt)
+        transport_capacity = limit_and_convert_transport_capacity(
+            transport_capacity,
+            q,
+            waterlevel,
+            width,
+            length,
+            dt,
+        )
+
+    else
+        transport_capacity = 0.0
+    end
+
+    return transport_capacity
+end
+
+"""
+    function trasnport_capacity_yang(
+        q,
+        waterlevel,
+        density,
+        d50,
+        width,
+        length,
+        slope,
+        dt,
+    )
+
+Total sediment transport capacity based on Yang sand and gravel equations.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `d50` (median grain size [m])
+- `width` (drain width [m])
+- `length` (drain length [m])
+- `slope` (slope [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_yang(q, waterlevel, density, d50, width, length, slope, dt)
+    # Transport capacity from Yang
+    omegas = 411 * d50^2 / 3600
+    # Hydraulic radius of the river [m] (rectangular channel)
+    hydrad = waterlevel * width / (width + 2 * waterlevel)
+    # Critical shear stress velocity
+    vshear = sqrt(9.81 * hydrad * slope)
+    var1 = vshear * d50 / 1000 / (1.16 * 1e-6)
+    var2 = omegas * d50 / 1000 / (1.16 * 1e-6)
+    vcr = ifelse(var1 >= 70.0, 2.05 * omegas, omegas * (2.5 / (log10(var1) - 0.06) + 0.66))
+    vcr = min(vcr, 0.0)
+
+    # Sand equation
+    if (width * waterlevel) > vcr && d50 < 2.0
+        logcppm = (
+            5.435 - 0.286 * log10(var2) - 0.457 * log10(vshear / omegas) + 1.799 -
+            0.409 * log10(var2) -
+            0.314 *
+            log10(vshear / omegas) *
+            log10((q / (width * waterlevel) - vcr) * slope / omegas)
+        )
+        # Gravel equation
+    elseif (width * waterlevel) > vcr && d50 >= 2.0
+        logcppm = (
+            6.681 - 0.633 * log10(var2) - 4.816 * log10(vshear / omegas) + 2.784 -
+            0.305 * log10(var2) -
+            0.282 *
+            log10(vshear / omegas) *
+            log10((q / (width * waterlevel) - vcr) * slope / omegas)
+        )
+    else
+        logcppm = 0.0
+    end
+
+    # Sediment concentration by weight
+    cw = 10^logcppm * 1e-6
+    # Transport capacity [tons/m3]
+    transport_capacity = cw / (cw + (1 - cw) * density / 1000) * density / 1000
+    transport_capacity = max(transport_capacity, 0.0)
+    # Transport capacity [tons]
+    transport_capacity = limit_and_convert_transport_capacity(
+        transport_capacity,
+        q,
+        waterlevel,
+        width,
+        length,
+        dt,
+    )
+
+    return transport_capacity
+end
+
+"""
+    function trasnport_capacity_molinas(
+        q,
+        waterlevel,
+        density,
+        d50,
+        width,
+        length,
+        slope,
+        dt,
+    )
+
+Total sediment transport capacity based on Molinas and Wu.
+
+# Arguments
+- `q` (discharge [m3 s-1])
+- `waterlevel` (water level [m])
+- `density` (sediment density [kg m-3])
+- `d50` (median grain size [m])
+- `width` (drain width [m])
+- `length` (drain length [m])
+- `slope` (slope [-])
+- `dt` (time step [s])
+
+# Output
+- `transport_capacity` (total sediment transport capacity [t dt-1])
+"""
+function transport_capacity_molinas(q, waterlevel, density, d50, width, length, slope, dt)
+    # Transport capacity from Molinas and Wu
+    if waterlevel > 0.0
+        # Flow velocity [m/s]
+        velocity = (q / (waterlevel * width))
+        omegas = 411 * d50^2 / 3600
+
+        # PSI parameter
+        psi = (
+            velocity^3 / (
+                (density / 1000 - 1) *
+                9.81 *
+                waterlevel *
+                omegas *
+                log10(1000 * waterlevel / d50)^2
+            )
+        )
+        # Concentration by weight
+        cw = 1430 * (0.86 + psi^0.5) * psi^1.5 / (0.016 + psi) * 1e-6
+        # Transport capacity [tons/m3]
+        transport_capacity = cw / (cw + (1 - cw) * density / 1000) * density / 1000
+        transport_capacity = max(transport_capacity, 0.0)
+        # Transport capacity [tons]
+        transport_capacity = limit_and_convert_transport_capacity(
+            transport_capacity,
+            q,
+            waterlevel,
+            width,
+            length,
+            dt,
+        )
+
+    else
+        transport_capacity = 0.0
+    end
+
+    return transport_capacity
+end
\ No newline at end of file
diff --git a/src/sediment_flux.jl b/src/sediment_flux.jl
new file mode 100644
index 000000000..4ba49e98d
--- /dev/null
+++ b/src/sediment_flux.jl
@@ -0,0 +1,177 @@
+"Sediment transport in overland flow model"
+@get_units @grid_loc @with_kw struct OverlandFlowSediment{TT, SF, TR}
+    hydrological_forcing::HydrologicalForcing
+    geometry::LandGeometry
+    transport_capacity::TT
+    sediment_flux::SF
+    to_river::TR
+    waterbodies::Vector{Bool} | "-"
+    rivers::Vector{Bool} | "-"
+end
+
+"Initialize the overland flow sediment transport model"
+function OverlandFlowSediment(dataset, config, indices, waterbodies, rivers)
+    n = length(indices)
+    hydrological_forcing = HydrologicalForcing(n)
+    geometry = LandGeometry(dataset, config, indices)
+    # Check what transport capacity equation will be used
+    do_river = get(config.model, "run_river_model", false)::Bool
+    # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
+    landtransportmethod = get(config.model, "land_transport", "yalinpart")::String
+
+    if do_river || landtransportmethod == "yalinpart"
+        transport_capacity_model =
+            TransportCapacityYalinDifferentiationModel(dataset, config, indices)
+    elseif landtransportmethod == "govers"
+        transport_capacity_model = TransportCapacityGoversModel(dataset, config, indices)
+    elseif landtransportmethod == "yalin"
+        transport_capacity_model = TransportCapacityYalinModel(dataset, config, indices)
+    else
+        error("Unknown land transport method: $landtransportmethod")
+    end
+
+    if do_river || landtransportmethod == "yalinpart"
+        sediment_flux_model = SedimentLandTransportDifferentiationModel(indices)
+        to_river_model = SedimentToRiverDifferentiationModel(indices)
+    else
+        sediment_flux_model = SedimentLandTransportModel(indices)
+        to_river_model = SedimentToRiverModel(indices)
+    end
+
+    overland_flow_sediment = OverlandFlowSediment{
+        typeof(transport_capacity_model),
+        typeof(sediment_flux_model),
+        typeof(to_river_model),
+    }(;
+        hydrological_forcing = hydrological_forcing,
+        geometry = geometry,
+        transport_capacity = transport_capacity_model,
+        sediment_flux = sediment_flux_model,
+        to_river = to_river_model,
+        waterbodies = waterbodies,
+        rivers = rivers,
+    )
+    return overland_flow_sediment
+end
+
+"Update the overland flow sediment transport model for a single timestep"
+function update!(model::OverlandFlowSediment, erosion_model::SoilErosionModel, network, dt)
+    # Transport capacity
+    update_boundary_conditions!(model.transport_capacity, model.hydrological_forcing, :land)
+    update!(model.transport_capacity, model.geometry, model.waterbodies, model.rivers, dt)
+
+    # Update boundary conditions before transport
+    update_boundary_conditions!(
+        model.sediment_flux,
+        erosion_model,
+        model.transport_capacity,
+    )
+    # Compute transport
+    update!(model.sediment_flux, network)
+
+    # Update boundary conditions before computing sediment reaching the river
+    update_boundary_conditions!(model.to_river, model.sediment_flux)
+    # Compute sediment reaching the river
+    update!(model.to_river, model.rivers)
+end
+
+### River ###
+"Sediment transport in river model"
+@get_units @grid_loc @with_kw struct RiverSediment{TTR, ER, SFR, CR}
+    hydrological_forcing::HydrologicalForcing
+    geometry::RiverGeometry
+    transport_capacity::TTR
+    potential_erosion::ER
+    sediment_flux::SFR
+    concentrations::CR
+    waterbodies::Vector{Bool} | "-"
+end
+
+"Initialize the river sediment transport model"
+function RiverSediment(dataset, config, indices, waterbodies)
+    n = length(indices)
+    hydrological_forcing = HydrologicalForcing(n)
+    geometry = RiverGeometry(dataset, config, indices)
+
+    # Check what transport capacity equation will be used
+    # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
+    transport_method = get(config.model, "river_transport", "bagnold")::String
+    if transport_method == "bagnold"
+        transport_capacity_model = TransportCapacityBagnoldModel(dataset, config, indices)
+    elseif transport_method == "engelund"
+        transport_capacity_model = TransportCapacityEngelundModel(dataset, config, indices)
+    elseif transport_method == "yang"
+        transport_capacity_model = TransportCapacityYangModel(dataset, config, indices)
+    elseif transport_method == "kodatie"
+        transport_capacity_model = TransportCapacityKodatieModel(dataset, config, indices)
+    elseif transport_method == "molinas"
+        transport_capacity_model = TransportCapacityMolinasModel(dataset, config, indices)
+    else
+        error("Unknown river transport method: $transport_method")
+    end
+
+    # Potential river erosion
+    potential_erosion_model = RiverErosionJulianTorresModel(dataset, config, indices)
+
+    # Sediment flux in river / mass balance
+    sediment_flux_model = SedimentRiverTransportModel(dataset, config, indices)
+
+    # Concentrations
+    concentrations_model = SedimentConcentrationsRiverModel(dataset, config, indices)
+
+    river_sediment = RiverSediment{
+        typeof(transport_capacity_model),
+        typeof(potential_erosion_model),
+        typeof(sediment_flux_model),
+        typeof(concentrations_model),
+    }(;
+        hydrological_forcing = hydrological_forcing,
+        geometry = geometry,
+        transport_capacity = transport_capacity_model,
+        potential_erosion = potential_erosion_model,
+        sediment_flux = sediment_flux_model,
+        concentrations = concentrations_model,
+        waterbodies = waterbodies,
+    )
+    return river_sediment
+end
+
+"Update the river sediment transport model for a single timestep"
+function update!(
+    model::RiverSediment,
+    to_river_model::SedimentToRiverDifferentiationModel,
+    network,
+    indices_river,
+    dt,
+)
+    # Transport capacity
+    update_boundary_conditions!(
+        model.transport_capacity,
+        model.hydrological_forcing,
+        :river,
+    )
+    update!(model.transport_capacity, model.geometry, dt)
+
+    # Potential maximum river erosion
+    update_boundary_conditions!(model.potential_erosion, model.hydrological_forcing)
+    update!(model.potential_erosion, model.geometry, dt)
+
+    # River transport
+    update_boundary_conditions!(
+        model.sediment_flux,
+        model.hydrological_forcing,
+        model.transport_capacity,
+        to_river_model,
+        model.potential_erosion,
+        indices_river,
+    )
+    update!(model.sediment_flux, network, model.geometry, model.waterbodies, dt)
+
+    # Concentrations
+    update_boundary_conditions!(
+        model.concentrations,
+        model.hydrological_forcing,
+        model.sediment_flux,
+    )
+    update!(model.concentrations, model.geometry, dt)
+end
\ No newline at end of file
diff --git a/src/sediment_model.jl b/src/sediment_model.jl
index 886640c19..1d43a4432 100644
--- a/src/sediment_model.jl
+++ b/src/sediment_model.jl
@@ -13,134 +13,131 @@ function initialize_sediment_model(config::Config)
 
     reader = prepare_reader(config)
     clock = Clock(config, reader)
-    dt = clock.dt
+    dataset = NCDataset(static_path)
 
-    do_river = get(config.model, "runrivermodel", false)::Bool
-
-    nc = NCDataset(static_path)
-    dims = dimnames(nc[param(config, "input.subcatchment")])
-
-    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)
-    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]
-
-    inds_riv, rev_inds_riv = active_indices(river_2d, 0)
-    nriv = length(inds_riv)
+    indices, rev_indices = active_indices(subcatch_2d, missing)
+    n = length(indices)
+
+    river_2d = ncread(
+        dataset,
+        config,
+        "river_location";
+        optional = false,
+        type = Bool,
+        fill = false,
+    )
+    river = river_2d[indices]
 
     # Needed to update the forcing
     reservoir = ()
     lake = ()
 
-    # 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)))
-
-    sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
-    xl, yl = cell_lengths(y, cellength, sizeinmetres)
-    riverfrac = get_river_fraction(river, riverlength, riverwidth, xl, yl)
-
-    eros = initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
+    soilloss = SoilLoss(dataset, config, indices)
+
+    # Get waterbodies mask
+    do_reservoirs = get(config.model, "doreservoir", false)::Bool
+    do_lakes = get(config.model, "dolake", false)::Bool
+    waterbodies = fill(0.0, n)
+    if do_reservoirs
+        reservoirs = ncread(
+            dataset,
+            config,
+            "reservoir_areas";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0,
+        )
+        waterbodies = waterbodies .+ reservoirs
+    end
+    if do_lakes
+        lakes = ncread(
+            dataset,
+            config,
+            "lake_areas";
+            optional = false,
+            sel = indices,
+            type = Float,
+            fill = 0,
+        )
+        waterbodies = waterbodies .+ lakes
+    end
+    waterbodies = waterbodies .> 0
 
-    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]
 
     # # lateral part sediment in overland flow
-    rivcell = float(river)
-    ols = OverlandFlowSediment{Float}(;
-        n = n,
-        rivcell = rivcell,
-        soilloss = fill(mv, n),
-        erosclay = fill(mv, n),
-        erossilt = fill(mv, n),
-        erossand = fill(mv, n),
-        erossagg = fill(mv, n),
-        eroslagg = fill(mv, n),
-        TCsed = fill(mv, n),
-        TCclay = fill(mv, n),
-        TCsilt = fill(mv, n),
-        TCsand = fill(mv, n),
-        TCsagg = fill(mv, n),
-        TClagg = fill(mv, n),
-        olsed = fill(mv, n),
-        olclay = fill(mv, n),
-        olsilt = fill(mv, n),
-        olsand = fill(mv, n),
-        olsagg = fill(mv, n),
-        ollagg = fill(mv, n),
-        inlandsed = fill(mv, n),
-        inlandclay = fill(mv, n),
-        inlandsilt = fill(mv, n),
-        inlandsand = fill(mv, n),
-        inlandsagg = fill(mv, n),
-        inlandlagg = fill(mv, n),
-    )
+    overland_flow_sediment =
+        OverlandFlowSediment(dataset, config, indices, waterbodies, river)
 
-    graph = flowgraph(ldd, inds, pcr_dir)
+    graph = flowgraph(ldd, indices, pcr_dir)
 
     # River processes
+    do_river = get(config.model, "run_river_model", false)::Bool
+    # TODO: see if we can skip init if the river model is not needed
+    # or if we leave it when we restructure the Wflow Model struct
 
-    # the indices of the river cells in the land(+river) cell vector
-    landslope =
-        ncread(nc, config, "lateral.land.slope"; optional = false, sel = inds, type = Float)
-    clamp!(landslope, 0.00001, Inf)
+    indices_riv, rev_indices_riv = active_indices(river_2d, 0)
 
-    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")
+    ldd_riv = ldd_2d[indices_riv]
+    graph_riv = flowgraph(ldd_riv, indices_riv, pcr_dir)
 
-    ldd_riv = ldd_2d[inds_riv]
-    graph_riv = flowgraph(ldd_riv, inds_riv, pcr_dir)
+    # Needed for frac_to_river?
+    landslope = ncread(
+        dataset,
+        config,
+        "vertical.land_parameter_set.slope";
+        optional = false,
+        sel = indices,
+        type = Float,
+    )
+    clamp!(landslope, 0.00001, Inf)
 
     index_river = filter(i -> !isequal(river[i], 0), 1:n)
     frac_toriver = fraction_runoff_to_river(graph, ldd, index_river, landslope)
 
-    rs = initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
+    river_sediment = RiverSediment(dataset, config, indices_riv, waterbodies)
 
-    modelmap = (vertical = eros, lateral = (land = ols, river = rs))
+    modelmap = (
+        vertical = soilloss,
+        lateral = (land = overland_flow_sediment, river = river_sediment),
+    )
     indices_reverse = (
-        land = rev_inds,
-        river = rev_inds_riv,
+        land = rev_indices,
+        river = rev_indices_riv,
         reservoir = isempty(reservoir) ? nothing : reservoir.reverse_indices,
         lake = isempty(lake) ? nothing : lake.reverse_indices,
     )
-    writer = prepare_writer(config, modelmap, indices_reverse, x_nc, y_nc, nc)
-    close(nc)
+    y_dataset = read_y_axis(dataset)
+    x_dataset = read_x_axis(dataset)
+    writer =
+        prepare_writer(config, modelmap, indices_reverse, x_dataset, y_dataset, dataset)
+    close(dataset)
 
     # for each domain save the directed acyclic graph, the traversion order,
     # and the indices that map it back to the two dimensional grid
     land = (
         graph = graph,
         order = topological_sort_by_dfs(graph),
-        indices = inds,
-        reverse_indices = rev_inds,
+        indices = indices,
+        reverse_indices = rev_indices,
     )
     river = (
         graph = graph_riv,
         order = topological_sort_by_dfs(graph_riv),
-        indices = inds_riv,
-        reverse_indices = rev_inds_riv,
+        indices = indices_riv,
+        reverse_indices = rev_indices_riv,
     )
 
     model = Model(
         config,
         (; land, river, reservoir, lake, index_river, frac_toriver),
-        (land = ols, river = rs),
-        eros,
+        (land = overland_flow_sediment, river = river_sediment),
+        soilloss,
         clock,
         reader,
         writer,
@@ -153,44 +150,28 @@ function initialize_sediment_model(config::Config)
     return model
 end
 
+"update sediment model for a single timestep"
 function update!(model::Model{N, L, V, R, W, T}) where {N, L, V, R, W, T <: SedimentModel}
-    (; lateral, vertical, network, config) = model
-
-    update_until_ols!(vertical, config)
-    update_until_oltransport!(vertical, config)
-
-    lateral.land.soilloss .= vertical.soilloss
-    lateral.land.erosclay .= vertical.erosclay
-    lateral.land.erossilt .= vertical.erossilt
-    lateral.land.erossand .= vertical.erossand
-    lateral.land.erossagg .= vertical.erossagg
-    lateral.land.eroslagg .= vertical.eroslagg
-
-    lateral.land.TCsed .= vertical.TCsed
-    lateral.land.TCclay .= vertical.TCclay
-    lateral.land.TCsilt .= vertical.TCsilt
-    lateral.land.TCsand .= vertical.TCsand
-    lateral.land.TCsagg .= vertical.TCsagg
-    lateral.land.TClagg .= vertical.TClagg
+    (; lateral, vertical, network, config, clock) = model
+    dt = tosecond(clock.dt)
 
-    update!(lateral.land, network.land, config)
+    # Soil erosion
+    update!(vertical, dt)
 
-    do_river = get(config.model, "runrivermodel", false)::Bool
+    # Overland flow sediment transport
+    update!(lateral.land, vertical.soil_erosion, network.land, dt)
 
+    # River sediment transport
+    do_river = get(config.model, "run_river_model", false)::Bool
     if do_river
-        inds_riv = network.index_river
-        lateral.river.inlandclay .= lateral.land.inlandclay[inds_riv]
-        lateral.river.inlandsilt .= lateral.land.inlandsilt[inds_riv]
-        lateral.river.inlandsand .= lateral.land.inlandsand[inds_riv]
-        lateral.river.inlandsagg .= lateral.land.inlandsagg[inds_riv]
-        lateral.river.inlandlagg .= lateral.land.inlandlagg[inds_riv]
-
-        update!(lateral.river, network.river, config)
+        indices_riv = network.index_river
+        update!(lateral.river, lateral.land.to_river, network.river, indices_riv, dt)
     end
 
     return nothing
 end
 
+"set the initial states of the sediment model"
 function set_states!(
     model::Model{N, L, V, R, W, T},
 ) where {N, L, V, R, W, T <: SedimentModel}
diff --git a/src/states.jl b/src/states.jl
index b0515c3a6..cc25adf5a 100644
--- a/src/states.jl
+++ b/src/states.jl
@@ -70,6 +70,30 @@ function get_soil_states(model_type::AbstractString; snow = false)
     return states
 end
 
+function get_sediment_states()
+    states = (
+        :leftover_clay,
+        :leftover_silt,
+        :leftover_sand,
+        :leftover_sagg,
+        :leftover_lagg,
+        :leftover_gravel,
+        :store_clay,
+        :store_silt,
+        :store_sand,
+        :store_sagg,
+        :store_lagg,
+        :store_gravel,
+        :clay,
+        :silt,
+        :sand,
+        :sagg,
+        :lagg,
+        :gravel,
+    )
+    return states
+end
+
 """
     add_to_required_states(required_states::Tuple, key_entry::Tuple, states::Tuple)
     add_to_required_states(required_states::Tuple, key_entry::Tuple, states::Nothing)
@@ -161,26 +185,7 @@ function extract_required_states(config::Config)
 
     # River states
     if model_type == "sediment"
-        river_states = (
-            :clayload,
-            :siltload,
-            :sandload,
-            :saggload,
-            :laggload,
-            :gravload,
-            :claystore,
-            :siltstore,
-            :sandstore,
-            :saggstore,
-            :laggstore,
-            :gravstore,
-            :outclay,
-            :outsilt,
-            :outsand,
-            :outsagg,
-            :outlagg,
-            :outgrav,
-        )
+        river_states = get_sediment_states()
     elseif model_type == "sbm" || model_type == "sbm_gwf"
         river_states = (:q, :h, :h_av)
     else
@@ -225,13 +230,20 @@ function extract_required_states(config::Config)
     # Add land states to dict
     required_states =
         add_to_required_states(required_states, (:lateral, :land, :variables), land_states)
-    # Add river states to dict
+    # Add sediment states to dict
     if model_type == "sediment"
-        key_entry = (:lateral, :river)
+        required_states = add_to_required_states(
+            required_states,
+            (:lateral, :river, :sediment_flux, :variables),
+            river_states,
+        )
     else
-        key_entry = (:lateral, :river, :variables)
+        required_states = add_to_required_states(
+            required_states,
+            (:lateral, :river, :variables),
+            river_states,
+        )
     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,


Test patch:
diff --git a/build/create_binaries/download_test_data.jl b/build/create_binaries/download_test_data.jl
index 339e19b91..4624d8d57 100644
--- a/build/create_binaries/download_test_data.jl
+++ b/build/create_binaries/download_test_data.jl
@@ -24,7 +24,7 @@ forcing_moselle_path = testdata(v"0.2.6", "forcing-moselle.nc", "forcing-moselle
 forcing_moselle_sed_path =
     testdata(v"0.2.3", "forcing-moselle-sed.nc", "forcing-moselle-sed.nc")
 staticmaps_moselle_sed_path =
-    testdata(v"0.2.3", "staticmaps-moselle-sed.nc", "staticmaps-moselle-sed.nc")
+    testdata(v"0.3.0", "staticmaps-moselle-sed.nc", "staticmaps-moselle-sed.nc")
 instates_moselle_sed_path =
     testdata(v"0.2", "instates-moselle-sed.nc", "instates-moselle-sed.nc")
 instates_moselle_path = testdata(v"0.2.6", "instates-moselle.nc", "instates-moselle.nc")
@@ -50,4 +50,5 @@ forcing_calendar_noleap_path =
 forcing_piave_path = testdata(v"0.2.9", "inmaps-era5-2010-piave.nc", "forcing-piave.nc")
 staticmaps_piave_path = testdata(v"0.2.9", "staticmaps-piave.nc", "staticmaps-piave.nc")
 instates_piave_path = testdata(v"0.2.9", "instates-piave.nc", "instates-piave.nc")
-instates_piave_gwf_path = testdata(v"0.2.9", "instates-piave-gwf.nc", "instates-piave-gwf.nc")
+instates_piave_gwf_path =
+    testdata(v"0.2.9", "instates-piave-gwf.nc", "instates-piave-gwf.nc")
diff --git a/test/run_sediment.jl b/test/run_sediment.jl
index 58c0096bc..603642e23 100644
--- a/test/run_sediment.jl
+++ b/test/run_sediment.jl
@@ -11,13 +11,17 @@ Wflow.run_timestep!(model)
 @testset "first timestep sediment model (vertical)" begin
     eros = model.vertical
 
-    @test eros.erosov[1] ≈ 0.9f0
+    @test eros.atmospheric_forcing.precipitation[1] ≈ 4.086122035980225f0
+    @test eros.hydrological_forcing.q_land[1] ≈ 0.0f0
+    @test eros.overland_flow_erosion.parameters.usle_k[1] ≈ 0.026510488241910934f0
+    @test eros.overland_flow_erosion.parameters.usle_c[1] ≈ 0.014194443821907043f0
+    @test eros.overland_flow_erosion.parameters.answers_k[1] ≈ 0.9f0
+    @test eros.overland_flow_erosion.variables.amount[1] ≈ 0.0f0
+    @test eros.rainfall_erosion.variables.amount[1] ≈ 0.00027245577922893746f0
     @test model.clock.iteration == 1
-    @test mean(eros.leaf_area_index) ≈ 1.7120018886212223f0
-    @test eros.dmsand[1] == 200.0f0
-    @test eros.dmlagg[1] == 500.0f0
-    @test mean(eros.interception) ≈ 0.4767846753916875f0
-    @test mean(eros.soilloss) ≈ 0.008596682196335555f0
+    @test mean(eros.overland_flow_erosion.variables.amount) ≈ 0.00861079076689589f0
+    @test mean(eros.rainfall_erosion.variables.amount) ≈ 0.00016326203201620437f0
+    @test mean(eros.soil_erosion.variables.amount) ≈ 0.008774052798912092f0
 end
 
 # run the second timestep
@@ -26,45 +30,222 @@ Wflow.run_timestep!(model)
 @testset "second timestep sediment model (vertical)" begin
     eros = model.vertical
 
-    @test mean(eros.soilloss) ≈ 0.07601393657280235f0
-    @test mean(eros.erosclay) ≈ 0.0022388720384961766f0
-    @test mean(eros.erossand) ≈ 0.02572046244407882f0
-    @test mean(eros.eroslagg) ≈ 0.022126541806118796f0
-    @test mean(eros.TCsed) == 0.0
-    @test mean(eros.TCsilt) ≈ 1.0988158364353527f6
-    @test mean(eros.TCsand) ≈ 1.0987090622888755f6
-    @test mean(eros.TCclay) ≈ 1.0992655197016734f6
-    @test mean(eros.TCsilt) ≈ 1.0988158364353527f6
+    @test mean(eros.soil_erosion.variables.amount) ≈ 0.07765800489746684f0
+    @test mean(eros.soil_erosion.variables.clay) ≈ 0.002287480354866626f0
+    @test mean(eros.soil_erosion.variables.silt) ≈ 0.0036164773521184896f0
+    @test mean(eros.soil_erosion.variables.sand) ≈ 0.026301393837924607f0
+    @test mean(eros.soil_erosion.variables.lagg) ≈ 0.022577957752547836f0
+    @test mean(eros.soil_erosion.variables.sagg) ≈ 0.022874695590802723f0
 end
 
 @testset "second timestep sediment model (lateral)" begin
-    lat = model.lateral
+    land = model.lateral.land
+    river = model.lateral.river
+
+    @test land.transport_capacity.parameters.dm_sand[1] == 200.0f0
+    @test land.transport_capacity.parameters.dm_lagg[1] == 500.0f0
+
+    @test mean(land.transport_capacity.boundary_conditions.q) ≈ 0.006879398771052133f0
+    @test mean(land.transport_capacity.variables.silt) ≈ 1.0988158364353527f6
+    @test mean(land.transport_capacity.variables.sand) ≈ 1.0987090622888755f6
+    @test mean(land.transport_capacity.variables.clay) ≈ 1.0992655197016734f6
+
+    @test mean(land.to_river.variables.amount) ≈ 0.07624135182616738f0
+    @test sum(land.to_river.variables.clay) ≈ 114.42704329506047f0
+    @test sum(land.to_river.variables.sand) ≈ 1289.4785484597958f0
+    @test mean(land.sediment_flux.variables.clay) ≈ 0.006578791733506439f0
 
-    @test mean(lat.land.inlandsed) ≈ 0.07463801685030906f0
-    @test mean(lat.land.inlandclay) ≈ 0.0022367786781657497f0
-    @test mean(lat.land.inlandsand) ≈ 0.02519222037812127f0
-    @test mean(lat.land.olclay) ≈ 0.006443036462118322f0
+    @test mean(river.hydrological_forcing.q_river) ≈ 0.6975180562953642f0
+    @test river.hydrological_forcing.waterlevel_river[network.river.order[end]] ≈
+          0.006103649735450745f0
+    @test mean(river.geometry.width) ≈ 22.628250814095523f0
 
-    @test mean(lat.river.SSconc) ≈ 0.8259993252994058f0
-    @test mean(lat.river.inlandclay) ≈ 0.01980468760667709f0
-    @test lat.river.h_riv[network.river.order[end]] ≈ 0.006103649735450745f0
-    @test lat.river.outclay[5649] ≈ 2.359031898208781f-9
+    @test mean(river.transport_capacity.boundary_conditions.q) ≈ 0.6975180562953642f0
+    @test mean(river.transport_capacity.variables.amount) ≈ 0.4458019733090582f0
+    @test mean(river.potential_erosion.variables.bed) ≈ 307.18492138827116f0
+
+    @test sum(river.sediment_flux.boundary_conditions.erosion_land_clay) ≈
+          114.42704329506047f0
+    @test sum(river.sediment_flux.boundary_conditions.erosion_land_sand) ≈
+          1289.4785484597958f0
+    @test mean(river.sediment_flux.boundary_conditions.transport_capacity) ≈
+          0.4458019733090582f0
+    @test mean(river.sediment_flux.variables.amount) ≈ 0.4333483865969662f0
+    @test mean(river.sediment_flux.variables.erosion) ≈ 0.019077695621351014f0
+    @test mean(river.sediment_flux.variables.deposition) ≈ 0.6941274181387916f0
+    @test river.sediment_flux.variables.clay[5649] ≈ 2.840979764480952f-9
+
+    @test mean(river.concentrations.variables.suspended) ≈ 0.8260083257660087f0
 end
 
-@testset "Exchange and grid location sediment" begin
-    @test Wflow.exchange(model.vertical.n) == false
-    @test Wflow.exchange(model.vertical.erosk) == true
-    @test Wflow.exchange(model.vertical.leaf_area_index) == true
-    @test Wflow.grid_loc(model.vertical, :n) == "none"
-    @test Wflow.grid_loc(model.vertical, :erosk) == "node"
-    @test Wflow.grid_loc(model.vertical, :leaf_area_index) == "node"
+Wflow.close_files(model)
+
+### Test the sediment model with a different configuration file ###
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+Wflow.run_timestep!(model)
+
+@testset "first timestep sediment model eurosem (vertical)" begin
+    eros = model.vertical
+
+    @test eros.atmospheric_forcing.precipitation[1] ≈ 4.086122035980225f0
+    @test eros.hydrological_forcing.interception[1] ≈ 0.6329902410507202f0
+    @test eros.hydrological_forcing.q_land[1] ≈ 0.0f0
+    @test eros.rainfall_erosion.parameters.soil_detachability[1] ≈ 2.0f0
+    @test eros.rainfall_erosion.parameters.eurosem_exponent[1] ≈ 2.0f0
+    @test eros.overland_flow_erosion.parameters.usle_c[1] ≈ 0.014194443821907043f0
+    @test eros.overland_flow_erosion.variables.amount[1] ≈ 0.0f0
+    @test eros.rainfall_erosion.variables.amount[1] ≈ 0.01232301374083337f0
+    @test mean(eros.overland_flow_erosion.variables.amount) ≈ 0.00861079076689589f0
+    @test mean(eros.rainfall_erosion.variables.amount) ≈ 0.0014726364432116048f0
+    @test mean(eros.soil_erosion.variables.amount) ≈ 0.010083427210107495f0
+end
+
+# run the second timestep
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model engelund (lateral)" begin
     land = model.lateral.land
-    @test Wflow.exchange(land.n) == false
-    @test Wflow.exchange(land.soilloss) == true
-    @test Wflow.exchange(land.inlandsed) == true
-    @test Wflow.grid_loc(land, :n) == "none"
-    @test Wflow.grid_loc(land, :soilloss) == "node"
-    @test Wflow.grid_loc(land, :inlandsed) == "node"
+    river = model.lateral.river
+
+    @test river.transport_capacity.parameters.d50[1] == 0.05000000074505806f0
+    @test mean(river.transport_capacity.boundary_conditions.q) ≈ 0.6975180562953642f0
+    @test mean(river.transport_capacity.variables.amount) ≈ 0.14184859055736687f0
+
+    @test mean(river.concentrations.variables.suspended) ≈ 0.24788001458305775f0
 end
 
 Wflow.close_files(model)
+
+### Test land only model configuration and transport capacity ###
+
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+# Update config to run only the land model
+config.model.run_river_model = false
+# Use govers equation for land transport capacity
+config.model.land_transport = "govers"
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+# run the first and second timestep
+Wflow.run_timestep!(model)
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model govers" begin
+    eros = model.vertical
+    land = model.lateral.land
+
+    @test mean(eros.soil_erosion.variables.amount) ≈ 0.0776983847440198f0
+    @test mean(land.transport_capacity.parameters.c_govers) ≈ 0.16393911236592437f0
+    @test mean(land.transport_capacity.variables.amount) ≈ 1.0988158364353527f6
+    @test mean(land.to_river.variables.amount) ≈ 0.07708434959918917f0
+end
+
+Wflow.close_files(model)
+
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+# Update config to run only the land model
+config.model.run_river_model = false
+# Use yalin equation for land transport capacity
+config.model.land_transport = "yalin"
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+# run the first and second timestep
+Wflow.run_timestep!(model)
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model yalin" begin
+    eros = model.vertical
+    land = model.lateral.land
+
+    @test mean(eros.soil_erosion.variables.amount) ≈ 0.0776983847440198f0
+    @test mean(land.transport_capacity.parameters.d50) ≈ 0.001534350291334408f0
+    @test mean(land.transport_capacity.variables.amount) ≈ 1.0988158364353527f6
+    @test mean(land.to_river.variables.amount) ≈ 0.07759391658531742f0
+end
+
+Wflow.close_files(model)
+
+### Test all river transport capacity ###
+
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+# Use yang equation for river transport capacity
+config.model.river_transport = "yang"
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+# run the first and second timestep
+Wflow.run_timestep!(model)
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model yang (lateral)" begin
+    river = model.lateral.river
+
+    @test river.transport_capacity.parameters.d50[1] == 0.05000000074505806f0
+    @test mean(river.transport_capacity.boundary_conditions.q) ≈ 0.6975180562953642f0
+    @test mean(river.transport_capacity.variables.amount) ≈ 39.959093179632234f0
+
+    @test mean(river.concentrations.variables.suspended) ≈ 0.004036947448899768f0
+end
+
+Wflow.close_files(model)
+
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+# Use kodatie equation for river transport capacity
+config.model.river_transport = "kodatie"
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+# run the first and second timestep
+Wflow.run_timestep!(model)
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model kodatie (lateral)" begin
+    river = model.lateral.river
+
+    @test river.transport_capacity.parameters.a_kodatie[1] == 2829.6
+    @test river.transport_capacity.parameters.b_kodatie[1] == 3.646
+    @test mean(river.transport_capacity.boundary_conditions.q) ≈ 0.6975180562953642f0
+    @test mean(river.transport_capacity.variables.amount) ≈ 30.332588671299625f0
+
+    @test mean(river.concentrations.variables.suspended) ≈ 54.75483621753514f0
+end
+
+Wflow.close_files(model)
+
+tomlpath = joinpath(@__DIR__, "sediment_eurosem_engelund_config.toml")
+config = Wflow.Config(tomlpath)
+# Use molinas equation for river transport capacity
+config.model.river_transport = "molinas"
+
+model = Wflow.initialize_sediment_model(config)
+(; network) = model
+
+# run the first and second timestep
+Wflow.run_timestep!(model)
+Wflow.run_timestep!(model)
+
+@testset "second timestep sediment model molinas (lateral)" begin
+    river = model.lateral.river
+
+    @test river.transport_capacity.parameters.d50[1] == 0.05000000074505806f0
+    @test mean(river.transport_capacity.boundary_conditions.q) ≈ 0.6975180562953642f0
+    @test mean(river.transport_capacity.variables.amount) ≈ 350.6483600591209f0
+
+    @test mean(river.concentrations.variables.suspended) ≈ 884.4249630198293f0
+end
+
+Wflow.close_files(model)
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index f768e4610..bc6d2764d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -41,7 +41,7 @@ forcing_moselle_path = testdata(v"0.2.6", "forcing-moselle.nc", "forcing-moselle
 forcing_moselle_sed_path =
     testdata(v"0.2.3", "forcing-moselle-sed.nc", "forcing-moselle-sed.nc")
 staticmaps_moselle_sed_path =
-    testdata(v"0.2.3", "staticmaps-moselle-sed.nc", "staticmaps-moselle-sed.nc")
+    testdata(v"0.3.0", "staticmaps-moselle-sed.nc", "staticmaps-moselle-sed.nc")
 instates_moselle_sed_path =
     testdata(v"0.2", "instates-moselle-sed.nc", "instates-moselle-sed.nc")
 instates_moselle_path = testdata(v"0.2.6", "instates-moselle.nc", "instates-moselle.nc")
diff --git a/test/sediment_config.toml b/test/sediment_config.toml
index 23c9fd4b5..6839ee4af 100644
--- a/test/sediment_config.toml
+++ b/test/sediment_config.toml
@@ -18,25 +18,25 @@ path_output = "outstates-moselle-sed.nc"
 # if listed, the variable must be present in the NetCDF or error
 # if not listed, the variable can get a default value if it has one
 
-[state.lateral.river]
-clayload = "clayload"
-claystore = "claystore"
-gravload = "gravload"
-gravstore = "gravstore"
-laggload = "laggload"
-laggstore = "laggstore"
-outclay = "outclay"
-outgrav = "outgrav"
-outlagg = "outlagg"
-outsagg = "outsagg"
-outsand = "outsand"
-outsilt = "outsilt"
-saggload = "saggload"
-saggstore = "saggstore"
-sandload = "sandload"
-sandstore = "sandstore"
-siltload = "siltload"
-siltstore = "siltstore"
+[state.lateral.river.sediment_flux.variables]
+leftover_clay = "clayload"
+store_clay = "claystore"
+leftover_gravel = "gravload"
+store_gravel = "gravstore"
+leftover_lagg = "laggload"
+store_lagg = "laggstore"
+clay = "outclay"
+gravel = "outgrav"
+lagg= "outlagg"
+sagg = "outsagg"
+sand = "outsand"
+silt = "outsilt"
+leftover_sagg= "saggload"
+store_sagg = "saggstore"
+leftover_sand = "sandload"
+store_sand = "sandstore"
+leftover_silt = "siltload"
+store_silt = "siltstore"
 
 [input]
 path_forcing = "forcing-moselle-sed.nc"
@@ -47,111 +47,168 @@ gauges = "wflow_gauges"
 ldd = "wflow_ldd"
 river_location = "wflow_river"
 subcatchment = "wflow_subcatch"
+reservoir_areas = "wflow_reservoirareas"
+lake_areas = "wflow_lakeareas"
 
 # specify the internal IDs of the parameters which vary over time
 # the external name mapping needs to be below together with the other mappings
 forcing = [
-  "vertical.h_land",
-  "vertical.interception",
-  "vertical.precipitation",
-  "vertical.q_land",
-  "lateral.river.h_riv",
-  "lateral.river.q_riv",
+  "vertical.atmospheric_forcing.precipitation",
+  "vertical.hydrological_forcing.waterlevel_land",
+  "lateral.land.hydrological_forcing.waterlevel_land",
+  "vertical.hydrological_forcing.q_land",
+  "lateral.land.hydrological_forcing.q_land",
+  "lateral.river.hydrological_forcing.waterlevel_river",
+  "lateral.river.hydrological_forcing.q_river",
 ]
 
-cyclic = ["vertical.leaf_area_index"]
-
-[input.vertical]
-altitude = "wflow_dem"
-canopyheight = "CanopyHeight"
-erosk = "ErosK"
-erosov = "eros_ov"
-erosspl = "eros_spl_EUROSEM"
-h_land = "levKinL"
-interception = "int"
-kext = "Kext"
-leaf_area_index = "LAI" # cyclic
-pathfrac = "PathFrac"
-pclay = "PercentClay"
+[input.vertical.atmospheric_forcing]
 precipitation = "P"
-psilt = "PercentSilt"
+
+[input.vertical.hydrological_forcing]
+waterlevel_land = "levKinL"
 q_land = "runL"
-rivcell = "wflow_river"
-slope = "Slope"
-specific_leaf = "Sl"
-storage_wood = "Swood"
-usleC = "USLE_C"
-usleK = "USLE_K"
-# Reservoir
-resareas = "wflow_reservoirareas"
-# Lake
-lakeareas = "wflow_lakeareas"
 
-[input.lateral.land]
+[input.vertical.land_parameter_set]
 slope = "Slope"
 
-[input.lateral.river]
-h_riv = "h"
-q_riv = "q"
-cbagnold = "c_Bagnold"
-d50 = "D50_River"
-d50engelund = "D50_River"
-ebagnold = "exp_Bagnold"
-fclayriv = "ClayF_River"
-fgravriv = "GravelF_River"
-fsandriv = "SandF_River"
-fsiltriv = "SiltF_River"
+[input.vertical.rainfall_erosion.parameters]
+soil_detachability = "soil_detachability"
+eurosem_exponent = "eros_spl_EUROSEM"
+canopygapfraction = "CanopyGapFraction"
+usle_k = "usle_k"
+usle_c = "USLE_C"
+pathfrac = "PathFrac"
+
+[input.vertical.overland_flow_erosion.parameters]
+usle_k = "usle_k"
+usle_c = "USLE_C"
+answers_k = "eros_ov"
+
+[input.vertical.soil_erosion.parameters]
+clay_fraction = "fclay_soil"
+silt_fraction = "fsilt_soil"
+sand_fraction = "fsand_soil"
+sagg_fraction = "fsagg_soil"
+lagg_fraction = "flagg_soil"
+
+[input.lateral.land.hydrological_forcing]
+waterlevel_land = "levKinL"
+q_land = "runL"
+
+[input.lateral.land.transport_capacity.parameters]
+density = "sediment_density"
+d50 = "d50_soil"
+c_govers = "c_govers"
+n_govers = "n_govers"
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+
+[input.lateral.river.hydrological_forcing]
+waterlevel_river = "h"
+q_river = "q"
+
+[input.lateral.river_parameter_set]
 length = "wflow_riverlength"
 slope = "RiverSlope"
 width = "wflow_riverwidth"
+
+[input.lateral.river.transport_capacity.parameters]
+density = "sediment_density"
+d50 = "D50_River"
+c_bagnold = "c_Bagnold"
+e_bagnold = "exp_Bagnold"
+a_kodatie = "a_kodatie"
+b_kodatie = "b_kodatie"
+c_kodatie = "c_kodatie"
+d_kodatie = "d_kodatie"
+
+[input.lateral.river.potential_erosion.parameters]
+d50 = "D50_River"
+
+[input.lateral.river.sediment_flux.parameters]
+clay_fraction = "ClayF_River"
+silt_fraction = "SiltF_River"
+sand_fraction = "SandF_River"
+gravel_fraction = "GravelF_River"
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+dm_gravel = "dm_gravel"
 # Reservoir
 resarea = "ResSimpleArea"
 restrapeff = "ResTrapEff"
-resareas = "wflow_reservoirareas"
 reslocs = "wflow_reservoirlocs"
 # Lake
 lakearea = "LakeArea"
-lakeareas = "wflow_lakeareas"
 lakelocs = "wflow_lakelocs"
 
+[input.lateral.river.concentrations.parameters]
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+dm_gravel = "dm_gravel"
+
 [model]
-dolake = false
-doreservoir = true
-landtransportmethod = "yalinpart" # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
-rainerosmethod = "answers" # Rainfall erosion equation: ["answers", "eurosem"]
 reinit = true
-rivtransportmethod = "bagnold" # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
-runrivermodel = true
+run_river_model = true
+dolake = false
+doreservoir = true # cannot use reservoirs as in sbm because then states/volume need to be added
+rainfall_erosion = "answers" # Rainfall erosion equation: ["answers", "eurosem"]
+overland_flow_erosion = "answers" # Overland flow erosion equation: ["answers"]
+land_transport = "yalinpart" # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
+river_transport = "bagnold" # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
 type = "sediment"
 
 [output]
 path = "output-moselle-sed.nc"
 
-[output.vertical]
-TCclay = "TCclay"
-TCsed = "TCsed"
-erosclay = "erosclay"
-pathfrac = "pathfrac"
-precipitation = "prec"
-sedov = "sedov"
-sedspl = "sedspl"
-soilloss = "soilloss"
-
-[output.lateral.land]
-inlandclay = "inlandclay"
-inlandsed = "inlandsed"
-olclay = "olclay"
-olsed = "olsed"
-
-[output.lateral.river]
-Bedconc = "Bedconc"
-SSconc = "SSconc"
-Sedconc = "Sedconc"
-clayload = "clayload"
-h_riv = "h_riv"
-inlandclay = "inlandclayriv"
-outclay = "outclay"
-width = "widthriv"
+[output.vertical.rainfall_erosion.variables]
+amount = "rainfall_erosion"
+
+[output.vertical.overland_flow_erosion.variables]
+amount = "overland_flow_erosion"
+
+[output.vertical.soil_erosion.variables]
+amount = "soilloss"
+clay = "erosclay"
+
+[output.lateral.land.transport_capacity.variables]
+clay = "TCclay"
+amount = "TCsed"
+
+[output.lateral.land.to_river.variables]
+clay = "inlandclay"
+amount = "inlandsed"
+
+[output.lateral.land.sediment_flux.variables]
+clay = "olclay"
+amount = "olsed"
+
+[output.lateral.river.concentrations.variables]
+bed = "Bedconc"
+suspended = "SSconc"
+total = "Sedconc"
+
+[output.lateral.river.hydrological_forcing]
+waterlevel_river = "h_riv"
+
+[output.lateral.river.sediment_flux.variables]
+leftover_clay = "clayload"
+clay = "outclay"
+erosion = "erosion_riv"
+deposition = "deposition_riv"
+
+[output.lateral.river.sediment_flux.boundary_conditions]
+erosion_land_clay = "inlandclayriv"
+transport_capacity = "TCsed_riv"
 
 [csv]
 path = "output-moselle-sediment.csv"
@@ -160,46 +217,46 @@ path = "output-moselle-sediment.csv"
 coordinate.x = 6.931
 coordinate.y = 48.085
 header = "SL"
-parameter = "vertical.soilloss"
+parameter = "vertical.soil_erosion.variables.amount"
 
 [[csv.column]]
 coordinate.x = 6.931
 coordinate.y = 48.085
 header = "SSPL"
-parameter = "vertical.sedspl"
+parameter = "vertical.rainfall_erosion.variables.amount"
 
 [[csv.column]]
 coordinate.x = 6.931
 coordinate.y = 48.085
 header = "SOV"
-parameter = "vertical.sedov"
+parameter = "vertical.overland_flow_erosion.variables.amount"
 
 [[csv.column]]
 coordinate.x = 6.931
 coordinate.y = 48.085
 header = "P"
-parameter = "vertical.precipitation"
+parameter = "vertical.atmospheric_forcing.precipitation"
 
 [[csv.column]]
 coordinate.x = 6.931
 coordinate.y = 48.085
 header = "ql"
-parameter = "vertical.q_land"
+parameter = "vertical.hydrological_forcing.q_land"
 
-[[csv.column]]
-coordinate.x = 6.931
-coordinate.y = 48.085
-header = "TCsed"
-parameter = "vertical.TCsed"
+# [[csv.column]]
+# coordinate.x = 6.931
+# coordinate.y = 48.085
+# header = "TCsed"
+# parameter = "lateral.land.transport_capacity.variables.amount"
 
-[[csv.column]]
-coordinate.x = 6.931
-coordinate.y = 48.085
-header = "TCclay"
-parameter = "vertical.TCclay"
+# [[csv.column]]
+# coordinate.x = 6.931
+# coordinate.y = 48.085
+# header = "TCclay"
+# parameter = "lateral.land.transport_capacity.variables.clay"
 
-[[csv.column]]
-coordinate.x = 6.931
-coordinate.y = 48.085
-header = "inlandsed"
-parameter = "lateral.land.inlandsed"
+# [[csv.column]]
+# coordinate.x = 6.931
+# coordinate.y = 48.085
+# header = "inlandsed"
+# parameter = "lateral.land.to_river.variables.amount"
diff --git a/test/sediment_eurosem_engelund_config.toml b/test/sediment_eurosem_engelund_config.toml
new file mode 100644
index 000000000..6c9fb3be8
--- /dev/null
+++ b/test/sediment_eurosem_engelund_config.toml
@@ -0,0 +1,170 @@
+# This is a TOML configuration file for Wflow.
+# Relative file paths are interpreted as being relative to this TOML file.
+# Wflow documentation https://deltares.github.io/Wflow.jl/dev/
+# TOML documentation: https://github.com/toml-lang/toml
+
+calendar = "proleptic_gregorian"
+endtime = 2000-01-03T00:00:00
+starttime = 1999-12-31T00:00:00
+time_units = "days since 1900-01-01 00:00:00"
+timestepsecs = 86400
+dir_input = "data/input"
+dir_output = "data/output"
+
+[state]
+path_input = "instates-moselle-sed.nc"
+path_output = "outstates-moselle-sed.nc"
+
+# if listed, the variable must be present in the NetCDF or error
+# if not listed, the variable can get a default value if it has one
+
+[state.lateral.river.sediment_flux.variables]
+leftover_clay = "clayload"
+store_clay = "claystore"
+leftover_gravel = "gravload"
+store_gravel = "gravstore"
+leftover_lagg = "laggload"
+store_lagg = "laggstore"
+clay = "outclay"
+gravel = "outgrav"
+lagg= "outlagg"
+sagg = "outsagg"
+sand = "outsand"
+silt = "outsilt"
+leftover_sagg= "saggload"
+store_sagg = "saggstore"
+leftover_sand = "sandload"
+store_sand = "sandstore"
+leftover_silt = "siltload"
+store_silt = "siltstore"
+
+[input]
+path_forcing = "forcing-moselle-sed.nc"
+path_static = "staticmaps-moselle-sed.nc"
+
+# these are not directly part of the model
+gauges = "wflow_gauges"
+ldd = "wflow_ldd"
+river_location = "wflow_river"
+subcatchment = "wflow_subcatch"
+reservoir_areas = "wflow_reservoirareas"
+lake_areas = "wflow_lakeareas"
+
+# specify the internal IDs of the parameters which vary over time
+# the external name mapping needs to be below together with the other mappings
+forcing = [
+  "vertical.atmospheric_forcing.precipitation",
+  "vertical.hydrological_forcing.interception",
+  "vertical.hydrological_forcing.waterlevel_land",
+  "lateral.land.hydrological_forcing.waterlevel_land",
+  "vertical.hydrological_forcing.q_land",
+  "lateral.land.hydrological_forcing.q_land",
+  "lateral.river.hydrological_forcing.waterlevel_river",
+  "lateral.river.hydrological_forcing.q_river",
+]
+
+[input.vertical.atmospheric_forcing]
+precipitation = "P"
+
+[input.vertical.hydrological_forcing]
+waterlevel_land = "levKinL"
+q_land = "runL"
+interception = "int"
+
+[input.vertical.land_parameter_set]
+slope = "Slope"
+
+[input.vertical.rainfall_erosion.parameters]
+soil_detachability = "soil_detachability"
+eurosem_exponent = "eros_spl_EUROSEM"
+canopyheight = "CanopyHeight"
+usle_k = "usle_k"
+usle_c = "USLE_C"
+pathfrac = "PathFrac"
+
+[input.vertical.overland_flow_erosion.parameters]
+usle_k = "usle_k"
+usle_c = "USLE_C"
+answers_k = "eros_ov"
+
+[input.vertical.soil_erosion.parameters]
+clay_fraction = "fclay_soil"
+silt_fraction = "fsilt_soil"
+sand_fraction = "fsand_soil"
+sagg_fraction = "fsagg_soil"
+lagg_fraction = "flagg_soil"
+
+[input.lateral.land.hydrological_forcing]
+waterlevel_land = "levKinL"
+q_land = "runL"
+
+[input.lateral.land.transport_capacity.parameters]
+density = "sediment_density"
+d50 = "d50_soil"
+c_govers = "c_govers"
+n_govers = "n_govers"
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+
+[input.lateral.river.hydrological_forcing]
+waterlevel_river = "h"
+q_river = "q"
+
+[input.lateral.river_parameter_set]
+length = "wflow_riverlength"
+slope = "RiverSlope"
+width = "wflow_riverwidth"
+
+[input.lateral.river.transport_capacity.parameters]
+density = "sediment_density"
+d50 = "D50_River"
+c_bagnold = "c_Bagnold"
+e_bagnold = "exp_Bagnold"
+a_kodatie.value = 2829.6
+b_kodatie.value = 3.646
+c_kodatie.value = 0.406
+d_kodatie.value = 0.412
+
+[input.lateral.river.potential_erosion.parameters]
+d50 = "D50_River"
+
+[input.lateral.river.sediment_flux.parameters]
+clay_fraction = "ClayF_River"
+silt_fraction = "SiltF_River"
+sand_fraction = "SandF_River"
+gravel_fraction = "GravelF_River"
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+dm_gravel = "dm_gravel"
+# Reservoir
+resarea = "ResSimpleArea"
+restrapeff = "ResTrapEff"
+reslocs = "wflow_reservoirlocs"
+# Lake
+lakearea = "LakeArea"
+lakelocs = "wflow_lakelocs"
+
+[input.lateral.river.concentrations.parameters]
+dm_clay = "dm_clay"
+dm_silt = "dm_silt"
+dm_sand = "dm_sand"
+dm_sagg = "dm_sagg"
+dm_lagg = "dm_lagg"
+dm_gravel = "dm_gravel"
+
+[model]
+reinit = true
+run_river_model = true
+dolake = false
+doreservoir = true # cannot use reservoirs as in sbm because then states/volume need to be added
+rainfall_erosion = "eurosem" # Rainfall erosion equation: ["answers", "eurosem"]
+overland_flow_erosion = "answers" # Overland flow erosion equation: ["answers"]
+land_transport = "yalinpart" # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
+river_transport = "engelund" # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
+type = "sediment"