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 3c27eacb..c4351bcd 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -178,6 +178,7 @@ def __init__(self): self._job_meta = None self._mpi_meta = None self.geo_meta = None + self.GeoMeta = HydrofabricGeoMeta self._grid_type = None self._grids = None self._grid_map = None @@ -290,7 +291,7 @@ def initialize(self, config_file: str, output_path: str | None = None) -> None: # Initialize MPI communication self._mpi_meta = MpiConfig(self._job_meta) - self.geo_meta = HydrofabricGeoMeta(self._job_meta, self._mpi_meta) + self.geo_meta = self.GeoMeta(self._job_meta, self._mpi_meta) try: comm = MPI.Comm.f2py(self._comm) if self._comm is not None else None @@ -305,7 +306,15 @@ def initialize(self, config_file: str, output_path: str | None = None) -> None: # LOG.debug(f"self._job_meta type: {type(self._job_meta)}") # Call ESMF mesh creation process 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) @@ -1665,7 +1674,7 @@ def set_var_names(self) -> None: # will support a BMI field for liquid fraction of precipitation self._output_var_names = BMI_MODEL["_output_var_names"] self._var_name_units_map = BMI_MODEL["_var_name_units_map"] - if self.config_options.include_lqfrac == 1: + if self._job_meta.include_lqfrac == 1: self._output_var_names += ["LQFRAC_ELEMENT"] self._var_name_units_map |= { "LQFRAC_ELEMENT": ["Liquid Fraction of Precipitation", "%"] 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 90a225cf..880b852b 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: @@ -131,12 +136,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) @@ -274,12 +300,13 @@ def nx_global(self) -> int: if self.mpi_config.rank == 0: try: if self.ndim_lat == 3: - return self.lat_var.shape[2] + nx = self.lat_var.shape[2] elif self.ndim_lat == 2: - return self.lat_var.shape[1] + nx = self.lat_var.shape[1] else: # NOTE Is this correct? using lon_var - return self.lon_var.shape[0] + nx = self.lon_var.shape[0] + return nx except Exception as e: self.config_options.errMsg = f"Unable to extract X dimension size from {self.config_options.lon_var} in: {self.config_options.geogrid}" log_critical(self.config_options, self.mpi_config) @@ -387,9 +414,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[:] @@ -421,9 +448,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 e0e48891..d63ed9a0 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/parallel.py @@ -239,6 +239,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) 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 5dd3a109..354af976 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/regrid.py @@ -11276,6 +11276,17 @@ 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": @@ -12017,6 +12028,11 @@ 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) + # Scatter global GFS latitude grid to processors.. if mpi_config.rank == 0: var_tmp = lat_tmp @@ -12396,13 +12412,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.