ngiab-data-preprocess 4.2.1__py3-none-any.whl → 4.3.0__py3-none-any.whl

data_processing/create_realization.py

@@ -1,30 +1,32 @@
 import json
+import logging
 import multiprocessing
+import shutil
 import sqlite3
 from datetime import datetime
 from pathlib import Path
-import shutil
-import requests
+from typing import Dict, Optional
 
 import pandas
+import requests
 import s3fs
 import xarray as xr
-import logging
-from collections import defaultdict
-from dask.distributed import Client, LocalCluster
+from data_processing.dask_utils import temp_cluster
 from data_processing.file_paths import file_paths
 from data_processing.gpkg_utils import (
     GeoPackage,
+    get_cat_to_nex_flowpairs,
     get_cat_to_nhd_feature_id,
     get_table_crs_short,
-    get_cat_to_nex_flowpairs,
 )
-from tqdm.rich import tqdm
 from pyproj import Transformer
+from tqdm.rich import tqdm
 
 logger = logging.getLogger(__name__)
 
-def get_approximate_gw_storage(paths: file_paths, start_date: datetime):
+
+@temp_cluster
+def get_approximate_gw_storage(paths: file_paths, start_date: datetime) -> Dict[str, int]:
     # get the gw levels from the NWM output on a given start date
     # this kind of works in place of warmstates for now
     year = start_date.strftime("%Y")
@@ -34,17 +36,10 @@ def get_approximate_gw_storage(paths: file_paths, start_date: datetime):
     fs = s3fs.S3FileSystem(anon=True)
     nc_url = f"s3://noaa-nwm-retrospective-3-0-pds/CONUS/netcdf/GWOUT/{year}/{formatted_dt}.GWOUT_DOMAIN1"
 
-    # make sure there's a dask cluster running
-    try:
-        client = Client.current()
-    except ValueError:
-        cluster = LocalCluster()
-        client = Client(cluster)
-
     with fs.open(nc_url) as file_obj:
-        ds = xr.open_dataset(file_obj)
+        ds = xr.open_dataset(file_obj)  # type: ignore
 
-        water_levels = dict()
+        water_levels: Dict[str, int] = dict()
         for cat, feature in tqdm(cat_to_feature.items()):
             # this value is in CM, we need meters to match max_gw_depth
             # xarray says it's in mm, with 0.1 scale factor. calling .values doesn't apply the scale
@@ -78,7 +73,9 @@ def make_cfe_config(
             slope=row["mean.slope_1km"],
             smcmax=row["mean.smcmax_soil_layers_stag=2"],
             smcwlt=row["mean.smcwlt_soil_layers_stag=2"],
-            max_gw_storage=row["mean.Zmax"]/1000 if row["mean.Zmax"] is not None else "0.011[m]", # mean.Zmax is in mm!
+            max_gw_storage=row["mean.Zmax"] / 1000
+            if row["mean.Zmax"] is not None
+            else "0.011[m]",  # mean.Zmax is in mm!
             gw_Coeff=row["mean.Coeff"] if row["mean.Coeff"] is not None else "0.0018[m h-1]",
             gw_Expon=row["mode.Expon"],
             gw_storage="{:.5}".format(gw_storage_ratio),
@@ -92,7 +89,6 @@ def make_cfe_config(
 def make_noahowp_config(
     base_dir: Path, divide_conf_df: pandas.DataFrame, start_time: datetime, end_time: datetime
 ) -> None:
-
     divide_conf_df.set_index("divide_id", inplace=True)
     start_datetime = start_time.strftime("%Y%m%d%H%M")
     end_datetime = end_time.strftime("%Y%m%d%H%M")
@@ -110,15 +106,15 @@ def make_noahowp_config(
                     end_datetime=end_datetime,
                     lat=divide_conf_df.loc[divide, "latitude"],
                     lon=divide_conf_df.loc[divide, "longitude"],
-                    terrain_slope= divide_conf_df.loc[divide, "mean.slope_1km"],
-                    azimuth= divide_conf_df.loc[divide, "circ_mean.aspect"],
-                    ISLTYP=int(divide_conf_df.loc[divide, "mode.ISLTYP"]),
-                    IVGTYP=int(divide_conf_df.loc[divide, "mode.IVGTYP"]),
+                    terrain_slope=divide_conf_df.loc[divide, "mean.slope_1km"],
+                    azimuth=divide_conf_df.loc[divide, "circ_mean.aspect"],
+                    ISLTYP=int(divide_conf_df.loc[divide, "mode.ISLTYP"]),  # type: ignore
+                    IVGTYP=int(divide_conf_df.loc[divide, "mode.IVGTYP"]),  # type: ignore
                 )
             )
 
 
-def get_model_attributes_modspatialite(hydrofabric: Path):
+def get_model_attributes_modspatialite(hydrofabric: Path) -> pandas.DataFrame:
     # modspatialite is faster than pyproj but can't be added as a pip dependency
     # This incantation took a while
     with GeoPackage(hydrofabric) as conn:
@@ -149,7 +145,7 @@ def get_model_attributes_modspatialite(hydrofabric: Path):
     return divide_conf_df
 
 
-def get_model_attributes_pyproj(hydrofabric: Path):
+def get_model_attributes_pyproj(hydrofabric: Path) -> pandas.DataFrame:
     # if modspatialite is not available, use pyproj
     with sqlite3.connect(hydrofabric) as conn:
         sql = """
@@ -182,7 +178,8 @@ def get_model_attributes_pyproj(hydrofabric: Path):
 
     return divide_conf_df
 
-def get_model_attributes(hydrofabric: Path):
+
+def get_model_attributes(hydrofabric: Path) -> pandas.DataFrame:
     try:
         with GeoPackage(hydrofabric) as conn:
             conf_df = pandas.read_sql_query(
@@ -205,30 +202,31 @@ def get_model_attributes(hydrofabric: Path):
             )
     except sqlite3.OperationalError:
         with sqlite3.connect(hydrofabric) as conn:
-            conf_df = pandas.read_sql_query("SELECT* FROM 'divide-attributes';", conn,)
+            conf_df = pandas.read_sql_query(
+                "SELECT* FROM 'divide-attributes';",
+                conn,
+            )
         source_crs = get_table_crs_short(hydrofabric, "divides")
         transformer = Transformer.from_crs(source_crs, "EPSG:4326", always_xy=True)
-        lon, lat = transformer.transform(
-            conf_df["centroid_x"].values, conf_df["centroid_y"].values
-        )
+        lon, lat = transformer.transform(conf_df["centroid_x"].values, conf_df["centroid_y"].values)
         conf_df["longitude"] = lon
         conf_df["latitude"] = lat
 
         conf_df.drop(columns=["centroid_x", "centroid_y"], axis=1, inplace=True)
     return conf_df
 
+
 def make_em_config(
     hydrofabric: Path,
     output_dir: Path,
     template_path: Path = file_paths.template_em_config,
 ):
-
     # test if modspatialite is available
     try:
         divide_conf_df = get_model_attributes_modspatialite(hydrofabric)
     except Exception as e:
         logger.warning(f"mod_spatialite not available, using pyproj instead: {e}")
-        logger.warning(f"Install mod_spatialite for improved performance")
+        logger.warning("Install mod_spatialite for improved performance")
         divide_conf_df = get_model_attributes_pyproj(hydrofabric)
 
     cat_config_dir = output_dir / "cat_config" / "empirical_model"
@@ -255,8 +253,7 @@ def make_em_config(
 
 def configure_troute(
     cat_id: str, config_dir: Path, start_time: datetime, end_time: datetime
-) -> int:
-
+) -> None:
     with open(file_paths.template_troute_config, "r") as file:
         troute_template = file.read()
     time_step_size = 300
@@ -269,7 +266,7 @@ def configure_troute(
         geo_file_path=f"./config/{cat_id}_subset.gpkg",
         start_datetime=start_time.strftime("%Y-%m-%d %H:%M:%S"),
         nts=nts,
-        max_loop_size=nts,        
+        max_loop_size=nts,
     )
 
     with open(config_dir / "troute.yaml", "w") as file:
@@ -301,9 +298,7 @@ def create_em_realization(cat_id: str, start_time: datetime, end_time: datetime)
         f.write(em_config)
 
     configure_troute(cat_id, paths.config_dir, start_time, end_time)
-    make_ngen_realization_json(
-        paths.config_dir, template_path, start_time, end_time
-    )
+    make_ngen_realization_json(paths.config_dir, template_path, start_time, end_time)
     make_em_config(paths.geopackage_path, paths.config_dir)
     # create some partitions for parallelization
     paths.setup_run_folders()
@@ -315,7 +310,7 @@ def create_realization(
     start_time: datetime,
     end_time: datetime,
     use_nwm_gw: bool = False,
-    gage_id: str = None,
+    gage_id: Optional[str] = None,
 ):
     paths = file_paths(cat_id)
 
@@ -324,15 +319,14 @@ def create_realization(
     if gage_id is not None:
         # try and download s3:communityhydrofabric/hydrofabrics/community/gage_parameters/gage_id
         # if it doesn't exist, use the default
-        try:
-            url = f"https://communityhydrofabric.s3.us-east-1.amazonaws.com/hydrofabrics/community/gage_parameters/{gage_id}.json"
-
+        url = f"https://communityhydrofabric.s3.us-east-1.amazonaws.com/hydrofabrics/community/gage_parameters/{gage_id}.json"
+        response = requests.get(url)
+        if response.status_code == 200:
             new_template = requests.get(url).json()
-            template_path = paths.config_dir / "calibrated_params.json"
+            template_path = paths.config_dir / "downloaded_params.json"
             with open(template_path, "w") as f:
                 json.dump(new_template, f)
-        except Exception as e:
-            logger.warning(f"Failed to download gage parameters")
+            logger.info(f"downloaded calibrated parameters for {gage_id}")
 
     conf_df = get_model_attributes(paths.geopackage_path)
 
@@ -347,21 +341,19 @@ def create_realization(
 
     configure_troute(cat_id, paths.config_dir, start_time, end_time)
 
-    make_ngen_realization_json(
-        paths.config_dir, template_path, start_time, end_time
-    )
+    make_ngen_realization_json(paths.config_dir, template_path, start_time, end_time)
 
     # create some partitions for parallelization
     paths.setup_run_folders()
     create_partitions(paths)
 
 
-def create_partitions(paths: Path, num_partitions: int = None) -> None:
+def create_partitions(paths: file_paths, num_partitions: Optional[int] = None) -> None:
     if num_partitions is None:
         num_partitions = multiprocessing.cpu_count()
 
     cat_to_nex_pairs = get_cat_to_nex_flowpairs(hydrofabric=paths.geopackage_path)
-    nexus = defaultdict(list)
+    # nexus = defaultdict(list)
 
     # for cat, nex in cat_to_nex_pairs:
     #     nexus[nex].append(cat)

data_processing/dask_utils.py

@@ -0,0 +1,92 @@
+import logging
+
+from dask.distributed import Client
+
+logger = logging.getLogger(__name__)
+
+
+def shutdown_cluster():
+    try:
+        client = Client.current()
+        client.shutdown()
+    except ValueError:
+        logger.debug("No cluster found to shutdown")
+
+
+def no_cluster(func):
+    """
+    Decorator that ensures the wrapped function runs with no active Dask cluster.
+
+    This decorator attempts to shut down any existing Dask cluster before
+    executing the wrapped function. If no cluster is found, it logs a debug message
+    and continues execution.
+
+    Parameters:
+        func: The function to be executed without a Dask cluster
+
+    Returns:
+        wrapper: The wrapped function that will be executed without a Dask cluster
+    """
+
+    def wrapper(*args, **kwargs):
+        shutdown_cluster()
+        result = func(*args, **kwargs)
+        return result
+
+    return wrapper
+
+
+def use_cluster(func):
+    """
+    Decorator that ensures the wrapped function has access to a Dask cluster.
+
+    If a Dask cluster is already running, it uses the existing one.
+    If no cluster is available, it creates a new one before executing the function.
+    The cluster remains active after the function completes.
+
+    Parameters:
+        func: The function to be executed with a Dask cluster
+
+    Returns:
+        wrapper: The wrapped function with access to a Dask cluster
+    """
+
+    def wrapper(*args, **kwargs):
+        try:
+            client = Client.current()
+        except ValueError:
+            client = Client()
+        result = func(*args, **kwargs)
+        return result
+
+    return wrapper
+
+
+def temp_cluster(func):
+    """
+    Decorator that provides a temporary Dask cluster for the wrapped function.
+
+    If a Dask cluster is already running, it uses the existing one and leaves it running.
+    If no cluster exists, it creates a temporary one and shuts it down after
+    the function completes.
+
+    Parameters:
+        func: The function to be executed with a Dask cluster
+
+    Returns:
+        wrapper: The wrapped function with access to a Dask cluster
+    """
+
+    def wrapper(*args, **kwargs):
+        cluster_was_running = True
+        try:
+            client = Client.current()
+        except ValueError:
+            cluster_was_running = False
+            client = Client()
+        result = func(*args, **kwargs)
+        if not cluster_was_running:
+            client.shutdown()
+        return result
+
+    return wrapper

data_processing/dataset_utils.py

@@ -1,19 +1,22 @@
 import logging
 import os
+from datetime import datetime
 from pathlib import Path
-from typing import Tuple, Union
+from typing import List, Literal, Optional, Tuple, Union
 
 import geopandas as gpd
 import numpy as np
 import xarray as xr
-from dask.distributed import Client,  progress
-import datetime
+from xarray.core.types import InterpOptions
+from dask.distributed import Client, progress, Future
+from data_processing.dask_utils import use_cluster
 
 logger = logging.getLogger(__name__)
 
 # known ngen variable names
 # https://github.com/CIROH-UA/ngen/blob/4fb5bb68dc397298bca470dfec94db2c1dcb42fe/include/forcing/AorcForcing.hpp#L77
 
+
 def validate_dataset_format(dataset: xr.Dataset) -> None:
     """
     Validate the format of the dataset.
@@ -41,8 +44,9 @@ def validate_dataset_format(dataset: xr.Dataset) -> None:
     if "name" not in dataset.attrs:
         raise ValueError("Dataset must have a name attribute to identify it")
 
+
 def validate_time_range(dataset: xr.Dataset, start_time: str, end_time: str) -> Tuple[str, str]:
-    '''
+    """
     Ensure that all selected times are in the passed dataset.
 
     Parameters
@@ -60,7 +64,7 @@ def validate_time_range(dataset: xr.Dataset, start_time: str, end_time: str) ->
         start_time, or if not available, earliest available timestep in dataset.
     str
         end_time, or if not available, latest available timestep in dataset.
-    '''
+    """
     end_time_in_dataset = dataset.time.isel(time=-1).values
     start_time_in_dataset = dataset.time.isel(time=0).values
     if np.datetime64(start_time) < start_time_in_dataset:
@@ -77,7 +81,10 @@ def validate_time_range(dataset: xr.Dataset, start_time: str, end_time: str) ->
 
 
 def clip_dataset_to_bounds(
-    dataset: xr.Dataset, bounds: Tuple[float, float, float, float], start_time: str, end_time: str
+    dataset: xr.Dataset,
+    bounds: Tuple[float, float, float, float] | np.ndarray[tuple[int], np.dtype[np.float64]],
+    start_time: str,
+    end_time: str,
 ) -> xr.Dataset:
     """
     Clip the dataset to specified geographical bounds.
@@ -86,14 +93,14 @@ def clip_dataset_to_bounds(
     ----------
     dataset : xr.Dataset
         Dataset to be clipped.
-    bounds : tuple[float, float, float, float]
-        Corners of bounding box. bounds[0] is x_min, bounds[1] is y_min, 
+    bounds : tuple[float, float, float, float] | np.ndarray[tuple[int], np.dtype[np.float64]]
+        Corners of bounding box. bounds[0] is x_min, bounds[1] is y_min,
         bounds[2] is x_max, bounds[3] is y_max.
     start_time : str
         Desired start time in YYYY/MM/DD HH:MM:SS format.
     end_time : str
         Desired end time in YYYY/MM/DD HH:MM:SS format.
-    
+
     Returns
     -------
     xr.Dataset
@@ -110,33 +117,103 @@ def clip_dataset_to_bounds(
     return dataset
 
 
-def save_to_cache(stores: xr.Dataset, cached_nc_path: Path) -> xr.Dataset:
-    """Compute the store and save it to a cached netCDF file. This is not required but will save time and bandwidth."""
-    logger.info("Downloading and caching forcing data, this may take a while")
+def interpolate_nan_values(
+    dataset: xr.Dataset,
+    variables: Optional[List[str]] = None,
+    dim: str = "time",
+    method: InterpOptions = "nearest",
+    fill_value: str = "extrapolate",
+) -> None:
+    """
+    Interpolates NaN values in specified (or all numeric time-dependent)
+    variables of an xarray.Dataset. Operates inplace on the dataset.
 
-    if not cached_nc_path.parent.exists():
-        cached_nc_path.parent.mkdir(parents=True)
+    Parameters
+    ----------
+    dataset : xr.Dataset
+        The input dataset.
+    variables : Optional[List[str]], optional
+        A list of variable names to process. If None (default),
+        all numeric variables containing the specified dimension will be processed.
+    dim : str, optional
+        The dimension along which to interpolate (default is "time").
+    method : str, optional
+        Interpolation method to use (e.g., "linear", "nearest", "cubic").
+        Default is "nearest".
+    fill_value : str, optional
+        Method for filling NaNs at the start/end of the series after interpolation.
+        Set to "extrapolate" to fill with the nearest valid value when using 'nearest' or 'linear'.
+        Default is "extrapolate".
+    """
+    for name, var in dataset.data_vars.items():
+        # if the variable is non-numeric, skip
+        if not np.issubdtype(var.dtype, np.number):
+            continue
+        # if there are no NANs, skip
+        if not var.isnull().any().compute():
+            continue
+
+        dataset[name] = var.interpolate_na(
+            dim=dim,
+            method=method,
+            fill_value=fill_value if method in ["nearest", "linear"] else None,
+        )
 
-    # sort of terrible work around for half downloaded files
-    temp_path = cached_nc_path.with_suffix(".downloading.nc")
-    if os.path.exists(temp_path):
-        os.remove(temp_path)
 
-    ## Cast every single variable to float32 to save space to save a lot of memory issues later
-    ## easier to do it now in this slow download step than later in the steps without dask
-    for var in stores.data_vars: 
-        stores[var] = stores[var].astype("float32")
+@use_cluster
+def save_dataset(
+    ds_to_save: xr.Dataset,
+    target_path: Path,
+    engine: Literal["netcdf4", "scipy", "h5netcdf"] = "h5netcdf",
+):
+    """
+    Helper function to compute and save an xarray.Dataset to a NetCDF file.
+    Uses a temporary file and rename for atomicity.
+    """
+    if not target_path.parent.exists():
+        target_path.parent.mkdir(parents=True, exist_ok=True)
+
+    temp_file_path = target_path.with_name(target_path.name + ".saving.nc")
+    if temp_file_path.exists():
+        os.remove(temp_file_path)
 
     client = Client.current()
-    future = client.compute(stores.to_netcdf(temp_path, compute=False))
-    # Display progress bar
+    future: Future = client.compute(
+        ds_to_save.to_netcdf(temp_file_path, engine=engine, compute=False)
+    )  # type: ignore
+    logger.debug(
+        f"NetCDF write task submitted to Dask. Waiting for completion to {temp_file_path}..."
+    )
     progress(future)
     future.result()
+    os.rename(str(temp_file_path), str(target_path))
+    logger.info(f"Successfully saved data to: {target_path}")
+
+
+@use_cluster
+def save_to_cache(
+    stores: xr.Dataset, cached_nc_path: Path, interpolate_nans: bool = True
+) -> xr.Dataset:
+    """
+    Compute the store and save it to a cached netCDF file. This is not required but will save time and bandwidth.
+    """
+    logger.info(f"Processing dataset for caching. Final cache target: {cached_nc_path}")
 
-    os.rename(temp_path, cached_nc_path)
+    # lasily cast all numbers to f32
+    for name, var in stores.data_vars.items():
+        if np.issubdtype(var.dtype, np.number):
+            stores[name] = var.astype("float32", casting="same_kind")
 
-    data = xr.open_mfdataset(cached_nc_path, parallel=True, engine="h5netcdf")
-    return data
+    # save dataset locally before manipulating it
+    save_dataset(stores, cached_nc_path)
+    stores = xr.open_mfdataset(cached_nc_path, parallel=True, engine="h5netcdf")
+
+    if interpolate_nans:
+        interpolate_nan_values(dataset=stores)
+        save_dataset(stores, cached_nc_path)
+        stores = xr.open_mfdataset(cached_nc_path, parallel=True, engine="h5netcdf")
+
+    return stores
 
 
 def check_local_cache(
@@ -144,9 +221,8 @@ def check_local_cache(
     start_time: str,
     end_time: str,
     gdf: gpd.GeoDataFrame,
-    remote_dataset: xr.Dataset
+    remote_dataset: xr.Dataset,
 ) -> Union[xr.Dataset, None]:
-
     merged_data = None
 
     if not os.path.exists(cached_nc_path):
@@ -155,9 +231,7 @@ def check_local_cache(
 
     logger.info("Found cached nc file")
     # open the cached file and check that the time range is correct
-    cached_data = xr.open_mfdataset(
-        cached_nc_path, parallel=True, engine="h5netcdf"
-    )
+    cached_data = xr.open_mfdataset(cached_nc_path, parallel=True, engine="h5netcdf")
 
     if "name" not in cached_data.attrs or "name" not in remote_dataset.attrs:
         logger.warning("No name attribute found to compare datasets")
@@ -166,9 +240,9 @@ def check_local_cache(
         logger.warning("Cached data from different source, .name attr doesn't match")
         return
 
-    range_in_cache = cached_data.time[0].values <= np.datetime64(
-        start_time
-    ) and cached_data.time[-1].values >= np.datetime64(end_time)
+    range_in_cache = cached_data.time[0].values <= np.datetime64(start_time) and cached_data.time[
+        -1
+    ].values >= np.datetime64(end_time)
 
     if not range_in_cache:
         # the cache does not contain the desired time range
@@ -186,10 +260,8 @@ def check_local_cache(
     if range_in_cache:
         logger.info("Time range is within cached data")
         logger.debug(f"Opened cached nc file: [{cached_nc_path}]")
-        merged_data = clip_dataset_to_bounds(
-            cached_data, gdf.total_bounds, start_time, end_time
-        )
-        logger.debug("Clipped stores")        
+        merged_data = clip_dataset_to_bounds(cached_data, gdf.total_bounds, start_time, end_time)
+        logger.debug("Clipped stores")
 
     return merged_data
 
@@ -197,16 +269,27 @@ def check_local_cache(
 def save_and_clip_dataset(
     dataset: xr.Dataset,
     gdf: gpd.GeoDataFrame,
-    start_time: datetime.datetime,
-    end_time: datetime.datetime,
+    start_time: datetime,
+    end_time: datetime,
     cache_location: Path,
 ) -> xr.Dataset:
     """convenience function clip the remote dataset, and either load from cache or save to cache if it's not present"""
     gdf = gdf.to_crs(dataset.crs)
 
-    cached_data = check_local_cache(cache_location, start_time, end_time, gdf, dataset)
+    cached_data = check_local_cache(
+        cache_location,
+        start_time,  # type: ignore
+        end_time,  # type: ignore
+        gdf,
+        dataset,
+    )
 
     if not cached_data:
-        clipped_data = clip_dataset_to_bounds(dataset, gdf.total_bounds, start_time, end_time)
+        clipped_data = clip_dataset_to_bounds(
+            dataset,
+            gdf.total_bounds,
+            start_time,  # type: ignore
+            end_time,  # type: ignore
+        )
         cached_data = save_to_cache(clipped_data, cache_location)
-    return cached_data
+    return cached_data