some optimizations

geb/HRUs.py

@@ -190,9 +190,31 @@ def M3toM(self, array: np.ndarray) -> np.ndarray:
         return array / self.cellArea
 
     def register_initial_data(self, name: str) -> None:
+        """Register initial data."""
         self.data.initial_conditions.append(name)
 
-    def load_initial(self, name, default=0.0, gpu=False):
+    def load_initial(self, name, default, gpu=False):
+        """Load initial data from disk when variable is set, otherwise use default value.
+
+        Args:
+            name: Name of variable.
+            default: Default value. This must be a callable function or lambda.
+            gpu: Whether to use GPU.
+
+        Returns:
+            data: Loaded data.
+
+        Raises:
+            AssertionError: If default is not a callable function or lambda.
+
+        Examples:
+            >>> def default():
+            ...     return np.zeros(self.data.grid.shape)
+            >>> self.var.load_initial("name", default=default)
+
+            >>> self.var.load_initial("name", default=lambda: np.zeros(self.data.grid.shape))
+        """
+        assert callable(default), "default must be a callable function or lambda"
         if self.model.load_initial_data:
             fp = os.path.join(self.data.get_save_state_path(), f"{name}.npz")
             if gpu:
@@ -201,7 +223,7 @@ def load_initial(self, name, default=0.0, gpu=False):
                 return np.load(fp)["data"]
         else:
             self.register_initial_data(name)
-            return default
+            return default()
 
 
 class Grid(BaseVariables):
@@ -344,7 +366,7 @@ def load(self, filepath, compress=True, layer=1):
             data = self.data.grid.compress(data)
         return data
 
-    def load_initial(self, name, default=0.0):
+    def load_initial(self, name, default):
         return super().load_initial("grid." + name, default=default)
 
     def load_forcing_ds(self, name):
@@ -842,7 +864,7 @@ def plot(self, HRU_array: np.ndarray, ax=None, show: bool = True):
         if show:
             plt.show()
 
-    def load_initial(self, name, default=0.0, gpu=None):
+    def load_initial(self, name, default, gpu=None):
         if gpu is None:
             gpu = self.model.use_gpu
         return super().load_initial("HRU." + name, default=default, gpu=gpu)
@@ -900,7 +922,7 @@ def __init__(self, data, model):
 
         BaseVariables.__init__(self)
 
-    def load_initial(self, name, default=0.0):
+    def load_initial(self, name, default):
         return super().load_initial("modflow." + name, default=default)
 
 

geb/agents/crop_farmers.py

@@ -858,25 +858,27 @@ def __init__(self, model, agents, reduncancy: float) -> None:
 
         self.var.actual_evapotranspiration_crop_life = self.var.load_initial(
             "actual_evapotranspiration_crop_life",
-            default=self.var.full_compressed(0, dtype=np.float32, gpu=False),
+            default=lambda: self.var.full_compressed(0, dtype=np.float32, gpu=False),
             gpu=False,
         )
         self.var.potential_evapotranspiration_crop_life = self.var.load_initial(
             "potential_evapotranspiration_crop_life",
-            default=self.var.full_compressed(0, dtype=np.float32, gpu=False),
+            default=lambda: self.var.full_compressed(0, dtype=np.float32, gpu=False),
             gpu=False,
         )
         self.var.crop_map = self.var.load_initial(
-            "crop_map", default=np.full_like(self.var.land_owners, -1), gpu=False
+            "crop_map",
+            default=lambda: np.full_like(self.var.land_owners, -1),
+            gpu=False,
         )
         self.var.crop_age_days_map = self.var.load_initial(
             "crop_age_days_map",
-            default=np.full_like(self.var.land_owners, -1),
+            default=lambda: np.full_like(self.var.land_owners, -1),
             gpu=False,
         )
         self.var.crop_harvest_age_days = self.var.load_initial(
             "crop_harvest_age_days",
-            default=np.full_like(self.var.land_owners, -1),
+            default=lambda: np.full_like(self.var.land_owners, -1),
             gpu=False,
         )
 

geb/hydrology/groundwater/__init__.py

@@ -82,7 +82,7 @@ def get_initial_head():
 
         self.var.heads = self.model.data.grid.load_initial(
             "heads",
-            default=get_initial_head(),
+            default=get_initial_head,
         )
 
         self.modflow = ModFlowSimulation(
@@ -100,7 +100,7 @@ def get_initial_head():
         )
 
         self.var.capillar = self.var.load_initial(
-            "capillar", default=self.var.full_compressed(0, dtype=np.float32)
+            "capillar", default=lambda: self.var.full_compressed(0, dtype=np.float32)
         )
 
     def step(self, groundwater_recharge, groundwater_abstraction_m3):

geb/hydrology/interception.py

@@ -70,7 +70,7 @@ def __init__(self, model):
 
         self.var.interceptStor = self.model.data.HRU.load_initial(
             "interceptStor",
-            default=self.model.data.HRU.full_compressed(0, dtype=np.float32),
+            default=lambda: self.model.data.HRU.full_compressed(0, dtype=np.float32),
         )
 
         minimum_intercept_capacity = {

geb/hydrology/lakes_reservoirs.py

@@ -244,7 +244,7 @@ def __init__(self, model):
 
         self.var.total_inflow_from_other_water_bodies = self.var.load_initial(
             "total_inflow_from_other_water_bodies",
-            default=np.zeros_like(self.var.volume, dtype=np.float32),
+            default=lambda: np.zeros_like(self.var.volume, dtype=np.float32),
         )
 
         # lake discharge at outlet to calculate alpha: parameter of channel width, gravity and weir coefficient
@@ -264,7 +264,7 @@ def __init__(self, model):
         )
 
         self.var.storage = self.var.load_initial(
-            "storage", default=self.var.volume.copy()
+            "storage", default=lambda: self.var.volume.copy()
         )
         self.var.outflow_height = estimate_outflow_height(
             self.var.volume, self.lake_factor, self.var.lake_area, average_discharge

geb/hydrology/routing/__init__.py

@@ -238,16 +238,16 @@ def __init__(self, model):
 
         # channel water volume [m3]
         # Initialise water volume in kinematic wave channels [m3]
-        channelStorageM3Ini = self.var.totalCrossSectionArea * self.var.chanLength * 0.1
         self.var.channelStorageM3 = self.var.load_initial(
-            "channelStorageM3", default=channelStorageM3Ini
+            "channelStorageM3",
+            default=lambda: self.var.totalCrossSectionArea * self.var.chanLength * 0.1,
         )
         # Initialise discharge at kinematic wave pixels (note that InvBeta is
         # simply 1/beta, computational efficiency!)
         # self.var.chanQKin = np.where(self.var.channelAlpha > 0, (self.var.totalCrossSectionArea / self.var.channelAlpha) ** self.var.invbeta, 0.)
         self.var.discharge = self.var.load_initial(
             "discharge",
-            default=(
+            default=lambda: (
                 self.var.channelStorageM3
                 * self.var.invchanLength
                 * self.var.invchannelAlpha
@@ -256,7 +256,7 @@ def __init__(self, model):
         )
         self.var.discharge_substep = self.var.load_initial(
             "discharge_substep",
-            default=np.full(
+            default=lambda: np.full(
                 (self.var.noRoutingSteps, self.var.discharge.size),
                 0,
                 dtype=self.var.discharge.dtype,
@@ -479,7 +479,7 @@ def step(self, openWaterEvap, channel_abstraction_m, return_flow):
                 prestorages=[pre_channel_storage_m3],
                 poststorages=[self.var.channelStorageM3],
                 name="routing_3",
-                tollerance=1e-8,
+                tollerance=100,
             )
             balance_check(
                 how="sum",
@@ -492,5 +492,5 @@ def step(self, openWaterEvap, channel_abstraction_m, return_flow):
                 prestorages=[pre_channel_storage_m3, pre_storage],
                 poststorages=[self.var.channelStorageM3, self.var.storage],
                 name="routing_4",
-                tollerance=1e-8,
+                tollerance=100,
             )

geb/hydrology/snow_frost.py

@@ -172,14 +172,15 @@ def __init__(self, model, elevation_std):
 
         # initialize snowcovers as many as snow layers -> read them as SnowCover1 , SnowCover2 ...
         # SnowCover1 is the highest zone
-        SnowCoverS = np.tile(
-            self.model.data.to_HRU(
-                data=self.model.data.grid.full_compressed(0, dtype=np.float32), fn=None
-            ),
-            (self.numberSnowLayers, 1),
-        )
         self.var.SnowCoverS = self.model.data.HRU.load_initial(
-            "SnowCoverS", default=SnowCoverS
+            "SnowCoverS",
+            default=lambda: np.tile(
+                self.model.data.to_HRU(
+                    data=self.model.data.grid.full_compressed(0, dtype=np.float32),
+                    fn=None,
+                ),
+                (self.numberSnowLayers, 1),
+            ),
         )
 
         # Pixel-average initial snow cover: average of values in 3 elevation
@@ -194,7 +195,7 @@ def __init__(self, model, elevation_std):
 
         self.var.frost_index = self.model.data.HRU.load_initial(
             "frost_index",
-            default=self.model.data.HRU.full_compressed(0, dtype=np.float32),
+            default=lambda: self.model.data.HRU.full_compressed(0, dtype=np.float32),
         )
 
         self.var.extfrost_index = False

geb/hydrology/soil.py

@@ -840,8 +840,7 @@ def vertical_water_transport(
             w[sink, i] = min(w[sink, i], ws[sink, i])
             w[source, i] = max(w[source, i], wres[source, i])
 
-        if is_bioarea[i]:
-            groundwater_recharge[i] = net_fluxes[-1, i] + preferential_flow[i]
+        groundwater_recharge[i] = net_fluxes[-1, i] + preferential_flow[i]
 
     return preferential_flow, direct_runoff, groundwater_recharge, net_fluxes
 
@@ -920,13 +919,18 @@ def __init__(self, model, elevation_std):
         self.wres = thetar * self.soil_layer_height
 
         # initial soil water storage between field capacity and wilting point
+        # set soil moisture to nan where land use is not bioarea
         self.var.w = self.model.data.HRU.load_initial(
             "w",
-            default=(self.wfc + self.wwp) / 2,
+            default=lambda: np.where(
+                self.var.land_use_type[np.newaxis, :] < SEALED,
+                (self.wfc + self.wwp) / 2,
+                np.nan,
+            ),
         )
         # for paddy irrigation flooded paddy fields
         self.var.topwater = self.model.data.HRU.load_initial(
-            "topwater", default=self.var.full_compressed(0, dtype=np.float32)
+            "topwater", default=lambda: self.var.full_compressed(0, dtype=np.float32)
         )
 
         lambda_pore_size_distribution = self.model.data.grid.load(
@@ -984,7 +988,7 @@ def __init__(self, model, elevation_std):
 
         self.var.aeration_days_counter = self.var.load_initial(
             "aeration_days_counter",
-            default=np.full(
+            default=lambda: np.full(
                 (N_SOIL_LAYERS, self.var.compressed_size), 0, dtype=np.int32
             ),
         )
@@ -1283,7 +1287,7 @@ def step(
 
             preferential_flow += preferential_flow_substep
             direct_runoff += direct_runoff_substep
-            groundwater_recharge += groundwater_recharge_substep
+            groundwater_recharge[bioarea] += groundwater_recharge_substep[bioarea]
 
         assert (self.var.w[:, bioarea] <= self.ws[:, bioarea]).all()
         assert (self.var.w[:, bioarea] >= self.wres[:, bioarea]).all()