diff --git a/docs/source/_static/images/4bucketmodel.svg b/docs/source/_static/images/4bucketmodel.svg
deleted file mode 100644
index 4d2508d7f..000000000
--- a/docs/source/_static/images/4bucketmodel.svg
+++ /dev/null
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/docs/source/_static/images/bucketmodel.svg b/docs/source/_static/images/bucketmodel.svg
new file mode 100644
index 000000000..896e9afc3
--- /dev/null
+++ b/docs/source/_static/images/bucketmodel.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/docs/source/_static/river_DEM.nc b/docs/source/_static/river_DEM.nc
new file mode 100644
index 000000000..dc0f9f64c
Binary files /dev/null and b/docs/source/_static/river_DEM.nc differ
diff --git a/docs/source/glossary.md b/docs/source/glossary.md
index f99d2baab..e2090a33c 100644
--- a/docs/source/glossary.md
+++ b/docs/source/glossary.md
@@ -110,4 +110,29 @@ gravitational head
reference level, influencing the direction and rate of water movement. It is typically
measured in units of length, such as meters (m), representing the height of a water
column.
+
+local runoff generation
+ The water generated within a grid cell during a timestep. Includes both surface runoff
+ (overland flow) and subsurface runoff (lateral soil flow and baseflow).
+
+surface runoff
+ The portion of local runoff that flows over the land surface into streams and rivers.
+ In this model, this is purely precipitation-derived.
+
+sub-surface runoff
+ The portion of local runoff that moves laterally through soil and groundwater pathways
+ before reaching the stream.
+
+total runoff
+ the depth of water produced from a drainage area during a particular time interval,
+ including surface and subsurface runoff (mm).
+
+inflow
+ Water entering a cell from all upstream cells during the same timestep. Can be divided
+ into surface inflow (overland) and subsurface inflow (through soils/groundwater).
+
+river discharge rate
+ The total volume of water moving through a river channel in a cell per unit time,
+ combining both surface and subsurface flows. The model returns this variable a rate
+ in cubic meters per second.
```
diff --git a/docs/source/virtual_ecosystem/implementation/hydrology_implementation.md b/docs/source/virtual_ecosystem/implementation/hydrology_implementation.md
index a26c10d2a..e108e76df 100644
--- a/docs/source/virtual_ecosystem/implementation/hydrology_implementation.md
+++ b/docs/source/virtual_ecosystem/implementation/hydrology_implementation.md
@@ -93,24 +93,24 @@ surface runoff out of the grid cell.
The [below ground](../../api/models/hydrology/below_ground.md) component considers
infiltration, bypass flow, percolation (= vertical flow), {term}`soil moisture` and
{term}`soil matric potential`, horizontal
-sub-surface flow out of the grid cell, and changes in groundwater storage.
+sub-surface runoff out of the grid cell, and changes in groundwater storage.
-:::{figure} ../../_static/images/4bucketmodel.svg
+:::{figure} ../../_static/images/bucketmodel.svg
:name: bucket_model
:alt: 4-Bucket Model
:class: bg-primary
:width: 600px
-The 4-Bucket Model represented within grid cell hydrology processes in the Virtual
-Ecosystem. Red solid arrows indicate upward vertical flow, red dashed arrows show
-vertical downward flow. The blue arrows indicate horizontal flow out of the
+The 4-Bucket Model represents within grid cell hydrology processes in the Virtual
+Ecosystem. Red solid arrows indicate upward vertical flows, red dashed arrows show
+vertical downward flows. The blue arrows indicate horizontal flows out of the
grid cell with solid lines representing water that flows out of each layer in the
current time step and dashed lines representing water that originates from upstream
-grid cells (previous time step) and flows through the grid cell directly to the stream.
+grid cells and flows through the grid cell directly to the stream.
**NOTE!** Top soil and middle soil are currently treated as one layer in the model.
-Subsurface runoff (Q2) and interflow (Q3) are currently not implemented; the
-river discharge is calculated as the sum of surface runoff (Q1) and the flows out of the
-groundwater buckets (Q4+Q5).
+Subsurface stormflow (Q2) and interflow (Q3) are currently not implemented; the
+total runoff is calculated as the sum of surface runoff (Q1) and the flows out of the
+groundwater buckets (=subsurface runoff, Q4+Q5).
:::
### Canopy interception
@@ -305,10 +305,10 @@ We do currently NOT include any horizontal flows from the soil layers towards th
(Q2 and Q3 in {numref}`bucket_model`).
```
-### Belowground horizontal flow and groundwater storage
+### Belowground runoff and groundwater storage
-Groundwater storage and horizontal transport are modelled using two parallel linear
-reservoirs, similar to the approach used in the HBV-96 model
+Groundwater storage and runoff towards the channel are modelled using two parallel
+linear reservoirs, similar to the approach used in the HBV-96 model
{cite}`lindstrom_development_1997` and the LISFLOOD
{cite}`van_der_knijff_lisflood_2010` (see for full documentation).
@@ -316,7 +316,7 @@ The upper zone represents a quick runoff component, which includes fast groundwa
and (vertical) subsurface flow through macro-pores in the soil. The lower zone
represents the slow groundwater component that generates the base flow.
-The outflow from the upper zone to the channel, $Q_{uz}$, (mm),
+The runoff from the upper zone to the channel, $Q_{uz}$, (mm),
(Q4 in {numref}`bucket_model`) equals:
$$Q_{uz} = \frac{1}{T_{uz}} \cdot UZ \cdot \Delta t$$
@@ -339,8 +339,8 @@ follows:
$$D_{uz,lz} = min(GW_{perc} \cdot \Delta t, UZ)$$
-where $GW_{perc}$, [mm day], is the maximum percolation rate from the upper to
-the lower groundwater zone. The outflow from the lower zone to the channel $Q_{lz}$,
+where $GW_{perc}$, [mm day-1], is the maximum percolation rate from the upper to
+the lower groundwater zone. The runoff from the lower zone to the channel $Q_{lz}$,
(mm), (Q5 in {numref}`bucket_model`) is then computed by:
$$Q_{lz} = \frac{1}{T_{lz}} \cdot LZ \cdot \Delta t$$
@@ -362,8 +362,15 @@ the value of $GW_{loss}$, the larger the amount of water that leaves the system.
## Across grid hydrology
The second part of the hydrology model calculates the horizontal water movement across
-the full model grid including accumulated surface runoff and sub-surface flow, and river
-discharge rate.
+the full model grid including {term}`surface runoff` and {term}`sub-surface runoff`,
+combined into {term}`local runoff generation` per grid cell and {term}`total runoff`
+across the grid.
+
+Hereinafter, we refer to {term}`river discharge rate` to indicate the amount of water
+passing a particular point of a river (m3 s−1), whereas runoff $R$ is regarded as the
+depth of water produced from a drainage area during a particular time interval (mm).
+
+### Drainage map
The flow direction of water above and below ground is based on a digital elevation model
which needs to be provided as a NetCDF file at the start of the simulation.
@@ -371,31 +378,18 @@ Here an description of the steps that happen during the hydrology model
initialisation (plotting only for illustration):
```{code-cell} ipython3
-# # Read elevation datafrom NetCDF
+# Read elevation data from NetCDF
import numpy as np
import xarray as xr
from xarray import DataArray
-input_file = "../../../../virtual_ecosystem/example_data/data/example_elevation_data.nc"
+input_file = "../../_static/river_DEM.nc"
digital_elevation_model = xr.open_dataset(input_file)
elevation = digital_elevation_model["elevation"]
```
```{code-cell} ipython3
-# Plot the elevation data
-import matplotlib.pyplot as plt
-from matplotlib import colors
-
-plt.figure(figsize=(10, 6))
-elevation.plot(cmap="terrain")
-plt.title("Elevation, m")
-plt.xlabel("x")
-plt.ylabel("y")
-plt.show()
-```
-
-```{code-cell} ipython3
-# Create Grid and Data objects and add elevation data (this happens automatically)
+# Create Grid and Data objects and add elevation data
from virtual_ecosystem.core.grid import Grid
from virtual_ecosystem.core.data import Data
@@ -404,6 +398,22 @@ grid = Grid(
)
data = Data(grid=grid)
data["elevation"] = elevation
+
+# Plot elevation data on grid
+import matplotlib.pyplot as plt
+from matplotlib import colors
+
+ele_plot = DataArray(
+ data["elevation"].to_numpy().reshape((9, 9)),
+ dims=("x", "y"),
+ coords={"x": np.arange(9), "y": np.arange(9)},
+)
+plt.figure(figsize=(10, 6))
+ele_plot.plot(cmap="terrain")
+plt.title("Elevation, [m]")
+plt.xlabel("x")
+plt.ylabel("y")
+plt.show()
```
The initialisation step of the hydrology model finds all the neighbours for each grid
@@ -417,8 +427,8 @@ grid.neighbours[56]
Based on that relationship, the model determines all upstream neighbours
for each grid cell and creates a drainage map, i.e. a dictionary that contains for each
-grid cell all upstream grid cells. For example, `cell_id = 56` has four upstream cells
-with the indices `[47, 56, 57, 65]`.
+grid cell all upstream grid cells. For example, `cell_id = 34` has upstream cells
+with the indices `[4, 13, 22]`. This gives the flow direction.
```{code-cell} ipython3
from virtual_ecosystem.models.hydrology.above_ground import calculate_drainage_map
@@ -429,39 +439,97 @@ drainage_map = calculate_drainage_map(
)
```
-The accumulated surface runoff is the calculated in each grid cell as the sum of current
-runoff and the runoff from upstream cells at the previous time step.
+### Runoff and river discharge rate
+
+We track horizontal water fluxes in two pathways — surface and subsurface runoff — and
+then combine them into total runoff, which can be converted to river discharge rate.
+
+#### Surface runoff ($R_{surface}$)
+
+Water moving over the land surface into the river channel. For each cell, this includes:
+
+* Local surface runoff: water generated within the cell during the current timestep.
+* Upstream surface runoff: surface runoff generated in all upstream cells during the
+ same timestep.
+
+#### Subsurface runoff ($R_{subsurface}$)
+
+Water moving laterally through soil and groundwater pathways towards the river channel.
+For each cell, this includes:
+
+* Local subsurface runoff: lateral + baseflow generated in the cell during the current
+ timestep.
+* Upstream subsurface runoff: subsurface runoff generated in upstream cells during the
+ same timestep.
+
+#### Total runoff ($R_{total}$) and river discharge rate
+
+The total water volume passing through a cell is the sum of these two pathways:
+
+$$R_{total} = R_{surface} + R_{subsurface}$$
+
+This total volume can then be converted to a river discharge rate in cubic meters per
+second (m3 s-1) using cell area and unit conversions.
```{code-cell} ipython3
-from virtual_ecosystem.models.hydrology.above_ground import accumulate_horizontal_flow
+from virtual_ecosystem.models.hydrology.above_ground import route_horizontal_flow
-previous_accumulated_runoff = DataArray(np.full(81, 10), dims="cell_id")
-surface_runoff = DataArray(np.full(81, 1), dims="cell_id")
+# Local runoff
+subsurface_runoff = DataArray(np.full_like(data["elevation"], 1.0), dims="cell_id")
+surface_runoff = DataArray(np.full_like(data["elevation"], 12), dims="cell_id")
-accumulated_runoff = accumulate_horizontal_flow(
+# Total runoff = local runoff generation + upstream inflow
+total_runoff = route_horizontal_flow(
drainage_map=drainage_map,
- current_flow=surface_runoff,
- previous_accumulated_flow=previous_accumulated_runoff,
+ surface_runoff=surface_runoff,
+ subsurface_runoff=subsurface_runoff,
+)
+
+# Reshape to 9x9 grid and plot total runoff map
+reshaped_runoff = DataArray(
+ total_runoff.reshape((9, 9)),
+ dims=("x", "y"),
+ coords={"x": np.arange(9), "y": np.arange(9)},
)
-# Plot accumulated runoff map
-reshaped_data = DataArray(accumulated_runoff.to_numpy().reshape(9, 9))
plt.figure(figsize=(10, 6))
-reshaped_data.plot(cmap="Blues")
-plt.title("Accumulated runoff, mm")
+reshaped_runoff.plot(cmap="Blues")
+plt.title("Total runoff, [mm]")
plt.xlabel("x")
plt.ylabel("y")
plt.show()
```
-Total river discharge is calculated as the sum of above- and below ground horizontal
-flow and converted to river discharge rate in m3/s.
+```{code-cell} ipython3
+from virtual_ecosystem.models.hydrology.above_ground import (
+ convert_mm_flow_to_m3_per_second,
+)
-```{note}
-The hydrology model requires a spinup period to establish a steady state flow of
-accumulated above and below ground flow - each simulation time step then represents the
-total flow through a grid cell. This is currently not implemented.
+# Convert total runoff [mm] to river discharge rate [m3 s-1]
+river_discharge_rate = convert_mm_flow_to_m3_per_second(
+ river_discharge_mm=total_runoff,
+ area=grid.cell_area,
+ days=1,
+ seconds_to_day=86400,
+ meters_to_millimeters=1000,
+)
+
+# Reshape to 9x9 grid and plot river discharge rate
+reshaped_discharge_rate = DataArray(
+ river_discharge_rate.reshape((9, 9)),
+ dims=("x", "y"),
+ coords={"x": np.arange(9), "y": np.arange(9)},
+)
+plt.figure(figsize=(10, 6))
+reshaped_discharge_rate.plot(cmap="Blues")
+plt.title("River discharge rate, [m3 s-1]")
+plt.xlabel("x")
+plt.ylabel("y")
+plt.show()
+```
+
+```{note}
To close the water balance, water needs to enter and leave the grid at some point. These
boundaries are currently not implemented.
```
diff --git a/tests/models/hydrology/test_above_ground.py b/tests/models/hydrology/test_above_ground.py
index e52ea2426..dd6cef4b1 100644
--- a/tests/models/hydrology/test_above_ground.py
+++ b/tests/models/hydrology/test_above_ground.py
@@ -170,34 +170,12 @@ def test_find_upstream_cells():
assert result == exp_result
-@pytest.mark.parametrize(
- "acc_runoff,raises,expected_log_entries",
- [
- (
- np.array([100, 100, 100, 100, 100, 100, 100, 100]),
- does_not_raise(),
- {},
- ),
- (
- np.array([-100, 100, 100, 100, 100, 100, 100, 100]),
- pytest.raises(ValueError),
- (
- (
- ERROR,
- "The accumulated flow should not be negative!",
- ),
- ),
- ),
- ],
-)
-def accumulate_horizontal_flow(caplog, acc_runoff, raises, expected_log_entries):
- """Test."""
+def test_route_horizontal_flow_basic():
+ """Test horizontal flow routing."""
- from virtual_ecosystem.models.hydrology.above_ground import (
- accumulate_horizontal_flow,
- )
+ from virtual_ecosystem.models.hydrology.above_ground import route_horizontal_flow
- upstream_ids = {
+ drainage_map = {
0: [],
1: [0],
2: [1, 2],
@@ -205,17 +183,43 @@ def accumulate_horizontal_flow(caplog, acc_runoff, raises, expected_log_entries)
4: [],
5: [3],
6: [],
- 7: [4, 5, 6, 7],
+ 7: [4, 5, 6],
}
- surface_runoff = np.array([100, 100, 100, 100, 100, 100, 100, 100])
- exp_result = np.array([100, 200, 300, 100, 100, 200, 100, 500])
- with raises:
- result = accumulate_horizontal_flow(upstream_ids, surface_runoff, acc_runoff)
- np.testing.assert_array_equal(result, exp_result)
+ surface_runoff = np.array([1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 2.0, 1.0])
+ subsurface_runoff = np.array([0.5, 0.0, 1.0, 0.5, 0.0, 1.0, 0.5, 1.0])
- # Final check that expected logging entries are produced
- log_check(caplog, expected_log_entries)
+ result = route_horizontal_flow(drainage_map, surface_runoff, subsurface_runoff)
+
+ expected_channel_inflow = np.array([1.5, 3.5, 10.0, 1.5, 2.0, 5.5, 2.5, 10.5])
+
+ np.testing.assert_array_almost_equal(result, expected_channel_inflow)
+
+
+def test_route_horizontal_flow_no_upstream():
+ """Test horizontal flow routing with no upstream cells."""
+ from virtual_ecosystem.models.hydrology.above_ground import route_horizontal_flow
+
+ # Single cell with no upstream
+ drainage_map = {0: []}
+ surface_runoff = np.array([5.0])
+ subsurface_runoff = np.array([2.0])
+
+ result = route_horizontal_flow(drainage_map, surface_runoff, subsurface_runoff)
+ expected = np.array([7.0]) # 5 + 2
+ np.testing.assert_array_equal(result, expected)
+
+
+def test_route_horizontal_flow_raises_on_negative():
+ """Test horizontal flow routing raises on negative input."""
+ from virtual_ecosystem.models.hydrology.above_ground import route_horizontal_flow
+
+ drainage_map = {0: [], 1: [0]}
+ surface_runoff = np.array([-1.0, 0.0])
+ subsurface_runoff = np.array([0.0, 0.0])
+
+ with pytest.raises(ValueError, match="The river discharge should not be negative"):
+ route_horizontal_flow(drainage_map, surface_runoff, subsurface_runoff)
@pytest.mark.parametrize(
@@ -279,7 +283,7 @@ def test_calculate_drainage_map(caplog, grid_type, raises, expected_log_entries)
result = calculate_drainage_map(grid, elevation)
assert len(result) == grid.n_cells
- assert result[1] == [2, 6]
+ assert result[1] == [2, 3, 4, 6, 7, 8, 9, 11, 12, 13, 14]
# Final check that expected logging entries are produced
log_check(caplog, expected_log_entries)
diff --git a/tests/models/hydrology/test_hydrology_model.py b/tests/models/hydrology/test_hydrology_model.py
index b1f62e111..0882458b7 100644
--- a/tests/models/hydrology/test_hydrology_model.py
+++ b/tests/models/hydrology/test_hydrology_model.py
@@ -107,7 +107,7 @@ def test_hydrology_model_initialization(
assert model.initial_soil_moisture == ini_soil_moisture
assert model.initial_groundwater_saturation == ini_groundwater_sat
# TODO: not sure on the value below, test with more expansive drainage maps
- assert model.drainage_map == {0: [], 1: [], 2: [0, 2, 3], 3: [1]}
+ assert model.drainage_map == {0: [], 1: [], 2: [0, 1, 2, 3], 3: [1]}
# Final check that expected logging entries are produced
if expected_log_entries:
@@ -243,9 +243,19 @@ def test_generate_hydrology_model(
],
},
{
- "total_river_discharge": [0, 0, 67002, 22095],
+ "total_runoff": [
+ 1477.363339,
+ 1476.076772,
+ 7371.383915,
+ 2945.792003,
+ ],
"surface_runoff": [20.343781, 6.316444, 2.721491, 1.192358],
- "surface_runoff_accumulated": [0, 0, 420, 90],
+ "surface_runoff_routed_plus_local": [
+ 20.343781,
+ 6.316444,
+ 33.295566,
+ 7.508803,
+ ],
"soil_evaporation": [5.93727, 12.359247, 25.50399, 51.620636],
},
id="1 month",
@@ -268,9 +278,19 @@ def test_generate_hydrology_model(
],
},
{
- "total_river_discharge": [0, 0, 5668, 1894],
+ "total_runoff": [
+ 483.741646,
+ 480.053134,
+ 2365.561823,
+ 936.948501,
+ ],
"surface_runoff": [163.019971, 158.780549, 150.030499, 132.191482],
- "surface_runoff_accumulated": [0, 0, 3094, 1085],
+ "surface_runoff_routed_plus_local": [
+ 163.019971,
+ 158.780549,
+ 754.053001,
+ 290.972032,
+ ],
"soil_evaporation": [1.388223, 2.904608, 6.066225, 12.622505],
},
id="1 week",
@@ -288,6 +308,7 @@ def test_setup(
):
"""Test set up and update."""
from virtual_ecosystem.core.core_components import CoreComponents
+ from virtual_ecosystem.models.hydrology import hydrology_tools
from virtual_ecosystem.models.hydrology.hydrology_model import HydrologyModel
# Build the config object and core components
@@ -364,3 +385,16 @@ def test_setup(
rtol=1e-2,
atol=1e-2,
)
+ # Mass balance check for the month
+ hydrology_tools.check_monthly_mass_balance(
+ drainage_map=model.drainage_map,
+ surface_channel_inflow_mm=model.data[
+ "surface_runoff_routed_plus_local"
+ ].to_numpy(),
+ monthly_precipitation_mm=dummy_climate_data_varying_canopy[
+ "precipitation"
+ ]
+ .isel(time_index=1)
+ .to_numpy(),
+ monthly_evaporation_mm=model.data["soil_evaporation"].to_numpy(),
+ )
diff --git a/tests/models/hydrology/test_hydrology_tools.py b/tests/models/hydrology/test_hydrology_tools.py
index b5fdd3d5e..ca2e5e1ee 100644
--- a/tests/models/hydrology/test_hydrology_tools.py
+++ b/tests/models/hydrology/test_hydrology_tools.py
@@ -74,8 +74,6 @@ def test_setup_hydrology_input_current_timestep(
"current_transpiration",
"top_soil_moisture_saturation",
"top_soil_moisture_residual",
- "previous_accumulated_runoff",
- "previous_subsurface_flow_accumulated",
"groundwater_storage",
"current_soil_moisture",
]
@@ -196,3 +194,43 @@ def test_calculate_effective_saturation():
)
assert np.allclose(actual_sats, expected_sats)
+
+
+def test_mass_balance_pass():
+ """Test a case where total outlet streamflow <= total precipitation."""
+ from virtual_ecosystem.models.hydrology.hydrology_tools import (
+ check_monthly_mass_balance,
+ )
+
+ drainage_map = {0: [], 1: [0], 2: [0, 1]}
+ surface_channel_inflow = np.array([50.0, 40.0, 30.0]) # mm per cell
+ monthly_precipitation = np.array([100.0, 100.0, 100.0]) # mm per cell
+ monthly_evaporation = np.array([10.0, 10.0, 10.0]) # mm per cell
+
+ # Should not raise an error
+ check_monthly_mass_balance(
+ drainage_map=drainage_map,
+ surface_channel_inflow_mm=surface_channel_inflow,
+ monthly_precipitation_mm=monthly_precipitation,
+ monthly_evaporation_mm=monthly_evaporation,
+ )
+
+
+def test_mass_balance_fail():
+ """Test a case where total outlet streamflow > total precipitation."""
+ from virtual_ecosystem.models.hydrology.hydrology_tools import (
+ check_monthly_mass_balance,
+ )
+
+ drainage_map = {0: [], 1: [0], 2: [0, 1]}
+ surface_channel_inflow = np.array([150.0, 140.0, 400.0]) # mm per cell
+ monthly_precipitation = np.array([100.0, 100.0, 100.0]) # total precip = 300
+ monthly_evaporation = np.array([10.0, 10.0, 10.0]) # mm per cell
+
+ with pytest.raises(AssertionError, match="Mass balance violated"):
+ check_monthly_mass_balance(
+ drainage_map=drainage_map,
+ surface_channel_inflow_mm=surface_channel_inflow,
+ monthly_precipitation_mm=monthly_precipitation,
+ monthly_evaporation_mm=monthly_evaporation,
+ )
diff --git a/virtual_ecosystem/data_variables.toml b/virtual_ecosystem/data_variables.toml
index 1d81c263a..cc3989452 100644
--- a/virtual_ecosystem/data_variables.toml
+++ b/virtual_ecosystem/data_variables.toml
@@ -413,8 +413,8 @@ variable_type = "float"
[[variable]]
axis = ["spatial"]
-description = "Accumulated subsurface flow"
-name = "subsurface_flow_accumulated"
+description = "Subsurface flow contributing to river discharge in each grid cell"
+name = "subsurface_runoff_routed_plus_local"
unit = "mm"
variable_type = "float"
@@ -490,8 +490,8 @@ variable_type = "float"
[[variable]]
axis = ["spatial"]
-description = "Accumulated surface runoff"
-name = "surface_runoff_accumulated"
+description = "Surface runoff contributing to river discharge in each grid cell"
+name = "surface_runoff_routed_plus_local"
unit = "mm"
variable_type = "float"
@@ -504,8 +504,8 @@ variable_type = "float"
[[variable]]
axis = ["spatial"]
-description = "Total river discharge"
-name = "total_river_discharge"
+description = "Total runoff through each grid cell"
+name = "total_runoff"
unit = "mm"
variable_type = "float"
diff --git a/virtual_ecosystem/models/hydrology/above_ground.py b/virtual_ecosystem/models/hydrology/above_ground.py
index 11274966b..12b6545fa 100644
--- a/virtual_ecosystem/models/hydrology/above_ground.py
+++ b/virtual_ecosystem/models/hydrology/above_ground.py
@@ -380,38 +380,52 @@ def find_upstream_cells(lowest_neighbour: list[int]) -> list[list[int]]:
return upstream_ids
-def accumulate_horizontal_flow(
+def route_horizontal_flow(
drainage_map: dict[int, list[int]],
- current_flow: np.ndarray,
- previous_accumulated_flow: np.ndarray,
+ surface_runoff: np.ndarray,
+ subsurface_runoff: np.ndarray,
) -> np.ndarray:
- """Calculate accumulated above-/belowground horizontal flow for each grid cell.
+ """Route horizontal flow for each grid cell (instantaneous channel routing).
- This function takes the accumulated above-/belowground horizontal flow from the
- previous timestep and adds all (sub-)surface flow of the current time step from
- upstream cell IDs.
+ This function calculates the total river discharge at each grid cell
+ for the current timestep by combining:
- The function currently raises a `ValueError` if accumulated flow is negative.
+ 1. Local generation: the water generated in the cell itself during the timestep,
+ including surface runoff and subsurface (lateral + baseflow) runoff.
+ 2. Inflow from upstream cells: contributions from all cells that drain into the
+ current cell, using their local generation from the same timestep.
+
+ No flows from previous timesteps are included, avoiding double-counting, and there
+ is also no time delay, so all the water runs through the whole grid in one time step
+ .
Args:
- drainage_map: Dict of all upstream IDs for each grid cell
- current_flow: (Sub-)surface flow of the current time step, [mm]
- previous_accumulated_flow: Accumulated flow from previous time step, [mm]
+ drainage_map: Dict mapping each cell ID -> list of upstream cell IDs
+ surface_runoff: Surface runoff for this timestep, [mm]
+ subsurface_runoff: Subsurface runoff for this timestep, [mm]
Returns:
- accumulated (sub-)surface flow, [mm]
+ Total river discharge at each grid cell, [mm]
"""
+ # local generation in this cell (surface + subsurface)
+ local_generation = np.nan_to_num(surface_runoff, nan=0.0) + np.nan_to_num(
+ subsurface_runoff, nan=0.0
+ )
+
+ inflow_from_upstream = np.zeros_like(local_generation)
- current_flow_true = np.nan_to_num(current_flow, nan=0.0)
for cell_id, upstream_ids in enumerate(drainage_map.values()):
- previous_accumulated_flow[cell_id] += np.sum(current_flow_true[upstream_ids])
+ if upstream_ids:
+ inflow_from_upstream[cell_id] = np.sum(local_generation[upstream_ids])
- if (previous_accumulated_flow < 0.0).any():
- to_raise = ValueError("The accumulated flow should not be negative!")
+ total_river_discharge = local_generation + inflow_from_upstream
+
+ if (total_river_discharge < 0.0).any():
+ to_raise = ValueError("The river discharge should not be negative!")
LOGGER.error(to_raise)
raise to_raise
- return previous_accumulated_flow
+ return total_river_discharge
def calculate_drainage_map(grid: Grid, elevation: np.ndarray) -> dict[int, list[int]]:
@@ -436,11 +450,33 @@ def calculate_drainage_map(grid: Grid, elevation: np.ndarray) -> dict[int, list[
LOGGER.error(to_raise)
raise to_raise
+ # Establish neighbour relationships
grid.set_neighbours(distance=sqrt(grid.cell_area))
+
+ # Find flow direction: each cell -> lowest neighbor
lowest_neighbours = find_lowest_neighbour(grid.neighbours, elevation)
- upstream_ids = find_upstream_cells(lowest_neighbours)
+ n_cells = len(lowest_neighbours)
+
+ # Build reverse graph: for each cell, who drains into it
+ direct_upstream: dict[int, list[int]] = {i: [] for i in range(n_cells)}
+ for cell, ln in enumerate(lowest_neighbours):
+ if ln is not None: # sink cells have no lowest neighbor
+ direct_upstream[ln].append(cell)
+
+ # Recursive collection of all upstream cells
+ def collect_upstream(cell: int, visited=None) -> list[int]:
+ if visited is None:
+ visited = set()
+ for up in direct_upstream[cell]:
+ if up not in visited:
+ visited.add(up)
+ collect_upstream(up, visited)
+ return list(visited)
+
+ # Compute upstream IDs for all cells
+ upstream_ids = {cell: collect_upstream(cell) for cell in range(n_cells)}
- return dict(enumerate(upstream_ids))
+ return upstream_ids
def calculate_interception(
diff --git a/virtual_ecosystem/models/hydrology/hydrology_model.py b/virtual_ecosystem/models/hydrology/hydrology_model.py
index dfe746d21..33f38773e 100644
--- a/virtual_ecosystem/models/hydrology/hydrology_model.py
+++ b/virtual_ecosystem/models/hydrology/hydrology_model.py
@@ -8,7 +8,7 @@ class as a child of the :class:`~virtual_ecosystem.core.base_model.BaseModel` cl
.. TODO:: processes
- * spin up soil moisture and accumulated runoff
+ * spin up soil moisture
* set boundaries for river discharge
* update infiltration process
@@ -65,12 +65,12 @@ class HydrologyModel(
"surface_runoff",
"vertical_flow",
"soil_evaporation",
- "surface_runoff_accumulated",
- "subsurface_flow_accumulated",
+ "surface_runoff_routed_plus_local",
+ "subsurface_runoff_routed_plus_local",
+ "total_runoff",
+ "river_discharge_rate",
"matric_potential",
"groundwater_storage",
- "river_discharge_rate",
- "total_river_discharge",
"subsurface_flow",
"baseflow",
"bypass_flow",
@@ -87,8 +87,6 @@ class HydrologyModel(
"layer_heights",
"soil_moisture",
"transpiration",
- "surface_runoff_accumulated",
- "subsurface_flow_accumulated",
"density_air",
"aerodynamic_resistance_canopy",
"specific_heat_air",
@@ -98,8 +96,6 @@ class HydrologyModel(
vars_populated_by_init=(
"soil_moisture",
"groundwater_storage",
- "surface_runoff_accumulated",
- "subsurface_flow_accumulated",
"aerodynamic_resistance_surface",
"aerodynamic_resistance_canopy",
"specific_heat_air",
@@ -116,8 +112,10 @@ class HydrologyModel(
"matric_potential",
"subsurface_flow",
"baseflow",
- "total_river_discharge",
+ "surface_runoff_routed_plus_local",
+ "subsurface_runoff_routed_plus_local",
"river_discharge_rate",
+ "total_runoff",
"canopy_evaporation",
),
):
@@ -161,7 +159,7 @@ def __init__(
self.model_constants: HydroConsts
"""Set of constants for the hydrology model"""
self.drainage_map: dict
- """Upstream neighbours for the calculation of accumulated horizontal flow."""
+ """Upstream neighbours for the calculation of horizontal flow."""
self.soil_layer_thickness_mm: np.ndarray
"""Soil layer thickness in mm."""
self.surface_layer_index: int
@@ -225,10 +223,6 @@ def _setup(
atmospheric variables are initialised to ensure they are available for update
when the Virtual Ecosystem is run with the `abiotic_simple` model.
- For the hydrology across the grid, this function initialises the accumulated
- surface runoff variable and the subsurface accumulated flow variable. Both
- require a spinup which is currently not implemented.
-
Args:
initial_soil_moisture: The initial volumetric relative water content
[unitless] for all layers. This will be converted to soil moisture in
@@ -300,15 +294,6 @@ def _setup(
name="groundwater_storage",
)
- # Set initial above-ground accumulated runoff and sub-surface flow to zero
- for var in ["surface_runoff_accumulated", "subsurface_flow_accumulated"]:
- self.data[var] = DataArray(
- np.zeros_like(self.data["elevation"]),
- dims="cell_id",
- name=var,
- coords={"cell_id": self.grid.cell_id},
- )
-
# Initialise atmospheric variables required for update
atmosphere_setup = hydrology_tools.initialise_atmosphere_for_hydrology(
data=self.data,
@@ -333,15 +318,15 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
* soil_moisture, [mm]
* matric_potential, [kPa]
* surface_runoff, [mm]
- * surface_runoff_accumulated, [mm]
* subsurface_flow, [mm]
- * subsurface_flow_accumulated, [mm]
* soil_evaporation, [mm]
* vertical_flow, [mm d-1]
* groundwater_storage, [mm]
* subsurface_flow, [mm]
* baseflow, [mm]
- * total_river_discharge, [mm]
+ * surface_runoff_routed_plus_local, [mm]
+ * subsurface_runoff_routed_plus_local, [mm]
+ * total_runoff, [mm]
* river_discharge_rate, [m3 s-1]
* bypass flow, [mm]
* aerodynamic_resistance_surface, [s m-1]
@@ -366,10 +351,8 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
, the excess water is added to runoff and top soil moisture is set to soil
moisture capacity value; if the top soil is not saturated, precipitation is
added to the current topsoil moisture level and runoff is set to zero.
- The accumulated surface runoff is calculated as the sum of current runoff and
- the runoff from upstream cells at the previous time step, see
- :func:`~virtual_ecosystem.models.hydrology.above_ground.accumulate_horizontal_flow`
- .
+ The local contribution of a cell to the river channel is calculated as its own
+ surface and subsurface runoff for the current timestep.
Potential soil evaporation is calculated with classical bulk aerodynamic
formulation, following the so-called ':math:`\alpha` method', see
@@ -394,8 +377,8 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
. The horizontal flow between grid cells currently uses the same function as the
above ground runoff.
- Total river discharge is calculated as the sum of above- and below ground
- horizontal flow and converted to river discharge rate in m3/s.
+ Total runoff is calculated as the sum of above- and below ground
+ runoff and converted to river discharge rate in [m3 s-1].
The function requires the following input variables from the data object:
@@ -409,8 +392,6 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
* layer heights, [m]
* Soil moisture (previous time step), [mm]
* transpiration (current time step), [mm]
- * accumulated surface runoff (previous time step), [mm]
- * accumulated subsurface flow (previous time step), [mm]
* aerodynamic_resistance_canopy, [s m-1]
* net radiation, [W m-2]
@@ -511,20 +492,20 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
daily_lists["precipitation_surface"].append(precipitation_surface)
# Calculate daily surface runoff of each grid cell, [mm]
- surface_runoff = above_ground.calculate_surface_runoff(
+ surface_runoff_local = above_ground.calculate_surface_runoff(
precipitation_surface=precipitation_surface,
top_soil_moisture=hydro_input["current_soil_moisture"][0],
top_soil_moisture_saturation=hydro_input[
"top_soil_moisture_saturation"
],
)
- daily_lists["surface_runoff"].append(surface_runoff)
+ daily_lists["surface_runoff"].append(surface_runoff_local)
# Calculate preferential bypass flow, [mm]
bypass_flow = above_ground.calculate_bypass_flow(
top_soil_moisture=hydro_input["current_soil_moisture"][0],
sat_top_soil_moisture=hydro_input["top_soil_moisture_saturation"],
- available_water=precipitation_surface - surface_runoff,
+ available_water=precipitation_surface - surface_runoff_local,
bypass_flow_coefficient=(self.model_constants.bypass_flow_coefficient),
)
daily_lists["bypass_flow"].append(bypass_flow)
@@ -534,7 +515,7 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
(
hydro_input["current_soil_moisture"][0]
+ precipitation_surface
- - surface_runoff
+ - surface_runoff_local
- bypass_flow,
),
0,
@@ -669,38 +650,41 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
for var in ["groundwater_storage", "subsurface_flow", "baseflow"]:
daily_lists[var].append(below_ground_flow[var])
- # Calculate horizontal flow
- # Calculate accumulated runoff for each cell (me+sum of upstream neighbours)
- new_accumulated_runoff = above_ground.accumulate_horizontal_flow(
+ # Calculate horizontal flows between grid cells individually for output
+ # Surface runoff routed to each cell + local surface runoff
+ surface_runoff_routed_plus_local = above_ground.route_horizontal_flow(
drainage_map=self.drainage_map,
- current_flow=surface_runoff,
- previous_accumulated_flow=hydro_input["previous_accumulated_runoff"],
+ surface_runoff=surface_runoff_local,
+ subsurface_runoff=np.zeros_like(
+ surface_runoff_local
+ ), # only surface here
+ )
+ daily_lists["surface_runoff_routed_plus_local"].append(
+ surface_runoff_routed_plus_local
)
- daily_lists["surface_runoff_accumulated"].append(new_accumulated_runoff)
- # Calculate subsurface accumulated flow, [mm]
- new_subsurface_flow_accumulated = above_ground.accumulate_horizontal_flow(
+ # Subsurface runoff routed to each cell + local subsurface runoff
+ subsurface_flow = np.array(
+ below_ground_flow["subsurface_flow"] + below_ground_flow["baseflow"]
+ )
+ subsurface_runoff_routed_plus_local = above_ground.route_horizontal_flow(
drainage_map=self.drainage_map,
- current_flow=np.array(
- below_ground_flow["subsurface_flow"] + below_ground_flow["baseflow"]
- ),
- previous_accumulated_flow=(
- hydro_input["previous_subsurface_flow_accumulated"]
- ),
+ surface_runoff=np.zeros_like(subsurface_flow), # only subsurface here
+ subsurface_runoff=subsurface_flow,
)
- daily_lists["subsurface_flow_accumulated"].append(
- new_subsurface_flow_accumulated
+ daily_lists["subsurface_runoff_routed_plus_local"].append(
+ subsurface_runoff_routed_plus_local
)
- # Calculate total river discharge as sum of above- and below-ground flow
- total_river_discharge = (
- new_accumulated_runoff + new_subsurface_flow_accumulated
+ # Total runoff at each cell = surface + subsurface contributions
+ total_runoff = (
+ surface_runoff_routed_plus_local + subsurface_runoff_routed_plus_local
)
- daily_lists["total_river_discharge"].append(total_river_discharge)
+ daily_lists["total_runoff"].append(total_runoff)
- # Convert total discharge to river discharge rate, [m3 s-1]
+ # Convert total runoff [mm] to river discharge rate [m³/s]
river_discharge_rate = above_ground.convert_mm_flow_to_m3_per_second(
- river_discharge_mm=total_river_discharge,
+ river_discharge_mm=total_runoff,
area=self.grid.cell_area,
days=days,
seconds_to_day=self.core_constants.seconds_to_day,
@@ -708,13 +692,11 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
)
daily_lists["river_discharge_rate"].append(river_discharge_rate)
- # update inputs for next day
+ # Update other model states for next day
hydro_input["current_soil_moisture"] = soil_moisture_updated
hydro_input["groundwater_storage"] = below_ground_flow[
"groundwater_storage"
]
- hydro_input["previous_accumulated_runoff"] = new_accumulated_runoff
- hydro_input["subsurface_flow_accumulated"] = new_subsurface_flow_accumulated
# create output dict as intermediate step to not overwrite data directly
soil_hydrology = {}
@@ -727,9 +709,9 @@ def _update(self, time_index: int, **kwargs: Any) -> None:
"subsurface_flow",
"baseflow",
"bypass_flow",
- "surface_runoff_accumulated",
- "subsurface_flow_accumulated",
- "total_river_discharge",
+ "surface_runoff_routed_plus_local",
+ "subsurface_runoff_routed_plus_local",
+ "total_runoff",
]:
soil_hydrology[var] = DataArray(
np.nansum(np.stack(daily_lists[var], axis=1), axis=1),
diff --git a/virtual_ecosystem/models/hydrology/hydrology_tools.py b/virtual_ecosystem/models/hydrology/hydrology_tools.py
index f7121d0a7..7bf3fef9b 100644
--- a/virtual_ecosystem/models/hydrology/hydrology_tools.py
+++ b/virtual_ecosystem/models/hydrology/hydrology_tools.py
@@ -135,8 +135,6 @@ def setup_hydrology_input_current_timestep(
* current_soil_moisture
* top_soil_moisture_saturation
* top_soil_moisture_residual
- * previous_accumulated_runoff
- * previous_subsurface_flow_accumulated
* groundwater_storage
Args:
@@ -193,13 +191,7 @@ def setup_hydrology_input_current_timestep(
data["soil_moisture"][layer_structure.index_all_soil]
).to_numpy()
- # Get accumulated runoff/flow and ground water level from previous time step
- output["previous_accumulated_runoff"] = data[
- "surface_runoff_accumulated"
- ].to_numpy()
- output["previous_subsurface_flow_accumulated"] = data[
- "subsurface_flow_accumulated"
- ].to_numpy()
+ # Get ground water level
output["groundwater_storage"] = data["groundwater_storage"].to_numpy()
return output
@@ -305,3 +297,68 @@ def calculate_effective_saturation(
return (soil_moisture - soil_moisture_residual) / (
soil_moisture_saturation - soil_moisture_residual
)
+
+
+def check_monthly_mass_balance(
+ drainage_map: dict[int, list[int]],
+ surface_channel_inflow_mm: NDArray[np.floating],
+ monthly_precipitation_mm: NDArray[np.floating],
+ monthly_evaporation_mm: NDArray[np.floating],
+) -> None:
+ """Check that total monthly streamflow at outlet(s) does not exceed total precip.
+
+ The function identifies the outlet cells (cells with no downstream connections) from
+ the drainage map. It then sums the surface channel inflow at these outlet cells and
+ compares it to the total catchment precipitation minus total evaporation. If the
+ streamflow exceeds the available water, an AssertionError is raised.
+
+ If no true outlet cells exist, the flow from the lowest cells (cells with fewest
+ upstream connections) is used for the check.
+
+ Args:
+ drainage_map: Dict mapping each cell ID -> list of upstream cell IDs
+ surface_channel_inflow_mm: Monthly total surface channel inflow per cell, [mm]
+ monthly_precipitation_mm: Monthly total precipitation per cell, [mm]
+ monthly_evaporation_mm: Monthly total evaporation per cell, [mm]
+
+ Raises:
+ AssertionError: if monthly streamflow exceeds total catchment precipitation.
+ """
+ n_cells = len(drainage_map)
+ all_cells = set(range(n_cells))
+
+ # Cells that are upstream to others
+ cells_with_downstream = set()
+ for upstream_ids in drainage_map.values():
+ cells_with_downstream.update(upstream_ids)
+
+ # Outlet cells = cells with no downstream
+ outlet_cells = list(all_cells - cells_with_downstream)
+
+ # If no outlet, pick the "lowest" cells (fewest upstream neighbors)
+ if not outlet_cells:
+ upstream_counts = {
+ cell: len(upstream_ids) for cell, upstream_ids in drainage_map.items()
+ }
+ min_upstream = min(upstream_counts.values())
+ outlet_cells = [
+ cell for cell, count in upstream_counts.items() if count == min_upstream
+ ]
+
+ # Total streamflow at outlet(s)
+ monthly_outlet_flow_mm = surface_channel_inflow_mm[outlet_cells].sum()
+
+ # Total catchment precipitation and soil evaporation
+ total_catchment_precip_mm = np.sum(monthly_precipitation_mm)
+ total_catchment_evaporation_mm = np.sum(monthly_evaporation_mm)
+
+ # Mass balance check
+ if (
+ monthly_outlet_flow_mm
+ > total_catchment_precip_mm - total_catchment_evaporation_mm
+ ):
+ raise AssertionError(
+ f"Mass balance violated: total streamflow ({monthly_outlet_flow_mm:.2f} mm)"
+ f" exceeds total catchment precip ({total_catchment_precip_mm:.2f} mm). "
+ f"Outlet cells: {outlet_cells}"
+ )