From 9a95a63b804a3844965b5f319c87c815b74efe92 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 24 Apr 2026 15:30:47 -0400 Subject: [PATCH 01/36] Simplify var passing within NWMv3ForcingEngineModel.run --- .../NextGen_Forcings_Engine/model.py | 172 ++++++------------ tests/test_utils.py | 40 +--- 2 files changed, 65 insertions(+), 147 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 2c27669e..84d2973a 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -12,6 +12,7 @@ disaggregateMod, downscale, err_handler, + forcingInputMod, layeringMod, ) from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.config import ( @@ -123,39 +124,12 @@ def run( :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ - ( - future_time, - config_options, - ) = self.determine_forecast( - future_time, - config_options, - ) - ( - config_options, - input_forcing_mod, - mpi_config, - ) = self.adjust_precip( - config_options, - input_forcing_mod, - mpi_config, - ) - ( - config_options, - mpi_config, - ) = self.log_forecast( - config_options, - mpi_config, - ) - ( - future_time, - config_options, - wrf_hydro_geo_meta, - input_forcing_mod, - supp_pcp_mod, - mpi_config, - output_obj, - input_forcings, - ) = self.loop_through_forcing_products( + + self.determine_forecast(future_time, config_options) + self.adjust_precip(config_options, input_forcing_mod, mpi_config) + self.log_forecast(config_options, mpi_config) + # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. + input_forcings = self.loop_through_forcing_products( future_time, config_options, wrf_hydro_geo_meta, @@ -164,13 +138,7 @@ def run( mpi_config, output_obj, ) - ( - config_options, - wrf_hydro_geo_meta, - supp_pcp_mod, - mpi_config, - output_obj, - ) = self.process_suplemental_precip( + self.process_suplemental_precip( config_options, wrf_hydro_geo_meta, supp_pcp_mod, @@ -178,23 +146,13 @@ def run( output_obj, input_forcings, ) - ( - config_options, - wrf_hydro_geo_meta, - mpi_config, - output_obj, - ) = self.write_output( + self.write_output( config_options, wrf_hydro_geo_meta, mpi_config, output_obj, ) - ( - model, - config_options, - wrf_hydro_geo_meta, - output_obj, - ) = self.update_dict( + self.update_dict( model, config_options, wrf_hydro_geo_meta, @@ -209,8 +167,13 @@ def determine_forecast( self, future_time: float, config_options: ConfigOptions, - ): - """Determine the forecast for the given future time and configuration.""" + ) -> None: + """Determine the forecast for the given future time and configuration. + + Warnings + -------- + Modifies mutable arguments in-place. + """ # Assign the future time to the configuration config_options.bmi_time = future_time self.disaggregate_fun = disaggregateMod.disaggregate_factory(config_options) @@ -264,38 +227,38 @@ def determine_forecast( if config_options.first_fcst_cycle is None: config_options.first_fcst_cycle = config_options.current_fcst_cycle - return ( - future_time, - config_options, - ) - @time_function def adjust_precip( self, config_options: ConfigOptions, input_forcing_mod: dict, mpi_config: MpiConfig, - ): - """Adjust precipitation for the given forecast cycle.""" + ) -> None: + """Adjust precipitation for the given forecast cycle. + + Warnings + -------- + Modifies mutable arguments in-place. + """ if not config_options.precip_only_flag: # reset skips if present for force_key in config_options.input_forcings: input_forcing_mod[force_key].skip = False err_handler.check_program_status(config_options, mpi_config) - return ( - config_options, - input_forcing_mod, - mpi_config, - ) @time_function def log_forecast( self, config_options: ConfigOptions, mpi_config: MpiConfig, - ): - """Log information about the current forecast cycle.""" + ) -> None: + """Log information about the current forecast cycle. + + Warnings + -------- + Modifies mutable arguments in-place. + """ # Log information about this forecast cycle if mpi_config.rank == 0: config_options.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" @@ -313,11 +276,6 @@ def log_forecast( err_handler.log_msg(config_options, mpi_config, True) # mpi_config.comm.barrier() - return ( - config_options, - mpi_config, - ) - @time_function def loop_through_forcing_products( self, @@ -328,8 +286,13 @@ def loop_through_forcing_products( supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, - ): - """Loop through each forcing product and process it for the current forecast cycle.""" + ) -> forcingInputMod.InputForcingsHydrofabric: + """Loop through each forcing product and process it for the current forecast cycle. + + Warnings + -------- + Modifies mutable arguments in-place. + """ # Loop through each output timestep. Perform the following functions: # 1.) Calculate all necessary input files per user options. # 2.) Read in input forcings from GRIB/NetCDF files. @@ -625,16 +588,7 @@ def loop_through_forcing_products( # err_handler.check_program_status(config_options, mpi_config) ############################################################################################## - return ( - future_time, - config_options, - wrf_hydro_geo_meta, - input_forcing_mod, - supp_pcp_mod, - mpi_config, - output_obj, - input_forcings, - ) + return input_forcings @time_function def process_suplemental_precip( @@ -645,8 +599,13 @@ def process_suplemental_precip( mpi_config: MpiConfig, output_obj: OutputObj, input_forcings: dict, - ): - """Process supplemental precipitation for the current forecast cycle.""" + ) -> None: + """Process supplemental precipitation for the current forecast cycle. + + Warnings + -------- + Modifies mutable arguments in-place. + """ if config_options.customSuppPcpFreq is not None: # Process supplemental precipitation if we specified in the configuration file. if config_options.number_supp_pcp > 0: @@ -700,14 +659,6 @@ def process_suplemental_precip( ) err_handler.check_program_status(config_options, mpi_config) - return ( - config_options, - wrf_hydro_geo_meta, - supp_pcp_mod, - mpi_config, - output_obj, - ) - @time_function def write_output( self, @@ -715,8 +666,13 @@ def write_output( wrf_hydro_geo_meta: GeoMeta, mpi_config: MpiConfig, output_obj: OutputObj, - ): - """Write the output for the current forecast cycle.""" + ) -> None: + """Write the output for the current forecast cycle. + + Warnings + -------- + Modifies mutable arguments in-place. + """ # If user requests output for given domain, then call # the I/O module to update opened netcdf file with forcing fields if ( @@ -726,12 +682,6 @@ def write_output( output_obj.gather_global_outputs( config_options, wrf_hydro_geo_meta, mpi_config ) - return ( - config_options, - wrf_hydro_geo_meta, - mpi_config, - output_obj, - ) """##################Step 6: flatten and update dict##########################################################################""" @@ -742,8 +692,13 @@ def update_dict( config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, output_obj: OutputObj, - ): - """Flatten the Forcings Engine output object and update the BMI dictionary.""" + ) -> None: + """Flatten the Forcings Engine output object and update the BMI dictionary. + + Warnings + -------- + Modifies mutable arguments in-place. + """ # Now loop through Forcings Engine output object # and flatten the 2D forcing array and append to # the BMI object to advertise to BMIinterface @@ -797,10 +752,3 @@ def update_dict( count, : ].flatten() model["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global - - return ( - model, - config_options, - wrf_hydro_geo_meta, - output_obj, - ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 513f4a76..11496c02 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -540,43 +540,13 @@ def pre_regrid(self) -> None: ) model = self.bmi_model._model - ### NOTE with the exception of setting the skip flag, the below - ### block is copied verbatim from NWMv3ForcingEngineModel.run() - ( - future_time, - config_options, - ) = model.determine_forecast( - future_time, - config_options, - ) - ( - config_options, - input_forcing_mod, - mpi_config, - ) = model.adjust_precip( - config_options, - input_forcing_mod, - mpi_config, - ) - ( - config_options, - mpi_config, - ) = model.log_forecast( - config_options, - mpi_config, - ) + ### NOTE this should mimic NWMv3ForcingEngineModel.run() with the exception of setting the skip flag + model.determine_forecast(future_time, config_options) + model.adjust_precip(config_options, input_forcing_mod, mpi_config) + model.log_forecast(config_options, mpi_config) ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() - ( - future_time, - config_options, - geo_meta, - input_forcing_mod, - supp_pcp_mod, - mpi_config, - output_obj, - input_forcings, - ) = model.loop_through_forcing_products( + model.loop_through_forcing_products( future_time, config_options, geo_meta, From 5a6e0dfbc9591751a8f2bb66957cb6bbffab992a Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 10:10:43 -0400 Subject: [PATCH 02/36] Raise error on unexpected grid_type --- .../NextGen_Forcings_Engine/model.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 84d2973a..fba52479 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -312,6 +312,10 @@ def loop_through_forcing_products( elif config_options.grid_type == "hydrofabric": # Reset out final grids to missing values. output_obj.output_local[:, :] = config_options.globalNdv + else: + raise ValueError( + f"Unexpected grid_type: {repr(config_options.grid_type)}" + ) # Increment or initialize output step count if config_options.current_output_step is None: @@ -482,6 +486,10 @@ def loop_through_forcing_products( input_forcings.regridded_forcings1[:, :] = ( input_forcings.regridded_forcings2[:, :] ) + else: + raise ValueError( + f"Unexpected grid_type: {repr(config_options.grid_type)}" + ) # Re-calculate the neighbor files. input_forcings.calc_neighbor_files( config_options, output_obj.outDate, mpi_config @@ -752,3 +760,5 @@ def update_dict( count, : ].flatten() model["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global + else: + raise ValueError(f"Unexpected grid_type: {repr(config_options.grid_type)}") From 2b0e9ce471787603643d8fefbd8bfaedab13a4f2 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 12:10:35 -0400 Subject: [PATCH 03/36] Use `time.perf_counter` instead of `time.time` --- .../NextGen_Forcings_Engine/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index fba52479..c7121a91 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,7 +1,7 @@ import datetime import os from contextlib import contextmanager -from time import time +from time import time, perf_counter import ewts import numpy as np @@ -42,9 +42,9 @@ def timing_block(step_str: str): step_str: Description of the step being timed. """ - start = time() + start = perf_counter() yield - end = time() + end = perf_counter() LOG.debug(f" Execution time for {step_str}: {round(end - start, 2)} seconds") From dc04de14f1eef937129a1b25ecb6997ebc6bfa56 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 13:37:59 -0400 Subject: [PATCH 04/36] Start encapsulating BMI model into model.py NWMv3ForcingEngineModel starting with _values dict --- .../NextGen_Forcings_Engine/bmi_model.py | 4 +-- .../NextGen_Forcings_Engine/model.py | 33 +++++++++++-------- 2 files changed, 21 insertions(+), 16 deletions(-) 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 c1b2f692..eb096861 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -353,7 +353,7 @@ def initialize(self, config_file: str, output_path: str | None = None) -> None: self._values["time_step_size"] = self.cfg_bmi["time_step_seconds"] # Initialize the Forcings Engine model - self._model = NWMv3ForcingEngineModel() + self._model = NWMv3ForcingEngineModel(self) # Set catchment ids if using hydrofabric if self._grid_type == "hydrofabric": @@ -470,7 +470,6 @@ def update_until(self, future_time: float): == self.cfg_bmi["initial_time"] ): self._model.run( - self._values, future_time, self._job_meta, self.geo_meta, @@ -487,7 +486,6 @@ def update_until(self, future_time: float): self._values["current_model_time"] += self._values["time_step_size"] # Run the model for the new current time and update the state. self._model.run( - self._values, self._values["current_model_time"], self._job_meta, self.geo_meta, diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index c7121a91..4b48e137 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,7 +1,9 @@ +from __future__ import annotations import datetime import os from contextlib import contextmanager from time import time, perf_counter +from typing import TYPE_CHECKING import ewts import numpy as np @@ -31,6 +33,12 @@ NWMV3OConusProcessor, ) +if TYPE_CHECKING: + # To allow type hint without circular import error + from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.bmi_model import ( + NWMv3_Forcing_Engine_BMI_model_Base, + ) + LOG = ewts.get_logger(ewts.FORCING_ID) @@ -62,9 +70,10 @@ def wrapper(*args, **kwargs): class NWMv3ForcingEngineModel: """NextGen Forcings Engine BMI model class for NWMv3 forcings.""" - def __init__(self): + def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): """Initialize the NWMv3 Forcing Engine Model.""" self.source_data_processor = None + self._bmi = bmi_model # TODO: refactor the bmi_model.py file and this to have this type maintain its own state. # def __init__(self): @@ -77,7 +86,6 @@ def __init__(self): def run( self, - model: dict, future_time: float, config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, @@ -88,10 +96,10 @@ def run( ) -> None: """Execute the full forcings engine BMI pipeline for a given future timestep. - This method updates the `model` state dictionary with atmospheric forcings computed from + This method updates the `self._bmi._values` state dictionary with atmospheric forcings computed from available input datasets. It handles initialization, AWS Zarr loading, regridding, temporal interpolation, bias correction, downscaling, supplemental precipitation processing, and output - population into the model structure. + population into the self._bmi._values structure. The following steps are performed: @@ -110,10 +118,9 @@ def run( b. Disaggregate and interpolate. c. Layer into the final output. 5. Write output to NetCDF forcing files if requested. - 6. Update the model state dictionary with flattened arrays. + 6. Update the self._bmi._values state dictionary with flattened arrays. 7. Advance the BMI time index. - :param model: The model state dictionary that will be updated with new forcing data. :param future_time: The number of seconds into the future to advance the model. :param config_options: Configuration object containing all model options, flags, and paths. :param wrf_hydro_geo_meta: Geospatial metadata needed for regridding and interpolation. @@ -153,7 +160,6 @@ def run( output_obj, ) self.update_dict( - model, config_options, wrf_hydro_geo_meta, output_obj, @@ -696,7 +702,6 @@ def write_output( @time_function def update_dict( self, - model: dict, config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, output_obj: OutputObj, @@ -745,20 +750,22 @@ def update_dict( ] if config_options.grid_type == "gridded": for count, variable in enumerate(variables): - model[variable + "_ELEMENT"] = output_obj.output_local[ + self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local[ count, :, : ].flatten() elif config_options.grid_type == "unstructured": for count, variable in enumerate(variables): - model[variable + "_ELEMENT"] = output_obj.output_local_elem[ + self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local_elem[ + count, : + ].flatten() + self._bmi._values[variable + "_NODE"] = output_obj.output_local[ count, : ].flatten() - model[variable + "_NODE"] = output_obj.output_local[count, :].flatten() elif config_options.grid_type == "hydrofabric": for count, variable in enumerate(variables): - model[variable + "_ELEMENT"] = output_obj.output_global[ + self._bmi._values[variable + "_ELEMENT"] = output_obj.output_global[ count, : ].flatten() - model["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global + self._bmi._values["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global else: raise ValueError(f"Unexpected grid_type: {repr(config_options.grid_type)}") From ffb1f55eb160a63d78548de103381c2d3a7d6326 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 14:34:31 -0400 Subject: [PATCH 05/36] Encapsulate ConfigOptions instance --- .../NextGen_Forcings_Engine/bmi_model.py | 2 - .../NextGen_Forcings_Engine/model.py | 367 ++++++++++-------- tests/test_utils.py | 8 +- 3 files changed, 205 insertions(+), 172 deletions(-) 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 eb096861..6a033761 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -471,7 +471,6 @@ def update_until(self, future_time: float): ): self._model.run( future_time, - self._job_meta, self.geo_meta, self._input_forcing_mod, self._supp_pcp_mod, @@ -487,7 +486,6 @@ def update_until(self, future_time: float): # Run the model for the new current time and update the state. self._model.run( self._values["current_model_time"], - self._job_meta, self.geo_meta, self._input_forcing_mod, self._supp_pcp_mod, diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 4b48e137..53a81729 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -87,7 +87,6 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): def run( self, future_time: float, - config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, input_forcing_mod: dict, supp_pcp_mod: dict, @@ -101,6 +100,8 @@ def run( interpolation, bias correction, downscaling, supplemental precipitation processing, and output population into the self._bmi._values structure. + `self._bmi._job_meta`, an instance of ConfigOptions is also updated in-place, for example for time handling. + The following steps are performed: 1. Determine the current forecast and output times based on the future timestamp @@ -122,7 +123,6 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param config_options: Configuration object containing all model options, flags, and paths. :param wrf_hydro_geo_meta: Geospatial metadata needed for regridding and interpolation. :param input_forcing_mod: Dictionary of initialized input forcing modules indexed by forcing key. :param supp_pcp_mod: Dictionary of supplemental precipitation modules indexed by key. @@ -132,13 +132,12 @@ def run( :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ - self.determine_forecast(future_time, config_options) - self.adjust_precip(config_options, input_forcing_mod, mpi_config) - self.log_forecast(config_options, mpi_config) + self.determine_forecast(future_time) + self.adjust_precip(input_forcing_mod, mpi_config) + self.log_forecast(mpi_config) # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - config_options, wrf_hydro_geo_meta, input_forcing_mod, supp_pcp_mod, @@ -146,7 +145,6 @@ def run( output_obj, ) self.process_suplemental_precip( - config_options, wrf_hydro_geo_meta, supp_pcp_mod, mpi_config, @@ -154,25 +152,22 @@ def run( input_forcings, ) self.write_output( - config_options, wrf_hydro_geo_meta, mpi_config, output_obj, ) self.update_dict( - config_options, wrf_hydro_geo_meta, output_obj, ) ## Update BMI model time index to next iteration - config_options.bmi_time_index += 1 + self._bmi._job_meta.bmi_time_index += 1 @time_function def determine_forecast( self, future_time: float, - config_options: ConfigOptions, ) -> None: """Determine the forecast for the given future time and configuration. @@ -181,62 +176,67 @@ def determine_forecast( Modifies mutable arguments in-place. """ # Assign the future time to the configuration - config_options.bmi_time = future_time - self.disaggregate_fun = disaggregateMod.disaggregate_factory(config_options) + self._bmi._job_meta.bmi_time = future_time + self.disaggregate_fun = disaggregateMod.disaggregate_factory( + self._bmi._job_meta + ) # Calculate current time stamp based on operational configuration - if config_options.ana_flag: + if self._bmi._job_meta.ana_flag: # If we're in an AnA configuration, then must offset the BMI future # timestamp to account for the "lookback" period being properly iterated # over between 3-28 hour look back time period and operation configuration - if config_options.input_forcings[0] in [20, 22]: - config_options.current_fcst_cycle = ( - config_options.b_date_proc + if self._bmi._job_meta.input_forcings[0] in [20, 22]: + self._bmi._job_meta.current_fcst_cycle = ( + self._bmi._job_meta.b_date_proc + pd.TimedeltaIndex( np.array([future_time - 7200.0], dtype=float), "s" )[0] ) - config_options.current_time = ( - config_options.b_date_proc + self._bmi._job_meta.current_time = ( + self._bmi._job_meta.b_date_proc + pd.TimedeltaIndex( np.array([future_time - 7200.0], dtype=float), "s" )[0] ) - config_options.future_time = future_time + self._bmi._job_meta.future_time = future_time else: # Puerto Rico / Hawaii AnA: 1-hour lookback (based on 6-hourly forecast cycles) - config_options.current_fcst_cycle = ( - config_options.b_date_proc + self._bmi._job_meta.current_fcst_cycle = ( + self._bmi._job_meta.b_date_proc + pd.TimedeltaIndex( np.array([future_time - 3600.0], dtype=float), "s" )[0] ) - config_options.current_time = ( - config_options.b_date_proc + self._bmi._job_meta.current_time = ( + self._bmi._job_meta.b_date_proc + pd.TimedeltaIndex( np.array([future_time - 3600.0], dtype=float), "s" )[0] ) else: # Forecast-only mode — use BMI timestamp as-is - config_options.current_fcst_cycle = config_options.b_date_proc - config_options.current_time = pd.Timestamp( - config_options.b_date_proc + self._bmi._job_meta.current_fcst_cycle = self._bmi._job_meta.b_date_proc + self._bmi._job_meta.current_time = pd.Timestamp( + self._bmi._job_meta.b_date_proc ) + pd.to_timedelta(future_time, unit="s") LOG.debug( "NextGen Forcings Engine processing meteorological forcings for BMI timestamp" ) - LOG.debug(f"Model.py current time: {config_options.current_time}") - LOG.debug(f"Model.py current fcst cycle: {config_options.current_fcst_cycle}") + LOG.debug(f"Model.py current time: {self._bmi._job_meta.current_time}") + LOG.debug( + f"Model.py current fcst cycle: {self._bmi._job_meta.current_fcst_cycle}" + ) - if config_options.first_fcst_cycle is None: - config_options.first_fcst_cycle = config_options.current_fcst_cycle + if self._bmi._job_meta.first_fcst_cycle is None: + self._bmi._job_meta.first_fcst_cycle = ( + self._bmi._job_meta.current_fcst_cycle + ) @time_function def adjust_precip( self, - config_options: ConfigOptions, input_forcing_mod: dict, mpi_config: MpiConfig, ) -> None: @@ -246,17 +246,16 @@ def adjust_precip( -------- Modifies mutable arguments in-place. """ - if not config_options.precip_only_flag: + if not self._bmi._job_meta.precip_only_flag: # reset skips if present - for force_key in config_options.input_forcings: + for force_key in self._bmi._job_meta.input_forcings: input_forcing_mod[force_key].skip = False - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) @time_function def log_forecast( self, - config_options: ConfigOptions, mpi_config: MpiConfig, ) -> None: """Log information about the current forecast cycle. @@ -267,26 +266,25 @@ def log_forecast( """ # Log information about this forecast cycle if mpi_config.rank == 0: - config_options.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" - err_handler.log_msg(config_options, mpi_config, True) - config_options.statusMsg = ( + self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" + err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + self._bmi._job_meta.statusMsg = ( "Processing Forecast Cycle: " - + config_options.current_fcst_cycle.strftime("%Y-%m-%d %H:%M") + + self._bmi._job_meta.current_fcst_cycle.strftime("%Y-%m-%d %H:%M") ) - err_handler.log_msg(config_options, mpi_config, True) - config_options.statusMsg = ( + err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + self._bmi._job_meta.statusMsg = ( "Forecast Cycle Length is: " - + str(config_options.cycle_length_minutes) + + str(self._bmi._job_meta.cycle_length_minutes) + " minutes" ) - err_handler.log_msg(config_options, mpi_config, True) + err_handler.log_msg(self._bmi._job_meta, mpi_config, True) # mpi_config.comm.barrier() @time_function def loop_through_forcing_products( self, future_time: float, - config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, input_forcing_mod: dict, supp_pcp_mod: dict, @@ -305,93 +303,99 @@ def loop_through_forcing_products( # 3.) Regrid the forcings, and temporally interpolate. # 4.) Downscale. # 5.) Layer, and output as necessary. - ana_factor = 1 if config_options.ana_flag is False else 0 + ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 show_message = True - if not config_options.precip_only_flag: - if config_options.grid_type == "gridded": + if not self._bmi._job_meta.precip_only_flag: + if self._bmi._job_meta.grid_type == "gridded": # Reset out final grids to missing values. - output_obj.output_local[:, :, :] = config_options.globalNdv - elif config_options.grid_type == "unstructured": + output_obj.output_local[:, :, :] = self._bmi._job_meta.globalNdv + elif self._bmi._job_meta.grid_type == "unstructured": # Reset out final grids to missing values. - output_obj.output_local[:, :] = config_options.globalNdv - output_obj.output_local_elem[:, :] = config_options.globalNdv - elif config_options.grid_type == "hydrofabric": + output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv + output_obj.output_local_elem[:, :] = self._bmi._job_meta.globalNdv + elif self._bmi._job_meta.grid_type == "hydrofabric": # Reset out final grids to missing values. - output_obj.output_local[:, :] = config_options.globalNdv + output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv else: raise ValueError( - f"Unexpected grid_type: {repr(config_options.grid_type)}" + f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" ) # Increment or initialize output step count - if config_options.current_output_step is None: - config_options.current_output_step = 1 + if self._bmi._job_meta.current_output_step is None: + self._bmi._job_meta.current_output_step = 1 else: - config_options.current_output_step += 1 + self._bmi._job_meta.current_output_step += 1 # Optional sub-output timestamp - if config_options.sub_output_hour is not None: + if self._bmi._job_meta.sub_output_hour is not None: # TODO This is not used - subOutDate = config_options.first_fcst_cycle + datetime.timedelta( - hours=config_options.sub_output_hour + subOutDate = self._bmi._job_meta.first_fcst_cycle + datetime.timedelta( + hours=self._bmi._job_meta.sub_output_hour ) # Compute the output timestamp for this step - if config_options.ana_flag: + if self._bmi._job_meta.ana_flag: output_obj.outDate = ( - config_options.current_fcst_cycle - + datetime.timedelta(seconds=config_options.output_freq * 60) + self._bmi._job_meta.current_fcst_cycle + + datetime.timedelta(seconds=self._bmi._job_meta.output_freq * 60) ) else: output_obj.outDate = ( - config_options.current_fcst_cycle + self._bmi._job_meta.current_fcst_cycle + datetime.timedelta(seconds=future_time) ) - config_options.current_output_date = output_obj.outDate + self._bmi._job_meta.current_output_date = output_obj.outDate # Adjust file_date for AnA if needed file_date = ( output_obj.outDate - - datetime.timedelta(seconds=config_options.output_freq * 60) - if config_options.ana_flag + - datetime.timedelta(seconds=self._bmi._job_meta.output_freq * 60) + if self._bmi._job_meta.ana_flag else output_obj.outDate ) # Compute previous output date (used for downscaling logic) - if config_options.current_output_step == ana_factor: - config_options.prev_output_date = config_options.current_output_date + if self._bmi._job_meta.current_output_step == ana_factor: + self._bmi._job_meta.prev_output_date = ( + self._bmi._job_meta.current_output_date + ) else: - config_options.prev_output_date = ( - config_options.current_output_date + self._bmi._job_meta.prev_output_date = ( + self._bmi._job_meta.current_output_date - datetime.timedelta(seconds=future_time) ) # Print message on log file indicating the timestamp # we are currently processing for forcings if mpi_config.rank == 0 and show_message: - config_options.statusMsg = "=========================================" - err_handler.log_msg(config_options, mpi_config, True) - config_options.statusMsg = f"Processing for output timestep: {file_date.strftime('%Y-%m-%d %H:%M')}" - err_handler.log_msg(config_options, mpi_config, True) - - config_options.currentForceNum = 0 - config_options.currentCustomForceNum = 0 - LOG.debug(f"config_options.input_forcings: {config_options.input_forcings}") + self._bmi._job_meta.statusMsg = ( + "=========================================" + ) + err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + self._bmi._job_meta.statusMsg = f"Processing for output timestep: {file_date.strftime('%Y-%m-%d %H:%M')}" + err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + + self._bmi._job_meta.currentForceNum = 0 + self._bmi._job_meta.currentCustomForceNum = 0 + LOG.debug( + f"config_options.input_forcings: {self._bmi._job_meta.input_forcings}" + ) # Loop over each of the input forcings specified. LOG.debug( - f"Model.py forcing loop: {len(config_options.input_forcings)} forcings configured: {config_options.input_forcings}" + f"Model.py forcing loop: {len(self._bmi._job_meta.input_forcings)} forcings configured: {self._bmi._job_meta.input_forcings}" ) - for force_key in config_options.input_forcings: + for force_key in self._bmi._job_meta.input_forcings: LOG.debug(f"force_key: {force_key}") - LOG.debug(f"config_options.aws: {config_options.aws}") + LOG.debug(f"config_options.aws: {self._bmi._job_meta.aws}") # Pass these methods for AORC data is ERA5-Interim blend is requested # so we can finish filling in the missing gaps if ( force_key == 23 - and 12 in config_options.input_forcings - and 21 in config_options.input_forcings + and 12 in self._bmi._job_meta.input_forcings + and 21 in self._bmi._job_meta.input_forcings ): input_forcings = input_forcing_mod[force_key] @@ -401,55 +405,63 @@ def loop_through_forcing_products( else: input_forcings = input_forcing_mod[force_key] input_forcings.calc_neighbor_files( - config_options, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, mpi_config ) if force_key in [12, 21, 27]: - if config_options.aws is None: + if self._bmi._job_meta.aws is None: # Calculate the previous and next input cycle files from the inputs. input_forcings.calc_neighbor_files( - config_options, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) else: # Flag to indicate the AWS .zarr AORC method if force_key == 12: if self.source_data_processor is None: self.source_data_processor = AORCConusProcessor( - config_options, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, mpi_config, wrf_hydro_geo_meta ) elif force_key == 21: if self.source_data_processor is None: self.source_data_processor = AORCAlaskaProcessor( - config_options, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, mpi_config, wrf_hydro_geo_meta ) # Flag to indicate the AWS .zarr NWMv3 Forcing file method elif force_key == 27: if self.source_data_processor is None: - if config_options.nwm_domain == "CONUS": + if self._bmi._job_meta.nwm_domain == "CONUS": self.source_data_processor = NWMV3ConusProcessor( - config_options, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, + mpi_config, + wrf_hydro_geo_meta, ) - elif config_options.nwm_domain in [ + elif self._bmi._job_meta.nwm_domain in [ "Hawaii", "PR", ]: self.source_data_processor = NWMV3OConusProcessor( - config_options, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, + mpi_config, + wrf_hydro_geo_meta, ) - elif config_options.nwm_domain == "Alaska": + elif self._bmi._job_meta.nwm_domain == "Alaska": self.source_data_processor = NWMV3AlaskaProcessor( - config_options, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, + mpi_config, + wrf_hydro_geo_meta, ) else: raise ValueError( - f"Unsupported domain type ({config_options.nwm_domain} for forcing type: {force_key} )" + f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" ) - config_options.aws_obj = ( + self._bmi._job_meta.aws_obj = ( self.source_data_processor.process_historical_data( - config_options.current_time + self._bmi._job_meta.current_time ) ) @@ -459,15 +471,15 @@ def loop_through_forcing_products( break # Regrid forcings. input_forcings.regrid_inputs( - config_options, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_forcing_bounds( - config_options, input_forcings, mpi_config + self._bmi._job_meta, input_forcings, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the # next set of forcings as the previous step just regridded the previous forcing. @@ -477,83 +489,89 @@ def loop_through_forcing_products( and input_forcings.regridded_forcings2 is not None ): # Set the forcings back to reflect we just regridded the previous set of inputs, not the next. - if config_options.grid_type == "gridded": + if self._bmi._job_meta.grid_type == "gridded": input_forcings.regridded_forcings1[:, :, :] = ( input_forcings.regridded_forcings2[:, :, :] ) - elif config_options.grid_type == "unstructured": + elif self._bmi._job_meta.grid_type == "unstructured": input_forcings.regridded_forcings1[:, :] = ( input_forcings.regridded_forcings2[:, :] ) input_forcings.regridded_forcings1_elem[:, :] = ( input_forcings.regridded_forcings2_elem[:, :] ) - elif config_options.grid_type == "hydrofabric": + elif self._bmi._job_meta.grid_type == "hydrofabric": input_forcings.regridded_forcings1[:, :] = ( input_forcings.regridded_forcings2[:, :] ) else: raise ValueError( - f"Unexpected grid_type: {repr(config_options.grid_type)}" + f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" ) # Re-calculate the neighbor files. input_forcings.calc_neighbor_files( - config_options, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Regrid the forcings for the end of the window. input_forcings.regrid_inputs( - config_options, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) input_forcings.rstFlag = 0 # Run temporal interpolation on the grids. - input_forcings.temporal_interpolate_inputs(config_options, mpi_config) - err_handler.check_program_status(config_options, mpi_config) + input_forcings.temporal_interpolate_inputs( + self._bmi._job_meta, mpi_config + ) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run bias correction. bias_correction.run_bias_correction( - input_forcings, config_options, wrf_hydro_geo_meta, mpi_config + input_forcings, self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run downscaling on grids for this output timestep. downscale.run_downscaling( - input_forcings, config_options, wrf_hydro_geo_meta, mpi_config + input_forcings, self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Layer in forcings from this product. layeringMod.layer_final_forcings( - output_obj, input_forcings, config_options, mpi_config + output_obj, input_forcings, self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, mpi_config) - config_options.currentForceNum += 1 + self._bmi._job_meta.currentForceNum += 1 if force_key == 10: - config_options.currentCustomForceNum += 1 + self._bmi._job_meta.currentCustomForceNum += 1 LOG.debug(f"End of loop for force_key {force_key}") # Process supplemental precipitation if we specified in the configuration file. - if config_options.number_supp_pcp > 0: - for supp_pcp_key in config_options.supp_precip_forcings: + if self._bmi._job_meta.number_supp_pcp > 0: + for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key != 13: # Like with input forcings, calculate the neighboring files to use. supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - config_options, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) # Regrid the supplemental precipitation. supp_pcp_mod[supp_pcp_key].regrid_inputs( - config_options, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) if ( supp_pcp_mod[supp_pcp_key].regridded_precip1 is not None @@ -561,45 +579,53 @@ def loop_through_forcing_products( ): # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_supp_pcp_bounds( - config_options, + self._bmi._job_meta, supp_pcp_mod[supp_pcp_key], mpi_config, wrf_hydro_geo_meta, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen self.disaggregate_fun( input_forcings, supp_pcp_mod[supp_pcp_key], - config_options, + self._bmi._job_meta, mpi_config, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) # Run temporal interpolation on the grids. supp_pcp_mod[supp_pcp_key].temporal_interpolate_inputs( - config_options, mpi_config + self._bmi._job_meta, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( output_obj, supp_pcp_mod[supp_pcp_key], - config_options, + self._bmi._job_meta, mpi_config, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) # Call the output routines # adjust date for AnA if necessary - if config_options.ana_flag: + if self._bmi._job_meta.ana_flag: output_obj.outDate = file_date ################ Commenting this out to bypass NWM forcing file output functionality ######### - # output_obj.output_final_ldasin(config_options, wrf_hydro_geo_meta, mpi_config) - # err_handler.check_program_status(config_options, mpi_config) + # output_obj.output_final_ldasin(self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config) + # err_handler.check_program_status(self._bmi._job_meta, mpi_config) ############################################################################################## return input_forcings @@ -607,7 +633,6 @@ def loop_through_forcing_products( @time_function def process_suplemental_precip( self, - config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, supp_pcp_mod: dict, mpi_config: MpiConfig, @@ -620,22 +645,26 @@ def process_suplemental_precip( -------- Modifies mutable arguments in-place. """ - if config_options.customSuppPcpFreq is not None: + if self._bmi._job_meta.customSuppPcpFreq is not None: # Process supplemental precipitation if we specified in the configuration file. - if config_options.number_supp_pcp > 0: - for supp_pcp_key in config_options.supp_precip_forcings: + if self._bmi._job_meta.number_supp_pcp > 0: + for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key == 14: # Like with input forcings, calculate the neighboring files to use. supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - config_options, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) # Regrid the supplemental precipitation. supp_pcp_mod[supp_pcp_key].regrid_inputs( - config_options, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) if ( supp_pcp_mod[supp_pcp_key].regridded_precip1 is not None @@ -643,40 +672,47 @@ def process_suplemental_precip( ): # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_supp_pcp_bounds( - config_options, + self._bmi._job_meta, supp_pcp_mod[supp_pcp_key], mpi_config, wrf_hydro_geo_meta, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) self.disaggregate_fun( input_forcings, supp_pcp_mod[supp_pcp_key], - config_options, + self._bmi._job_meta, mpi_config, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) # Run temporal interpolation on the grids. supp_pcp_mod[supp_pcp_key].temporal_interpolate_inputs( - config_options, mpi_config + self._bmi._job_meta, mpi_config + ) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config ) - err_handler.check_program_status(config_options, mpi_config) # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( output_obj, supp_pcp_mod[supp_pcp_key], - config_options, + self._bmi._job_meta, mpi_config, ) - err_handler.check_program_status(config_options, mpi_config) + err_handler.check_program_status( + self._bmi._job_meta, mpi_config + ) @time_function def write_output( self, - config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, mpi_config: MpiConfig, output_obj: OutputObj, @@ -690,11 +726,11 @@ def write_output( # If user requests output for given domain, then call # the I/O module to update opened netcdf file with forcing fields if ( - config_options.forcing_output == 1 - or config_options.grid_type == "hydrofabric" + self._bmi._job_meta.forcing_output == 1 + or self._bmi._job_meta.grid_type == "hydrofabric" ): output_obj.gather_global_outputs( - config_options, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config ) """##################Step 6: flatten and update dict##########################################################################""" @@ -702,7 +738,6 @@ def write_output( @time_function def update_dict( self, - config_options: ConfigOptions, wrf_hydro_geo_meta: GeoMeta, output_obj: OutputObj, ) -> None: @@ -725,7 +760,7 @@ def update_dict( # 7.) Surface incoming shortwave radiation flux (W/m^2) # 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations - if config_options.include_lqfrac == 1: + if self._bmi._job_meta.include_lqfrac == 1: variables = [ "U2D", "V2D", @@ -748,12 +783,12 @@ def update_dict( "PSFC", "SWDOWN", ] - if config_options.grid_type == "gridded": + if self._bmi._job_meta.grid_type == "gridded": for count, variable in enumerate(variables): self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local[ count, :, : ].flatten() - elif config_options.grid_type == "unstructured": + elif self._bmi._job_meta.grid_type == "unstructured": for count, variable in enumerate(variables): self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local_elem[ count, : @@ -761,11 +796,13 @@ def update_dict( self._bmi._values[variable + "_NODE"] = output_obj.output_local[ count, : ].flatten() - elif config_options.grid_type == "hydrofabric": + elif self._bmi._job_meta.grid_type == "hydrofabric": for count, variable in enumerate(variables): self._bmi._values[variable + "_ELEMENT"] = output_obj.output_global[ count, : ].flatten() self._bmi._values["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global else: - raise ValueError(f"Unexpected grid_type: {repr(config_options.grid_type)}") + raise ValueError( + f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" + ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 11496c02..b15c6a36 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -527,7 +527,6 @@ def pre_regrid(self) -> None: f"In pre_regrid, expected state to be either None or 'post_ran' but got {repr(self._state)}. The test is set up incorrectly." ) - config_options = self.config_options mpi_config = self.mpi_config geo_meta = self.geo_meta supp_pcp_mod = self.bmi_model._supp_pcp_mod @@ -541,14 +540,13 @@ def pre_regrid(self) -> None: model = self.bmi_model._model ### NOTE this should mimic NWMv3ForcingEngineModel.run() with the exception of setting the skip flag - model.determine_forecast(future_time, config_options) - model.adjust_precip(config_options, input_forcing_mod, mpi_config) - model.log_forecast(config_options, mpi_config) + model.determine_forecast(future_time) + model.adjust_precip(input_forcing_mod, mpi_config) + model.log_forecast(mpi_config) ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() model.loop_through_forcing_products( future_time, - config_options, geo_meta, input_forcing_mod, supp_pcp_mod, From b04da44fb09f17e752473f526d5564178f43b367 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 14:53:12 -0400 Subject: [PATCH 06/36] Encapsulate GeoMeta instance --- .../NextGen_Forcings_Engine/bmi_model.py | 2 - .../NextGen_Forcings_Engine/model.py | 52 +++++++++---------- tests/test_utils.py | 2 - 3 files changed, 26 insertions(+), 30 deletions(-) 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 6a033761..8c74f258 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -471,7 +471,6 @@ def update_until(self, future_time: float): ): self._model.run( future_time, - self.geo_meta, self._input_forcing_mod, self._supp_pcp_mod, self._mpi_meta, @@ -486,7 +485,6 @@ def update_until(self, future_time: float): # Run the model for the new current time and update the state. self._model.run( self._values["current_model_time"], - self.geo_meta, self._input_forcing_mod, self._supp_pcp_mod, self._mpi_meta, diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 53a81729..1ac1a57f 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -87,7 +87,6 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): def run( self, future_time: float, - wrf_hydro_geo_meta: GeoMeta, input_forcing_mod: dict, supp_pcp_mod: dict, mpi_config: MpiConfig, @@ -123,7 +122,6 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param wrf_hydro_geo_meta: Geospatial metadata needed for regridding and interpolation. :param input_forcing_mod: Dictionary of initialized input forcing modules indexed by forcing key. :param supp_pcp_mod: Dictionary of supplemental precipitation modules indexed by key. :param mpi_config: Object containing MPI communication settings such as rank and communicator. @@ -138,26 +136,22 @@ def run( # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - wrf_hydro_geo_meta, input_forcing_mod, supp_pcp_mod, mpi_config, output_obj, ) self.process_suplemental_precip( - wrf_hydro_geo_meta, supp_pcp_mod, mpi_config, output_obj, input_forcings, ) self.write_output( - wrf_hydro_geo_meta, mpi_config, output_obj, ) self.update_dict( - wrf_hydro_geo_meta, output_obj, ) @@ -285,7 +279,6 @@ def log_forecast( def loop_through_forcing_products( self, future_time: float, - wrf_hydro_geo_meta: GeoMeta, input_forcing_mod: dict, supp_pcp_mod: dict, mpi_config: MpiConfig, @@ -422,12 +415,16 @@ def loop_through_forcing_products( if force_key == 12: if self.source_data_processor is None: self.source_data_processor = AORCConusProcessor( - self._bmi._job_meta, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, + mpi_config, + self._bmi.geo_meta, ) elif force_key == 21: if self.source_data_processor is None: self.source_data_processor = AORCAlaskaProcessor( - self._bmi._job_meta, mpi_config, wrf_hydro_geo_meta + self._bmi._job_meta, + mpi_config, + self._bmi.geo_meta, ) # Flag to indicate the AWS .zarr NWMv3 Forcing file method @@ -437,7 +434,7 @@ def loop_through_forcing_products( self.source_data_processor = NWMV3ConusProcessor( self._bmi._job_meta, mpi_config, - wrf_hydro_geo_meta, + self._bmi.geo_meta, ) elif self._bmi._job_meta.nwm_domain in [ "Hawaii", @@ -446,13 +443,13 @@ def loop_through_forcing_products( self.source_data_processor = NWMV3OConusProcessor( self._bmi._job_meta, mpi_config, - wrf_hydro_geo_meta, + self._bmi.geo_meta, ) elif self._bmi._job_meta.nwm_domain == "Alaska": self.source_data_processor = NWMV3AlaskaProcessor( self._bmi._job_meta, mpi_config, - wrf_hydro_geo_meta, + self._bmi.geo_meta, ) else: raise ValueError( @@ -471,7 +468,7 @@ def loop_through_forcing_products( break # Regrid forcings. input_forcings.regrid_inputs( - self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status(self._bmi._job_meta, mpi_config) @@ -516,7 +513,7 @@ def loop_through_forcing_products( # Regrid the forcings for the end of the window. input_forcings.regrid_inputs( - self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status(self._bmi._job_meta, mpi_config) @@ -530,13 +527,19 @@ def loop_through_forcing_products( # Run bias correction. bias_correction.run_bias_correction( - input_forcings, self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + input_forcings, + self._bmi._job_meta, + self._bmi.geo_meta, + mpi_config, ) err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run downscaling on grids for this output timestep. downscale.run_downscaling( - input_forcings, self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + input_forcings, + self._bmi._job_meta, + self._bmi.geo_meta, + mpi_config, ) err_handler.check_program_status(self._bmi._job_meta, mpi_config) @@ -567,7 +570,7 @@ def loop_through_forcing_products( # Regrid the supplemental precipitation. supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status( self._bmi._job_meta, mpi_config @@ -582,7 +585,7 @@ def loop_through_forcing_products( self._bmi._job_meta, supp_pcp_mod[supp_pcp_key], mpi_config, - wrf_hydro_geo_meta, + self._bmi.geo_meta, ) err_handler.check_program_status( self._bmi._job_meta, mpi_config @@ -624,7 +627,7 @@ def loop_through_forcing_products( output_obj.outDate = file_date ################ Commenting this out to bypass NWM forcing file output functionality ######### - # output_obj.output_final_ldasin(self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config) + # output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, mpi_config) # err_handler.check_program_status(self._bmi._job_meta, mpi_config) ############################################################################################## @@ -633,7 +636,6 @@ def loop_through_forcing_products( @time_function def process_suplemental_precip( self, - wrf_hydro_geo_meta: GeoMeta, supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, @@ -660,7 +662,7 @@ def process_suplemental_precip( # Regrid the supplemental precipitation. supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status( self._bmi._job_meta, mpi_config @@ -675,7 +677,7 @@ def process_suplemental_precip( self._bmi._job_meta, supp_pcp_mod[supp_pcp_key], mpi_config, - wrf_hydro_geo_meta, + self._bmi.geo_meta, ) err_handler.check_program_status( self._bmi._job_meta, mpi_config @@ -713,7 +715,6 @@ def process_suplemental_precip( @time_function def write_output( self, - wrf_hydro_geo_meta: GeoMeta, mpi_config: MpiConfig, output_obj: OutputObj, ) -> None: @@ -730,7 +731,7 @@ def write_output( or self._bmi._job_meta.grid_type == "hydrofabric" ): output_obj.gather_global_outputs( - self._bmi._job_meta, wrf_hydro_geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) """##################Step 6: flatten and update dict##########################################################################""" @@ -738,7 +739,6 @@ def write_output( @time_function def update_dict( self, - wrf_hydro_geo_meta: GeoMeta, output_obj: OutputObj, ) -> None: """Flatten the Forcings Engine output object and update the BMI dictionary. @@ -801,7 +801,7 @@ def update_dict( self._bmi._values[variable + "_ELEMENT"] = output_obj.output_global[ count, : ].flatten() - self._bmi._values["CAT-ID"] = wrf_hydro_geo_meta.element_ids_global + self._bmi._values["CAT-ID"] = self._bmi.geo_meta.element_ids_global else: raise ValueError( f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" diff --git a/tests/test_utils.py b/tests/test_utils.py index b15c6a36..65f55b5a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -528,7 +528,6 @@ def pre_regrid(self) -> None: ) mpi_config = self.mpi_config - geo_meta = self.geo_meta supp_pcp_mod = self.bmi_model._supp_pcp_mod output_obj = self.bmi_model._output_obj input_forcing_mod = self.bmi_model._input_forcing_mod @@ -547,7 +546,6 @@ def pre_regrid(self) -> None: self.set_input_forcings_skip_flags() model.loop_through_forcing_products( future_time, - geo_meta, input_forcing_mod, supp_pcp_mod, mpi_config, From d2ca18db00a5e607d5886b3e87a52a7653dcc5de Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 15:36:11 -0400 Subject: [PATCH 07/36] Encapsulate input forcing mod dict --- .../NextGen_Forcings_Engine/bmi_model.py | 2 -- .../NextGen_Forcings_Engine/model.py | 13 ++++--------- tests/test_utils.py | 6 ++---- 3 files changed, 6 insertions(+), 15 deletions(-) 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 8c74f258..50d46ad1 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -471,7 +471,6 @@ def update_until(self, future_time: float): ): self._model.run( future_time, - self._input_forcing_mod, self._supp_pcp_mod, self._mpi_meta, self._output_obj, @@ -485,7 +484,6 @@ def update_until(self, future_time: float): # Run the model for the new current time and update the state. self._model.run( self._values["current_model_time"], - self._input_forcing_mod, self._supp_pcp_mod, self._mpi_meta, self._output_obj, diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 1ac1a57f..69c3259c 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -87,7 +87,6 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): def run( self, future_time: float, - input_forcing_mod: dict, supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, @@ -122,7 +121,6 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param input_forcing_mod: Dictionary of initialized input forcing modules indexed by forcing key. :param supp_pcp_mod: Dictionary of supplemental precipitation modules indexed by key. :param mpi_config: Object containing MPI communication settings such as rank and communicator. :param output_obj: Output object that stores the generated atmospheric forcing arrays. @@ -131,12 +129,11 @@ def run( """ self.determine_forecast(future_time) - self.adjust_precip(input_forcing_mod, mpi_config) + self.adjust_precip(mpi_config) self.log_forecast(mpi_config) # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - input_forcing_mod, supp_pcp_mod, mpi_config, output_obj, @@ -231,7 +228,6 @@ def determine_forecast( @time_function def adjust_precip( self, - input_forcing_mod: dict, mpi_config: MpiConfig, ) -> None: """Adjust precipitation for the given forecast cycle. @@ -243,7 +239,7 @@ def adjust_precip( if not self._bmi._job_meta.precip_only_flag: # reset skips if present for force_key in self._bmi._job_meta.input_forcings: - input_forcing_mod[force_key].skip = False + self._bmi._input_forcing_mod[force_key].skip = False err_handler.check_program_status(self._bmi._job_meta, mpi_config) @@ -279,7 +275,6 @@ def log_forecast( def loop_through_forcing_products( self, future_time: float, - input_forcing_mod: dict, supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, @@ -390,13 +385,13 @@ def loop_through_forcing_products( and 12 in self._bmi._job_meta.input_forcings and 21 in self._bmi._job_meta.input_forcings ): - input_forcings = input_forcing_mod[force_key] + input_forcings = self._bmi._input_forcing_mod[force_key] # These are not used # AORC_mask = input_forcings.regridded_mask_AORC # AORC_elem_mask = input_forcings.regridded_mask_elem_AORC else: - input_forcings = input_forcing_mod[force_key] + input_forcings = self._bmi._input_forcing_mod[force_key] input_forcings.calc_neighbor_files( self._bmi._job_meta, output_obj.outDate, mpi_config ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 65f55b5a..b5ed53a8 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -530,7 +530,6 @@ def pre_regrid(self) -> None: mpi_config = self.mpi_config supp_pcp_mod = self.bmi_model._supp_pcp_mod output_obj = self.bmi_model._output_obj - input_forcing_mod = self.bmi_model._input_forcing_mod future_time = ( self.bmi_model._values["current_model_time"] @@ -540,13 +539,12 @@ def pre_regrid(self) -> None: ### NOTE this should mimic NWMv3ForcingEngineModel.run() with the exception of setting the skip flag model.determine_forecast(future_time) - model.adjust_precip(input_forcing_mod, mpi_config) + model.adjust_precip(mpi_config) model.log_forecast(mpi_config) ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() model.loop_through_forcing_products( future_time, - input_forcing_mod, supp_pcp_mod, mpi_config, output_obj, @@ -558,7 +556,7 @@ def pre_regrid(self) -> None: def set_input_forcings_skip_flags(self) -> None: """Set the `skip` flag on the InputForcings object so that forcing regrid will not occur during loop_through_forcing_products().""" logging.debug( - "Setting input_forcing.skip = True for each value in dict self.input_forcing_mod" + "Setting input_forcing.skip = True for each value in dict self.bmi_model._input_forcing_mod" ) for force_key, input_forcing in self.bmi_model._input_forcing_mod.items(): input_forcing.skip = True From c120dd17cd633122c91292e24c7576108a6d76d6 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 15:43:48 -0400 Subject: [PATCH 08/36] Encapsulate supp pcp mod --- .../NextGen_Forcings_Engine/bmi_model.py | 2 - .../NextGen_Forcings_Engine/model.py | 46 ++++++++++--------- tests/test_utils.py | 2 - 3 files changed, 24 insertions(+), 26 deletions(-) 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 50d46ad1..4a8585dd 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -471,7 +471,6 @@ def update_until(self, future_time: float): ): self._model.run( future_time, - self._supp_pcp_mod, self._mpi_meta, self._output_obj, ) @@ -484,7 +483,6 @@ def update_until(self, future_time: float): # Run the model for the new current time and update the state. self._model.run( self._values["current_model_time"], - self._supp_pcp_mod, self._mpi_meta, self._output_obj, ) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 69c3259c..7043c024 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -87,7 +87,6 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): def run( self, future_time: float, - supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, ) -> None: @@ -121,7 +120,6 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param supp_pcp_mod: Dictionary of supplemental precipitation modules indexed by key. :param mpi_config: Object containing MPI communication settings such as rank and communicator. :param output_obj: Output object that stores the generated atmospheric forcing arrays. @@ -134,12 +132,10 @@ def run( # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - supp_pcp_mod, mpi_config, output_obj, ) self.process_suplemental_precip( - supp_pcp_mod, mpi_config, output_obj, input_forcings, @@ -275,7 +271,6 @@ def log_forecast( def loop_through_forcing_products( self, future_time: float, - supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, ) -> forcingInputMod.InputForcingsHydrofabric: @@ -556,7 +551,7 @@ def loop_through_forcing_products( for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key != 13: # Like with input forcings, calculate the neighboring files to use. - supp_pcp_mod[supp_pcp_key].calc_neighbor_files( + self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( self._bmi._job_meta, output_obj.outDate, mpi_config ) err_handler.check_program_status( @@ -564,7 +559,7 @@ def loop_through_forcing_products( ) # Regrid the supplemental precipitation. - supp_pcp_mod[supp_pcp_key].regrid_inputs( + self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status( @@ -572,13 +567,15 @@ def loop_through_forcing_products( ) if ( - supp_pcp_mod[supp_pcp_key].regridded_precip1 is not None - and supp_pcp_mod[supp_pcp_key].regridded_precip2 is not None + self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 + is not None + and self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip2 + is not None ): # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_supp_pcp_bounds( self._bmi._job_meta, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], mpi_config, self._bmi.geo_meta, ) @@ -589,7 +586,7 @@ def loop_through_forcing_products( # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen self.disaggregate_fun( input_forcings, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, mpi_config, ) @@ -598,7 +595,9 @@ def loop_through_forcing_products( ) # Run temporal interpolation on the grids. - supp_pcp_mod[supp_pcp_key].temporal_interpolate_inputs( + self._bmi._supp_pcp_mod[ + supp_pcp_key + ].temporal_interpolate_inputs( self._bmi._job_meta, mpi_config ) err_handler.check_program_status( @@ -608,7 +607,7 @@ def loop_through_forcing_products( # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( output_obj, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, mpi_config, ) @@ -631,7 +630,6 @@ def loop_through_forcing_products( @time_function def process_suplemental_precip( self, - supp_pcp_mod: dict, mpi_config: MpiConfig, output_obj: OutputObj, input_forcings: dict, @@ -648,7 +646,7 @@ def process_suplemental_precip( for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key == 14: # Like with input forcings, calculate the neighboring files to use. - supp_pcp_mod[supp_pcp_key].calc_neighbor_files( + self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( self._bmi._job_meta, output_obj.outDate, mpi_config ) err_handler.check_program_status( @@ -656,7 +654,7 @@ def process_suplemental_precip( ) # Regrid the supplemental precipitation. - supp_pcp_mod[supp_pcp_key].regrid_inputs( + self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, mpi_config ) err_handler.check_program_status( @@ -664,13 +662,15 @@ def process_suplemental_precip( ) if ( - supp_pcp_mod[supp_pcp_key].regridded_precip1 is not None - and supp_pcp_mod[supp_pcp_key].regridded_precip2 is not None + self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 + is not None + and self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip2 + is not None ): # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_supp_pcp_bounds( self._bmi._job_meta, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], mpi_config, self._bmi.geo_meta, ) @@ -680,7 +680,7 @@ def process_suplemental_precip( self.disaggregate_fun( input_forcings, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, mpi_config, ) @@ -689,7 +689,9 @@ def process_suplemental_precip( ) # Run temporal interpolation on the grids. - supp_pcp_mod[supp_pcp_key].temporal_interpolate_inputs( + self._bmi._supp_pcp_mod[ + supp_pcp_key + ].temporal_interpolate_inputs( self._bmi._job_meta, mpi_config ) err_handler.check_program_status( @@ -699,7 +701,7 @@ def process_suplemental_precip( # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( output_obj, - supp_pcp_mod[supp_pcp_key], + self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, mpi_config, ) diff --git a/tests/test_utils.py b/tests/test_utils.py index b5ed53a8..13879226 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -528,7 +528,6 @@ def pre_regrid(self) -> None: ) mpi_config = self.mpi_config - supp_pcp_mod = self.bmi_model._supp_pcp_mod output_obj = self.bmi_model._output_obj future_time = ( @@ -545,7 +544,6 @@ def pre_regrid(self) -> None: self.set_input_forcings_skip_flags() model.loop_through_forcing_products( future_time, - supp_pcp_mod, mpi_config, output_obj, ) From 783672b8cabfbf3beba2f85a0ea17cabdb74429f Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 15:50:57 -0400 Subject: [PATCH 09/36] Encapsulate MpiConfig --- .../NextGen_Forcings_Engine/bmi_model.py | 2 - .../NextGen_Forcings_Engine/model.py | 162 +++++++++--------- tests/test_utils.py | 11 +- 3 files changed, 84 insertions(+), 91 deletions(-) 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 4a8585dd..0a010a24 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -471,7 +471,6 @@ def update_until(self, future_time: float): ): self._model.run( future_time, - self._mpi_meta, self._output_obj, ) else: @@ -483,7 +482,6 @@ def update_until(self, future_time: float): # Run the model for the new current time and update the state. self._model.run( self._values["current_model_time"], - self._mpi_meta, self._output_obj, ) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 7043c024..859cde61 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -87,7 +87,6 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): def run( self, future_time: float, - mpi_config: MpiConfig, output_obj: OutputObj, ) -> None: """Execute the full forcings engine BMI pipeline for a given future timestep. @@ -120,28 +119,24 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param mpi_config: Object containing MPI communication settings such as rank and communicator. :param output_obj: Output object that stores the generated atmospheric forcing arrays. :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ self.determine_forecast(future_time) - self.adjust_precip(mpi_config) - self.log_forecast(mpi_config) + self.adjust_precip() + self.log_forecast() # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - mpi_config, output_obj, ) self.process_suplemental_precip( - mpi_config, output_obj, input_forcings, ) self.write_output( - mpi_config, output_obj, ) self.update_dict( @@ -222,10 +217,7 @@ def determine_forecast( ) @time_function - def adjust_precip( - self, - mpi_config: MpiConfig, - ) -> None: + def adjust_precip(self) -> None: """Adjust precipitation for the given forecast cycle. Warnings @@ -237,13 +229,10 @@ def adjust_precip( for force_key in self._bmi._job_meta.input_forcings: self._bmi._input_forcing_mod[force_key].skip = False - err_handler.check_program_status(self._bmi._job_meta, mpi_config) + err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) @time_function - def log_forecast( - self, - mpi_config: MpiConfig, - ) -> None: + def log_forecast(self) -> None: """Log information about the current forecast cycle. Warnings @@ -251,28 +240,25 @@ def log_forecast( Modifies mutable arguments in-place. """ # Log information about this forecast cycle - if mpi_config.rank == 0: + if self._bmi._mpi_meta.rank == 0: self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" - err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) self._bmi._job_meta.statusMsg = ( "Processing Forecast Cycle: " + self._bmi._job_meta.current_fcst_cycle.strftime("%Y-%m-%d %H:%M") ) - err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) self._bmi._job_meta.statusMsg = ( "Forecast Cycle Length is: " + str(self._bmi._job_meta.cycle_length_minutes) + " minutes" ) - err_handler.log_msg(self._bmi._job_meta, mpi_config, True) - # mpi_config.comm.barrier() + err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) + # self._bmi._mpi_meta.comm.barrier() @time_function def loop_through_forcing_products( - self, - future_time: float, - mpi_config: MpiConfig, - output_obj: OutputObj, + self, future_time: float, output_obj: OutputObj ) -> forcingInputMod.InputForcingsHydrofabric: """Loop through each forcing product and process it for the current forecast cycle. @@ -352,13 +338,13 @@ def loop_through_forcing_products( # Print message on log file indicating the timestamp # we are currently processing for forcings - if mpi_config.rank == 0 and show_message: + if self._bmi._mpi_meta.rank == 0 and show_message: self._bmi._job_meta.statusMsg = ( "=========================================" ) - err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) self._bmi._job_meta.statusMsg = f"Processing for output timestep: {file_date.strftime('%Y-%m-%d %H:%M')}" - err_handler.log_msg(self._bmi._job_meta, mpi_config, True) + err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) self._bmi._job_meta.currentForceNum = 0 self._bmi._job_meta.currentCustomForceNum = 0 @@ -388,17 +374,17 @@ def loop_through_forcing_products( else: input_forcings = self._bmi._input_forcing_mod[force_key] input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta ) if force_key in [12, 21, 27]: if self._bmi._job_meta.aws is None: # Calculate the previous and next input cycle files from the inputs. input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) else: # Flag to indicate the AWS .zarr AORC method @@ -406,14 +392,14 @@ def loop_through_forcing_products( if self.source_data_processor is None: self.source_data_processor = AORCConusProcessor( self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) elif force_key == 21: if self.source_data_processor is None: self.source_data_processor = AORCAlaskaProcessor( self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) @@ -423,7 +409,7 @@ def loop_through_forcing_products( if self._bmi._job_meta.nwm_domain == "CONUS": self.source_data_processor = NWMV3ConusProcessor( self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) elif self._bmi._job_meta.nwm_domain in [ @@ -432,13 +418,13 @@ def loop_through_forcing_products( ]: self.source_data_processor = NWMV3OConusProcessor( self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) elif self._bmi._job_meta.nwm_domain == "Alaska": self.source_data_processor = NWMV3AlaskaProcessor( self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) else: @@ -458,15 +444,19 @@ def loop_through_forcing_products( break # Regrid forcings. input_forcings.regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_forcing_bounds( - self._bmi._job_meta, input_forcings, mpi_config + self._bmi._job_meta, input_forcings, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the # next set of forcings as the previous step just regridded the previous forcing. @@ -497,47 +487,59 @@ def loop_through_forcing_products( ) # Re-calculate the neighbor files. input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Regrid the forcings for the end of the window. input_forcings.regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) input_forcings.rstFlag = 0 # Run temporal interpolation on the grids. input_forcings.temporal_interpolate_inputs( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run bias correction. bias_correction.run_bias_correction( input_forcings, self._bmi._job_meta, self._bmi.geo_meta, - mpi_config, + self._bmi._mpi_meta, + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Run downscaling on grids for this output timestep. downscale.run_downscaling( input_forcings, self._bmi._job_meta, self._bmi.geo_meta, - mpi_config, + self._bmi._mpi_meta, + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) # Layer in forcings from this product. layeringMod.layer_final_forcings( - output_obj, input_forcings, self._bmi._job_meta, mpi_config + output_obj, input_forcings, self._bmi._job_meta, self._bmi._mpi_meta + ) + err_handler.check_program_status( + self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status(self._bmi._job_meta, mpi_config) self._bmi._job_meta.currentForceNum += 1 @@ -552,18 +554,18 @@ def loop_through_forcing_products( if supp_pcp_key != 13: # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Regrid the supplemental precipitation. self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) if ( @@ -576,11 +578,11 @@ def loop_through_forcing_products( err_handler.check_supp_pcp_bounds( self._bmi._job_meta, self._bmi._supp_pcp_mod[supp_pcp_key], - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen @@ -588,20 +590,20 @@ def loop_through_forcing_products( input_forcings, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Run temporal interpolation on the grids. self._bmi._supp_pcp_mod[ supp_pcp_key ].temporal_interpolate_inputs( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Layer in the supplemental precipitation into the current output object. @@ -609,10 +611,10 @@ def loop_through_forcing_products( output_obj, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Call the output routines @@ -621,8 +623,8 @@ def loop_through_forcing_products( output_obj.outDate = file_date ################ Commenting this out to bypass NWM forcing file output functionality ######### - # output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, mpi_config) - # err_handler.check_program_status(self._bmi._job_meta, mpi_config) + # output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta) + # err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) ############################################################################################## return input_forcings @@ -630,7 +632,6 @@ def loop_through_forcing_products( @time_function def process_suplemental_precip( self, - mpi_config: MpiConfig, output_obj: OutputObj, input_forcings: dict, ) -> None: @@ -647,18 +648,18 @@ def process_suplemental_precip( if supp_pcp_key == 14: # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, mpi_config + self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Regrid the supplemental precipitation. self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) if ( @@ -671,31 +672,31 @@ def process_suplemental_precip( err_handler.check_supp_pcp_bounds( self._bmi._job_meta, self._bmi._supp_pcp_mod[supp_pcp_key], - mpi_config, + self._bmi._mpi_meta, self._bmi.geo_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) self.disaggregate_fun( input_forcings, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Run temporal interpolation on the grids. self._bmi._supp_pcp_mod[ supp_pcp_key ].temporal_interpolate_inputs( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) # Layer in the supplemental precipitation into the current output object. @@ -703,16 +704,15 @@ def process_suplemental_precip( output_obj, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, - mpi_config, + self._bmi._mpi_meta, ) err_handler.check_program_status( - self._bmi._job_meta, mpi_config + self._bmi._job_meta, self._bmi._mpi_meta ) @time_function def write_output( self, - mpi_config: MpiConfig, output_obj: OutputObj, ) -> None: """Write the output for the current forecast cycle. @@ -728,7 +728,7 @@ def write_output( or self._bmi._job_meta.grid_type == "hydrofabric" ): output_obj.gather_global_outputs( - self._bmi._job_meta, self._bmi.geo_meta, mpi_config + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) """##################Step 6: flatten and update dict##########################################################################""" diff --git a/tests/test_utils.py b/tests/test_utils.py index 13879226..f78c4474 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -527,7 +527,6 @@ def pre_regrid(self) -> None: f"In pre_regrid, expected state to be either None or 'post_ran' but got {repr(self._state)}. The test is set up incorrectly." ) - mpi_config = self.mpi_config output_obj = self.bmi_model._output_obj future_time = ( @@ -538,15 +537,11 @@ def pre_regrid(self) -> None: ### NOTE this should mimic NWMv3ForcingEngineModel.run() with the exception of setting the skip flag model.determine_forecast(future_time) - model.adjust_precip(mpi_config) - model.log_forecast(mpi_config) + model.adjust_precip() + model.log_forecast() ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() - model.loop_through_forcing_products( - future_time, - mpi_config, - output_obj, - ) + model.loop_through_forcing_products(future_time, output_obj) # Update test fixture status self._state = "pre_ran" From 07daebaa04cc661da8a5ff63a190641ed49f7bc6 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 16:01:33 -0400 Subject: [PATCH 10/36] Encapsulate OutputObj instance --- .../NextGen_Forcings_Engine/bmi_model.py | 10 +- .../NextGen_Forcings_Engine/model.py | 125 ++++++++---------- tests/test_utils.py | 4 +- 3 files changed, 61 insertions(+), 78 deletions(-) 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 0a010a24..196395d5 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/bmi_model.py @@ -469,10 +469,7 @@ def update_until(self, future_time: float): == future_time == self.cfg_bmi["initial_time"] ): - self._model.run( - future_time, - self._output_obj, - ) + self._model.run(future_time) else: # Start a while loop to iterate the model time step by step until the # current model time reaches or exceeds the future_time. @@ -480,10 +477,7 @@ def update_until(self, future_time: float): # Advance the model time by the defined time step size. self._values["current_model_time"] += self._values["time_step_size"] # Run the model for the new current time and update the state. - self._model.run( - self._values["current_model_time"], - self._output_obj, - ) + self._model.run(self._values["current_model_time"]) # ------------------------------------------------------------ def finalize(self): diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 859cde61..b47e1af6 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -84,11 +84,7 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): # def aws_obj(files): # return xr.open_mfdataset(files, engine="zarr", parallel=True, consolidated=True) - def run( - self, - future_time: float, - output_obj: OutputObj, - ) -> None: + def run(self, future_time: float) -> None: """Execute the full forcings engine BMI pipeline for a given future timestep. This method updates the `self._bmi._values` state dictionary with atmospheric forcings computed from @@ -119,7 +115,6 @@ def run( 7. Advance the BMI time index. :param future_time: The number of seconds into the future to advance the model. - :param output_obj: Output object that stores the generated atmospheric forcing arrays. :raises RuntimeError: If the model fails to initialize or if required arguments are missing. """ @@ -130,27 +125,16 @@ def run( # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. input_forcings = self.loop_through_forcing_products( future_time, - output_obj, - ) - self.process_suplemental_precip( - output_obj, - input_forcings, - ) - self.write_output( - output_obj, - ) - self.update_dict( - output_obj, ) + self.process_suplemental_precip(input_forcings) + self.write_output() + self.update_dict() ## Update BMI model time index to next iteration self._bmi._job_meta.bmi_time_index += 1 @time_function - def determine_forecast( - self, - future_time: float, - ) -> None: + def determine_forecast(self, future_time: float) -> None: """Determine the forecast for the given future time and configuration. Warnings @@ -258,7 +242,7 @@ def log_forecast(self) -> None: @time_function def loop_through_forcing_products( - self, future_time: float, output_obj: OutputObj + self, future_time: float ) -> forcingInputMod.InputForcingsHydrofabric: """Loop through each forcing product and process it for the current forecast cycle. @@ -277,14 +261,18 @@ def loop_through_forcing_products( if not self._bmi._job_meta.precip_only_flag: if self._bmi._job_meta.grid_type == "gridded": # Reset out final grids to missing values. - output_obj.output_local[:, :, :] = self._bmi._job_meta.globalNdv + self._bmi._output_obj.output_local[:, :, :] = ( + self._bmi._job_meta.globalNdv + ) elif self._bmi._job_meta.grid_type == "unstructured": # Reset out final grids to missing values. - output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv - output_obj.output_local_elem[:, :] = self._bmi._job_meta.globalNdv + self._bmi._output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv + self._bmi._output_obj.output_local_elem[:, :] = ( + self._bmi._job_meta.globalNdv + ) elif self._bmi._job_meta.grid_type == "hydrofabric": # Reset out final grids to missing values. - output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv + self._bmi._output_obj.output_local[:, :] = self._bmi._job_meta.globalNdv else: raise ValueError( f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" @@ -305,24 +293,24 @@ def loop_through_forcing_products( # Compute the output timestamp for this step if self._bmi._job_meta.ana_flag: - output_obj.outDate = ( + self._bmi._output_obj.outDate = ( self._bmi._job_meta.current_fcst_cycle + datetime.timedelta(seconds=self._bmi._job_meta.output_freq * 60) ) else: - output_obj.outDate = ( + self._bmi._output_obj.outDate = ( self._bmi._job_meta.current_fcst_cycle + datetime.timedelta(seconds=future_time) ) - self._bmi._job_meta.current_output_date = output_obj.outDate + self._bmi._job_meta.current_output_date = self._bmi._output_obj.outDate # Adjust file_date for AnA if needed file_date = ( - output_obj.outDate + self._bmi._output_obj.outDate - datetime.timedelta(seconds=self._bmi._job_meta.output_freq * 60) if self._bmi._job_meta.ana_flag - else output_obj.outDate + else self._bmi._output_obj.outDate ) # Compute previous output date (used for downscaling logic) @@ -374,14 +362,18 @@ def loop_through_forcing_products( else: input_forcings = self._bmi._input_forcing_mod[force_key] input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, ) if force_key in [12, 21, 27]: if self._bmi._job_meta.aws is None: # Calculate the previous and next input cycle files from the inputs. input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, ) err_handler.check_program_status( self._bmi._job_meta, self._bmi._mpi_meta @@ -487,7 +479,9 @@ def loop_through_forcing_products( ) # Re-calculate the neighbor files. input_forcings.calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, ) err_handler.check_program_status( self._bmi._job_meta, self._bmi._mpi_meta @@ -535,7 +529,10 @@ def loop_through_forcing_products( # Layer in forcings from this product. layeringMod.layer_final_forcings( - output_obj, input_forcings, self._bmi._job_meta, self._bmi._mpi_meta + self._bmi._output_obj, + input_forcings, + self._bmi._job_meta, + self._bmi._mpi_meta, ) err_handler.check_program_status( self._bmi._job_meta, self._bmi._mpi_meta @@ -554,7 +551,9 @@ def loop_through_forcing_products( if supp_pcp_key != 13: # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, ) err_handler.check_program_status( self._bmi._job_meta, self._bmi._mpi_meta @@ -608,7 +607,7 @@ def loop_through_forcing_products( # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( - output_obj, + self._bmi._output_obj, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, self._bmi._mpi_meta, @@ -620,21 +619,17 @@ def loop_through_forcing_products( # Call the output routines # adjust date for AnA if necessary if self._bmi._job_meta.ana_flag: - output_obj.outDate = file_date + self._bmi._output_obj.outDate = file_date ################ Commenting this out to bypass NWM forcing file output functionality ######### - # output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta) + # self._bmi._output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta) # err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) ############################################################################################## return input_forcings @time_function - def process_suplemental_precip( - self, - output_obj: OutputObj, - input_forcings: dict, - ) -> None: + def process_suplemental_precip(self, input_forcings: dict) -> None: """Process supplemental precipitation for the current forecast cycle. Warnings @@ -648,7 +643,9 @@ def process_suplemental_precip( if supp_pcp_key == 14: # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, output_obj.outDate, self._bmi._mpi_meta + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, ) err_handler.check_program_status( self._bmi._job_meta, self._bmi._mpi_meta @@ -701,7 +698,7 @@ def process_suplemental_precip( # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( - output_obj, + self._bmi._output_obj, self._bmi._supp_pcp_mod[supp_pcp_key], self._bmi._job_meta, self._bmi._mpi_meta, @@ -711,10 +708,7 @@ def process_suplemental_precip( ) @time_function - def write_output( - self, - output_obj: OutputObj, - ) -> None: + def write_output(self) -> None: """Write the output for the current forecast cycle. Warnings @@ -727,17 +721,14 @@ def write_output( self._bmi._job_meta.forcing_output == 1 or self._bmi._job_meta.grid_type == "hydrofabric" ): - output_obj.gather_global_outputs( + self._bmi._output_obj.gather_global_outputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) """##################Step 6: flatten and update dict##########################################################################""" @time_function - def update_dict( - self, - output_obj: OutputObj, - ) -> None: + def update_dict(self) -> None: """Flatten the Forcings Engine output object and update the BMI dictionary. Warnings @@ -782,22 +773,22 @@ def update_dict( ] if self._bmi._job_meta.grid_type == "gridded": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local[ - count, :, : - ].flatten() + self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._output_obj.output_local[count, :, :].flatten() + ) elif self._bmi._job_meta.grid_type == "unstructured": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = output_obj.output_local_elem[ - count, : - ].flatten() - self._bmi._values[variable + "_NODE"] = output_obj.output_local[ - count, : - ].flatten() + self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._output_obj.output_local_elem[count, :].flatten() + ) + self._bmi._values[variable + "_NODE"] = ( + self._bmi._output_obj.output_local[count, :].flatten() + ) elif self._bmi._job_meta.grid_type == "hydrofabric": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = output_obj.output_global[ - count, : - ].flatten() + self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._output_obj.output_global[count, :].flatten() + ) self._bmi._values["CAT-ID"] = self._bmi.geo_meta.element_ids_global else: raise ValueError( diff --git a/tests/test_utils.py b/tests/test_utils.py index f78c4474..0a16ed73 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -527,8 +527,6 @@ def pre_regrid(self) -> None: f"In pre_regrid, expected state to be either None or 'post_ran' but got {repr(self._state)}. The test is set up incorrectly." ) - output_obj = self.bmi_model._output_obj - future_time = ( self.bmi_model._values["current_model_time"] + self.bmi_model._values["time_step_size"] @@ -541,7 +539,7 @@ def pre_regrid(self) -> None: model.log_forecast() ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() - model.loop_through_forcing_products(future_time, output_obj) + model.loop_through_forcing_products(future_time) # Update test fixture status self._state = "pre_ran" From 381b9e4271c9a8f7ba397fccf2eeb9510fc65a23 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 16:03:13 -0400 Subject: [PATCH 11/36] Remove unused imports --- .../NextGen_Forcings_Engine/model.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index b47e1af6..661e57ec 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,8 +1,7 @@ from __future__ import annotations import datetime -import os from contextlib import contextmanager -from time import time, perf_counter +from time import perf_counter from typing import TYPE_CHECKING import ewts @@ -17,14 +16,6 @@ forcingInputMod, layeringMod, ) -from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.config import ( - ConfigOptions, -) -from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.geoMod import ( - GeoMeta, -) -from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.ioMod import OutputObj -from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.parallel import MpiConfig from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.historical_forcing import ( AORCAlaskaProcessor, AORCConusProcessor, From 4701c284358e9a4d1cbacf0a38c14967a13c9f7f Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 16:09:41 -0400 Subject: [PATCH 12/36] Remove hard-coded msg control --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 661e57ec..26953f87 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -248,7 +248,6 @@ def loop_through_forcing_products( # 4.) Downscale. # 5.) Layer, and output as necessary. ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 - show_message = True if not self._bmi._job_meta.precip_only_flag: if self._bmi._job_meta.grid_type == "gridded": # Reset out final grids to missing values. @@ -317,7 +316,7 @@ def loop_through_forcing_products( # Print message on log file indicating the timestamp # we are currently processing for forcings - if self._bmi._mpi_meta.rank == 0 and show_message: + if self._bmi._mpi_meta.rank == 0: self._bmi._job_meta.statusMsg = ( "=========================================" ) From a4e23a05064955b14bffa26162c3cb548268896b Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 16:57:44 -0400 Subject: [PATCH 13/36] DRYify some AnA deltas --- .../NextGen_Forcings_Engine/model.py | 27 ++++++++----------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 26953f87..621a01f5 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -143,33 +143,28 @@ def determine_forecast(self, future_time: float) -> None: # If we're in an AnA configuration, then must offset the BMI future # timestamp to account for the "lookback" period being properly iterated # over between 3-28 hour look back time period and operation configuration + # TODO confirm these codes, and should they consider all input_forcings not just [0]? if self._bmi._job_meta.input_forcings[0] in [20, 22]: + delta = pd.TimedeltaIndex( + np.array([future_time - 7200.0], dtype=float), "s" + )[0] self._bmi._job_meta.current_fcst_cycle = ( - self._bmi._job_meta.b_date_proc - + pd.TimedeltaIndex( - np.array([future_time - 7200.0], dtype=float), "s" - )[0] + self._bmi._job_meta.b_date_proc + delta ) self._bmi._job_meta.current_time = ( - self._bmi._job_meta.b_date_proc - + pd.TimedeltaIndex( - np.array([future_time - 7200.0], dtype=float), "s" - )[0] + self._bmi._job_meta.b_date_proc + delta ) self._bmi._job_meta.future_time = future_time else: # Puerto Rico / Hawaii AnA: 1-hour lookback (based on 6-hourly forecast cycles) + delta = pd.TimedeltaIndex( + np.array([future_time - 3600.0], dtype=float), "s" + )[0] self._bmi._job_meta.current_fcst_cycle = ( - self._bmi._job_meta.b_date_proc - + pd.TimedeltaIndex( - np.array([future_time - 3600.0], dtype=float), "s" - )[0] + self._bmi._job_meta.b_date_proc + delta ) self._bmi._job_meta.current_time = ( - self._bmi._job_meta.b_date_proc - + pd.TimedeltaIndex( - np.array([future_time - 3600.0], dtype=float), "s" - )[0] + self._bmi._job_meta.b_date_proc + delta ) else: # Forecast-only mode — use BMI timestamp as-is From 93fec41fef99a08f85b56ff9e43bd0d7aeea8e80 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 17:25:16 -0400 Subject: [PATCH 14/36] DRYify calls to check_program_status --- .../NextGen_Forcings_Engine/model.py | 92 ++++++------------- 1 file changed, 27 insertions(+), 65 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 621a01f5..de08be49 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -75,6 +75,10 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): # def aws_obj(files): # return xr.open_mfdataset(files, engine="zarr", parallel=True, consolidated=True) + def check_program_status(self) -> None: + """Call err_handler.check_program_status""" + err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) + def run(self, future_time: float) -> None: """Execute the full forcings engine BMI pipeline for a given future timestep. @@ -199,7 +203,7 @@ def adjust_precip(self) -> None: for force_key in self._bmi._job_meta.input_forcings: self._bmi._input_forcing_mod[force_key].skip = False - err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) + self.check_program_status() @time_function def log_forecast(self) -> None: @@ -360,9 +364,7 @@ def loop_through_forcing_products( self._bmi._output_obj.outDate, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() else: # Flag to indicate the AWS .zarr AORC method if force_key == 12: @@ -423,17 +425,13 @@ def loop_through_forcing_products( input_forcings.regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Run check on regridded fields for reasonable values that are not missing values. err_handler.check_forcing_bounds( self._bmi._job_meta, input_forcings, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the # next set of forcings as the previous step just regridded the previous forcing. @@ -468,17 +466,13 @@ def loop_through_forcing_products( self._bmi._output_obj.outDate, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Regrid the forcings for the end of the window. input_forcings.regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() input_forcings.rstFlag = 0 @@ -486,9 +480,7 @@ def loop_through_forcing_products( input_forcings.temporal_interpolate_inputs( self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Run bias correction. bias_correction.run_bias_correction( @@ -497,9 +489,7 @@ def loop_through_forcing_products( self._bmi.geo_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Run downscaling on grids for this output timestep. downscale.run_downscaling( @@ -508,9 +498,7 @@ def loop_through_forcing_products( self._bmi.geo_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Layer in forcings from this product. layeringMod.layer_final_forcings( @@ -519,9 +507,7 @@ def loop_through_forcing_products( self._bmi._job_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() self._bmi._job_meta.currentForceNum += 1 @@ -540,17 +526,13 @@ def loop_through_forcing_products( self._bmi._output_obj.outDate, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Regrid the supplemental precipitation. self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() if ( self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 @@ -565,9 +547,7 @@ def loop_through_forcing_products( self._bmi._mpi_meta, self._bmi.geo_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen self.disaggregate_fun( @@ -576,9 +556,7 @@ def loop_through_forcing_products( self._bmi._job_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Run temporal interpolation on the grids. self._bmi._supp_pcp_mod[ @@ -586,9 +564,7 @@ def loop_through_forcing_products( ].temporal_interpolate_inputs( self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( @@ -597,9 +573,7 @@ def loop_through_forcing_products( self._bmi._job_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Call the output routines # adjust date for AnA if necessary @@ -608,7 +582,7 @@ def loop_through_forcing_products( ################ Commenting this out to bypass NWM forcing file output functionality ######### # self._bmi._output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta) - # err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) + # self.check_program_status() ############################################################################################## return input_forcings @@ -632,17 +606,13 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: self._bmi._output_obj.outDate, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Regrid the supplemental precipitation. self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() if ( self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 @@ -657,9 +627,7 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: self._bmi._mpi_meta, self._bmi.geo_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() self.disaggregate_fun( input_forcings, @@ -667,9 +635,7 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: self._bmi._job_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Run temporal interpolation on the grids. self._bmi._supp_pcp_mod[ @@ -677,9 +643,7 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: ].temporal_interpolate_inputs( self._bmi._job_meta, self._bmi._mpi_meta ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() # Layer in the supplemental precipitation into the current output object. layeringMod.layer_supplemental_forcing( @@ -688,9 +652,7 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: self._bmi._job_meta, self._bmi._mpi_meta, ) - err_handler.check_program_status( - self._bmi._job_meta, self._bmi._mpi_meta - ) + self.check_program_status() @time_function def write_output(self) -> None: From b28735c10815665034afca00f9c6e7de10823017 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Mon, 27 Apr 2026 17:32:54 -0400 Subject: [PATCH 15/36] Clean up comments and docstrings --- .../NextGen_Forcings_Engine/model.py | 63 +++++++++---------- 1 file changed, 29 insertions(+), 34 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index de08be49..6449f87f 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,3 +1,5 @@ +"""NWMv3ForcingEngineModel, to be constructed and managed by inheritors of NWMv3_Forcing_Engine_BMI_model_Base from bmi_model.py""" + from __future__ import annotations import datetime from contextlib import contextmanager @@ -25,7 +27,6 @@ ) if TYPE_CHECKING: - # To allow type hint without circular import error from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.bmi_model import ( NWMv3_Forcing_Engine_BMI_model_Base, ) @@ -59,22 +60,15 @@ def wrapper(*args, **kwargs): class NWMv3ForcingEngineModel: - """NextGen Forcings Engine BMI model class for NWMv3 forcings.""" + """NextGen Forcings Engine BMI model class for NWMv3 forcings. + To be constructed and managed by inheritors of NWMv3_Forcing_Engine_BMI_model_Base from bmi_model.py. + """ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): """Initialize the NWMv3 Forcing Engine Model.""" self.source_data_processor = None self._bmi = bmi_model - # TODO: refactor the bmi_model.py file and this to have this type maintain its own state. - # def __init__(self): - # super(ngen_model, self).__init__() - # #self._model = model - - # @dask.delayed - # def aws_obj(files): - # return xr.open_mfdataset(files, engine="zarr", parallel=True, consolidated=True) - def check_program_status(self) -> None: """Call err_handler.check_program_status""" err_handler.check_program_status(self._bmi._job_meta, self._bmi._mpi_meta) @@ -213,7 +207,6 @@ def log_forecast(self) -> None: -------- Modifies mutable arguments in-place. """ - # Log information about this forecast cycle if self._bmi._mpi_meta.rank == 0: self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) @@ -236,16 +229,17 @@ def loop_through_forcing_products( ) -> forcingInputMod.InputForcingsHydrofabric: """Loop through each forcing product and process it for the current forecast cycle. + Loop through each output timestep. Perform the following functions: + 1.) Calculate all necessary input files per user options. + 2.) Read in input forcings from GRIB/NetCDF files. + 3.) Regrid the forcings, and temporally interpolate. + 4.) Downscale. + 5.) Layer, and output as necessary. + Warnings -------- Modifies mutable arguments in-place. """ - # Loop through each output timestep. Perform the following functions: - # 1.) Calculate all necessary input files per user options. - # 2.) Read in input forcings from GRIB/NetCDF files. - # 3.) Regrid the forcings, and temporally interpolate. - # 4.) Downscale. - # 5.) Layer, and output as necessary. ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 if not self._bmi._job_meta.precip_only_flag: if self._bmi._job_meta.grid_type == "gridded": @@ -366,6 +360,7 @@ def loop_through_forcing_products( ) self.check_program_status() else: + # TODO assert one force_key? # Flag to indicate the AWS .zarr AORC method if force_key == 12: if self.source_data_processor is None: @@ -511,6 +506,7 @@ def loop_through_forcing_products( self._bmi._job_meta.currentForceNum += 1 + # TODO what is this? if force_key == 10: self._bmi._job_meta.currentCustomForceNum += 1 @@ -657,13 +653,13 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: @time_function def write_output(self) -> None: """Write the output for the current forecast cycle. + If user requests output for given domain, then call + the I/O module to update opened netcdf file with forcing fields. Warnings -------- Modifies mutable arguments in-place. """ - # If user requests output for given domain, then call - # the I/O module to update opened netcdf file with forcing fields if ( self._bmi._job_meta.forcing_output == 1 or self._bmi._job_meta.grid_type == "hydrofabric" @@ -672,28 +668,27 @@ def write_output(self) -> None: self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta ) - """##################Step 6: flatten and update dict##########################################################################""" - @time_function def update_dict(self) -> None: """Flatten the Forcings Engine output object and update the BMI dictionary. + Loop through Forcings Engine output object + and flatten the 2D forcing array and append to + the BMI object to advertise to BMIinterface. + 0.) U-Wind (m/s) + 1.) V-Wind (m/s) + 2.) Surface incoming longwave radiation flux (W/m^2) + 3.) Precipitation rate (mm/s) + 4.) 2-meter temperature (K) + 5.) 2-meter specific humidity (kg/kg) + 6.) Surface pressure (Pa) + 7.) Surface incoming shortwave radiation flux (W/m^2) + 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations + Warnings -------- Modifies mutable arguments in-place. """ - # Now loop through Forcings Engine output object - # and flatten the 2D forcing array and append to - # the BMI object to advertise to BMIinterface - # 0.) U-Wind (m/s) - # 1.) V-Wind (m/s) - # 2.) Surface incoming longwave radiation flux (W/m^2) - # 3.) Precipitation rate (mm/s) - # 4.) 2-meter temperature (K) - # 5.) 2-meter specific humidity (kg/kg) - # 6.) Surface pressure (Pa) - # 7.) Surface incoming shortwave radiation flux (W/m^2) - # 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations if self._bmi._job_meta.include_lqfrac == 1: variables = [ From ca9cb1257406eabb45a6360ab1fcc68121b8464e Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 07:30:16 -0400 Subject: [PATCH 16/36] Add forcing key count assertion --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 6449f87f..d0383820 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -360,7 +360,10 @@ def loop_through_forcing_products( ) self.check_program_status() else: - # TODO assert one force_key? + if len(self._bmi._job_meta.input_forcings) != 1: + raise ValueError( + f"Expected to have 1 forcing key, but have {len(self._bmi._job_meta.input_forcings)}: {list(self._bmi._job_meta.input_forcings)}" + ) # Flag to indicate the AWS .zarr AORC method if force_key == 12: if self.source_data_processor is None: From e5e769a82d7b2751462504f7bbd44de8b92c1bc2 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 07:32:27 -0400 Subject: [PATCH 17/36] Update docstrings --- .../NextGen_Forcings_Engine/model.py | 26 ++----------------- 1 file changed, 2 insertions(+), 24 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index d0383820..b4702a60 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -186,12 +186,7 @@ def determine_forecast(self, future_time: float) -> None: @time_function def adjust_precip(self) -> None: - """Adjust precipitation for the given forecast cycle. - - Warnings - -------- - Modifies mutable arguments in-place. - """ + """Adjust precipitation for the given forecast cycle.""" if not self._bmi._job_meta.precip_only_flag: # reset skips if present for force_key in self._bmi._job_meta.input_forcings: @@ -201,12 +196,7 @@ def adjust_precip(self) -> None: @time_function def log_forecast(self) -> None: - """Log information about the current forecast cycle. - - Warnings - -------- - Modifies mutable arguments in-place. - """ + """Log information about the current forecast cycle.""" if self._bmi._mpi_meta.rank == 0: self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) @@ -235,10 +225,6 @@ def loop_through_forcing_products( 3.) Regrid the forcings, and temporally interpolate. 4.) Downscale. 5.) Layer, and output as necessary. - - Warnings - -------- - Modifies mutable arguments in-place. """ ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 if not self._bmi._job_meta.precip_only_flag: @@ -658,10 +644,6 @@ def write_output(self) -> None: """Write the output for the current forecast cycle. If user requests output for given domain, then call the I/O module to update opened netcdf file with forcing fields. - - Warnings - -------- - Modifies mutable arguments in-place. """ if ( self._bmi._job_meta.forcing_output == 1 @@ -687,10 +669,6 @@ def update_dict(self) -> None: 6.) Surface pressure (Pa) 7.) Surface incoming shortwave radiation flux (W/m^2) 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations - - Warnings - -------- - Modifies mutable arguments in-place. """ if self._bmi._job_meta.include_lqfrac == 1: From 52e0d08b7f6d9c3afc3e48616107197b68d04ac6 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 07:48:47 -0400 Subject: [PATCH 18/36] Move block using `rstFlag` to new private method --- .../NextGen_Forcings_Engine/model.py | 89 ++++++++++--------- 1 file changed, 49 insertions(+), 40 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index b4702a60..8a6bb2c9 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -419,46 +419,7 @@ def loop_through_forcing_products( # If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the # next set of forcings as the previous step just regridded the previous forcing. - if input_forcings.rstFlag == 1: - if ( - input_forcings.regridded_forcings1 is not None - and input_forcings.regridded_forcings2 is not None - ): - # Set the forcings back to reflect we just regridded the previous set of inputs, not the next. - if self._bmi._job_meta.grid_type == "gridded": - input_forcings.regridded_forcings1[:, :, :] = ( - input_forcings.regridded_forcings2[:, :, :] - ) - elif self._bmi._job_meta.grid_type == "unstructured": - input_forcings.regridded_forcings1[:, :] = ( - input_forcings.regridded_forcings2[:, :] - ) - input_forcings.regridded_forcings1_elem[:, :] = ( - input_forcings.regridded_forcings2_elem[:, :] - ) - elif self._bmi._job_meta.grid_type == "hydrofabric": - input_forcings.regridded_forcings1[:, :] = ( - input_forcings.regridded_forcings2[:, :] - ) - else: - raise ValueError( - f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" - ) - # Re-calculate the neighbor files. - input_forcings.calc_neighbor_files( - self._bmi._job_meta, - self._bmi._output_obj.outDate, - self._bmi._mpi_meta, - ) - self.check_program_status() - - # Regrid the forcings for the end of the window. - input_forcings.regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta - ) - self.check_program_status() - - input_forcings.rstFlag = 0 + self.__use_rstFlag(input_forcings) # Run temporal interpolation on the grids. input_forcings.temporal_interpolate_inputs( @@ -572,6 +533,54 @@ def loop_through_forcing_products( return input_forcings + def __use_rstFlag(self, input_forcings): + """ + If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the + next set of forcings as the previous step just regridded the previous forcing. + + This code block was cut and pasted from method `loop_through_forcing_products` during refactor. + """ + if input_forcings.rstFlag == 1: + if ( + input_forcings.regridded_forcings1 is not None + and input_forcings.regridded_forcings2 is not None + ): + # Set the forcings back to reflect we just regridded the previous set of inputs, not the next. + if self._bmi._job_meta.grid_type == "gridded": + input_forcings.regridded_forcings1[:, :, :] = ( + input_forcings.regridded_forcings2[:, :, :] + ) + elif self._bmi._job_meta.grid_type == "unstructured": + input_forcings.regridded_forcings1[:, :] = ( + input_forcings.regridded_forcings2[:, :] + ) + input_forcings.regridded_forcings1_elem[:, :] = ( + input_forcings.regridded_forcings2_elem[:, :] + ) + elif self._bmi._job_meta.grid_type == "hydrofabric": + input_forcings.regridded_forcings1[:, :] = ( + input_forcings.regridded_forcings2[:, :] + ) + else: + raise ValueError( + f"Unexpected grid_type: {repr(self._bmi._job_meta.grid_type)}" + ) + # Re-calculate the neighbor files. + input_forcings.calc_neighbor_files( + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, + ) + self.check_program_status() + + # Regrid the forcings for the end of the window. + input_forcings.regrid_inputs( + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta + ) + self.check_program_status() + + input_forcings.rstFlag = 0 + @time_function def process_suplemental_precip(self, input_forcings: dict) -> None: """Process supplemental precipitation for the current forecast cycle. From d490a67fb9be98d3cc1b9f533ddb66c67e41cb0a Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 08:25:03 -0400 Subject: [PATCH 19/36] Move block for supplemental precip handling to new private method --- .../NextGen_Forcings_Engine/model.py | 166 +++++++----------- 1 file changed, 59 insertions(+), 107 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 8a6bb2c9..2c207373 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -466,60 +466,9 @@ def loop_through_forcing_products( if self._bmi._job_meta.number_supp_pcp > 0: for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key != 13: - # Like with input forcings, calculate the neighboring files to use. - self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, - self._bmi._output_obj.outDate, - self._bmi._mpi_meta, - ) - self.check_program_status() - - # Regrid the supplemental precipitation. - self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta - ) - self.check_program_status() - - if ( - self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 - is not None - and self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip2 - is not None - ): - # Run check on regridded fields for reasonable values that are not missing values. - err_handler.check_supp_pcp_bounds( - self._bmi._job_meta, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - self.check_program_status() - - # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen - self.disaggregate_fun( - input_forcings, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._job_meta, - self._bmi._mpi_meta, - ) - self.check_program_status() - - # Run temporal interpolation on the grids. - self._bmi._supp_pcp_mod[ - supp_pcp_key - ].temporal_interpolate_inputs( - self._bmi._job_meta, self._bmi._mpi_meta - ) - self.check_program_status() - - # Layer in the supplemental precipitation into the current output object. - layeringMod.layer_supplemental_forcing( - self._bmi._output_obj, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._job_meta, - self._bmi._mpi_meta, - ) - self.check_program_status() + # Below comment copied from earlier code, the comment had been just above the call to `disaggregate_fun`. + # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen + self.__process_supp_precip_key(input_forcings, supp_pcp_key) # Call the output routines # adjust date for AnA if necessary @@ -533,6 +482,61 @@ def loop_through_forcing_products( return input_forcings + def __process_supp_precip_key(self, input_forcings: dict, supp_pcp_key: int): + """Process supplemental precipitation for one supplemental precipitation key. + + This code block was cut and pasted from methods `loop_through_forcing_products` and `process_suplemental_precip` during refactor. + """ + # Like with input forcings, calculate the neighboring files to use. + self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, + ) + self.check_program_status() + + # Regrid the supplemental precipitation. + self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( + self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta + ) + self.check_program_status() + + if ( + self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 is not None + and self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip2 is not None + ): + # Run check on regridded fields for reasonable values that are not missing values. + err_handler.check_supp_pcp_bounds( + self._bmi._job_meta, + self._bmi._supp_pcp_mod[supp_pcp_key], + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + self.check_program_status() + + self.disaggregate_fun( + input_forcings, + self._bmi._supp_pcp_mod[supp_pcp_key], + self._bmi._job_meta, + self._bmi._mpi_meta, + ) + self.check_program_status() + + # Run temporal interpolation on the grids. + self._bmi._supp_pcp_mod[supp_pcp_key].temporal_interpolate_inputs( + self._bmi._job_meta, self._bmi._mpi_meta + ) + self.check_program_status() + + # Layer in the supplemental precipitation into the current output object. + layeringMod.layer_supplemental_forcing( + self._bmi._output_obj, + self._bmi._supp_pcp_mod[supp_pcp_key], + self._bmi._job_meta, + self._bmi._mpi_meta, + ) + self.check_program_status() + def __use_rstFlag(self, input_forcings): """ If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the @@ -594,59 +598,7 @@ def process_suplemental_precip(self, input_forcings: dict) -> None: if self._bmi._job_meta.number_supp_pcp > 0: for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key == 14: - # Like with input forcings, calculate the neighboring files to use. - self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( - self._bmi._job_meta, - self._bmi._output_obj.outDate, - self._bmi._mpi_meta, - ) - self.check_program_status() - - # Regrid the supplemental precipitation. - self._bmi._supp_pcp_mod[supp_pcp_key].regrid_inputs( - self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta - ) - self.check_program_status() - - if ( - self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip1 - is not None - and self._bmi._supp_pcp_mod[supp_pcp_key].regridded_precip2 - is not None - ): - # Run check on regridded fields for reasonable values that are not missing values. - err_handler.check_supp_pcp_bounds( - self._bmi._job_meta, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - self.check_program_status() - - self.disaggregate_fun( - input_forcings, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._job_meta, - self._bmi._mpi_meta, - ) - self.check_program_status() - - # Run temporal interpolation on the grids. - self._bmi._supp_pcp_mod[ - supp_pcp_key - ].temporal_interpolate_inputs( - self._bmi._job_meta, self._bmi._mpi_meta - ) - self.check_program_status() - - # Layer in the supplemental precipitation into the current output object. - layeringMod.layer_supplemental_forcing( - self._bmi._output_obj, - self._bmi._supp_pcp_mod[supp_pcp_key], - self._bmi._job_meta, - self._bmi._mpi_meta, - ) - self.check_program_status() + self.__process_supp_precip_key(input_forcings, supp_pcp_key) @time_function def write_output(self) -> None: From 2eb5db36ceebfa25b54a0ccfab2f5b6fec4faa22 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 08:28:11 -0400 Subject: [PATCH 20/36] docstrings and type hints --- .../NextGen_Forcings_Engine/model.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 2c207373..81479602 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -482,10 +482,16 @@ def loop_through_forcing_products( return input_forcings - def __process_supp_precip_key(self, input_forcings: dict, supp_pcp_key: int): + def __process_supp_precip_key( + self, input_forcings: dict, supp_pcp_key: int + ) -> None: """Process supplemental precipitation for one supplemental precipitation key. This code block was cut and pasted from methods `loop_through_forcing_products` and `process_suplemental_precip` during refactor. + + Warnings + -------- + Modifies mutable arguments in-place. """ # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( @@ -537,12 +543,16 @@ def __process_supp_precip_key(self, input_forcings: dict, supp_pcp_key: int): ) self.check_program_status() - def __use_rstFlag(self, input_forcings): + def __use_rstFlag(self, input_forcings: dict) -> None: """ If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the next set of forcings as the previous step just regridded the previous forcing. This code block was cut and pasted from method `loop_through_forcing_products` during refactor. + + Warnings + -------- + Modifies mutable arguments in-place. """ if input_forcings.rstFlag == 1: if ( From 9de57cccfb1575d8f3f28cdba8d6d24e3cb592e7 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 08:42:10 -0400 Subject: [PATCH 21/36] Move block for AORC and NWM handling to new private method --- .../NextGen_Forcings_Engine/model.py | 140 ++++++++++-------- 1 file changed, 76 insertions(+), 64 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 81479602..d23beb11 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -336,70 +336,8 @@ def loop_through_forcing_products( self._bmi._mpi_meta, ) - if force_key in [12, 21, 27]: - if self._bmi._job_meta.aws is None: - # Calculate the previous and next input cycle files from the inputs. - input_forcings.calc_neighbor_files( - self._bmi._job_meta, - self._bmi._output_obj.outDate, - self._bmi._mpi_meta, - ) - self.check_program_status() - else: - if len(self._bmi._job_meta.input_forcings) != 1: - raise ValueError( - f"Expected to have 1 forcing key, but have {len(self._bmi._job_meta.input_forcings)}: {list(self._bmi._job_meta.input_forcings)}" - ) - # Flag to indicate the AWS .zarr AORC method - if force_key == 12: - if self.source_data_processor is None: - self.source_data_processor = AORCConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - elif force_key == 21: - if self.source_data_processor is None: - self.source_data_processor = AORCAlaskaProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - - # Flag to indicate the AWS .zarr NWMv3 Forcing file method - elif force_key == 27: - if self.source_data_processor is None: - if self._bmi._job_meta.nwm_domain == "CONUS": - self.source_data_processor = NWMV3ConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - elif self._bmi._job_meta.nwm_domain in [ - "Hawaii", - "PR", - ]: - self.source_data_processor = NWMV3OConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - elif self._bmi._job_meta.nwm_domain == "Alaska": - self.source_data_processor = NWMV3AlaskaProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) - else: - raise ValueError( - f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" - ) - - self._bmi._job_meta.aws_obj = ( - self.source_data_processor.process_historical_data( - self._bmi._job_meta.current_time - ) - ) + # Handle AORC and NWM force keys + self.__handle_aorc_and_nwm_force_keys(input_forcings, force_key) # If skipping this forcing, continue early if input_forcings.skip is True: @@ -482,6 +420,80 @@ def loop_through_forcing_products( return input_forcings + def __handle_aorc_and_nwm_force_keys(self, input_forcings, force_key: int) -> None: + """During `loop_through_forcing_products`, handle the case of the force key being AORC or NWM. + + This code block was cut and pasted from methods `loop_through_forcing_products` during refactor. + + Warnings + -------- + Modifies mutable arguments in-place. + """ + if force_key in [12, 21, 27]: + if self._bmi._job_meta.aws is None: + # Calculate the previous and next input cycle files from the inputs. + input_forcings.calc_neighbor_files( + self._bmi._job_meta, + self._bmi._output_obj.outDate, + self._bmi._mpi_meta, + ) + self.check_program_status() + else: + if len(self._bmi._job_meta.input_forcings) != 1: + raise ValueError( + f"Expected to have 1 forcing key, but have {len(self._bmi._job_meta.input_forcings)}: {list(self._bmi._job_meta.input_forcings)}" + ) + # Flag to indicate the AWS .zarr AORC method + if force_key == 12: + if self.source_data_processor is None: + self.source_data_processor = AORCConusProcessor( + self._bmi._job_meta, + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + elif force_key == 21: + if self.source_data_processor is None: + self.source_data_processor = AORCAlaskaProcessor( + self._bmi._job_meta, + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + + # Flag to indicate the AWS .zarr NWMv3 Forcing file method + elif force_key == 27: + if self.source_data_processor is None: + if self._bmi._job_meta.nwm_domain == "CONUS": + self.source_data_processor = NWMV3ConusProcessor( + self._bmi._job_meta, + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + elif self._bmi._job_meta.nwm_domain in [ + "Hawaii", + "PR", + ]: + self.source_data_processor = NWMV3OConusProcessor( + self._bmi._job_meta, + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + elif self._bmi._job_meta.nwm_domain == "Alaska": + self.source_data_processor = NWMV3AlaskaProcessor( + self._bmi._job_meta, + self._bmi._mpi_meta, + self._bmi.geo_meta, + ) + else: + raise ValueError( + f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" + ) + + self._bmi._job_meta.aws_obj = ( + self.source_data_processor.process_historical_data( + self._bmi._job_meta.current_time + ) + ) + def __process_supp_precip_key( self, input_forcings: dict, supp_pcp_key: int ) -> None: From db7b07e27d916a7d113220df59e7e21d9f3cbb64 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 08:46:40 -0400 Subject: [PATCH 22/36] Comments and type hints --- .../NextGen_Forcings_Engine/model.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index d23beb11..fa34117c 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -340,9 +340,11 @@ def loop_through_forcing_products( self.__handle_aorc_and_nwm_force_keys(input_forcings, force_key) # If skipping this forcing, continue early + # NOTE this is used by the esmf regrid pytests, to halt the loop before "manually" calling a particular regrid function. if input_forcings.skip is True: LOG.debug(f"Breaking loop for force_key {force_key}") break + # Regrid forcings. input_forcings.regrid_inputs( self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta @@ -495,7 +497,7 @@ def __handle_aorc_and_nwm_force_keys(self, input_forcings, force_key: int) -> No ) def __process_supp_precip_key( - self, input_forcings: dict, supp_pcp_key: int + self, input_forcings: forcingInputMod.InputForcings, supp_pcp_key: int ) -> None: """Process supplemental precipitation for one supplemental precipitation key. @@ -555,7 +557,7 @@ def __process_supp_precip_key( ) self.check_program_status() - def __use_rstFlag(self, input_forcings: dict) -> None: + def __use_rstFlag(self, input_forcings: forcingInputMod.InputForcings) -> None: """ If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the next set of forcings as the previous step just regridded the previous forcing. @@ -608,7 +610,9 @@ def __use_rstFlag(self, input_forcings: dict) -> None: input_forcings.rstFlag = 0 @time_function - def process_suplemental_precip(self, input_forcings: dict) -> None: + def process_suplemental_precip( + self, input_forcings: forcingInputMod.InputForcings + ) -> None: """Process supplemental precipitation for the current forecast cycle. Warnings From 1827ae3f590b48bf637ee492ab6d72ce93ca4273 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 08:59:01 -0400 Subject: [PATCH 23/36] Move model.py BMI variables to consts.py --- .../NextGen_Forcings_Engine/core/consts.py | 16 +++++++++++ .../NextGen_Forcings_Engine/model.py | 28 ++++--------------- 2 files changed, 22 insertions(+), 22 deletions(-) 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 aad7cb71..31f4cd4c 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py @@ -955,6 +955,22 @@ "precipBiasCorrectOpt", ], } + +MODEL = { + # Used by method `model.NWMv3ForcingEngineModel.update_dict` + "update_dict_base_vars": [ + "U2D", + "V2D", + "LWDOWN", + "RAINRATE", + "T2D", + "Q2D", + "PSFC", + "SWDOWN", + ], + "update_dict_var_include_lqfraq": "LQFRAC", +} + TEST_UTILS = { "OLD_NEW_VAR_MAP": { "q2dBiasCorrectOpt": "q2BiasCorrectOpt", diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index fa34117c..66830511 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -18,6 +18,9 @@ forcingInputMod, layeringMod, ) +from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.core.consts import ( + MODEL as model_consts, +) from NextGen_Forcings_Engine_BMI.NextGen_Forcings_Engine.historical_forcing import ( AORCAlaskaProcessor, AORCConusProcessor, @@ -658,29 +661,10 @@ def update_dict(self) -> None: 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations """ + variables = model_consts["update_dict_base_vars"] if self._bmi._job_meta.include_lqfrac == 1: - variables = [ - "U2D", - "V2D", - "LWDOWN", - "RAINRATE", - "T2D", - "Q2D", - "PSFC", - "SWDOWN", - "LQFRAC", - ] - else: - variables = [ - "U2D", - "V2D", - "LWDOWN", - "RAINRATE", - "T2D", - "Q2D", - "PSFC", - "SWDOWN", - ] + variables.append(model_consts["update_dict_var_include_lqfraq"]) + if self._bmi._job_meta.grid_type == "gridded": for count, variable in enumerate(variables): self._bmi._values[variable + "_ELEMENT"] = ( From 421312c669b297522bf05344258e812c5a0e89ee Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 09:17:45 -0400 Subject: [PATCH 24/36] Fix usage of new consts.MODEL list --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 66830511..13cecbb2 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,6 +1,7 @@ """NWMv3ForcingEngineModel, to be constructed and managed by inheritors of NWMv3_Forcing_Engine_BMI_model_Base from bmi_model.py""" from __future__ import annotations +import copy import datetime from contextlib import contextmanager from time import perf_counter @@ -661,7 +662,7 @@ def update_dict(self) -> None: 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations """ - variables = model_consts["update_dict_base_vars"] + variables = copy.deepcopy(model_consts["update_dict_base_vars"]) if self._bmi._job_meta.include_lqfrac == 1: variables.append(model_consts["update_dict_var_include_lqfraq"]) From 38a56d3e81e5862e8cad1411fb24cce2d92aaaee Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Wed, 29 Apr 2026 13:24:28 -0400 Subject: [PATCH 25/36] Comments --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 13cecbb2..7f994d1d 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -147,6 +147,7 @@ def determine_forecast(self, future_time: float) -> None: # over between 3-28 hour look back time period and operation configuration # TODO confirm these codes, and should they consider all input_forcings not just [0]? if self._bmi._job_meta.input_forcings[0] in [20, 22]: + # NOTE This appears to be intending to operate on Alaska-only AnA. delta = pd.TimedeltaIndex( np.array([future_time - 7200.0], dtype=float), "s" )[0] @@ -158,6 +159,7 @@ def determine_forecast(self, future_time: float) -> None: ) self._bmi._job_meta.future_time = future_time else: + # NOTE below comment was original, but this appears to be operating on all non-Alaska AnA, not just Puerto Rico / Hawaii AnA. # Puerto Rico / Hawaii AnA: 1-hour lookback (based on 6-hourly forecast cycles) delta = pd.TimedeltaIndex( np.array([future_time - 3600.0], dtype=float), "s" From 77eff0a38252fe5199c17e2de25bc8861e5d6849 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Thu, 30 Apr 2026 16:19:38 -0400 Subject: [PATCH 26/36] Add NotImplementedError for SubOutputHour --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 7f994d1d..97f2616e 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -261,6 +261,9 @@ def loop_through_forcing_products( # Optional sub-output timestamp if self._bmi._job_meta.sub_output_hour is not None: + raise NotImplementedError( + f"sub_output_hour (config SubOutputHour) is {repr(self._bmi._job_meta.sub_output_hour)} (not None) but is not used." + ) # TODO This is not used subOutDate = self._bmi._job_meta.first_fcst_cycle + datetime.timedelta( hours=self._bmi._job_meta.sub_output_hour From ea4923f2bdfd2e1f54f6db9452c1c6694427182c Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 14:04:18 -0400 Subject: [PATCH 27/36] Type hints and docstrings --- .../NextGen_Forcings_Engine/model.py | 75 +++++++++++++++---- 1 file changed, 60 insertions(+), 15 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 97f2616e..1252ff84 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -40,11 +40,12 @@ @contextmanager def timing_block(step_str: str): - """Context manager for timing code execution. - - Args: - step_str: Description of the step being timed. + """Context manager for timing code execution. Used by decorator `time_function`. + Parameters + ---------- + step_str : str + Description of the step being timed. """ start = perf_counter() yield @@ -53,7 +54,7 @@ def timing_block(step_str: str): def time_function(func): - """Measure the execution time of a function.""" + """Decorator for measuring the execution time of a function.""" def wrapper(*args, **kwargs): with timing_block(f"Executing {func.__name__}"): @@ -65,11 +66,17 @@ def wrapper(*args, **kwargs): class NWMv3ForcingEngineModel: """NextGen Forcings Engine BMI model class for NWMv3 forcings. + To be constructed and managed by inheritors of NWMv3_Forcing_Engine_BMI_model_Base from bmi_model.py. """ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): - """Initialize the NWMv3 Forcing Engine Model.""" + """Initialize the NWMv3 Forcing Engine Model. + + Parameters + ---------- + bmi_model : NWMv3_Forcing_Engine_BMI_model_Base + """ self.source_data_processor = None self._bmi = bmi_model @@ -107,9 +114,17 @@ def run(self, future_time: float) -> None: 6. Update the self._bmi._values state dictionary with flattened arrays. 7. Advance the BMI time index. - :param future_time: The number of seconds into the future to advance the model. - - :raises RuntimeError: If the model fails to initialize or if required arguments are missing. + Parameters + ---------- + future_time : float + Timestamp, represented as *seconds relative to overall start time*, to advance to before returning. + Since this is relative to overall start time, it is unaware of the actual UTC datetimestamp of the start. + For example, since 1-hour timesteps are typical, the first value of this would typically be 3600, the second value 7200, etc. + + Raises + ------ + RuntimeError + If the model fails to initialize or if required arguments are missing. """ self.determine_forecast(future_time) @@ -132,7 +147,7 @@ def determine_forecast(self, future_time: float) -> None: Warnings -------- - Modifies mutable arguments in-place. + Modifies mutable arguments in-place. """ # Assign the future time to the configuration self._bmi._job_meta.bmi_time = future_time @@ -231,6 +246,15 @@ def loop_through_forcing_products( 3.) Regrid the forcings, and temporally interpolate. 4.) Downscale. 5.) Layer, and output as necessary. + + Parameters + ---------- + future_time : float + See description in `self.run`. + + Returns + ---------- + input_forcings: forcingInputMod.InputForcings """ ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 if not self._bmi._job_meta.precip_only_flag: @@ -431,14 +455,21 @@ def loop_through_forcing_products( return input_forcings - def __handle_aorc_and_nwm_force_keys(self, input_forcings, force_key: int) -> None: + def __handle_aorc_and_nwm_force_keys( + self, input_forcings: forcingInputMod.InputForcings, force_key: int + ) -> None: """During `loop_through_forcing_products`, handle the case of the force key being AORC or NWM. This code block was cut and pasted from methods `loop_through_forcing_products` during refactor. + Parameters + ---------- + input_forcings : forcingInputMod.InputForcings + force_key : int + Warnings -------- - Modifies mutable arguments in-place. + Modifies mutable arguments in-place. """ if force_key in [12, 21, 27]: if self._bmi._job_meta.aws is None: @@ -512,9 +543,14 @@ def __process_supp_precip_key( This code block was cut and pasted from methods `loop_through_forcing_products` and `process_suplemental_precip` during refactor. + Parameters + ---------- + input_forcings : forcingInputMod.InputForcings + supp_pcp_key : int + Warnings -------- - Modifies mutable arguments in-place. + Modifies mutable arguments in-place. """ # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( @@ -573,9 +609,13 @@ def __use_rstFlag(self, input_forcings: forcingInputMod.InputForcings) -> None: This code block was cut and pasted from method `loop_through_forcing_products` during refactor. + Parameters + ---------- + input_forcings : forcingInputMod.InputForcings + Warnings -------- - Modifies mutable arguments in-place. + Modifies mutable arguments in-place. """ if input_forcings.rstFlag == 1: if ( @@ -624,9 +664,13 @@ def process_suplemental_precip( ) -> None: """Process supplemental precipitation for the current forecast cycle. + Parameters + ---------- + input_forcings : forcingInputMod.InputForcings + Warnings -------- - Modifies mutable arguments in-place. + Modifies mutable arguments in-place. """ if self._bmi._job_meta.customSuppPcpFreq is not None: # Process supplemental precipitation if we specified in the configuration file. @@ -638,6 +682,7 @@ def process_suplemental_precip( @time_function def write_output(self) -> None: """Write the output for the current forecast cycle. + If user requests output for given domain, then call the I/O module to update opened netcdf file with forcing fields. """ From 936ef112de3c4659a11fcfeda409e670b055b941 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 19:42:16 -0400 Subject: [PATCH 28/36] Rename methods --- .../NextGen_Forcings_Engine/core/consts.py | 2 +- .../NextGen_Forcings_Engine/model.py | 23 ++++++++----------- tests/test_utils.py | 9 ++++---- 3 files changed, 16 insertions(+), 18 deletions(-) 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 31f4cd4c..9444dc8c 100644 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/core/consts.py @@ -957,7 +957,7 @@ } MODEL = { - # Used by method `model.NWMv3ForcingEngineModel.update_dict` + # Used by method `model.NWMv3ForcingEngineModel.update_bmi_output_dict` "update_dict_base_vars": [ "U2D", "V2D", diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 1252ff84..943dfc24 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -127,22 +127,19 @@ def run(self, future_time: float) -> None: If the model fails to initialize or if required arguments are missing. """ - self.determine_forecast(future_time) - self.adjust_precip() - self.log_forecast() - # TODO look into input_forcings usage in `process_suplemental_precip` and in `loop_through_forcing_products` at `disaggregate_fun`. - input_forcings = self.loop_through_forcing_products( - future_time, - ) + self.set_cycle_timing_attrs(future_time) + self.set_skip_flags() + self.log_cycle() + input_forcings = self.loop_through_forcing_products(future_time) self.process_suplemental_precip(input_forcings) self.write_output() - self.update_dict() + self.update_bmi_output_dict() ## Update BMI model time index to next iteration self._bmi._job_meta.bmi_time_index += 1 @time_function - def determine_forecast(self, future_time: float) -> None: + def set_cycle_timing_attrs(self, future_time: float) -> None: """Determine the forecast for the given future time and configuration. Warnings @@ -206,7 +203,7 @@ def determine_forecast(self, future_time: float) -> None: ) @time_function - def adjust_precip(self) -> None: + def set_skip_flags(self) -> None: """Adjust precipitation for the given forecast cycle.""" if not self._bmi._job_meta.precip_only_flag: # reset skips if present @@ -216,7 +213,7 @@ def adjust_precip(self) -> None: self.check_program_status() @time_function - def log_forecast(self) -> None: + def log_cycle(self) -> None: """Log information about the current forecast cycle.""" if self._bmi._mpi_meta.rank == 0: self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" @@ -254,7 +251,7 @@ def loop_through_forcing_products( Returns ---------- - input_forcings: forcingInputMod.InputForcings + input_forcings: forcingInputMod.InputForcings | None """ ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 if not self._bmi._job_meta.precip_only_flag: @@ -695,7 +692,7 @@ def write_output(self) -> None: ) @time_function - def update_dict(self) -> None: + def update_bmi_output_dict(self) -> None: """Flatten the Forcings Engine output object and update the BMI dictionary. Loop through Forcings Engine output object diff --git a/tests/test_utils.py b/tests/test_utils.py index 0a16ed73..16ce4458 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -533,10 +533,11 @@ def pre_regrid(self) -> None: ) model = self.bmi_model._model - ### NOTE this should mimic NWMv3ForcingEngineModel.run() with the exception of setting the skip flag - model.determine_forecast(future_time) - model.adjust_precip() - model.log_forecast() + # NOTE this should mimic NWMv3ForcingEngineModel.run() + # with the exception of externally setting the skip flags within this class. + model.set_cycle_timing_attrs(future_time) + model.set_skip_flags() + model.log_cycle() ### NOTE setting the flag causes the regrid step to be skipped self.set_input_forcings_skip_flags() model.loop_through_forcing_products(future_time) From 9410efa2fb3c14272f5b5d80ac273cd7ac4764a5 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 19:43:47 -0400 Subject: [PATCH 29/36] Fix input_forcings reference (return None conditionally) --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 943dfc24..5e62f153 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -234,7 +234,7 @@ def log_cycle(self) -> None: @time_function def loop_through_forcing_products( self, future_time: float - ) -> forcingInputMod.InputForcingsHydrofabric: + ) -> forcingInputMod.InputForcingsHydrofabric | None: """Loop through each forcing product and process it for the current forecast cycle. Loop through each output timestep. Perform the following functions: @@ -449,6 +449,8 @@ def loop_through_forcing_products( # self._bmi._output_obj.output_final_ldasin(self._bmi._job_meta, self._bmi.geo_meta, self._bmi._mpi_meta) # self.check_program_status() ############################################################################################## + else: + input_forcings = None return input_forcings From 03106704c5bc700572b41b1f7d5995190473672e Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 20:02:45 -0400 Subject: [PATCH 30/36] Use partials for log calls and use MPI-aware log methods --- .../NextGen_Forcings_Engine/model.py | 60 +++++++++---------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 5e62f153..d3aa22d6 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -4,6 +4,7 @@ import copy import datetime from contextlib import contextmanager +from functools import partial from time import perf_counter from typing import TYPE_CHECKING @@ -50,7 +51,7 @@ def timing_block(step_str: str): start = perf_counter() yield end = perf_counter() - LOG.debug(f" Execution time for {step_str}: {round(end - start, 2)} seconds") + LOG.debug(msg=f" Execution time for {step_str}: {round(end - start, 2)} seconds") def time_function(func): @@ -79,6 +80,13 @@ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): """ self.source_data_processor = None self._bmi = bmi_model + # Partials + self.log_info = partial( + err_handler.log_msg, self._bmi._job_meta, self._bmi._mpi_meta, False + ) + self.log_debug = partial( + err_handler.log_msg, self._bmi._job_meta, self._bmi._mpi_meta, True + ) def check_program_status(self) -> None: """Call err_handler.check_program_status""" @@ -189,12 +197,12 @@ def set_cycle_timing_attrs(self, future_time: float) -> None: self._bmi._job_meta.b_date_proc ) + pd.to_timedelta(future_time, unit="s") - LOG.debug( - "NextGen Forcings Engine processing meteorological forcings for BMI timestamp" + self.log_debug( + msg="NextGen Forcings Engine processing meteorological forcings for BMI timestamp" ) - LOG.debug(f"Model.py current time: {self._bmi._job_meta.current_time}") - LOG.debug( - f"Model.py current fcst cycle: {self._bmi._job_meta.current_fcst_cycle}" + self.log_debug(msg=f"Model.py current time: {self._bmi._job_meta.current_time}") + self.log_debug( + msg=f"Model.py current fcst cycle: {self._bmi._job_meta.current_fcst_cycle}" ) if self._bmi._job_meta.first_fcst_cycle is None: @@ -216,19 +224,13 @@ def set_skip_flags(self) -> None: def log_cycle(self) -> None: """Log information about the current forecast cycle.""" if self._bmi._mpi_meta.rank == 0: - self._bmi._job_meta.statusMsg = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" - err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) - self._bmi._job_meta.statusMsg = ( - "Processing Forecast Cycle: " - + self._bmi._job_meta.current_fcst_cycle.strftime("%Y-%m-%d %H:%M") + self.log_debug(msg="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX") + self.log_debug( + msg=f"Processing Forecast Cycle: {self._bmi._job_meta.current_fcst_cycle.strftime('%Y-%m-%d %H:%M')}" ) - err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) - self._bmi._job_meta.statusMsg = ( - "Forecast Cycle Length is: " - + str(self._bmi._job_meta.cycle_length_minutes) - + " minutes" + self.log_debug( + msg=f"Forecast Cycle Length is: {self._bmi._job_meta.cycle_length_minutes!s} minutes" ) - err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) # self._bmi._mpi_meta.comm.barrier() @time_function @@ -326,26 +328,24 @@ def loop_through_forcing_products( # Print message on log file indicating the timestamp # we are currently processing for forcings if self._bmi._mpi_meta.rank == 0: - self._bmi._job_meta.statusMsg = ( - "=========================================" + self.log_debug(msg="=========================================") + self.log_debug( + msg=f"Processing for output timestep: {file_date.strftime('%Y-%m-%d %H:%M')}" ) - err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) - self._bmi._job_meta.statusMsg = f"Processing for output timestep: {file_date.strftime('%Y-%m-%d %H:%M')}" - err_handler.log_msg(self._bmi._job_meta, self._bmi._mpi_meta, True) self._bmi._job_meta.currentForceNum = 0 self._bmi._job_meta.currentCustomForceNum = 0 - LOG.debug( - f"config_options.input_forcings: {self._bmi._job_meta.input_forcings}" + self.log_debug( + msg=f"config_options.input_forcings: {self._bmi._job_meta.input_forcings}" ) # Loop over each of the input forcings specified. - LOG.debug( - f"Model.py forcing loop: {len(self._bmi._job_meta.input_forcings)} forcings configured: {self._bmi._job_meta.input_forcings}" + self.log_debug( + msg=f"Model.py forcing loop: {len(self._bmi._job_meta.input_forcings)} forcings configured: {self._bmi._job_meta.input_forcings}" ) for force_key in self._bmi._job_meta.input_forcings: - LOG.debug(f"force_key: {force_key}") - LOG.debug(f"config_options.aws: {self._bmi._job_meta.aws}") + self.log_debug(msg=f"force_key: {force_key}") + self.log_debug(msg=f"config_options.aws: {self._bmi._job_meta.aws}") # Pass these methods for AORC data is ERA5-Interim blend is requested # so we can finish filling in the missing gaps if ( @@ -372,7 +372,7 @@ def loop_through_forcing_products( # If skipping this forcing, continue early # NOTE this is used by the esmf regrid pytests, to halt the loop before "manually" calling a particular regrid function. if input_forcings.skip is True: - LOG.debug(f"Breaking loop for force_key {force_key}") + self.log_debug(msg=f"Breaking loop for force_key {force_key}") break # Regrid forcings. @@ -430,7 +430,7 @@ def loop_through_forcing_products( if force_key == 10: self._bmi._job_meta.currentCustomForceNum += 1 - LOG.debug(f"End of loop for force_key {force_key}") + self.log_debug(msg=f"End of loop for force_key {force_key}") # Process supplemental precipitation if we specified in the configuration file. if self._bmi._job_meta.number_supp_pcp > 0: From 8afcc8d991485ecd372a9e5fd75e339ec901a974 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 20:03:09 -0400 Subject: [PATCH 31/36] Run flynt -tc -ll 9999 --- .../NextGen_Forcings_Engine/model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index d3aa22d6..687bed42 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -717,20 +717,20 @@ def update_bmi_output_dict(self) -> None: if self._bmi._job_meta.grid_type == "gridded": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._values[f"{variable}_ELEMENT"] = ( self._bmi._output_obj.output_local[count, :, :].flatten() ) elif self._bmi._job_meta.grid_type == "unstructured": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._values[f"{variable}_ELEMENT"] = ( self._bmi._output_obj.output_local_elem[count, :].flatten() ) - self._bmi._values[variable + "_NODE"] = ( + self._bmi._values[f"{variable}_NODE"] = ( self._bmi._output_obj.output_local[count, :].flatten() ) elif self._bmi._job_meta.grid_type == "hydrofabric": for count, variable in enumerate(variables): - self._bmi._values[variable + "_ELEMENT"] = ( + self._bmi._values[f"{variable}_ELEMENT"] = ( self._bmi._output_obj.output_global[count, :].flatten() ) self._bmi._values["CAT-ID"] = self._bmi.geo_meta.element_ids_global From 1b9fc8a925cb35a52e285a56d8016e8e2d877d54 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 20:19:10 -0400 Subject: [PATCH 32/36] Format docstrings to reST --- .../NextGen_Forcings_Engine/model.py | 181 ++++++++---------- 1 file changed, 80 insertions(+), 101 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 687bed42..336dfb06 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -1,6 +1,7 @@ """NWMv3ForcingEngineModel, to be constructed and managed by inheritors of NWMv3_Forcing_Engine_BMI_model_Base from bmi_model.py""" from __future__ import annotations + import copy import datetime from contextlib import contextmanager @@ -41,12 +42,9 @@ @contextmanager def timing_block(step_str: str): - """Context manager for timing code execution. Used by decorator `time_function`. + """Context manager for timing code execution. Used by the decorator ``time_function``. - Parameters - ---------- - step_str : str - Description of the step being timed. + :param str step_str: Description of the step being timed. """ start = perf_counter() yield @@ -72,11 +70,9 @@ class NWMv3ForcingEngineModel: """ def __init__(self, bmi_model: NWMv3_Forcing_Engine_BMI_model_Base): - """Initialize the NWMv3 Forcing Engine Model. + """Initialize the NWMv3 Forcing Engine model. - Parameters - ---------- - bmi_model : NWMv3_Forcing_Engine_BMI_model_Base + :param bmi_model NWMv3_Forcing_Engine_BMI_model_Base: BMI model instance to initialize. """ self.source_data_processor = None self._bmi = bmi_model @@ -95,44 +91,47 @@ def check_program_status(self) -> None: def run(self, future_time: float) -> None: """Execute the full forcings engine BMI pipeline for a given future timestep. - This method updates the `self._bmi._values` state dictionary with atmospheric forcings computed from - available input datasets. It handles initialization, AWS Zarr loading, regridding, temporal - interpolation, bias correction, downscaling, supplemental precipitation processing, and output - population into the self._bmi._values structure. + This method updates the ``self._bmi._values`` state dictionary with atmospheric + forcings computed from available input datasets. It handles initialization, + AWS Zarr loading, regridding, temporal interpolation, bias correction, + downscaling, supplemental precipitation processing, and output population into + the ``self._bmi._values`` structure. - `self._bmi._job_meta`, an instance of ConfigOptions is also updated in-place, for example for time handling. + ``self._bmi._job_meta``, an instance of ``ConfigOptions``, is also updated + in-place, for example for forecast time handling. The following steps are performed: 1. Determine the current forecast and output times based on the future timestamp - and analysis mode (AnA or forecast). + and analysis mode (AnA or forecast). 2. Initialize or reset output grids and step counters. 3. Loop over each input forcing product: - a. Calculate neighboring input files. - b. Load AWS-hosted Zarr datasets if needed. - c. Regrid input forcings to the model grid. - d. Perform temporal interpolation. - e. Apply bias correction and downscaling. - f. Layer final forcings into the output object. + + a. Calculate neighboring input files. + b. Load AWS-hosted Zarr datasets if needed. + c. Regrid input forcings to the model grid. + d. Perform temporal interpolation. + e. Apply bias correction and downscaling. + f. Layer final forcings into the output object. + 4. Optionally process supplemental precipitation forcings: - a. Regrid and validate. - b. Disaggregate and interpolate. - c. Layer into the final output. + + a. Regrid and validate. + b. Disaggregate and interpolate. + c. Layer into the final output. + 5. Write output to NetCDF forcing files if requested. - 6. Update the self._bmi._values state dictionary with flattened arrays. + 6. Update the ``self._bmi._values`` state dictionary with flattened arrays. 7. Advance the BMI time index. - Parameters - ---------- - future_time : float - Timestamp, represented as *seconds relative to overall start time*, to advance to before returning. - Since this is relative to overall start time, it is unaware of the actual UTC datetimestamp of the start. - For example, since 1-hour timesteps are typical, the first value of this would typically be 3600, the second value 7200, etc. - - Raises - ------ - RuntimeError - If the model fails to initialize or if required arguments are missing. + :param float future_time: Timestamp, represented as *seconds relative to overall + start time*, to advance to before returning. Since this value is relative + to the overall start time, it is unaware of the actual UTC datetimestamp of + the start. For example, since 1-hour timesteps are typical, the first value + would typically be 3600, the second value 7200, etc. + + :raises RuntimeError: If the model fails to initialize or if required arguments + are missing. """ self.set_cycle_timing_attrs(future_time) @@ -150,9 +149,7 @@ def run(self, future_time: float) -> None: def set_cycle_timing_attrs(self, future_time: float) -> None: """Determine the forecast for the given future time and configuration. - Warnings - -------- - Modifies mutable arguments in-place. + :warning: Modifies mutable arguments in-place """ # Assign the future time to the configuration self._bmi._job_meta.bmi_time = future_time @@ -239,21 +236,17 @@ def loop_through_forcing_products( ) -> forcingInputMod.InputForcingsHydrofabric | None: """Loop through each forcing product and process it for the current forecast cycle. - Loop through each output timestep. Perform the following functions: - 1.) Calculate all necessary input files per user options. - 2.) Read in input forcings from GRIB/NetCDF files. - 3.) Regrid the forcings, and temporally interpolate. - 4.) Downscale. - 5.) Layer, and output as necessary. - - Parameters - ---------- - future_time : float - See description in `self.run`. - - Returns - ---------- - input_forcings: forcingInputMod.InputForcings | None + Loop through each output timestep and perform the following steps: + + 1. Calculate all necessary input files per user options. + 2. Read input forcings from GRIB/NetCDF files. + 3. Regrid the forcings and perform temporal interpolation. + 4. Downscale. + 5. Layer and write output as necessary. + + :param float future_time: See description in ``self.run``. + :returns: Processed input forcings for the current timestep. + :rtype: forcingInputMod.InputForcings | None """ ana_factor = 1 if self._bmi._job_meta.ana_flag is False else 0 if not self._bmi._job_meta.precip_only_flag: @@ -457,18 +450,14 @@ def loop_through_forcing_products( def __handle_aorc_and_nwm_force_keys( self, input_forcings: forcingInputMod.InputForcings, force_key: int ) -> None: - """During `loop_through_forcing_products`, handle the case of the force key being AORC or NWM. + """During ``loop_through_forcing_products``, handle the case where the force key is AORC or NWM. - This code block was cut and pasted from methods `loop_through_forcing_products` during refactor. + This code block was cut and pasted from the method ``loop_through_forcing_products`` during refactor. - Parameters - ---------- - input_forcings : forcingInputMod.InputForcings - force_key : int + :param input_forcings forcingInputMod.InputForcings: Input forcings object to be modified. + :param int force_key: Identifier for the forcing type. - Warnings - -------- - Modifies mutable arguments in-place. + :warning: Modifies mutable arguments in-place. """ if force_key in [12, 21, 27]: if self._bmi._job_meta.aws is None: @@ -538,18 +527,15 @@ def __handle_aorc_and_nwm_force_keys( def __process_supp_precip_key( self, input_forcings: forcingInputMod.InputForcings, supp_pcp_key: int ) -> None: - """Process supplemental precipitation for one supplemental precipitation key. + """Process supplemental precipitation for a single supplemental precipitation key. - This code block was cut and pasted from methods `loop_through_forcing_products` and `process_suplemental_precip` during refactor. + This code block was cut and pasted from the methods + ``loop_through_forcing_products`` and ``process_suplemental_precip`` during refactor. - Parameters - ---------- - input_forcings : forcingInputMod.InputForcings - supp_pcp_key : int + :param input_forcings forcingInputMod.InputForcings: Input forcings object to be modified. + :param int supp_pcp_key: Identifier for the supplemental precipitation forcing. - Warnings - -------- - Modifies mutable arguments in-place. + :warning: Modifies mutable arguments in-place. """ # Like with input forcings, calculate the neighboring files to use. self._bmi._supp_pcp_mod[supp_pcp_key].calc_neighbor_files( @@ -602,19 +588,15 @@ def __process_supp_precip_key( self.check_program_status() def __use_rstFlag(self, input_forcings: forcingInputMod.InputForcings) -> None: - """ - If we are restarting a forecast cycle, re-calculate the neighboring files, and regrid the - next set of forcings as the previous step just regridded the previous forcing. + """If restarting a forecast cycle, re-calculate neighboring files and regrid the + next set of forcings, as the previous step regridded the prior forcing. - This code block was cut and pasted from method `loop_through_forcing_products` during refactor. + This code block was cut and pasted from the method + ``loop_through_forcing_products`` during refactor. - Parameters - ---------- - input_forcings : forcingInputMod.InputForcings + :param input_forcings forcingInputMod.InputForcings: Input forcings object to be modified. - Warnings - -------- - Modifies mutable arguments in-place. + :warning: Modifies mutable arguments in-place. """ if input_forcings.rstFlag == 1: if ( @@ -663,13 +645,9 @@ def process_suplemental_precip( ) -> None: """Process supplemental precipitation for the current forecast cycle. - Parameters - ---------- - input_forcings : forcingInputMod.InputForcings + :param input_forcings forcingInputMod.InputForcings: Input forcings object to be modified. - Warnings - -------- - Modifies mutable arguments in-place. + :warning: Modifies mutable arguments in-place. """ if self._bmi._job_meta.customSuppPcpFreq is not None: # Process supplemental precipitation if we specified in the configuration file. @@ -697,20 +675,21 @@ def write_output(self) -> None: def update_bmi_output_dict(self) -> None: """Flatten the Forcings Engine output object and update the BMI dictionary. - Loop through Forcings Engine output object - and flatten the 2D forcing array and append to - the BMI object to advertise to BMIinterface. - 0.) U-Wind (m/s) - 1.) V-Wind (m/s) - 2.) Surface incoming longwave radiation flux (W/m^2) - 3.) Precipitation rate (mm/s) - 4.) 2-meter temperature (K) - 5.) 2-meter specific humidity (kg/kg) - 6.) Surface pressure (Pa) - 7.) Surface incoming shortwave radiation flux (W/m^2) - 8.) Liquid Precipitation Fraction (%), Only available in certain operational configurations - """ + Loop through the Forcings Engine output object, flatten the 2D forcing arrays, + and append them to the BMI object for advertisement through the BMI interface. + + The flattened variables are ordered as follows: + 0. U-wind (m/s) + 1. V-wind (m/s) + 2. Surface incoming longwave radiation flux (W/m²) + 3. Precipitation rate (mm/s) + 4. 2-meter air temperature (K) + 5. 2-meter specific humidity (kg/kg) + 6. Surface pressure (Pa) + 7. Surface incoming shortwave radiation flux (W/m²) + 8. Liquid precipitation fraction (%), available only in certain operational configurations + """ variables = copy.deepcopy(model_consts["update_dict_base_vars"]) if self._bmi._job_meta.include_lqfrac == 1: variables.append(model_consts["update_dict_var_include_lqfraq"]) From 6528a83c597fe06eb8083a77c707a2db94ad192e Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 20:35:50 -0400 Subject: [PATCH 33/36] DRYify --- .../NextGen_Forcings_Engine/model.py | 35 +++++-------------- 1 file changed, 9 insertions(+), 26 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 336dfb06..11491aa5 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -419,7 +419,7 @@ def loop_through_forcing_products( self._bmi._job_meta.currentForceNum += 1 - # TODO what is this? + # NOTE currentCustomForceNum does not appear to be used. if force_key == 10: self._bmi._job_meta.currentCustomForceNum += 1 @@ -459,6 +459,8 @@ def __handle_aorc_and_nwm_force_keys( :warning: Modifies mutable arguments in-place. """ + proc_args = (self._bmi._job_meta, self._bmi._mpi_meta, self._bmi.geo_meta) + if force_key in [12, 21, 27]: if self._bmi._job_meta.aws is None: # Calculate the previous and next input cycle files from the inputs. @@ -476,48 +478,29 @@ def __handle_aorc_and_nwm_force_keys( # Flag to indicate the AWS .zarr AORC method if force_key == 12: if self.source_data_processor is None: - self.source_data_processor = AORCConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) + proc_cls = AORCConusProcessor elif force_key == 21: if self.source_data_processor is None: - self.source_data_processor = AORCAlaskaProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) + proc_cls = AORCAlaskaProcessor # Flag to indicate the AWS .zarr NWMv3 Forcing file method elif force_key == 27: if self.source_data_processor is None: if self._bmi._job_meta.nwm_domain == "CONUS": - self.source_data_processor = NWMV3ConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) + proc_cls = NWMV3ConusProcessor elif self._bmi._job_meta.nwm_domain in [ "Hawaii", "PR", ]: - self.source_data_processor = NWMV3OConusProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) + proc_cls = NWMV3OConusProcessor elif self._bmi._job_meta.nwm_domain == "Alaska": - self.source_data_processor = NWMV3AlaskaProcessor( - self._bmi._job_meta, - self._bmi._mpi_meta, - self._bmi.geo_meta, - ) + proc_cls = NWMV3AlaskaProcessor else: raise ValueError( f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" ) + self.source_data_processor = proc_cls(*proc_args) self._bmi._job_meta.aws_obj = ( self.source_data_processor.process_historical_data( self._bmi._job_meta.current_time From cca2cf1a4c62d0e4240b107641bae7befc061d33 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 20:36:50 -0400 Subject: [PATCH 34/36] Add assertion --- NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 11491aa5..452f44b0 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -499,6 +499,8 @@ def __handle_aorc_and_nwm_force_keys( raise ValueError( f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" ) + else: + raise ValueError(f"Unexpected force_key: {force_key}") self.source_data_processor = proc_cls(*proc_args) self._bmi._job_meta.aws_obj = ( From 547f9d8567deeb09367ff51f50bc28920fd086fd Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 21:04:34 -0400 Subject: [PATCH 35/36] Fix source_data_processor sets --- .../NextGen_Forcings_Engine/model.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 452f44b0..9ef3c123 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -475,17 +475,14 @@ def __handle_aorc_and_nwm_force_keys( raise ValueError( f"Expected to have 1 forcing key, but have {len(self._bmi._job_meta.input_forcings)}: {list(self._bmi._job_meta.input_forcings)}" ) - # Flag to indicate the AWS .zarr AORC method - if force_key == 12: - if self.source_data_processor is None: + if self.source_data_processor is None: + # Flag to indicate the AWS .zarr AORC method + if force_key == 12: proc_cls = AORCConusProcessor - elif force_key == 21: - if self.source_data_processor is None: + elif force_key == 21: proc_cls = AORCAlaskaProcessor - - # Flag to indicate the AWS .zarr NWMv3 Forcing file method - elif force_key == 27: - if self.source_data_processor is None: + # Flag to indicate the AWS .zarr NWMv3 Forcing file method + elif force_key == 27: if self._bmi._job_meta.nwm_domain == "CONUS": proc_cls = NWMV3ConusProcessor elif self._bmi._job_meta.nwm_domain in [ @@ -499,10 +496,10 @@ def __handle_aorc_and_nwm_force_keys( raise ValueError( f"Unsupported domain type ({self._bmi._job_meta.nwm_domain} for forcing type: {force_key} )" ) - else: - raise ValueError(f"Unexpected force_key: {force_key}") + else: + raise ValueError(f"Unexpected force_key: {force_key}") + self.source_data_processor = proc_cls(*proc_args) - self.source_data_processor = proc_cls(*proc_args) self._bmi._job_meta.aws_obj = ( self.source_data_processor.process_historical_data( self._bmi._job_meta.current_time From 270f8115f33b2e0da09b69e944e232c3027af7a4 Mon Sep 17 00:00:00 2001 From: Max Kipp Date: Fri, 1 May 2026 21:07:30 -0400 Subject: [PATCH 36/36] Docstrings --- .../NextGen_Forcings_Engine/model.py | 24 ++++++++----------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py index 9ef3c123..16b1c4a7 100755 --- a/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py +++ b/NextGen_Forcings_Engine_BMI/NextGen_Forcings_Engine/model.py @@ -106,20 +106,16 @@ def run(self, future_time: float) -> None: and analysis mode (AnA or forecast). 2. Initialize or reset output grids and step counters. 3. Loop over each input forcing product: - - a. Calculate neighboring input files. - b. Load AWS-hosted Zarr datasets if needed. - c. Regrid input forcings to the model grid. - d. Perform temporal interpolation. - e. Apply bias correction and downscaling. - f. Layer final forcings into the output object. - + a. Calculate neighboring input files. + b. Load AWS-hosted Zarr datasets if needed. + c. Regrid input forcings to the model grid. + d. Perform temporal interpolation. + e. Apply bias correction and downscaling. + f. Layer final forcings into the output object. 4. Optionally process supplemental precipitation forcings: - - a. Regrid and validate. - b. Disaggregate and interpolate. - c. Layer into the final output. - + a. Regrid and validate. + b. Disaggregate and interpolate. + c. Layer into the final output. 5. Write output to NetCDF forcing files if requested. 6. Update the ``self._bmi._values`` state dictionary with flattened arrays. 7. Advance the BMI time index. @@ -429,7 +425,7 @@ def loop_through_forcing_products( if self._bmi._job_meta.number_supp_pcp > 0: for supp_pcp_key in self._bmi._job_meta.supp_precip_forcings: if supp_pcp_key != 13: - # Below comment copied from earlier code, the comment had been just above the call to `disaggregate_fun`. + # Below comment copied from earlier code, the comment had been just above the call to ``disaggregate_fun``. # TODO input_forcings has not yet been initialized, so this is a bug waiting to happen self.__process_supp_precip_key(input_forcings, supp_pcp_key)