From fd67437df22801faf9322a66a24d4d56ce66b5f5 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 18 Sep 2026 12:18:54 -0400 Subject: [PATCH 1/2] Mirror coastal gridded forcing changes for PW, using command `git diff --binary origin/development..origin/coastalforcing | git apply --3way --index` --- .../NextGen_Forcings_Engine/bmi_model.py | 48 ++- .../NextGen_Forcings_Engine/core/config.py | 51 +++- .../NextGen_Forcings_Engine/core/consts.py | 22 +- .../NextGen_Forcings_Engine/core/geoMod.py | 72 ++++- .../NextGen_Forcings_Engine/core/parallel.py | 9 +- .../NextGen_Forcings_Engine/core/regrid.py | 44 ++- NextGen_Forcings_Engine_BMI/README.md | 4 + .../coastal/__init__.py | 1 + NextGen_Forcings_Engine_BMI/coastal/consts.py | 73 +++++ .../coastal/gridded/__init__.py | 1 + .../coastal/gridded/config_manager.py | 191 ++++++++++++ .../coastal/gridded/warm_regrid_weights.py | 287 ++++++++++++++++++ NextGen_Forcings_Engine_BMI/run_bmi_model.py | 287 ++++++++++++++++-- pyproject.toml | 4 +- 14 files changed, 1024 insertions(+), 70 deletions(-) create mode 100644 NextGen_Forcings_Engine_BMI/coastal/__init__.py create mode 100644 NextGen_Forcings_Engine_BMI/coastal/consts.py create mode 100644 NextGen_Forcings_Engine_BMI/coastal/gridded/__init__.py create mode 100644 NextGen_Forcings_Engine_BMI/coastal/gridded/config_manager.py create mode 100644 NextGen_Forcings_Engine_BMI/coastal/gridded/warm_regrid_weights.py diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py index 895ac14c..afff96bd 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -140,6 +140,9 @@ class NWMv3_Forcing_Engine_BMI_model_Base(Bmi): It includes methods for initializing the model, updating it, accessing model variables, and managing model configuration. This class is responsible for interacting with geospatial data and forcing inputs for the model simulation. + + For the init arguments that are optional, they act as overrides on the equivalent + values set in the provided configuration file (provided when calling ``initialize()``). """ def __init__( @@ -147,14 +150,29 @@ def __init__( b_date: str = None, geogrid: str = None, output_path: str = None, + output_steps: int = None, ) -> None: """Create a model that is ready for initialization. Initializes the model with default values for time, variables, and grid types. + + Args: + b_date: Optional processing-cycle timestamp in ``YYYYMMDDHHMM`` format. + Overrides the ``RefcstBDateProc`` configuration value. + geogrid: Optional path to the target geogrid. + Overrides the ``GeogridIn`` configuration value. + output_path: Optional path for gridded forcing output. + Overrides the output location otherwise derived from ``ScratchDir``. + output_steps: Optional positive override for the number of output steps. + When not provided, the output-step count is derived dynamically from + the forcing configuration's timing settings. This must be ``None`` + for AnA aka Analysis & Assimilation runs, whose output count is + controlled by their configured lookback window. """ self.output_path = output_path self._geogrid = geogrid self._b_date = b_date + self._output_steps = output_steps self._values = {} self._start_time = 0.0 @@ -236,7 +254,10 @@ def _job_meta(self, value: ConfigOptions) -> None: if value is None: try: value = ConfigOptions( - self.cfg_bmi, b_date=self._b_date, geogrid=self._geogrid + self.cfg_bmi, + b_date=self._b_date, + geogrid=self._geogrid, + output_steps=self._output_steps, ) except KeyboardInterrupt as e: err_handler.err_out_screen("User keyboard interrupt", e) @@ -273,10 +294,12 @@ def geo_meta(self) -> GeoMeta: https://github.com/NGWPC/ngen-forcing/pull/212 -- Coastal Forcing """ if self._geo_meta is None: - assert self._job_meta.grid_type == "hydrofabric", ( - f"Only 'hydrofabric' grid type is currently supported; got '{self._job_meta.grid_type}'. See docstrings for discretization types." + assert self._job_meta.grid_type in ["gridded", "hydrofabric"], ( + f"Only 'gridded' and 'hydrofabric' grid types are currently supported. Got '{self._job_meta.grid_type}'. See docstrings for discretization types." + ) + self._geo_meta = GeoMeta.for_grid_type( + self._job_meta.grid_type, self._job_meta, self._mpi_meta ) - self._geo_meta = HydrofabricGeoMeta(self._job_meta, self._mpi_meta) return self._geo_meta @geo_meta.setter @@ -303,7 +326,17 @@ def init_scratch_dir(self) -> None: def create_esmf_mesh(self) -> None: """Create the ESMF mesh for the model and set ``self._cat_ids`` (later used as BMI variable "CAT-ID").""" if self._mpi_meta.rank == 0: - cat_ids = esmf_creation.create_mesh(self._job_meta) + if self._job_meta.grid_type == "gridded": + # Gridded output regrids directly onto the target grid and has no + # catchments/hydrofabric divides -- ESMF mesh creation (which + # requires a hydrofabric geopackage) does not apply here. + cat_ids = np.array([], dtype=np.int64) + elif self._job_meta.grid_type == "hydrofabric": + cat_ids = esmf_creation.create_mesh(self._job_meta) + else: + raise ValueError( + f"Unsupported grid_type {self._job_meta.grid_type} for this function" + ) cat_count = np.array( [len(cat_ids) if self._mpi_meta.rank == 0 else 0], dtype=np.intc ) @@ -501,7 +534,7 @@ def _configure_output_path(self) -> None: ) self._output_obj.init_forcing_file( - self._job_meta, self.geo_meta, self._mpi_meta, self._values["CAT-ID"] + self._job_meta, self.geo_meta, self._mpi_meta, self._cat_ids ) self._output_configured = True @@ -1563,12 +1596,13 @@ def __init__( b_date: str = None, geogrid: str = None, output_path: str = None, + output_steps: int = None, ): """Create a model that is ready for initialization. Initializes the model with default values for time, variables, and grid types. """ - super().__init__(b_date, geogrid, output_path) + super().__init__(b_date, geogrid, output_path, output_steps) def grid_ranks(self) -> list[int]: """Get the grid ranks for the gridded domain.""" diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py index a6c302ab..70a111c0 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py @@ -31,7 +31,13 @@ class ConfigOptions: """Configuration abstract class for configuration options read in from the file specified by the user.""" - def __init__(self, cfg_bmi: dict, b_date: str = None, geogrid: str = None) -> None: + def __init__( + self, + cfg_bmi: dict, + b_date: str = None, + geogrid: str = None, + output_steps: int = None, + ) -> None: """Initialize the configuration class to empty None attributes. The attributes of this class are populated by the validate_config function, which reads in the configuration file and checks that all necessary options are provided and properly formatted. The attributes of this class are used to control the flow of the program and the processing of input forcings. @@ -40,12 +46,18 @@ def __init__(self, cfg_bmi: dict, b_date: str = None, geogrid: str = None) -> No cfg_bmi (dict): The configuration dictionary read in from the configuration file specified by the user. This should be read in using the config_utils.read_config function, which also handles any necessary preprocessing of the configuration file. b_date (str, optional): The beginning date of processing in the format YYYYMMDDHHMM. This is used to calculate the processing window for realtime simulations. If not provided, it will be read from the configuration file. geogrid (str, optional): The filepath to the geogrid file to be used for processing. This is used to specify the grid information for regridding input forcings. If not provided, it will be read from the configuration file. + output_steps (int, optional): Override the configured output count. When set, + this must be at least one and takes precedence over analysis and + forecast-cycle output-step calculations. """ - if geogrid is not None: - self.user_provided_geogrid_flag = True - else: - self.user_provided_geogrid_flag = False + self._b_date_proc = None + self._geogrid = None + self.user_provided_geogrid_flag = geogrid is not None + + if output_steps is not None and output_steps < 1: + raise ValueError(f"output_steps must be at least 1, but got {output_steps}") + self.output_steps = output_steps # If b_date not provided, try to read from config file if b_date is None: @@ -58,6 +70,11 @@ def __init__(self, cfg_bmi: dict, b_date: str = None, geogrid: str = None) -> No self.b_date_proc = b_date self.cfg_bmi = cfg_bmi self.geogrid = geogrid + self.reuse_regrid_weights = cfg_bmi.get("ReuseRegridWeights", False) + if not isinstance(self.reuse_regrid_weights, bool): + raise TypeError( + f"ReuseRegridWeights must be a boolean, but got: {self.reuse_regrid_weights}" + ) self.bmi_time_index = 0 self.globalNdv = -9999.0 @@ -75,7 +92,6 @@ def __init__(self, cfg_bmi: dict, b_date: str = None, geogrid: str = None) -> No # These must exist (as None) before the properties are accessed self._supp_precip_forcings = None - self._b_date_proc = None self._input_forcings = None self._nwm_geogrid = None self._output_freq = None @@ -89,7 +105,6 @@ def __init__(self, cfg_bmi: dict, b_date: str = None, geogrid: str = None) -> No self._fcst_input_horizons = None self._spatial_meta = None self._geopackage = None - self._geogrid = None self._grid_type = None # Backing vars for setters that guard on precip_only_flag and do not unconditionally assign self._fcst_input_offsets = None @@ -1203,13 +1218,31 @@ def actual_output_steps(self) -> int: whether the user has chosen to run a reforecast simulation with a specified processing window, which will only output time steps for which input forcings are available based on the processing window and forecast time horizons - specified by the user in the configuration file. - """ + specified by the user in the configuration file. An explicit ``output_steps`` + override takes precedence over forecast-cycle calculations, but is not + compatible with analysis configurations, whose output count is determined by + their lookback window. + """ + if self.output_steps is not None: + if self.ana_flag: + raise ValueError( + "output_steps is not supported for analysis configurations" + ) + return np.int32(self.output_steps) if self.ana_flag: return np.int32(self.nFcsts) else: return np.int32(self.num_output_steps) + @property + def should_randomize_weight_file_name(self) -> bool: + """Boolean, if resolves to True then the intermediary regridding weights file + will include a random string. This is used so that "gridded" workflows, for + example those used to produce coastal forcing inputs, can reuse regrid weights.""" + if self.grid_type == "gridded" and self.reuse_regrid_weights: + return False + return True + @property def grid_type(self) -> str: """Tells the NextGen Forcings Engine BMI which grid type the engine is diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py index cc99f943..4898cde3 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py @@ -9,6 +9,26 @@ weighted_average, ) +# These are the variable names as stored in the input geogrid files, +# which need to be read when using the "gridded" grid_type aka discretization type. +GEOGRID_FILE_INPUT_VARIABLE_NAME_ATTRS = ( + "lat_var", + "lon_var", + "hgt_var", + "cosalpha_var", + "sinalpha_var", + "slope_var", + "slope_azimuth_var", + "slope_var_elem", + "slope_azimuth_var_elem", + "nodecoords_var", + "elemcoords_var", + "elemconn_var", + "numelemconn_var", + "element_id_var", + "hgt_elem_var", +) + BMI_MODEL = { "NWMv3_Forcing_Engine_BMI_model_Base": [ "_model", @@ -139,8 +159,6 @@ "nodeCoords", "centerCoords", "inds", - "esmf_lat", - "esmf_lon", ], "UnstructuredGeoMeta": [ "x_lower_bound", diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py index 7f6dd909..91822785 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -24,7 +24,10 @@ from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.config import ( ConfigOptions, ) -from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.consts import GEOMOD +from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.consts import ( + GEOGRID_FILE_INPUT_VARIABLE_NAME_ATTRS, + GEOMOD, +) from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.err_handler import ( log_critical, ) @@ -92,7 +95,14 @@ def wrapper(self) -> Any: assert isinstance(post_slice, bool) assert isinstance(name, str) assert isinstance(config_options, ConfigOptions) - assert isinstance(var, np.ndarray) + # var is only populated on rank 0 -- every other rank gets its + # share via scatter_array() below, which is None-safe on + # non-zero ranks (confirmed in parallel.py). Asserting here + # unconditionally rejected every rank but 0. + if self.mpi_config.rank == 0: + assert isinstance(var, np.ndarray) + else: + assert isinstance(var, (np.ndarray, type(None))) var = self.mpi_config.scatter_array(self, var, config_options) if post_slice: @@ -123,6 +133,24 @@ def __init__(self, config_options: ConfigOptions, mpi_config: MpiConfig) -> None for attr in GEOMOD[__class__.__name__]: setattr(self, attr, None) + @classmethod + def for_grid_type( + cls, grid_type: str, config_options: ConfigOptions, mpi_config: MpiConfig + ) -> GeoMeta: + """Create the geometry metadata implementation for a grid type.""" + lookup = { + "gridded": GriddedGeoMeta, + "unstructured": UnstructuredGeoMeta, + "hydrofabric": HydrofabricGeoMeta, + } + try: + _class = lookup[grid_type] + except KeyError as error: + raise ValueError( + f"Unsupported grid type: {grid_type}, expected one of: {list(lookup.keys())}" + ) from error + return _class(config_options, mpi_config) + @cached_property def spatial_metadata_exists(self) -> bool: """Check to make sure the geospatial metadata file exists in the config_options.""" @@ -133,10 +161,22 @@ def spatial_metadata_exists(self) -> bool: @cached_property def geogrid_ds(self) -> xr.Dataset: - """Open the geogrid file and return the xarray dataset object.""" + """Open the geogrid file and load only the variables needed by this run. + + geo_em_CONUS.nc is ~8.9 GB; loading the whole file with ds.load() exhausts + RAM on standard instances. We open lazily, select only the variables + referenced by config_options, then load that subset into memory. + Global attributes (DX, DY, etc.) are preserved on the subset dataset. + """ try: with xr.open_dataset(self.config_options.geogrid) as ds: - return ds.load() + needed = [ + getattr(self.config_options, attr) + for attr in GEOGRID_FILE_INPUT_VARIABLE_NAME_ATTRS + if getattr(self.config_options, attr, None) is not None + and getattr(self.config_options, attr) in ds + ] + return ds[needed].load() except Exception as e: self.config_options.errMsg = "Unable to open geogrid file with xarray" log_critical(self.config_options, self.mpi_config) @@ -282,12 +322,13 @@ def nx_global(self) -> int: if self.mpi_config.rank == 0: try: if self.ndim_lat == 3: - return self.lat_var.shape[2] + nx = self.lat_var.shape[2] elif self.ndim_lat == 2: - return self.lat_var.shape[1] + nx = self.lat_var.shape[1] else: # NOTE Is this correct? using lon_var - return self.lon_var.shape[0] + nx = self.lon_var.shape[0] + return nx except Exception as e: self.config_options.errMsg = f"Unable to extract X dimension size from {self.config_options.lon_var} in: {self.config_options.geogrid}" log_critical(self.config_options, self.mpi_config) @@ -365,7 +406,7 @@ def dx_meters(self) -> float: @cached_property def esmf_grid(self) -> ESMF.Grid: """Create the ESMF grid object for the gridded domain.""" - return esmf_grid_retry( + esmf_grid = esmf_grid_retry( self.mpi_config, self.config_options, err_handler, @@ -373,6 +414,13 @@ def esmf_grid(self) -> ESMF.Grid: staggerloc=ESMF.StaggerLoc.CENTER, coord_sys=ESMF.CoordSys.SPH_DEG, ) + # NOTE Populating the destination coordinates before creating fields or regridding. + # Coordinate scattering reads this grid's local bounds. Need to cache the grid + # first so that lookup does not re-enter this cached property (avoid infinite recursion) + self.esmf_grid = esmf_grid + esmf_grid.get_coords(1)[:, :] = self.latitude_grid + esmf_grid.get_coords(0)[:, :] = self.longitude_grid + return esmf_grid @cached_property def esmf_lat(self) -> np.ndarray: @@ -395,9 +443,9 @@ def latitude_grid(self) -> np.ndarray: # Scatter global XLAT_M grid to processors.. if self.mpi_config.rank == 0: if self.ndim_lat == 3: - var_tmp = self.lat_var[0, :, :] + var_tmp = np.asarray(self.lat_var[0, :, :]) elif self.ndim_lat == 2: - var_tmp = self.lat_var[:, :] + var_tmp = np.asarray(self.lat_var[:, :]) elif self.ndim_lat == 1: lat = self.lat_var[:] lon = self.lon_var[:] @@ -429,9 +477,9 @@ def longitude_grid(self) -> np.ndarray: if ( self.ndim_lat == 3 ): # NOTE The original code has lat here... should it maybe be lon instead? - var_tmp = self.lon_var[0, :, :] + var_tmp = np.asarray(self.lon_var[0, :, :]) elif self.ndim_lon == 2: - var_tmp = self.lon_var[:, :] + var_tmp = np.asarray(self.lon_var[:, :]) elif self.ndim_lon == 1: lat = self.lat_var[:] lon = self.lon_var[:] diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py index e221758a..b12ef445 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -161,11 +161,18 @@ def _cleanup_scratch_dir(self) -> None: self.try_remove_empty_dir_no_reraise(self.config_options.scratch_dir) def _cleanup_geogrid(self) -> None: - """Remove temporary geogrid file if it exists.""" + """Remove engine-created geogrids while preserving caller-supplied files. + + Coastal forcing workflows provide their target geogrid to the BMI, so it + must remain available after finalization. + """ self.log_debug("Cleanup: starting geogrid cleanup") if self.config_options is None: self.log_debug("Cleanup: config_options is not set") return + if self.config_options.user_provided_geogrid_flag: + self.log_debug("Cleanup: preserving caller-supplied geogrid") + return geogrid = getattr(self.config_options, "geogrid", None) if geogrid is not None: self.try_delete_file_no_reraise(geogrid) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py index fa751dc9..407e8241 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -11638,7 +11638,8 @@ def get_weight_file_names( grid_key = input_forcings.product_name file_key = f"{grid_key}_{config_options.geogrid}" hash_key = hashlib.md5(file_key.encode()).hexdigest()[:8] - hash_key += f"_{mpi_config.uid64}" + if config_options.should_randomize_weight_file_name: + hash_key += f"_{mpi_config.uid64}" weight_file = os.path.join(config_options.weightsDir, f"ESMF_weight_{hash_key}.nc4") @@ -12020,6 +12021,10 @@ def calculate_weights( err_handler.check_program_status(config_options, mpi_config) + if mpi_config.rank == 0 and lat_tmp is not None and lon_tmp is not None: + # Normalize source longitudes to [-180, 180) to match geo_em geogrids. + lon_tmp = ((lon_tmp + 180) % 360) - 180 + # Scatter global GFS latitude grid to processors.. if mpi_config.rank == 0: var_tmp = lat_tmp @@ -12399,13 +12404,38 @@ def calculate_supp_pcp_weights( supplemental_precip.esmf_field_in.data[:] = var_sub_tmp # mpi_config.comm.barrier() - supplemental_precip.regridObj = pt.esmf_regrid_retry_partial( - supplemental_precip.esmf_field_in, - supplemental_precip.esmf_field_out, - src_mask_values=np.array([0]), - regrid_method=ESMF.RegridMethod.BILINEAR, - unmapped_action=ESMF.UnmappedAction.IGNORE, + # Reuse a cached weight file if one exists (see get_weight_file_names / + # load_weight_file, also used by calculate_weights() for InputForcings) -- + # this code path previously always recomputed weights in memory, with no + # on-disk caching, unlike the InputForcings path. + weight_file, _ = get_weight_file_names( + mpi_config, config_options, supplemental_precip ) + if config_options.weightsDir is not None: + if not os.path.exists(weight_file): + supplemental_precip.regridObj = pt.esmf_regrid_retry_partial( + supplemental_precip.esmf_field_in, + supplemental_precip.esmf_field_out, + src_mask_values=np.array([0]), + regrid_method=ESMF.RegridMethod.BILINEAR, + unmapped_action=ESMF.UnmappedAction.IGNORE, + filename=weight_file, + ) + load_weight_file( + mpi_config, + config_options, + supplemental_precip, + weight_file, + element_mode=False, + ) + else: + supplemental_precip.regridObj = pt.esmf_regrid_retry_partial( + supplemental_precip.esmf_field_in, + supplemental_precip.esmf_field_out, + src_mask_values=np.array([0]), + regrid_method=ESMF.RegridMethod.BILINEAR, + unmapped_action=ESMF.UnmappedAction.IGNORE, + ) # Run the regridding object on this test dataset. Check the output grid for # any 0 values. diff --git a/NextGen_Forcings_Engine_BMI/README.md b/NextGen_Forcings_Engine_BMI/README.md index 96a0357f..58d373fe 100644 --- a/NextGen_Forcings_Engine_BMI/README.md +++ b/NextGen_Forcings_Engine_BMI/README.md @@ -57,6 +57,10 @@ 2. Within the “NextGen_Forcings_Engine_BMI” directory, there is a sub-directory called “BMI_NextGen_Configs” that contains all the BMI configuration files needed for a Medium range forecast ran by GFS that would support a gridded model (./NextGen_Forcings_Engine_BMI/BMI_NextGen_Configs/Medium_Range/Gridded/config.yml), an unstructured mesh coastal model (./NextGen_Forcings_Engine_BMI/BMI_NextGen_Configs/Medium_Range/gridded/config.yml), and a hydrofabric unstructured mesh ((./NextGen_Forcings_Engine_BMI/BMI_NextGen_Configs/Medium_Range/hydrofabric/config.yml) as a supported domain to extract regridded forcings in a BMI-complaint fashion. Copy over one of those config.yml files to the “NextGen_Forcings_Engine_BMI” main directory for BMI execution of the NextGen Forcings Engine. Inside each config.yml contains a set of standard NWMv3.0 Forcing Engine variables that allow a given user to utilize a variety of methods, which are further described in the “./NextGen_Forcings_Engine_BMI/BMI_NextGen_Configs/README.md” file. 3. To test out the NextGen Forcings Engine BMI functionality, please see the “./NextGen_Forcings_Engine_BMI/README.md” file that highlights the utility of each of the BMI Python scripts as well as supporting sub-directories contained within the repository. +## Geogrid cleanup + +The standalone BMI runner accepts a caller-owned geogrid with `-geogrid /path/to/geogrid.nc`; it forwards this path to the BMI model constructor. During finalization, the engine preserves a geogrid supplied this way. This supports coastal forcing generation, where the caller-owned target geogrid may be reused after a run. When `-geogrid` is omitted, the engine uses `GeogridIn` from the forcing configuration, creates its per-run copy, and removes that temporary file during cleanup. + # Overview bullet points for modifying the original NWM Forcings Engine into a BMI complaint NextGen Forcings Engine capable of handling any domain types • To streamline the NWMv3.0 Forcings engine into a Basic Model Interface application, we’ve had to initialize the BMI model using the same approach highlighted in the “genForcing.py” module and then directly reconfigure the forecast module (forecastMod.py) workflow to streamline the ability to update and produce gridded forcings for the WRFHydro domain based on a specified time stamp within the standard BMI functionality (“model_update_until”). This is all completed within the “model.py” module, which essentially mimics the “forecastMod.py” module within the “core” directory as a BMI-compliant module. diff --git a/NextGen_Forcings_Engine_BMI/coastal/__init__.py b/NextGen_Forcings_Engine_BMI/coastal/__init__.py new file mode 100644 index 00000000..9a6c07f3 --- /dev/null +++ b/NextGen_Forcings_Engine_BMI/coastal/__init__.py @@ -0,0 +1 @@ +"""Coastal forcing utilities.""" diff --git a/NextGen_Forcings_Engine_BMI/coastal/consts.py b/NextGen_Forcings_Engine_BMI/coastal/consts.py new file mode 100644 index 00000000..e4693832 --- /dev/null +++ b/NextGen_Forcings_Engine_BMI/coastal/consts.py @@ -0,0 +1,73 @@ +"""Constants shared by coastal forcing utilities.""" + +from pathlib import Path + +ANALYSIS_FLAG = "AnAFlag" +FORECAST_INPUT_HORIZONS = "ForecastInputHorizons" +REFERENCE_FORECAST_DATE = "RefcstBDateProc" +LOOKBACK = "LookBack" +GRID_TYPE = "GRID_TYPE" +GRIDDED_GRID_TYPE = "gridded" +OUTPUT = "Output" +GEOPACKAGE = "Geopackage" +GEOGRID_INPUT = "GeogridIn" +SPATIAL_METADATA_INPUT = "SpatialMetaIn" +LONGITUDE_VARIABLE = "LONVAR" +LATITUDE_VARIABLE = "LATVAR" +REUSE_REGRID_WEIGHTS = "ReuseRegridWeights" +REGRID_WEIGHTS_DIR = "RegridWeightsDir" +SCRATCH_DIRECTORY = "ScratchDir" +INPUT_FORCINGS = "InputForcings" +SUPPLEMENTAL_PRECIPITATION = "SuppPcp" + +GRID_LONGITUDE_VARIABLE = "XLONG_M" +GRID_LATITUDE_VARIABLE = "XLAT_M" +DEFAULT_FORCING_ROOT_DIR = Path("/ngen-app/data") +DEFAULT_FORCING_TEMPLATE_DIR = ( + Path(__file__).resolve().parents[1] / "BMI_NextGen_Configs" / "config_templates" +) +DEFAULT_GEOGRID_DIR = DEFAULT_FORCING_ROOT_DIR / "esmf_mesh" / "NWM" / "domain" +DEFAULT_REGRID_WEIGHTS_DIR = DEFAULT_FORCING_ROOT_DIR / "esmf_mesh" / "regrid_weights" +DEBUG_CONUS_SUBSET = "CONUS_debug_gauge_01123000" +SPATIAL_METADATA_SUFFIX = { + "CONUS": "CONUS", + "Alaska": "AK", + "Hawaii": "HI", + "Puerto_Rico": "PRVI", +} + +INPUT_PRODUCT_FIELDS = ( + INPUT_FORCINGS, + "InputForcingDirectories", + "InputForcingTypes", + "InputMandatory", + "ForecastInputHorizons", + "ForecastInputOffsets", + "IgnoredBorderWidths", + "RegridOpt", + "ForcingTemporalInterpolation", + "TemperatureBiasCorrection", + "PressureBiasCorrection", + "HumidityBiasCorrection", + "WindBiasCorrection", + "SwBiasCorrection", + "LwBiasCorrection", + "PrecipBiasCorrection", + "TemperatureDownscaling", + "ShortwaveDownscaling", + "PressureDownscaling", + "PrecipDownscaling", + "HumidityDownscaling", + "DownscalingParamDirs", + "custom_input_fcst_freq", +) + +SUPPLEMENTAL_PRODUCT_FIELDS = ( + SUPPLEMENTAL_PRECIPITATION, + "SuppPcpForcingTypes", + "SuppPcpDirectories", + "RegridOptSuppPcp", + "SuppPcpTemporalInterpolation", + "SuppPcpInputOffsets", + "SuppPcpMandatory", +) diff --git a/NextGen_Forcings_Engine_BMI/coastal/gridded/__init__.py b/NextGen_Forcings_Engine_BMI/coastal/gridded/__init__.py new file mode 100644 index 00000000..b174d21e --- /dev/null +++ b/NextGen_Forcings_Engine_BMI/coastal/gridded/__init__.py @@ -0,0 +1 @@ +"""Gridded coastal forcing utilities.""" diff --git a/NextGen_Forcings_Engine_BMI/coastal/gridded/config_manager.py b/NextGen_Forcings_Engine_BMI/coastal/gridded/config_manager.py new file mode 100644 index 00000000..8650e656 --- /dev/null +++ b/NextGen_Forcings_Engine_BMI/coastal/gridded/config_manager.py @@ -0,0 +1,191 @@ +"""Build effective forcing configurations for gridded coastal domains. + +This was added when adding support for coastal workflows. +The original business logic was in the ``nwm-rte`` repo in ``nwm-rte/run_coastal.py`` in the ``coastalforcing` branch. +""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from datetime import datetime +from pathlib import Path +from typing import Any + +import yaml + +from NextGen_Forcings_Engine_BMI.coastal import consts +from NextGen_Forcings_Engine_BMI.run_bmi_model import REFERENCE_TIME_FORMAT + +LOG = logging.getLogger("FORCING") + + +def analysis_num_updates(config_path: Path) -> int | None: + """Return the hourly AnA iteration count implied by input horizons.""" + config = yaml.safe_load(config_path.read_text()) + if config.get(consts.ANALYSIS_FLAG, 0) != 1: + return None + horizons = config.get(consts.FORECAST_INPUT_HORIZONS, []) + if not horizons: + raise ValueError( + f"{consts.FORECAST_INPUT_HORIZONS} is required for analysis cycles" + ) + num_updates = int(max(horizons) / 60) + if num_updates < 1: + raise ValueError( + f"{consts.FORECAST_INPUT_HORIZONS} must cover at least one hour" + ) + return num_updates + + +class CoastalConfigManager: + """Build a gridded forcing config for one coastal domain.""" + + def __init__( + self, + forcing_configuration: str, + global_domain: str, + cycle_datetime: datetime | None = None, + start_time: datetime | None = None, + lookback: int | None = None, + forecast_input_horizons: int | None = None, + *, + forcing_root_dir: str | Path = consts.DEFAULT_FORCING_ROOT_DIR, + forcing_template_dir: str | Path = consts.DEFAULT_FORCING_TEMPLATE_DIR, + geogrid_dir: str | Path = consts.DEFAULT_GEOGRID_DIR, + regrid_weights_dir: str | Path = consts.DEFAULT_REGRID_WEIGHTS_DIR, + spatial_metadata_suffix: Mapping[str, str] = consts.SPATIAL_METADATA_SUFFIX, + debug_subset_name: str = consts.DEBUG_CONUS_SUBSET, + debug_conus_subset: bool = False, + ) -> None: + if debug_conus_subset and global_domain != "CONUS": + raise ValueError("debug_conus_subset requires global_domain='CONUS'") + self.forcing_configuration = forcing_configuration + self.global_domain = global_domain + self.cycle_datetime = cycle_datetime + self.start_time = start_time + self.lookback = lookback + self.forecast_input_horizons = forecast_input_horizons + self.forcing_root_dir = Path(forcing_root_dir) + self.forcing_template_dir = Path(forcing_template_dir) + self.geogrid_dir = Path(geogrid_dir) + self.regrid_weights_dir = Path(regrid_weights_dir) + self.spatial_metadata_suffix = spatial_metadata_suffix + self.debug_subset_name = debug_subset_name + self.debug_conus_subset = debug_conus_subset + + @property + def target_domain(self) -> str: + if self.debug_conus_subset: + return self.debug_subset_name + return self.global_domain + + @property + def template_path(self) -> Path: + return self.forcing_template_dir / f"{self.forcing_configuration}_config.yml" + + @property + def geogrid(self) -> Path: + return self.geogrid_dir / f"geo_em_{self.target_domain}.nc" + + @property + def spatial_metadata(self) -> Path: + if self.debug_conus_subset: + spatial_suffix = self.debug_subset_name + else: + spatial_suffix = self.spatial_metadata_suffix.get( + self.global_domain, self.global_domain + ) + return ( + self.geogrid_dir / f"GEOGRID_LDASOUT_Spatial_Metadata_{spatial_suffix}.nc" + ) + + @property + def output_dir(self) -> Path: + subset_suffix = f"_{self.target_domain}" if self.debug_conus_subset else "" + return ( + self.forcing_root_dir + / "scratch" + / f"{self.forcing_configuration}_coastal{subset_suffix}" + ) + + @property + def file_date(self) -> datetime: + file_date = self.cycle_datetime or self.start_time + if file_date is None: + raise ValueError("Provide cycle_datetime or start_time and end_time") + return file_date + + @property + def output_path(self) -> Path: + timestamp = self.file_date.strftime(REFERENCE_TIME_FORMAT) + return self.output_dir / f"{self.target_domain}_{timestamp}.nc" + + @property + def config_path(self) -> Path: + return self.output_path.with_suffix(".yml") + + def _validate_paths(self) -> None: + if not self.template_path.is_file(): + raise FileNotFoundError( + f"No forcing template found for {self.forcing_configuration!r}: " + f"{self.template_path}" + ) + for required_path in (self.geogrid, self.spatial_metadata): + if not required_path.is_file(): + raise FileNotFoundError( + f"Required coastal domain file not found: {required_path}" + ) + + def build_config(self) -> dict[str, Any]: + """Load the installed template and apply coastal gridded overrides.""" + self._validate_paths() + template = self.template_path.read_text() + template = template.replace("{root_dir}", str(self.forcing_root_dir)) + template = template.replace("{global_domain}", self.global_domain) + config = yaml.safe_load(template) + config[consts.REFERENCE_FORECAST_DATE] = self.file_date.strftime( + REFERENCE_TIME_FORMAT + ) + + if ( + self.lookback is not None or self.forecast_input_horizons is not None + ) and config.get(consts.ANALYSIS_FLAG, 0) != 1: + raise ValueError( + "lookback and forecast_input_horizons are only supported for " + f"analysis configurations. {self.forcing_configuration} has " + f"{consts.ANALYSIS_FLAG}=" + f"{config.get(consts.ANALYSIS_FLAG, 0)}" + ) + if self.lookback is not None: + config[consts.LOOKBACK] = self.lookback + if self.forecast_input_horizons is not None: + configured_horizons = config.get(consts.FORECAST_INPUT_HORIZONS, [1]) + config[consts.FORECAST_INPUT_HORIZONS] = [ + self.forecast_input_horizons + ] * len(configured_horizons) + + config.update( + { + consts.GRID_TYPE: consts.GRIDDED_GRID_TYPE, + consts.OUTPUT: 1, + consts.GEOPACKAGE: "", + consts.GEOGRID_INPUT: str(self.geogrid), + consts.SPATIAL_METADATA_INPUT: str(self.spatial_metadata), + consts.LONGITUDE_VARIABLE: consts.GRID_LONGITUDE_VARIABLE, + consts.LATITUDE_VARIABLE: consts.GRID_LATITUDE_VARIABLE, + consts.REGRID_WEIGHTS_DIR: str(self.regrid_weights_dir), + consts.REUSE_REGRID_WEIGHTS: True, + consts.SCRATCH_DIRECTORY: str(self.output_dir), + } + ) + return config + + def write_config(self) -> Path: + """Write the effective config beside the final grid and return its path.""" + self.output_dir.mkdir(parents=True, exist_ok=True) + self.regrid_weights_dir.mkdir(parents=True, exist_ok=True) + config = self.build_config() + with self.config_path.open("w") as config_file: + yaml.safe_dump(config, config_file, sort_keys=False) + return self.config_path diff --git a/NextGen_Forcings_Engine_BMI/coastal/gridded/warm_regrid_weights.py b/NextGen_Forcings_Engine_BMI/coastal/gridded/warm_regrid_weights.py new file mode 100644 index 00000000..a582e8be --- /dev/null +++ b/NextGen_Forcings_Engine_BMI/coastal/gridded/warm_regrid_weights.py @@ -0,0 +1,287 @@ +"""Module for recreating cached ESMF regrid weights for a gridded configuration. + +This was added to support large domains which were consuming too much memory when +the existing ESMF weight-generation operations were called against the full range of variables +for very large domains. + +This was added when adding support for coastal workflows. +The original business logic was in the ``nwm-rte`` repo in ``nwm-rte/run_coastal.py`` in the ``coastalforcing` branch. +""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import logging +import subprocess +import sys +import tempfile +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Literal + +import yaml +from NextGen_Forcings_Engine.bmi_model import BMIMODEL + +from NextGen_Forcings_Engine_BMI.coastal import consts +from NextGen_Forcings_Engine_BMI.coastal.gridded.config_manager import ( + CoastalConfigManager, +) +from NextGen_Forcings_Engine_BMI.run_bmi_model import REFERENCE_TIME_FORMAT + +ProductKind = Literal["input", "supplemental"] + +LOG = logging.getLogger("FORCING") + + +@dataclass(frozen=True) +class Product: + """One configured forcing product to warm in an isolated process.""" + + kind: ProductKind + key: int + + +def _read_config(config_path: Path) -> dict[str, Any]: + with config_path.open() as config_file: + return yaml.safe_load(config_file) + + +def _products_to_warm(config: dict[str, Any]) -> list[Product]: + input_products = [ + Product("input", product_key) + for product_key in config.get(consts.INPUT_FORCINGS, []) + ] + supplemental_products = [ + Product("supplemental", product_key) + for product_key in config.get(consts.SUPPLEMENTAL_PRECIPITATION, []) + ] + return input_products + supplemental_products + + +def _remove_cached_weights(model, product: Product) -> None: + forcing_modules = ( + model._input_forcing_mod if product.kind == "input" else model._supp_pcp_mod + ) + forcing = forcing_modules[product.key] + file_key = f"{forcing.product_name}_{model._job_meta.geogrid}" + hash_key = hashlib.md5(file_key.encode()).hexdigest()[:8] + weight_file = Path(model._job_meta.weightsDir) / f"ESMF_weight_{hash_key}.nc4" + weight_file.unlink(missing_ok=True) + + +def _select_product_at_index( + config: dict[str, Any], + product_field: str, + related_fields: tuple[str, ...], + product_index: int, +) -> None: + """Keep one entry from each list belonging to a configured product.""" + product_count = len(config[product_field]) + for field in related_fields: + values = config.get(field) + if not values: + continue + if not isinstance(values, list) or len(values) != product_count: + raise ValueError( + f"{field} must contain one entry per {product_field} product" + ) + config[field] = [values[product_index]] + + +def _remove_supplemental_products(config: dict[str, Any]) -> None: + """Remove supplemental products while warming an input forcing.""" + for field in consts.SUPPLEMENTAL_PRODUCT_FIELDS: + if field in config: + config[field] = [] + + +def _config_for_product(config: dict[str, Any], product: Product) -> dict[str, Any]: + """Create a no-output config containing only the product being warmed.""" + product_config = copy.deepcopy(config) + product_config[consts.OUTPUT] = 0 + + if product.kind == "input": + product_index = product_config[consts.INPUT_FORCINGS].index(product.key) + _select_product_at_index( + product_config, + consts.INPUT_FORCINGS, + consts.INPUT_PRODUCT_FIELDS, + product_index, + ) + _remove_supplemental_products(product_config) + else: + if not product_config.get(consts.INPUT_FORCINGS): + raise ValueError("A supplemental product requires an input forcing") + _select_product_at_index( + product_config, + consts.INPUT_FORCINGS, + consts.INPUT_PRODUCT_FIELDS, + 0, + ) + product_index = product_config[consts.SUPPLEMENTAL_PRECIPITATION].index( + product.key + ) + _select_product_at_index( + product_config, + consts.SUPPLEMENTAL_PRECIPITATION, + consts.SUPPLEMENTAL_PRODUCT_FIELDS, + product_index, + ) + + return product_config + + +def _create_gridded_model(cycle_datetime: datetime, geogrid: Path): + return BMIMODEL[consts.GRIDDED_GRID_TYPE]( + b_date=cycle_datetime.strftime(REFERENCE_TIME_FORMAT), + geogrid=str(geogrid), + ) + + +def _warm_product( + config_path: Path, + cycle_datetime: datetime, + geogrid: Path, + product: Product, +) -> None: + config = _config_for_product(_read_config(config_path), product) + model = None + initialized = False + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yml") as config_file: + yaml.safe_dump(config, config_file, sort_keys=False) + config_file.flush() + + try: + model = _create_gridded_model(cycle_datetime, geogrid) + model.initialize(config_file.name) + initialized = True + _remove_cached_weights(model, product) + model.update() + finally: + if initialized: + model.finalize() + + +def warm_coastal_regrid_weights( + forcing_configuration: str, + global_domain: str, + cycle_datetime: datetime, + lookback: int | None = None, + forecast_input_horizons: int | None = None, + *, + forcing_root_dir: str | Path = consts.DEFAULT_FORCING_ROOT_DIR, + forcing_template_dir: str | Path = consts.DEFAULT_FORCING_TEMPLATE_DIR, + geogrid_dir: str | Path = consts.DEFAULT_GEOGRID_DIR, + regrid_weights_dir: str | Path = consts.DEFAULT_REGRID_WEIGHTS_DIR, + spatial_metadata_suffix: Mapping[str, str] = consts.SPATIAL_METADATA_SUFFIX, + debug_subset_name: str = consts.DEBUG_CONUS_SUBSET, + debug_conus_subset: bool = False, +) -> None: + """Build a coastal config and recreate each product's regrid weights.""" + config_manager = CoastalConfigManager( + forcing_configuration, + global_domain, + cycle_datetime=cycle_datetime, + lookback=lookback, + forecast_input_horizons=forecast_input_horizons, + forcing_root_dir=forcing_root_dir, + forcing_template_dir=forcing_template_dir, + geogrid_dir=geogrid_dir, + regrid_weights_dir=regrid_weights_dir, + spatial_metadata_suffix=spatial_metadata_suffix, + debug_subset_name=debug_subset_name, + debug_conus_subset=debug_conus_subset, + ) + config_path = config_manager.write_config() + + LOG.info( + "Warming coastal regrid weights: configuration=%s domain=%s", + forcing_configuration, + global_domain, + ) + warm_regrid_weights( + config_path=config_path, + cycle_datetime=cycle_datetime, + geogrid=config_manager.geogrid, + ) + + +def warm_regrid_weights( + config_path: str | Path, + cycle_datetime: datetime, + geogrid: str | Path, +) -> None: + """Recreate cached weights for every product, one subprocess per product.""" + config_path = Path(config_path).resolve() + geogrid = Path(geogrid).resolve() + config = _read_config(config_path) + + if config.get(consts.GRID_TYPE) != consts.GRIDDED_GRID_TYPE: + raise ValueError( + f"Regrid-weight warmup requires {consts.GRID_TYPE}: " + f"{consts.GRIDDED_GRID_TYPE}" + ) + if not config.get(consts.REUSE_REGRID_WEIGHTS): + raise ValueError( + f"Regrid-weight warmup requires {consts.REUSE_REGRID_WEIGHTS}: true" + ) + + weights_dir = config.get(consts.REGRID_WEIGHTS_DIR) + if not weights_dir: + raise ValueError(f"Regrid-weight warmup requires {consts.REGRID_WEIGHTS_DIR}") + Path(weights_dir).mkdir(parents=True, exist_ok=True) + + products = _products_to_warm(config) + if not products: + raise ValueError("The forcing configuration contains no products to warm") + + for product in products: + subprocess.run( + [ + sys.executable, + "-m", + "NextGen_Forcings_Engine_BMI.coastal.gridded.warm_regrid_weights", + "--config-path", + str(config_path), + "--cycle-datetime", + cycle_datetime.isoformat(), + "--geogrid", + str(geogrid), + "--product-kind", + product.kind, + "--product-key", + str(product.key), + ], + check=True, + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config-path", required=True, type=Path) + parser.add_argument("--cycle-datetime", required=True, type=datetime.fromisoformat) + parser.add_argument("--geogrid", required=True, type=Path) + parser.add_argument( + "--product-kind", required=True, choices=("input", "supplemental") + ) + parser.add_argument("--product-key", required=True, type=int) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + _warm_product( + args.config_path, + args.cycle_datetime, + args.geogrid, + Product(args.product_kind, args.product_key), + ) + + +if __name__ == "__main__": + main() diff --git a/NextGen_Forcings_Engine_BMI/run_bmi_model.py b/NextGen_Forcings_Engine_BMI/run_bmi_model.py index 553b6cf6..69c1835e 100755 --- a/NextGen_Forcings_Engine_BMI/run_bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/run_bmi_model.py @@ -1,7 +1,7 @@ import argparse import datetime import pathlib -from pathlib import Path +from abc import ABC, abstractmethod import numpy as np import pandas as pd @@ -14,6 +14,8 @@ parse_config, ) +REFERENCE_TIME_FORMAT = "%Y%m%d%H%M" + def get_date_times(start_time: str, end_time: str) -> tuple: """Convert start and end time strings to a list of datetimes at hourly intervals.""" @@ -35,7 +37,7 @@ def print_init(model: NWMv3_Forcing_Engine_BMI_model, num_iterations: int) -> No def print_pre_update(num, num_iterations, timestamp) -> None: """Print the current iteration number and timestamp before updating the model.""" print("\n---------------------------------------------------") - print(f"Iteration #{num} of {num_iterations} for {timestamp}") + print(f"Iteration #{num} of {num_iterations} for timestamp={timestamp}") def print_post_update(model, start_time) -> None: @@ -310,13 +312,248 @@ def get_options(): return parser.parse_args() +class ForcingRunner(ABC): + """Define the shared forcing-engine execution lifecycle.""" + + def __init__( + self, + config_path: pathlib.Path | None = None, + cycle_datetime: datetime.datetime | None = None, + start_time: str | datetime.datetime | None = None, + end_time: str | datetime.datetime | None = None, + b_date: str | None = None, + geogrid: str = None, + output_path: pathlib.Path = None, + config: dict | None = None, + num_updates: int | None = None, + ) -> None: + self.config_path = self.resolve_config_path(config_path) + self.cycle_datetime = cycle_datetime + self.start_time = start_time + self.end_time = end_time + self.b_date = b_date + self.geogrid = geogrid + self.output_path = output_path + self.config = config + self.requested_num_updates = num_updates + self.model = None + self.output_steps = None + self.ngen_datetimes = None + self.is_general = False + self.num_updates = None + + @classmethod + def create( + cls, + config_path: pathlib.Path | None = None, + **kwargs, + ) -> "ForcingRunner": + """Create the runner child class appropriate for the configured grid type.""" + config_path = cls.resolve_config_path(config_path) + config = cls._read_config(config_path) + + if config["GRID_TYPE"] == "gridded": + child = ForcingRunnerGridded + else: + child = ForcingRunnerGeneral + + return child(config_path, config=config, **kwargs) + + @staticmethod + def resolve_config_path(config_path: pathlib.Path | None) -> pathlib.Path: + """Resolve a supplied config path or use the bundled default.""" + if config_path is not None: + return pathlib.Path(config_path) + return pathlib.Path(__file__).parent.resolve() / "config.yml" + + @staticmethod + def _read_config(config_path: pathlib.Path) -> dict: + """Read and parse a forcing configuration file.""" + with config_path.open("r") as config_file: + return parse_config(yaml.safe_load(config_file)) + + def validate_args(self) -> None: + """Validate arguments shared by every runner.""" + if (self.start_time is None) != (self.end_time is None): + raise ValueError("start_time and end_time must be provided together") + if self.requested_num_updates is not None and self.requested_num_updates < 1: + raise ValueError("num_updates must be a positive integer") + + def load_config(self) -> None: + """Load and parse the forcing configuration.""" + if self.config is None: + self.config = self._read_config(self.config_path) + + def validate_config(self) -> None: + """Validate configuration shared by every runner.""" + if self.config["GRID_TYPE"] not in BMIMODEL: + raise ValueError(f"Unsupported GRID_TYPE: {self.config['GRID_TYPE']}") + + @abstractmethod + def _create_model(self): + """Construct the BMI model for this runner type.""" + raise NotImplementedError + + @abstractmethod + def _get_update_plan(self) -> int: + """Return the number of model updates to run.""" + raise NotImplementedError + + def initialize_model(self) -> None: + """Construct and initialize the BMI model.""" + print("Creating an instance of the BMI model object") + self.model = self._create_model() + self.model.initialize(str(self.config_path)) + self.num_updates = self._get_update_plan() + print_init(self.model, self.num_updates) + + def _update_model(self, num_updates: int) -> None: + """Run model updates with standard progress output when available.""" + for num in range(num_updates): + if self.is_general: + print_pre_update(num, num_updates, self.ngen_datetimes[num]) + self.model.update() + if self.is_general: + print_post_update(self.model, self.start_time) + + def _run_updates(self) -> None: + """Run the update plan supplied by the concrete runner.""" + self._update_model(self.num_updates) + + def run_model(self) -> None: + """Initialize, run, and finalize the BMI model.""" + try: + self.initialize_model() + self._run_updates() + finally: + if self.model is not None: + print("\nFinalizing the BMI model") + self.model.finalize() + + +class ForcingRunnerGeneral(ForcingRunner): + """Run a hydrofabric or unstructured forcing configuration.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.is_general = True + + def validate_args(self) -> None: + """Validate and normalize the standard start and end times.""" + super().validate_args() + if self.cycle_datetime is not None: + raise ValueError( + "cycle_datetime is only supported for gridded configurations" + ) + if self.start_time is None or self.end_time is None: + raise ValueError("start_time and end_time are required") + self.ngen_datetimes, self.start_time, self.end_time = get_date_times( + self.start_time, self.end_time + ) + self.num_updates = len(self.ngen_datetimes) + + def validate_config(self) -> None: + """Validate a non-gridded forcing configuration.""" + super().validate_config() + if self.config["GRID_TYPE"] == "gridded": + raise ValueError("ForcingRunnerGeneral does not support GRID_TYPE: gridded") + + def _create_model(self): + """Construct the configured BMI model.""" + return BMIMODEL[self.config["GRID_TYPE"]]( + self.b_date, + self.geogrid, + output_path=str(self.output_path) if self.output_path else None, + ) + + def _get_update_plan(self) -> int: + """Run one update for every requested timestamp.""" + return self.num_updates + + +class ForcingRunnerGridded(ForcingRunner): + """Run a gridded forcing cycle or explicit retrospective window.""" + + @property + def explicit_window(self) -> bool: + """Whether the runner was given a fixed start/end time window.""" + return self.start_time is not None or self.end_time is not None + + def validate_args(self) -> None: + """Validate and normalize timing inputs.""" + super().validate_args() + self.validate_gridded_args() + + def validate_gridded_args(self) -> None: + """Validate and normalize gridded timing inputs.""" + if self.cycle_datetime is not None and self.explicit_window: + raise ValueError( + "cycle_datetime cannot be combined with start_time/end_time" + ) + if self.requested_num_updates is not None and self.cycle_datetime is None: + raise ValueError("num_updates requires cycle_datetime") + if self.cycle_datetime is None and ( + self.start_time is None or self.end_time is None + ): + raise ValueError("Provide cycle_datetime or both start_time and end_time") + if isinstance(self.start_time, str) and isinstance(self.end_time, str): + _, self.start_time, self.end_time = get_date_times( + self.start_time, self.end_time + ) + if self.explicit_window and self.end_time < self.start_time: + raise ValueError("end_time must not be earlier than start_time") + + def validate_config(self) -> None: + """Load configuration and derive the gridded output count.""" + super().validate_config() + self.validate_gridded_config() + + def validate_gridded_config(self) -> None: + """Validate gridded configuration and derive its output count.""" + if self.config["GRID_TYPE"] != "gridded": + raise ValueError("ForcingRunnerGridded requires GRID_TYPE: gridded") + if self.explicit_window and self.config["AnAFlag"]: + raise ValueError( + "Explicit windows are not supported for analysis configurations; " + "provide cycle_datetime so LookBack controls the window" + ) + if self.explicit_window: + duration_seconds = (self.end_time - self.start_time).total_seconds() + step_seconds = self.config["time_step_seconds"] + if duration_seconds % step_seconds: + raise ValueError( + "The explicit window must be evenly divisible by time_step_seconds" + ) + self.output_steps = int(duration_seconds / step_seconds) + 1 + + def _create_model(self): + """Construct the gridded BMI model with timing overrides.""" + reference_time = self.cycle_datetime or self.start_time + return BMIMODEL["gridded"]( + b_date=self.b_date or reference_time.strftime(REFERENCE_TIME_FORMAT), + geogrid=self.geogrid, + output_path=str(self.output_path) if self.output_path else None, + output_steps=self.output_steps, + ) + + def _get_update_plan(self) -> int: + """Use the configured or requested update count.""" + if self.explicit_window: + return self.output_steps + if self.requested_num_updates is not None: + return self.requested_num_updates + return self.model._job_meta.actual_output_steps + + def run_bmi( - start_time: str, - end_time: str, + start_time: str | None = None, + end_time: str | None = None, config_path: pathlib.Path = None, b_date: str = None, geogrid: str = None, output_path: pathlib.Path = None, + cycle_datetime: datetime.datetime | None = None, + num_updates: int | None = None, ): """Execute the NextGen Forcings Engine BMI model. @@ -325,40 +562,30 @@ def run_bmi( :param start_time: The start time for the simulation, in the format 'YYYY-MM-DD HH:mm:ss'. :param end_time: The end time for the simulation, in the format 'YYYY-MM-DD HH:mm:ss'. - :param config_path: Optional path to the configuration file. Defaults to './config.yml' if not provided. + :param config_path: Optional path to the configuration file. Defaults to the config.yml bundled with this module. :param b_date: The begin date for the forecast cycle, in the format 'YYYYMMDDHHmm'. If omitted, reads from config. :param geogrid: Path to the geospatial grid file. If omitted, reads from the config file. :param output_path: Path to the output file. If omitted, a default output path is generated. + :param cycle_datetime: Optional gridded forecast cycle time. Cannot be combined with start_time or end_time. + :param num_updates: Optional positive update count for a gridded cycle. This limits iteration count without changing an analysis lookback window. :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ print("Initializing the BMI model") - # Set the path for the config file, using the default if none is provided - cfg_path = ( - str(config_path) - if config_path is not None - else str(Path(__file__).parent.resolve() / "config.yml") - ) - with open(cfg_path, "r") as fp: - config = parse_config(yaml.safe_load(fp)) - - print("Creating an instance of the BMI model object") - model = BMIMODEL[config.get("GRID_TYPE")]( - b_date, geogrid, output_path=str(output_path) if output_path else None + runner = ForcingRunner.create( + config_path, + cycle_datetime=cycle_datetime, + start_time=start_time, + end_time=end_time, + b_date=b_date, + geogrid=geogrid, + output_path=output_path, + num_updates=num_updates, ) - model.initialize(cfg_path) - - ngen_datetimes, start_time, end_time = get_date_times(start_time, end_time) - num_iterations = len(ngen_datetimes) - print_init(model, num_iterations) - - for num, timestamp in enumerate(ngen_datetimes): - print_pre_update(num, num_iterations, timestamp) - model.update() - print_post_update(model, start_time) - - print("\nFinalizing the BMI model") - model.finalize() + runner.validate_args() + runner.load_config() + runner.validate_config() + runner.run_model() def main(): diff --git a/pyproject.toml b/pyproject.toml index 8c04182c..7438b0fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,8 +52,8 @@ build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] where = ["NextGen_Forcings_Engine_BMI","."] -include = ["NextGen_Forcings_Engine*", - "Forcing_Extraction_Scripts*", +include = ["NextGen_Forcings_Engine*", + "Forcing_Extraction_Scripts*", "ESMF_Mesh_Domain_Configuration_Production*"] namespaces = false From 7f1d0e515d4364b32a66bbd44c7538609d7ebda0 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 18 Sep 2026 17:46:59 -0400 Subject: [PATCH 2/2] Coastal support: Add optional T0 output for gridded configuration --- .../NextGen_Forcings_Engine/bmi_model.py | 24 +++++++++++-- .../NextGen_Forcings_Engine/core/config.py | 33 ++++++++++++----- .../NextGen_Forcings_Engine/model.py | 31 ++++++++++++++++ NextGen_Forcings_Engine_BMI/run_bmi_model.py | 36 +++++++++++++++++-- 4 files changed, 110 insertions(+), 14 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py index afff96bd..04ef773e 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -151,6 +151,7 @@ def __init__( geogrid: str = None, output_path: str = None, output_steps: int = None, + output_t0: bool = False, ) -> None: """Create a model that is ready for initialization. @@ -168,11 +169,14 @@ def __init__( the forcing configuration's timing settings. This must be ``None`` for AnA aka Analysis & Assimilation runs, whose output count is controlled by their configured lookback window. + output_t0: Whether T0 precedes the configured output steps. By default + (when this is False), T0 will not exist in the output timesteps. """ self.output_path = output_path self._geogrid = geogrid self._b_date = b_date self._output_steps = output_steps + self._output_t0 = output_t0 self._values = {} self._start_time = 0.0 @@ -258,6 +262,7 @@ def _job_meta(self, value: ConfigOptions) -> None: b_date=self._b_date, geogrid=self._geogrid, output_steps=self._output_steps, + output_t0=self._output_t0, ) except KeyboardInterrupt as e: err_handler.err_out_screen("User keyboard interrupt", e) @@ -1597,12 +1602,19 @@ def __init__( geogrid: str = None, output_path: str = None, output_steps: int = None, + output_t0: bool = False, ): """Create a model that is ready for initialization. Initializes the model with default values for time, variables, and grid types. """ - super().__init__(b_date, geogrid, output_path, output_steps) + super().__init__( + b_date, + geogrid, + output_path, + output_steps, + output_t0, + ) def grid_ranks(self) -> list[int]: """Get the grid ranks for the gridded domain.""" @@ -1667,12 +1679,15 @@ def __init__( b_date: str = None, geogrid: str = None, output_path: str = None, + output_t0: bool = False, ): """Create a model that is ready for initialization. Initializes the model with default values for time, variables, and grid types. """ - super().__init__(b_date, geogrid, output_path) + if output_t0 is not False: + raise ValueError("output_t0 is only supported for gridded configurations") + super().__init__(b_date, geogrid, output_path, output_t0=output_t0) def grid_ranks(self) -> list[int]: """Get the grid ranks for the hydrofabric domain.""" @@ -1733,12 +1748,15 @@ def __init__( b_date: str = None, geogrid: str = None, output_path: str = None, + output_t0: bool = False, ): """Create a model that is ready for initialization. Initializes the model with default values for time, variables, and grid types. """ - super().__init__(b_date, geogrid, output_path) + if output_t0 is not False: + raise ValueError("output_t0 is only supported for gridded configurations") + super().__init__(b_date, geogrid, output_path, output_t0=output_t0) def grid_ranks(self) -> list[int]: """Get the grid ranks for the unstructured domain.""" diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py index 70a111c0..8b587582 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py @@ -37,6 +37,7 @@ def __init__( b_date: str = None, geogrid: str = None, output_steps: int = None, + output_t0: bool = False, ) -> None: """Initialize the configuration class to empty None attributes. @@ -47,17 +48,30 @@ def __init__( b_date (str, optional): The beginning date of processing in the format YYYYMMDDHHMM. This is used to calculate the processing window for realtime simulations. If not provided, it will be read from the configuration file. geogrid (str, optional): The filepath to the geogrid file to be used for processing. This is used to specify the grid information for regridding input forcings. If not provided, it will be read from the configuration file. output_steps (int, optional): Override the configured output count. When set, - this must be at least one and takes precedence over analysis and - forecast-cycle output-step calculations. + this must be nonnegative when including T0 and otherwise at least + one. It takes precedence over forecast-cycle output-step calculations. + output_t0 (bool, optional): + Write an additional T0 record before the normal gridded forecast outputs, with identical forcing values as T1. + Not supported for analysis configurations. + Only supported for ``gridded`` configurations. + An independent T0 state cannot be reliably computed because several forecast products omit required fields at hour zero and substitute hour one. Examples: + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1202-L1206 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1444-L1448 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L2043-L2047 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L4129-L4135 """ self._b_date_proc = None self._geogrid = None self.user_provided_geogrid_flag = geogrid is not None - if output_steps is not None and output_steps < 1: - raise ValueError(f"output_steps must be at least 1, but got {output_steps}") + minimum_output_steps = 0 if output_t0 else 1 + if output_steps is not None and output_steps < minimum_output_steps: + raise ValueError( + f"output_steps must be at least {minimum_output_steps}, but got {output_steps}" + ) self.output_steps = output_steps + self._output_t0 = output_t0 # If b_date not provided, try to read from config file if b_date is None: @@ -1223,16 +1237,19 @@ def actual_output_steps(self) -> int: compatible with analysis configurations, whose output count is determined by their lookback window. """ + if self._output_t0 and self.ana_flag: + raise ValueError("output_t0 is not supported for analysis configurations") if self.output_steps is not None: if self.ana_flag: raise ValueError( "output_steps is not supported for analysis configurations" ) - return np.int32(self.output_steps) - if self.ana_flag: - return np.int32(self.nFcsts) + output_steps = self.output_steps + elif self.ana_flag: + output_steps = self.nFcsts else: - return np.int32(self.num_output_steps) + output_steps = self.num_output_steps + return np.int32(output_steps + int(self._output_t0)) @property def should_randomize_weight_file_name(self) -> bool: diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 769f4a58..39b6d8a1 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -139,6 +139,7 @@ def run(self, future_time: float) -> None: self.log_cycle() input_forcings = self.loop_through_forcing_products(future_time) self.process_suplemental_precip(input_forcings) + self.write_t0_output() self.write_output() self.update_bmi_output_dict() @@ -639,6 +640,36 @@ def process_suplemental_precip( if supp_pcp_key == 14: self.__process_supp_precip_key(input_forcings, supp_pcp_key) + def write_t0_output(self) -> None: + """Write an additional T0 record before the first forecast output, with identical forcing values as T1. + + An independent T0 state cannot be reliably computed because several forecast + products omit required fields at hour zero and substitute hour one. Examples: + + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1202-L1206 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1444-L1448 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L2043-L2047 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L4129-L4135 + """ + job_meta = self._bmi._job_meta + if not job_meta._output_t0 or job_meta.bmi_time_index != 0: + return + + time_step = datetime.timedelta(seconds=self._bmi.get_time_step()) + saved_time = job_meta.current_time + saved_output_date = job_meta.current_output_date + saved_out_date = self._bmi._output_obj.outDate + try: + job_meta.current_time -= time_step + job_meta.current_output_date -= time_step + self._bmi._output_obj.outDate -= time_step + self.write_output() + job_meta.bmi_time_index += 1 + finally: + job_meta.current_time = saved_time + job_meta.current_output_date = saved_output_date + self._bmi._output_obj.outDate = saved_out_date + @time_function def write_output(self) -> None: """Write the output for the current forecast cycle. diff --git a/NextGen_Forcings_Engine_BMI/run_bmi_model.py b/NextGen_Forcings_Engine_BMI/run_bmi_model.py index 69c1835e..fcf6db60 100755 --- a/NextGen_Forcings_Engine_BMI/run_bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/run_bmi_model.py @@ -326,6 +326,7 @@ def __init__( output_path: pathlib.Path = None, config: dict | None = None, num_updates: int | None = None, + output_t0: bool = False, ) -> None: self.config_path = self.resolve_config_path(config_path) self.cycle_datetime = cycle_datetime @@ -336,6 +337,7 @@ def __init__( self.output_path = output_path self.config = config self.requested_num_updates = num_updates + self.output_t0 = output_t0 self.model = None self.output_steps = None self.ngen_datetimes = None @@ -457,6 +459,8 @@ def validate_config(self) -> None: super().validate_config() if self.config["GRID_TYPE"] == "gridded": raise ValueError("ForcingRunnerGeneral does not support GRID_TYPE: gridded") + if self.output_t0 is not False: + raise ValueError("output_t0 is only supported for gridded configurations") def _create_model(self): """Construct the configured BMI model.""" @@ -512,6 +516,8 @@ def validate_gridded_config(self) -> None: """Validate gridded configuration and derive its output count.""" if self.config["GRID_TYPE"] != "gridded": raise ValueError("ForcingRunnerGridded requires GRID_TYPE: gridded") + if self.output_t0 and self.config["AnAFlag"]: + raise ValueError("output_t0 is not supported for analysis configurations") if self.explicit_window and self.config["AnAFlag"]: raise ValueError( "Explicit windows are not supported for analysis configurations; " @@ -529,20 +535,33 @@ def validate_gridded_config(self) -> None: def _create_model(self): """Construct the gridded BMI model with timing overrides.""" reference_time = self.cycle_datetime or self.start_time + + # Do some math in case the user asked for the T0 output to be included. + forecast_output_steps = self.output_steps + if self.output_t0 and forecast_output_steps is not None: + forecast_output_steps -= 1 + elif ( + self.output_t0 + and self.requested_num_updates is not None + and not self.config["AnAFlag"] + ): + forecast_output_steps = self.requested_num_updates + return BMIMODEL["gridded"]( b_date=self.b_date or reference_time.strftime(REFERENCE_TIME_FORMAT), geogrid=self.geogrid, output_path=str(self.output_path) if self.output_path else None, - output_steps=self.output_steps, + output_steps=forecast_output_steps, + output_t0=self.output_t0, ) def _get_update_plan(self) -> int: """Use the configured or requested update count.""" if self.explicit_window: - return self.output_steps + return self.output_steps - int(self.output_t0) if self.requested_num_updates is not None: return self.requested_num_updates - return self.model._job_meta.actual_output_steps + return self.model._job_meta.actual_output_steps - int(self.output_t0) def run_bmi( @@ -554,6 +573,7 @@ def run_bmi( output_path: pathlib.Path = None, cycle_datetime: datetime.datetime | None = None, num_updates: int | None = None, + output_t0: bool = False, ): """Execute the NextGen Forcings Engine BMI model. @@ -568,6 +588,15 @@ def run_bmi( :param output_path: Path to the output file. If omitted, a default output path is generated. :param cycle_datetime: Optional gridded forecast cycle time. Cannot be combined with start_time or end_time. :param num_updates: Optional positive update count for a gridded cycle. This limits iteration count without changing an analysis lookback window. + :param output_t0: + Write an additional T0 record before the normal gridded forecast outputs, with identical forcing values as T1. + Not supported for analysis configurations. + Only supported for ``gridded`` configurations. + An independent T0 state cannot be reliably computed because several forecast products omit required fields at hour zero and substitute hour one. Examples: + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1202-L1206 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L1444-L1448 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L2043-L2047 + https://github.com/NGWPC/ngen-forcing/blob/27e03ba138478dd449ce957b1c3ba4c36fc33d8f/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/time_handling.py#L4129-L4135 :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ @@ -581,6 +610,7 @@ def run_bmi( geogrid=geogrid, output_path=output_path, num_updates=num_updates, + output_t0=output_t0, ) runner.validate_args() runner.load_config()