diff --git a/.github/workflows/core_tests.yml b/.github/workflows/core_tests.yml
index a5020269f..50015322d 100644
--- a/.github/workflows/core_tests.yml
+++ b/.github/workflows/core_tests.yml
@@ -1,7 +1,7 @@
 name: Julia Tests
 on:
   push:
-    branches: [main]
+    branches: [main, update/pixi-lock]
     paths-ignore: [".teamcity/**"]
     tags: ["*"]
   pull_request:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index ef6eea4c1..4497abdda 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -1,7 +1,7 @@
 name: Docs
 on:
   push:
-    branches: [main]
+    branches: [main, update/pixi-lock]
     paths-ignore: [".teamcity/**"]
   pull_request:
   merge_group:
diff --git a/.github/workflows/pixi_auto_update.yml b/.github/workflows/pixi_auto_update.yml
new file mode 100644
index 000000000..96e7c513a
--- /dev/null
+++ b/.github/workflows/pixi_auto_update.yml
@@ -0,0 +1,32 @@
+name: Pixi auto update
+
+on:
+  schedule:
+  # At 03:00 on day 1 of the month
+    - cron: "0 3 1 * *"
+  # on demand
+  workflow_dispatch:
+
+jobs:
+  auto-update:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ssh-key: ${{ secrets.SSH_PRIVATE_KEY }}
+      - uses: prefix-dev/setup-pixi@v0.4.3
+        with:
+          pixi-version: "latest"
+          cache: false
+      - name: Update pixi lock file
+        run: |
+          rm pixi.lock
+          pixi install
+      - uses: peter-evans/create-pull-request@v5
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          branch: update/pixi-lock
+          title: Update pixi lock file
+          commit-message: "Update `pixi.lock`"
+          body: Update pixi dependencies to the latest version.
+          author: "GitHub <noreply@github.com>"
diff --git a/.github/workflows/pre-commit_check.yml b/.github/workflows/pre-commit_check.yml
index 17d1ae7a2..fb5f88949 100644
--- a/.github/workflows/pre-commit_check.yml
+++ b/.github/workflows/pre-commit_check.yml
@@ -4,7 +4,7 @@ on:
   pull_request:
   merge_group:
   push:
-    branches: [main, update/pre-commit-hooks]
+    branches: [main, update/pre-commit-hooks, update/pixi-lock]
 
 jobs:
   check:
diff --git a/.github/workflows/python_lint.yml b/.github/workflows/python_lint.yml
index 47310710d..a5a08c9a5 100644
--- a/.github/workflows/python_lint.yml
+++ b/.github/workflows/python_lint.yml
@@ -1,7 +1,7 @@
 name: Python Lint
 on:
   push:
-    branches: [main]
+    branches: [main, update/pixi-lock]
     paths-ignore: [".teamcity/**"]
     tags: ["*"]
   pull_request:
diff --git a/.github/workflows/python_tests.yml b/.github/workflows/python_tests.yml
index 07022ddde..6089785e8 100644
--- a/.github/workflows/python_tests.yml
+++ b/.github/workflows/python_tests.yml
@@ -1,7 +1,7 @@
 name: Ribasim Python Tests
 on:
   push:
-    branches: [main]
+    branches: [main, update/pixi-lock]
     paths-ignore: [".teamcity/**"]
     tags: ["*"]
   pull_request:
diff --git a/.github/workflows/qgis.yml b/.github/workflows/qgis.yml
index 67f159a35..608820001 100644
--- a/.github/workflows/qgis.yml
+++ b/.github/workflows/qgis.yml
@@ -2,7 +2,7 @@ name: QGIS Tests
 
 on:
     push:
-      branches: [main]
+      branches: [main, update/pixi-lock]
       paths-ignore: [".teamcity/**"]
       tags: ["*"]
     pull_request:
diff --git a/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_BuildRibasimCli.xml b/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_BuildRibasimCli.xml
index 2e30ad363..f203da184 100644
--- a/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_BuildRibasimCli.xml
+++ b/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_BuildRibasimCli.xml
@@ -90,12 +90,6 @@ pixi run build-ribasim-cli]]></param>
           <param name="secure:github_access_token" value="credentialsJSON:71420214-373c-4ccd-ba32-2ea886843f62" />
         </parameters>
       </extension>
-      <extension id="PROVIDE_AWS_CREDS" type="PROVIDE_AWS_CREDS">
-        <parameters>
-          <param name="awsConnectionId" value="awsConnection_1" />
-          <param name="awsSessionDuration" value="60" />
-        </parameters>
-      </extension>
     </build-extensions>
     <cleanup />
   </settings>
diff --git a/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_TestRibasimApi.xml b/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_TestRibasimApi.xml
index 7849428d2..8ba314363 100644
--- a/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_TestRibasimApi.xml
+++ b/.teamcity/Ribasim_Linux/buildTypes/Ribasim_Linux_TestRibasimApi.xml
@@ -53,7 +53,7 @@ pixi run test-ribasim-api]]></param>
         <parameters>
           <param name="afterSuccessfulBuildOnly" value="true" />
           <param name="branchFilter" value="+:&lt;default&gt;" />
-          <param name="dependsOn" value="Ribasim_RibasimLinux_BuildLibribasimWindows" />
+          <param name="dependsOn" value="Ribasim_Linux_BuildLibribasim" />
         </parameters>
       </build-trigger>
     </build-triggers>
diff --git a/Manifest.toml b/Manifest.toml
index 43112888e..326f174c5 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "212d113fb58ab2e43d8fe0283d717923e72a9a88"
+project_hash = "77dde6fd5d6ca1581c6694925c34d2bb0d69830e"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
diff --git a/Project.toml b/Project.toml
index ea6fed355..2bd113459 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,6 +5,7 @@ description = "Meta-project used to share the Manifest of Ribasim and its depend
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
+Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
@@ -16,6 +17,7 @@ Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Ribasim = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b"
 TimeZones = "f269a46b-ccf7-5d73-abea-4c690281aa53"
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index f41b3dd28..f8b939fab 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -73,10 +73,8 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
     end
     @debug "Read database into memory."
 
-    storage, errors = get_storages_from_levels(parameters.basin, state.node_id, state.level)
-    if errors
-        error("Encountered errors while parsing the initial levels of basins.")
-    end
+    storage = get_storages_from_levels(parameters.basin, state.level)
+
     # Synchronize level with storage
     set_current_basin_properties!(parameters.basin, storage)
 
diff --git a/core/src/utils.jl b/core/src/utils.jl
index 8898fc141..9da01ce32 100644
--- a/core/src/utils.jl
+++ b/core/src/utils.jl
@@ -381,33 +381,28 @@ function get_storage_from_level(basin::Basin, state_idx::Int, level::Float64)::F
 end
 
 """Compute the storages of the basins based on the water level of the basins."""
-function get_storages_from_levels(
-    basin::Basin,
-    state_node_id::Vector{Int},
-    state_level::Vector{Float64},
-)::Tuple{Vector{Float64}, Bool}
-    (; node_id) = basin
-
-    storages = fill(1.0, length(node_id))
-    n_specified_states = length(state_node_id)
-
-    if n_specified_states > 0
-        basin_state_index = 1
-        basin_state_node_id = state_node_id[1]
-
-        for (i, id) in enumerate(node_id)
-            if basin_state_node_id == id.value
-                storages[i] =
-                    get_storage_from_level(basin, i, state_level[basin_state_index])
-                basin_state_index += 1
-                if basin_state_index > n_specified_states
-                    break
-                end
-                basin_state_node_id = state_node_id[basin_state_index]
-            end
+function get_storages_from_levels(basin::Basin, levels::Vector)::Vector{Float64}
+    errors = false
+    state_length = length(levels)
+    basin_length = length(basin.level)
+    if state_length != basin_length
+        @error "Unexpected 'Basin / state' length." state_length basin_length
+        errors = true
+    end
+    storages = zeros(state_length)
+
+    for (i, level) in enumerate(levels)
+        storage = get_storage_from_level(basin, i, level)
+        if isnan(storage)
+            errors = true
         end
+        storages[i] = storage
+    end
+    if errors
+        error("Encountered errors while parsing the initial levels of basins.")
     end
-    return storages, any(isnan.(storages))
+
+    return storages
 end
 
 """
diff --git a/core/src/validation.jl b/core/src/validation.jl
index d49dc3506..922b02fe2 100644
--- a/core/src/validation.jl
+++ b/core/src/validation.jl
@@ -338,30 +338,6 @@ function variable_nt(s::Any)
     NamedTuple{names}((getfield(s, x) for x in names))
 end
 
-function is_consistent(node, edge, state, static, profile, time)
-
-    # Check that node ids exist
-    # TODO Do we need to check the reverse as well? All ids in use?
-    ids = node.fid
-    @assert edge.from_node_id ⊆ ids "Edge from_node_id not in node ids"
-    @assert edge.to_node_id ⊆ ids "Edge to_node_id not in node ids"
-    @assert state.node_id ⊆ ids "State id not in node ids"
-    @assert static.node_id ⊆ ids "Static id not in node ids"
-    @assert profile.node_id ⊆ ids "Profile id not in node ids"
-    @assert time.node_id ⊆ ids "Time id not in node ids"
-
-    # Check edges for uniqueness
-    @assert allunique(edge, [:from_node_id, :to_node_id]) "Duplicate edge found"
-
-    # TODO Check states
-
-    # TODO Check statics
-
-    # TODO Check forcings
-
-    true
-end
-
 # functions used by sort(x; by)
 sort_by_fid(row) = row.fid
 sort_by_id(row) = row.node_id
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 847af18e9..f0ac72763 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -179,6 +179,7 @@ end
     db_path = Ribasim.input_path(cfg, cfg.database)
     db = SQLite.DB(db_path)
     p = Ribasim.Parameters(db, cfg)
+    close(db)
     (; allocation, graph) = p
     (; main_network_connections) = allocation
 
diff --git a/core/test/bmi_test.jl b/core/test/bmi_test.jl
index 8c6c95dc7..39599eeac 100644
--- a/core/test/bmi_test.jl
+++ b/core/test/bmi_test.jl
@@ -49,7 +49,7 @@ end
     toml_path = normpath(@__DIR__, "../../generated_testmodels/basic/ribasim.toml")
     model = BMI.initialize(Ribasim.Model, toml_path)
     storage0 = BMI.get_value_ptr(model, "volume")
-    @test storage0 == ones(4)
+    @test storage0 ≈ ones(4)
     @test_throws "Unknown variable foo" BMI.get_value_ptr(model, "foo")
     BMI.update_until(model, 86400.0)
     storage = BMI.get_value_ptr(model, "volume")
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 995b07032..227ec010f 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -73,7 +73,7 @@
         @test flow.from_node_id[1:4] == [6, typemax(Int), 0, 6]
         @test flow.to_node_id[1:4] == [6, typemax(Int), typemax(Int), 0]
 
-        @test basin.storage[1] == 1.0
+        @test basin.storage[1] ≈ 1.0
         @test basin.level[1] ≈ 0.044711584
 
         # The exporter interpolates 1:1 for three subgrid elements, but shifted by 1.0 meter.
diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl
index b9076459f..249eb05d0 100644
--- a/core/test/utils_test.jl
+++ b/core/test/utils_test.jl
@@ -125,16 +125,14 @@ end
 
     logger = TestLogger()
     with_logger(logger) do
-        storages, errors = Ribasim.get_storages_from_levels(basin, [1], [-1.0])
-        @test isnan(storages[1])
-        @test errors
-        # Storage for basin with unspecified level is set to 1.0
-        @test storages[2] == 1.0
+        @test_throws ErrorException Ribasim.get_storages_from_levels(basin, [-1.0])
     end
 
-    @test length(logger.logs) == 1
+    @test length(logger.logs) == 2
     @test logger.logs[1].level == Error
-    @test logger.logs[1].message ==
+    @test logger.logs[1].message == "Unexpected 'Basin / state' length."
+    @test logger.logs[2].level == Error
+    @test logger.logs[2].message ==
           "The level -1.0 of basin #1 is lower than the bottom of this basin 0.0."
 
     # Converting from storages to levels and back should return the same storages
diff --git a/python/ribasim/ribasim/config.py b/python/ribasim/ribasim/config.py
index e992e811b..167f8fb64 100644
--- a/python/ribasim/ribasim/config.py
+++ b/python/ribasim/ribasim/config.py
@@ -79,154 +79,132 @@ class Logging(ChildModel):
 
 class Terminal(NodeModel):
     static: TableModel[TerminalStaticSchema] = Field(
-        default_factory=TableModel[TerminalStaticSchema]
+        default_factory=TableModel[TerminalStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id"]}
-
 
 class PidControl(NodeModel):
     static: TableModel[PidControlStaticSchema] = Field(
-        default_factory=TableModel[PidControlStaticSchema]
+        default_factory=TableModel[PidControlStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
     time: TableModel[PidControlTimeSchema] = Field(
-        default_factory=TableModel[PidControlTimeSchema]
+        default_factory=TableModel[PidControlTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id", "control_state"],
-        "time": ["node_id", "time"],
-    }
-
 
 class LevelBoundary(NodeModel):
     static: TableModel[LevelBoundaryStaticSchema] = Field(
-        default_factory=TableModel[LevelBoundaryStaticSchema]
+        default_factory=TableModel[LevelBoundaryStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
     time: TableModel[LevelBoundaryTimeSchema] = Field(
-        default_factory=TableModel[LevelBoundaryTimeSchema]
+        default_factory=TableModel[LevelBoundaryTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id"],
-        "time": ["node_id", "time"],
-    }
-
 
 class Pump(NodeModel):
     static: TableModel[PumpStaticSchema] = Field(
-        default_factory=TableModel[PumpStaticSchema]
+        default_factory=TableModel[PumpStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id", "control_state"]}
-
 
 class TabulatedRatingCurve(NodeModel):
     static: TableModel[TabulatedRatingCurveStaticSchema] = Field(
-        default_factory=TableModel[TabulatedRatingCurveStaticSchema]
+        default_factory=TableModel[TabulatedRatingCurveStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state", "level"]},
     )
     time: TableModel[TabulatedRatingCurveTimeSchema] = Field(
-        default_factory=TableModel[TabulatedRatingCurveTimeSchema]
+        default_factory=TableModel[TabulatedRatingCurveTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "time", "level"]},
     )
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id", "control_state", "level"],
-        "time": ["node_id", "time", "level"],
-    }
 
 
 class User(NodeModel):
     static: TableModel[UserStaticSchema] = Field(
-        default_factory=TableModel[UserStaticSchema]
+        default_factory=TableModel[UserStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "priority"]},
+    )
+    time: TableModel[UserTimeSchema] = Field(
+        default_factory=TableModel[UserTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "priority", "time"]},
     )
-    time: TableModel[UserTimeSchema] = Field(default_factory=TableModel[UserTimeSchema])
-
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id", "priority"],
-        "time": ["node_id", "priority", "time"],
-    }
 
 
 class FlowBoundary(NodeModel):
     static: TableModel[FlowBoundaryStaticSchema] = Field(
-        default_factory=TableModel[FlowBoundaryStaticSchema]
+        default_factory=TableModel[FlowBoundaryStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
     time: TableModel[FlowBoundaryTimeSchema] = Field(
-        default_factory=TableModel[FlowBoundaryTimeSchema]
+        default_factory=TableModel[FlowBoundaryTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id"],
-        "time": ["node_id", "time"],
-    }
-
 
 class Basin(NodeModel):
     profile: TableModel[BasinProfileSchema] = Field(
-        default_factory=TableModel[BasinProfileSchema]
+        default_factory=TableModel[BasinProfileSchema],
+        json_schema_extra={"sort_keys": ["node_id", "level"]},
     )
     state: TableModel[BasinStateSchema] = Field(
-        default_factory=TableModel[BasinStateSchema]
+        default_factory=TableModel[BasinStateSchema],
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
     static: TableModel[BasinStaticSchema] = Field(
-        default_factory=TableModel[BasinStaticSchema]
+        default_factory=TableModel[BasinStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id"]},
     )
     time: TableModel[BasinTimeSchema] = Field(
-        default_factory=TableModel[BasinTimeSchema]
+        default_factory=TableModel[BasinTimeSchema],
+        json_schema_extra={"sort_keys": ["node_id", "time"]},
     )
     subgrid: TableModel[BasinSubgridSchema] = Field(
-        default_factory=TableModel[BasinSubgridSchema]
+        default_factory=TableModel[BasinSubgridSchema],
+        json_schema_extra={"sort_keys": ["subgrid_id", "basin_level"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {
-        "static": ["node_id"],
-        "state": ["node_id"],
-        "profile": ["node_id", "level"],
-        "time": ["node_id", "time"],
-        "subgrid": ["subgrid_id", "basin_level"],
-    }
-
 
 class ManningResistance(NodeModel):
     static: TableModel[ManningResistanceStaticSchema] = Field(
-        default_factory=TableModel[ManningResistanceStaticSchema]
+        default_factory=TableModel[ManningResistanceStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id", "control_state"]}
-
 
 class DiscreteControl(NodeModel):
     condition: TableModel[DiscreteControlConditionSchema] = Field(
-        default_factory=TableModel[DiscreteControlConditionSchema]
+        default_factory=TableModel[DiscreteControlConditionSchema],
+        json_schema_extra={
+            "sort_keys": ["node_id", "listen_feature_id", "variable", "greater_than"]
+        },
     )
     logic: TableModel[DiscreteControlLogicSchema] = Field(
-        default_factory=TableModel[DiscreteControlLogicSchema]
+        default_factory=TableModel[DiscreteControlLogicSchema],
+        json_schema_extra={"sort_keys": ["node_id", "truth_state"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {
-        "condition": ["node_id", "listen_feature_id", "variable", "greater_than"],
-        "logic": ["node_id", "truth_state"],
-    }
-
 
 class Outlet(NodeModel):
     static: TableModel[OutletStaticSchema] = Field(
-        default_factory=TableModel[OutletStaticSchema]
+        default_factory=TableModel[OutletStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id", "control_state"]}
-
 
 class LinearResistance(NodeModel):
     static: TableModel[LinearResistanceStaticSchema] = Field(
-        default_factory=TableModel[LinearResistanceStaticSchema]
+        default_factory=TableModel[LinearResistanceStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
 
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id", "control_state"]}
-
 
 class FractionalFlow(NodeModel):
     static: TableModel[FractionalFlowStaticSchema] = Field(
-        default_factory=TableModel[FractionalFlowStaticSchema]
+        default_factory=TableModel[FractionalFlowStaticSchema],
+        json_schema_extra={"sort_keys": ["node_id", "control_state"]},
     )
-
-    _sort_keys: dict[str, list[str]] = {"static": ["node_id", "control_state"]}
diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
index 2e1a9e067..7db32c1f8 100644
--- a/python/ribasim/ribasim/input_base.py
+++ b/python/ribasim/ribasim/input_base.py
@@ -10,6 +10,7 @@
     Any,
     Generic,
     TypeVar,
+    cast,
 )
 
 import geopandas as gpd
@@ -22,6 +23,8 @@
     ConfigDict,
     DirectoryPath,
     Field,
+    PrivateAttr,
+    ValidationInfo,
     field_validator,
     model_serializer,
     model_validator,
@@ -160,6 +163,7 @@ def _load(cls, filepath: Path | None) -> dict[str, Any]:
 
 class TableModel(FileModel, Generic[TableT]):
     df: DataFrame[TableT] | None = Field(default=None, exclude=True, repr=False)
+    _sort_keys: list[str] = PrivateAttr(default=[])
 
     def __eq__(self, other) -> bool:
         if not type(self) == type(other):
@@ -309,15 +313,14 @@ def _save(
         self,
         directory: DirectoryPath,
         input_dir: DirectoryPath,
-        sort_keys: list[str] = ["node_id"],
     ) -> None:
         # TODO directory could be used to save an arrow file
         db_path = context_file_loading.get().get("database")
         if self.df is not None and self.filepath is not None:
-            self.sort(sort_keys)
+            self.sort()
             self._write_arrow(self.filepath, directory, input_dir)
         elif self.df is not None and db_path is not None:
-            self.sort(sort_keys)
+            self.sort()
             self._write_table(db_path)
 
     def _write_table(self, temp_path: Path) -> None:
@@ -379,13 +382,13 @@ def _from_arrow(cls, path: FilePath) -> pd.DataFrame:
         directory = context_file_loading.get().get("directory", Path("."))
         return pd.read_feather(directory / path)
 
-    def sort(self, sort_keys: list[str]):
+    def sort(self):
         """Sort the table as required.
 
         Sorting is done automatically before writing the table.
         """
         if self.df is not None:
-            self.df.sort_values(sort_keys, ignore_index=True, inplace=True)
+            self.df.sort_values(self._sort_keys, ignore_index=True, inplace=True)
 
     @classmethod
     def tableschema(cls) -> TableT:
@@ -464,7 +467,7 @@ def _write_table(self, path: FilePath) -> None:
 
         gdf.to_file(path, layer=self.tablename(), driver="GPKG")
 
-    def sort(self, sort_keys: list[str]):
+    def sort(self):
         self.df.sort_index(inplace=True)
 
 
@@ -482,8 +485,6 @@ def check_parent(self) -> "ChildModel":
 class NodeModel(ChildModel):
     """Base class to handle combining the tables for a single node type."""
 
-    _sort_keys: dict[str, list[str]] = {}
-
     @model_serializer(mode="wrap")
     def set_modeld(
         self, serializer: Callable[[type["NodeModel"]], dict[str, Any]]
@@ -491,6 +492,19 @@ def set_modeld(
         content = serializer(self)
         return dict(filter(lambda x: x[1], content.items()))
 
+    @field_validator("*")
+    @classmethod
+    def set_sort_keys(cls, v: Any, info: ValidationInfo) -> Any:
+        """Set sort keys for all TableModels if present in FieldInfo."""
+        if isinstance(v, (TableModel,)):
+            field = cls.model_fields[getattr(info, "field_name")]
+            extra = field.json_schema_extra
+            if extra is not None and isinstance(extra, dict):
+                # We set sort_keys ourselves as list[str] in json_schema_extra
+                # but mypy doesn't know.
+                v._sort_keys = cast(list[str], extra.get("sort_keys", []))
+        return v
+
     @classmethod
     def get_input_type(cls):
         return cls.__name__
@@ -579,7 +593,6 @@ def _save(self, directory: DirectoryPath, input_dir: DirectoryPath, **kwargs):
             getattr(self, field)._save(
                 directory,
                 input_dir,
-                sort_keys=self._sort_keys[field],
             )
 
     def _repr_content(self) -> str:
diff --git a/python/ribasim/tests/test_io.py b/python/ribasim/tests/test_io.py
index b1ac52056..e86b9456b 100644
--- a/python/ribasim/tests/test_io.py
+++ b/python/ribasim/tests/test_io.py
@@ -111,9 +111,13 @@ def test_sort(level_setpoint_with_minmax, tmp_path):
     # apply a wrong sort, then call the sort method to restore order
     table.df.sort_values("greater_than", ascending=False, inplace=True)
     assert table.df.iloc[0]["greater_than"] == 15.0
-    sort_keys = model.discrete_control._sort_keys["condition"]
-    assert sort_keys == ["node_id", "listen_feature_id", "variable", "greater_than"]
-    table.sort(sort_keys)
+    assert table._sort_keys == [
+        "node_id",
+        "listen_feature_id",
+        "variable",
+        "greater_than",
+    ]
+    table.sort()
     assert table.df.iloc[0]["greater_than"] == 5.0
 
     # re-apply wrong sort, then check if it gets sorted on write
diff --git a/python/ribasim_testmodels/ribasim_testmodels/backwater.py b/python/ribasim_testmodels/ribasim_testmodels/backwater.py
index e9adc4f07..42e950ec9 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/backwater.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/backwater.py
@@ -52,16 +52,17 @@ def backwater_model():
     )
 
     # Rectangular profile, width of 1.0 m.
+    basin_ids = ids[node_type == "Basin"]
     profile = pd.DataFrame(
         data={
-            "node_id": np.repeat(ids[node_type == "Basin"], 2),
+            "node_id": np.repeat(basin_ids, 2),
             "area": [20.0, 20.0] * n_basin,
             "level": [0.0, 1.0] * n_basin,
         }
     )
     static = pd.DataFrame(
         data={
-            "node_id": ids[node_type == "Basin"],
+            "node_id": basin_ids,
             "drainage": 0.0,
             "potential_evaporation": 0.0,
             "infiltration": 0.0,
@@ -69,7 +70,8 @@ def backwater_model():
             "urban_runoff": 0.0,
         }
     )
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(data={"node_id": basin_ids, "level": 0.05})
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     manning_resistance = ribasim.ManningResistance(
         static=pd.DataFrame(
diff --git a/python/ribasim_testmodels/ribasim_testmodels/basic.py b/python/ribasim_testmodels/ribasim_testmodels/basic.py
index 7a3982173..657314b95 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/basic.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/basic.py
@@ -36,8 +36,10 @@ def basic_model() -> ribasim.Model:
     )
     static = static.iloc[[0, 0, 0, 0]]
     static["node_id"] = [1, 3, 6, 9]
-
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(
+        data={"node_id": static["node_id"], "level": 0.04471158417652035}
+    )
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup linear resistance:
     linear_resistance = ribasim.LinearResistance(
@@ -303,8 +305,11 @@ def tabulated_rating_curve_model() -> ribasim.Model:
             "urban_runoff": 0.0,
         }
     )
+    state = pd.DataFrame(
+        data={"node_id": static["node_id"], "level": 0.04471158417652035}
+    )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Set up a rating curve node:
     # Discharge: lose 1% of storage volume per day at storage = 1000.0.
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 6481214ad..e82b606a0 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -511,7 +511,9 @@ def tabulated_rating_curve_control_model() -> ribasim.Model:
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(data={"node_id": [1], "level": 0.04471158417652035})
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Set up a rating curve node:
     # Discharge: lose 1% of storage volume per day at storage = 100.0.
diff --git a/python/ribasim_testmodels/ribasim_testmodels/invalid.py b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
index 19cee6613..7301f7f2a 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/invalid.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/invalid.py
@@ -61,7 +61,9 @@ def invalid_qh_model():
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(data={"node_id": [3], "level": 1.4112729908597084})
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     rating_curve_static = pd.DataFrame(
         # Invalid: levels must not be repeated
@@ -171,7 +173,14 @@ def invalid_fractional_flow_model():
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(
+        data={
+            "node_id": [1, 2],
+            "level": 1.4112729908597084,
+        }
+    )
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup terminal:
     terminal = ribasim.Terminal(static=pd.DataFrame(data={"node_id": [5, 6]}))
@@ -263,7 +272,14 @@ def invalid_discrete_control_model():
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(
+        data={
+            "node_id": [1, 3],
+            "level": 1.4112729908597084,
+        }
+    )
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup pump:
     pump = ribasim.Pump(
@@ -391,7 +407,14 @@ def invalid_edge_types_model():
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(
+        data={
+            "node_id": [1, 3],
+            "level": 0.04471158417652035,
+        }
+    )
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup pump:
     pump = ribasim.Pump(
diff --git a/python/ribasim_testmodels/ribasim_testmodels/time.py b/python/ribasim_testmodels/ribasim_testmodels/time.py
index bb34dfea7..42a0bdd6a 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/time.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/time.py
@@ -66,7 +66,14 @@ def flow_boundary_time_model():
         }
     )
 
-    basin = ribasim.Basin(profile=profile, static=static)
+    state = pd.DataFrame(
+        data={
+            "node_id": [2],
+            "level": 0.04471158417652035,
+        }
+    )
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     n_times = 100
     time = pd.date_range(
diff --git a/python/ribasim_testmodels/ribasim_testmodels/trivial.py b/python/ribasim_testmodels/ribasim_testmodels/trivial.py
index 664b69197..99b5067cf 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/trivial.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/trivial.py
@@ -74,6 +74,8 @@ def trivial_model() -> ribasim.Model:
         }
     )
 
+    state = pd.DataFrame(data={"node_id": [6], "level": 0.04471158417652035})
+
     # Create a subgrid level interpolation from one basin to three elements. Scale one to one, but:
     #
     # 22. start at -1.0
@@ -88,7 +90,7 @@ def trivial_model() -> ribasim.Model:
             "subgrid_level": [-1.0, 0.0, 0.0, 1.0, 1.0, 2.0],
         }
     )
-    basin = ribasim.Basin(profile=profile, static=static, subgrid=subgrid)
+    basin = ribasim.Basin(profile=profile, static=static, state=state, subgrid=subgrid)
 
     # Set up a rating curve node:
     # Discharge: lose 1% of storage volume per day at storage = 1000.0.