diff --git a/src/itzi_core/__init__.py b/src/itzi_core/__init__.py index 955ba93..da7efed 100644 --- 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 index 1ea4547..ebca0c2 100644 --- 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 index 94e66ba..460ded6 100644 --- 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 index 31496bc..c440545 100644 --- 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]