From ea604d5446403a1eb0bfb9ce891f94ec4c7b39e4 Mon Sep 17 00:00:00 2001 From: Lauren Schambach <157057315+laurscham@users.noreply.github.com> Date: Thu, 4 Jun 2026 12:51:11 -0400 Subject: [PATCH 01/25] updates to run run_coastal.sh via the RTE --- .../NextGen_Forcings_Engine/bmi_model.py | 2 +- .../NextGen_Forcings_Engine/core/geoMod.py | 8 ++++---- .../NextGen_Forcings_Engine/core/parallel.py | 3 +++ 3 files changed, 8 insertions(+), 5 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 895ac14c..2c6b31f3 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -276,7 +276,7 @@ def geo_meta(self) -> GeoMeta: 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." ) - self._geo_meta = HydrofabricGeoMeta(self._job_meta, self._mpi_meta) + self._geo_meta = GeoMeta(self._job_meta, self._mpi_meta) return self._geo_meta @geo_meta.setter 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..d7d6135b 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -395,9 +395,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 +429,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..4eef6293 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -166,6 +166,9 @@ def _cleanup_geogrid(self) -> None: if self.config_options is None: self.log_debug("Cleanup: config_options is not set") return + if getattr(self.config_options, "grid_type", None) == "gridded": + self.log_debug("Cleanup: skipping geogrid cleanup for gridded mode") + return geogrid = getattr(self.config_options, "geogrid", None) if geogrid is not None: self.try_delete_file_no_reraise(geogrid) From cce1a4687d9d319b29f7331fc41cd091f4505ce0 Mon Sep 17 00:00:00 2001 From: Lauren Schambach <157057315+laurscham@users.noreply.github.com> Date: Thu, 4 Jun 2026 15:14:51 -0400 Subject: [PATCH 02/25] adding debug statements to geoMod.py and regrid.py to debug no data values coming out of forcing engine on regular grids --- .../NextGen_Forcings_Engine/core/geoMod.py | 18 ++++++-- .../NextGen_Forcings_Engine/core/regrid.py | 41 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) 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 d7d6135b..7fc12282 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -282,12 +282,15 @@ 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] + print(f"[DEBUG GeoMeta] geogrid={self.config_options.geogrid}") + print(f"[DEBUG GeoMeta] ndim_lat={self.ndim_lat}, ny_global={self.ny_global}, nx_global={nx}") + 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) @@ -406,6 +409,11 @@ def latitude_grid(self) -> np.ndarray: # Flag to grab entire array for AWS slicing if self.config_options.aws: self.lat_bounds = var_tmp + print(f"[DEBUG GeoMeta] lat_var name: {self.config_options.lat_var}") + print(f"[DEBUG GeoMeta] latitude_grid shape: {var_tmp.shape}") + print(f"[DEBUG GeoMeta] latitude_grid min/max: {var_tmp.min():.4f} / {var_tmp.max():.4f}") + print(f"[DEBUG GeoMeta] latitude_grid corners: [0,0]={var_tmp[0,0]:.4f}, [-1,-1]={var_tmp[-1,-1]:.4f}") + print(f"[DEBUG GeoMeta] latitude_grid first row mean: {var_tmp[0,:].mean():.4f}, last row mean: {var_tmp[-1,:].mean():.4f}") else: var_tmp = None return var_tmp, "latitude_grid", self.config_options, False @@ -440,6 +448,10 @@ def longitude_grid(self) -> np.ndarray: # Flag to grab entire array for AWS slicing if self.config_options.aws: self.lon_bounds = var_tmp + print(f"[DEBUG GeoMeta] lon_var name: {self.config_options.lon_var}") + print(f"[DEBUG GeoMeta] longitude_grid shape: {var_tmp.shape}") + print(f"[DEBUG GeoMeta] longitude_grid min/max: {var_tmp.min():.4f} / {var_tmp.max():.4f}") + print(f"[DEBUG GeoMeta] longitude_grid corners: [0,0]={var_tmp[0,0]:.4f}, [-1,-1]={var_tmp[-1,-1]:.4f}") else: var_tmp = None 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..0881a065 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -1934,6 +1934,14 @@ def regrid_conus_hrrr(input_forcings, config_options, wrf_hydro_geo_meta, mpi_co ) err_handler.check_program_status(config_options, mpi_config) + if mpi_config.rank == 0: + out_data = input_forcings.esmf_field_out.data + n_nodata = int(np.sum(out_data == config_options.globalNdv)) + n_total = out_data.size + print(f"[DEBUG regrid out] var={input_forcings.netcdf_var_names[force_count]}, esmf_field_out shape: {out_data.shape}") + print(f"[DEBUG regrid out] min/max: {out_data.min():.4f} / {out_data.max():.4f}") + print(f"[DEBUG regrid out] no-data count: {n_nodata}/{n_total} ({100*n_nodata/n_total:.1f}%)") + try: input_forcings.regridded_forcings2[ input_forcings.input_map_output[force_count], :, : @@ -11283,6 +11291,21 @@ def check_regrid_status( """ pt = Partials(mpi_config, config_options) + # For gridded domains, esmf_lat/esmf_lon on GriddedGeoMeta are shadowed to None by + # GeoMeta.__init__ (via GEOMOD["GeoMeta"]), so the @cached_property never runs and the + # destination ESMF Grid's coordinate arrays are never populated. Populate them once here. + if config_options.grid_type == "gridded" and wrf_hydro_geo_meta.esmf_lat is None: + esmf_lat = wrf_hydro_geo_meta.esmf_grid.get_coords(1) + esmf_lat[:, :] = wrf_hydro_geo_meta.latitude_grid + wrf_hydro_geo_meta.esmf_lat = esmf_lat + esmf_lon = wrf_hydro_geo_meta.esmf_grid.get_coords(0) + esmf_lon[:, :] = wrf_hydro_geo_meta.longitude_grid + wrf_hydro_geo_meta.esmf_lon = esmf_lon + if mpi_config.rank == 0: + print(f"[DEBUG dest grid] populated ESMF destination grid coords") + print(f"[DEBUG dest grid] esmf_lat min/max: {esmf_lat.min():.4f} / {esmf_lat.max():.4f}") + print(f"[DEBUG dest grid] esmf_lon min/max: {esmf_lon.min():.4f} / {esmf_lon.max():.4f}") + # If the destination ESMF field hasn't been created, create it here. if not input_forcings.esmf_field_out: if config_options.grid_type == "gridded": @@ -12020,6 +12043,19 @@ 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: + print(f"[DEBUG calc_weights] product={input_forcings.product_name}, lat_var='{lat_var}', lon_var='{lon_var}'") + print(f"[DEBUG calc_weights] raw lat_tmp shape: {lat_tmp.shape}, lon_tmp shape: {lon_tmp.shape}") + print(f"[DEBUG calc_weights] lat_tmp min/max: {lat_tmp.min():.4f} / {lat_tmp.max():.4f}") + print(f"[DEBUG calc_weights] lon_tmp min/max: {lon_tmp.min():.4f} / {lon_tmp.max():.4f}") + print(f"[DEBUG calc_weights] lat_tmp corners: [0,0]={lat_tmp[0,0]:.4f}, [-1,-1]={lat_tmp[-1,-1]:.4f}") + print(f"[DEBUG calc_weights] lon_tmp corners: [0,0]={lon_tmp[0,0]:.4f}, [-1,-1]={lon_tmp[-1,-1]:.4f}") + print(f"[DEBUG calc_weights] input_forcings.ny_global={input_forcings.ny_global}, nx_global={input_forcings.nx_global}") + # Normalize source longitudes from 0-360 to -180/+180 to match geo_em geogrid convention + if lon_tmp.max() > 180: + lon_tmp = np.where(lon_tmp > 180, lon_tmp - 360, lon_tmp) + print(f"[DEBUG calc_weights] lon_tmp normalized to -180/+180: min/max={lon_tmp.min():.4f} / {lon_tmp.max():.4f}") + # Scatter global GFS latitude grid to processors.. if mpi_config.rank == 0: var_tmp = lat_tmp @@ -12053,6 +12089,11 @@ def calculate_weights( input_forcings.esmf_lats[:, :] = var_sub_lat_tmp input_forcings.esmf_lons[:, :] = var_sub_lon_tmp + if mpi_config.rank == 0: + print(f"[DEBUG ESMF src] esmf_lats (get_coords(1)) shape={input_forcings.esmf_lats.shape}") + print(f"[DEBUG ESMF src] esmf_lats min/max: {input_forcings.esmf_lats.min():.4f} / {input_forcings.esmf_lats.max():.4f}") + print(f"[DEBUG ESMF src] esmf_lons (get_coords(0)) shape={input_forcings.esmf_lons.shape}") + print(f"[DEBUG ESMF src] esmf_lons min/max: {input_forcings.esmf_lons.min():.4f} / {input_forcings.esmf_lons.max():.4f}") del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp From 65c72e6d5b77b3df5ed760a766f443c6de883a15 Mon Sep 17 00:00:00 2001 From: Lauren Schambach <157057315+laurscham@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:47:32 -0400 Subject: [PATCH 03/25] Attempt to reduce memory usage - only read in fields needed from the geogrids --- .../NextGen_Forcings_Engine/core/geoMod.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) 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 7fc12282..f25ce145 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -131,12 +131,33 @@ def spatial_metadata_exists(self) -> bool: else: return True + # All config_options attributes that hold geogrid variable names. + # Used to select only the needed variables when opening the geogrid file. + _GEOGRID_VAR_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", + ) + @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 self._GEOGRID_VAR_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) From 32523567a967b4c909e326300864b8c3976173af Mon Sep 17 00:00:00 2001 From: Lauren Schambach <157057315+laurscham@users.noreply.github.com> Date: Tue, 14 Jul 2026 21:18:31 -0400 Subject: [PATCH 04/25] changes needed to make sure all files are downloaded that are needed for the different config files. the bmi_model change is related to a crash that was being experienced when this wasn't commented out. the pyproject.toml change might not be needed, it was so the config templates get installed when building the rte --- NextGen_Forcings_Engine_BMI/forcing_extraction.py | 9 ++++++++- pyproject.toml | 4 ++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/forcing_extraction.py b/NextGen_Forcings_Engine_BMI/forcing_extraction.py index 84249418..9e182898 100644 --- a/NextGen_Forcings_Engine_BMI/forcing_extraction.py +++ b/NextGen_Forcings_Engine_BMI/forcing_extraction.py @@ -110,7 +110,14 @@ def retrieve_forcing(config_options: ConfigOptions): forcing_start_time = forcing_start_time + timedelta(hours=1) if supp_forcing_hours is not None: - forcing_start_time += timedelta(hours=supp_forcing_hours) + if supp_forcing_hours > 0 and input_horizons[i] > 0: + # Anchor end of window to the forecast horizon so all supp files + # are covered. ForecastDownloader extends forward via input_horizon, + # but FixedFileDownloader (e.g. MRMS) only looks backward. + forcing_start_time = refcstbdate + timedelta(minutes=input_horizons[i]) + else: + forcing_start_time += timedelta(hours=supp_forcing_hours) + lookback_hours += supp_forcing_hours if ana_flag == 1: lookback_hours = lookback_hours + 1 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 c18c9aa2794bd9bca03cfcc3572638781c9c28a4 Mon Sep 17 00:00:00 2001 From: Lauren Schambach Date: Wed, 5 Aug 2026 12:14:53 -0400 Subject: [PATCH 05/25] remove debug statements --- .../NextGen_Forcings_Engine/core/geoMod.py | 11 -------- .../NextGen_Forcings_Engine/core/regrid.py | 25 ------------------- 2 files changed, 36 deletions(-) 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 f25ce145..9f7629eb 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -309,8 +309,6 @@ def nx_global(self) -> int: else: # NOTE Is this correct? using lon_var nx = self.lon_var.shape[0] - print(f"[DEBUG GeoMeta] geogrid={self.config_options.geogrid}") - print(f"[DEBUG GeoMeta] ndim_lat={self.ndim_lat}, ny_global={self.ny_global}, nx_global={nx}") 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}" @@ -430,11 +428,6 @@ def latitude_grid(self) -> np.ndarray: # Flag to grab entire array for AWS slicing if self.config_options.aws: self.lat_bounds = var_tmp - print(f"[DEBUG GeoMeta] lat_var name: {self.config_options.lat_var}") - print(f"[DEBUG GeoMeta] latitude_grid shape: {var_tmp.shape}") - print(f"[DEBUG GeoMeta] latitude_grid min/max: {var_tmp.min():.4f} / {var_tmp.max():.4f}") - print(f"[DEBUG GeoMeta] latitude_grid corners: [0,0]={var_tmp[0,0]:.4f}, [-1,-1]={var_tmp[-1,-1]:.4f}") - print(f"[DEBUG GeoMeta] latitude_grid first row mean: {var_tmp[0,:].mean():.4f}, last row mean: {var_tmp[-1,:].mean():.4f}") else: var_tmp = None return var_tmp, "latitude_grid", self.config_options, False @@ -469,10 +462,6 @@ def longitude_grid(self) -> np.ndarray: # Flag to grab entire array for AWS slicing if self.config_options.aws: self.lon_bounds = var_tmp - print(f"[DEBUG GeoMeta] lon_var name: {self.config_options.lon_var}") - print(f"[DEBUG GeoMeta] longitude_grid shape: {var_tmp.shape}") - print(f"[DEBUG GeoMeta] longitude_grid min/max: {var_tmp.min():.4f} / {var_tmp.max():.4f}") - print(f"[DEBUG GeoMeta] longitude_grid corners: [0,0]={var_tmp[0,0]:.4f}, [-1,-1]={var_tmp[-1,-1]:.4f}") else: var_tmp = None 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 0881a065..dc0fe18e 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -1934,14 +1934,6 @@ def regrid_conus_hrrr(input_forcings, config_options, wrf_hydro_geo_meta, mpi_co ) err_handler.check_program_status(config_options, mpi_config) - if mpi_config.rank == 0: - out_data = input_forcings.esmf_field_out.data - n_nodata = int(np.sum(out_data == config_options.globalNdv)) - n_total = out_data.size - print(f"[DEBUG regrid out] var={input_forcings.netcdf_var_names[force_count]}, esmf_field_out shape: {out_data.shape}") - print(f"[DEBUG regrid out] min/max: {out_data.min():.4f} / {out_data.max():.4f}") - print(f"[DEBUG regrid out] no-data count: {n_nodata}/{n_total} ({100*n_nodata/n_total:.1f}%)") - try: input_forcings.regridded_forcings2[ input_forcings.input_map_output[force_count], :, : @@ -11301,10 +11293,6 @@ def check_regrid_status( esmf_lon = wrf_hydro_geo_meta.esmf_grid.get_coords(0) esmf_lon[:, :] = wrf_hydro_geo_meta.longitude_grid wrf_hydro_geo_meta.esmf_lon = esmf_lon - if mpi_config.rank == 0: - print(f"[DEBUG dest grid] populated ESMF destination grid coords") - print(f"[DEBUG dest grid] esmf_lat min/max: {esmf_lat.min():.4f} / {esmf_lat.max():.4f}") - print(f"[DEBUG dest grid] esmf_lon min/max: {esmf_lon.min():.4f} / {esmf_lon.max():.4f}") # If the destination ESMF field hasn't been created, create it here. if not input_forcings.esmf_field_out: @@ -12044,17 +12032,9 @@ 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: - print(f"[DEBUG calc_weights] product={input_forcings.product_name}, lat_var='{lat_var}', lon_var='{lon_var}'") - print(f"[DEBUG calc_weights] raw lat_tmp shape: {lat_tmp.shape}, lon_tmp shape: {lon_tmp.shape}") - print(f"[DEBUG calc_weights] lat_tmp min/max: {lat_tmp.min():.4f} / {lat_tmp.max():.4f}") - print(f"[DEBUG calc_weights] lon_tmp min/max: {lon_tmp.min():.4f} / {lon_tmp.max():.4f}") - print(f"[DEBUG calc_weights] lat_tmp corners: [0,0]={lat_tmp[0,0]:.4f}, [-1,-1]={lat_tmp[-1,-1]:.4f}") - print(f"[DEBUG calc_weights] lon_tmp corners: [0,0]={lon_tmp[0,0]:.4f}, [-1,-1]={lon_tmp[-1,-1]:.4f}") - print(f"[DEBUG calc_weights] input_forcings.ny_global={input_forcings.ny_global}, nx_global={input_forcings.nx_global}") # Normalize source longitudes from 0-360 to -180/+180 to match geo_em geogrid convention if lon_tmp.max() > 180: lon_tmp = np.where(lon_tmp > 180, lon_tmp - 360, lon_tmp) - print(f"[DEBUG calc_weights] lon_tmp normalized to -180/+180: min/max={lon_tmp.min():.4f} / {lon_tmp.max():.4f}") # Scatter global GFS latitude grid to processors.. if mpi_config.rank == 0: @@ -12089,11 +12069,6 @@ def calculate_weights( input_forcings.esmf_lats[:, :] = var_sub_lat_tmp input_forcings.esmf_lons[:, :] = var_sub_lon_tmp - if mpi_config.rank == 0: - print(f"[DEBUG ESMF src] esmf_lats (get_coords(1)) shape={input_forcings.esmf_lats.shape}") - print(f"[DEBUG ESMF src] esmf_lats min/max: {input_forcings.esmf_lats.min():.4f} / {input_forcings.esmf_lats.max():.4f}") - print(f"[DEBUG ESMF src] esmf_lons (get_coords(0)) shape={input_forcings.esmf_lons.shape}") - print(f"[DEBUG ESMF src] esmf_lons min/max: {input_forcings.esmf_lons.min():.4f} / {input_forcings.esmf_lons.max():.4f}") del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp From 105537838a5dec257ddaf318f0b343648320cbe9 Mon Sep 17 00:00:00 2001 From: Lauren Schambach Date: Wed, 5 Aug 2026 17:06:29 -0400 Subject: [PATCH 06/25] when type is gridded, does not delete the esmf mesh --- .../NextGen_Forcings_Engine/bmi_model.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 2c6b31f3..1ef179f8 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -303,6 +303,15 @@ 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. See + # consts.py's COASTAL_CONFIG_OVERRIDES in nwm-rte, which + # deliberately leaves "Geopackage" empty for gridded runs. + cat_ids = np.array([], dtype=np.int64) + else: cat_ids = esmf_creation.create_mesh(self._job_meta) cat_count = np.array( [len(cat_ids) if self._mpi_meta.rank == 0 else 0], dtype=np.intc From dfb0563300a09b3423426b8a370bf75579d5d329 Mon Sep 17 00:00:00 2001 From: Lauren Schambach Date: Wed, 5 Aug 2026 21:41:37 +0000 Subject: [PATCH 07/25] messed up syntax in last commit, fixing --- .../NextGen_Forcings_Engine/bmi_model.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 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 1ef179f8..910c9e3b 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -303,19 +303,18 @@ 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. See - # consts.py's COASTAL_CONFIG_OVERRIDES in nwm-rte, which - # deliberately leaves "Geopackage" empty for gridded runs. - cat_ids = np.array([], dtype=np.int64) - else: - cat_ids = esmf_creation.create_mesh(self._job_meta) - cat_count = np.array( - [len(cat_ids) if self._mpi_meta.rank == 0 else 0], dtype=np.intc - ) + 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. See + # consts.py's COASTAL_CONFIG_OVERRIDES in nwm-rte, which + # deliberately leaves "Geopackage" empty for gridded runs. + cat_ids = np.array([], dtype=np.int64) + else: + cat_ids = esmf_creation.create_mesh(self._job_meta) + cat_count = np.array([ + len(cat_ids) if self._mpi_meta.rank == 0 else 0 + ], dtype=np.intc) self._mpi_meta.comm.Bcast(cat_count, root=0) if self._mpi_meta.rank != 0: cat_ids = np.empty(cat_count[0], dtype=np.int64) From 6899f4bcc74a473772b575c55ff32918e6394126 Mon Sep 17 00:00:00 2001 From: Lauren Schambach Date: Tue, 11 Aug 2026 19:50:52 +0000 Subject: [PATCH 08/25] updates for forcing type gridded so that esmf weight files can be saved and re-used between coastal forecasts --- .../NextGen_Forcings_Engine/core/geoMod.py | 7 +++- .../NextGen_Forcings_Engine/core/regrid.py | 33 +++++++++++++++---- 2 files changed, 32 insertions(+), 8 deletions(-) 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 9f7629eb..3c953820 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -92,7 +92,12 @@ 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) var = self.mpi_config.scatter_array(self, var, config_options) if post_slice: 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 dc0fe18e..e801cdae 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -12415,13 +12415,32 @@ 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. From 7cd0c4480a8e3ece97a0d2475cb6b367aeddf6cb Mon Sep 17 00:00:00 2001 From: Lauren Schambach Date: Thu, 27 Aug 2026 13:44:47 -0400 Subject: [PATCH 09/25] reverting untested change so that it matches the tested coastalforcing-pw --- NextGen_Forcings_Engine_BMI/forcing_extraction.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/forcing_extraction.py b/NextGen_Forcings_Engine_BMI/forcing_extraction.py index 9e182898..84249418 100644 --- a/NextGen_Forcings_Engine_BMI/forcing_extraction.py +++ b/NextGen_Forcings_Engine_BMI/forcing_extraction.py @@ -110,14 +110,7 @@ def retrieve_forcing(config_options: ConfigOptions): forcing_start_time = forcing_start_time + timedelta(hours=1) if supp_forcing_hours is not None: - if supp_forcing_hours > 0 and input_horizons[i] > 0: - # Anchor end of window to the forecast horizon so all supp files - # are covered. ForecastDownloader extends forward via input_horizon, - # but FixedFileDownloader (e.g. MRMS) only looks backward. - forcing_start_time = refcstbdate + timedelta(minutes=input_horizons[i]) - else: - forcing_start_time += timedelta(hours=supp_forcing_hours) - lookback_hours += supp_forcing_hours + forcing_start_time += timedelta(hours=supp_forcing_hours) if ana_flag == 1: lookback_hours = lookback_hours + 1 From b86ec2c79da66054190ff64716dbdd9425c784b0 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 18:13:06 -0400 Subject: [PATCH 10/25] Allow gridded and hydrofabric --- .../NextGen_Forcings_Engine/bmi_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 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 910c9e3b..f6a7b077 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -273,8 +273,8 @@ 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(self._job_meta, self._mpi_meta) return self._geo_meta From 9bcf0f1236d18aa41d498dedfa9c4cf9bcd7c5c9 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 20:42:48 -0400 Subject: [PATCH 11/25] Fix CAT-ID usage --- .../NextGen_Forcings_Engine/bmi_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f6a7b077..4716df31 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -509,7 +509,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 From f55348df74c821b7f20390f19ec3e6cfed32bf64 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 20:43:57 -0400 Subject: [PATCH 12/25] Fix GeoMod constructor --- .../NextGen_Forcings_Engine/bmi_model.py | 4 +++- .../NextGen_Forcings_Engine/core/geoMod.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) 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 4716df31..f8f3406c 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -276,7 +276,9 @@ def geo_meta(self) -> GeoMeta: 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(self._job_meta, self._mpi_meta) + self._geo_meta = GeoMeta.for_grid_type( + self._job_meta.grid_type, self._job_meta, self._mpi_meta + ) return self._geo_meta @geo_meta.setter 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 3c953820..3d2065a6 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -128,6 +128,22 @@ 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.""" From 179007dd9a8c15d216002e91338e65a60440c366 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 20:45:35 -0400 Subject: [PATCH 13/25] Format with ruff line length 88 --- .../NextGen_Forcings_Engine/bmi_model.py | 6 ++--- .../NextGen_Forcings_Engine/core/geoMod.py | 23 +++++++++++++++---- .../NextGen_Forcings_Engine/core/regrid.py | 10 ++++++-- 3 files changed, 29 insertions(+), 10 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 f8f3406c..9c7d17a9 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -314,9 +314,9 @@ def create_esmf_mesh(self) -> None: cat_ids = np.array([], dtype=np.int64) else: cat_ids = esmf_creation.create_mesh(self._job_meta) - cat_count = np.array([ - len(cat_ids) if self._mpi_meta.rank == 0 else 0 - ], dtype=np.intc) + cat_count = np.array( + [len(cat_ids) if self._mpi_meta.rank == 0 else 0], dtype=np.intc + ) self._mpi_meta.comm.Bcast(cat_count, root=0) if self._mpi_meta.rank != 0: cat_ids = np.empty(cat_count[0], dtype=np.int64) 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 3d2065a6..b7d3d96c 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -141,7 +141,9 @@ def for_grid_type( 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 + raise ValueError( + f"Unsupported grid type: {grid_type}, expected one of: {list(lookup.keys())}" + ) from error return _class(config_options, mpi_config) @cached_property @@ -155,10 +157,21 @@ def spatial_metadata_exists(self) -> bool: # All config_options attributes that hold geogrid variable names. # Used to select only the needed variables when opening the geogrid file. _GEOGRID_VAR_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", + "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", ) @cached_property 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 e801cdae..d9ca3302 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -12419,7 +12419,9 @@ def calculate_supp_pcp_weights( # 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) + 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( @@ -12431,7 +12433,11 @@ def calculate_supp_pcp_weights( filename=weight_file, ) load_weight_file( - mpi_config, config_options, supplemental_precip, weight_file, element_mode=False + mpi_config, + config_options, + supplemental_precip, + weight_file, + element_mode=False, ) else: supplemental_precip.regridObj = pt.esmf_regrid_retry_partial( From 9ebbf112f75169baf8044bc4e7f6ffabfde8c980 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 21:29:10 -0400 Subject: [PATCH 14/25] Fix None-shadowing of GeoMeta esmf_lat and esmf_lon --- .../NextGen_Forcings_Engine/core/consts.py | 2 -- 1 file changed, 2 deletions(-) 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..ede385c1 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py @@ -139,8 +139,6 @@ "nodeCoords", "centerCoords", "inds", - "esmf_lat", - "esmf_lon", ], "UnstructuredGeoMeta": [ "x_lower_bound", From 4536abec81f5617fad896ba85ba0121746cd6219 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 21:30:49 -0400 Subject: [PATCH 15/25] Revert regrid.py changes that had been patching a deeper None-shadowing issue. --- .../NextGen_Forcings_Engine/core/regrid.py | 11 ----------- 1 file changed, 11 deletions(-) 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 d9ca3302..7166d2b2 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -11283,17 +11283,6 @@ def check_regrid_status( """ pt = Partials(mpi_config, config_options) - # For gridded domains, esmf_lat/esmf_lon on GriddedGeoMeta are shadowed to None by - # GeoMeta.__init__ (via GEOMOD["GeoMeta"]), so the @cached_property never runs and the - # destination ESMF Grid's coordinate arrays are never populated. Populate them once here. - if config_options.grid_type == "gridded" and wrf_hydro_geo_meta.esmf_lat is None: - esmf_lat = wrf_hydro_geo_meta.esmf_grid.get_coords(1) - esmf_lat[:, :] = wrf_hydro_geo_meta.latitude_grid - wrf_hydro_geo_meta.esmf_lat = esmf_lat - esmf_lon = wrf_hydro_geo_meta.esmf_grid.get_coords(0) - esmf_lon[:, :] = wrf_hydro_geo_meta.longitude_grid - wrf_hydro_geo_meta.esmf_lon = esmf_lon - # If the destination ESMF field hasn't been created, create it here. if not input_forcings.esmf_field_out: if config_options.grid_type == "gridded": From 250e83cb51d9028fecc53edc7572b2330dd6827b Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 21:38:36 -0400 Subject: [PATCH 16/25] Assert types --- .../NextGen_Forcings_Engine/core/geoMod.py | 2 ++ 1 file changed, 2 insertions(+) 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 b7d3d96c..83a67840 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -98,6 +98,8 @@ def wrapper(self) -> Any: # 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: From 4d3cdbdd51a8468d7057ccba8b66b609c31ff4eb Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 21:45:07 -0400 Subject: [PATCH 17/25] Preserve caller-supplied geogrids during cleanup (don't just look at grid type to control this conditional) and update README about this. --- .../NextGen_Forcings_Engine/core/parallel.py | 10 +++++++--- NextGen_Forcings_Engine_BMI/README.md | 4 ++++ 2 files changed, 11 insertions(+), 3 deletions(-) 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 4eef6293..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,13 +161,17 @@ 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 getattr(self.config_options, "grid_type", None) == "gridded": - self.log_debug("Cleanup: skipping geogrid cleanup for gridded mode") + 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: 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. From da9658fefaa2c56da3d9d14bde0f357c9f6bd846 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Tue, 15 Sep 2026 21:48:15 -0400 Subject: [PATCH 18/25] Harden longitude normalization --- .../NextGen_Forcings_Engine/core/regrid.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 7166d2b2..2825ebb7 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -12021,9 +12021,8 @@ 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 from 0-360 to -180/+180 to match geo_em geogrid convention - if lon_tmp.max() > 180: - lon_tmp = np.where(lon_tmp > 180, lon_tmp - 360, lon_tmp) + # 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: From a4eb884c1c13f438fb68ed4d95087c851e14821d Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 16 Sep 2026 09:56:10 -0400 Subject: [PATCH 19/25] Move input data variables list needed for gridded mode out to consts.py and rename it --- .../NextGen_Forcings_Engine/core/consts.py | 20 ++++++++++++++ .../NextGen_Forcings_Engine/core/geoMod.py | 27 ++++--------------- 2 files changed, 25 insertions(+), 22 deletions(-) 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 ede385c1..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", 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 83a67840..9872f4f2 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, ) @@ -156,26 +159,6 @@ def spatial_metadata_exists(self) -> bool: else: return True - # All config_options attributes that hold geogrid variable names. - # Used to select only the needed variables when opening the geogrid file. - _GEOGRID_VAR_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", - ) - @cached_property def geogrid_ds(self) -> xr.Dataset: """Open the geogrid file and load only the variables needed by this run. @@ -189,7 +172,7 @@ def geogrid_ds(self) -> xr.Dataset: with xr.open_dataset(self.config_options.geogrid) as ds: needed = [ getattr(self.config_options, attr) - for attr in self._GEOGRID_VAR_ATTRS + 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 ] From 24cc84028fd4f8b6b155093f26f47bf984db084e Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 16 Sep 2026 10:05:56 -0400 Subject: [PATCH 20/25] Add defensiveness to esmf mesh creation logic --- .../NextGen_Forcings_Engine/bmi_model.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 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 9c7d17a9..b3eac916 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -308,12 +308,14 @@ def create_esmf_mesh(self) -> None: 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. See - # consts.py's COASTAL_CONFIG_OVERRIDES in nwm-rte, which - # deliberately leaves "Geopackage" empty for gridded runs. + # requires a hydrofabric geopackage) does not apply here. cat_ids = np.array([], dtype=np.int64) - else: + 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 ) From 27d5dfa301868e48e65c7330dd63097bee624c48 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 16 Sep 2026 15:55:42 -0400 Subject: [PATCH 21/25] Fix gridded esmf output by assigning missing coordinates --- .../NextGen_Forcings_Engine/core/geoMod.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 9872f4f2..048f2e87 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -406,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, @@ -414,6 +414,10 @@ 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. + 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: From 310ed39775c06d9774c938da71524d6d775bcbe9 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Thu, 17 Sep 2026 10:47:52 -0400 Subject: [PATCH 22/25] Gridded forcing: to support coastal use case, add overridable output_steps argument to ConfigOptions and BMI model classes, and for scalability, avoid randomizing name of intermediary weights file to allow re-use. --- .../NextGen_Forcings_Engine/bmi_model.py | 26 +++++++++- .../NextGen_Forcings_Engine/core/config.py | 51 +++++++++++++++---- .../NextGen_Forcings_Engine/core/regrid.py | 3 +- 3 files changed, 68 insertions(+), 12 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 b3eac916..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) @@ -1575,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/regrid.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py index 2825ebb7..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") From 0d7c8ac6dbb5126d7f2045b5a1e89c5443319fc5 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Thu, 17 Sep 2026 17:17:20 -0400 Subject: [PATCH 23/25] Fix infinite recursion --- .../NextGen_Forcings_Engine/core/geoMod.py | 3 +++ 1 file changed, 3 insertions(+) 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 048f2e87..91822785 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py @@ -415,6 +415,9 @@ def esmf_grid(self) -> ESMF.Grid: 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 From 27e03ba138478dd449ce957b1c3ba4c36fc33d8f Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 18 Sep 2026 11:19:54 -0400 Subject: [PATCH 24/25] Add support for coastal gridded forcing (edit run_bmi_model.py, add new coastal gridded module) --- .../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 ++++++++++++++++-- 6 files changed, 810 insertions(+), 30 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/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(): From b5293d20720d09aa42c7012f4a2d38d936e88c2e Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 18 Sep 2026 17:46:59 -0400 Subject: [PATCH 25/25] 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()