NWM_data_download.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download retrospective data from the National Water Model (NWM) amazon bucket\n",
    "- Inputs: Excel files or shapefile that consists the information of the COMIDs that you want to download the NWM retrospective data for\n",
    "- Output: Streamflows from NWM retrospective analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import fsspec\n",
    "import pandas as pd\n",
    "import dask\n",
    "\n",
    "# Define the path to the excel or base directory\n",
    "comid_stn = pd.read_excel('usgs_with_comid.xlsx')\n",
    "all_sites = comid_stn['COMID']\n",
    "\n",
    "# URI for NWM 3.0 data source\n",
    "nwm_loc = 's3://noaa-nwm-retrospective-3-0-pds/CONUS/zarr/chrtout.zarr'\n",
    "# URI for NWM 2.1 data source\n",
    "# nwm_loc = 's3://noaa-nwm-retrospective-2-1-zarr-pds/chrtout.zarr'\n",
    "\n",
    "# Open the NWM dataset with fsspec \n",
    "ds = xr.open_zarr(fsspec.get_mapper(nwm_loc, anon=True, requester_pays=True), chunks='auto')\n",
    "\n",
    "# Define the time range for which the data is to be downloaded. NWM 2.1 has data from 1979-2020 and NWM 3.0 has data from 1979-2023 \n",
    "start_time = '2018-03-01T00:00:00'\n",
    "end_time = '2019-01-31T00:00:00'\n",
    "\n",
    "# Initialize an empty DataFrame to store the results\n",
    "nwm = pd.DataFrame()\n",
    "\n",
    "# Select the data for all COMIDs and the time range\n",
    "sub_ds = ds.sel(feature_id=all_sites.values, time=slice(start_time, end_time))\n",
    "\n",
    "# Computation on the streamflow\n",
    "streamflow_data = sub_ds['streamflow'].compute()\n",
    "\n",
    "# Extract the time index \n",
    "nwm['timestamp'] = pd.to_datetime(sub_ds['time'].values)\n",
    "\n",
    "#Store streamflow data in a  dataframe\n",
    "for i, reach_id in enumerate(all_sites):\n",
    "    try:\n",
    "        reach_id_numeric = pd.to_numeric(reach_id)\n",
    "        nwm[str(reach_id_numeric)] = streamflow_data[:, i].values\n",
    "    except KeyError as e:\n",
    "        print(f\"COMID {reach_id_numeric} not found in NWM dataset. Error: {e}\")\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred while processing COMID {reach_id_numeric}: {e}\")\n",
    "print('c')\n",
    "# Save the DataFrame to a CSV file or anyfile that you ant to work with\n",
    "output_file = 'nwm_discharge_test2.csv'\n",
    "nwm.to_csv(output_file)\n",
    "\n",
    "print(f\"Data saved to {output_file}\")\n",
    "print(nwm.head())  \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.18 ('assimilation')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4bc09726d3eb0feb4d73d3636dad7388c6ae5ba645f208395a639ba59e5d5f59"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}