adjustments to support gridded forcing

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py

@@ -2,14 +2,15 @@
 # This is needed for get_var_bytes
 import gc
 import hashlib
+import logging
 import os
 
 # time debugging
 import time
 from collections import defaultdict
-from pathlib import Path
 from datetime import datetime, timezone
-import logging
+from pathlib import Path
+
 import netCDF4 as nc
 
 # import data_tools
@@ -68,15 +69,14 @@
 try:
     from ewts.helper import getenv_any
     from ewts.logger import configure_existing_logger
+
     FORCING_USE_EWTS = True
 except ImportError:
     FORCING_USE_EWTS = False
 
-class StdoutStyleFormatter(logging.Formatter):
 
-    INFO_FORMAT = (
-        "%(asctime)s %(name)-8s %(levelname)-7s %(message)s"
-    )
+class StdoutStyleFormatter(logging.Formatter):
+    INFO_FORMAT = "%(asctime)s %(name)-8s %(levelname)-7s %(message)s"
 
     DETAILED_FORMAT = (
         "%(asctime)s %(name)-8s %(levelname)-7s "
@@ -91,7 +91,7 @@ def format(self, record):
             self._style._fmt = self.DETAILED_FORMAT
 
         return super().format(record)
-    
+
     def formatTime(self, record, datefmt=None):
         dt = datetime.fromtimestamp(record.created, tz=timezone.utc)
         return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
@@ -108,6 +108,7 @@ def _configure_stdout_logging():
 
     LOG.propagate = False
 
+
 # If less than 0, then ESMF.__version__ is greater than 8.7.0
 if ESMF.version_compare("8.7.0", ESMF.__version__) < 0:
     manager = ESMF.api.esmpymanager.Manager(endFlag=ESMF.constants.EndAction.KEEP_MPI)
@@ -158,7 +159,9 @@ def __init__(self):
                 configure_existing_logger(LOG)
             else:
                 _configure_stdout_logging()
-                LOG.warning("ewts package installed but EWTS_USE_NGEN_BRIDGE not on. Falling back to default logging.")
+                LOG.warning(
+                    "ewts package installed but EWTS_USE_NGEN_BRIDGE not on. Falling back to default logging."
+                )
         else:
             _configure_stdout_logging()
 
@@ -236,7 +239,6 @@ def initialize(self, config_file: str, output_path: str | None = None) -> None:
         :param config_file: The path to the configuration file for model initialization.
         :raises RuntimeError: If the configuration file is invalid or missing.
         """
-
         LOG.info("---------------------------")
         LOG.info(
             f"BMI Forcing Engine initializing with {config_file}{Pld(St.INITTING, modnm=MODNM)}"
@@ -290,7 +292,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
@@ -304,15 +306,7 @@ 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)
-        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)
-        self._mpi_meta.comm.Bcast(cat_ids, root=0)
+        self._values["CAT-ID"] = self.cat_ids
 
         # Call forcing_extraction process
         if self._job_meta.nwmConfig not in ["AORC", "NWM"]:
@@ -408,8 +402,6 @@ def initialize(self, config_file: str, output_path: str | None = None) -> None:
         # Set initial time, step, and true catchment IDs
         self._values["current_model_time"] = self.cfg_bmi["initial_time"]
         self._values["time_step_size"] = self.cfg_bmi["time_step_seconds"]
-        self._values["CAT-ID"] = cat_ids
-
         # Initialize the Forcings Engine model
         self._model = NWMv3ForcingEngineModel()
 
@@ -485,7 +477,10 @@ def _configure_output_path(self, output_path: str | None = None) -> 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._values.get("CAT-ID"),
             )
             self._output_configured = True
 
@@ -782,7 +777,7 @@ def get_value_ptr(self, var_name: str) -> NDArray[Any]:
 
         # Ensure dtype is float64 (C double), except for CAT-ID
         if var_name == "CAT-ID":
-            return arr # allow CAT-ID to pass on whatever the dtype is based on the input data
+            return arr  # allow CAT-ID to pass on whatever the dtype is based on the input data
         elif arr.dtype != np.float64:
             LOG.warning(
                 f"[BMI] Array for '{var_name}' has dtype {arr.dtype}, expected float64; converting."
@@ -1637,6 +1632,7 @@ def __init__(self):
         """
         super().__init__()
         self.GeoMeta = GriddedGeoMeta
+        self.cat_ids = None
 
     def grid_ranks(self) -> list[int]:
         """Get the grid ranks for the gridded domain."""
@@ -1665,7 +1661,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", "%"]
@@ -1704,6 +1700,20 @@ def __init__(self):
         super().__init__()
         self.GeoMeta = HydrofabricGeoMeta
 
+    @property
+    def cat_ids(self) -> NDArray[np.int_]:
+        """Get the catchment IDs for the hydrofabric domain."""
+        if self._mpi_meta.rank == 0:
+            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)
+        self._mpi_meta.comm.Bcast(cat_ids, root=0)
+        return cat_ids
+
     def grid_ranks(self) -> list[int]:
         """Get the grid ranks for the hydrofabric domain."""
         return [self.grid_4.rank]
@@ -1765,6 +1775,7 @@ def __init__(self):
         """
         super().__init__()
         self.GeoMeta = UnstructuredGeoMeta
+        self.cat_ids = None
 
     def grid_ranks(self) -> list[int]:
         """Get the grid ranks for the unstructured domain."""
@@ -1858,12 +1869,14 @@ def set_var_names(self) -> None:
         self._grids = [self.grid_2, self.grid_3]
 
 
-BMIMODEL = {
-    "gridded": NWMv3_Forcing_Engine_BMI_model_Gridded,
-    "unstructured": NWMv3_Forcing_Engine_BMI_model_Unstructured,
-    "hydrofabric": NWMv3_Forcing_Engine_BMI_model_HydroFabric,
-}
+class NWMv3_Forcing_Engine_BMI_model:
+    """Factory class for creating instances of the NWMv3 Forcing Engine BMI model based on the specified grid type."""
 
-### NOTE patch so ngen always accesses the Hydrofabric child for now.
-### Other discretization modes currently do not have a ngen workflow.
-NWMv3_Forcing_Engine_BMI_model = NWMv3_Forcing_Engine_BMI_model_HydroFabric
+    def __new__(cls, grid_type: str = "hydrofabric"):
+        """Create an instance of the NWMv3 Forcing Engine BMI model based on the specified grid type."""
+        bmi_model = {
+            "gridded": NWMv3_Forcing_Engine_BMI_model_Gridded,
+            "unstructured": NWMv3_Forcing_Engine_BMI_model_Unstructured,
+            "hydrofabric": NWMv3_Forcing_Engine_BMI_model_HydroFabric,
+        }
+        return bmi_model[grid_type]()

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py

@@ -120,8 +120,6 @@
         "nodeCoords",
         "centerCoords",
         "inds",
-        "esmf_lat",
-        "esmf_lon",
     ],
     "UnstructuredGeoMeta": [
         "x_lower_bound",

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/downscale.py

@@ -291,7 +291,7 @@ def param_lapse(
         if mpi_config.rank == 0:
             while True:
                 # First ensure we have a parameter directory
-                if input_forcings.paramDir == "NONE":
+                if input_forcings.dScaleParamDirs == "NONE":
                     config_options.errMsg = (
                         "User has specified spatial temperature lapse rate "
                         "downscaling while no downscaling parameter directory "
@@ -301,7 +301,7 @@ def param_lapse(
                     break
 
                 # Compose the path to the lapse rate grid file.
-                lapsePath = f"{input_forcings.paramDir}/lapse_param.nc"
+                lapsePath = f"{input_forcings.dScaleParamDirs}/lapse_param.nc"
                 if not os.path.isfile(lapsePath):
                     ConfigOptions.errMsg = f"Expected lapse rate parameter file: {lapsePath} does not exist."
                     err_handler.log_critical(config_options, mpi_config)
@@ -717,13 +717,13 @@ def nwm_monthly_PRISM_downscale(
 
             if mmVersion == 1:
                 # Compose paths to the expected files.
-                numeratorPath = f"{input_forcings.paramDir}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_Grid.nc"
-                denominatorPath = f"{input_forcings.paramDir}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_to_{keyValueStr!s}_Grid.nc"
+                numeratorPath = f"{input_forcings.dScaleParamDirs}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_Grid.nc"
+                denominatorPath = f"{input_forcings.dScaleParamDirs}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_to_{keyValueStr!s}_Grid.nc"
 
             elif mmVersion == 2:
                 # Compose paths to the expected files.
-                numeratorPath = f"{input_forcings.paramDir}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_Grid.nc"
-                denominatorPath = f"{input_forcings.paramDir}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_{keyValueStr!s}_to_NWM_Grid.nc"
+                numeratorPath = f"{input_forcings.dScaleParamDirs}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_NWM_Grid.nc"
+                denominatorPath = f"{input_forcings.dScaleParamDirs}/PRISM_Precip_Clim_{ConfigOptions.current_output_date.strftime('%b')}_{keyValueStr!s}_to_NWM_Grid.nc"
 
             # Make sure files exist.
             if not os.path.isfile(numeratorPath):

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py

@@ -89,10 +89,13 @@ def wrapper(self) -> Any:
         """
         try:
             var, name, config_options, post_slice = prop.func(self)
+            if isinstance(var, (xr.core.variable.Variable, xr.Dataset)):
+                var = var.values
+
             assert isinstance(post_slice, bool)
             assert isinstance(name, str)
             assert isinstance(config_options, ConfigOptions)
-            assert isinstance(var, np.ndarray)
+            assert isinstance(var, (np.ndarray, type(None)))
 
             var = self.mpi_config.scatter_array(self, var, config_options)
             if post_slice:
@@ -136,7 +139,7 @@ def geogrid_ds(self) -> xr.Dataset:
         """Open the geogrid file and return the xarray dataset object."""
         try:
             with xr.open_dataset(self.config_options.geogrid) as ds:
-                return ds.load()
+                return ds
         except Exception as e:
             self.config_options.errMsg = "Unable to open geogrid file with xarray"
             log_critical(self.config_options, self.mpi_config)
@@ -148,7 +151,7 @@ def esmf_ds(self) -> xr.Dataset:
         """Open the geospatial metadata file and return the xarray dataset object."""
         try:
             with xr.open_dataset(self.config_options.spatial_meta) as ds:
-                esmf_ds = ds.load()
+                esmf_ds = ds
         except Exception as e:
             self.config_options.errMsg = (
                 f"Unable to open esmf file: {self.config_options.spatial_meta}"
@@ -967,7 +970,7 @@ class UnstructuredGeoMeta(GeoMeta):
     """Class for handling information about the hydrofabric domain forcing."""
 
     def __init__(self, config_options: ConfigOptions, mpi_config: MpiConfig) -> None:
-        """Initialize HydrofabricGeoMeta class variables.
+        """Initialize Unstructured GeoMeta class variables.
 
         Initialization function to initialize ESMF through ESMPy,
         calculate the global parameters of the unstructured mesh

NextGen_Forcings_Engine_BMI/esmf_creation.py

@@ -1,9 +1,9 @@
 import argparse
 from pathlib import Path
 from types import SimpleNamespace
-import numpy as np
-import geopandas as gpd
 
+import geopandas as gpd
+import numpy as np
 import yaml
 from NextGen_Forcings_Engine.core.config import ConfigOptions
 
@@ -29,8 +29,8 @@ def create_mesh(cfg: ConfigOptions):
         # The generation will sort the IDs,
         # so return the sorted IDs from the geopackage
         # to maintain the true->false ID indexing
-        hyfab = gpd.read_file(hyfab_name, layer='divides')
-        return np.sort(hyfab.div_id.values, copy=True, dtype=np.int64)
+        hyfab = gpd.read_file(hyfab_name, layer="divides")
+        return np.array(np.sort(hyfab.div_id.values), copy=True,dtype=np.int64)
     return convert_hyfab_to_esmf(hyfab_gpkg=hyfab_name, esmf_mesh_output=mesh_out_path)
 
 

NextGen_Forcings_Engine_BMI/run_bmi_model.py

@@ -9,7 +9,6 @@
 
 # This is the NextGen Forcings Engine BMI instance to execute
 from NextGen_Forcings_Engine.bmi_model import (
-    BMIMODEL,
     NWMv3_Forcing_Engine_BMI_model,
     parse_config,
 )
@@ -343,7 +342,7 @@ def run_bmi(
         config = parse_config(yaml.safe_load(fp))
 
     print("Creating an instance of the BMI model object")
-    model = BMIMODEL[config.get("GRID_TYPE")]()
+    model = NWMv3_Forcing_Engine_BMI_model(config.get("GRID_TYPE"))
 
     # IMPORTANT: We are not calling initialize() directly here.
     # Instead, we call initialize_with_params(), which handles