NWM_data_download_short_range.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is to download the NWM shortrange data from the google archive. Originally compiled by: Mark Wang, passaH2O group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import python libraries\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pandas as pd\n",
    "import requests\n",
    "import sys\n",
    "import xarray as xr\n",
    "import re\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib.dates import DateFormatter\n",
    "from time import sleep\n",
    "\n",
    "# select input parameters\n",
    "\n",
    "begindate = '20201103' # data is avail. beginning 20180917\n",
    "enddate = '20201104'\n",
    "# enddate = '20190930' # inclusive\n",
    "product = 'channel_rt' # 'channel_rt' or 'forcing'\n",
    "destfolder_name = 'november2020' # directory will be created in current working directory if it does not already exist\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define functions\n",
    "\n",
    "\n",
    "def get_netcdf(filetype,begin_date,end_date,output_folder_name):\n",
    "    \n",
    "    output_folder = os.path.join(os.getcwd(),output_folder_name)\n",
    "    \n",
    "    if not os.path.exists(output_folder):\n",
    "        os.mkdir(output_folder)\n",
    "    \n",
    "    if filetype =='channel_rt':\n",
    "        prodstr = ''\n",
    "    elif filetype == 'forcing':\n",
    "        prodstr = 'forcing_'\n",
    "    else:\n",
    "        print(\"Product error. Choose 'channel_rt' or 'forcing'.\")\n",
    "        sys.exit()\n",
    "    \n",
    "        \n",
    "    \n",
    "    t1 = pd.to_datetime(begin_date)\n",
    "    t2 = pd.to_datetime(end_date)\n",
    "    dates = pd.date_range(t1,t2)\n",
    "    \n",
    "    for i in range(len(dates)):\n",
    "        date = dates[i]\n",
    "        save_dir = date.strftime('%Y%m%d')\n",
    "        \n",
    "        if not os.path.exists(output_folder+'/'+save_dir):\n",
    "            os.mkdir(output_folder+'/'+save_dir)\n",
    "            \n",
    "        for hr in range(24):\n",
    "            url = f'https://storage.googleapis.com/national-water-model/' \\\n",
    "                  f'nwm.{date.strftime(\"%Y%m%d\")}/{prodstr}short_range/' \\\n",
    "                  f'nwm.t{str(hr).zfill(2)}z.short_range.{filetype}.f001.conus.nc'\n",
    "            filename = os.path.basename(url)\n",
    "            write_file = os.path.join(output_folder, save_dir, filename)\n",
    "\n",
    "            for attempt in range(30):\n",
    "                try:\n",
    "                    r = requests.get(url)\n",
    "                    with open(write_file, 'wb') as f:\n",
    "                        f.write(r.content)\n",
    "                    break\n",
    "                except Exception as ex:\n",
    "                    if attempt != max_attempts - 1:\n",
    "                        sleep(0.5)  # Give NOAA time to wake up\n",
    "                    else:\n",
    "                        m = 'Could not download file.\\n' + str(ex)\n",
    "                        raise Exception(m)\n",
    "                        \n",
    "        print(f'{save_dir} done')\n",
    "       \n",
    "        \n",
    "def get_series(comid, begin_date, end_date, datafolder_name, filetype):\n",
    "    \n",
    "    data_folder = os.path.join(os.getcwd(), datafolder_name)\n",
    "    \n",
    "    t1 = pd.to_datetime(begin_date)\n",
    "    t2 = pd.to_datetime(end_date)\n",
    "    dates = pd.date_range(t1,t2)\n",
    "    \n",
    "    df = pd.DataFrame()\n",
    "    \n",
    "    for i in range(len(dates)):\n",
    "        date = dates[i]\n",
    "        date_dir = date.strftime('%Y%m%d')\n",
    "            \n",
    "        for hr in range(24):\n",
    "            filename = f'nwm.t{str(hr).zfill(2)}z.short_range.{filetype}.f001.conus.nc'\n",
    "            nc_file = os.path.join(data_folder, date_dir, filename)\n",
    "            try:\n",
    "                data = xr.open_dataset(nc_file, engine=\"netcdf4\")\n",
    "                Q = float(data.sel(feature_id=comid).streamflow.values)\n",
    "            except OSError as e:\n",
    "                if re.match('.*NetCDF: Unknown file format:.*', str(e)):\n",
    "                    Q=np.nan\n",
    "                else:\n",
    "                        raise e\n",
    "            timestamp = pd.to_datetime(f'{date_dir} {hr}:00')\n",
    "            df.loc[timestamp,'Q'] = Q\n",
    "    \n",
    "    return df\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download files\n",
    "get_netcdf(product,begindate,enddate,destfolder_name)"
   ]
  }
 ],
 "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
}