fix bug with drop overlapping assim flag (#597)

Author: samlamont

src/teehr/fetching/nwm/grid_utils.py

@@ -183,7 +183,8 @@ def fetch_and_format_nwm_grids(
     overwrite_output: bool,
     location_id_prefix: Union[str, None],
     variable_mapper: Dict[str, Dict[str, str]],
-    timeseries_type: TimeseriesTypeEnum
+    timeseries_type: TimeseriesTypeEnum,
+    drop_overlapping_assimilation_values: bool
 ):
     """Compute weighted average, grouping by reference time.
 
@@ -243,6 +244,11 @@ def fetch_and_format_nwm_grids(
             Path(output_parquet_dir), f"{ref_time_str}.parquet"
         )
         z_hour_df.sort_values([LOCATION_ID, VALUE_TIME], inplace=True)
+
+        if drop_overlapping_assimilation_values and "assim" in nwm_configuration_name:
+            # Set reference_time to NaT for assimilation values
+            z_hour_df.loc[:, REFERENCE_TIME] = pd.NaT
+
         write_timeseries_parquet_file(
             filepath=parquet_filepath,
             overwrite_output=overwrite_output,

src/teehr/fetching/nwm/nwm_grids.py

@@ -382,5 +382,6 @@ def nwm_grids_to_parquet(
             overwrite_output=overwrite_output,
             location_id_prefix=location_id_prefix,
             variable_mapper=variable_mapper,
-            timeseries_type=timeseries_type
+            timeseries_type=timeseries_type,
+            drop_overlapping_assimilation_values=drop_overlapping_assimilation_values
         )

src/teehr/fetching/nwm/nwm_points.py

@@ -348,5 +348,6 @@ def nwm_to_parquet(
             overwrite_output,
             nwm_version,
             variable_mapper,
-            timeseries_type
+            timeseries_type,
+            drop_overlapping_assimilation_values
         )

src/teehr/fetching/nwm/point_utils.py

@@ -102,7 +102,8 @@ def process_chunk_of_files(
     overwrite_output: bool,
     nwm_version: str,
     variable_mapper: Dict[str, Dict[str, str]],
-    timeseries_type: TimeseriesTypeEnum
+    timeseries_type: TimeseriesTypeEnum,
+    drop_overlapping_assimilation_values: bool
 ):
     """Assemble a table for a chunk of NWM files."""
     location_ids = np.array(location_ids).astype(int)
@@ -156,6 +157,12 @@ def process_chunk_of_files(
         end = f"{end_json[1]}T{end_json[3][1:3]}F{end_json[6][1:]}"
         filename = f"{start}_{end}.parquet"
 
+    if drop_overlapping_assimilation_values and "assim" in configuration:
+        # Set reference_time to NaT for assimilation values
+        df_output = output_table.to_pandas()
+        df_output.loc[:, REFERENCE_TIME] = pd.NaT
+        output_table = pa.Table.from_pandas(df_output, schema=schema)
+
     write_timeseries_parquet_file(
         Path(output_parquet_dir, filename),
         overwrite_output,
@@ -176,7 +183,8 @@ def fetch_and_format_nwm_points(
     overwrite_output: bool,
     nwm_version: str,
     variable_mapper: Dict[str, Dict[str, str]],
-    timeseries_type: TimeseriesTypeEnum
+    timeseries_type: TimeseriesTypeEnum,
+    drop_overlapping_assimilation_values: bool
 ):
     """Fetch NWM point data and save as parquet files.
 
@@ -211,6 +219,12 @@ def fetch_and_format_nwm_points(
         they already exist.  True = overwrite; False = fail.
     nwm_version : str
         Specified NWM version.
+    variable_mapper : Dict[str, Dict[str, str]]
+        A mapping dictionary for variable names and units.
+    timeseries_type : TimeseriesTypeEnum
+        The type of timeseries being processed.
+    drop_overlapping_assimilation_values : bool
+        Whether to drop assimilation values that overlap in value_time.
     """
     output_parquet_dir = Path(output_parquet_dir)
     if not output_parquet_dir.exists():
@@ -241,5 +255,6 @@ def fetch_and_format_nwm_points(
             overwrite_output,
             nwm_version,
             variable_mapper,
-            timeseries_type
+            timeseries_type,
+            drop_overlapping_assimilation_values
         )