{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "change the river discharge source point of Fraser and create new Fraser River flow file "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Re-write monthly & yearly rivers file with different freshwater grid cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "from salishsea_tools import rivertools\n",
    "from salishsea_tools import nc_tools\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import netCDF4 as nc\n",
    "import arrow\n",
    "import numpy.ma as ma\n",
    "import sys\n",
    "sys.path.append('/ocean/klesouef/meopar/tools/I_ForcingFiles/Rivers')\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "filename = '/ocean/jieliu/research/meopar/nemo-forcing/rivers/rivers_month.nc'\n",
    "clim_rivers = nc.Dataset(filename, 'r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<type 'netCDF4.Dimension'>: name = 'x', size = 398\n",
      "\n",
      "<type 'netCDF4.Dimension'>: name = 'y', size = 898\n",
      "\n",
      "<type 'netCDF4.Dimension'> (unlimited): name = 'time_counter', size = 12\n",
      "\n",
      "[u'nav_lat', u'nav_lon', u'time_counter', u'rorunoff', u'rodepth', u'rotemper']\n"
     ]
    }
   ],
   "source": [
    "nc_tools.show_dimensions(clim_rivers)\n",
    "nc_tools.show_variables(clim_rivers)\n",
    "criverflow = clim_rivers.variables['rorunoff']\n",
    "# get other variables so we can put them in new files\n",
    "lat = clim_rivers.variables['nav_lat']\n",
    "lon = clim_rivers.variables['nav_lon']\n",
    "riverdepth = clim_rivers.variables['rodepth']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "rivertype = 'constant' ## monthly or constant(yearly)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if rivertype == 'monthly':\n",
    "    fluxfile = nc.Dataset('/ocean/sallen/allen/research/MEOPAR/nemo-forcing/rivers/Salish_allrivers_monthly.nc','r')\n",
    "    #inialise the runoff and run_depth arrays\n",
    "    runoff, run_depth, run_temp = rivertools.init_runoff_array_monthly()\n",
    "#get river fluxes from netcdf file\n",
    "if rivertype == 'constant':\n",
    "    fluxfile = nc.Dataset('/ocean/sallen/allen/research/MEOPAR/nemo-forcing/rivers/Salish_allrivers_cnst.nc','r')\n",
    "    #inialise the runoff and run_depth arrays\n",
    "    runoff, run_depth, run_temp = rivertools.init_runoff_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#list of watersheds we are including\n",
    "names = ['skagit','fraser','evi_n','howe','bute','puget','jdf','evi_s','jervis','toba']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "skagit has 10 rivers\n",
      "(908.05509995059549, 937.3491821289062)\n",
      "fraser has 10 rivers\n",
      "(3578.4171797583772, 3549.736083984375)\n",
      "evi_n has 21 rivers\n",
      "(255.62751641603765, 638.462158203125)\n",
      "howe has 2 rivers\n",
      "(588.99650864276578, 572.8439331054688)\n",
      "bute has 3 rivers\n",
      "(609.26456547060798, 550.5059204101562)\n",
      "puget has 43 rivers\n",
      "(480.71352669585275, 502.76959228515625)\n",
      "jdf has 27 rivers\n",
      "(399.59934187151299, 409.8137512207031)\n",
      "evi_s has 17 rivers\n",
      "(330.71661134751884, 329.566162109375)\n",
      "jervis has 17 rivers\n",
      "(307.47534529284559, 296.90533447265625)\n",
      "toba has 1 rivers\n",
      "(285.36386496863395, 270.2106018066406)\n"
     ]
    }
   ],
   "source": [
    "for name in range(0,len(names)):\n",
    "    watershedname = names[name]\n",
    "    Flux = fluxfile.variables[watershedname][:]\n",
    "    if rivertype == 'constant':\n",
    "        Flux = float(Flux)\n",
    "    runoff_orig = np.copy(runoff)\n",
    "    runoff, run_depth, run_temp = rivertools.put_watershed_into_runoff(rivertype,\n",
    "                                            watershedname, Flux, runoff, run_depth, run_temp)\n",
    "    if rivertype == 'monthly':\n",
    "        rivertools.check_sum_monthly(runoff_orig, runoff, Flux)\n",
    "    if rivertype == 'constant':\n",
    "        rivertools.check_sum(runoff_orig, runoff, Flux)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1.0\n"
     ]
    }
   ],
   "source": [
    "print run_depth[414,334]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if rivertype == 'monthly':\n",
    "    \n",
    "    nemo = nc.Dataset('/ocean/jieliu/research/meopar/river-treatment/rivers_month_edit.nc', 'w') \n",
    "    nemo.description = 'Monthly Averages, All Rivers, modify on depth and runoff grid point' \n",
    "    \n",
    "    # dimensions\n",
    "    nemo.createDimension('x', 398) \n",
    "    nemo.createDimension('y', 898)\n",
    "    nemo.createDimension('time_counter', None)\n",
    "    \n",
    "    # variables\n",
    "    # latitude and longitude\n",
    "    nav_lat = nemo.createVariable('nav_lat','float32',('y','x'),zlib=True)\n",
    "    nav_lat = lat\n",
    "    x = nemo.createVariable('nav_lon','float32',('y','x'),zlib=True)\n",
    "    nav_lon = lon\n",
    "    # time\n",
    "    time_counter = nemo.createVariable('time_counter', 'float32', ('time_counter'),zlib=True)\n",
    "    time_counter.units = 'non-dim'\n",
    "    time_counter[0:12] = range(1,13)\n",
    "    # runoff\n",
    "    rorunoff = nemo.createVariable('rorunoff', 'float32', ('time_counter','y','x'), zlib=True)\n",
    "    rorunoff._Fillvalue = 0.\n",
    "    rorunoff._missing_value = 0.\n",
    "    rorunoff._units = 'kg m-2 s-1'\n",
    "    rorunoff[0:12,:] = runoff\n",
    "    # depth\n",
    "    rodepth = nemo.createVariable('rodepth','float32',('y','x'),zlib=True)\n",
    "    rodepth._Fillvalue = -1.\n",
    "    rodepth.missing_value = -1.\n",
    "    rodepth.units = 'm'\n",
    "    rodepth[:] = run_depth[0,:,:]\n",
    "    # temperature\n",
    "    rotemper = nemo.createVariable('rotemper','float32',('time_counter','y','x'),zlib=True)\n",
    "    rotemper._Fillvalue = -99.\n",
    "    rotemper.missing_value = -99.\n",
    "    rotemper.units = 'deg C'\n",
    "    rotemper[0:12,:] = run_temp\n",
    "    nemo.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if rivertype == 'constant':\n",
    "\n",
    "    nemo = nc.Dataset('/ocean/jieliu/research/meopar/river-treatment/rivers_cnst_edit.nc', 'w')\n",
    "    nemo.description = 'Constant Yearly Average, All Rivers, modify on depth and runoff grid point' \n",
    "    \n",
    "    # dimensions\n",
    "    nemo.createDimension('x', 398)\n",
    "    nemo.createDimension('y', 898)\n",
    "    nemo.createDimension('time_counter', None)\n",
    "    \n",
    "    # variables\n",
    "    # latitude and longitude\n",
    "    nav_lat = nemo.createVariable('nav_lat','float32',('y','x'),zlib=True)\n",
    "    nav_lat = lat\n",
    "    x = nemo.createVariable('nav_lon','float32',('y','x'),zlib=True)\n",
    "    nav_lon = lon\n",
    "    # time\n",
    "    time_counter = nemo.createVariable('time_counter', 'float32', ('time_counter'),zlib=True)\n",
    "    time_counter.units = 'non-dim'\n",
    "    time_counter[0] = 1\n",
    "    # runoff\n",
    "    rorunoff = nemo.createVariable('rorunoff', 'float32', ('time_counter','y','x'), zlib=True)\n",
    "    rorunoff._Fillvalue = 0.\n",
    "    rorunoff._missing_value = 0.\n",
    "    rorunoff._units = 'kg m-2 s-1'\n",
    "    rorunoff[0,:] = runoff\n",
    "    # depth\n",
    "    rodepth = nemo.createVariable('rodepth','float32',('y','x'),zlib=True)\n",
    "    rodepth._Fillvalue = -1.\n",
    "    rodepth.missing_value = -1.\n",
    "    rodepth.units = 'm'\n",
    "    rodepth[:] = run_depth\n",
    "    nemo.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Re-write daily Fraser flow file from May 14, 2015-June 14, 2015"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 304,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2015-06-14T00:00:00+00:00 2015-06-14T00:00:00+00:00\n"
     ]
    }
   ],
   "source": [
    "# Constant and data ranges etc\n",
    "year = 2015\n",
    "smonth = 06\n",
    "emonth = 06\n",
    "startdate = arrow.get(year,smonth,14)\n",
    "enddate = arrow.get(year,emonth,14)\n",
    "print startdate, enddate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 2013.       10.       17.     1709.085]\n",
      " [ 2013.       10.       18.     1676.078]\n",
      " [ 2013.       10.       19.     1653.68 ]\n",
      " ..., \n",
      " [ 2015.        6.       28.     4622.91 ]\n",
      " [ 2015.        6.       29.     4586.816]\n",
      " [ 2015.        6.       30.     4717.231]]\n"
     ]
    }
   ],
   "source": [
    "# get Fraser Flow data\n",
    "filename = '/data/dlatorne/SOG-projects/SOG-forcing/ECget/Fraser_flow'\n",
    "fraserflow = np.loadtxt(filename)\n",
    "print fraserflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 306,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fraser has 10 rivers\n"
     ]
    }
   ],
   "source": [
    "#Fraser watershed\n",
    "pd = rivertools.get_watershed_prop_dict('fraser')\n",
    "totalfraser = (pd['Fraser1']['prop'] + pd['Fraser2']['prop'] + \n",
    "               pd['Fraser3']['prop'] + pd['Fraser4']['prop'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 307,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "418 397\n"
     ]
    }
   ],
   "source": [
    "# Climatology, Fraser Watershed\n",
    "fluxfile = nc.Dataset('/ocean/jieliu/research/meopar/nemo-forcing/rivers/Salish_allrivers_monthly.nc','r')\n",
    "climFraserWaterShed = fluxfile.variables['fraser'][:]\n",
    "# Fraser River at Hope Seasonal Climatology (found in matlab using Mark's mean daily data)\n",
    "climFraseratHope = (931, 878, 866, 1814, 4097, 6970, 5538, 3539, 2372, 1937, 1595, 1119)\n",
    "NonHope = climFraserWaterShed - climFraseratHope\n",
    "otherratio = 0.016\n",
    "fraserratio = 1-otherratio\n",
    "\n",
    "nonFraser = (otherratio * climFraserWaterShed.sum()/NonHope.sum()) * NonHope\n",
    "afterHope = NonHope - nonFraser\n",
    "print pd['Fraser1']['i'],pd['Fraser1']['j']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 308,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_daily_flow(r,criverflow):\n",
    "    '''interpolate the daily values from the monthly values'''\n",
    "    print r.day, r.month\n",
    "    if r.day < 16:\n",
    "        prevmonth = r.month-1\n",
    "        if prevmonth == 0:\n",
    "            prevmonth = 12\n",
    "        nextmonth = r.month\n",
    "    else:\n",
    "        prevmonth = r.month\n",
    "        nextmonth = r.month + 1\n",
    "        if nextmonth == 13:\n",
    "            nextmonth = 1\n",
    "    fp = r - arrow.get(year,prevmonth,15)\n",
    "    fn = arrow.get(year,nextmonth,15) - r\n",
    "    ft = fp+fn\n",
    "    fp = fp.days/ft.days\n",
    "    fn = fn.days/ft.days\n",
    "    print ft, fp, fn\n",
    "    driverflow = fn*criverflow[prevmonth-1] + fp*criverflow[nextmonth-1]\n",
    "    return driverflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 309,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def write_file(r,flow,lat,lon,riverdepth):\n",
    "    ''' given the flow and the riverdepth and the date, write the nc file'''\n",
    "    directory = '.'\n",
    "    # set up filename to follow NEMO conventions\n",
    "    filename = 'RFraserCElse_y'+str(year)+'m'+'{:0=2}'.format(r.month)+'d'+'{:0=2}'.format(r.day)+'.nc'\n",
    "   # print directory+'/'+filename\n",
    "    nemo = nc.Dataset(directory+'/'+filename, 'w')\n",
    "    nemo.description = 'Real Fraser Values, Daily Climatology for Other Rivers' \n",
    "    \n",
    "    # dimensions\n",
    "    ymax, xmax = lat.shape\n",
    "    nemo.createDimension('x', xmax)\n",
    "    nemo.createDimension('y', ymax)\n",
    "    nemo.createDimension('time_counter', None)\n",
    "    \n",
    "    # variables\n",
    "    # latitude and longitude\n",
    "    nav_lat = nemo.createVariable('nav_lat','float32',('y','x'),zlib=True)\n",
    "    nav_lat = lat\n",
    "    x = nemo.createVariable('nav_lon','float32',('y','x'),zlib=True)\n",
    "    nav_lon = lon\n",
    "    # time\n",
    "    time_counter = nemo.createVariable('time_counter', 'float32', ('time_counter'),zlib=True)\n",
    "    time_counter.units = 'non-dim'\n",
    "    time_counter[0:1] = range(1,2)\n",
    "    # runoff\n",
    "    rorunoff = nemo.createVariable('rorunoff', 'float32', ('time_counter','y','x'), zlib=True)\n",
    "    rorunoff._Fillvalue = 0.\n",
    "    rorunoff._missing_value = 0.\n",
    "    rorunoff._units = 'kg m-2 s-1'\n",
    "    rorunoff[0,:] = flow\n",
    "    # depth\n",
    "    rodepth = nemo.createVariable('rodepth','float32',('y','x'),zlib=True)\n",
    "    rodepth._Fillvalue = -1.\n",
    "    rodepth.missing_value = -1.\n",
    "    rodepth.units = 'm'\n",
    "    rodepth = riverdepth\n",
    "    nemo.close()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 310,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def fraser_correction(pd, fraserflux, r, afterHope, NonFraser, fraserratio, otherratio,\n",
    "                      runoff):\n",
    "    ''' for the Fraser Basin only, replace basic values with the new climatology after Hope and the\n",
    "     observed values for Hope.  Note, we are changing runoff only and not using/changing river\n",
    "     depth '''\n",
    "    for key in pd:\n",
    "        if \"Fraser\" in key:\n",
    "            flux = calculate_daily_flow(r,afterHope) + fraserflux\n",
    "            subarea = fraserratio\n",
    "        else:\n",
    "            flux = calculate_daily_flow(r,NonFraser)\n",
    "            subarea = otherratio\n",
    "        \n",
    "        river = pd[key]\n",
    "        runoff = rivertools.fill_runoff_array(flux*river['prop']/subarea,river['i'],\n",
    "                          river['di'],river['j'],river['dj'],river['depth'],\n",
    "                          runoff,np.empty_like(runoff))[0]\n",
    "    return runoff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# load the re-written climatology river files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 311,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "##open climatolgy file with modified fresh water point source \n",
    "clim_rivers_edit = nc.Dataset('rivers_month_edit.nc','r' )\n",
    "criverflow_edit = clim_rivers_edit.variables['rorunoff']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 312,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2015-06-14T00:00:00+00:00\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "14 6\n",
      "31 days, 0:00:00 0.967741935484 0.0322580645161\n",
      "[ 14.77230835  10.16569614   8.63957405] 26.5695\n",
      "27.8402 26.5695\n",
      "0.0237942 0.0143087\n",
      "0.203718 0.203718\n"
     ]
    }
   ],
   "source": [
    "for r in arrow.Arrow.range('day', startdate, enddate):\n",
    "    print r\n",
    "    driverflow = calculate_daily_flow(r, criverflow_edit)\n",
    "    storeflow = calculate_daily_flow(r, criverflow_edit)\n",
    "    step1 = fraserflow[fraserflow[:,0] == r.year]\n",
    "    step2 = step1[step1[:,1] == r.month]\n",
    "    step3 = step2[step2[:,2] == r.day]\n",
    "#    print r.year, r.month, r.day, step3[0,3]\n",
    "    runoff = fraser_correction(pd, step3[0,3] , r, afterHope, nonFraser, fraserratio, otherratio,\n",
    "                              driverflow)\n",
    "    write_file(r,runoff,lat,lon,riverdepth)\n",
    "ig = 418\n",
    "jg = 397\n",
    "print criverflow_edit[7:10,418,397], driverflow[ig,jg]\n",
    "print storeflow[ig,jg], driverflow[ig,jg]\n",
    "ig = 351; jg = 345\n",
    "print storeflow[ig,jg], driverflow[ig,jg]\n",
    "ig = 749; jg=123\n",
    "print storeflow[ig,jg], driverflow[ig,jg]\n",
    "\n",
    "# jan 0, feb 1, mar 2, apr 3, may 4, jun 5\n",
    "# jul 6, aug 7, sep 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}