ngiab-data-preprocess 4.2.1__tar.gz → 4.3.0__tar.gz

{ngiab_data_preprocess-4.2.1 → ngiab_data_preprocess-4.3.0}/.gitignore

@@ -14,3 +14,5 @@ dist
 *.dat
 uv.lock
 /build
+*.pmtiles
+*style.json

{ngiab_data_preprocess-4.2.1 → ngiab_data_preprocess-4.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ngiab_data_preprocess
-Version: 4.2.1
+Version: 4.3.0
 Summary: Graphical Tools for creating Next Gen Water model input data.
 Author-email: Josh Cunningham <jcunningham8@ua.edu>
 Project-URL: Homepage, https://github.com/CIROH-UA/NGIAB_data_preprocess
@@ -15,7 +15,7 @@ Requires-Dist: pyogrio>=0.7.2
 Requires-Dist: pyproj>=3.6.1
 Requires-Dist: Flask==3.0.2
 Requires-Dist: geopandas>=1.0.0
-Requires-Dist: requests==2.32.2
+Requires-Dist: requests==2.32.4
 Requires-Dist: igraph==0.11.4
 Requires-Dist: s3fs==2024.3.1
 Requires-Dist: xarray==2024.2.0
@@ -23,8 +23,6 @@ Requires-Dist: zarr==2.17.1
 Requires-Dist: netCDF4>=1.6.5
 Requires-Dist: dask==2024.4.1
 Requires-Dist: dask[distributed]==2024.4.1
-Requires-Dist: black==24.3.0
-Requires-Dist: isort==5.13.2
 Requires-Dist: h5netcdf==1.3.0
 Requires-Dist: exactextract==0.2.0
 Requires-Dist: numpy>=1.26.4
@@ -49,15 +47,19 @@ This repository contains tools for preparing data to run a [next gen](https://gi
 ## Table of Contents
 
 1. [What does this tool do?](#what-does-this-tool-do)
-2. [Requirements](#requirements)
-3. [Installation and Running](#installation-and-running)
-4. [Development Installation](#development-installation)
-5. [Usage](#usage)
-6. [CLI Documentation](#cli-documentation)
+2. [What does it not do?](#what-does-it-not-do)
+   - [Evaluation](#evaluation)
+   - [Visualisation](#visualisation)
+3. [Requirements](#requirements)
+4. [Installation and Running](#installation-and-running)
+   - [Running without install](#running-without-install)
+5. [For legacy pip installation](#for-legacy-pip-installation)
+6. [Development Installation](#development-installation)
+7. [Usage](#usage)
+8. [CLI Documentation](#cli-documentation)
    - [Arguments](#arguments)
+   - [Usage Notes](#usage-notes)
    - [Examples](#examples)
-   - [File Formats](#file-formats)
-   - [Output](#output)
 
 ## What does this tool do?
 
@@ -231,12 +233,12 @@ Once all the steps are finished, you can run NGIAB on the folder shown underneat
 
 3. Create realization using a lat/lon pair and output to a named folder:
    ```bash
-   python -m ngiab_data_cli -i 54.33,-69.4 -l -r --start 2022-01-01 --end 2022-02-28 -o custom_output
+   python -m ngiab_data_cli -i 33.22,-87.54 -l -r --start 2022-01-01 --end 2022-02-28 -o custom_output
    ```
 
 4. Perform all operations using a lat/lon pair:
    ```bash
-   python -m ngiab_data_cli -i 54.33,-69.4 -l -s -f -r --start 2022-01-01 --end 2022-02-28
+   python -m ngiab_data_cli -i 33.22,-87.54 -l -s -f -r --start 2022-01-01 --end 2022-02-28
    ```
 
 5. Subset hydrofabric using gage ID:

{ngiab_data_preprocess-4.2.1 → ngiab_data_preprocess-4.3.0}/README.md

@@ -7,15 +7,19 @@ This repository contains tools for preparing data to run a [next gen](https://gi
 ## Table of Contents
 
 1. [What does this tool do?](#what-does-this-tool-do)
-2. [Requirements](#requirements)
-3. [Installation and Running](#installation-and-running)
-4. [Development Installation](#development-installation)
-5. [Usage](#usage)
-6. [CLI Documentation](#cli-documentation)
+2. [What does it not do?](#what-does-it-not-do)
+   - [Evaluation](#evaluation)
+   - [Visualisation](#visualisation)
+3. [Requirements](#requirements)
+4. [Installation and Running](#installation-and-running)
+   - [Running without install](#running-without-install)
+5. [For legacy pip installation](#for-legacy-pip-installation)
+6. [Development Installation](#development-installation)
+7. [Usage](#usage)
+8. [CLI Documentation](#cli-documentation)
    - [Arguments](#arguments)
+   - [Usage Notes](#usage-notes)
    - [Examples](#examples)
-   - [File Formats](#file-formats)
-   - [Output](#output)
 
 ## What does this tool do?
 
@@ -189,12 +193,12 @@ Once all the steps are finished, you can run NGIAB on the folder shown underneat
 
 3. Create realization using a lat/lon pair and output to a named folder:
    ```bash
-   python -m ngiab_data_cli -i 54.33,-69.4 -l -r --start 2022-01-01 --end 2022-02-28 -o custom_output
+   python -m ngiab_data_cli -i 33.22,-87.54 -l -r --start 2022-01-01 --end 2022-02-28 -o custom_output
    ```
 
 4. Perform all operations using a lat/lon pair:
    ```bash
-   python -m ngiab_data_cli -i 54.33,-69.4 -l -s -f -r --start 2022-01-01 --end 2022-02-28
+   python -m ngiab_data_cli -i 33.22,-87.54 -l -s -f -r --start 2022-01-01 --end 2022-02-28
    ```
 
 5. Subset hydrofabric using gage ID:

{ngiab_data_preprocess-4.2.1 → ngiab_data_preprocess-4.3.0}/modules/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)

ngiab_data_preprocess-4.3.0/modules/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