WRF_FCST_Examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Examples of Plotting WRF and FV3 Forecast Files\n",
    "\n",
    "FV3 forecasts are written out in two separate files at each output time. Files prefixed `dynf` contain 3D dynamics variables, while files prefixed `phyf` contain physics forecast information, mainly at the surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import Namespace\n",
    "import math\n",
    "import os\n",
    "import sys\n",
    "import warnings\n",
    "import yaml\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "from mpl_toolkits.basemap import shiftgrid\n",
    "import numpy as np\n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, interplevel, to_np, latlon_coords\n",
    "\n",
    "warnings.filterwarnings('ignore') # Quiet the Jupyter Notebook reports of warnings in matplotlib/basemap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Provide a file path to a forecast directory. \n",
    "# The example below creates a dictionary containing 0hr forecast files for the HRRR and FV3 experiments\n",
    "\n",
    "\n",
    "file_path = '/scratch1/BMC/wrfruc/chunhua/fv3sar-testing/data/2019091912.hrrr.arw/wrfprd.old'\n",
    "fn_hrrr = os.path.join(file_path, 'wrfout_d01_2019-09-19_18_00_00')\n",
    "\n",
    "file_path = '/scratch2/BMC/wrfruc/cholt/work/fv3sar_data/chunhua/2019091912.GSD'\n",
    "datasets = {\n",
    "    'HRRR': Dataset(fn_hrrr, 'r'),\n",
    "    'GSDFV3': {fn:  Dataset(os.path.join(file_path, fn + f'{0:03d}.nc'), 'r') for fn in ['dynf', 'phyf']},\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Print out the available variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for v, info in datasets['HRRR'].variables.items():\n",
    "    try:\n",
    "        print(v, ':', info.description, info.shape, info.units)\n",
    "    except:\n",
    "        print(v, ':', info.name)    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot subplots with experiments and diffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eq_contours(data, data2=None):\n",
    "    '''\n",
    "    Returns a balanced set of contours for data that has negative values.\n",
    "    Also returns default colorbar to use for balanced, vs all positive values.\n",
    "    '''\n",
    "    \n",
    "    minval = np.amin(data) if data2 is None else min(np.amin(data), np.amin(data2))\n",
    "    maxval = np.amax(data) if data2 is None else max(np.amax(data), np.amax(data2))\n",
    "    \n",
    "    minval = math.floor(minval) if abs(minval) > 1 else minval\n",
    "    maxval = math.ceil(maxval) if maxval > 1 else maxval\n",
    "    if minval == maxval:\n",
    "        return np.linspace(-1, 1, 5), 'seismic'\n",
    "    if np.amin(data) < 0:\n",
    "        # Set balanced contours. Choose an odd number in linspace below\n",
    "        maxval = max(abs(minval), abs(maxval))\n",
    "        return np.linspace(-maxval, maxval, 21), 'seismic'\n",
    "    else:\n",
    "        return np.linspace(minval, maxval, 21), 'jet'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data(data, lat, lon, title, expt, fig, ax, contours=None, cm=None):\n",
    "    \n",
    "    '''\n",
    "    Input parameters:\n",
    "    \n",
    "        dataL: 2D Numpy array to be plotted in Left column\n",
    "        dataR: 2D Numpy array to be plotted in Right column\n",
    "        lat: 2D Numpy array of latitude\n",
    "        lon: 2D Numpy array of longitude\n",
    "        title: String describing the variable being plotted.\n",
    "        \n",
    "    Draws a Basemap representation with the contoured data overlayed, \n",
    "    with a colorbar for each experiment, and the difference between the two.\n",
    "        \n",
    "    '''\n",
    "    \n",
    "    def trim_grid():\n",
    "        '''\n",
    "        The u, v, and H data from analysis are all on grids either one column, or one row smaller than lat/lon. \n",
    "        Return the smaller lat, lon grids, given the shape of the data to be plotted.\n",
    "        Has no effect when all grids are the same size.\n",
    "        '''\n",
    "        y, x = np.shape(data)\n",
    "        return lat[:y, :x], lon[:y, :x]\n",
    "                   \n",
    "    \n",
    "    lat_trim, lon_trim = trim_grid()\n",
    "\n",
    "    # Check out this link for all cmap options: https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html\n",
    "    # A good redwhiteblue cmap for increments is seismic, and for full fields with rainbow, change to jet\n",
    "\n",
    "    m = Basemap(projection='mill', \n",
    "                llcrnrlon=lon.min()-2,\n",
    "                urcrnrlon=lon.max()+2,\n",
    "                llcrnrlat=lat.min()-2,\n",
    "                urcrnrlat=lat.max()+2,\n",
    "                resolution='l',\n",
    "                ax=ax,\n",
    "               )\n",
    "    x, y = m(lon_trim, lat_trim)\n",
    "\n",
    "    # Use the same contour values for both experiments.\n",
    "    contours, cm = eq_contours(data) if contours is None else (contours, cm)\n",
    "    \n",
    "    # Draw the contoured data over the map\n",
    "    cs = m.contourf(x, y, data, contours, cmap=cm, ax=ax)\n",
    "    m.drawcoastlines();\n",
    "    m.drawmapboundary();\n",
    "    m.drawparallels(np.arange(-90.,120.,5),labels=[1,0,0,0]);\n",
    "    m.drawmeridians(np.arange(-180.,180.,5),labels=[0,0,0,1]);\n",
    "    fig.colorbar(cs, ax=ax, orientation='vertical', shrink=0.25);\n",
    "    ax.set_title(f\"{expt}: {title}\")\n",
    "    return contours, cm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load in the variable information\n",
    "\n",
    "The mapping of variables between WRF and FV3 is not 1:1. The file `variable_mapping.yaml` contains information about atmospheric variables as they relate to each model. The `yaml` file is read into a dictionary in the following cell.\n",
    "\n",
    "The keys in that file generally follow this heirarchy:\n",
    "```\n",
    "   generic variable name:  # Name to be referenced in this script\n",
    "       description:        # Brief description of the variable\n",
    "       fv3:                # Section containing FV3-specific settings\n",
    "           file:           # Prefix of file that contains this variable [phyf, dynf]\n",
    "           name:           # Name of NetCDF variable\n",
    "       unit:               # Preferred unit\n",
    "       wrf:                # Secction containing WRF-specific settings\n",
    "           method:         # The retrieval method for the variable [direct, getvar] \n",
    "           name:           # The name of the variable to retrieve, if \"direct\" method, \n",
    "                           # the NetCDF variable, if \"getvar\" method, then the variable \n",
    "                           # name it expects\n",
    "           \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('variable_mapping.yaml') as fn:\n",
    "    var_map = yaml.load(fn, Loader=yaml.FullLoader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot the variables\n",
    "\n",
    "## A little info about dictionary references\n",
    "\n",
    "Referencing dictionary entries can be done with the following syntax:\n",
    "\n",
    "```\n",
    "     my_dict = {'a': 2, 'b': 3, 'd': {'apple': 'red', 'banana': 'orange'}}\n",
    "     my_dict['a']            # returns 2\n",
    "     my_dict['d']['apple']   # returns 'red'\n",
    "     my_dict['c']            # Run time error!\n",
    "```\n",
    "\n",
    "To avoid the run time error when entries don't exist, we should make use of the dictionary attribute `get()`. By default, `get()` returns `None` if the key does not exist. If we have a nested dictionary, this can also be problematic for retrieving missing values. To ensure we move forward in the event we don't care about missing values we can use the following syntax:\n",
    "\n",
    "```\n",
    "    my_dict.get('c')                   # Returns None\n",
    "    my_dict.get('c', {})               # Returns an empty dict\n",
    "    my_dict.get('c', {}).get('apple')  # Returns None\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_fraction(arr):\n",
    "    \n",
    "    '''Given an array of values, returns the array in percent'''\n",
    "    \n",
    "    return arr / 100.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fun_call_by_name(val):\n",
    "    \n",
    "    ''' Given an input string, val, returns the corresponding callable function.'''\n",
    "    \n",
    "    if '.' in val:\n",
    "        module_name, fun_name = val.rsplit('.', 1)\n",
    "        # you should restrict which modules may be loaded here\n",
    "        assert module_name.startswith('my.')\n",
    "    else:\n",
    "        module_name = '__main__'\n",
    "        fun_name = val\n",
    "    try:\n",
    "        __import__(module_name)\n",
    "    except ImportError as exc:\n",
    "        raise ConstructorError(\n",
    "            \"while constructing a Python object\", mark,\n",
    "            \"cannot find module %r (%s)\" % (utf8(module_name), exc), mark)\n",
    "    module = sys.modules[module_name]\n",
    "    fun = getattr(module, fun_name)\n",
    "    return fun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform(data, model, var, var_map):\n",
    "    \n",
    "    '''Applies a trasnformation on the input data given settings in var_map.'''\n",
    "    \n",
    "    # Check to see if a transformation is needed.\n",
    "    transform = var_map.get(var, {}).get(model, {}).get('transform', False)\n",
    "    \n",
    "    if transform:\n",
    "        f = fun_call_by_name(transform)\n",
    "        return f(data)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fhr = 0\n",
    "\n",
    "# Load in the latitude/longitude variables from each model.\n",
    "# Since these are different, we won't remap here, just plot on the NetCDF map provided.\n",
    "lat_f = datasets['GSDFV3']['dynf']['grid_yt'][::] * 180 / math.pi\n",
    "lon_f = datasets['GSDFV3']['dynf']['grid_xt'][::] * 180 / math.pi\n",
    "\n",
    "lat_h = np.squeeze(datasets['HRRR']['XLAT'][::])\n",
    "lon_h = np.squeeze(datasets['HRRR']['XLONG'][::])\n",
    "\n",
    "\n",
    "# List of two experiments, and diff for labeling plots\n",
    "expt = ['GSDFV3', 'HRRR']\n",
    "\n",
    "# Variables to plot from dynf and phyf files\n",
    "# To add additional variables, corresponding entries must be made in variable_mapping.yaml\n",
    "vars_ = {\n",
    "    'dynf': ['u', 'v', 'temp', 'specific_humidity', 'delz'],\n",
    "    'phyf': ['temp_sfc', 'soil_temp', 'albedo_ave', 'ground_flux', 'soil_type', 'orog', 'pbl_height'],\n",
    "}\n",
    "\n",
    "\n",
    "# Vertical level from FV3 (surface is last index, -1; TOA is first index, 0)\n",
    "lev_fv3 = -1\n",
    "lev_wrf = 0\n",
    "\n",
    "\n",
    "for file, varlist in vars_.items():\n",
    "    nvar = len(varlist)\n",
    "\n",
    "    # Make a figure with subplots for each variable (nvar rows)\n",
    "    fig, ax = plt.subplots(nvar, 2, figsize=(24, 8*nvar))\n",
    "\n",
    "\n",
    "    for i, var in enumerate(varlist):\n",
    "\n",
    "        #### Get/plot data from FV3 ###\n",
    "        \n",
    "        # Ensure that \"None\" returns if the variable does not exist in the var_map\n",
    "        fv3_var = var_map.get(var, {}).get('fv3', {}).get('name')\n",
    "        titleL = f'{fv3_var} at {fhr} hr fcst'\n",
    "        data = np.squeeze(datasets[expt[0]][file][fv3_var][::])\n",
    "        \n",
    "        # Transform data, if needed\n",
    "        data = transform(data, 'fv3', var, var_map)\n",
    "        dataL = data[lev_fv3] if data.ndim == 3 else data\n",
    "\n",
    "        #### Get data from WRF ###\n",
    "        wrf_var = var_map.get(var, {}).get('wrf', {}).get('name')\n",
    "        titleR = f'{wrf_var} at {fhr} hr fcst'\n",
    "\n",
    "        if wrf_var:\n",
    "            data = getvar(datasets[expt[1]], wrf_var, squeeze=True)\n",
    "            # Transform data, if needed\n",
    "            data = transform(data, 'wrf', var, var_map)\n",
    "                \n",
    "        else:\n",
    "            print(f'Plotting zeros for HRRR {var}')\n",
    "            data = np.zeros_like(lat_h)\n",
    "\n",
    "        dataR = data[lev_wrf] if data.ndim == 3 else data        \n",
    "        \n",
    "        # Determine contours based on both datasets and plot\n",
    "        contours, cm = eq_contours(dataL, dataR)\n",
    "        plot_data(dataL, lat_f, lon_f, titleL, expt[0], fig, ax[i, 0], contours, cm)\n",
    "        plot_data(dataR, lat_h, lon_h, titleR, expt[1], fig, ax[i, 1], contours, cm)\n"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}