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

map_app/static/js/main.js

@@ -1,34 +1,34 @@
 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() {
-    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);
-    }
+  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"
+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;
 }
@@ -41,121 +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 = '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 = '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
 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);
-        });
+  // 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);
+document.addEventListener("DOMContentLoaded", initializeToggleSwitches);
+
+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;
+  }
+});
 
+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

@@ -22,7 +22,12 @@
 
     <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">
@@ -52,6 +57,10 @@
                 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;
 

map_app/views.py

@@ -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

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

ngiab_data_cli/forcing_cli.py

@@ -1,19 +1,17 @@
-from data_sources.source_validation import validate_all
-from ngiab_data_cli.custom_logging import setup_logging
-from data_processing.forcings import compute_zonal_stats
-from data_processing.dataset_utils import check_local_cache, save_to_cache, clip_dataset_to_bounds
-from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
-from data_processing.file_paths import file_paths
 import argparse
 import logging
+import shutil
 import time
-import xarray as xr
-import geopandas as gpd
 from datetime import datetime
 from pathlib import Path
-import shutil
 
-from dask.distributed import Client, LocalCluster
+import geopandas as gpd
+from data_processing.dask_utils import shutdown_cluster
+from data_processing.dataset_utils import check_local_cache, clip_dataset_to_bounds, save_to_cache
+from data_processing.datasets import load_aorc_zarr, load_v3_retrospective_zarr
+from data_processing.forcings import compute_zonal_stats
+from data_sources.source_validation import validate_all
+from ngiab_data_cli.custom_logging import setup_logging
 
 # Constants
 DATE_FORMAT = "%Y-%m-%d"  # used for datetime parsing
@@ -69,6 +67,7 @@ def parse_arguments() -> argparse.Namespace:
 
     return parser.parse_args()
 
+
 def main() -> None:
     time.sleep(0.01)
     setup_logging()
@@ -111,15 +110,7 @@ def main() -> None:
     # remove the working directory
     shutil.rmtree(forcing_working_dir)
 
-    try:
-        client = Client.current()
-    except ValueError:
-        cluster = LocalCluster()
-        client = Client(cluster)
-        cluster.close()
-
-    # shut down the client and cluster
-    client.close()
+    shutdown_cluster()
 
 
 if __name__ == "__main__":