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

map_app/static/js/main.js

@@ -1,26 +1,38 @@
 var colorDict = {
-    selectedCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--selected-cat-outline'),
-    selectedCatFill: getComputedStyle(document.documentElement).getPropertyValue('--selected-cat-fill'),
-    upstreamCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--upstream-cat-outline'),
-    upstreamCatFill: getComputedStyle(document.documentElement).getPropertyValue('--upstream-cat-fill'),
-    flowlineToCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--flowline-to-cat-outline'),
-    flowlineToNexusOutline: getComputedStyle(document.documentElement).getPropertyValue('--flowline-to-nexus-outline'),
-    nexusOutline: getComputedStyle(document.documentElement).getPropertyValue('--nexus-outline'),
-    nexusFill: getComputedStyle(document.documentElement).getPropertyValue('--nexus-fill'),
-    clearFill: getComputedStyle(document.documentElement).getPropertyValue('--clear-fill')
+  selectedCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--selected-cat-outline'),
+  selectedCatFill: getComputedStyle(document.documentElement).getPropertyValue('--selected-cat-fill'),
+  upstreamCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--upstream-cat-outline'),
+  upstreamCatFill: getComputedStyle(document.documentElement).getPropertyValue('--upstream-cat-fill'),
+  flowlineToCatOutline: getComputedStyle(document.documentElement).getPropertyValue('--flowline-to-cat-outline'),
+  flowlineToNexusOutline: getComputedStyle(document.documentElement).getPropertyValue('--flowline-to-nexus-outline'),
+  nexusOutline: getComputedStyle(document.documentElement).getPropertyValue('--nexus-outline'),
+  nexusFill: getComputedStyle(document.documentElement).getPropertyValue('--nexus-fill'),
+  clearFill: getComputedStyle(document.documentElement).getPropertyValue('--clear-fill')
 };
 
 // A function that creates a cli command from the interface
 function create_cli_command() {
-    var selected_basins = $('#selected-basins').text();
-    var start_time = document.getElementById('start-time').value.split('T')[0];
-    var end_time = document.getElementById('end-time').value.split('T')[0];
-    var command = `python -m ngiab_data_cli -i ${selected_basins} --subset --start ${start_time} --end ${end_time} --forcings --realization --run`;
-    var command_all = `python -m ngiab_data_cli -i ${selected_basins} --start ${start_time} --end ${end_time} --all`;
-    if (selected_basins != "None - get clicking!") {
-        $('#cli-command').text(command);
-    }
+  const cliPrefix = document.getElementById("cli-prefix");
+  cliPrefix.style.opacity = 1;
+  var selected_basins = $("#selected-basins").text();
+  var start_time = document.getElementById("start-time").value.split("T")[0];
+  var end_time = document.getElementById("end-time").value.split("T")[0];
+  var command = `-i ${selected_basins} --subset --start ${start_time} --end ${end_time} --forcings --realization --run`;
+  var command_all = `-i ${selected_basins} --start ${start_time} --end ${end_time} --all`;
+  if (selected_basins != "None - get clicking!") {
+    $("#cli-command").text(command);
+  }
+}
+
+function updateCommandPrefix() {
+  const toggleInput = document.getElementById("runcmd-toggle");
+  const cliPrefix = document.getElementById("cli-prefix");
+  const uvxText = "uvx --from ngiab_data_preprocess cli";
+  const pythonText = "python -m ngiab_data_cli";
+  // Set initial handle text based on the default state using data attribute
+  cliPrefix.textContent = toggleInput.checked ? pythonText : uvxText;
 }
+document.getElementById("runcmd-toggle").addEventListener('change', updateCommandPrefix);
 
 // These functions are exported by data_processing.js
 document.getElementById('map').addEventListener('click', create_cli_command);
@@ -29,110 +41,240 @@ document.getElementById('end-time').addEventListener('change', create_cli_comman
 
 
 // add the PMTiles plugin to the maplibregl global.
-let protocol = new pmtiles.Protocol({metadata: true});
+let protocol = new pmtiles.Protocol({ metadata: true });
 maplibregl.addProtocol("pmtiles", protocol.tile);
 
 // select light-style if the browser is in light mode
 // select dark-style if the browser is in dark mode
-var style = 'static/resources/light-style.json';
+var style = 'https://communityhydrofabric.s3.us-east-1.amazonaws.com/map/styles/light-style.json';
+var colorScheme = "light";
 if (window.matchMedia && window.matchMedia('(prefers-color-scheme: dark)').matches) {
-    style = 'static/resources/dark-style.json';
+  style = 'https://communityhydrofabric.s3.us-east-1.amazonaws.com/map/styles/dark-style.json';
+  colorScheme = "dark";
 }
 var map = new maplibregl.Map({
-    container: 'map', // container id
-    style: style, // style URL
-    center: [-96, 40], // starting position [lng, lat]
-    zoom: 4 // starting zoom
+  container: "map", // container id
+  style: style, // style URL
+  center: [-96, 40], // starting position [lng, lat]
+  zoom: 4, // starting zoom
+});
+
+map.on("load", () => {
+  map.addSource("camels_basins", {
+    type: "vector",
+    url: "pmtiles://https://communityhydrofabric.s3.us-east-1.amazonaws.com/map/camels.pmtiles",
+  });
+  map.addLayer({
+    id: "camels",
+    type: "line",
+    source: "camels_basins",
+    "source-layer": "camels_basins",
+    layout: {},
+    filter: ["any", ["==", "hru_id", ""]],
+    paint: {
+      "line-width": 1.5,
+      "line-color": ["rgba", 134, 30, 232, 1],
+    },
+  });
+});
+
+if (colorScheme == "light") {
+  nwm_paint = {
+    "line-width": 1,
+    "line-color": ["rgba", 0, 0, 0, 1],
+  };
+  aorc_paint = {
+    "line-width": 1,
+    "line-color": ["rgba", 71, 58, 222, 1],
+  };
+}
+if (colorScheme == "dark") {    
+  nwm_paint = {
+    "line-width": 1,
+    "line-color": ["rgba", 255, 255, 255, 1],
+  };
+  aorc_paint = {
+    "line-width": 1,
+    "line-color": ["rgba", 242, 252, 126, 1],
+  };
+}
+
+
+map.on("load", () => {
+  map.addSource("nwm_zarr_chunks", {
+    type: "vector",
+    url: "pmtiles://https://communityhydrofabric.s3.us-east-1.amazonaws.com/map/forcing_chunks/nwm_retro_v3_zarr_chunks.pmtiles",
+  });
+  map.addSource("aorc_zarr_chunks", {
+    type: "vector",
+    url: "pmtiles://https://communityhydrofabric.s3.us-east-1.amazonaws.com/map/forcing_chunks/aorc_zarr_chunks.pmtiles",
+  });
+  map.addLayer({
+    id: "nwm_zarr_chunks",
+    type: "line",
+    source: "nwm_zarr_chunks",
+    "source-layer": "nwm_zarr_chunks",
+    layout: {},
+    filter: ["any"],
+    paint: nwm_paint,
+  });
+  map.addLayer({
+    id: "aorc_zarr_chunks",
+    type: "line",
+    source: "aorc_zarr_chunks",
+    "source-layer": "aorc_zarr_chunks",
+    layout: {},
+    filter: ["any"],
+    paint: aorc_paint,
+  });
 });
+
 function update_map(cat_id, e) {
-    $('#selected-basins').text(cat_id)
-    map.setFilter('selected-catchments', ['any', ['in', 'divide_id', cat_id]]);
-    map.setFilter('upstream-catchments', ['any', ['in', 'divide_id', ""]])
-    
-    fetch('/get_upstream_catids', {
-        method: 'POST',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify(cat_id),
-    })
-    .then(response => response.json())
-    .then(data => {
-        map.setFilter('upstream-catchments', ['any', ['in', 'divide_id', ...data]]);
-        if (data.length === 0) { 
-            new maplibregl.Popup()
-            .setLngLat(e.lngLat)
-            .setHTML('No upstreams')
-            .addTo(map);
-        }
-    });
+  $('#selected-basins').text(cat_id)
+  map.setFilter('selected-catchments', ['any', ['in', 'divide_id', cat_id]]);
+  map.setFilter('upstream-catchments', ['any', ['in', 'divide_id', ""]])
+
+  fetch('/get_upstream_catids', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify(cat_id),
+  })
+  .then(response => response.json())
+  .then(data => {
+    map.setFilter('upstream-catchments', ['any', ['in', 'divide_id', ...data]]);
+    if (data.length === 0) {
+      new maplibregl.Popup()
+        .setLngLat(e.lngLat)
+        .setHTML('No upstreams')
+        .addTo(map);
+    }
+  });
 }
 map.on('click', 'catchments', (e) => {
-    cat_id = e.features[0].properties.divide_id;    
-    update_map(cat_id, e);
+  cat_id = e.features[0].properties.divide_id;
+  update_map(cat_id, e);
 });
 
 // Create a popup, but don't add it to the map yet.
 const popup = new maplibregl.Popup({
-    closeButton: false,
-    closeOnClick: false
+  closeButton: false,
+  closeOnClick: false
 });
 
 map.on('mouseenter', 'conus_gages', (e) => {
-    // Change the cursor style as a UI indicator.
-    map.getCanvas().style.cursor = 'pointer';
+  // Change the cursor style as a UI indicator.
+  map.getCanvas().style.cursor = 'pointer';
 
-    const coordinates = e.features[0].geometry.coordinates.slice();
-    const description = e.features[0].properties.hl_uri + "<br> click for more info";
+  const coordinates = e.features[0].geometry.coordinates.slice();
+  const description = e.features[0].properties.hl_uri + "<br> click for more info";
 
-    // Ensure that if the map is zoomed out such that multiple
-    // copies of the feature are visible, the popup appears
-    // over the copy being pointed to.
-    while (Math.abs(e.lngLat.lng - coordinates[0]) > 180) {
-        coordinates[0] += e.lngLat.lng > coordinates[0] ? 360 : -360;
-    }
+  // Ensure that if the map is zoomed out such that multiple
+  // copies of the feature are visible, the popup appears
+  // over the copy being pointed to.
+  while (Math.abs(e.lngLat.lng - coordinates[0]) > 180) {
+    coordinates[0] += e.lngLat.lng > coordinates[0] ? 360 : -360;
+  }
 
-    // Populate the popup and set its coordinates
-    // based on the feature found.
-    popup.setLngLat(coordinates).setHTML(description).addTo(map);
+  // Populate the popup and set its coordinates
+  // based on the feature found.
+  popup.setLngLat(coordinates).setHTML(description).addTo(map);
 });
 
-map.on('mouseleave', 'conus_gages', () => {
-    map.getCanvas().style.cursor = '';
-    popup.remove();
+map.on("mouseleave", "conus_gages", () => {
+  map.getCanvas().style.cursor = "";
+  popup.remove();
 });
 
-map.on('click', 'conus_gages', (e) => {
-    //  https://waterdata.usgs.gov/monitoring-location/02465000 
-    window.open("https://waterdata.usgs.gov/monitoring-location/" + e.features[0].properties.hl_link, '_blank');
-}
-);
+map.on("click", "conus_gages", (e) => {
+  //  https://waterdata.usgs.gov/monitoring-location/02465000
+  window.open(
+    "https://waterdata.usgs.gov/monitoring-location/" +
+    e.features[0].properties.hl_link,
+    "_blank",
+  );
+});
 show = false;
 
 // TOGGLE BUTTON LOGIC
-const toggleInput = document.querySelector('.toggle-input');
-const toggleHandle = document.querySelector('.toggle-handle');
+function initializeToggleSwitches() {
+  // Find all toggle switches
+  const toggleSwitches = document.querySelectorAll(".toggle-switch");
+  // Process each toggle switch
+  toggleSwitches.forEach((toggleSwitch) => {
+    const toggleInput = toggleSwitch.querySelector(".toggle-input");
+    const toggleHandle = toggleSwitch.querySelector(".toggle-handle");
+    const leftText =
+      toggleSwitch.querySelector(".toggle-text-left").textContent;
+    const rightText =
+      toggleSwitch.querySelector(".toggle-text-right").textContent;
+    // Set initial handle text based on the default state using data attribute
+    toggleHandle.textContent = toggleInput.checked ? rightText : leftText;
+    // Add event listener
+    toggleInput.addEventListener("change", function () {
+      setTimeout(() => {
+        if (this.checked) {
+          toggleHandle.textContent = rightText;
+        } else {
+          toggleHandle.textContent = leftText;
+        }
+      }, 180);
+    });
+  });
+}
+document.addEventListener("DOMContentLoaded", initializeToggleSwitches);
 
-// Set initial text based on the default state
-toggleHandle.textContent = toggleInput.checked ? 'AORC' : 'NWM';
+showGages = false;
+const toggleButtonGages = document.querySelector("#toggle-button-gages");
+toggleButtonGages.addEventListener("click", () => {
+  if (showGages) {
+    map.setFilter("conus_gages", ["any", ["==", "hl_uri", ""]]);
+    toggleButtonGages.innerText = "Show gages";
+    showGages = false;
+  } else {
+    map.setFilter("conus_gages", null);
+    toggleButtonGages.innerText = "Hide gages";
+    showGages = true;
+  }
+});
 
-toggleInput.addEventListener('change', function() {
-    if (this.checked) {
-        toggleHandle.textContent = 'AORC'; // Update handle text to AORC
-    } else {
-        toggleHandle.textContent = 'NWM'; // Update handle text to NWM
-    }
+showCamels = false;
+const toggleButtonCamels = document.querySelector("#toggle-button-camels");
+toggleButtonCamels.addEventListener("click", () => {
+  if (showCamels) {
+    map.setFilter("camels", ["any", ["==", "hru_id", ""]]);
+    toggleButtonCamels.innerText = "Show CAMELS basins";
+    showCamels = false;
+  } else {
+    map.setFilter("camels", null);
+    toggleButtonCamels.innerText = "Hide CAMELS basins";
+    showCamels = true;
+  }
 });
 
-show = false;
-const toggleButton = document.querySelector('#toggle-button');
-toggleButton.addEventListener('click', () => {
-    if (show) {
-        map.setFilter('conus_gages', ['any', ['==', 'hl_uri', ""]])
-        toggleButton.innerText = 'Show gages';
-        show = false;
-    } else {
-        map.setFilter('conus_gages', null)
-        toggleButton.innerText = 'Hide gages';
-        show = true;
-    }
+showNwm = false;
+const toggleButtonNwm = document.querySelector("#toggle-button-nwm");
+toggleButtonNwm.addEventListener("click", () => {
+  if (showNwm) {
+    map.setFilter("nwm_zarr_chunks", ["any"]);
+    toggleButtonNwm.innerText = "Overlay NWM chunks";
+    showNwm = false;
+  } else {
+    map.setFilter("nwm_zarr_chunks", null);
+    toggleButtonNwm.innerText = "Hide NWM chunks";
+    showNwm = true;
+  }
 });
 
+showAorc = false;
+const toggleButtonAorc = document.querySelector("#toggle-button-aorc");
+toggleButtonAorc.addEventListener("click", () => {
+  if (showAorc) {
+    map.setFilter("aorc_zarr_chunks", ["any"]);
+    toggleButtonAorc.innerText = "Overlay AORC chunks";
+    showAorc = false;
+  } else {
+    map.setFilter("aorc_zarr_chunks", null);
+    toggleButtonAorc.innerText = "Hide AORC chunks";
+    showAorc = true;
+  }
+});

map_app/templates/index.html

@@ -17,15 +17,28 @@
 
 <body>
     <header>
-        <h1>Select catchments by clicking!</h1>        
+        <h1>Select catchments by clicking!</h1>
     </header>
 
     <main>
         <section id="map-container">
-            <div id="map"><button id="toggle-button">Show gages</button></div>  
-                      
+            <div id="map">
+                <button id="toggle-button-gages">Show gages</button> 
+                <button id="toggle-button-camels">Show CAMELS basins</button>
+                <button id="toggle-button-nwm">Overlay NWM chunks</button>
+                <button id="toggle-button-aorc">Overlay AORC chunks</button>
+            </div>
+
             <div class="command-container">
                 <div class="command-header">
+                        <div class="toggle-switch">
+                            <input type="checkbox" id="runcmd-toggle" class="toggle-input">
+                            <label for="runcmd-toggle" class="toggle-label">
+                                <span class="toggle-text toggle-text-left">UVX</span>
+                                <span class="toggle-text toggle-text-right">Pip</span>
+                                <span class="toggle-handle"></span>
+                            </label>
+                        </div>
                     <span>Command</span>
                     <button id="copy-to-clip" class="copy-button" aria-label="Copy to clipboard">
                         <svg class="copy-icon" viewBox="0 0 16 16" width="16" height="16">
@@ -35,22 +48,31 @@
                         <span class="copy-text">Copy</span>
                     </button>
                 </div>
-                <div id="cli-command" class="command-content"></div>
+                <div id="command-builder">
+                    <div id="cli-prefix" class="command-content">uvx --from ngiab_data_preprocess cli </div>
+                    <div id="cli-command" class="command-content"></div>
+                </div>
             </div>
             <script>
                 document.getElementById('copy-to-clip').addEventListener('click', async function() {
+                    const commandPrefix = document.getElementById('cli-prefix').textContent;
                     const commandText = document.getElementById('cli-command').textContent;
+                    if (commandPrefix && !commandPrefix.endsWith(' ')) {
+                        // Ensure there's a space after the prefix
+                        commandPrefix += ' ';
+                    }
+                    const full_command = commandPrefix + commandText;
                     const button = this;
-                    
+
                     try {
-                        await navigator.clipboard.writeText(commandText);
-                        
+                        await navigator.clipboard.writeText(full_command);
+
                         // Update button state
                         button.classList.add('copied');
                         const copyText = button.querySelector('.copy-text');
                         const originalText = copyText.textContent;
                         copyText.textContent = 'Copied!';
-                        
+
                         // Reset button state after 2 seconds
                         setTimeout(() => {
                             button.classList.remove('copied');
@@ -62,7 +84,7 @@
                 });
             </script>
             <h2>For selection by gage id please use the cli.</h2>
-        </section> 
+        </section>
         <section id="basins-container">
             <h2>Selected Basins</h2>
             <div id="selected-basins">None - get clicking!</div>

map_app/views.py

@@ -3,22 +3,22 @@ import logging
 from datetime import datetime
 from pathlib import Path
 
+import geopandas as gpd
 from data_processing.create_realization import create_realization
+from data_processing.dataset_utils import save_and_clip_dataset
+from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
 from data_processing.file_paths import file_paths
 from data_processing.forcings import create_forcings
-from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
-from data_processing.dataset_utils import save_and_clip_dataset
 from data_processing.graph_utils import get_upstream_cats, get_upstream_ids
 from data_processing.subset import subset
 from flask import Blueprint, jsonify, render_template, request
 
-import geopandas as gpd
-
 main = Blueprint("main", __name__)
 intra_module_db = {}
 
 logger = logging.getLogger(__name__)
 
+
 @main.route("/")
 def index():
     return render_template("index.html")
@@ -86,8 +86,8 @@ def get_forcings():
     app = intra_module_db["app"]
     debug_enabled = app.debug
     app.debug = False
-    logger.info(f"get_forcings() disabled debug mode at {datetime.now()}")
-    print(f"forcing_dir: {output_folder}")
+    logger.debug(f"get_forcings() disabled debug mode at {datetime.now()}")
+    logger.debug(f"forcing_dir: {output_folder}")
     try:
         if data_source == "aorc":
             data = load_aorc_zarr(start_time.year, end_time.year)
@@ -96,7 +96,7 @@ def get_forcings():
         gdf = gpd.read_file(paths.geopackage_path, layer="divides")
         cached_data = save_and_clip_dataset(data, gdf, start_time, end_time, paths.cached_nc_file)
 
-        create_forcings(cached_data, paths.output_dir.stem)
+        create_forcings(cached_data, paths.output_dir.stem)  # type: ignore
     except Exception as e:
         logger.info(f"get_forcings() failed with error: {str(e)}")
         return jsonify({"error": str(e)}), 500
@@ -138,5 +138,5 @@ def get_logs():
                 if len(reversed_lines) > 100:
                     break
             return jsonify({"logs": reversed_lines}), 200
-    except Exception as e:
+    except Exception:
         return jsonify({"error": "unable to fetch logs"})

ngiab_data_cli/__main__.py

@@ -1,34 +1,36 @@
+from typing import Tuple
 import rich.status
 
 # add a status bar for these imports so the cli feels more responsive
 with rich.status.Status("Initializing...") as status:
-    from data_sources.source_validation import validate_all
-    from ngiab_data_cli.custom_logging import setup_logging, set_logging_to_critical_only
-    from ngiab_data_cli.arguments import parse_arguments
-    from data_processing.file_paths import file_paths
     import argparse
     import logging
-    import time
-    from typing import List
     import subprocess
     import time
-    from dask.distributed import Client
-    from data_processing.gpkg_utils import get_catid_from_point, get_cat_from_gage_id
+    from typing import List
+
+    import geopandas as gpd
+    from data_processing.create_realization import create_em_realization, create_realization
+    from data_processing.dask_utils import shutdown_cluster
+    from data_processing.dataset_utils import save_and_clip_dataset
+    from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
+    from data_processing.file_paths import file_paths
+    from data_processing.forcings import create_forcings
+    from data_processing.gpkg_utils import get_cat_from_gage_id, get_catid_from_point
     from data_processing.graph_utils import get_upstream_cats
     from data_processing.subset import subset, subset_vpu
-    from data_processing.forcings import create_forcings
-    from data_processing.create_realization import create_realization, create_em_realization
-    from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
-    from data_processing.dataset_utils import save_and_clip_dataset
-    import geopandas as gpd
+    from data_sources.source_validation import validate_output_dir, validate_hydrofabric
+    from ngiab_data_cli.arguments import parse_arguments
+    from ngiab_data_cli.custom_logging import set_logging_to_critical_only, setup_logging
 
 
-def validate_input(args: argparse.Namespace) -> None:
+def validate_input(args: argparse.Namespace) -> Tuple[str, str]:
     """Validate input arguments."""
 
     if args.vpu:
         if not args.output_name:
             args.output_name = f"vpu-{args.vpu}"
+            validate_output_dir()
         return args.vpu, args.output_name
 
     input_feature = args.input_feature.replace("_", "-")
@@ -51,9 +53,11 @@ def validate_input(args: argparse.Namespace) -> None:
         raise ValueError("Cannot use both --latlon and --gage options at the same time.")
 
     if args.latlon:
+        validate_hydrofabric()
         feature_name = get_cat_id_from_lat_lon(input_feature)
         logging.info(f"Found {feature_name} from {input_feature}")
     elif args.gage:
+        validate_hydrofabric()
         feature_name = get_cat_from_gage_id(input_feature)
         logging.info(f"Found {feature_name} from {input_feature}")
     else:
@@ -61,15 +65,18 @@ def validate_input(args: argparse.Namespace) -> None:
 
     if args.output_name:
         output_folder = args.output_name
+        validate_output_dir()
     elif args.gage:
         output_folder = input_feature
+        validate_output_dir()
     else:
         output_folder = feature_name
+        validate_output_dir()
 
     return feature_name, output_folder
 
 
-def get_cat_id_from_lat_lon(input_feature: str) -> List[str]:
+def get_cat_id_from_lat_lon(input_feature: str) -> str:
     """Read catchment IDs from input file or return single ID."""
     if "," in input_feature:
         coords = input_feature.split(",")
@@ -79,7 +86,6 @@ def get_cat_id_from_lat_lon(input_feature: str) -> List[str]:
 
 
 def set_dependent_flags(args, paths: file_paths):
-
     # if validate is set, run everything that is missing
     if args.validate:
         logging.info("Running all missing steps required to run ngiab.")
@@ -122,7 +128,6 @@ def validate_run_directory(args, paths: file_paths):
 
 def main() -> None:
     setup_logging()
-    validate_all()
     try:
         args = parse_arguments()
         if args.debug:
@@ -146,11 +151,15 @@ def main() -> None:
                 subset_vpu(args.vpu, output_gpkg_path=paths.geopackage_path)
                 logging.info("Subsetting complete.")
             else:
-                logging.info(f"Subsetting hydrofabric")
+                logging.info("Subsetting hydrofabric")
                 include_outlet = True
                 if args.gage:
                     include_outlet = False
-                subset(feature_to_subset, output_gpkg_path=paths.geopackage_path, include_outlet=include_outlet)
+                subset(
+                    feature_to_subset,
+                    output_gpkg_path=paths.geopackage_path,
+                    include_outlet=include_outlet,
+                )
                 logging.info("Subsetting complete.")
 
         if args.forcings:
@@ -189,14 +198,6 @@ def main() -> None:
                 )
             logging.info("Realization creation complete.")
 
-        # check if the dask client is still running and close it
-        try:
-            client = Client.current()
-            client.shutdown()
-        except ValueError:
-            # value error is raised if no client is running
-            pass
-
         if args.run:
             logging.info("Running Next Gen using NGIAB...")
             # open the partitions.json file and get the number of partitions
@@ -218,7 +219,8 @@ def main() -> None:
         if args.eval:
             plot = False
             try:
-                import seaborn, matplotlib
+                import matplotlib
+                import seaborn
 
                 plot = True
             except ImportError:
@@ -254,6 +256,7 @@ def main() -> None:
     except Exception as e:
         logging.error(f"An error occurred: {str(e)}")
         raise
+    shutdown_cluster()
 
 
 if __name__ == "__main__":

ngiab_data_cli/arguments.py

@@ -1,6 +1,5 @@
 import argparse
 from datetime import datetime
-from data_processing.file_paths import file_paths
 
 # Constants
 DATE_FORMAT = "%Y-%m-%d"  # used for datetime parsing