allow exporting of various values for water balance

geb/hydrology/landcover.py

@@ -433,16 +433,21 @@ def step(self):
         assert not np.isnan(channel_abstraction_m).any()
         assert not np.isnan(open_water_evaporation).any()
 
-        if __debug__:
-            total_evapotranspiration = (
-                self.var.actual_evapotranspiration
-                + actual_bare_soil_evaporation
-                + open_water_evaporation
-                + self.var.interception_evaporation
-                + self.var.snowEvap  # ice should be included in the future
-                + irrigation_loss_to_evaporation_m
-            )
+        total_evapotranspiration = (
+            self.var.actual_evapotranspiration
+            + actual_bare_soil_evaporation
+            + open_water_evaporation
+            + self.var.interception_evaporation
+            + self.var.snowEvap  # ice should be included in the future
+            + irrigation_loss_to_evaporation_m
+        )
 
+        self.model.data.grid.total_evapotranspiration_m3 = self.model.data.to_grid(
+            HRU_data=self.model.data.HRU.MtoM3(total_evapotranspiration),
+            fn="sum",
+        )
+
+        if __debug__:
             balance_check(
                 name="landcover_1",
                 how="cellwise",

geb/hydrology/routing/__init__.py

@@ -104,7 +104,7 @@ class Routing(object):
     QDelta
     act_bigLakeResAbst
     act_smallLakeResAbst
-    returnFlow
+    return_flow
     sumsideflow
     inflowDt
     ====================  ================================================================================  =========
@@ -267,7 +267,7 @@ def __init__(self, model):
             + self.model.config["parameters"]["lakeEvaFactor"]
         )
 
-    def step(self, openWaterEvap, channel_abstraction_m, returnFlow):
+    def step(self, openWaterEvap, channel_abstraction_m, return_flow):
         """
         Dynamic part of the routing module
 
@@ -368,9 +368,9 @@ def step(self, openWaterEvap, channel_abstraction_m, returnFlow):
         # WDAddM3Dt = self.var.act_SurfaceWaterAbstract.copy() #MS CWatM edit Shouldn't this only be from the river abstractions? Currently includes the larger reservoir...
         WDAddMDt = channel_abstraction_m
         # return flow from (m) non irrigation water demand
-        # WDAddM3Dt = WDAddM3Dt - self.var.nonIrrReturnFlowFraction * self.var.act_nonIrrDemand
+        # WDAddM3Dt = WDAddM3Dt - self.var.nonIrrreturn_flowFraction * self.var.act_nonIrrDemand
         WDAddMDt = (
-            WDAddMDt - returnFlow
+            WDAddMDt - return_flow
         )  # Couldn't this be negative? If return flow is mainly coming from gw? Fine, then more water going in.
         WDAddM3Dt = WDAddMDt * self.var.cellArea / self.var.noRoutingSteps
 
@@ -397,7 +397,7 @@ def step(self, openWaterEvap, channel_abstraction_m, returnFlow):
             # sideflowChanM3 -= riverbedExchangeDt
 
             sideflowChanM3 -= WDAddM3Dt
-            # minus waterdemand + returnflow
+            # minus waterdemand + return_flow
 
             outflow_to_river_network, waterbody_evaporation = (
                 self.model.lakes_reservoirs.routing(

geb/hydrology/water_demand.py

@@ -254,14 +254,15 @@ def step(self, potential_evapotranspiration):
         domestic_water_demand_m3 = self.model.data.grid.MtoM3(domestic_water_demand)
         del domestic_water_demand
 
-        domestic_withdrawal_m3 = self.withdraw(
+        self.model.data.grid.domestic_withdrawal_m3 = self.withdraw(
             available_channel_storage_m3, domestic_water_demand_m3
         )  # withdraw from surface water
-        domestic_withdrawal_m3 += self.withdraw(
+        self.model.data.grid.domestic_withdrawal_m3 += self.withdraw(
             available_groundwater_m3, domestic_water_demand_m3
         )  # withdraw from groundwater
         domestic_return_flow_m = self.model.data.grid.M3toM(
-            domestic_withdrawal_m3 * (1 - domestic_water_efficiency)
+            self.model.data.grid.domestic_withdrawal_m3
+            * (1 - domestic_water_efficiency)
         )
 
         # 2. industry (surface + ground)
@@ -271,14 +272,15 @@ def step(self, potential_evapotranspiration):
         industry_water_demand_m3 = self.model.data.grid.MtoM3(industry_water_demand)
         del industry_water_demand
 
-        industry_withdrawal_m3 = self.withdraw(
+        self.model.data.grid.industry_withdrawal_m3 = self.withdraw(
             available_channel_storage_m3, industry_water_demand_m3
         )  # withdraw from surface water
-        industry_withdrawal_m3 += self.withdraw(
+        self.model.data.grid.industry_withdrawal_m3 += self.withdraw(
             available_groundwater_m3, industry_water_demand_m3
         )  # withdraw from groundwater
         industry_return_flow_m = self.model.data.grid.M3toM(
-            industry_withdrawal_m3 * (1 - industry_water_efficiency)
+            self.model.data.grid.industry_withdrawal_m3
+            * (1 - industry_water_efficiency)
         )
 
         # 3. livestock (surface)
@@ -288,11 +290,12 @@ def step(self, potential_evapotranspiration):
         livestock_water_demand_m3 = self.model.data.grid.MtoM3(livestock_water_demand)
         del livestock_water_demand
 
-        livestock_withdrawal_m3 = self.withdraw(
+        self.model.data.grid.livestock_withdrawal_m3 = self.withdraw(
             available_channel_storage_m3, livestock_water_demand_m3
         )  # withdraw from surface water
         livestock_return_flow_m = self.model.data.grid.M3toM(
-            livestock_withdrawal_m3 * (1 - livestock_water_efficiency)
+            self.model.data.grid.livestock_withdrawal_m3
+            * (1 - livestock_water_efficiency)
         )
         timer.new_split("Water withdrawal")
 
@@ -351,6 +354,10 @@ def step(self, potential_evapotranspiration):
 
         assert (self.var.actual_irrigation_consumption + 1e-5 >= 0).all()
 
+        self.model.data.grid.irrigation_consumption_m3 = self.model.data.to_grid(
+            HRU_data=self.var.MtoM3(self.var.actual_irrigation_consumption), fn="sum"
+        )
+
         groundwater_abstraction_m3 = (
             available_groundwater_m3_pre - available_groundwater_m3
         )
@@ -402,9 +409,9 @@ def step(self, potential_evapotranspiration):
                 how="sum",
                 influxes=[],
                 outfluxes=[
-                    domestic_withdrawal_m3,
-                    industry_withdrawal_m3,
-                    livestock_withdrawal_m3,
+                    self.model.data.grid.domestic_withdrawal_m3,
+                    self.model.data.grid.industry_withdrawal_m3,
+                    self.model.data.grid.livestock_withdrawal_m3,
                     (
                         irrigation_water_withdrawal_m * self.var.cellArea.get()
                         if self.model.use_gpu
@@ -429,6 +436,6 @@ def step(self, potential_evapotranspiration):
         return (
             groundwater_abstraction_m3,
             channel_abstraction_m3 / self.model.data.grid.cellArea,
-            return_flow,  # from all sources
+            return_flow,  # from all sources, re-added in routing
             irrigation_loss_to_evaporation_m,
         )