move GA infiltration amount to the rasterdomain

Author: lrntct

src/itzi/array_definitions.py

@@ -15,7 +15,6 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
-from collections import Counter
 
 import numpy as np
 
@@ -209,6 +208,17 @@ class ArrayDefinition:
         cf_unit="",
         var_loc="face",
     ),
+    ArrayDefinition(
+        key="total_infiltration",
+        csdms_name="",
+        cf_name="",
+        category=[ArrayCategory.INTERNAL],
+        description="Total calculated infiltration. Used for Green-Ampt model",
+        unit="m",
+        cf_unit="",
+        var_loc="face",
+        fill_value=1e-4,
+    ),
     ArrayDefinition(
         key="eff_precip",
         csdms_name="land_surface_water__effective_precipitation_leq-volume_flux",
@@ -590,19 +600,3 @@ class ArrayDefinition:
     + _ACCUM_ARRAY_DEFINITIONS
     + _OUTPUT_ARRAY_DEFINITIONS
 )
-
-number_of_array_definitions = len(ARRAY_DEFINITIONS)
-
-# Some sanity check
-for attr in ["key", "csdms_name", "cf_name", "description"]:
-    all_values = [getattr(arr_def, attr) for arr_def in ARRAY_DEFINITIONS]
-    # No empty name
-    if "" in all_values and not attr == "cf_name":
-        raise ValueError(f"Found empty names in <{attr}>.")
-    # Make sure there is no duplicates
-    values_counts = Counter(all_values)
-    duplicates = [item for item, count in values_counts.items() if count > 1]
-    if 0 < len(duplicates):
-        if attr == "cf_name" and len(duplicates) == 1 and "" in duplicates:
-            continue
-        raise ValueError(f"Found duplicates in <{attr}>: {duplicates}")

src/itzi/infiltration.py

@@ -13,7 +13,6 @@
 """
 
 from datetime import timedelta
-import numpy as np
 
 import itzi.flow as flow
 from itzi.itzi_error import DtError
@@ -69,22 +68,14 @@ def step(self):
 class InfGreenAmpt(InfiltrationModel):
     """Calculate infiltration using Green-Ampt formula"""
 
-    def __init__(self, raster_domain, dt_inf):
-        InfiltrationModel.__init__(self, raster_domain, dt_inf)
-        # Initial cumulative infiltration set to tiny value
-        # (prevent division by zero)
-        self.infiltration_amount = np.full(
-            shape=self.dom.shape, fill_value=(1e-4), dtype=self.dom.dtype
-        )
-
     def step(self):
         """update infiltration rate map in m/s."""
         flow.infiltration_ga(
             arr_h=self.dom.get_array("water_depth"),
             arr_eff_por=self.dom.get_array("effective_porosity"),
             arr_pressure=self.dom.get_array("capillary_pressure"),
             arr_conduct=self.dom.get_array("hydraulic_conductivity"),
-            arr_inf_amount=self.infiltration_amount,
+            arr_inf_amount=self.dom.get_array("total_infiltration"),
             arr_water_soil_content=self.dom.get_array("soil_water_content"),
             arr_inf_out=self.dom.get_array("computed_infiltration"),
             dt=self._dt,

src/itzi/rasterdomain.py

@@ -104,12 +104,8 @@ def __init__(self, dtype, arr_mask, cell_shape):
 
         # slice for a simple padding (allow stencil calculation on boundary)
         self.simple_pad = (slice(1, -1), slice(1, -1))
-        # Fill values for input arrays
-        self.input_fill_values = {
-            arr_def.key: arr_def.fill_value
-            for arr_def in ARRAY_DEFINITIONS
-            if ArrayCategory.INPUT in arr_def.category
-        }
+        # Fill values
+        self.fill_values = {arr_def.key: arr_def.fill_value for arr_def in ARRAY_DEFINITIONS}
 
         # all keys that will be used for the arrays
         self.k_input = [
@@ -130,14 +126,7 @@ def __init__(self, dtype, arr_mask, cell_shape):
         # Instantiate arrays and padded arrays filled with zeros
         self.arr = dict.fromkeys(self.k_all)
         self.arrp = dict.fromkeys(self.k_all)
-        self.create_arrays()
-
-    def zeros_array(self):
-        """return a np array of the domain dimension, filled with zeros.
-        dtype is set to object's dtype.
-        Intended to be used as default for the input model maps.
-        """
-        return np.zeros(shape=self.shape, dtype=self.dtype)
+        self._create_arrays()
 
     def pad_array(self, arr):
         """Return the original input array
@@ -147,12 +136,13 @@ def pad_array(self, arr):
         arr = arr_p[self.simple_pad]
         return arr, arr_p
 
-    def create_arrays(self):
+    def _create_arrays(self):
         """Instantiate masked arrays and padded arrays
         the unpadded arrays are a slice of the padded ones
         """
         for k in self.arr.keys():
-            self.arr[k], self.arrp[k] = self.pad_array(self.zeros_array())
+            arr = np.full(fill_value=self.fill_values[k], shape=self.shape, dtype=self.dtype)
+            self.arr[k], self.arrp[k] = self.pad_array(arr)
         return self
 
     def update_mask(self, arr):
@@ -196,14 +186,13 @@ def swap_arrays(self, k1, k2):
 
     def update_array(self, arr_key, arr):
         """Update the values of an array with those of a given array."""
-        fill_value = self.input_fill_values[arr_key]
         if arr.shape != self.shape:
-            return ValueError
+            return ValueError(f"Updated values for array '{arr_key}' do not match domain size.")
         if arr_key == "water_surface_elevation":
             # Calculate actual depth and update the internal depth array
             arr = rastermetrics.calculate_h_from_wse(arr_wse=arr, arr_dem=self.get_array("dem"))
             arr_key = "water_depth"
-        self.mask_array(arr, fill_value)
+        self.mask_array(arr, self.fill_values[arr_key])
         self.arr[arr_key][:], self.arrp[arr_key][:] = self.pad_array(arr)
         return self
 

tests/test_array_definitions.py

@@ -0,0 +1,41 @@
+from collections import Counter
+
+from itzi.array_definitions import ARRAY_DEFINITIONS, ArrayCategory
+
+
+def test_array_definitions():
+    # Empty names is acceptable in internal arrays
+    for attr in ["csdms_name", "cf_name"]:
+        all_values = [
+            getattr(arr_def, attr)
+            for arr_def in ARRAY_DEFINITIONS
+            if ArrayCategory.INTERNAL not in arr_def.category
+        ]
+        print(all_values)
+        # No empty name
+        if "" in all_values and not attr == "cf_name":
+            print(attr)
+            assert False, f"Found empty names in <{attr}>."
+        # Make sure there is no duplicates
+        values_counts = Counter(all_values)
+        duplicates = [item for item, count in values_counts.items() if count > 1]
+        if 0 < len(duplicates):
+            if attr == "cf_name" and len(duplicates) == 1 and "" in duplicates:
+                continue
+            assert False, f"Found duplicates in <{attr}>: {duplicates}"
+
+    # All arrays must have a unique key and description
+    for attr in ["key", "description"]:
+        all_values = [getattr(arr_def, attr) for arr_def in ARRAY_DEFINITIONS]
+        print(all_values)
+        # No empty name
+        if "" in all_values and not attr == "cf_name":
+            print(attr)
+            assert False, f"Found empty names in <{attr}>."
+        # Make sure there is no duplicates
+        values_counts = Counter(all_values)
+        duplicates = [item for item, count in values_counts.items() if count > 1]
+        if 0 < len(duplicates):
+            if attr == "cf_name" and len(duplicates) == 1 and "" in duplicates:
+                continue
+            assert False, f"Found duplicates in <{attr}>: {duplicates}"

tests/test_infiltration.py

@@ -102,7 +102,7 @@ def infiltration_sim(infiltration_parameters):
     while elapsed_time < inf_params.time:
         inf_sim.step()
         elapsed_time += inf_sim._dt
-    return inf_sim.infiltration_amount.max()
+    return raster_domain.get_array("total_infiltration").max()
 
 
 def test_infiltration(reference_infiltration, infiltration_sim):