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..3c2d814e 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/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]() 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 2fd37fe1..b1fd02db 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py +++ b/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", diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/downscale.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/downscale.py index cdec3320..51282ebb 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/downscale.py +++ b/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): 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..49404a77 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/geoMod.py +++ b/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 diff --git a/NextGen_Forcings_Engine_BMI/esmf_creation.py b/NextGen_Forcings_Engine_BMI/esmf_creation.py index 2dd82ef3..8bc51653 100644 --- a/NextGen_Forcings_Engine_BMI/esmf_creation.py +++ b/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) diff --git a/NextGen_Forcings_Engine_BMI/run_bmi_model.py b/NextGen_Forcings_Engine_BMI/run_bmi_model.py index 70c56d52..5ffa3eb1 100755 --- a/NextGen_Forcings_Engine_BMI/run_bmi_model.py +++ b/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