Rename UnifiedModelSetup to UnifiedCalibrator and add ET output support

Author: zhuanglaihong

configs/example_config.yaml

@@ -3,35 +3,37 @@
 
 data:
   # Data source configuration
-  dataset: "camels_us"              # Dataset type: camels_us, selfmadehydrodataset, etc.
+  dataset: "camels"              # Dataset type: camels_us, selfmadehydrodataset, etc.
   path: null                        # Data path (null = auto-detect from hydro_setting.yml or use ~/hydromodel_data/)
                                     # If specified, provide the datasets-origin directory only
                                     # Example: "D:/project/data" (NOT "D:/project/data/CAMELS_US")
                                     # aqua_fetch will automatically append "CAMELS_US" directory name
-  basin_ids: ["01013500","01022500","11532500"]           # List of basin IDs to calibrate
+  basin_ids: ["12025000"]           # List of basin IDs to calibrate
   warmup_length: 365                # Warmup period in days
-  variables: ["precipitation", "potential_evapotranspiration", "streamflow"]  # Input variables
+  variables: ["prcp", "PET", "dayl", "srad", "tmax", "tmin", "vp", "streamflow"]  # Input variables
 
   # Time periods
-  train_period: ["1985-10-01", "1995-09-30"]   # Training period
-  valid_period: ["1995-10-01", "2005-09-30"]   # Validation period
-  test_period: ["2005-10-01", "2014-09-30"]    # Testing period
+  train_period: ["1981-01-01", "2004-12-31"]   # Training period
+  valid_period: ["2005-01-01", "2009-12-31"]   # Validation period
+  test_period: ["2010-01-01", "2014-12-31"]    # Testing period
 
   # Output configuration
-  output_dir: "results"             # Output directory for results
+  output_dir: "results/12025000"             # Output directory for results
 
 model:
   # Model configuration
-  name: "xaj"                    # Model type: xaj, xaj_mz, gr4j
+  name: "xaj"                    # Model type: xaj, xaj_mz, gr4j, gr5j, gr6j, etc.
   params:                           # Model parameters
     source_type: "sources"          # XAJ source type
     source_book: "HF"               # HF (Hydrological Forecast) computation method
     kernel_size: 15                 # Kernel size for Muskingum routing
 
+  # Output variable: "qsim" (streamflow), "es" (ET for XAJ), "et" (ET for GR)
+  output_variable: "qsim"    
+
 training:
   # Training configuration
   algorithm: "SCE_UA"                  # Calibration algorithm: SCE_UA (or sceua), GA (or genetic_algorithm), scipy (or scipy_minimize)
-
   # Algorithm-specific parameters (uncomment the section for your chosen algorithm)
 
   # SCE-UA (Shuffled Complex Evolution) parameters

docs/changelog.md

@@ -14,8 +14,8 @@
 
 - **Fixed Data Access Issues in Calibration**:
   - Fixed `AttributeError` when calibrating multi-basin flood event data
-  - Previous behavior: `UnifiedModelSetup` tried to access `p_and_e.shape[1]` when `p_and_e=None` for separate basin data
-  - Added proper detection of `basin_data_separate` mode in `UnifiedModelSetup.__init__()`
+  - Previous behavior: `UnifiedCalibrator` tried to access `p_and_e.shape[1]` when `p_and_e=None` for separate basin data
+  - Added proper detection of `basin_data_separate` mode in `UnifiedCalibrator.__init__()`
   - Modified `set_basin_index()` to extract basin-specific data from `basin_data_separate` when available
   - Each basin now uses its own time series without padding during calibration
   - Modified `unified_calibrate.py` (lines 116-253, 377-386)

docs/usage.md

@@ -225,7 +225,7 @@ calibrate()
   ↓
 2. UnifiedDataLoader.load_data()
   ↓
-3. Create UnifiedModelSetup (wraps MODEL_DICT)
+3. Create UnifiedCalibrator (wraps MODEL_DICT)
   ↓
 4. Select algorithm (SCE_UA, GA, scipy)
   ↓

hydromodel/configs/config_manager.py

@@ -448,6 +448,8 @@ def load_simplified_config(
         "model_cfgs": {
             "model_name": model_cfg["name"],
             **model_cfg.get("params", {}),
+            # Add output_variable configuration if specified
+            **({"output_variable": model_cfg["output_variable"]} if "output_variable" in model_cfg else {}),
         },
         "training_cfgs": {
             "algorithm_name": training_cfg["algorithm"],
@@ -462,6 +464,8 @@ def load_simplified_config(
             "experiment_name": f"{model_cfg['name']}_{training_cfg['algorithm']}",
             "random_seed": training_cfg.get("random_seed", 1234),
             "save_config": training_cfg.get("save_config", True),
+            # Add output_variable configuration if specified in training (highest priority)
+            **({"output_variable": training_cfg["output_variable"]} if "output_variable" in training_cfg else {}),
         },
         "evaluation_cfgs": {
             "metrics": eval_cfg["metrics"],

hydromodel/trainers/__init__.py

@@ -16,7 +16,7 @@
 from .unified_calibrate import (
     calibrate,
     ModelSetupBase,
-    UnifiedModelSetup,
+    UnifiedCalibrator,
     DEAP_AVAILABLE,
 )
 
@@ -31,7 +31,7 @@
     "calibrate",
     "evaluate",
     "ModelSetupBase",
-    "UnifiedModelSetup",
+    "UnifiedCalibrator",
     "UnifiedEvaluator",
     "DEAP_AVAILABLE",
 ]

hydromodel/trainers/unified_calibrate.py

@@ -75,7 +75,7 @@ def calculate_objective(
         pass
 
 
-class UnifiedModelSetup(ModelSetupBase):
+class UnifiedCalibrator(ModelSetupBase):
     """
     Completely unified setup for all hydrological models using the MODEL_DICT interface.
 
@@ -109,6 +109,20 @@ def __init__(
         self.training_config = training_config or {}
         self.warmup_length = warmup_length
 
+        # Get output variable configuration (default: qsim for streamflow)
+        # Priority: training_config > model_config > default
+        self.output_variable = (
+            self.training_config.get("output_variable") or
+            model_config.get("output_variable") or
+            "qsim"
+        )
+
+        # Validate output variable
+        valid_outputs = ["qsim", "es", "et", "ets"]
+        if self.output_variable not in valid_outputs:
+            print(f"Warning: output_variable '{self.output_variable}' not recognized. Using 'qsim'.")
+            self.output_variable = "qsim"
+
         # Load data using unified data loader
         self.data_loader = UnifiedDataLoader(data_config)
         self.p_and_e_full, qobs_full = self.data_loader.load_data()
@@ -275,14 +289,24 @@ def simulate(self, params: np.ndarray) -> np.ndarray:
             is_event_data=self.is_event_data,
         )
 
-        # Extract simulation output (qsim) - most calibration only needs this
-        # This avoids returning the full results dict every time
-        if isinstance(results, dict) and "qsim" in results:
-            return results["qsim"]
+        # Extract simulation output based on configured variable
+        # self.output_variable can be "qsim" (streamflow), "es"/"et"/"ets" (evapotranspiration)
+        if isinstance(results, dict):
+            # First try the configured output variable
+            if self.output_variable in results:
+                return results[self.output_variable]
+            # Fallback to qsim if configured variable not available
+            elif "qsim" in results:
+                if self.output_variable != "qsim":
+                    print(f"Warning: '{self.output_variable}' not in results, using 'qsim' instead")
+                return results["qsim"]
+            # Last resort: return first available result
+            else:
+                return list(results.values())[0]
         elif isinstance(results, np.ndarray):
             return results
         else:
-            # Fallback: return first available result
+            # Fallback for unexpected result types
             return (
                 list(results.values())[0]
                 if isinstance(results, dict)
@@ -368,7 +392,7 @@ def calibrate(config, **kwargs) -> Dict[str, Any]:
     os.makedirs(output_dir, exist_ok=True)
 
     # Create unified model setup
-    model_setup = UnifiedModelSetup(
+    model_setup = UnifiedCalibrator(
         data_config=data_config,
         model_config=model_config,
         loss_config=loss_config,
@@ -506,7 +530,7 @@ def _save_calibration_config(
 
 
 def _calibrate_model(
-    model_setup: UnifiedModelSetup,
+    model_setup: UnifiedCalibrator,
     algorithm_config: Dict,
     output_dir: str,
     basin_id: str,
@@ -799,7 +823,7 @@ class UnifiedSpotSetup(SpotSetup):
         def __init__(self, model_setup, total_iterations):
             # Do not call super().__init__ to avoid reloading or storing redundant data
             self.model_setup = model_setup
-            # Reuse parameter definitions created by UnifiedModelSetup
+            # Reuse parameter definitions created by UnifiedCalibrator
             self.params = model_setup.params
             self.parameter_names = model_setup.parameter_names
             # Progress tracking