Mirror PR #212 for PW (Add Coastal Forcing)

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py

@@ -140,21 +140,43 @@ 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__(
         self,
         b_date: str = None,
         geogrid: str = None,
         output_path: str = None,
+        output_steps: int = None,
+        output_t0: bool = False,
     ) -> 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.
+            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
@@ -236,7 +258,11 @@ 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,
+                    output_t0=self._output_t0,
                 )
             except KeyboardInterrupt as e:
                 err_handler.err_out_screen("User keyboard interrupt", e)
@@ -273,10 +299,12 @@ 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.for_grid_type(
+                self._job_meta.grid_type, self._job_meta, self._mpi_meta
             )
-            self._geo_meta = HydrofabricGeoMeta(self._job_meta, self._mpi_meta)
         return self._geo_meta
 
     @geo_meta.setter
@@ -303,7 +331,17 @@ 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.
+                cat_ids = np.array([], dtype=np.int64)
+            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
         )
@@ -501,7 +539,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
 
@@ -1563,12 +1601,20 @@ def __init__(
         b_date: str = None,
         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)
+        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."""
@@ -1633,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."""
@@ -1699,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."""

NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/config.py

@@ -31,7 +31,14 @@
 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,
+        output_t0: bool = False,
+    ) -> 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 +47,31 @@ 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 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
 
         """
-        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
+
+        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:
@@ -58,6 +84,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 +106,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 +119,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,12 +1232,33 @@ 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.
-        """
-        if self.ana_flag:
-            return np.int32(self.nFcsts)
+        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_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"
+                )
+            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:
+        """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:

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",
@@ -139,8 +159,6 @@
         "nodeCoords",
         "centerCoords",
         "inds",
-        "esmf_lat",
-        "esmf_lon",
     ],
     "UnstructuredGeoMeta": [
         "x_lower_bound",