lrntct · Oct 1, 2026 · Sep 30, 2026
diff --git a/src/itzi_core/__init__.py b/src/itzi_core/__init__.py
@@ -4,6 +4,15 @@
 from itzi_core.array_definitions import (
     INPUT_ARRAY_KEYS as INPUT_ARRAY_KEYS,
 )
+from itzi_core.array_definitions import (
+    INPUT_GREEN_AMPT_KEYS as INPUT_GREEN_AMPT_KEYS,
+)
+from itzi_core.array_definitions import (
+    INPUT_HYDROLOGY_KEYS as INPUT_HYDROLOGY_KEYS,
+)
+from itzi_core.array_definitions import (
+    INPUT_STAGE_KEYS as INPUT_STAGE_KEYS,
+)
 from itzi_core.array_definitions import (
     OUTPUT_ARRAY_KEYS as OUTPUT_ARRAY_KEYS,
 )
@@ -38,6 +47,9 @@
 __all__ = [
     "ARRAY_DEFINITIONS",
     "INPUT_ARRAY_KEYS",
+    "INPUT_GREEN_AMPT_KEYS",
+    "INPUT_HYDROLOGY_KEYS",
+    "INPUT_STAGE_KEYS",
     "OUTPUT_ARRAY_KEYS",
     "ArrayCategory",
     "ArrayDefinition",

diff --git a/src/itzi_core/array_definitions.py b/src/itzi_core/array_definitions.py
@@ -24,6 +24,9 @@ class ArrayCategory(Enum):
     """Categories of arrays in the simulation"""
 
     INPUT = "INPUT"  # Read from external sources (maps, time series)
+    INPUT_STAGE = "STAGE_INPUT_KEYS"
+    INPUT_GREEN_AMPT = "INPUT_GREEN_AMPT"
+    INPUT_HYDROLOGY = "INPUT_HYDROLOGY"
     INTERNAL = "INTERNAL"  # Computed during simulation (state variables)
     ACCUMULATION = "ACCUMULATION"  # Time-integrated values for reporting
     OUTPUT = "OUTPUT"  # Derived arrays for output/reporting
@@ -45,7 +48,7 @@ class ArrayDefinition:
     description: str  # Human-readable description
     unit: str  # Physical units of the array
     cf_unit: str  # The unit expected by the CF convention
-    var_loc: Literal["face", "edge"]  # Location of the value.
+    var_loc: Literal["face", "edge"]  # Location of the value in the cell.
     fill_value: float | int = 0.0  # Fill value (replace NaN)
     computes_from: str | None = None  # For accumulation arrays
     dtype: DTypeLike | None = None  # Optional storage dtype override
@@ -82,7 +85,11 @@ class ArrayDefinition:
         key="water_depth",
         csdms_name="land_surface_water__depth",
         cf_name="flood_water_thickness",
-        category=(ArrayCategory.INPUT, ArrayCategory.OUTPUT),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.OUTPUT,
+            ArrayCategory.INPUT_STAGE,
+        ),
         description="Water depth.",
         unit="m",
         cf_unit="m",
@@ -93,7 +100,11 @@ class ArrayDefinition:
         key="effective_porosity",
         csdms_name="soil_water__effective_porosity",
         cf_name="",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+            ArrayCategory.INPUT_GREEN_AMPT,
+        ),
         description="Porosity available to contribute to fluid flow.",
         unit="m/m",
         cf_unit="1",
@@ -104,7 +115,11 @@ class ArrayDefinition:
         key="capillary_pressure",
         csdms_name="soil_water__pressure_head",
         cf_name="",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+            ArrayCategory.INPUT_GREEN_AMPT,
+        ),
         description="Soil capillary pressure. Also called suction head.",
         unit="m",
         cf_unit="m",
@@ -115,7 +130,11 @@ class ArrayDefinition:
         key="hydraulic_conductivity",
         csdms_name="soil_water__hydraulic_conductivity",
         cf_name="soil_hydraulic_conductivity_at_saturation",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+            ArrayCategory.INPUT_GREEN_AMPT,
+        ),
         description="Soil’s ability to transmit water through its pores during infiltration.",
         unit="m s-1",
         cf_unit="m s-1",
@@ -126,7 +145,11 @@ class ArrayDefinition:
         key="soil_water_content",
         csdms_name="soil_water__volume_fraction",
         cf_name="volume_fraction_of_condensed_water_in_soil",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+            ArrayCategory.INPUT_GREEN_AMPT,
+        ),
         description="Relative soil water content.",
         unit="m/m",
         cf_unit="1",
@@ -137,7 +160,10 @@ class ArrayDefinition:
         key="infiltration",
         csdms_name="soil_surface_water__infiltration_leq-volume_flux",
         cf_name="",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+        ),
         description="User-defined infiltration rate.",
         unit="m s-1",
         cf_unit="",
@@ -148,7 +174,10 @@ class ArrayDefinition:
         key="losses",
         csdms_name="land_surface_water__losses_leq-volume_flux",
         cf_name="",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+        ),
         description="User-defined water losses.",
         unit="m s-1",
         cf_unit="",
@@ -159,7 +188,10 @@ class ArrayDefinition:
         key="rainfall_rate",
         csdms_name="atmosphere_water__precipitation_leq-volume_flux",
         cf_name="rainfall_rate",
-        category=(ArrayCategory.INPUT,),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.INPUT_HYDROLOGY,
+        ),
         description="User-provided precipitation rate.",
         unit="m s-1",
         cf_unit="",
@@ -475,7 +507,11 @@ class ArrayDefinition:
         key="water_surface_elevation",
         csdms_name="land_surface_water_surface__elevation",
         cf_name="water_surface_height_above_reference_datum",
-        category=(ArrayCategory.INPUT, ArrayCategory.OUTPUT),
+        category=(
+            ArrayCategory.INPUT,
+            ArrayCategory.OUTPUT,
+            ArrayCategory.INPUT_STAGE,
+        ),
         description="Water surface elevation. Terrain elevation + water depth",
         unit="m",
         cf_unit="",
@@ -601,6 +637,19 @@ class ArrayDefinition:
 INPUT_ARRAY_KEYS = frozenset(
     arr_def.key for arr_def in ARRAY_DEFINITIONS if ArrayCategory.INPUT in arr_def.category
 )
+INPUT_STAGE_KEYS = frozenset(
+    arr_def.key for arr_def in ARRAY_DEFINITIONS if ArrayCategory.INPUT_STAGE in arr_def.category
+)
+INPUT_HYDROLOGY_KEYS = frozenset(
+    arr_def.key
+    for arr_def in ARRAY_DEFINITIONS
+    if ArrayCategory.INPUT_HYDROLOGY in arr_def.category
+)
+INPUT_GREEN_AMPT_KEYS = frozenset(
+    arr_def.key
+    for arr_def in ARRAY_DEFINITIONS
+    if ArrayCategory.INPUT_GREEN_AMPT in arr_def.category
+)
 OUTPUT_ARRAY_KEYS = frozenset(
     arr_def.key for arr_def in ARRAY_DEFINITIONS if ArrayCategory.OUTPUT in arr_def.category
 )
diff --git a/src/itzi_core/simulation_builder.py b/src/itzi_core/simulation_builder.py
@@ -26,7 +26,12 @@
 from numpy.typing import ArrayLike, DTypeLike
 
 from itzi_core import infiltration
-from itzi_core.array_definitions import ARRAY_DEFINITIONS, ArrayCategory
+from itzi_core.array_definitions import (
+    ARRAY_DEFINITIONS,
+    INPUT_HYDROLOGY_KEYS,
+    INPUT_STAGE_KEYS,
+    ArrayCategory,
+)
 from itzi_core.const import InfiltrationModelType
 from itzi_core.data_containers import DrainageNodeCouplingData
 from itzi_core.drainage import CouplingTypes, DrainageLink, DrainageNode, DrainageSimulation
@@ -63,19 +68,6 @@
 class SimulationBuilder:
     """Builder for creating Simulation objects with different provider configurations."""
 
-    _STAGE_INPUT_KEYS = frozenset({"water_depth", "water_surface_elevation"})
-    _HYDROLOGY_INPUT_KEYS = frozenset(
-        {
-            "rainfall_rate",
-            "losses",
-            "infiltration",
-            "effective_porosity",
-            "capillary_pressure",
-            "hydraulic_conductivity",
-            "soil_water_content",
-        }
-    )
-
     def __init__(
         self,
         sim_config: SimulationConfig,
@@ -202,7 +194,7 @@ def _validate_resume_config_congruence(
             )
 
         changed_input_keys = self._changed_input_keys(resume_config)
-        changed_stage_keys = changed_input_keys & self._STAGE_INPUT_KEYS
+        changed_stage_keys = changed_input_keys & INPUT_STAGE_KEYS
         if changed_stage_keys:
             raise HotstartError(
                 "Hotstart input map changes are not supported for evolved stage inputs: "
@@ -432,7 +424,7 @@ def _restore_hotstart(
             )
             for array_key, array in updates:
                 simulation.set_array(array_key, array, simulation.sim_time)
-            if changed_input_keys & self._HYDROLOGY_INPUT_KEYS:
+            if changed_input_keys & INPUT_HYDROLOGY_KEYS:
                 simulation.schedule.set_deadline("hydrology", simulation.sim_time)
             if (
                 not changed_input_keys

diff --git a/tests/test_array_definitions.py b/tests/test_array_definitions.py
@@ -17,7 +17,14 @@
 
 import pytest
 
-from itzi_core.array_definitions import ARRAY_DEFINITIONS, ArrayCategory
+from itzi_core.array_definitions import (
+    ARRAY_DEFINITIONS,
+    INPUT_ARRAY_KEYS,
+    INPUT_GREEN_AMPT_KEYS,
+    INPUT_HYDROLOGY_KEYS,
+    INPUT_STAGE_KEYS,
+    ArrayCategory,
+)
 
 
 def test_array_definitions():
@@ -52,6 +59,12 @@ def test_computed_from_arrays_are_defined():
             assert arr_def.computes_from != arr_def.key
 
 
+def test_input_key_groups_are_consistent():
+    assert INPUT_STAGE_KEYS <= INPUT_ARRAY_KEYS
+    assert INPUT_HYDROLOGY_KEYS <= INPUT_ARRAY_KEYS
+    assert INPUT_GREEN_AMPT_KEYS <= INPUT_HYDROLOGY_KEYS
+
+
 def test_array_catalog_is_immutable() -> None:
     assert isinstance(ARRAY_DEFINITIONS, tuple)
     definition = ARRAY_DEFINITIONS[0]