aggregate_precip_to_HRU.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aggregate gridded data to HRU using areal averaging\n",
    "\n",
    "### Output Format:\n",
    "\n",
    "year month day hour minute second hru1 hru2 hru3\n",
    "\n",
    "1996 10 1 0 0 0 0.004 0.050 0.070\n",
    "\n",
    "1996 10 2 0 0 0 0.500 0.040 0.100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import rasterio as rs\n",
    "import geopandas as gpd\n",
    "from netCDF4 import Dataset\n",
    "import os\n",
    "import datetime\n",
    "import progressbar as bar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def process_tiffs(df):\n",
    "    reg = df.reg\n",
    "    nhru = df.nhruID\n",
    "    \n",
    "    fl = './data/nhrus/AEA_tiffs/HUC_%s_nhruID_%s.tiff'%(reg,nhru)\n",
    "    \n",
    "    if os.path.isfile(fl): # only proceed if the tiff exists\n",
    "        with rs.open(fl) as ds:\n",
    "            rast = ds.read(1)\n",
    "\n",
    "        n,m = rast.shape\n",
    "        rast.shape = n*m\n",
    "        rast = rast[rast!=0] # remove no data cells\n",
    "        k = float(len(rast))\n",
    "\n",
    "        cells = np.unique(rast)\n",
    "        #print(len(cells))\n",
    "        percents = []\n",
    "        for cell in cells:\n",
    "            percents.append(len(rast[rast==cell])/k) # divide by the total cells in the basin to get the propotion of each cell in the basin\n",
    "\n",
    "        cells = list(cells)\n",
    "        return cells,percents\n",
    "    \n",
    "    else:\n",
    "        return [],[]\n",
    "    \n",
    "def process_tiffs_12(df):\n",
    "    reg = df.reg\n",
    "    nhru = df.reg_hruID\n",
    "    \n",
    "    fl = './data/nhrus/AEA_tiffs/HUC_%s_nhruID_%s.tiff'%(reg,nhru)\n",
    "    \n",
    "    if os.path.isfile(fl): # only proceed if the tiff exists\n",
    "        with rs.open(fl) as ds:\n",
    "            rast = ds.read(1)\n",
    "\n",
    "        n,m = rast.shape\n",
    "        rast.shape = n*m\n",
    "        rast = rast[rast!=0] # remove no data cells\n",
    "        k = float(len(rast))\n",
    "\n",
    "        cells = np.unique(rast)\n",
    "        #print(len(cells))\n",
    "        percents = []\n",
    "        for cell in cells:\n",
    "            percents.append(len(rast[rast==cell])/k) # divide by the total cells in the basin to get the propotion of each cell in the basin\n",
    "\n",
    "        cells = list(cells)\n",
    "        return cells,percents\n",
    "    \n",
    "    else:\n",
    "        return [],[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_contributions(fl,test=False):\n",
    "    tmp = gpd.read_file(fl)\n",
    "    dat = pd.DataFrame() # first generate a list of hrus and their grid cell contributions\n",
    "    dat['nhruID'] = tmp.hru_id_nat\n",
    "    dat['reg_hruID'] = tmp.hru_id_reg\n",
    "    reg = fl.split('_')[-2]\n",
    "    dat['reg'] = tmp.region\n",
    "    cells,percents = zip(*dat.apply(process_tiffs,axis=1)) # run the aggregation function\n",
    "    dat['cells'] = cells # insert results back into the dataframe\n",
    "    dat['percents'] = percents\n",
    "\n",
    "    if (reg == '08') | (reg == '10U') | (reg == '04'): # if either of these regions occur\n",
    "        dat2 = pd.read_pickle('./data/reg%s_unclipped.pcl'%reg) # load the missing data\n",
    "        \n",
    "        for nhru in dat2.nhruID: # remove the overlapping rows from the data frame\n",
    "            dat = dat[dat.nhruID != nhru]\n",
    "        \n",
    "        dat = dat.append(dat2) # merge the two dataframes\n",
    "    \n",
    "    # write some tests\n",
    "    if test:\n",
    "        # do all the percentages equal very close to 1\n",
    "        def test_percent(df):\n",
    "            perc = np.sum(df.percents)\n",
    "            if 1-perc > 0.0001:\n",
    "                print('percent does not sum: %s: %s'%(1-perc,df.nhruID))\n",
    "            \n",
    "        dat.apply(test_percent,axis=1)\n",
    "        # is the length of the original df the same as the produced one\n",
    "        if len(tmp) - len(dat) > 0:\n",
    "            print('data frames are different lengths')\n",
    "    \n",
    "    dat.to_pickle('./data/nhru_contrib/huc_%s_cell_contrib.pcl'%reg)\n",
    "    print('%s Complete!'%reg)\n",
    "    \n",
    "def compute_contributions_12(fl,test=False):\n",
    "    '''for region 12'''\n",
    "    tmp = gpd.read_file(fl)\n",
    "    dat = pd.DataFrame() # first generate a list of hrus and their grid cell contributions\n",
    "    dat['reg_hruID'] = tmp.hru_id\n",
    "    reg = fl.split('_')[-2]\n",
    "    dat['reg'] = tmp.region\n",
    "    cells,percents = zip(*dat.apply(process_tiffs_12,axis=1)) # run the aggregation function\n",
    "    dat['cells'] = cells # insert results back into the dataframe\n",
    "    dat['percents'] = percents\n",
    "    \n",
    "    # write some tests\n",
    "    if test:\n",
    "        # do all the percentages equal very close to 1\n",
    "        def test_percent(df):\n",
    "            perc = np.sum(df.percents)\n",
    "            if 1-perc > 0.0001:\n",
    "                print('percent does not sum: %s: %s'%(1-perc,df.nhruID))\n",
    "            \n",
    "        dat.apply(test_percent,axis=1)\n",
    "        # is the length of the original df the same as the produced one\n",
    "        if len(tmp) - len(dat) > 0:\n",
    "            print('data frames are different lengths')\n",
    "    \n",
    "    dat.to_pickle('./data/nhru_contrib/huc_%s_cell_contrib.pcl'%reg)\n",
    "    print('%s Complete!'%reg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# list regions\n",
    "regions = glob.glob('./data/nhrus/clean_AEA/nhru_*_clean.shp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12 Complete!\n"
     ]
    }
   ],
   "source": [
    "# run region 12\n",
    "compute_contributions_12(regions[-1],test=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "07 Complete!\n",
      "08 Complete!\n",
      "09 Complete!\n",
      "10L Complete!\n",
      "10U Complete!\n",
      "11 Complete!\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[None, None, None, None, None, None]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[compute_contributions(reg,test=True) for reg in regions]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Crop the radar data by each HRU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# load the precip data\n",
    "fl = './stage4_map_daily_20041220-20150107.nc'\n",
    "ds = Dataset(fl,'r')\n",
    "m,k,l = ds.variables['Total_precipitation_surface_1_Hour_Accumulation'].shape # get the dimensions of the precip data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# unsure if this is needed\n",
    "\n",
    "# load the index raster\n",
    "#idx = np.load('./data/hrap_grid_index.npy')\n",
    "#idx.shape = k*l "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# compute the dates\n",
    "time = ds.variables['time']\n",
    "#print('Time Units: %s'%time.units)\n",
    "timeoffset = time.units[-20:] # strip the string\n",
    "strt = pd.to_datetime(timeoffset) # convert string into datetime object\n",
    "time = np.array(ds.variables['time'])\n",
    "\n",
    "def compute_time(time,offset):\n",
    "    dt = datetime.timedelta(hours=time)\n",
    "    time = offset+dt\n",
    "    return str(time.date())\n",
    "\n",
    "times = np.vectorize(compute_time)(time,strt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def year(index): return index.year\n",
    "def month(index): return index.month\n",
    "def day(index): return index.day"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_precip(df,datetime=[],rast=[],out=[]):\n",
    "    '''\n",
    "    Compute precip for an hru based on its contributing grid cells\n",
    "    '''\n",
    "    \n",
    "    precip = rast[df.cells]\n",
    "    percents = np.array(df.percents)\n",
    "    \n",
    "    weighted_precip = np.sum(precip*percents) # precip in mm\n",
    "    weighted_precip *= 0.0393701 # mm > inches\n",
    "\n",
    "    out.loc[datetime,'hru_%s'%df.reg_hruID] = weighted_precip # insert into the out data frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def generate_output(fl):\n",
    "    reg = fl.split('_')[-2] # extract the region\n",
    "    print('Starting region %s...'%reg)\n",
    "    dat = pd.read_pickle('./data/nhru_contrib/huc_%s_cell_contrib.pcl'%reg) # load the contributing cells and percentages\n",
    "    dat.sort_values('reg_hruID',inplace=True,ascending=True) # sort by regional hru\n",
    "    \n",
    "    # prepair the output data frame\n",
    "    out = pd.DataFrame()\n",
    "    out['datetime'] = pd.DatetimeIndex(times)\n",
    "    out.index = pd.DatetimeIndex(out.datetime)\n",
    "    out['year'] = out.index.map(year)\n",
    "    out['month'] = out.index.map(month)\n",
    "    out['day'] = out.index.map(day)\n",
    "    out['hour'] = 0\n",
    "    out['minute'] = 0\n",
    "    out['second'] = 0\n",
    "\n",
    "    for hru in dat.reg_hruID: # create space for each HRU\n",
    "        out['hru_%s'%hru] = -999\n",
    "\n",
    "    del out['datetime'] # clean up\n",
    "    \n",
    "    pb = bar.ProgressBar(min_value=0,max_value=m)\n",
    "\n",
    "    for i in range(m): # iterate through slices of the dataset\n",
    "        rast = np.array(ds.variables['Total_precipitation_surface_1_Hour_Accumulation'][i,:,:]) # pull a slice\n",
    "        rast.shape = (k*l) # reshape the dataset in the say way as the index values\n",
    "\n",
    "        dat.apply(compute_precip,axis=1,datetime=times[i],rast=rast,out=out) # compute precip for each hru for the time slice\n",
    "        pb.update(i)\n",
    "    \n",
    "    \n",
    "    out.to_csv('./data/hru_%s_stage_4_precip.cbh'%reg,sep=' ',header=False,index=False,float_format='%.2f')\n",
    "    out.to_pickle('./data/hru_%s_stage_4_precip.pcl'%reg)\n",
    "    \n",
    "    print('Region %s complete!'%reg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting region 12...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 99% (3668 of 3669) |############################################################################################################################ | Elapsed Time: 7:54:59 ETA: 0:00:07"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Region 12 complete!\n"
     ]
    }
   ],
   "source": [
    "# run only region 12\n",
    "generate_output(regions[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting region 11...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 99% (3668 of 3669) |############################################################################ | Elapsed Time: 7:03:55 ETA: 0:00:06"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Region 11 complete!\n"
     ]
    }
   ],
   "source": [
    "for fl in regions[-1:]:\n",
    "    generate_output(fl)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "out.to_csv('./data/hru_%s_stage_4_precip.cbh'%reg,sep=' ',header=False,index=False,float_format='%.2f')\n",
    "out.to_pickle('./data/hru_%s_stage_4_precip.pcl'%reg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['./data/nhrus/clean_AEA/nhru_11_clean.shp']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "regions[-1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}