Coastalforcing-pw

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", "%"]

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[:]

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)

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.