FV3_FCST_Examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Examples of Plotting 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": [
    "import math\n",
    "import os\n",
    "from argparse import Namespace\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"
   ]
  },
  {
   "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 2 experiments, expt, through the first 4 forecast hours (including 0)\n",
    "\n",
    "\n",
    "file_path = '/scratch2/BMC/wrfruc/cholt/work/fv3sar_data/chunhua/2019091912.{expt}'\n",
    "\n",
    "last_fhr = 4\n",
    "files = {expt: {x: [os.path.join(file_path.format(expt=expt), x + f'{i:03d}.nc') for i in range(5)]  for x in ['dynf', 'phyf']} for expt in ['GSD', 'GFS']}\n",
    "\n",
    "print('The files dict looks like this:')\n",
    "files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load each file in the files dict into a NetCDF Dataset\n",
    "\n",
    "def load_Dataset(files: dict, ds: dict):\n",
    "    assert(isinstance(files, dict))\n",
    "    for k, v in files.items():\n",
    "        if isinstance(v, dict):\n",
    "            ds[k] = {}\n",
    "            load_Dataset(v, ds[k])\n",
    "        else:\n",
    "            ds[k] = [Dataset(f, 'r') for f in list(v)]\n",
    "\n",
    "ds_dict = {}\n",
    "load_Dataset(files, ds_dict)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the dictionary into a Namespace data structure.\n",
    "# This step is not necessary, but cuts down the syntax needed to reference each item in the dict.\n",
    "#\n",
    "# Example: Retrieve the 0 hr forecast Dataset from GFS Dynamics\n",
    "#            dict: ds_dict['GFS']['dynf'][0]\n",
    "#       Namespace: datasets.GFS.dynf[0]\n",
    "\n",
    "\n",
    "def make_namespace(ns: Namespace(), d: dict):\n",
    "    assert(isinstance(d, dict))\n",
    "    for k, v in d.items():\n",
    "        if isinstance(v, dict):\n",
    "            leaf_ns = Namespace()\n",
    "            ns.__dict__[k] = leaf_ns\n",
    "            make_namespace(leaf_ns, v)\n",
    "        else:\n",
    "            ns.__dict__[k] = v\n",
    "   \n",
    "    \n",
    "datasets = Namespace()\n",
    "make_namespace(datasets, ds_dict)\n",
    "datasets.GFS.dynf[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Print out the available variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('~~~~~~~~~ DYNAM FILE ~~~~~~~~~~~~~~~')\n",
    "for v, info in datasets.GFS.dynf[0].variables.items():\n",
    "    print(v, ':', info.long_name, info.shape)\n",
    "    \n",
    "print('~~~~~~~~~ PHYSICS FILE ~~~~~~~~~~~~~~~')\n",
    "for v, info in datasets.GFS.phyf[0].variables.items():\n",
    "    print(v, ':', info.long_name, info.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot subplots with experiments and diffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data(dataL, dataR, lat, lon, title, expt):\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(dataL)\n",
    "        return lat[:y, :x], lon[:y, :x]\n",
    "    \n",
    "    def eq_contours(indata):\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",
    "        minval = np.amin(indata)\n",
    "        maxval = np.amax(indata)\n",
    "        if minval == maxval:\n",
    "            return np.linspace(-1, 1, 5), 'seismic'\n",
    "        if np.amin(indata) < 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), 'hsv'\n",
    "                     \n",
    "    \n",
    "    lat_trim, lon_trim = trim_grid()\n",
    "\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 3, figsize=(24, 12))\n",
    "    \n",
    "    for i, data in enumerate([dataL, dataR, dataL-dataR]):\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 hsv\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='c',\n",
    "                    ax=ax[i],\n",
    "                   )\n",
    "        x, y = m(lon_trim, lat_trim)\n",
    "       \n",
    "        # Use the same contour values for both experiments.\n",
    "        if i < 2:\n",
    "            contours, cm = eq_contours(dataL)\n",
    "        else:\n",
    "            contours, cm = eq_contours(data)\n",
    "            \n",
    "        # Draw the contoured data over the map\n",
    "        cs = m.contourf(x, y, data, contours, cmap=cm, ax=ax[i])\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[i], orientation='vertical', shrink=0.25);\n",
    "        ax[i].set_title(f\"{expt[i]}: {title}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot the variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fhr = 0\n",
    "lat = datasets.GSD.dynf[fhr]['grid_yt'][::] * 180 / math.pi\n",
    "lon = datasets.GSD.dynf[fhr]['grid_xt'][::] * 180 / math.pi\n",
    "\n",
    "# List of two experiments, and diff for labeling plots\n",
    "expt = ['GSD', 'GFS', 'GSD-GFS']\n",
    "\n",
    "# Variables to plot from dynf and phyf files\n",
    "vars_ = {\n",
    "    'dynf': ['ugrd', 'vgrd', 'tmp', 'spfh', 'delz'],\n",
    "    'phyf': ['tmpsfc', 'soilt1', 'soilt2', 'soilt3', 'albdo_ave',  'gflux', 'sotyp', 'orog', 'hpbl'],\n",
    "}\n",
    "\n",
    "# Vertical level (surface is last index, -1; TOA is first index, 0)\n",
    "lev = -1\n",
    "\n",
    "\n",
    "file = 'dynf'\n",
    "for var in vars_[file]:\n",
    "    dataL = np.squeeze(ds_dict[expt[0]][file][fhr][var][::])[lev]\n",
    "    dataR = np.squeeze(ds_dict[expt[1]][file][fhr][var][::])[lev]\n",
    "    title = f'{var} at {fhr} hr fcst'\n",
    "    plot_data(dataL, dataR, lat, lon, title, expt)\n",
    "\n",
    "file = 'phyf'\n",
    "for var in vars_[file]:\n",
    "    dataL = np.squeeze(ds_dict[expt[0]][file][fhr][var][::])\n",
    "    dataR = np.squeeze(ds_dict[expt[1]][file][fhr][var][::])\n",
    "    title = f'{var} at {fhr} hr fcst'\n",
    "    plot_data(dataL, dataR, lat, lon, title, expt)\n",
    "        \n",
    "\n",
    "\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
}