diff --git a/README.md b/README.md index d7b1712..3fa6377 100644 --- a/README.md +++ b/README.md @@ -133,21 +133,23 @@ from nwm_fcst_mgr.forecast import run_hindcast run_hindcast( valid_yaml='/path/to/valid.yaml', - input_path='/path/to/input.config', + config='/path/to/input.config', fcst_run_name='my_hindcast_run', cycle_interval=3, num_iterations=10, - cold_start_state='/path/to/cold_start_state/' + cold_start_state='/path/to/cold_start_state/', + yield_realizations=False, ) ``` #### Arguments - `valid_yaml` - Path to validation yaml file from previous calibration run (from nwm-cal-mgr) -- `input_path` - Path to forecast input configuration file (from nwm-msw-mgr) +- `config` - Path to forecast input configuration file (as str), or instance of InputConfig (both from nwm-msw-mgr) - `my_hindcast_run` - Name for the hindcast run folder - `cycle_interval` - Cycle interval in hours (spacing between hindcast cycles) - `num_iterations` - Number of hindcast cycles to perform -- `--cold_start_state` - (Optional) Path to cold start state to initialize hindcasting workflow +- `cold_start_state` - (Optional) Path to cold start state to initialize hindcasting workflow +- `yield_realizations` - (Optional) Default False. If True, then this function will act as a generator and will yield each RealizationBuilder instance after constructing it and calling its build_fcst_realization() method. If False, this function itself will execute each ngen realization of the hindcast sequence as they become built. #### Hindcast Example diff --git a/python/nwm_fcst_mgr/forecast.py b/python/nwm_fcst_mgr/forecast.py index 66a901c..3d4bc54 100644 --- a/python/nwm_fcst_mgr/forecast.py +++ b/python/nwm_fcst_mgr/forecast.py @@ -12,6 +12,7 @@ from datetime import datetime, timedelta from enum import Enum, auto from pathlib import Path +from typing import Generator import geopandas as gpd import matplotlib.pyplot as plt @@ -20,7 +21,8 @@ import yaml from ewts import Payload, Status from ewts.modules import ModuleKey -from mswm.manager import build_fcst +from mswm.manager import RealizationBuilder +from mswm.utils.input_configuration import InputConfig from nwm_fcst_mgr.consts import PARTITION_CONFIG_FILE_NAME_SUFFIX from nwm_fcst_mgr.exceptions import ( @@ -31,8 +33,8 @@ from nwm_fcst_mgr.utils import ( OS_ENV_KEY_NGEN_LOG_FILE_PREFIX, OS_ENV_KEY_RESULTS_DIR, - initialize_logger, initialize_hindcast_logger, + initialize_logger, set_os_env_key, ) @@ -316,7 +318,7 @@ def poll_ngen_flush_log(self) -> None: self._stop_ngen() self._check_process_returncode() - def preprocess(self) -> None: + def preprocess(self, do_override_log_file_prefix: bool = False) -> None: """Preprocess an ngen run, validate some inputs, and set the execution status.""" # Use cached config values @@ -331,7 +333,7 @@ def preprocess(self) -> None: OS_ENV_KEY_RESULTS_DIR, str(self.out_dir), override=False ) set_os_env_key( - OS_ENV_KEY_NGEN_LOG_FILE_PREFIX, self.out_dir.name, override=False + OS_ENV_KEY_NGEN_LOG_FILE_PREFIX, self.out_dir.name, override=do_override_log_file_prefix ) self._status = RunStatus.PREPROCESSED @@ -673,19 +675,26 @@ def read_troute_output( return output -def check_hind_intervals(input_path: str, hind_interval: list) -> None: +def check_hind_intervals(config: str | InputConfig, hind_interval: list) -> None: """ Check that all hindcast intervals fall on valid cycle hours for a given forcing configuration Parameters ---------- - input_path: str - Path to input.config file + config: str | InputConfig + Path to input.config file or an instance of InputConfig hind_interval: list List of hindcast intervals in hours """ # Load config file - config = load_config(input_path) + if isinstance(config, str): + config = load_config(config) + elif isinstance(config, InputConfig): + config = config.model_dump() + else: + msg = f"Expected config to be either str or InputConfig, but got {type(config)}" + logger.critical(msg) + raise TypeError(msg) # Read values from config file try: @@ -756,15 +765,30 @@ def run_forecast( logger.info("Ngen run completed") -def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iterations, cold_start_state=None): +def run_hindcast( + config: str | InputConfig, + valid_yaml, + fcst_run_name, + cycle_interval, + num_iterations, + cold_start_state=None, + yield_realizations: bool = False, + ) -> None | Generator[RealizationBuilder, None, None]: """ Run hindcast workflow with warm start runs, initial cold start should be run separately - Accepts cycle interval and number of intervals for repeated hindcasts + Accepts cycle interval and number of intervals for repeated hindcasts. + + If yield_realizations is True, then this function acts as a generator of (built) RealizationBuilder + instances, with the assumption that the caller will execute each realization as it is generated, + before the next one is generated. + + If yield_realizations is False, then this function builds and runs the realizations sequence itself + (which takes significant time to return). Parameters --------- - input_path : str - Path to input.config file for hindcast + config : str | InputConfig + Path to input.config file, or an instance of InputConfig valid_yaml : str Path to validation yaml file from previous run of nwm-cal-mgr fcst_run_name : str @@ -777,11 +801,14 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter Path to directory containing state files to load at start of first hindcast If provided, will be used for first hindcast cycle (hind_cycle=0) Subsequent cycles will use warm start states + yield_realizations: bool + If True, then this function will act as a generator and will yield each RealizationBuilder + instance after constructing it and calling its build_fcst_realization() method. """ # Set up hindcast orchestration logger, initialized once the hindcast root directory is known hindcast_logger = None - # Buffer messages logged before orchestration log's directory is resolvable (before first build_fcst() call) + # Buffer messages logged before orchestration log's directory is resolvable (before first rb.build_fcst_realization() call) pending_logs = [f'Initializing hindcast runs from: {valid_yaml}'] if valid_yaml is None: @@ -796,7 +823,7 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter hind_interval = list(range(0, num_iterations * cycle_interval, cycle_interval)) # Validate that all hindcast intervals fall on valid cycle hours for this configuration - check_hind_intervals(input_path, hind_interval) + check_hind_intervals(config, hind_interval) pending_logs.append(f"Initializing hindcast runs at intervals: {hind_interval}") @@ -806,23 +833,47 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter # Initialize previous state to be loaded for warm start prev_warm_start_state = cold_start_state + hind_kwargs_base = { + # vvv One of these is replaced later depending on the type of ``config``. + "input_path": None, + "config_overrides": None, + # ^^^ One of these is replaced later depending on the type of ``config``. + "valid_yaml": valid_yaml, + "fcst_run_name": fcst_run_name, + } + if isinstance(config, str): + hind_kwargs_base["input_path"] = config + elif isinstance(config, InputConfig): + hind_kwargs_base["config_overrides"] = config + else: + msg = f"Expected config to be either str (path to config file) or InputConfig instance, but got {type(config)}" + logger.critical(msg) + raise TypeError(msg) + # Loop through hindcast intervals for hind_cycle in hind_interval: # Skip warm start for first hindcast, which will use the cold start state if hind_cycle != 0: - hindcast_logger.info(f"Initializing warm start AnA run for hindcast iteration at {hind_cycle} hours") + warmstart_kwargs = hind_kwargs_base | { + "use_warm_start": True, + "hind_cycle": hind_cycle, + "prev_hind_cycle": prev_hind_cycle, + "save_state": True, + "load_state_from": prev_warm_start_state, + } # Generate msw-mgr inputs for warm start run for hindcast iteration - warm_start_real_path, warm_start_state = build_fcst(input_path=input_path, valid_yaml=valid_yaml, - fcst_run_name=fcst_run_name, use_warm_start=True, - hind_cycle=hind_cycle, prev_hind_cycle=prev_hind_cycle, - save_state=True, load_state_from=prev_warm_start_state) + rb = RealizationBuilder(**warmstart_kwargs) + warm_start_real_path, warm_start_state = rb.build_fcst_realization() hindcast_logger.info(f"Warm start realization file for hindcast iteration at {hind_cycle} hours written to: {warm_start_real_path}") # Execute warm start ngen run to generate hindcasting model states - run_workflow(warm_start_real_path, config_cache, suppress_output=True) + if yield_realizations: + yield rb + else: + run_workflow(warm_start_real_path, config_cache, suppress_output=True) hindcast_logger.info(f"Warm start run for hindcast iteration at {hind_cycle} hours completed") hindcast_logger.info(f"Warm start state saved to {warm_start_state}") @@ -830,10 +881,7 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter prev_warm_start_state = warm_start_state # Create hindcast input files - hind_kwargs = { - 'input_path': input_path, - 'valid_yaml': valid_yaml, - 'fcst_run_name': fcst_run_name, + hind_kwargs = hind_kwargs_base | { 'use_hindcast': True, 'hind_cycle': hind_cycle } @@ -854,9 +902,10 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter hind_kwargs['load_state_from'] = warm_start_state hindcast_logger.info(f"Hindcast iteration at {hind_cycle} hours loading state from: {warm_start_state}") - hind_real_path, _ = build_fcst(**hind_kwargs) + rb = RealizationBuilder(**hind_kwargs) + hind_real_path, _ = rb.build_fcst_realization() - # Initialize hindcast orchestration logger after first build_fcst() call; hindcast root directory now resolvable + # Initialize hindcast orchestration logger after first rb.build_fcst_realization() call; hindcast root directory now resolvable if hindcast_logger is None: hindcast_root = Path(hind_real_path).parent.parent hindcast_logger = initialize_hindcast_logger(str(hindcast_root)) @@ -867,7 +916,10 @@ def run_hindcast(input_path, valid_yaml, fcst_run_name, cycle_interval, num_iter hindcast_logger.info(f'Hindcast realization file for iteration at {hind_cycle} hours written to: {hind_real_path}') # Run hindcasting period - run_workflow(hind_real_path, config_cache) + if yield_realizations: + yield rb + else: + run_workflow(hind_real_path, config_cache) hindcast_logger.info(f"Hindcast run for iteration at {hind_cycle} hours completed") # Store previous hindcast cycle value to set next warm start duration @@ -910,7 +962,7 @@ def main(): if args.command == "run_forecast": run_forecast(real_path=args.real_path, valid_yaml=args.valid_yaml, no_valid=args.no_valid, partition_file=args.partition_file) elif args.command == "run_hindcast": - run_hindcast(valid_yaml=args.valid_yaml, input_path=args.input_path, + run_hindcast(valid_yaml=args.valid_yaml, config=args.input_path, fcst_run_name=args.fcst_run_name, cycle_interval=args.cycle_interval, num_iterations=args.num_iterations, cold_start_state=args.cold_start_state) else: