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

Issue:
Check if all states are provided when using `reinit=false`
### Wflow version checks

- [X] I have checked that this issue has not already been reported.
- [X] I have checked that this bug exists on the latest version of Wflow.


### Reproducible Example

If you start a wflow_sbm model with `glacier=true`, and don't pass the glacierstore in the state section, there is no error (or warning thrown.



### Current behaviour

See above

### Desired behaviour

When using `reinit=false`, it might be worth that Wflow checks whether all expected states are provided by the instate file. If not, at least a warning should be thrown (and perhaps even an error). 

### Additional Context

_No response_


Discussion / hints:

Fix patch:
diff --git a/build/create_binaries/add_metadata.jl b/build/create_binaries/add_metadata.jl
index a508499a1..64e8fd441 100644
--- a/build/create_binaries/add_metadata.jl
+++ b/build/create_binaries/add_metadata.jl
@@ -20,18 +20,18 @@ function add_metadata(project_dir, license_file, output_dir, git_repo)
     cp(
         normpath(project_dir, "Project.toml"),
         normpath(output_dir, "share/julia/Project.toml");
-        force=true
+        force = true,
     )
     # the Manifest.toml always gives the exact version of Wflow that was built
     cp(
         normpath(project_dir, "Manifest.toml"),
         normpath(output_dir, "share/julia/Manifest.toml");
-        force=true
+        force = true,
     )
 
     # put the LICENSE in the top level directory
-    cp(license_file, normpath(output_dir, "LICENSE"); force=true)
-    cp(normpath(project_dir, "README.md"), normpath(output_dir, "README.md"); force=true)
+    cp(license_file, normpath(output_dir, "LICENSE"); force = true)
+    cp(normpath(project_dir, "README.md"), normpath(output_dir, "README.md"); force = true)
     open(normpath(output_dir, "README.md"), "a") do io
         # since the exact Wflow version may be hard to find in the Manifest.toml file
         # we can also extract that information, and add it to the README.md
diff --git a/build/create_binaries/create_app.jl b/build/create_binaries/create_app.jl
index 4509ef8bb..28d8010a9 100644
--- a/build/create_binaries/create_app.jl
+++ b/build/create_binaries/create_app.jl
@@ -15,10 +15,10 @@ create_app(
     project_dir,
     output_dir;
     # map from binary name to julia function name
-    executables=["wflow_cli" => "julia_main"],
-    precompile_execution_file="precompile.jl",
-    filter_stdlibs=false,
-    force=true,
+    executables = ["wflow_cli" => "julia_main"],
+    precompile_execution_file = "precompile.jl",
+    filter_stdlibs = false,
+    force = true,
 )
 
 include("add_metadata.jl")
diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 0de9dadd4..0b1f22c54 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -29,6 +29,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Total water storage as an export variable for `SBM` concept. This is the total water stored
   per grid cell in millimeters. Excluded from this variable are the floodplain, lakes and
   reservoirs.
+- Checks to see if all states are covered in the .toml file. If not all states are covered,
+  an error is thrown. If there are more states specified than required, these states are
+  ignored (with a warning in the logging), and the simulation will continue.
 
 ## v0.7.3 - 2024-01-12
 
@@ -47,7 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   terms of the quadratic equation (and solution) were fixed.
 - Use `kvfrac` for the computation of vertical saturated hydraulic conductivity at the
   bottom of the soil layer, since `kvfrac` is also used for the computation of vertical
-  unsaturated flow. 
+  unsaturated flow.
 
 ### Changed
 - For cyclic parameters different cyclic time inputs are supported (only one common cyclic
diff --git a/src/Wflow.jl b/src/Wflow.jl
index 3bf4b03d8..7b511596f 100644
--- a/src/Wflow.jl
+++ b/src/Wflow.jl
@@ -52,7 +52,7 @@ end
 
 function Clock(config, reader)
     nctimes = reader.dataset["time"][:]
-    
+
     # if the timestep is not given, use the difference between NetCDF time 1 and 2
     timestepsecs = get(config, "timestepsecs", nothing)
     if timestepsecs === nothing
@@ -60,7 +60,7 @@ function Clock(config, reader)
         config.timestepsecs = timestepsecs
     end
     Δt = Second(timestepsecs)
-    
+
     # if the config file does not have a start or endtime, follow the NetCDF times
     # and add them to the config
     starttime = get(config, "starttime", nothing)
@@ -80,7 +80,7 @@ function Clock(config, reader)
         config.starttime = starttime + Δt
     end
     starttime = cftime(config.starttime, calendar)
-    
+
     Clock(starttime, 0, Δt)
 end
 
@@ -133,6 +133,7 @@ include("utils.jl")
 include("bmi.jl")
 include("subdomains.jl")
 include("logging.jl")
+include("states.jl")
 
 """
     run(tomlpath::AbstractString; silent=false)
@@ -218,8 +219,8 @@ function run(model::Model; close_files = true)
     # determine timesteps to run
     calendar = get(config, "calendar", "standard")::String
     @warn string(
-        "The definition of `starttime` has changed (equal to model state time).\n Please", 
-        " update your settings TOML file by subtracting one model timestep Δt from the", 
+        "The definition of `starttime` has changed (equal to model state time).\n Please",
+        " update your settings TOML file by subtracting one model timestep Δt from the",
         " `starttime`, if it was used with a Wflow version up to v0.6.3.",
     )
     starttime = clock.time
diff --git a/src/flextopo.jl b/src/flextopo.jl
index d02d79713..ea2db00d4 100644
--- a/src/flextopo.jl
+++ b/src/flextopo.jl
@@ -6,17 +6,17 @@
     # Number of cells
     n::Int | "-" | 0 | "none" | "none"
     #dictionary with all possible functions for each store
-    dic_function::Dict |"-"| 0 | "none" | "none"
+    dic_function::Dict | "-" | 0 | "none" | "none"
     #current class
-    kclass::Vector{Int64} |"-"| 0 | "none" | "none"
-    classes::Vector{String} |"-"| 0 | "none" | "none"
-    select_snow::Vector{String} |"-"| 0 | "none" | "none"
-    select_interception::Vector{String} |"-"| 0 | "none" | "none"
-    select_hortonponding::Vector{String} |"-"| 0 | "none" | "none"
-    select_hortonrunoff::Vector{String} |"-"| 0 | "none" | "none"
-    select_rootzone::Vector{String} |"-"| 0 | "none" | "none"
-    select_fast::Vector{String} |"-"| 0 | "none" | "none"
-    select_slow::Vector{String} |"-"| 0 | "none" | "none"
+    kclass::Vector{Int64} | "-" | 0 | "none" | "none"
+    classes::Vector{String} | "-" | 0 | "none" | "none"
+    select_snow::Vector{String} | "-" | 0 | "none" | "none"
+    select_interception::Vector{String} | "-" | 0 | "none" | "none"
+    select_hortonponding::Vector{String} | "-" | 0 | "none" | "none"
+    select_hortonrunoff::Vector{String} | "-" | 0 | "none" | "none"
+    select_rootzone::Vector{String} | "-" | 0 | "none" | "none"
+    select_fast::Vector{String} | "-" | 0 | "none" | "none"
+    select_slow::Vector{String} | "-" | 0 | "none" | "none"
 
     #fraction of each class
     hrufrac::Vector{SVector{N,T}} | "-"
@@ -223,16 +223,6 @@
 
 end
 
-statevars(::FLEXTOPO) = (
-    :snow,
-    :snowwater,
-    :interceptionstorage,
-    :hortonpondingstorage,
-    :hortonrunoffstorage,
-    :rootzonestorage,
-    :faststorage,
-    :slowstorage,
-)
 
 function common_snow_hbv(flextopo::FLEXTOPO)
     for i = 1:flextopo.n
diff --git a/src/flextopo_model.jl b/src/flextopo_model.jl
index 048359e7e..26ba13d6b 100644
--- a/src/flextopo_model.jl
+++ b/src/flextopo_model.jl
@@ -722,8 +722,7 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:FlextopoModel
     # read and set states in model object if reinit=true
     if reinit == false
         instate_path = input_path(config, config.state.path_input)
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames; type = Float, dimname = :classes)
+        set_states(instate_path, model; type = Float, dimname = :classes)
 
         # update kinematic wave volume for river and land domain
         lateral.land.volume .= lateral.land.h .* lateral.land.width .* lateral.land.dl
@@ -738,4 +737,4 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:FlextopoModel
         end
     end
     return model
-end
\ No newline at end of file
+end
diff --git a/src/flow.jl b/src/flow.jl
index df3aec6d2..c0c0a1515 100644
--- a/src/flow.jl
+++ b/src/flow.jl
@@ -1,7 +1,8 @@
 
 abstract type SurfaceFlow end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowRiver{T,R,L} <: SurfaceFlow
+@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowRiver{T,R,L} <:
+                                                               SurfaceFlow
     β::T | "-" | 0 | "scalar"                    # constant in Manning's equation
     sl::Vector{T} | "m m-1"                      # Slope [m m⁻¹]
     n::Vector{T} | "s m-1/3"                     # Manning's roughness [s m⁻⅓]
@@ -36,8 +37,9 @@ abstract type SurfaceFlow end
     # end
 end
 
-@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowLand{T} <: SurfaceFlow
-    β::T | "-"  | 0 | "scalar"                      # constant in Manning's equation
+@get_units @exchange @grid_type @grid_location @with_kw struct SurfaceFlowLand{T} <:
+                                                               SurfaceFlow
+    β::T | "-" | 0 | "scalar"                       # constant in Manning's equation
     sl::Vector{T} | "m m-1"                         # Slope [m m⁻¹]
     n::Vector{T} | "s m-1/3"                        # Manning's roughness [s m⁻⅓]
     dl::Vector{T} | "m"                             # Drain length [m]
@@ -57,7 +59,7 @@ end
     α::Vector{T} | "s3/5 m1/5"                      # Constant in momentum equation A = αQᵝ, based on Manning's equation
     cel::Vector{T} | "m s-1"                        # Celerity of the kinematic wave
     to_river::Vector{T} | "m3 s-1"                  # Part of overland flow [m³ s⁻¹] that flows to the river
-    kinwave_it::Bool | "-" | 0 |  "none" | "none"   # Boolean for iterations kinematic wave
+    kinwave_it::Bool | "-" | 0 | "none" | "none"    # Boolean for iterations kinematic wave
 end
 
 function initialize_surfaceflow_land(nc, config, inds; sl, dl, width, iterate, tstep, Δt)
@@ -177,8 +179,6 @@ function initialize_surfaceflow_river(
     return sf_river
 end
 
-statevars(::SurfaceFlowRiver) = (:q, :h, :h_av)
-statevars(::SurfaceFlowLand) = (:q, :h, :h_av)
 
 function update(sf::SurfaceFlowLand, network, frac_toriver)
     @unpack graph, subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes =
@@ -388,7 +388,7 @@ function stable_timestep(sf::S) where {S<:SurfaceFlow}
     return Δt, its
 end
 
-@get_units @exchange @grid_type  @grid_location @with_kw struct LateralSSF{T}
+@get_units @exchange @grid_type @grid_location @with_kw struct LateralSSF{T}
     kh₀::Vector{T} | "m d-1"               # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
     f::Vector{T} | "m-1"                   # A scaling parameter [m⁻¹] (controls exponential decline of kh₀)
     soilthickness::Vector{T} | "m"         # Soil thickness [m]
@@ -412,7 +412,6 @@ end
     end
 end
 
-statevars(::LateralSSF) = (:ssf,)
 
 function update(ssf::LateralSSF, network, frac_toriver)
     @unpack subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes = network
@@ -475,7 +474,7 @@ end
     h_thresh::T | "m" | 0 | "scalar"                        # depth threshold for calculating flow
     Δt::T | "s" | 0 | "none" | "none"                       # model time step [s]
     q::Vector{T} | "m3 s-1" | _ | "edge"                    # river discharge (subgrid channel)
-    q0::Vector{T} | "m3 s-1"| _ | "edge"                    # river discharge (subgrid channel) at previous time step
+    q0::Vector{T} | "m3 s-1" | _ | "edge"                   # river discharge (subgrid channel) at previous time step
     q_av::Vector{T} | "m3 s-1" | _ | "edge"                 # average river channel (+ floodplain) discharge [m³ s⁻¹]
     q_channel_av::Vector{T} | "m3 s-1"                      # average river channel discharge [m³ s⁻¹]
     zb_max::Vector{T} | "m"                                 # maximum channel bed elevation
@@ -489,7 +488,7 @@ end
     h_av::Vector{T} | "m"                                   # average water depth
     dl::Vector{T} | "m"                                     # river length
     dl_at_link::Vector{T} | "m" | _ | "edge"                # river length at edge/link
-    width::Vector{T} | "m"                      	        # river width
+    width::Vector{T} | "m"                                  # river width
     width_at_link::Vector{T} | "m" | _ | "edge"             # river width at edge/link
     a::Vector{T} | "m2" | _ | "edge"                        # flow area at edge/link
     r::Vector{T} | "m" | _ | "edge"                         # wetted perimeter at edge/link
@@ -972,7 +971,7 @@ const dirs = (:yd, :xd, :xu, :yu)
     runoff::Vector{T} | "m3 s-1"                            # runoff from hydrological model
     h::Vector{T} | "m"                                      # water depth of cell (for river cells the reference is the river bed elevation `zb`)
     z::Vector{T} | "m"                                      # elevation of cell
-    froude_limit::Bool | "-" | 0  |"none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
+    froude_limit::Bool | "-" | 0 | "none" | "none"          # if true a check is performed if froude number > 1.0 (algorithm is modified)
     rivercells::Vector{Bool} | "-"                          # river cells
     h_av::Vector{T} | "m"                                   # average water depth (for river cells the reference is the river bed elevation `zb`)
 end
@@ -1319,7 +1318,7 @@ function update(sw::ShallowWaterLand{T}, swr::ShallowWaterRiver{T}, network, Δt
     end
 end
 
-""" 
+"""
     FloodPlainProfile
 
 Floodplain `volume` is a function of `depth` (flood depth intervals). Based on the
diff --git a/src/groundwater/boundary_conditions.jl b/src/groundwater/boundary_conditions.jl
index 581bc9070..f38e6d321 100644
--- a/src/groundwater/boundary_conditions.jl
+++ b/src/groundwater/boundary_conditions.jl
@@ -41,7 +41,8 @@ function flux!(Q, river::River, aquifer)
 end
 
 
-@get_units @exchange @grid_type @grid_location struct Drainage{T} <: AquiferBoundaryCondition
+@get_units @exchange @grid_type @grid_location struct Drainage{T} <:
+                                                      AquiferBoundaryCondition
     elevation::Vector{T} | "m"
     conductance::Vector{T} | "m2 d-1"
     flux::Vector{T} | "m3 d-1"
@@ -79,7 +80,8 @@ function flux!(Q, headboundary::HeadBoundary, aquifer)
 end
 
 
-@get_units @exchange @grid_type @grid_location struct Recharge{T} <: AquiferBoundaryCondition
+@get_units @exchange @grid_type @grid_location struct Recharge{T} <:
+                                                      AquiferBoundaryCondition
     rate::Vector{T} | "m d-1"
     flux::Vector{T} | "m3 d-1"
     index::Vector{Int} | "-"
diff --git a/src/hbv.jl b/src/hbv.jl
index df69ad78b..dd22bb309 100644
--- a/src/hbv.jl
+++ b/src/hbv.jl
@@ -57,14 +57,6 @@
     runoff::Vector{T}                   # Total specific runoff per cell [mm Δt⁻¹]
 end
 
-statevars(::HBV) = (
-    :soilmoisture,
-    :snow,
-    :snowwater,
-    :upperzonestorage,
-    :lowerzonestorage,
-    :interceptionstorage,
-)
 
 function update_until_snow(hbv::HBV, config)
 
diff --git a/src/hbv_model.jl b/src/hbv_model.jl
index a03298d71..877127bab 100644
--- a/src/hbv_model.jl
+++ b/src/hbv_model.jl
@@ -421,8 +421,7 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:HbvModel}
     if reinit == false
         instate_path = input_path(config, config.state.path_input)
         @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames; type = Float)
+        set_states(instate_path, model; type = Float)
         # update kinematic wave volume for river and land domain
         @unpack lateral = model
         # makes sure land cells with zero flow width are set to zero q and h
@@ -447,4 +446,4 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:HbvModel}
     end
 
     return model
-end
\ No newline at end of file
+end
diff --git a/src/io.jl b/src/io.jl
index 648fe2f49..56447090d 100644
--- a/src/io.jl
+++ b/src/io.jl
@@ -955,7 +955,7 @@ function prepare_writer(
     # create a separate state output NetCDF that will hold the last timestep of all states
     # but only if config.state.path_output has been set
     if haskey(config, "state") && haskey(config.state, "path_output")
-        state_ncnames = ncnames(config.state)
+        state_ncnames = check_states(config)
         state_map = out_map(state_ncnames, modelmap)
         nc_state_path = output_path(config, config.state.path_output)
         @info "Create a state output NetCDF file `$nc_state_path`."
diff --git a/src/reservoir_lake.jl b/src/reservoir_lake.jl
index 110d2edff..1daf39d69 100644
--- a/src/reservoir_lake.jl
+++ b/src/reservoir_lake.jl
@@ -22,7 +22,6 @@
     end
 end
 
-statevars(::SimpleReservoir) = (:volume,)
 
 function initialize_simple_reservoir(config, nc, inds_riv, nriv, pits, Δt)
     # read only reservoir data if reservoirs true
@@ -461,18 +460,6 @@ function maximum_storage(storfunc, outflowfunc, area, sh, hq)
     return maxstorage
 end
 
-statevars(::Lake) = (:waterlevel,)
-
-"""
-    statevars(::SBM; snow=false)
-    statevars(::LateralSSF)
-    statevars(::SurfaceFlow)
-    statevars(::NaturalLake)
-    statevars(::SimpleReservoir)
-
-Get a tuple of symbols representing the fields that are model states.
-"""
-function statevars end
 
 function interpolate_linear(x, xp, fp)
     if x <= minimum(xp)
diff --git a/src/sbm.jl b/src/sbm.jl
index 9cc336074..77c44347c 100644
--- a/src/sbm.jl
+++ b/src/sbm.jl
@@ -203,9 +203,6 @@
     end
 end
 
-statevars(::SBM; snow = false) =
-    snow ? (:satwaterdepth, :snow, :tsoil, :ustorelayerdepth, :snowwater, :canopystorage) :
-    (:satwaterdepth, :ustorelayerdepth, :canopystorage)
 
 function initialize_canopy(nc, config, inds)
     n = length(inds)
@@ -599,7 +596,7 @@ function initialize_sbm(nc, config, riverfrac, inds)
         leaf_area_index = fill(mv, n),
         waterlevel_land = fill(mv, n),
         waterlevel_river = zeros(Float, n), #set to zero to account for cells outside river domain
-        total_storage = zeros(Float, n) # Set the total water storage from initialized values
+        total_storage = zeros(Float, n), # Set the total water storage from initialized values
     )
 
     return sbm
@@ -613,7 +610,7 @@ function update_until_snow(sbm::SBM, config)
     modelglacier = get(config.model, "glacier", false)::Bool
     modelsnow = get(config.model, "snow", false)::Bool
 
-    threaded_foreach(1:sbm.n, basesize=1000) do i
+    threaded_foreach(1:sbm.n, basesize = 1000) do i
         if do_lai
             cmax = sbm.sl[i] * sbm.leaf_area_index[i] + sbm.swood[i]
             canopygapfraction = exp(-sbm.kext[i] * sbm.leaf_area_index[i])
@@ -696,7 +693,7 @@ function update_until_recharge(sbm::SBM, config)
     transfermethod = get(config.model, "transfermethod", false)::Bool
     ust = get(config.model, "whole_ust_available", false)::Bool # should be removed from optional setting and code?
 
-    threaded_foreach(1:sbm.n, basesize=250) do i
+    threaded_foreach(1:sbm.n, basesize = 250) do i
         if modelsnow
             rainfallplusmelt = sbm.rainfallplusmelt[i]
             if modelglacier
@@ -977,7 +974,7 @@ end
 
 function update_after_subsurfaceflow(sbm::SBM, zi, exfiltsatwater)
 
-    threaded_foreach(1:sbm.n, basesize=1000) do i
+    threaded_foreach(1:sbm.n, basesize = 1000) do i
         usl, n_usl = set_layerthickness(zi[i], sbm.sumlayers[i], sbm.act_thickl[i])
         # exfiltration from ustore
         usld = sbm.ustorelayerdepth[i]
@@ -1071,12 +1068,12 @@ Takes the following parameters:
     The land routing struct, i.e. model.lateral.land
 """
 function update_total_water_storage(
-    sbm::SBM, 
-    river_network, 
-    cell_xsize, 
-    cell_ysize, 
-    river_routing, 
-    land_routing
+    sbm::SBM,
+    river_network,
+    cell_xsize,
+    cell_ysize,
+    river_routing,
+    land_routing,
 )
     # Get length active river cells
     nriv = length(river_network)
@@ -1086,29 +1083,29 @@ function update_total_water_storage(
 
     # Burn the river routing values
     sbm.total_storage[river_network] = (
-        ( river_routing.h_av[1:nriv] .* river_routing.width[1:nriv] .* 
-        river_routing.dl[1:nriv] ) ./
-        ( cell_xsize[river_network] .* cell_ysize[river_network] ) * 
-        1000 # Convert to mm
+        (
+            river_routing.h_av[1:nriv] .* river_routing.width[1:nriv] .*
+            river_routing.dl[1:nriv]
+        ) ./ (cell_xsize[river_network] .* cell_ysize[river_network]) * 1000 # Convert to mm
     )
 
     # Chunk the data for parallel computing
-    threaded_foreach(1:sbm.n, basesize=1000) do i
+    threaded_foreach(1:sbm.n, basesize = 1000) do i
 
         # Cumulate per vertical type
         # Maybe re-categorize in the future
         surface = (
             sbm.glacierstore[i] * sbm.glacierfrac[i] +
-            sbm.snow[i] + sbm.snowwater[i] + sbm.canopystorage[i]
+            sbm.snow[i] +
+            sbm.snowwater[i] +
+            sbm.canopystorage[i]
         )
         sub_surface = sbm.ustoredepth[i] + sbm.satwaterdepth[i]
         lateral = (
             land_routing.h_av[i] * (1 - sbm.riverfrac[i]) * 1000 # convert to mm
         )
-        
+
         # Add everything to the total water storage
-        sbm.total_storage[i] += (
-            surface + sub_surface + lateral
-        )
+        sbm.total_storage[i] += (surface + sub_surface + lateral)
     end
 end
diff --git a/src/sbm_gwf_model.jl b/src/sbm_gwf_model.jl
index f723bc5df..c8be4c8db 100644
--- a/src/sbm_gwf_model.jl
+++ b/src/sbm_gwf_model.jl
@@ -471,8 +471,7 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
     if reinit == false
         instate_path = input_path(config, config.state.path_input)
         @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames, type = Float, dimname = :layer)
+        set_states(instate_path, model, type = Float, dimname = :layer)
         # update kinematic wave volume for river and land domain
         @unpack lateral = model
         # makes sure land cells with zero flow width are set to zero q and h
@@ -496,4 +495,4 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
         @info "Set initial conditions from default values."
     end
     return model
-end
\ No newline at end of file
+end
diff --git a/src/sbm_model.jl b/src/sbm_model.jl
index 3189677be..085fca0d4 100644
--- a/src/sbm_model.jl
+++ b/src/sbm_model.jl
@@ -444,9 +444,7 @@ Update of the total water storage at the end of each timestep per model cell.
 
 This is done here at model level.
 """
-function update_total_water_storage(
-    model::Model{N,L,V,R,W,T},
-) where {N,L,V,R,W,T<:SbmModel}
+function update_total_water_storage(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
     @unpack lateral, vertical, network, clock, config = model
 
     # Update the total water storage based on vertical states
@@ -477,12 +475,11 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
         nriv = length(network.river.indices)
         instate_path = input_path(config, config.state.path_input)
         @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
         @warn string(
             "The unit of `ssf` (lateral subsurface flow) is now m3 d-1. Please update your",
             " input state file if it was produced with a Wflow version up to v0.5.2.",
         )
-        set_states(instate_path, model, state_ncnames; type = Float, dimname = :layer)
+        set_states(instate_path, model; type = Float, dimname = :layer)
         # update zi for vertical sbm and kinematic wave volume for river and land domain
         zi =
             max.(
@@ -534,7 +531,7 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmModel}
         if floodplain_1d
             initialize_volume!(lateral.river, nriv)
         end
-    
+
         if do_lakes
             # storage must be re-initialized after loading the state with the current
             # waterlevel otherwise the storage will be based on the initial water level
diff --git a/src/sediment.jl b/src/sediment.jl
index 773f6ee5a..f77ec9039 100644
--- a/src/sediment.jl
+++ b/src/sediment.jl
@@ -99,7 +99,6 @@
     end
 end
 
-statevars(::LandSediment) = ()
 
 function initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
     # Initialize parameters for the soil loss part
@@ -636,7 +635,6 @@ end
     end
 end
 
-statevars(::OverlandFlowSediment) = ()
 
 function partial_update!(inland, rivcell, eroded)
     no_erosion = zero(eltype(eroded))
@@ -785,26 +783,6 @@ end
     # end
 end
 
-statevars(::RiverSediment) = (
-    :clayload,
-    :siltload,
-    :sandload,
-    :saggload,
-    :laggload,
-    :gravload,
-    :claystore,
-    :siltstore,
-    :sandstore,
-    :saggstore,
-    :laggstore,
-    :gravstore,
-    :outclay,
-    :outsilt,
-    :outsand,
-    :outsagg,
-    :outlagg,
-    :outgrav,
-)
 
 function initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
     # Initialize river parameters
diff --git a/src/sediment_model.jl b/src/sediment_model.jl
index a1944263b..aea2ceb83 100644
--- a/src/sediment_model.jl
+++ b/src/sediment_model.jl
@@ -199,8 +199,7 @@ function set_states(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SedimentModel
     if reinit == false
         instate_path = input_path(config, config.state.path_input)
         @info "Set initial conditions from state file `$instate_path`."
-        state_ncnames = ncnames(config.state)
-        set_states(instate_path, model, state_ncnames; type = Float)
+        set_states(instate_path, model; type = Float)
     else
         @info "Set initial conditions from default values."
     end
diff --git a/src/states.jl b/src/states.jl
new file mode 100644
index 000000000..777754c3a
--- /dev/null
+++ b/src/states.jl
@@ -0,0 +1,250 @@
+"""
+    get_vertical_states(model_type; snow = false, glacier = false)
+
+Function to extract all required vertical states, given a certain model type. Passes the snow
+and glacier options only for the `sbm` model_type. Returns a tuple with the required states
+(internal names as symbols)
+"""
+function get_vertical_states(model_type::AbstractString; snow = false, glacier = false)
+    if model_type == "sbm" || model_type == "sbm_gwf"
+        if snow && glacier
+            vertical_states = (
+                :satwaterdepth,
+                :snow,
+                :tsoil,
+                :ustorelayerdepth,
+                :snowwater,
+                :canopystorage,
+                :glacierstore,
+            )
+        elseif snow
+            vertical_states = (
+                :satwaterdepth,
+                :snow,
+                :tsoil,
+                :ustorelayerdepth,
+                :snowwater,
+                :canopystorage,
+            )
+        else
+            vertical_states = (:satwaterdepth, :ustorelayerdepth, :canopystorage)
+        end
+    elseif model_type == "hbv"
+        vertical_states = (
+            :soilmoisture,
+            :snow,
+            :snowwater,
+            :upperzonestorage,
+            :lowerzonestorage,
+            :interceptionstorage,
+        )
+    elseif model_type == "flextopo"
+        vertical_states = (
+            :snow,
+            :snowwater,
+            :interceptionstorage,
+            :hortonpondingstorage,
+            :hortonrunoffstorage,
+            :rootzonestorage,
+            :faststorage,
+            :slowstorage,
+        )
+    elseif model_type == "sediment"
+        vertical_states = ()
+    else
+        throw(ArgumentError("Unknown model_type provided (`$model_type`)"))
+    end
+    return vertical_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)
+
+Function to iterate through the list of `states` and add these to the `required_states`
+variable. This is a tuple of tuples, that contains the correct "prefix" to the state
+category (provided by the `key_entry` variable). This is added in front of each element in
+the list of `states`, and the `required_states` tuple is returned.
+
+When `nothing` is passed as the `states`, the `required_states` variable is simply returned.
+"""
+function add_to_required_states(required_states::Tuple, key_entry::Tuple, states::Tuple)
+    for state in states
+        required_states = (required_states..., (key_entry..., state))
+    end
+
+    return required_states
+end
+
+add_to_required_states(required_states::Tuple, key_entry::Tuple, states::Nothing) =
+    required_states
+
+"""
+    extract_required_states(config::Config)
+
+Function to retrieve the required states, given the model configuration. The required states
+are inferred from the model settings (mainly model_type, model options and routing types).
+Returns as list of required states, in the same formats as the keys that are returned from the
+`ncnames` function.
+"""
+function extract_required_states(config::Config)
+    # Extract model type
+    model_type = config.model.type::String
+
+    # Extract model settings
+    do_snow = get(config.model, "snow", false)::Bool
+    do_glaciers = get(config.model, "glacier", false)::Bool
+    do_lakes = get(config.model, "lakes", false)::Bool
+    do_reservoirs = get(config.model, "reservoirs", false)::Bool
+    do_floodplains = get(config.model, "floodplain_1d", false)::Bool
+
+    # Extract required stated based on model configuration file
+    vertical_states = get_vertical_states(model_type; snow = do_snow, glacier = do_glaciers)
+
+    # Subsurface states
+    if model_type == "sbm_gwf"
+        ssf_states = (:head,)
+    else
+        ssf_states = haskey(config.input.lateral, "subsurface") ? (:ssf,) : nothing
+    end
+
+    # Land states
+    if model_type == "sediment"
+        land_states = ()
+    else
+        routing_options = ("kinematic-wave", "local-inertial")
+        land_routing = get_options(
+            config.model,
+            "land_routing",
+            routing_options,
+            "kinematic-wave",
+        )::String
+        if land_routing == "local-inertial"
+            land_states = (:qx, :qy, :h, :h_av)
+        elseif land_routing == "kinematic-wave"
+            if model_type == "sbm"
+                land_states = (:q, :h, :h_av)
+            else
+                land_states = (:q, :h)
+            end
+        end
+    end
+
+    # 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,
+        )
+    elseif model_type == "sbm"
+        river_states = (:q, :h, :h_av)
+    else
+        river_states = (:q, :h)
+    end
+
+    # Floodplain states
+    floodplain_states = do_floodplains ? (:q, :h) : nothing
+
+    # Lake and reservoir states
+    lake_states = do_lakes ? (:waterlevel,) : nothing
+    reservoir_states = do_reservoirs ? (:volume,) : nothing
+
+    # Build required states in a tuple, similar to the keys in the output of
+    # `ncnames(config.state)`
+    required_states = ()
+    # Add vertical states to dict
+    required_states = add_to_required_states(required_states, (:vertical,), vertical_states)
+    # Add subsurface states to dict
+    if model_type == "sbm_gwf"
+        key_entry = (:lateral, :subsurface, :flow, :aquifer)
+    else
+        key_entry = (:lateral, :subsurface)
+    end
+    required_states = add_to_required_states(required_states, key_entry, ssf_states)
+    # Add land states to dict
+    required_states =
+        add_to_required_states(required_states, (:lateral, :land), land_states)
+    # Add river states to dict
+    required_states =
+        add_to_required_states(required_states, (:lateral, :river), river_states)
+    # Add floodplain states to dict
+    required_states = add_to_required_states(
+        required_states,
+        (:lateral, :river, :floodplain),
+        floodplain_states,
+    )
+    # Add lake states to dict
+    required_states =
+        add_to_required_states(required_states, (:lateral, :river, :lake), lake_states)
+    # Add reservoir states to dict
+    required_states = add_to_required_states(
+        required_states,
+        (:lateral, :river, :reservoir),
+        reservoir_states,
+    )
+
+    return required_states
+end
+
+"""
+    check_states(config::Config)
+
+This function checks whether the states provided in the model configuration are covering all
+required states. If not all states are covered, an error is thrown to warn the user
+(including logging messages to explain which variables are missing). If more states are
+provided than required, a warning is logged, and the unnecessary state is removed.
+"""
+function check_states(config::Config)
+    state_ncnames = ncnames(config.state)
+
+    # Get required states
+    required_states = extract_required_states(config)
+
+    # Flag to keep track of whether to throw an error
+    error = false
+    missing_states = ()
+    # Check if all states are covered
+    for state in required_states
+        if !haskey(state_ncnames, state)
+            @error string(
+                "State variable `$state` not provided, please ensure that it is set ",
+                "correctly in the model configuration",
+            )
+            error = true
+            missing_states = (missing_states..., state)
+        end
+    end
+    # Throw error when not all states are covered
+    if error
+        throw(
+            ArgumentError(
+                "Not all required states are provided, these states are missing: `$missing_states`",
+            ),
+        )
+    end
+
+    # Check if more states are passed than required, and warn user if this is the case
+    for (state, _) in state_ncnames
+        if !(state in required_states)
+            @warn "State variable `$state` provided, but is not used in model setup, skipping."
+            delete!(state_ncnames, state)
+        end
+    end
+    return state_ncnames
+end
diff --git a/src/utils.jl b/src/utils.jl
index c2ee0b607..b229be67d 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -136,8 +136,12 @@ and set states in `model` object. Active cells are selected with the correspondi
 # Arguments
 - `type = nothing`: type to convert data to after reading. By default no conversion is done.
 """
-function set_states(instate_path, model, state_ncnames; type = nothing, dimname = nothing)
-    @unpack network = model
+function set_states(instate_path, model; type = nothing, dimname = nothing)
+    @unpack network, config = model
+
+    # Check if required states are covered
+    state_ncnames = check_states(config)
+
     # states in NetCDF include dim time (one value) at index 3 or 4, 3 or 4 dims are allowed
     NCDataset(instate_path) do ds
         for (state, ncname) in state_ncnames
diff --git a/src/vertical_process.jl b/src/vertical_process.jl
index f3fe22c1d..971cc6221 100644
--- a/src/vertical_process.jl
+++ b/src/vertical_process.jl
@@ -154,7 +154,7 @@ function acttransp_unsat_sbm(
     θₛ,
     θᵣ,
     hb,
-    ust::Bool=false,
+    ust::Bool = false,
 )
 
     # AvailCap is fraction of unsat zone containing roots
@@ -363,9 +363,9 @@ function snowpack_hbv(
     ttm,
     cfmax,
     whc;
-    rfcf=1.0,
-    sfcf=1.0,
-    cfr=0.05
+    rfcf = 1.0,
+    sfcf = 1.0,
+    cfr = 0.05,
 )
 
     # fraction of precipitation which falls as rain


Test patch:
diff --git a/test/io.jl b/test/io.jl
index f9a408e34..27fd54054 100644
--- a/test/io.jl
+++ b/test/io.jl
@@ -33,7 +33,7 @@ config = Wflow.Config(tomlpath)
     # modifiers can also be applied
     kvconf = Wflow.get_alias(config.input.vertical, "kv_0", "kv₀", nothing)
     @test kvconf isa Wflow.Config
-    ncname, modifier = Wflow.ncvar_name_modifier(kvconf, config=config)
+    ncname, modifier = Wflow.ncvar_name_modifier(kvconf, config = config)
     @test ncname === "KsatVer"
     @test modifier.scale == 1.0
     @test modifier.offset == 0.0
@@ -63,7 +63,7 @@ end
     # mock a NCReader object
     ncpath = Wflow.input_path(config, config.input.path_forcing)
     ds = NCDataset(ncpath)
-    reader = (; dataset=ds)
+    reader = (; dataset = ds)
 
     # if these keys are missing, they are derived from the NetCDF
     pop!(Dict(config), "starttime")
@@ -292,7 +292,7 @@ Wflow.load_dynamic_input!(model)
     ]
 end
 
-Wflow.close_files(model, delete_output=false)
+Wflow.close_files(model, delete_output = false)
 
 @testset "NetCDF creation" begin
     path = Base.Filesystem.tempname()
@@ -321,17 +321,17 @@ end
         @test Wflow.internal_dim_name(:latitude) == :y
         @test Wflow.internal_dim_name(:time) == :time
 
-        @test_throws ArgumentError Wflow.read_dims(ds["c"], (x=:, y=:))
-        @test_throws ArgumentError Wflow.read_dims(ds["LAI"], (x=:, y=:))
-        data, data_dim_order = Wflow.read_dims(ds["wflow_dem"], (x=:, y=:))
+        @test_throws ArgumentError Wflow.read_dims(ds["c"], (x = :, y = :))
+        @test_throws ArgumentError Wflow.read_dims(ds["LAI"], (x = :, y = :))
+        data, data_dim_order = Wflow.read_dims(ds["wflow_dem"], (x = :, y = :))
         @test data isa Matrix{Union{Float32,Missing}}
         @test data[end, end] === missing
         @test data[125, 1] ≈ 647.187f0
         @test data_dim_order == (:x, :y)
 
-        @test Wflow.dim_directions(ds, (:x, :y)) === (x=true, y=false)
+        @test Wflow.dim_directions(ds, (:x, :y)) === (x = true, y = false)
         @test Wflow.dim_directions(ds, (:y, :x, :layer)) ===
-              (y=false, x=true, layer=true)
+              (y = false, x = true, layer = true)
 
         data, dims = Wflow.permute_data(zeros(1, 2, 3), (:layer, :y, :x))
         @test size(data) == (3, 2, 1)
@@ -343,17 +343,17 @@ end
 
         data = collect(reshape(1:6, (2, 3)))
         # flip y, which is the second dimension
-        @test Wflow.reverse_data!(data, (y=false, x=true))[1, :] == [5, 3, 1]
+        @test Wflow.reverse_data!(data, (y = false, x = true))[1, :] == [5, 3, 1]
         # and mutate it back, the NamedTuple order should not matter
-        @test Wflow.reverse_data!(data, (x=true, y=false))[1, :] == [1, 3, 5]
+        @test Wflow.reverse_data!(data, (x = true, y = false))[1, :] == [1, 3, 5]
         # flip both dimensions at the same time
-        data = Wflow.reverse_data!(data, (x=false, y=false))
+        data = Wflow.reverse_data!(data, (x = false, y = false))
         @test data[1, :] == [6, 4, 2]
         @test data[:, 1] == [6, 5]
 
-        data = Wflow.read_standardized(ds, "wflow_dem", (x=:, y=:))
+        data = Wflow.read_standardized(ds, "wflow_dem", (x = :, y = :))
         # since in this case only the second dimension needs reversing, we can easily do it manually
-        manual_fix = reverse(ds["wflow_dem"]; dims=2)
+        manual_fix = reverse(ds["wflow_dem"]; dims = 2)
         @test all(data .=== manual_fix)
     end
 end
@@ -364,7 +364,7 @@ end
     @test Wflow.parse_loglevel(0) == Logging.Info
 
     tomlpath = joinpath(@__DIR__, "sbm_simple.toml")
-    Wflow.run(tomlpath; silent=true)
+    Wflow.run(tomlpath; silent = true)
 
     config = Wflow.Config(tomlpath)
     output = normpath(abspath(Wflow.get(config, "dir_output", ".")))
@@ -441,3 +441,33 @@ end
     Wflow.advance!(clock)
     @test clock.time == DateTimeNoLeap(2000, 3, 1)
 end
+
+@testset "State checking" begin
+    tomlpath = joinpath(@__DIR__, "sbm_config.toml")
+    config = Wflow.Config(tomlpath)
+    config = Wflow.Config(tomlpath)
+
+    # Extracting required states and test if some are covered (not all are tested!)
+    required_states = Wflow.extract_required_states(config)
+    @test (:vertical, :satwaterdepth) in required_states
+    @test (:vertical, :ustorelayerdepth) in required_states
+    @test (:vertical, :canopystorage) in required_states
+    @test (:lateral, :subsurface, :ssf) in required_states
+    @test (:lateral, :river, :q) in required_states
+    @test !((:lateral, :river, :lake, :waterlevel) in required_states)
+
+    # Adding an unused state the see if the right warning message is thrown
+    config.state.vertical.additional_state = "additional_state"
+    @test_logs (
+        :warn,
+        string(
+            "State variable `(:vertical, :additional_state)` provided, but is not used in ",
+            "model setup, skipping.",
+        ),
+    ) Wflow.check_states(config)
+
+    # Removing the unused and required state, to test the exception being thrown
+    delete!(config.state["vertical"], "additional_state")
+    delete!(config.state["vertical"], "snow")
+    @test_throws ArgumentError Wflow.check_states(config)
+end
diff --git a/test/reservoir_lake.jl b/test/reservoir_lake.jl
index 6e17dfe77..9570ccb63 100644
--- a/test/reservoir_lake.jl
+++ b/test/reservoir_lake.jl
@@ -80,7 +80,7 @@ sh = [
             [2, 1],
             [472461536.0, 60851088.0],
             sh,
-            [missing, Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))]
+            [missing, Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))],
         ),
         threshold = [393.7, 0.0],
         storfunc = [2, 2],
@@ -132,7 +132,7 @@ end
             [1],
             [200_000_000],
             [Wflow.read_sh_csv(joinpath(datadir, "input", "lake_sh_2.csv"))],
-            [Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))]
+            [Wflow.read_hq_csv(joinpath(datadir, "input", "lake_hq_2.csv"))],
         ),
         threshold = [0.0],
         storfunc = [2],
diff --git a/test/run.jl b/test/run.jl
index 7da9b006c..084b11b93 100644
--- a/test/run.jl
+++ b/test/run.jl
@@ -100,4 +100,4 @@ varnames = setdiff(keys(endstate_restart), keys(endstate_restart.dim))
         b = endstate_restart[varname][:]
         @test all(skipmissing(a) .== skipmissing(b))
     end
-end
\ No newline at end of file
+end
diff --git a/test/run_sbm.jl b/test/run_sbm.jl
index f9a9c2976..5f22f35a0 100644
--- a/test/run_sbm.jl
+++ b/test/run_sbm.jl
@@ -282,7 +282,7 @@ model = Wflow.run_timestep(model)
 end
 Wflow.close_files(model, delete_output = false)
 
-# test local-inertial option for river flow including 1D floodplain schematization 
+# test local-inertial option for river flow including 1D floodplain schematization
 tomlpath = joinpath(@__DIR__, "sbm_config.toml")
 config = Wflow.Config(tomlpath)
 
@@ -290,6 +290,8 @@ config.model.floodplain_1d = true
 config.model.river_routing = "local-inertial"
 config.model.land_routing = "kinematic-wave"
 Dict(config.input.lateral.river)["floodplain"] = Dict("volume" => "floodplain_volume")
+Dict(config.state.lateral.river)["floodplain"] =
+    Dict("q" => "q_floodplain", "h" => "h_floodplain")
 
 model = Wflow.initialize_sbm_model(config)