Source code for pyiwfm.components.small_watershed
"""
Small Watershed component classes for IWFM models.
This module provides classes for representing small watersheds, including
watershed units with root zone and aquifer parameters, and the main
application class. It mirrors IWFM's Package_AppSmallWatershed.
"""
from __future__ import annotations
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from pyiwfm.core.base_component import BaseComponent
from pyiwfm.core.exceptions import ComponentError
if TYPE_CHECKING:
from pyiwfm.io.small_watershed import SmallWatershedMainConfig
[docs]
@dataclass
class WatershedGWNode:
"""Groundwater node connection for a small watershed.
Attributes:
gw_node_id: Groundwater node ID (1-based)
max_perc_rate: Maximum percolation rate (positive) or baseflow layer (negative)
is_baseflow: Whether this is a baseflow node
layer: Baseflow layer (if is_baseflow)
"""
gw_node_id: int = 0
max_perc_rate: float = 0.0
is_baseflow: bool = False
layer: int = 0
[docs]
@dataclass
class WatershedUnit:
"""A single small watershed unit.
Attributes:
id: Watershed unit ID (1-based)
area: Watershed area
dest_stream_node: Destination stream node for outflow (1-based)
gw_nodes: Connected groundwater nodes
precip_col: Precipitation time-series column index
precip_factor: Precipitation conversion factor
et_col: ET time-series column index
wilting_point: Soil wilting point
field_capacity: Soil field capacity
total_porosity: Total porosity
lambda_param: Pore size distribution parameter
root_depth: Root zone depth
hydraulic_cond: Hydraulic conductivity
kunsat_method: Unsaturated K method code
curve_number: SCS curve number
gw_threshold: Groundwater storage threshold
max_gw_storage: Maximum groundwater storage
surface_flow_coeff: Surface flow recession coefficient
baseflow_coeff: Baseflow recession coefficient
"""
id: int = 0
area: float = 0.0
dest_stream_node: int = 0
gw_nodes: list[WatershedGWNode] = field(default_factory=list)
# Root zone parameters
precip_col: int = 0
precip_factor: float = 1.0
et_col: int = 0
wilting_point: float = 0.0
field_capacity: float = 0.0
total_porosity: float = 0.0
lambda_param: float = 0.0
root_depth: float = 0.0
hydraulic_cond: float = 0.0
kunsat_method: int = 0
curve_number: float = 0.0
# Aquifer parameters
gw_threshold: float = 0.0
max_gw_storage: float = 0.0
surface_flow_coeff: float = 0.0
baseflow_coeff: float = 0.0
# Initial conditions
initial_soil_moisture: float = 0.0
initial_gw_storage: float = 0.0
@property
def n_gw_nodes(self) -> int:
"""Return number of connected groundwater nodes."""
return len(self.gw_nodes)
def __repr__(self) -> str:
return f"WatershedUnit(id={self.id}, area={self.area:.1f})"
[docs]
@dataclass
class AppSmallWatershed(BaseComponent):
"""Small Watershed application component.
This class manages all small watersheds in the model domain including
watershed units with root zone and aquifer parameters. It mirrors
IWFM's Package_AppSmallWatershed.
Attributes:
watersheds: Dictionary mapping watershed ID to WatershedUnit
area_factor: Area conversion factor
flow_factor: Flow rate conversion factor
flow_time_unit: Time unit for flow rates
rz_solver_tolerance: Root zone solver tolerance
rz_max_iterations: Root zone solver max iterations
rz_length_factor: Root zone length conversion factor
rz_cn_factor: Curve number conversion factor
rz_k_factor: Hydraulic conductivity conversion factor
rz_k_time_unit: Time unit for hydraulic conductivity
aq_gw_factor: GW conversion factor
aq_time_factor: Time conversion factor
aq_time_unit: Time unit for recession coefficients
budget_output_file: Path to budget output file
final_results_file: Path to final simulation results file
"""
watersheds: dict[int, WatershedUnit] = field(default_factory=dict)
# Geospatial conversion factors
area_factor: float = 1.0
flow_factor: float = 1.0
flow_time_unit: str = ""
# Root zone solver parameters
rz_solver_tolerance: float = 1e-8
rz_max_iterations: int = 2000
rz_length_factor: float = 1.0
rz_cn_factor: float = 1.0
rz_k_factor: float = 1.0
rz_k_time_unit: str = ""
# Aquifer conversion factors
aq_gw_factor: float = 1.0
aq_time_factor: float = 1.0
aq_time_unit: str = ""
# Initial conditions conversion factor
ic_factor: float = 1.0
# Output files
budget_output_file: str = ""
final_results_file: str = ""
@property
def n_items(self) -> int:
"""Return number of watersheds (primary entities)."""
return len(self.watersheds)
@property
def n_watersheds(self) -> int:
"""Return number of watersheds."""
return len(self.watersheds)
[docs]
def add_watershed(self, ws: WatershedUnit) -> None:
"""Add a watershed unit to the component."""
self.watersheds[ws.id] = ws
[docs]
def get_watershed(self, ws_id: int) -> WatershedUnit:
"""Get a watershed unit by ID."""
return self.watersheds[ws_id]
[docs]
def iter_watersheds(self) -> Iterator[WatershedUnit]:
"""Iterate over watersheds in ID order."""
for wid in sorted(self.watersheds.keys()):
yield self.watersheds[wid]
[docs]
def validate(self) -> None:
"""Validate the small watershed component.
Raises:
ComponentError: If component is invalid
"""
for ws in self.watersheds.values():
if ws.area <= 0:
raise ComponentError(f"Watershed {ws.id} has non-positive area: {ws.area}")
if ws.dest_stream_node <= 0:
raise ComponentError(
f"Watershed {ws.id} has invalid destination stream node: {ws.dest_stream_node}"
)
if ws.n_gw_nodes == 0:
raise ComponentError(f"Watershed {ws.id} has no connected GW nodes")
[docs]
@classmethod
def from_config(cls, config: SmallWatershedMainConfig) -> AppSmallWatershed:
"""Create component from a parsed SmallWatershedMainConfig.
Args:
config: Parsed configuration from the reader
Returns:
AppSmallWatershed instance
"""
comp = cls(
area_factor=config.area_factor,
flow_factor=config.flow_factor,
flow_time_unit=config.flow_time_unit,
rz_solver_tolerance=config.rz_solver_tolerance,
rz_max_iterations=config.rz_max_iterations,
rz_length_factor=config.rz_length_factor,
rz_cn_factor=config.rz_cn_factor,
rz_k_factor=config.rz_k_factor,
rz_k_time_unit=config.rz_k_time_unit,
aq_gw_factor=config.aq_gw_factor,
aq_time_factor=config.aq_time_factor,
aq_time_unit=config.aq_time_unit,
ic_factor=config.ic_factor,
budget_output_file=(
str(config.budget_output_file) if config.budget_output_file else ""
),
final_results_file=(
str(config.final_results_file) if config.final_results_file else ""
),
)
# Build lookup dicts for rootzone, aquifer, and IC params by ID
rz_by_id = {rz.id: rz for rz in config.rootzone_params}
aq_by_id = {aq.id: aq for aq in config.aquifer_params}
ic_by_id = {ic.id: ic for ic in config.initial_conditions}
for spec in config.watershed_specs:
rz = rz_by_id.get(spec.id)
aq = aq_by_id.get(spec.id)
ws = WatershedUnit(
id=spec.id,
area=spec.area,
dest_stream_node=spec.dest_stream_node,
gw_nodes=[
WatershedGWNode(
gw_node_id=gn.gw_node_id,
max_perc_rate=gn.max_perc_rate,
is_baseflow=gn.is_baseflow,
layer=gn.layer,
)
for gn in spec.gw_nodes
],
)
if rz is not None:
ws.precip_col = rz.precip_col
ws.precip_factor = rz.precip_factor
ws.et_col = rz.et_col
ws.wilting_point = rz.wilting_point
ws.field_capacity = rz.field_capacity
ws.total_porosity = rz.total_porosity
ws.lambda_param = rz.lambda_param
ws.root_depth = rz.root_depth
ws.hydraulic_cond = rz.hydraulic_cond
ws.kunsat_method = rz.kunsat_method
ws.curve_number = rz.curve_number
if aq is not None:
ws.gw_threshold = aq.gw_threshold
ws.max_gw_storage = aq.max_gw_storage
ws.surface_flow_coeff = aq.surface_flow_coeff
ws.baseflow_coeff = aq.baseflow_coeff
ic = ic_by_id.get(spec.id)
if ic is not None:
ws.initial_soil_moisture = ic.soil_moisture
ws.initial_gw_storage = ic.gw_storage
comp.add_watershed(ws)
return comp
def __repr__(self) -> str:
return f"AppSmallWatershed(n_watersheds={self.n_watersheds})"