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

data_processing/forcings.py

@@ -7,14 +7,14 @@ from functools import partial
 from math import ceil
 from multiprocessing import shared_memory
 from pathlib import Path
-from typing import Tuple
+from typing import List, Tuple
 
 import geopandas as gpd
 import numpy as np
 import pandas as pd
 import psutil
 import xarray as xr
-from dask.distributed import Client, LocalCluster
+from data_processing.dask_utils import no_cluster, use_cluster
 from data_processing.dataset_utils import validate_dataset_format
 from data_processing.file_paths import file_paths
 from exactextract import exact_extract
@@ -92,21 +92,21 @@ def get_cell_weights(raster: xr.Dataset, gdf: gpd.GeoDataFrame, wkt: str) -> pd.
         DataFrame indexed by divide_id that contains information about coverage
         for each raster cell in gridded forcing file.
     """
-    xmin = raster.x[0]
-    xmax = raster.x[-1]
-    ymin = raster.y[0]
-    ymax = raster.y[-1]
+    xmin = min(raster.x)
+    xmax = max(raster.x)
+    ymin = min(raster.y)
+    ymax = max(raster.y)
     data_vars = list(raster.data_vars)
     rastersource = NumPyRasterSource(
         raster[data_vars[0]], srs_wkt=wkt, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax
     )
-    output = exact_extract(
+    output: pd.DataFrame = exact_extract(
         rastersource,
         gdf,
         ["cell_id", "coverage"],
         include_cols=["divide_id"],
         output="pandas",
-    )
+    )  # type: ignore
     return output.set_index("divide_id")
 
 
@@ -164,8 +164,8 @@ def get_index_chunks(data: xr.DataArray) -> list[tuple[int, int]]:
 
 
 def create_shared_memory(
-    lazy_array: xr.Dataset,
-) -> Tuple[shared_memory.SharedMemory, np.dtype, np.dtype]:
+    lazy_array: xr.DataArray,
+) -> Tuple[shared_memory.SharedMemory, Tuple[int, ...], np.dtype]:
     """
     Create a shared memory object so that multiple processes can access loaded
     data.
@@ -180,7 +180,7 @@ def create_shared_memory(
     shared_memory.SharedMemory
         A specific block of memory allocated by the OS of the size of
         lazy_array.
-    np.dtype.shape
+    Tuple[int, ...]
         A shape object with dimensions (# timesteps, # of raster cells) in
         reference to lazy_array.
     np.dtype
@@ -205,9 +205,9 @@ def process_chunk_shared(
     variable: str,
     times: np.ndarray,
     shm_name: str,
-    shape: np.dtype.shape,
+    shape: Tuple[int, ...],
     dtype: np.dtype,
-    chunk: gpd.GeoDataFrame,
+    chunk: pd.DataFrame,
 ) -> xr.DataArray:
     """
     Process the gridded forcings chunk loaded into a SharedMemory block.
@@ -276,7 +276,7 @@ def get_cell_weights_parallel(
         for each raster cell and each timestep in gridded forcing file.
     """
     gdf_chunks = np.array_split(gdf, num_partitions)
-    wkt = gdf.crs.to_wkt()
+    wkt = gdf.crs.to_wkt()  # type: ignore
     one_timestep = input_forcings.isel(time=0).compute()
     with multiprocessing.Pool() as pool:
         args = [(one_timestep, gdf_chunk, wkt) for gdf_chunk in gdf_chunks]
@@ -305,6 +305,7 @@ def get_units(dataset: xr.Dataset) -> dict:
     return units
 
 
+@no_cluster
 def compute_zonal_stats(
     gdf: gpd.GeoDataFrame, gridded_data: xr.Dataset, forcings_dir: Path
 ) -> None:
@@ -331,7 +332,7 @@ def compute_zonal_stats(
     catchments = get_cell_weights_parallel(gdf, gridded_data, num_partitions)
     units = get_units(gridded_data)
 
-    cat_chunks = np.array_split(catchments, num_partitions)
+    cat_chunks: List[pd.DataFrame] = np.array_split(catchments, num_partitions)  # type: ignore
 
     progress = Progress(
         TextColumn("[progress.description]{task.description}"),
@@ -350,25 +351,28 @@ def compute_zonal_stats(
         "[cyan]Processing variables...", total=len(gridded_data.data_vars), elapsed=0
     )
     progress.start()
-    for variable in list(gridded_data.data_vars):
+    for data_var_name in list(gridded_data.data_vars):
+        data_var_name: str
         progress.update(variable_task, advance=1)
-        progress.update(variable_task, description=f"Processing {variable}")
+        progress.update(variable_task, description=f"Processing {data_var_name}")
 
         # to make sure this fits in memory, we need to chunk the data
-        time_chunks = get_index_chunks(gridded_data[variable])
+        time_chunks = get_index_chunks(gridded_data[data_var_name])
         chunk_task = progress.add_task("[purple] processing chunks", total=len(time_chunks))
         for i, times in enumerate(time_chunks):
             progress.update(chunk_task, advance=1)
             start, end = times
             # select the chunk of time we want to process
-            data_chunk = gridded_data[variable].isel(time=slice(start, end))
+            data_chunk = gridded_data[data_var_name].isel(time=slice(start, end))
             # put it in shared memory
             shm, shape, dtype = create_shared_memory(data_chunk)
             times = data_chunk.time.values
             # create a partial function to pass to the multiprocessing pool
-            partial_process_chunk = partial(process_chunk_shared,variable,times,shm.name,shape,dtype)
+            partial_process_chunk = partial(
+                process_chunk_shared, data_var_name, times, shm.name, shape, dtype
+            )
 
-            logger.debug(f"Processing variable: {variable}")
+            logger.debug(f"Processing variable: {data_var_name}")
             # process the chunks of catchments in parallel
             with multiprocessing.Pool(num_partitions) as pool:
                 variable_data = pool.map(partial_process_chunk, cat_chunks)
@@ -376,24 +380,24 @@ def compute_zonal_stats(
             # clean up the shared memory
             shm.close()
             shm.unlink()
-            logger.debug(f"Processed variable: {variable}")
+            logger.debug(f"Processed variable: {data_var_name}")
             concatenated_da = xr.concat(variable_data, dim="catchment")
             # delete the data to free up memory
             del variable_data
-            logger.debug(f"Concatenated variable: {variable}")
+            logger.debug(f"Concatenated variable: {data_var_name}")
             # write this to disk now to save memory
             # xarray will monitor memory usage, but it doesn't account for the shared memory used to store the raster
             # This reduces memory usage by about 60%
-            concatenated_da.to_dataset(name=variable).to_netcdf(
-                forcings_dir / "temp" / f"{variable}_timechunk_{i}.nc"
+            concatenated_da.to_dataset(name=data_var_name).to_netcdf(
+                forcings_dir / "temp" / f"{data_var_name}_timechunk_{i}.nc"
             )
         # Merge the chunks back together
         datasets = [
-            xr.open_dataset(forcings_dir / "temp" / f"{variable}_timechunk_{i}.nc")
+            xr.open_dataset(forcings_dir / "temp" / f"{data_var_name}_timechunk_{i}.nc")
             for i in range(len(time_chunks))
         ]
         result = xr.concat(datasets, dim="time")
-        result.to_netcdf(forcings_dir / "temp" / f"{variable}.nc")
+        result.to_netcdf(forcings_dir / "temp" / f"{data_var_name}.nc")
         # close the datasets
         result.close()
         _ = [dataset.close() for dataset in datasets]
@@ -411,6 +415,7 @@ def compute_zonal_stats(
     write_outputs(forcings_dir, units)
 
 
+@use_cluster
 def write_outputs(forcings_dir: Path, units: dict) -> None:
     """
     Write outputs to disk in the form of a NetCDF file, using dask clusters to
@@ -428,13 +433,6 @@ def write_outputs(forcings_dir: Path, units: dict) -> None:
         units. Differs from variables, as this dictionary depends on the gridded
         forcing dataset.
     """
-
-    # start a dask cluster if there isn't one already running
-    try:
-        client = Client.current()
-    except ValueError:
-        cluster = LocalCluster()
-        client = Client(cluster)
     temp_forcings_dir = forcings_dir / "temp"
     # Combine all variables into a single dataset using dask
     results = [xr.open_dataset(file, chunks="auto") for file in temp_forcings_dir.glob("*.nc")]
@@ -471,14 +469,18 @@ def write_outputs(forcings_dir: Path, units: dict) -> None:
         time_array = (
             final_ds.time.astype("datetime64[s]").astype(np.int64).values // 10**9
         )  ## convert from ns to s
-    time_array = time_array.astype(np.int32) ## convert to int32 to save space
-    final_ds = final_ds.drop_vars(["catchment", "time"]) ## drop the original time and catchment vars
-    final_ds = final_ds.rename_dims({"catchment": "catchment-id"}) # rename the catchment dimension
+    time_array = time_array.astype(np.int32)  ## convert to int32 to save space
+    final_ds = final_ds.drop_vars(
+        ["catchment", "time"]
+    )  ## drop the original time and catchment vars
+    final_ds = final_ds.rename_dims({"catchment": "catchment-id"})  # rename the catchment dimension
     # add the time as a 2d data var, yes this is wasting disk space.
     final_ds["Time"] = (("catchment-id", "time"), [time_array for _ in range(len(final_ds["ids"]))])
     # set the time unit
     final_ds["Time"].attrs["units"] = "s"
-    final_ds["Time"].attrs["epoch_start"] = "01/01/1970 00:00:00" # not needed but suppresses the ngen warning
+    final_ds["Time"].attrs["epoch_start"] = (
+        "01/01/1970 00:00:00"  # not needed but suppresses the ngen warning
+    )
 
     final_ds.to_netcdf(forcings_dir / "forcings.nc", engine="netcdf4")
     # close the datasets

data_processing/gpkg_utils.py

@@ -2,11 +2,12 @@ import logging
 import sqlite3
 import struct
 from pathlib import Path
-from typing import List, Tuple, Union
+from typing import List, Tuple, Dict
 
 import pyproj
 from data_processing.file_paths import file_paths
-from shapely.geometry import Point, Polygon
+from shapely.geometry import Point
+from shapely.geometry.base import BaseGeometry
 from shapely.ops import transform
 from shapely.wkb import loads
 
@@ -27,7 +28,7 @@ class GeoPackage:
         self.conn.close()
 
 
-def verify_indices(gpkg: str = file_paths.conus_hydrofabric) -> None:
+def verify_indices(gpkg: Path = file_paths.conus_hydrofabric) -> None:
     """
     Verify that the indices in the specified geopackage are correct.
     If they are not, create the correct indices.
@@ -92,12 +93,9 @@ def add_triggers_to_gpkg(gpkg: Path) -> None:
     logger.debug(f"Added triggers to subset gpkg {gpkg}")
 
 
-# whenever this is imported, check if the indices are correct
-if file_paths.conus_hydrofabric.is_file():
-    verify_indices()
 
 
-def blob_to_geometry(blob: bytes) -> Union[Point, Polygon]:
+def blob_to_geometry(blob: bytes) -> BaseGeometry | None:
     """
     Convert a blob to a geometry.
     from http://www.geopackage.org/spec/#gpb_format
@@ -120,7 +118,7 @@ def blob_to_geometry(blob: bytes) -> Union[Point, Polygon]:
     return geometry
 
 
-def blob_to_centre_point(blob: bytes) -> Point:
+def blob_to_centre_point(blob: bytes) -> Point | None:
     """
     Convert a blob to a geometry.
     from http://www.geopackage.org/spec/#gpb_format
@@ -151,7 +149,7 @@ def blob_to_centre_point(blob: bytes) -> Point:
     return Point(x, y)
 
 
-def convert_to_5070(shapely_geometry):
+def convert_to_5070(shapely_geometry: Point) -> Point:
     # convert to web mercator
     if shapely_geometry.is_empty:
         return shapely_geometry
@@ -164,7 +162,7 @@ def convert_to_5070(shapely_geometry):
     return new_geometry
 
 
-def get_catid_from_point(coords):
+def get_catid_from_point(coords: Dict[str, float]) -> str:
     """
     Retrieves the watershed boundary ID (catid) of the watershed that contains the given point.
 
@@ -196,6 +194,8 @@ def get_catid_from_point(coords):
         # check the geometries to see which one contains the point
         for result in results:
             geom = blob_to_geometry(result[1])
+            if geom is None:
+                continue
             if geom.contains(point):
                 return result[0]
     return results[0][0]
@@ -398,7 +398,7 @@ def subset_table(table: str, ids: List[str], hydrofabric: Path, subset_gpkg_name
     dest_db.close()
 
 
-def get_table_crs_short(gpkg: str, table: str) -> str:
+def get_table_crs_short(gpkg: str | Path, table: str) -> str:
     """
     Gets the CRS of the specified table in the specified geopackage as a short string. e.g. EPSG:5070
 
@@ -518,7 +518,7 @@ def get_available_tables(gpkg: Path) -> List[str]:
     return tables
 
 
-def get_cat_to_nhd_feature_id(gpkg: Path = file_paths.conus_hydrofabric) -> dict:
+def get_cat_to_nhd_feature_id(gpkg: Path = file_paths.conus_hydrofabric) -> Dict[str, int]:
     available_tables = get_available_tables(gpkg)
     possible_tables = ["flowpath_edge_list", "network"]
 
@@ -535,7 +535,7 @@ def get_cat_to_nhd_feature_id(gpkg: Path = file_paths.conus_hydrofabric) -> dict
     sql_query = f"SELECT divide_id, hf_id FROM {table_name} WHERE divide_id IS NOT NULL AND hf_id IS NOT NULL"
 
     with sqlite3.connect(gpkg) as conn:
-        result = conn.execute(sql_query).fetchall()
+        result: List[Tuple[str, str]] = conn.execute(sql_query).fetchall()
 
     mapping = {}
     for cat, feature in result:

data_processing/graph_utils.py

@@ -2,7 +2,7 @@ import logging
 import sqlite3
 from functools import cache
 from pathlib import Path
-from typing import List, Set, Union
+from typing import List, Optional, Set, Union
 
 import igraph as ig
 from data_processing.file_paths import file_paths
@@ -11,7 +11,7 @@ logger = logging.getLogger(__name__)
 
 
 def get_from_to_id_pairs(
-    hydrofabric: Path = file_paths.conus_hydrofabric, ids: Set = None
+    hydrofabric: Path = file_paths.conus_hydrofabric, ids: Optional[Set | List] = None
 ) -> List[tuple]:
     """
     Retrieves the from and to IDs from the specified hydrofabric.
@@ -112,7 +112,7 @@ def get_graph() -> ig.Graph:
     return network_graph
 
 
-def get_outlet_id(wb_or_cat_id: str) -> str:
+def get_outlet_id(wb_or_cat_id: str) -> str | None:
     """
     Retrieves the ID of the node downstream of the given node in the hydrological network.
 
@@ -209,7 +209,7 @@ def get_upstream_ids(names: Union[str, List[str]], include_outlet: bool = True)
         if name in parent_ids:
             continue
         try:
-            if "cat" in name:
+            if "cat" in name:  # type: ignore # If name is None, this will raise an error, which is handled below
                 node_index = graph.vs.find(cat=name).index
             else:
                 node_index = graph.vs.find(name=name).index

data_processing/s3fs_utils.py

@@ -32,7 +32,7 @@ class S3ParallelFileSystem(S3FileSystem):
                     "head_object", Bucket=bucket, Key=key, **version_kw, **self.req_kw
                 )
             )["ContentLength"]
-        except Exception as e:
+        except Exception:
             # Fall back to single request if HEAD fails
             return await self._download_chunk(bucket, key, {}, version_kw)
 

data_processing/subset.py

@@ -64,7 +64,7 @@ def subset_vpu(
 
 
 def subset(
-    cat_ids: List[str],
+    cat_ids: str | List[str],
     hydrofabric: Path = file_paths.conus_hydrofabric,
     output_gpkg_path: Path = Path(),
     include_outlet: bool = True,

data_sources/source_validation.py

@@ -1,19 +1,29 @@
 import gzip
+import json
 import os
 import tarfile
 import warnings
-import json
-import requests
-from data_processing.file_paths import file_paths
-from tqdm import TqdmExperimentalWarning
 from time import sleep
+
 import boto3
-from botocore.exceptions import ClientError
+import psutil
+import requests
 from boto3.s3.transfer import TransferConfig
+from botocore.exceptions import ClientError
+import botocore
+from data_processing.file_paths import file_paths
 from rich.console import Console
+from rich.progress import (Progress, 
+                           SpinnerColumn, 
+                           TextColumn, 
+                           TimeElapsedColumn, 
+                           BarColumn, 
+                           DownloadColumn, 
+                           TransferSpeedColumn)
 from rich.prompt import Prompt
-from rich.progress import Progress, TextColumn, TimeElapsedColumn, SpinnerColumn
-import psutil
+from tqdm import TqdmExperimentalWarning
+from data_processing.gpkg_utils import verify_indices
+import sqlite3
 
 warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
 
@@ -23,25 +33,22 @@ S3_KEY = "hydrofabrics/community/conus_nextgen.tar.gz"
 S3_REGION = "us-east-1"
 hydrofabric_url = f"https://{S3_BUCKET}.s3.{S3_REGION}.amazonaws.com/{S3_KEY}"
 
+
 def decompress_gzip_tar(file_path, output_dir):
-    # use rich to display "decompressing" message with a progress bar that just counts down from 30s
-    # actually measuring this is hard and it usually takes ~20s to decompress
     console.print("Decompressing Hydrofabric...", style="bold green")
     progress = Progress(
         SpinnerColumn(),
         TextColumn("[progress.description]{task.description}"),
-        TimeElapsedColumn(),        
+        TimeElapsedColumn(),
     )
     task = progress.add_task("Decompressing", total=1)
-    progress.start()
-    with gzip.open(file_path, "rb") as f_in:
-        with tarfile.open(fileobj=f_in) as tar:
-            # Extract all contents
-            for member in tar:
-                tar.extract(member, path=output_dir)
-                # Update the progress bar
-    progress.update(task, completed=1)
-    progress.stop()
+    with progress:
+        with gzip.open(file_path, "rb") as f_in:
+            with tarfile.open(fileobj=f_in) as tar:
+                # Extract all contents
+                for member in tar:
+                    tar.extract(member, path=output_dir)
+                    progress.update(task, advance=1 / len(tar.getmembers()))
 
 
 def download_from_s3(save_path, bucket=S3_BUCKET, key=S3_KEY, region=S3_REGION):
@@ -53,10 +60,13 @@ def download_from_s3(save_path, bucket=S3_BUCKET, key=S3_KEY, region=S3_REGION):
     if os.path.exists(save_path):
         console.print(f"File already exists: {save_path}", style="bold yellow")
         os.remove(save_path)
-
+    
+    client_config = botocore.config.Config(
+        max_pool_connections=75
+    )
     # Initialize S3 client
     s3_client = boto3.client(
-        "s3", aws_access_key_id="", aws_secret_access_key="", region_name=region
+        "s3", aws_access_key_id="", aws_secret_access_key="", region_name=region, config=client_config
     )
     # Disable request signing for public buckets
     s3_client._request_signer.sign = lambda *args, **kwargs: None
@@ -92,19 +102,15 @@ def download_from_s3(save_path, bucket=S3_BUCKET, key=S3_KEY, region=S3_REGION):
         use_threads=True,
     )
 
-    console.print(f"Downloading {key} to {save_path}...", style="bold green")
-    console.print(
-        f"The file downloads faster with no progress indicator, this should take around 30s",
-        style="bold yellow",
-    )
-    console.print(
-        f"Please use network monitoring on your computer if you wish to track the download",
-        style="green",
-    )
 
     try:
+        dl_progress = Progress(BarColumn(), DownloadColumn(), TransferSpeedColumn())
         # Download file using optimized transfer config
-        s3_client.download_file(Bucket=bucket, Key=key, Filename=save_path, Config=config)
+        with dl_progress:
+            task = dl_progress.add_task("Downloading...", total=total_size)
+            s3_client.download_file(Bucket=bucket, Key=key, Filename=save_path, Config=config,
+                                    Callback=lambda bytes_downloaded: dl_progress.update(
+                                        task, advance=bytes_downloaded))
         return True
     except Exception as e:
         console.print(f"Error downloading file: {e}", style="bold red")
@@ -122,6 +128,14 @@ def get_headers():
 
 
 def download_and_update_hf():
+
+    if file_paths.conus_hydrofabric.is_file():
+        console.print(
+            f"Hydrofabric already exists at {file_paths.conus_hydrofabric}, removing it to download the latest version.",
+            style="bold yellow",
+        )
+        file_paths.conus_hydrofabric.unlink()
+        
     download_from_s3(
         file_paths.conus_hydrofabric.with_suffix(".tar.gz"),
         bucket="communityhydrofabric",
@@ -207,6 +221,17 @@ def validate_hydrofabric():
             )
             sleep(2)
             return
+    
+    # moved this from gpkg_utils to here to avoid potential nested rich live displays
+    if file_paths.conus_hydrofabric.is_file():
+        valid_hf = False
+        while not valid_hf:
+            try:
+                verify_indices()
+                valid_hf = True
+            except sqlite3.DatabaseError:
+                console.print(f"Hydrofabric {file_paths.conus_hydrofabric} is corrupted. Redownloading...", style="red")
+                download_and_update_hf()
 
 
 def validate_output_dir():
@@ -220,7 +245,7 @@ def validate_output_dir():
             response = Prompt.ask("Enter the path to the working directory")
         if response == "" or response.lower() == "y":
             response = "~/ngiab_preprocess_output/"
-        file_paths.set_working_dir(response)
+        file_paths.set_working_dir(response)  # type: ignore
 
 
 def validate_all():

map_app/__main__.py

@@ -39,12 +39,13 @@ def main():
         Timer(2, set_logs_to_warning).start()
         with open("app.log", "a") as f:
             f.write("Running in debug mode\n")
-        app.run(debug=True, host="0.0.0.0", port="8080")
+        app.run(debug=True, host="0.0.0.0", port="8080")  # type: ignore
     else:
         Timer(1, open_browser).start()
         with open("app.log", "a") as f:
             f.write("Running in production mode\n")
-        app.run(host="0.0.0.0", port="0")
+        app.run(host="0.0.0.0", port="0")  # type: ignore
+
 
 if __name__ == "__main__":
     main()

map_app/static/css/main.css

@@ -28,10 +28,9 @@ main {
 
 .maplibregl-popup-content {
   background: var(--surface-color) !important;
-
 }
 
-#toggle-button {
+#toggle-button-gages, #toggle-button-camels, #toggle-button-nwm, #toggle-button-aorc  {
   position: relative;
   top: 20px;
   left: 20px;
@@ -46,11 +45,23 @@ main {
   z-index: 1;
 }
 
-#toggle-button:hover {
+#toggle-button-gages:hover, #toggle-button-camels:hover, #toggle-button-nwm:hover, #toggle-button-aorc:hover {
   scale: 1.1;
   box-shadow: var(--shadow-md);
 }
 
+#toggle-button-camels {
+  left: 30px;
+}
+
+#toggle-button-nwm {
+  left: 40px;
+}
+
+#toggle-button-aorc {
+  left: 50px;
+}
+
 body {
   font-family: 'Inter', system-ui, -apple-system, sans-serif;
   margin: 0;