subset.AR.ERA5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Subset ERA5 data needed to calculate atmospheric rivers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This program:\n",
    "\n",
    "Reads in U, V, Q at 14 levels from ERA5 for Nov-Dec-Jan-Feb-Mar 1986/87 to 2015/16\n",
    "\n",
    "Data are located in: \n",
    "\n",
    "    /glade/collections/rda/data/ds633.0/e5.oper.an.pl/yyyymm/e5.oper.an.pl.128_133_variable.ll025sc.yyyymmddhh_yyyymmddhh.grb\n",
    "\n",
    "Subsets to the PNA region\n",
    "    \n",
    "    Defined as: 150-300E; 20-80N\n",
    "    \n",
    "Choose whether to interpolate the data to a N32 (126x64) Gaussian Grid or an N128 (512 x 256) grid, or select the native grid.\n",
    "\n",
    "Writes datasets out to a netcdf files containing a single variable\n",
    "\n",
    "Output file is located in: `/glade/work/mkorendy/metisregimes/era5.AR.variable_pna_NDJFM.1986-2016.nc`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose interpolation: N32 or N128 or native\n",
    "interpolate = 'N128'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function to get the PNA region for the variable at the 14 levels, and interpolate to new grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preproc(ds):\n",
    "    \n",
    "    pressures = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 100, 50, 30, 10]\n",
    "\n",
    "    # Extract PNA region for every 12 hours for the 14 pressure levels\n",
    "    ds=ds.sel(time=ds['time'][0::12],isobaricInhPa = pressures,latitude=slice(80,20),longitude=slice(150,300))\n",
    "    \n",
    "    #Extract PNA region\n",
    "    #ds = ds.sel(latitude=slice(20,80),longitude=slice(150,300))\n",
    "\n",
    "    # Create a new xarray dataset for the variable\n",
    "    #ds_tmp=xr.DataArray(ds,\n",
    "    #                    coords={'time':ds['time'].values, \n",
    "    #                            'levs':ds['isobaricInhPa'].values,\n",
    "    #                            'lat':ds['latitude'].values,\n",
    "    #                            'lon': ds['longitude'].values},\n",
    "    #                        dims=['time','levs','latitude','longitude'])        \n",
    "    #ds_tmp=ds_tmp.to_dataset(name=variable)\n",
    "    \n",
    "    ds = ds.rename(latitude='lat',longitude='lon')\n",
    "    \n",
    "    # Interpolate data to coarser grid\n",
    "    if interpolate == 'N32':\n",
    "        ds_new_grid=xr.open_dataset('tempgaussian.128x64.nc')\n",
    "        ds_new_grid=ds_new_grid.rename({'temp':variable}).sel(lat=slice(80,20),lon=slice(150,300))\n",
    "        #ds=ds_tmp.interp_like(ds_new_grid)\n",
    "        ds=ds.interp_like(ds_new_grid)\n",
    "    elif interpolate == 'N128':\n",
    "        ds_new_grid=xr.open_dataset('/glade/work/mkorendy/metisregimes/metis.vqvi_pna_NDJFM.1986-2016.N128.nc')\n",
    "        lat = ds_new_grid.sel(lat=slice(20,80),lon=slice(150,300)).lat\n",
    "        lon = ds_new_grid.sel(lat=slice(20,80),lon=slice(150,300)).lon\n",
    "        go = np.empty([len(lat),len(lon)])\n",
    "        ds_new_grid = xr.DataArray(go,coords={'lat':lat,'lon':lon}, dims=['lat','lon']).to_dataset(name=variable)\n",
    "        #ds=ds_tmp.interp_like(ds_new_grid)\n",
    "        ds=ds.interp_like(ds_new_grid)\n",
    "    else:\n",
    "        print('native')\n",
    "        \n",
    "    return ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define months and years to get\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2015]\n"
     ]
    }
   ],
   "source": [
    "# Months list (Nov, Dec, Jan, Feb, Mar)\n",
    "mnums=['11', '12', '01', '02', '03']\n",
    "#mnums = ['04','05','06','07','08','09']\n",
    "#mnums = ['04']\n",
    "mnums = ['03']\n",
    "\n",
    "# Dates\n",
    "sdate='19861101'\n",
    "edate='20160331'\n",
    "\n",
    "# Years List\n",
    "yrs_list=np.arange(1986,2016)\n",
    "#yrs_list = np.arange(1987,1988)\n",
    "yrs_list = np.arange(2015,2016)\n",
    "print(yrs_list)\n",
    "\n",
    "#Variable\n",
    "variable = 'q'\n",
    "\n",
    "#Names for key-wargs\n",
    "if variable == 'u':\n",
    "    name = 'U component of wind'\n",
    "    file = '/*u.*uv.*grb'\n",
    "elif variable == 'v':\n",
    "    name = 'V component of wind'\n",
    "    file = '/*v.*uv.*grb'\n",
    "elif variable == 'q':\n",
    "    name = 'Specific humidity'\n",
    "    file = '/*q*grb'\n",
    "    #file = '/*q*nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define file path and names¶"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input\n",
    "path='/glade/collections/rda/data/ds633.0/e5.oper.an.pl/'\n",
    "\n",
    "# Output\n",
    "out_path='/glade/work/mkorendy/metisregimes/'\n",
    "\n",
    "if interpolate == 'N32':\n",
    "    ofname = 'era5.AR.' + variable + '_pna_NDJFM.1986-2016.nc'\n",
    "elif interpolate == 'native':\n",
    "    ofname = 'era5.AR.' + variable + '_pna_NDJFM.1986-2016.reg_grid.nc'\n",
    "else:\n",
    "    ofname = 'era5.AR.' + variable + '_pna_NDJFM.1986-2016.N128.nc'\n",
    "#ofname='era5.AR.' + variable + '_pna_AMJJAS.1986-2016.nc'\n",
    "#ofname = 'era5.AR.' + variable + '_pna_O.1986-2016.nc'\n",
    "outfile=out_path+ofname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/glade/work/mkorendy/metisregimes/era5.AR.q_pna_NDJFM.1986-2016.N128.nc\n"
     ]
    }
   ],
   "source": [
    "print(outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test = xr.open_dataset('/glade/collections/rda/data/ds633.0/e5.oper.an.pl/199710/e5.oper.an.pl.128_133_q.ll025sc.1997102500_1997102523.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "pressures = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 100, 50, 30, 10]\n",
    "#test.sel(time=test['time'][0::12], level = pressures,latitude=slice(80,20),longitude=slice(150,300))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in the files for all months and years, and prepocess to subset data for U, V, or Q in the PNA region & interpolate to coarser grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "unrecognized engine cfgrib must be one of: ['netcdf4', 'scipy', 'pydap', 'store', 'zarr']",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/glade/scratch/mkorendy/ipykernel_149704/2029367224.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     17\u001b[0m                            backend_kwargs={'indexpath':'',\n\u001b[1;32m     18\u001b[0m                                            'filter_by_keys':{'name': name}},\n\u001b[0;32m---> 19\u001b[0;31m                           preprocess=preproc)\n\u001b[0m\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0mds_months\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/mkorendy/conda-envs/wxregimes-npl/lib/python3.7/site-packages/xarray/backends/api.py\u001b[0m in \u001b[0;36mopen_mfdataset\u001b[0;34m(paths, chunks, concat_dim, compat, preprocess, engine, data_vars, coords, combine, parallel, join, attrs_file, combine_attrs, **kwargs)\u001b[0m\n\u001b[1;32m    906\u001b[0m         \u001b[0mgetattr_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    907\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 908\u001b[0;31m     \u001b[0mdatasets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mopen_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mopen_kwargs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpaths\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    909\u001b[0m     \u001b[0mclosers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mgetattr_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"_close\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mds\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdatasets\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    910\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mpreprocess\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/mkorendy/conda-envs/wxregimes-npl/lib/python3.7/site-packages/xarray/backends/api.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    906\u001b[0m         \u001b[0mgetattr_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    907\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 908\u001b[0;31m     \u001b[0mdatasets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mopen_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mopen_kwargs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpaths\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    909\u001b[0m     \u001b[0mclosers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mgetattr_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"_close\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mds\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdatasets\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    910\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mpreprocess\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/glade/work/mkorendy/conda-envs/wxregimes-npl/lib/python3.7/site-packages/xarray/backends/api.py\u001b[0m in \u001b[0;36mopen_dataset\u001b[0;34m(filename_or_obj, engine, chunks, cache, decode_cf, mask_and_scale, decode_times, decode_timedelta, use_cftime, concat_characters, decode_coords, drop_variables, backend_kwargs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    479\u001b[0m         \u001b[0mengine\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mplugins\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mguess_engine\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilename_or_obj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    480\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 481\u001b[0;31m     \u001b[0mbackend\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mplugins\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_backend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mengine\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    482\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    483\u001b[0m     decoders = _resolve_decoders_kwargs(\n",
      "\u001b[0;32m/glade/work/mkorendy/conda-envs/wxregimes-npl/lib/python3.7/site-packages/xarray/backends/plugins.py\u001b[0m in \u001b[0;36mget_backend\u001b[0;34m(engine)\u001b[0m\n\u001b[1;32m    162\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mengine\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mengines\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    163\u001b[0m             raise ValueError(\n\u001b[0;32m--> 164\u001b[0;31m                 \u001b[0;34mf\"unrecognized engine {engine} must be one of: {list(engines)}\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    165\u001b[0m             )\n\u001b[1;32m    166\u001b[0m         \u001b[0mbackend\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mengines\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mengine\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: unrecognized engine cfgrib must be one of: ['netcdf4', 'scipy', 'pydap', 'store', 'zarr']"
     ]
    }
   ],
   "source": [
    "ds_months = []\n",
    "\n",
    "for mnum in mnums:\n",
    "\n",
    "    fnames = [f'{path}{year}{mnum}' + file for year in yrs_list]\n",
    "\n",
    "    filesVar=[]\n",
    "    for files in fnames:\n",
    "        f2=glob.glob(files)\n",
    "        for f in f2:\n",
    "            filesVar.append(f)\n",
    "    \n",
    "    filesVar=sorted(filesVar)\n",
    "    \n",
    "    ds = xr.open_mfdataset(filesVar,engine='cfgrib',\n",
    "                           combine='nested',concat_dim=['time'],\n",
    "                           backend_kwargs={'indexpath':'',\n",
    "                                           'filter_by_keys':{'name': name}},\n",
    "                          preprocess=preproc)\n",
    "    \n",
    "    ds_months.append(ds)\n",
    "\n",
    "ds_months = xr.combine_nested(ds_months, concat_dim='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_months"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write to netcdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_months.to_netcdf(outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/glade/work/mkorendy/metisregimes/era5.AR.q_pna_NDJFM.1986-2016.N128.nc\n"
     ]
    }
   ],
   "source": [
    "print(outfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (wxregimes)",
   "language": "python",
   "name": "wxregimes"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}