Florence analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hurricane Florence PyDDA retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "## You are using the Python ARM Radar Toolkit (Py-ART), an open source\n",
      "## library for working with weather radar data. Py-ART is partly\n",
      "## supported by the U.S. Department of Energy as part of the Atmospheric\n",
      "## Radiation Measurement (ARM) Climate Research Facility, an Office of\n",
      "## Science user facility.\n",
      "##\n",
      "## If you use this software to prepare a publication, please cite:\n",
      "##\n",
      "##     JJ Helmus and SM Collis, JORS 2016, doi: 10.5334/jors.119\n",
      "\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import glob\n",
    "import pyart\n",
    "import pydda\n",
    "import datetime\n",
    "import cartopy.crs as ccrs\n",
    "import os\n",
    "import dask.bag as db\n",
    "import gc\n",
    "from boto.s3.connection import S3Connection\n",
    "\n",
    "from scipy.interpolate import NearestNDInterpolator, LinearNDInterpolator\n",
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "florence_path = '/lcrc/group/earthscience/radar/florance/'\n",
    "ltx_list = sorted(glob.glob(florence_path + '/**/KLTX*V06.ar2v', recursive=True))\n",
    "mhx_list = sorted(glob.glob(florence_path + '/**/KMHX*V06.ar2v', recursive=True))\n",
    "\n",
    "other_rads_path = '/lcrc/group/earthscience/rjackson/florence/'\n",
    "cae_list = sorted(glob.glob(other_rads_path + '/**/KLTX*V06', recursive=True))\n",
    "clx_list = sorted(glob.glob(other_rads_path + '/**/KCLX*V06', recursive=True))\n",
    "fcx_list = sorted(glob.glob(other_rads_path + '/**/KFCX*V06', recursive=True))\n",
    "gsp_list = sorted(glob.glob(other_rads_path + '/**/KGSP*V06', recursive=True))\n",
    "rax_list = sorted(glob.glob(other_rads_path + '/**/KRAX*V06', recursive=True))\n",
    "\n",
    "print(len(gsp_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mhx_list[0])\n",
    "\n",
    "def parse_dt(file_path):\n",
    "    return (datetime.datetime.strptime(file_path[-24:-9], '%Y%m%d_%H%M%S'))\n",
    "\n",
    "def parse_dt_no_ext(file_path):\n",
    "    return (datetime.datetime.strptime(file_path[-19:-4], '%Y%m%d_%H%M%S'))\n",
    "mhx_times = np.array([parse_dt(x) for x in ltx_list])\n",
    "ltx_times = np.array([parse_dt(x) for x in mhx_list])\n",
    "cae_times = np.array([parse_dt_no_ext(x) for x in cae_list])\n",
    "clx_times = np.array([parse_dt_no_ext(x) for x in clx_list])\n",
    "fcx_times = np.array([parse_dt_no_ext(x) for x in fcx_list])\n",
    "gsp_times = np.array([parse_dt_no_ext(x) for x in gsp_list])\n",
    "rax_times = np.array([parse_dt_no_ext(x) for x in rax_list])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "the_time = datetime.datetime(2018,9,14,6,50)\n",
    "the_ind_mhx = np.argmin(np.abs(mhx_times-the_time))\n",
    "the_ind_ltx = np.argmin(np.abs(ltx_times-the_time))\n",
    "mhx_radar = pyart.io.read(ltx_list[the_ind_mhx])\n",
    "ltx_radar = pyart.io.read(mhx_list[the_ind_ltx])\n",
    "gf_mhx = pyart.filters.GateFilter(mhx_radar)\n",
    "gf_mhx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "gf_mhx.exclude_below('reflectivity', -20)\n",
    "gf_ltx = pyart.filters.GateFilter(ltx_radar)\n",
    "gf_ltx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "gf_ltx.exclude_below('reflectivity', -20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_mhgx = pyart.graph.RadarMapDisplay(mhx_radar)\n",
    "display_mhgx.plot_ppi_map('reflectivity', resolution='l', gatefilter=gf_mhx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_ltx = pyart.graph.RadarMapDisplay(ltx_radar)\n",
    "display_ltx.plot_ppi_map('reflectivity', resolution='l', gatefilter=gf_ltx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9,9))\n",
    "display_mhgx.plot_ppi_map('velocity', sweep=1, resolution='l')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9,9))\n",
    "display_ltx = pyart.graph.RadarMapDisplay(ltx_radar)\n",
    "display_ltx.plot_ppi_map('velocity', sweep=1, resolution='l')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9,9))\n",
    "dealiased_vel_mhx = pyart.correct.dealias_region_based(mhx_radar)\n",
    "mhx_radar.add_field('corrected_velocity', dealiased_vel_mhx, replace_existing=True)\n",
    "display_mhgx.plot_ppi_map('corrected_velocity', sweep=1, resolution='l')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9,9))\n",
    "dealiased_vel_ltx = pyart.correct.dealias_region_based(ltx_radar)\n",
    "ltx_radar.add_field('corrected_velocity', dealiased_vel_ltx, replace_existing=True)\n",
    "display_ltx.plot_ppi_map('corrected_velocity', sweep=1, resolution='l', vmin=-100, vmax=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mhx_radar.fields.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_mhx = pyart.map.grid_from_radars(mhx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0.,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "grid_ltx = pyart.map.grid_from_radars(ltx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "pyart.io.write_grid('grid_mhx.nc', grid_mhx)\n",
    "pyart.io.write_grid('grid_ltx.nc', grid_ltx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_mhx = pyart.io.read_grid('grid_mhx.nc')\n",
    "grid_ltx = pyart.io.read_grid('grid_ltx.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,4))\n",
    "grid_disp_mhx = pyart.graph.GridMapDisplay(grid_mhx)\n",
    "grid_disp_mhx.plot_longitude_slice('corrected_velocity', lat=-78)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_disp_mhx.plot_grid('corrected_velocity', level=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid_mhx, (0.0, 0.0, 0.0))\n",
    "out_grids = pydda.retrieval.get_dd_wind_field([grid_mhx, grid_ltx], u_init, v_init, w_init, Co=1.0, Cm=100.0,\n",
    "                                             mask_outside_opt=True, vel_name='corrected_velocity')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyart.io.write_grid('grid0.nc', out_grids[0])\n",
    "pyart.io.write_grid('grid1.nc', out_grids[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_grids[0].projection_proj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_grids = [pyart.io.read_grid('grid0.nc'),\n",
    "             pyart.io.read_grid('grid1.nc')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(15,10)) \n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "out_grids[1].fields['u']['data'] = np.ma.masked_where(np.logical_or(out_grids[0].fields['corrected_velocity']['data'].mask,\n",
    "                                                                    out_grids[1].fields['corrected_velocity']['data'].mask),\n",
    "                                                      out_grids[1].fields['u']['data'])\n",
    "out_grids[1].fields['v']['data'] = np.ma.masked_where(np.logical_or(out_grids[0].fields['corrected_velocity']['data'].mask,\n",
    "                                                                    out_grids[1].fields['corrected_velocity']['data'].mask),\n",
    "                                                      out_grids[1].fields['v']['data'])\n",
    "ax = pydda.vis.plot_horiz_xsection_barbs_map(out_grids, ax=ax, bg_grid_no=-1, level=3, barb_spacing_x_km=20.0,\n",
    "                                             barb_spacing_y_km=20.0)\n",
    "\n",
    "plt.title(out_grids[0].time['units'][13:] + ' winds at 1.5 km')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_array = np.ma.stack([x.fields['reflectivity']['data'] for x in out_grids])\n",
    "plt.imshow(grid_array.max(axis=0)[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load historical HRRR data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cfgrib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hrrr_data_path = '/lcrc/group/earthscience/rjackson/florence_hrrr/20180914/hrrr.t06z.wrfprsf00.grib2'\n",
    "the_grib = cfgrib.Dataset.from_path(hrrr_data_path, filter_by_keys={'typeOfLevel': 'isobaricInhPa'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "the_grib.variables.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grb_u = the_grib.variables['u']\n",
    "grb_v = the_grib.variables['v']\n",
    "gh = the_grib.variables['gh']\n",
    "\n",
    "lat = the_grib.variables['latitude'].data[:,:]\n",
    "lon = the_grib.variables['longitude'].data[:,:]\n",
    "lon[lon > 180] = lon[lon>180]-360\n",
    "print(lon.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EARTH_MEAN_RADIUS = 6.3781e6\n",
    "gh = gh.data[:,:,:]\n",
    "height = (EARTH_MEAN_RADIUS*gh)/(EARTH_MEAN_RADIUS-gh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "the_grib.variables.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.interpolate import griddata\n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "u = grb_u.data[1,:,:]\n",
    "# need to shift data grid longitudes from (0..360) to (-180..180)\n",
    "v = grb_v.data[0,:,:]\n",
    "# need to shift data grid longitudes from (0..360) to (-180..180)\n",
    "\n",
    "r = the_grib.variables['gh']\n",
    "u = u[:,:]\n",
    "p = ax.pcolormesh(lon, lat, u, transform=ccrs.PlateCarree())\n",
    "ax.coastlines(resolution='10m')\n",
    "ax.set_xlim([-80, -75])\n",
    "ax.set_ylim([33, 36])\n",
    "plt.colorbar(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grid the HRRR data onto the analysis grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We do not need the entire box, just the radar domain\n",
    "radar_grid_lat = out_grids[0].point_latitude['data']\n",
    "radar_grid_lon = out_grids[0].point_longitude['data']\n",
    "radar_grid_alt = out_grids[0].point_z['data']\n",
    "lat_min = radar_grid_lat.min()\n",
    "lat_max = radar_grid_lat.max()\n",
    "lon_min = radar_grid_lon.min()\n",
    "lon_max = radar_grid_lon.max()\n",
    "lon_r = np.tile(lon, (height.shape[0],1,1))\n",
    "lat_r = np.tile(lat, (height.shape[0],1,1))\n",
    "lon_flattened = lon_r.flatten()\n",
    "lat_flattened = lat_r.flatten()\n",
    "height_flattened = gh.flatten()\n",
    "the_box = np.where(np.logical_and.reduce((lon_flattened >= lon_min, lat_flattened >= lat_min,\n",
    "                                          lon_flattened <= lon_max, lat_flattened <= lat_max)))[0]\n",
    "\n",
    "lon_flattened = lon_flattened[the_box]\n",
    "lat_flattened = lat_flattened[the_box]\n",
    "height_flattened = height_flattened[the_box]\n",
    "\n",
    "u_flattened = grb_u.data[:,:,:].flatten()\n",
    "u_flattened = u_flattened[the_box]\n",
    "u_interp = NearestNDInterpolator((height_flattened, lat_flattened, lon_flattened), u_flattened, rescale=True)\n",
    "u_new = u_interp(radar_grid_alt, radar_grid_lat, radar_grid_lon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(u_new.shape)\n",
    "u_dict = {'data': u_new, 'long_name': \"U from HRRR\", 'units': \"m/s\"}\n",
    "out_grids[0].add_field(\"U_hrrr\", u_dict, replace_existing=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disp = pyart.graph.GridMapDisplay(out_grids[0])\n",
    "disp.plot_grid('U_hrrr', level=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.pcolormesh(u_new[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.where(np.logical_and.reduce((lon_r.flatten() >= lon_min,\n",
    "                                          lon_r.flatten() <= lon_max)))[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(lon_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_grids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now do the retrieval with HRRR data, dude!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_mhx = pyart.io.read_grid('grid_mhx.nc')\n",
    "grid_ltx = pyart.io.read_grid('grid_ltx.nc')\n",
    "grid_mhx = pydda.initialization.add_hrrr_constraint_to_grid(grid_mhx,\n",
    "    '/lcrc/group/earthscience/rjackson/florence_hrrr/20180914/hrrr.t06z.wrfprsf00.grib2')\n",
    "disp = pyart.graph.GridMapDisplay(grid_mhx)\n",
    "disp.plot_grid('U_hrrr', level=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid_mhx, (0.0, 0.0, 0.0))\n",
    "out_grids = pydda.retrieval.get_dd_wind_field([grid_mhx, grid_ltx], u_init, v_init, w_init, Co=10.0, Cm=50.0,\n",
    "                                              Cmod=0.0, mask_outside_opt=True, vel_name='corrected_velocity',\n",
    "                                               \n",
    "                                              )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(15,10)) \n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "#out_grids[1].fields['u']['data'] = np.ma.masked_where(np.logical_or(out_grids[0].fields['corrected_velocity']['data'].mask,\n",
    "#                                                                    out_grids[1].fields['corrected_velocity']['data'].mask),\n",
    "#                                                      out_grids[1].fields['u']['data'])\n",
    "#out_grids[1].fields['v']['data'] = np.ma.masked_where(np.logical_or(out_grids[0].fields['corrected_velocity']['data'].mask,\n",
    "#                                                                    out_grids[1].fields['corrected_velocity']['data'].mask),\n",
    "#                                                      out_grids[1].fields['v']['data'])\n",
    "ax = pydda.vis.plot_horiz_xsection_barbs_map(out_grids, ax=ax, bg_grid_no=-1, level=1, barb_spacing_x_km=20.0,\n",
    "                                             barb_spacing_y_km=20.0)\n",
    "\n",
    "plt.title(out_grids[0].time['units'][13:] + ' winds at 0.5 km')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_w = out_grids[1].fields['u']['data'][2]\n",
    "plt.contourf(max_w)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hey there, it's time for Dask-jobqueue!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask_jobqueue\n",
    "import dask.bag as db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_img_path = '/lcrc/group/earthscience/rjackson/florence_winds/png/'\n",
    "out_grid_path = '/lcrc/group/earthscience/rjackson/florence_winds/grids/'\n",
    "\n",
    "def make_retrieved_grid(the_time, ltx_list, mhx_list, do_hrrr=True):\n",
    "    out_grid_dir = (out_grid_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    out_img_dir = (out_img_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    if(not os.path.isdir((out_img_dir))):\n",
    "        os.makedirs(out_img_dir)\n",
    "    if(not os.path.isdir((out_grid_dir))):\n",
    "        os.makedirs(out_grid_dir)\n",
    "    if(do_hrrr == True):\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.png')\n",
    "    else:\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc')\n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nohrr.png')\n",
    "        \n",
    "    if(os.path.isfile(out_grid_mhx_file_path) and os.path.isfile(out_grid_ltx_file_path)):\n",
    "        return\n",
    "    \n",
    "    print(\"## Loading data...\")\n",
    "    the_ind_mhx = np.argmin(np.abs(mhx_times-the_time))\n",
    "    the_ind_ltx = np.argmin(np.abs(ltx_times-the_time))\n",
    "    if(np.abs(mhx_times[the_ind_mhx]-ltx_times[the_ind_ltx]) > datetime.timedelta(minutes=5)):\n",
    "        print(\"No simultaneous coverage!\")\n",
    "        return\n",
    "    try:\n",
    "        mhx_radar = pyart.io.read(ltx_list[the_ind_mhx])\n",
    "        ltx_radar = pyart.io.read(mhx_list[the_ind_ltx])\n",
    "    except:\n",
    "        print(str(the_time) + \" Failed!\")\n",
    "        return\n",
    "        \n",
    "    gf_mhx = pyart.filters.GateFilter(mhx_radar)\n",
    "    gf_mhx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_mhx.exclude_below('reflectivity', -20)\n",
    "    gf_ltx = pyart.filters.GateFilter(ltx_radar)\n",
    "    gf_ltx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_ltx.exclude_below('reflectivity', -20)\n",
    "\n",
    "    print(\"## Dealiasing...\")\n",
    "    # Dealias\n",
    "    try:\n",
    "        dealiased_vel_mhx = pyart.correct.dealias_region_based(mhx_radar)\n",
    "        mhx_radar.add_field('corrected_velocity', dealiased_vel_mhx, replace_existing=True) \n",
    "        dealiased_vel_ltx = pyart.correct.dealias_region_based(ltx_radar)\n",
    "        ltx_radar.add_field('corrected_velocity', dealiased_vel_ltx, replace_existing=True)  \n",
    "    except KeyError:\n",
    "        print(\"No velocity information available!\")\n",
    "        return\n",
    "    \n",
    "    print(\"## Gridding...\")\n",
    "    # Grid\n",
    "    grid_mhx = pyart.map.grid_from_radars(mhx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0.,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    grid_ltx = pyart.map.grid_from_radars(ltx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    # Get HRRR data from nearest hour\n",
    "    if(do_hrrr == True):\n",
    "        print(\"## Processing HRRR data...\")\n",
    "        hrrr_date = datetime.datetime(the_time.year, the_time.month, the_time.day, the_time.hour)\n",
    "        if(the_time.minute > 30):\n",
    "            hrrr_date += datetime.timedelta(hours=1)\n",
    "    \n",
    "        hrrr_path = ('/lcrc/group/earthscience/rjackson/florence_hrrr/' + \n",
    "                     \"%04d\" % hrrr_date.year +\n",
    "                     \"%02d\" % hrrr_date.month +\n",
    "                     \"%02d\" % hrrr_date.day +\n",
    "                     '/hrrr.t' + \"%02d\" % hrrr_date.hour  + 'z.wrfprsf00.grib2')\n",
    "        grid_mhx = pydda.initialization.add_hrrr_constraint_to_grid(grid_mhx,\n",
    "            hrrr_path)\n",
    "        Cmod = 5e-6\n",
    "        model_fields=[\"hrrr\"]\n",
    "    else:\n",
    "        Cmod = 0.0\n",
    "        model_fields=None\n",
    "        \n",
    "    print(\"## Running PyDDA...\")\n",
    "    u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid_mhx, (0.0, 0.0, 0.0))\n",
    "    out_grids = pydda.retrieval.get_dd_wind_field([grid_mhx, grid_ltx], u_init, v_init, w_init, Co=10.0, Cm=50.0,\n",
    "                                              Cmod=Cmod, mask_outside_opt=True, vel_name='corrected_velocity',\n",
    "                                              model_fields=model_fields\n",
    "                                              )\n",
    "    print('## Making plot..')\n",
    "    fig = plt.figure(figsize=(15,10)) \n",
    "    ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ax = pydda.vis.plot_horiz_xsection_barbs_map(out_grids, ax=ax, bg_grid_no=-1, level=1, barb_spacing_x_km=20.0,\n",
    "                                             barb_spacing_y_km=20.0)\n",
    "\n",
    "    plt.title(out_grids[0].time['units'][13:] + ' winds at 0.5 km')\n",
    "    print(\"## Saving plot...\")\n",
    "       \n",
    "    plt.savefig(out_img_file_path)\n",
    " \n",
    "    pyart.io.write_grid(out_grid_mhx_file_path, out_grids[0])\n",
    "    pyart.io.write_grid(out_grid_ltx_file_path, out_grids[1])\n",
    "    del out_grids, grid_mhx, grid_ltx, mhx_radar, ltx_radar, u_init, v_init, w_init\n",
    "    gc.collect()\n",
    "    \n",
    "    \n",
    "def make_retrieved_grid_only_hrrr(the_time, ltx_list, mhx_list, do_hrrr=True):\n",
    "    out_grid_dir = (out_grid_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    out_img_dir = (out_img_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    if(not os.path.isdir((out_img_dir))):\n",
    "        os.makedirs(out_img_dir)\n",
    "    if(not os.path.isdir((out_grid_dir))):\n",
    "        os.makedirs(out_grid_dir)\n",
    "    if(do_hrrr == True):\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.onlyhrrr.nc') \n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.onlyhrrr.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.onlyhrrr.png')\n",
    "    else:\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'hrrronly.nc')\n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltx' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'hrrronly.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nohrr.png')\n",
    "        \n",
    "    if(os.path.isfile(out_grid_mhx_file_path) and os.path.isfile(out_grid_ltx_file_path)):\n",
    "        return\n",
    "    \n",
    "    print(\"## Loading data...\")\n",
    "    the_ind_mhx = np.argmin(np.abs(mhx_times-the_time))\n",
    "    the_ind_ltx = np.argmin(np.abs(ltx_times-the_time))\n",
    "    if(np.abs(mhx_times[the_ind_mhx]-ltx_times[the_ind_ltx]) > datetime.timedelta(minutes=5)):\n",
    "        print(\"No simultaneous coverage!\")\n",
    "        return\n",
    "    try:\n",
    "        mhx_radar = pyart.io.read(ltx_list[the_ind_mhx])\n",
    "        ltx_radar = pyart.io.read(mhx_list[the_ind_ltx])\n",
    "    except:\n",
    "        print(str(the_time) + \" Failed!\")\n",
    "        return\n",
    "        \n",
    "    gf_mhx = pyart.filters.GateFilter(mhx_radar)\n",
    "    gf_mhx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_mhx.exclude_below('reflectivity', -20)\n",
    "    gf_ltx = pyart.filters.GateFilter(ltx_radar)\n",
    "    gf_ltx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_ltx.exclude_below('reflectivity', -20)\n",
    "\n",
    "    print(\"## Dealiasing...\")\n",
    "    # Dealias\n",
    "    try:\n",
    "        dealiased_vel_mhx = pyart.correct.dealias_region_based(mhx_radar)\n",
    "        mhx_radar.add_field('corrected_velocity', dealiased_vel_mhx, replace_existing=True) \n",
    "        dealiased_vel_ltx = pyart.correct.dealias_region_based(ltx_radar)\n",
    "        ltx_radar.add_field('corrected_velocity', dealiased_vel_ltx, replace_existing=True)  \n",
    "    except KeyError:\n",
    "        print(\"No velocity information available!\")\n",
    "        return\n",
    "    \n",
    "    print(\"## Gridding...\")\n",
    "    # Grid\n",
    "    grid_mhx = pyart.map.grid_from_radars(mhx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0.,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    grid_ltx = pyart.map.grid_from_radars(ltx_radar,(31,351,401),\n",
    "                   ((0.,15000.),(-100000.,200000.),(-150000.,300000.)),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    # Get HRRR data from nearest hour\n",
    "    if(do_hrrr == True):\n",
    "        print(\"## Processing HRRR data...\")\n",
    "        hrrr_date = datetime.datetime(the_time.year, the_time.month, the_time.day, the_time.hour)\n",
    "        if(the_time.minute > 30):\n",
    "            hrrr_date += datetime.timedelta(hours=1)\n",
    "    \n",
    "        hrrr_path = ('/lcrc/group/earthscience/rjackson/florence_hrrr/' + \n",
    "                     \"%04d\" % hrrr_date.year +\n",
    "                     \"%02d\" % hrrr_date.month +\n",
    "                     \"%02d\" % hrrr_date.day +\n",
    "                     '/hrrr.t' + \"%02d\" % hrrr_date.hour  + 'z.wrfprsf00.grib2')\n",
    "        grid_mhx = pydda.initialization.add_hrrr_constraint_to_grid(grid_mhx,\n",
    "            hrrr_path)\n",
    "        Cmod = 5e-6\n",
    "        model_fields=[\"hrrr\"]\n",
    "    else:\n",
    "        Cmod = 0.0\n",
    "        model_fields=None\n",
    "        \n",
    "    print(\"## Running PyDDA...\")\n",
    "    u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid_mhx, (0.0, 0.0, 0.0))\n",
    "    out_grids = pydda.retrieval.get_dd_wind_field([grid_mhx, grid_ltx], u_init, v_init, w_init, Co=0.0, Cm=0.0,\n",
    "                                              Cmod=1e-3, mask_outside_opt=True, vel_name='corrected_velocity',\n",
    "                                              model_fields=model_fields\n",
    "                                              )\n",
    "    print('## Making plot..')\n",
    "    fig = plt.figure(figsize=(15,10)) \n",
    "    ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ax = pydda.vis.plot_horiz_xsection_barbs_map(out_grids, ax=ax, bg_grid_no=-1, level=1, barb_spacing_x_km=20.0,\n",
    "                                             barb_spacing_y_km=20.0)\n",
    "\n",
    "    plt.title(out_grids[0].time['units'][13:] + ' winds at 0.5 km')\n",
    "    print(\"## Saving plot...\")\n",
    "       \n",
    "    plt.savefig(out_img_file_path)\n",
    " \n",
    "    pyart.io.write_grid(out_grid_mhx_file_path, out_grids[0])\n",
    "    pyart.io.write_grid(out_grid_ltx_file_path, out_grids[1])\n",
    "    del out_grids, grid_mhx, grid_ltx, mhx_radar, ltx_radar, u_init, v_init, w_init\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_retrieved_grid(ltx_times[241], ltx_list, mhx_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(mhx_times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_jobqueue import SLURMCluster\n",
    "cluster = SLURMCluster(cores=2, project='rainfall', walltime='2:00:00', \n",
    "                       processes=2, memory='128GB')\n",
    "\n",
    "cluster.scale(8)         # Ask for ten workers\n",
    "\n",
    "from dask.distributed import Client\n",
    "client = Client(cluster)  # Connect this local process to remote workers\n",
    "\n",
    "# wait for jobs to arrive, depending on the queue, this may take some time\n",
    "\n",
    "import dask.array as da\n",
    "from distributed import wait"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.stop_all_jobs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_inds = np.where(np.logical_and(mhx_times >= datetime.datetime(2018, 9, 14, 8, 0, 1), \n",
    "                                    mhx_times <= datetime.datetime(2018, 9, 14, 10, 0, 1)))[0]\n",
    "make_grid = lambda x: make_retrieved_grid(x, ltx_list, mhx_list, True)\n",
    "make_grid_no_hrrr = lambda x: make_retrieved_grid(x, ltx_list, mhx_list, False)\n",
    "futures = client.map(make_grid, ltx_times[229:426])\n",
    "wait(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_grid = lambda x: make_retrieved_grid(x, ltx_list, mhx_list, True)\n",
    "make_grid_no_hrrr = lambda x: make_retrieved_grid(x, ltx_list, mhx_list, False)\n",
    "make_grid_hrrr_only = lambda x: make_retrieved_grid_only_hrrr(x, ltx_list, mhx_list, True)\n",
    "make_grid_hrrr_only(ltx_times[230])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import correlate2d\n",
    "grid_all_ltx = pyart.io.read_grid('/lcrc/group/earthscience/rjackson/florence_winds/grids/20180914/05kmwinds_gridltx20180914.0624.nc')\n",
    "grid_all_mtx = pyart.io.read_grid('/lcrc/group/earthscience/rjackson/florence_winds/grids/20180914/05kmwinds_gridmhx20180914.0624.nc')\n",
    "grid_only_hrrr  = pyart.io.read_grid('/lcrc/group/earthscience/rjackson/florence_winds/grids/20180914/05kmwinds_gridltx20180914.0624.onlyhrrr.nc')\n",
    "\n",
    "correlation_u = correlate2d(grid_all.fields[\"u\"][\"data\"][1], grid_only_hrrr.fields[\"u\"][\"data\"][1])\n",
    "correlation_v = correlate2d(grid_all.fields[\"v\"][\"data\"][1], grid_only_hrrr.fields[\"v\"][\"data\"][1])\n",
    "\n",
    "plt.pcolormesh(correlation_u)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ltx_list = sorted(glob.glob('/lcrc/group/earthscience/rjackson/florence_winds/grids/20180914/*ltx*.nc'))\n",
    "mhx_list = sorted(glob.glob('/lcrc/group/earthscience/rjackson/florence_winds/grids/20180914/*mhx*.nc'))\n",
    "from copy import deepcopy\n",
    "\n",
    "for i in range(len(ltx_list)):\n",
    "    ltx_grid = pyart.io.read_grid(ltx_list[i])\n",
    "    mhx_grid = pyart.io.read_grid(mhx_list[i])\n",
    "    fig = plt.figure(figsize=(35, 20)) \n",
    "    font = {'family' : 'normal',\n",
    "            'weight' : 'bold',\n",
    "            'size'   : 44}\n",
    "\n",
    "    plt.rc('font', **font)\n",
    "    ltx_grid.fields[\"rainfall_rate\"] = deepcopy(ltx_grid.fields[\"reflectivity\"])\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"standard_name\"] = \"rainfall_rate\"\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"long_name\"] = \"rainfall rate\"\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"units\"] = \"mm hr-1\"\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"data\"] = (10**(ltx_grid.fields[\"reflectivity\"][\"data\"]/10)/300)**(1/1.4)\n",
    "    \n",
    "    mhx_grid.fields[\"rainfall_rate\"] = deepcopy(mhx_grid.fields[\"reflectivity\"])\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"standard_name\"] = \"rainfall_rate\"\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"long_name\"] = \"rainfall rate\"\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"units\"] = \"mm hr-1\"\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"data\"] = (10**(mhx_grid.fields[\"reflectivity\"][\"data\"]/10)/300)**(1/1.4)\n",
    "    \n",
    "    the_mask = np.logical_and(ltx_grid.fields[\"rainfall_rate\"][\"data\"].mask,\n",
    "                              mhx_grid.fields[\"rainfall_rate\"][\"data\"].mask)\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"data\"] = ltx_grid.fields[\"rainfall_rate\"][\"data\"].filled(0)\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"data\"] = mhx_grid.fields[\"rainfall_rate\"][\"data\"].filled(0)\n",
    "    ltx_grid.fields[\"rainfall_rate\"][\"data\"] = np.ma.masked_where(the_mask, \n",
    "                                                                  ltx_grid.fields[\"rainfall_rate\"][\"data\"])\n",
    "    mhx_grid.fields[\"rainfall_rate\"][\"data\"] = np.ma.masked_where(the_mask, \n",
    "                                                                  mhx_grid.fields[\"rainfall_rate\"][\"data\"])\n",
    "    ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ax = pydda.vis.plot_horiz_xsection_streamlines_map([ltx_grid, mhx_grid], ax=ax, \n",
    "                                                       background_field='rainfall_rate',\n",
    "                                                       bg_grid_no=-1, level=2, \n",
    "                                                       vmin=0, vmax=50, show_lobes=False)\n",
    "    wind_speed = np.sqrt(ltx_grid.fields[\"u\"][\"data\"]**2 + ltx_grid.fields[\"v\"][\"data\"]**2)\n",
    "    wind_speed = wind_speed.filled(np.nan)\n",
    "    lons = ltx_grid.point_longitude[\"data\"]\n",
    "    lats = ltx_grid.point_latitude[\"data\"]\n",
    "    cs = ax.contour(lons[2, ::4, ::4], lats[2, ::4, ::4], wind_speed[2, ::4, ::4], levels=[28, 32], \n",
    "                   linewidths=8, colors=['b', 'r', 'k'])\n",
    "    plt.clabel(cs, ax=ax, inline=1, fontsize=15)\n",
    "    ax.set_xticks(np.arange(-80, -75, 0.5))\n",
    "    ax.set_yticks(np.arange(33, 35.8, 0.5))\n",
    "    ax.set_title(ltx_grid.time[\"units\"][-20:])\n",
    "    plt.savefig(ltx_list[i][-15:] + '.png')\n",
    "    del fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ltx_list = sorted(glob.glob('/lcrc/group/earthscience/rjackson/ddop_grids/grids/2006/20060120/berrwinds*'))\n",
    "mhx_list = sorted(glob.glob('/lcrc/group/earthscience/rjackson/ddop_grids/grids/2006/20060120/cpolwinds*'))\n",
    "ltx_grid = pyart.io.read_grid(ltx_list[0])\n",
    "mhx_grid = pyart.io.read_grid(mhx_list[0])\n",
    "fig = plt.figure(figsize=(10,5)) \n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "pydda.vis.plot_horiz_xsection_barbs_map([ltx_grid, mhx_grid], ax=ax, bg_grid_no=-1, level=2, barb_spacing_x_km=5.0,\n",
    "                                            barb_spacing_y_km=5.0, vmin=0, vmax=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pydda.vis.plot_horiz_xsection_barbs_map?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download NEXRAD level 2 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#first lets connect to the bucket\n",
    "conn = S3Connection(anon = True)\n",
    "bucket = conn.get_bucket('noaa-nexrad-level2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_pref = '2018/09/14/KRAX/'\n",
    "bucket_list = list(bucket.list(prefix = my_pref))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kvnx_download_path = '/lcrc/group/earthscience/rjackson/florence/'\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(bucket_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for item in bucket_list:\n",
    "    item.get_contents_to_filename(os.path.join(kvnx_download_path,item.key))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_img_path = '/lcrc/group/earthscience/rjackson/florence_winds/png/'\n",
    "out_grid_path = '/lcrc/group/earthscience/rjackson/florence_winds/grids/'\n",
    "\n",
    "from scipy.interpolate import griddata\n",
    "from copy import deepcopy\n",
    "import os\n",
    "\n",
    "def reduce_pyart_grid_res(Grid, skip_factor):\n",
    "    Grid2 = deepcopy(Grid)\n",
    "    field_dict = {}\n",
    "    for field_name in Grid2.fields.keys():\n",
    "        field_dict[field_name] = Grid2.fields[field_name].copy()\n",
    "        field_dict[field_name][\"data\"] = Grid2.fields[field_name][\"data\"][:, ::skip_factor, ::skip_factor]\n",
    "        \n",
    "    x = Grid2.x\n",
    "    x[\"data\"] = x[\"data\"][::skip_factor]\n",
    "    y = Grid2.y\n",
    "    y[\"data\"] = y[\"data\"][::skip_factor]\n",
    "    z = Grid2.z\n",
    "    z[\"data\"] = z[\"data\"]\n",
    "    metadata = Grid2.metadata\n",
    "    origin_latitude = Grid2.origin_latitude\n",
    "    origin_longitude = Grid2.origin_longitude\n",
    "    origin_altitude = Grid2.origin_altitude\n",
    "    projection = Grid2.projection\n",
    "    radar_latitude = Grid2.radar_latitude\n",
    "    radar_longitude = Grid2.radar_longitude\n",
    "    radar_altitude = Grid2.radar_altitude\n",
    "    radar_time = Grid2.radar_time\n",
    "    radar_name = Grid2.radar_name\n",
    "    gtime = Grid2.time\n",
    "    new_grid = pyart.core.Grid(gtime, field_dict, metadata, origin_latitude, origin_longitude, origin_altitude, \n",
    "                               x, y, z, projection, radar_latitude, radar_longitude, radar_altitude, \n",
    "                               radar_time, radar_name) \n",
    "    del Grid2\n",
    "    return new_grid\n",
    "\n",
    "def split_pyart_grid(Grid, split_factor, axis=1):\n",
    "    grid_splits = []\n",
    "    split_field = {}\n",
    "    Grid2 = deepcopy(Grid)\n",
    "    for field_name in Grid2.fields.keys():\n",
    "        if isinstance(Grid2.fields[field_name][\"data\"], np.ma.MaskedArray):\n",
    "            no_mask = Grid2.fields[field_name][\"data\"].filled(np.nan).copy() \n",
    "        else:\n",
    "            no_mask = Grid2.fields[field_name][\"data\"].copy()\n",
    "        split_field[field_name] = np.array_split(no_mask, split_factor, axis=axis)\n",
    "        if isinstance(Grid2.fields[field_name][\"data\"], np.ma.MaskedArray):\n",
    "            split_field[field_name] = [np.ma.masked_where(\n",
    "                np.isnan(arr), arr) for arr in split_field[field_name]]\n",
    "    x = Grid2.x\n",
    "    y = Grid2.y\n",
    "    z = Grid2.z\n",
    "    x_split = np.array_split(x[\"data\"], split_factor)\n",
    "    y_split = np.array_split(y[\"data\"], split_factor)\n",
    "    z_split = np.array_split(z[\"data\"], split_factor)\n",
    "    gtime = Grid2.time\n",
    "    metadata = Grid2.metadata\n",
    "    origin_latitude = Grid2.origin_latitude\n",
    "    origin_longitude = Grid2.origin_longitude\n",
    "    origin_altitude = Grid2.origin_altitude\n",
    "    projection = Grid2.projection\n",
    "    radar_latitude = Grid2.radar_latitude\n",
    "    radar_longitude = Grid2.radar_longitude\n",
    "    radar_altitude = Grid2.radar_altitude\n",
    "    radar_time = Grid2.radar_time\n",
    "    radar_name = Grid2.radar_name\n",
    "    for i in range(split_factor):\n",
    "        grid_dic = {}\n",
    "\n",
    "        for field_name in Grid2.fields.keys():\n",
    "            grid_dic[field_name] = Grid2.fields[field_name].copy()\n",
    "            grid_dic[field_name][\"data\"] = split_field[field_name][i]\n",
    "        x_dic = x.copy()\n",
    "        y_dic = y.copy()\n",
    "        z_dic = z.copy()\n",
    "        if(axis == 1):\n",
    "            y_dic[\"data\"] = y_split[i]\n",
    "        elif(axis == 2):\n",
    "            x_dic[\"data\"] = x_split[i]\n",
    "        elif(axis == 0):\n",
    "            z_dic[\"data\"] = z_split[i]\n",
    "        \n",
    "        new_grid = pyart.core.Grid(gtime, grid_dic, metadata, origin_latitude, origin_longitude, origin_altitude, \n",
    "                               x_dic, y_dic, z_dic, projection, radar_latitude, radar_longitude, radar_altitude, \n",
    "                               radar_time, radar_name) \n",
    "        grid_splits.append(new_grid)\n",
    "        \n",
    "        \n",
    "    return grid_splits\n",
    "\n",
    "def concatenate_pyart_grids(grid_list, axis=1):\n",
    "    new_grid = deepcopy(grid_list[0])\n",
    "    for field_name in new_grid.fields.keys():\n",
    "        new_grid.fields[field_name][\"data\"] = np.ma.concatenate([x.fields[field_name][\"data\"] for x in grid_list], axis=axis)\n",
    "    if(axis == 2):\n",
    "        new_grid.x[\"data\"] = np.ma.concatenate([x.x[\"data\"] for x in grid_list])\n",
    "        new_grid.nx = np.sum([x.nx for x in grid_list])\n",
    "    elif(axis == 1):\n",
    "        new_grid.y[\"data\"] = np.ma.concatenate([x.y[\"data\"] for x in grid_list])\n",
    "        new_grid.ny = np.sum([x.ny for x in grid_list])\n",
    "    elif(axis == 0):\n",
    "        new_grid.z[\"data\"] = np.ma.concatenate([x.z[\"data\"] for x in grid_list]) \n",
    "        new_grid.nz = np.sum([x.nz for x in grid_list])\n",
    "    return new_grid\n",
    "\n",
    "# Procedure: 1. Do first pass of retrieval on reduced resolution grid\n",
    "# 2. Then, we use the reduced resolution retrieval as an input to the\n",
    "# high resolution retrieval in each region\n",
    "# Finally, we check for continuity at the boundaries\n",
    "def do_dd_wind_field_nested(grid_list, u_init, v_init, w_init, reduction_factor=2,\n",
    "                            num_splits=2, **kwargs):\n",
    "    \"\"\"\n",
    "    This function performs a wind retrieval using a nested domain. This is useful for\n",
    "    grids that are larger than about 400 by 400 by 40 points, since the use of larger\n",
    "    grids on a single machine will exceed memory limitations. \n",
    "    \n",
    "    This procedure relies on a dask distributed cluster to be set up. The retrieval is \n",
    "    first performed at a resolution that is coarser than the analysis grid by \n",
    "    reduction_factor. This provides the initial state for the \n",
    "    \n",
    "    The domain is split into num_splits**2 sub-domains for the nested retrieval step, and\n",
    "    each nested retrieval is mapped onto a distributed worker for parallel processing. If\n",
    "    NumPy and SciPy are already set up to use parallel numerical analysis libraries, it is \n",
    "    recommended that a single machine be dedicated to each nest rather than a single core\n",
    "    for best peformance.\n",
    "    \n",
    "    Parameters\n",
    "    ==========\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    **kwargs: dict\n",
    "        This function will take the same keyword arguments as get_dd_wind_field, as these \n",
    "        arguments are passed into each call of get_dd_wind_field.\n",
    "    \n",
    "    \"\"\"\n",
    "    # First, we do retrieval on whole grid with fraction of resolution\n",
    "    grid_lo_res_list = [reduce_pyart_grid_res(G, reduction_factor) for G in grid_list]\n",
    "    \n",
    "    first_pass = pydda.retrieval.get_dd_wind_field(grid_lo_res_list, \n",
    "                                                   u_init[::, ::reduction_factor, ::reduction_factor], \n",
    "                                                   v_init[::, ::reduction_factor, ::reduction_factor], \n",
    "                                                   w_init[::, ::reduction_factor, ::reduction_factor], **kwargs)\n",
    "    \n",
    "    # Take the first pass field and regrid to analysis field\n",
    "    reduced_x = first_pass[0].point_x[\"data\"].flatten()\n",
    "    reduced_y = first_pass[0].point_y[\"data\"].flatten()\n",
    "    reduced_z = first_pass[0].point_z[\"data\"].flatten()\n",
    "    x = grid_list[0].point_x[\"data\"].flatten()\n",
    "    y = grid_list[0].point_y[\"data\"].flatten()\n",
    "    z = grid_list[0].point_z[\"data\"].flatten()\n",
    "    u_init_new = griddata((reduced_z, reduced_y, reduced_x),\n",
    "                                             first_pass[0].fields[\"u\"][\"data\"].flatten(), \n",
    "                                             (z, y, x), method='nearest')\n",
    "    v_init_new = griddata((reduced_z, reduced_y, reduced_x),\n",
    "                                             first_pass[0].fields[\"v\"][\"data\"].flatten(), \n",
    "                                             (z, y, x), method='nearest')\n",
    "    w_init_new = griddata((reduced_z, reduced_y, reduced_x),\n",
    "                                             first_pass[0].fields[\"w\"][\"data\"].flatten(), \n",
    "                                             (z, y, x), method='nearest')\n",
    "    u_init_new = np.reshape(u_init_new, u_init.shape)\n",
    "    v_init_new = np.reshape(v_init_new, v_init.shape)\n",
    "    w_init_new = np.reshape(w_init_new, w_init.shape)\n",
    "    \n",
    "    # Finally, split the analysis into num_splits**2 pieces and save\n",
    "    # as temporary files\n",
    "    tempfile_name_base = datetime.datetime.now().strftime('%y%m%d.%H%M%S')\n",
    "    tiny_grids = []\n",
    "    k = 0\n",
    "    for G in grid_list:\n",
    "        cur_list = []\n",
    "        split_grids_x = split_pyart_grid(G, num_splits, axis=2)\n",
    "        i = 0\n",
    "        for sgrid in split_grids_x:\n",
    "            g_list = split_pyart_grid(sgrid, num_splits)\n",
    "            grid_fns = []\n",
    "            j = 0\n",
    "            for g in g_list:\n",
    "                fn = tempfile_name_base + str(k) + '.' + str(i) + '.' + str(j) + '.nc'\n",
    "                pyart.io.write_grid(tempfile_name_base + str(k) + '.' + str(i) + '.' + str(j) + '.nc',\n",
    "                                    g)\n",
    "                j = j + 1\n",
    "                grid_fns.append(fn)\n",
    "            cur_list.append(grid_fns)\n",
    "            i = i + 1\n",
    "        del split_grids_x, g_list\n",
    "        \n",
    "        k = k + 1\n",
    "        tiny_grids.append(cur_list)\n",
    "    \n",
    "    # Temporarily save the tiny grids and free up memory...we want to load these when\n",
    "    # we are running it on the cluster\n",
    "    \n",
    "    u_init_split_x = np.array_split(u_init_new, num_splits, axis=2)\n",
    "    u_init_split = [np.array_split(ux, num_splits, axis=1) for ux in u_init_split_x]\n",
    "    w_init_split_x = np.array_split(w_init_new, num_splits, axis=2)\n",
    "    w_init_split = [np.array_split(wx, num_splits, axis=1) for wx in w_init_split_x]\n",
    "    v_init_split_x = np.array_split(v_init_new, num_splits, axis=2)\n",
    "    v_init_split = [np.array_split(vx, num_splits, axis=1) for vx in v_init_split_x]\n",
    "    \n",
    "    # Clear out unneeded variables (do not need lo-res grids in memory anymore)\n",
    "    del u_init_split_x, w_init_split_x, v_init_split_x\n",
    "    del first_pass, reduced_x, reduced_y, reduced_z, x, y, z, grid_lo_res_list\n",
    "    gc.collect()\n",
    "    \n",
    "    # Serial just for testing, need to use dask in future\n",
    "    tiny_retrieval = []\n",
    "    def do_tiny_retrieval(i,j):\n",
    "        \n",
    "        tgrids = [pyart.io.read_grid(tiny_grids[k][i][j]) for k in range(len(grid_list))]\n",
    "        print(tgrids)\n",
    "        new_grids = pydda.retrieval.get_dd_wind_field(tgrids, u_init_split[i][j], v_init_split[i][j], \n",
    "                                                 w_init_split[i][j], **kwargs)\n",
    "        del tgrids\n",
    "        gc.collect()\n",
    "        return new_grids\n",
    "    \n",
    "    futures_array = []\n",
    "    for i in range(num_splits):\n",
    "        for j in range(num_splits):\n",
    "            futures_array.append(client.submit(do_tiny_retrieval, i, j))\n",
    "    \n",
    "    print(\"Waiting for nested grid to be retrieved...\")\n",
    "    wait(futures_array)\n",
    "    \n",
    "    \n",
    "    tiny_retrieval2 = client.gather(futures_array)\n",
    "    \n",
    "    tiny_retrieval = []\n",
    "    \n",
    "    for i in range(num_splits):    \n",
    "        new_grid_list = []\n",
    "    \n",
    "        for j in range(len(grid_list)):\n",
    "            print([tiny_retrieval2[k+i*num_splits][j] for k in range(0, num_splits)])\n",
    "            \n",
    "            new_grid_list.append(concatenate_pyart_grids([tiny_retrieval2[k+i*num_splits][j] for k in range(0, num_splits)], \n",
    "                                                         axis=1))\n",
    "        tiny_retrieval.append(new_grid_list)\n",
    "    \n",
    "    new_grid_list = []\n",
    "    for i in range(len(grid_list)):\n",
    "        new_grid_list.append(concatenate_pyart_grids([tiny_retrieval[k][i] for k in range(num_splits)], axis=2))\n",
    "    \n",
    "    tempfile_list = glob.glob(tempfile_name_base + \"*\")\n",
    "    for fn in tempfile_list:\n",
    "        os.remove(fn)\n",
    "    return new_grid_list\n",
    "\n",
    "    # Then just tile the pieces back together\n",
    "    # Combine the split grids together into one!\n",
    "             \n",
    "def make_retrieved_grid_extended(the_time, ltx_list, mhx_list, cae_list,\n",
    "                                 clx_list, fcx_list, rax_list, gsp_list, do_hrrr=True):\n",
    "    out_grid_dir = (out_grid_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    out_img_dir = (out_img_path + '/' + \"%04d\" % the_time.year +\n",
    "                   \"%02d\" % the_time.month +\"%02d\" % the_time.day + '/')\n",
    "    if(not os.path.isdir((out_img_dir))):\n",
    "        os.makedirs(out_img_dir)\n",
    "    if(not os.path.isdir((out_grid_dir))):\n",
    "        os.makedirs(out_grid_dir)\n",
    "    if(do_hrrr == True):\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_cae_file_path = (out_grid_dir + '05kmwinds_gridcaeext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_clx_file_path = (out_grid_dir + '05kmwinds_gridclxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc')\n",
    "        out_grid_fcx_file_path = (out_grid_dir + '05kmwinds_gridfcxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_gsp_file_path = (out_grid_dir + '05kmwinds_gridgspext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_grid_rax_file_path = (out_grid_dir + '05kmwinds_gridraxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.png')\n",
    "    else:\n",
    "        out_grid_mhx_file_path = (out_grid_dir + '05kmwinds_gridmhxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc')\n",
    "        out_grid_ltx_file_path = (out_grid_dir + '05kmwinds_gridltxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_grid_cae_file_path = (out_grid_dir + '05kmwinds_gridcaeext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_grid_clx_file_path = (out_grid_dir + '05kmwinds_gridclxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc')\n",
    "        out_grid_fcx_file_path = (out_grid_dir + '05kmwinds_gridfcxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_grid_gsp_file_path = (out_grid_dir + '05kmwinds_gridgspext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_grid_rax_file_path = (out_grid_dir + '05kmwinds_gridraxext' + \"%04d\" % the_time.year + \n",
    "            \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "            \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + 'nohrrr.nc') \n",
    "        out_img_file_path = (out_img_dir + '05kmwinds' + \"%04d\" % the_time.year + \n",
    "                         \"%02d\" % the_time.month + \"%02d\" % the_time.day + '.' +\n",
    "                         \"%02d\" % the_time.hour + \"%02d\" % the_time.minute + '.nohrr.png')\n",
    "        \n",
    "    #if(os.path.isfile(out_grid_mhx_file_path) and os.path.isfile(out_grid_ltx_file_path)):\n",
    "    #    return\n",
    "    \n",
    "    print(\"## Loading data...\")\n",
    "    the_ind_mhx = np.argmin(np.abs(mhx_times-the_time))\n",
    "    the_ind_ltx = np.argmin(np.abs(ltx_times-the_time))\n",
    "    the_ind_clx = np.argmin(np.abs(clx_times-the_time))\n",
    "    the_ind_fcx = np.argmin(np.abs(fcx_times-the_time))\n",
    "    the_ind_gsp = np.argmin(np.abs(gsp_times-the_time))\n",
    "    the_ind_rax = np.argmin(np.abs(rax_times-the_time))\n",
    "\n",
    "    if(np.abs(mhx_times[the_ind_mhx]-ltx_times[the_ind_ltx]) > datetime.timedelta(minutes=5)):\n",
    "        print(\"No simultaneous coverage!\")\n",
    "        return\n",
    "    try:\n",
    "        mhx_radar = pyart.io.read(ltx_list[the_ind_mhx])\n",
    "        ltx_radar = pyart.io.read(mhx_list[the_ind_ltx])\n",
    "        clx_radar = pyart.io.read(clx_list[the_ind_clx])\n",
    "        fcx_radar = pyart.io.read(fcx_list[the_ind_fcx])\n",
    "        gsp_radar = pyart.io.read(gsp_list[the_ind_gsp])\n",
    "        rax_radar = pyart.io.read(rax_list[the_ind_rax])\n",
    "    except:\n",
    "        print(str(the_time) + \" Failed!\")\n",
    "        return\n",
    "        \n",
    "    gf_mhx = pyart.filters.GateFilter(mhx_radar)\n",
    "    gf_mhx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_mhx.exclude_below('reflectivity', -20)\n",
    "    gf_ltx = pyart.filters.GateFilter(ltx_radar)\n",
    "    gf_ltx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_ltx.exclude_below('reflectivity', -20)\n",
    "    gf_fcx = pyart.filters.GateFilter(fcx_radar)\n",
    "    gf_fcx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_fcx.exclude_below('reflectivity', -20)\n",
    "    gf_gsp = pyart.filters.GateFilter(gsp_radar)\n",
    "    gf_gsp.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_gsp.exclude_below('reflectivity', -20)\n",
    "    gf_rax = pyart.filters.GateFilter(rax_radar)\n",
    "    gf_rax.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_rax.exclude_below('reflectivity', -20)\n",
    "    gf_clx = pyart.filters.GateFilter(clx_radar)\n",
    "    gf_clx.exclude_below('cross_correlation_ratio', 0.5)\n",
    "    gf_clx.exclude_below('reflectivity', -20)\n",
    "    \n",
    "    print(\"## Dealiasing...\")\n",
    "    # Dealias\n",
    "    try:\n",
    "        dealiased_vel_mhx = pyart.correct.dealias_region_based(mhx_radar, gatefilter=gf_mhx)\n",
    "        mhx_radar.add_field('corrected_velocity', dealiased_vel_mhx, replace_existing=True) \n",
    "        dealiased_vel_ltx = pyart.correct.dealias_region_based(ltx_radar, gatefilter=gf_ltx)\n",
    "        ltx_radar.add_field('corrected_velocity', dealiased_vel_ltx, replace_existing=True)  \n",
    "        dealiased_vel_fcx = pyart.correct.dealias_region_based(fcx_radar, gatefilter=gf_fcx)\n",
    "        fcx_radar.add_field('corrected_velocity', dealiased_vel_fcx, replace_existing=True)  \n",
    "        dealiased_vel_gsp = pyart.correct.dealias_region_based(gsp_radar, gatefilter=gf_gsp)\n",
    "        gsp_radar.add_field('corrected_velocity', dealiased_vel_gsp, replace_existing=True) \n",
    "        dealiased_vel_rax = pyart.correct.dealias_region_based(rax_radar, gatefilter=gf_rax)\n",
    "        rax_radar.add_field('corrected_velocity', dealiased_vel_rax, replace_existing=True)  \n",
    "        dealiased_vel_clx = pyart.correct.dealias_region_based(clx_radar, gatefilter=gf_clx)\n",
    "        clx_radar.add_field('corrected_velocity', dealiased_vel_clx, replace_existing=True)  \n",
    "    except KeyError:\n",
    "        print(\"No velocity information available!\")\n",
    "        return\n",
    "    \n",
    "    print(\"## Gridding...\")\n",
    "    # Grid\n",
    "    grid_spec = (31, 1101, 1101)\n",
    "    grid_z = (0., 15000.)\n",
    "    grid_y = (-650000., 650000.)\n",
    "    grid_x = (-650000., 650000.)\n",
    "    grid_mhx = pyart.map.grid_from_radars(mhx_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0.,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    grid_ltx = pyart.map.grid_from_radars(ltx_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "\n",
    "    grid_fcx = pyart.map.grid_from_radars(fcx_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "    grid_gsp = pyart.map.grid_from_radars(gsp_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0.,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "    grid_rax = pyart.map.grid_from_radars(rax_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "    grid_clx = pyart.map.grid_from_radars(clx_radar,grid_spec,\n",
    "                   (grid_z, grid_y, grid_x),\n",
    "                   fields=['reflectivity','corrected_velocity'],\n",
    "                   refl_field='reflectivity',roi_func='dist_beam',\n",
    "                   h_factor=0,nb=0.6,bsp=1.,min_radius=200.,\n",
    "                   grid_origin=(mhx_radar.latitude['data'], mhx_radar.longitude['data'])\n",
    "                   )\n",
    "    \n",
    "    # Get HRRR data from nearest hour\n",
    "    if(do_hrrr == True):\n",
    "        print(\"## Processing HRRR data...\")\n",
    "        hrrr_date = datetime.datetime(the_time.year, the_time.month, the_time.day, the_time.hour)\n",
    "        if(the_time.minute > 30):\n",
    "            hrrr_date += datetime.timedelta(hours=1)\n",
    "    \n",
    "        hrrr_path = ('/lcrc/group/earthscience/rjackson/florence_hrrr/' + \n",
    "                     \"%04d\" % hrrr_date.year +\n",
    "                     \"%02d\" % hrrr_date.month +\n",
    "                     \"%02d\" % hrrr_date.day +\n",
    "                     '/hrrr.t' + \"%02d\" % hrrr_date.hour  + 'z.wrfprsf00.grib2')\n",
    "        grid_mhx = pydda.constraints.add_hrrr_constraint_to_grid(grid_mhx,\n",
    "            hrrr_path)\n",
    "        Cmod = 5e-6\n",
    "        model_fields=[\"hrrr\", \"erainterim\"]\n",
    "    else:\n",
    "        Cmod = 0.0\n",
    "        model_fields=[\"erainterim\"]\n",
    "    grid_mhx = pydda.constraints.make_constraint_from_era_interim(grid_mhx)    \n",
    "    print(\"## Running PyDDA...\")\n",
    "    u_init, v_init, w_init = pydda.initialization.make_constant_wind_field(grid_mhx, (0.0, 0.0, 0.0))\n",
    "    out_grids = pydda.retrieval.get_dd_wind_field_nested([grid_mhx, grid_ltx, grid_fcx, grid_gsp,\n",
    "                                                  grid_rax, grid_clx], \n",
    "                                                  u_init, v_init, w_init, Co=1.0, Cm=100.0,\n",
    "                                                  Cmod=Cmod, mask_outside_opt=True, vel_name='corrected_velocity',\n",
    "                                                  model_fields=model_fields, client=client)\n",
    "    print('## Making plot..')\n",
    "    print(out_grids[0].fields[\"u\"][\"data\"].shape)\n",
    "    print(out_grids[0].nz)\n",
    "    print(out_grids[0].ny)\n",
    "    print(out_grids[0].nx)\n",
    "    fig = plt.figure(figsize=(15,10)) \n",
    "    ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ax = pydda.vis.plot_horiz_xsection_barbs_map(out_grids, ax=ax, bg_grid_no=-1, level=1, barb_spacing_x_km=60.0,\n",
    "                                             barb_spacing_y_km=60.0)\n",
    "\n",
    "    plt.title(out_grids[0].time['units'][13:] + ' winds at 0.5 km')\n",
    "    print(\"## Saving plot...\")\n",
    "       \n",
    "    plt.savefig(out_img_file_path)\n",
    " \n",
    "    pyart.io.write_grid(out_grid_mhx_file_path, out_grids[0])\n",
    "    pyart.io.write_grid(out_grid_ltx_file_path, out_grids[1])\n",
    "    pyart.io.write_grid(out_grid_fcx_file_path, out_grids[2])\n",
    "    pyart.io.write_grid(out_grid_gsp_file_path, out_grids[3])\n",
    "    pyart.io.write_grid(out_grid_rax_file_path, out_grids[4])\n",
    "    pyart.io.write_grid(out_grid_clx_file_path, out_grids[5])\n",
    "\n",
    "    del out_grids, grid_mhx, grid_ltx,  grid_fcx, grid_gsp, grid_rax, grid_clx \n",
    "    del mhx_radar, ltx_radar,  fcx_radar, gsp_radar, rax_radar, clx_radar, u_init, v_init, w_init\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "the_time = datetime.datetime(2018,9,14,6,50)\n",
    "make_retrieved_grid_extended(the_time, ltx_list, mhx_list, cae_list,\n",
    "                                 clx_list, fcx_list, rax_list, gsp_list, do_hrrr=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_retrieved_grid_extended(the_time, ltx_list, mhx_list, cae_list,\n",
    "                                 clx_list, fcx_list, rax_list, gsp_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_bar(ax, length, location=(0.5, 0.05), linewidth=3):\n",
    "    \"\"\"\n",
    "    ax is the axes to draw the scalebar on.\n",
    "    location is center of the scalebar in axis coordinates ie. 0.5 is the middle of the plot\n",
    "    length is the length of the scalebar in km.\n",
    "    linewidth is the thickness of the scalebar.\n",
    "    \"\"\"\n",
    "    #Projection in metres, need to change this to suit your own figure\n",
    "    utm = ccrs.UTM(17)\n",
    "    #Get the extent of the plotted area in coordinates in metres\n",
    "    x0, x1, y0, y1 = ax.get_extent(utm)\n",
    "    #Turn the specified scalebar location into coordinates in metres\n",
    "    sbcx, sbcy = x0 + (x1 - x0) * location[0], y0 + (y1 - y0) * location[1]\n",
    "    #Generate the x coordinate for the ends of the scalebar\n",
    "    bar_xs = [sbcx - length * 500, sbcx + length * 500]\n",
    "    #Plot the scalebar\n",
    "    ax.plot(bar_xs, [sbcy, sbcy], transform=utm, color='k', linewidth=linewidth)\n",
    "    #Plot the scalebar label\n",
    "    ax.text(sbcx, sbcy, str(length) + ' km', transform=utm,\n",
    "            horizontalalignment='center', verticalalignment='bottom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_list = glob.glob('/lcrc/group/earthscience/rjackson/florence_winds/grids/test_extended/05kmwinds_grid*ext*4.0650.nc')\n",
    "\n",
    "grids = []\n",
    "for fn in grid_list:\n",
    "    grids.append(pyart.io.read_grid(fn))\n",
    "print(grid_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(15,10)) \n",
    "font = {'family' : 'monospace',\n",
    "        'weight' : 'bold',\n",
    "        'size'   : 20}\n",
    "\n",
    "plt.rc('font', **font)\n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "ax = pydda.vis.plot_horiz_xsection_barbs_map(grids, ax=ax, bg_grid_no=-1, level=1, barb_spacing_x_km=50.0,\n",
    "                                             barb_spacing_y_km=50.0, show_lobes=False)\n",
    "wind_speed = np.sqrt(grids[0].fields[\"u\"][\"data\"]**2 + grids[1].fields[\"v\"][\"data\"]**2)\n",
    "wind_speed = wind_speed.filled(np.nan)\n",
    "lons = grids[0].point_longitude[\"data\"]\n",
    "lats = grids[0].point_latitude[\"data\"]\n",
    "cs = ax.contour(lons[2, ::2, ::2], lats[1, ::2, ::2], wind_speed[2, ::2, ::2], levels=[28, 32], \n",
    "                linewidths=3, colors=['b', 'r', 'k'])\n",
    "rad_list = [\"GSP\", \"RAX\", \"FCX\", \"MHX\", \"CLX\", \"LTX\"]\n",
    "for i in range(len(grids)):\n",
    "    ax.text(grids[i].radar_longitude[\"data\"], grids[i].radar_latitude[\"data\"], rad_list[i],\n",
    "        fontsize=20, horizontalalignment=\"center\")\n",
    "\n",
    "ax.set_xticks(np.arange(-85, -70, 1))\n",
    "ax.set_yticks(np.arange(27, 40, 1))\n",
    "ax.set_xlim([-84, -73])\n",
    "ax.set_ylim([31, 39])\n",
    "scale_bar(ax, length=100, location=(0.5, 0.05))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print([i in enumerate(grids)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_grid = reduce_pyart_grid_res(grids[0], 2)\n",
    "disp = pyart.graph.GridMapDisplay(new_grid)\n",
    "#new_grid.y\n",
    "disp.plot_grid('velocity', level=7, vmin=0, vmax=60)\n",
    "#new_grid.fields['reflectivity'][\"data\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grids[0].fields[\"reflectivity\"][\"data\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split_grids_x= split_pyart_grid(grids[0], 3, axis=2)\n",
    "#fig, ax = plt.subplots(3, 1, figsize=(30,30))\n",
    "#for i in range(3):\n",
    "disp = pyart.graph.GridMapDisplay(split_grids_x[1])\n",
    "disp.plot_grid('reflectivity', level=7, vmin=0, vmax=60)\n",
    "\n",
    "#for i in range(3):\n",
    "#    disp = pyart.graph.GridMapDisplay(split_grids_y[i])\n",
    "#    disp.plot_grid('reflectivity', ax=ax[i,1], level=7, vmin=0, vmax=60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array_split([[2,2,2,4], [3,3,3,4]], 2, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}