gcc_vs_gcc_comparison.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare two GEOS-Chem Classic datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview of this Notebook\n",
    "\n",
    "* Import dependencies\n",
    "* Define functions: utility, grid, and plotting\n",
    "* Define data directories and global variables\n",
    "* Compare restart files\n",
    "* Store areas for normalization\n",
    "* Compare SpeciesConc diagnostic collection\n",
    "* Create PDFs for subset of species\n",
    "* Other diagnostic collections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import cubedsphere as cs\n",
    "import xesmf as xe\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib import ticker\n",
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "from matplotlib.colors import ListedColormap\n",
    "from cartopy import crs\n",
    "from cartopy.mpl.geoaxes import GeoAxes\n",
    "\n",
    "%matplotlib inline\n",
    "import warnings; warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Also define colormap. Colormap file source: https://bitbucket.org/gcst/gcpy\n",
    "rgb_WhGrYlRd = np.genfromtxt('/n/home08/elundgren/GC/python/WhGrYlRd.txt',delimiter=' ')\n",
    "WhGrYlRd = ListedColormap(rgb_WhGrYlRd/255.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define general utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gcc_filepath(outputdir, collection, day, time):\n",
    "    if collection == 'Emissions':\n",
    "        filepath = os.path.join(outputdir, 'HEMCO_diagnostics.{}{}.nc'.format(day,time))\n",
    "    else:\n",
    "        filepath = os.path.join(outputdir, 'GEOSChem.{}.{}_{}z.nc4'.format(collection,day,time))\n",
    "    return filepath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_paths( refpath, devpath):\n",
    "    if not os.path.exists(refpath):\n",
    "        print('ERROR! Path 1 does not exist: {}'.format(refpath))\n",
    "    else:\n",
    "        print('Path 1 exists: {}'.format(refpath))\n",
    "    if not os.path.exists(devpath):\n",
    "        print('ERROR! Path 2 does not exist: {}'.format(devpath))\n",
    "    else:\n",
    "        print('Path 2 exists: {}'.format(devpath))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def compare_varnames(refdata, refstr, devdata, devstr):\n",
    "    \n",
    "    # Find common variables in collection by generating lists and list overlap\n",
    "    refvars = [k for k in refdata.data_vars.keys()]\n",
    "    devvars= [k for k in devdata.data_vars.keys()]\n",
    "    commonvars = sorted(list(set(refvars).intersection(set(devvars))))\n",
    "    refonly = [v for v in refvars if v not in devvars]\n",
    "    devonly = [v for v in devvars if v not in refvars]\n",
    "    dimmismatch = [v for v in commonvars if refdata[v].ndim != devdata[v].ndim]\n",
    "    commonvars2D = [v for v in commonvars if refdata[v].ndim == 3]\n",
    "    commonvars3D = [v for v in commonvars if devdata[v].ndim == 4]\n",
    "    \n",
    "    # Print information on common and mismatching variables, as well as dimensions\n",
    "    print('{} common variables ({} are 2-dim and {} are 3-dim)'.format(len(commonvars), len(commonvars2D), len(commonvars3D)))\n",
    "    if len(refonly) > 0:\n",
    "        print('{} variables in {} only (skip)'.format(len(refnly)),refstr)\n",
    "        print('   Variable names: {}'.format(refonly))\n",
    "    if len(devonly) > 0:\n",
    "        print('{} variables in {} only (skip)'.format(len(devonly)),devstr)\n",
    "        print('   Variable names: {}'.format(devonly))\n",
    "    if len(dimmismatch) > 0:\n",
    "        print('{} common variables have different dimensions'.format(len(dimmismatch)))\n",
    "        print('   Variable names: {}'.format(dimmismatch))\n",
    "        \n",
    "    return [commonvars, commonvars2D, commonvars3D]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stats(refdata, refstr, devdata, devstr, varname):\n",
    "    refvar = refdata[varname]\n",
    "    devvar = devdata[varname]\n",
    "    units = refdata[varname].units\n",
    "    print('Data units: {}'.format(units))\n",
    "    print('Array sizes:')\n",
    "    print('    {}:  {}'.format(refstr,refvar.shape))\n",
    "    print('    {}:  {}'.format(devstr,devvar.shape))\n",
    "    print('Global stats:')\n",
    "    print('  Mean:')\n",
    "    print('    {}:  {}'.format(refstr,np.round(refvar.values.mean(),20)))\n",
    "    print('    {}:  {}'.format(devstr,np.round(devvar.values.mean(),20)))\n",
    "    print('  Min:')\n",
    "    print('    {}:  {}'.format(refstr,np.round(refvar.values.min(),20)))\n",
    "    print('    {}:  {}'.format(devstr,np.round(devvar.values.min(),20)))\n",
    "    print('  Max:')\n",
    "    print('    {}:  {}'.format(refstr,np.round(refvar.values.max(),20)))\n",
    "    print('    {}:  {}'.format(devstr,np.round(devvar.values.max(),20)))\n",
    "    print('  Sum:')\n",
    "    print('    {}:  {}'.format(refstr,np.round(refvar.values.sum(),20)))\n",
    "    print('    {}:  {}'.format(devstr,np.round(devvar.values.sum(),20)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_collection_data(refdir, devdir, collection, day, time):\n",
    "    reffile = get_gcc_filepath(refdir, collection, day, time)\n",
    "    devfile = get_gcc_filepath(devdir, collection, day, time)\n",
    "    check_paths(reffile, devfile)\n",
    "    refdata = xr.open_dataset(reffile)\n",
    "    devdata = xr.open_dataset(devfile)\n",
    "    return [refdata, devdata]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define grid functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_grid_LL(llres):\n",
    "    [dlat,dlon] = list(map(float, llres.split('x')))\n",
    "    lon_b = np.linspace(-180 - dlon/2, 180 - dlon/2, int(360/dlon) + 1, endpoint=True)\n",
    "    lat_b = np.linspace(-90 - dlat/2, 90 + dlat/2, int(180/dlat) + 2, endpoint=True).clip(-90,90)\n",
    "    lat = (lat_b[1:] + lat_b[:-1]) / 2\n",
    "    lon = (lon_b[1:] + lon_b[:-1]) / 2\n",
    "    llgrid = {'lat': lat, \n",
    "              'lon': lon, \n",
    "              'lat_b': lat_b, \n",
    "              'lon_b': lon_b}\n",
    "    return llgrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define functions for plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to plot and compare single vertical level of GC Classic\n",
    "\n",
    "Notes: The two GCX files must be at the same grid resolution. You can use this function to plot interactively or to generate a multi-page pdf of plots. It takes a list of variable names to plot for a single collection only. You can plot for any level and any time slice in the file. By default the colorbars for the plots will have the same range, but you can turn this feature off. Also by default the colorbar of the fractional difference between the model outputs will be limited to +/-2, but you can change this as well via the passed parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_single_level(refdata, refstr, devdata, devstr, llres, varlist=None,\n",
    "                         ilev=0, itime=0, savepdf=False, pdfname='gcc_vs_gcc_map.pdf', match_cbar=True, \n",
    "                         full_ratio_range=False, normalize_by_area=False, area_ref=None, \n",
    "                         area_dev=None, check_units=True):\n",
    "    \n",
    "    # If no varlist is passed, plot all (surface only for 3D)\n",
    "    if varlist == None:\n",
    "        [varlist, commonvars2D, commonvars3D] = compare_varnames(refdata, devdata)\n",
    "        print('Plotting all common variables (surface only if 3D)')\n",
    "    n_var = len(varlist)\n",
    "\n",
    "    # Get lat-lon grid\n",
    "    llgrid = make_grid_LL(llres)\n",
    "\n",
    "    # Get lat/lon extents\n",
    "    [minlon, maxlon] = [min(llgrid['lon_b']), max(llgrid['lon_b'])]\n",
    "    [minlat, maxlat] = [min(llgrid['lat_b']), max(llgrid['lat_b'])]\n",
    "    \n",
    "    # Create pdf (if saving)\n",
    "    if savepdf:\n",
    "        print('\\nCreating {} for {} variables'.format(pdfname,n_var))\n",
    "        pdf = PdfPages(pdfname)\n",
    "\n",
    "    # Loop over variables\n",
    "    for ivar in range(n_var):\n",
    "        if savepdf: print('{} '.format(ivar), end='')\n",
    "        varname = varlist[ivar]\n",
    "        \n",
    "        # Do some checks: dimensions and units\n",
    "        varndim_ref = refdata[varname].ndim\n",
    "        varndim_dev = devdata[varname].ndim\n",
    "        if check_units: \n",
    "            assert varndim_ref == varndim_dev, 'Dimensions do not agree for {}!'.format(varname)\n",
    "        units_ref = refdata[varname].units\n",
    "        units_dev = devdata[varname].units\n",
    "        if check_units: \n",
    "            assert units_ref == units_dev, 'Units do not match for {}!'.format(varname)\n",
    "            \n",
    "        # if normalizing by area, adjust units to be per m2, and adjust title string\n",
    "        units = units_ref\n",
    "        varndim = varndim_ref\n",
    "        subtitle_extra = ''\n",
    "                    \n",
    "        # Slice the data\n",
    "        if varndim == 4: \n",
    "            ds_ref = refdata[varname].isel(time=itime,lev=ilev)\n",
    "            ds_dev = devdata[varname].isel(time=itime,lev=ilev)\n",
    "        elif varndim == 3: \n",
    "            ds_ref = refdata[varname].isel(time=itime)\n",
    "            ds_dev = devdata[varname].isel(time=itime)\n",
    "            \n",
    "        # if normalizing by area, transform on the native grid and adjust units and subtitle string\n",
    "        exclude_list = ['WetLossConvFrac','Prod_','Loss_']\n",
    "        if normalize_by_area and not any(s in varname for s in exclude_list):\n",
    "            ds_ref.values = ds_ref.values / refdata['AREAM2'].values\n",
    "            ds_dev.values = ds_dev.values / devdata['AREAM2'].values\n",
    "            units = '{} m-2'.format(units)\n",
    "            subtitle_extra = ', Normalized by Area'\n",
    "            \n",
    "        # Regrid the slices (skip for gchp vs gchp for now)\n",
    "        vmin_ref = ds_ref.min()\n",
    "        vmin_dev = ds_dev.min()\n",
    "        vmin_cmn = np.min([vmin_ref, vmin_dev])\n",
    "        vmax_ref = ds_ref.max()\n",
    "        vmax_dev = ds_dev.max()\n",
    "        vmax_cmn = np.max([vmax_ref, vmax_dev])\n",
    "        if match_cbar: [vmin, vmax] = [vmin_cmn, vmax_cmn]\n",
    "        \n",
    "        # Create 2x2 figure\n",
    "        figs, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2, 2, figsize=[12,9], \n",
    "                                                      subplot_kw={'projection': crs.PlateCarree()})\n",
    "        # Give the figure a title\n",
    "        offset = 0.96\n",
    "        fontsize=25\n",
    "        if varndim == 4:\n",
    "            if ilev == 0: levstr = 'Surface'\n",
    "            elif ilev == 22: levstr = '500 hPa'\n",
    "            else: levstr = 'Level ' +  str(ilev-1)\n",
    "            figs.suptitle('{}, {}'.format(varname,levstr), fontsize=fontsize, y=offset)\n",
    "        elif varndim == 3: \n",
    "            figs.suptitle('{}'.format(varname), fontsize=fontsize, y=offset)\n",
    "        else:\n",
    "            print('varndim is 2 for {}! Must be 3 or 4.'.format(varname))   \n",
    "        \n",
    "        # Subplot (0,0): Ref\n",
    "        ax0.coastlines()\n",
    "        if not match_cbar: [vmin, vmax] = [vmin_ref, vmax_ref]\n",
    "        plot0 = ax0.imshow(ds_ref, extent=(minlon, maxlon, minlat, maxlat), \n",
    "                           cmap=WhGrYlRd,vmin=vmin, vmax=vmax)\n",
    "        ax0.set_title('{} (Ref){}\\n{}'.format(refstr,subtitle_extra,llres)) \n",
    "        cb = plt.colorbar(plot0, ax=ax0, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.1 or (vmax-vmin) > 100:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "        \n",
    "        # Subplot (0,1): Dev\n",
    "        ax1.coastlines()\n",
    "        if not match_cbar: [vmin, vmax] = [vmin_dev, vmax_dev]\n",
    "        plot1 = ax1.imshow(ds_dev, extent=(minlon, maxlon, minlat, maxlat), \n",
    "                           cmap=WhGrYlRd,vmin=vmin, vmax=vmax)\n",
    "        ax1.set_title('{} (Dev){}\\n{}'.format(devstr,subtitle_extra,llres)) \n",
    "        cb = plt.colorbar(plot1, ax=ax1, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.1 or (vmax-vmin) > 100:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "            \n",
    "        # Subplot (1,0): Difference\n",
    "        absdiff = ds_dev - ds_ref\n",
    "        diffabsmax = max([np.abs(absdiff.min()), np.abs(absdiff.max())])\n",
    "        [vmin, vmax] = [-diffabsmax, diffabsmax]\n",
    "        ax2.coastlines()\n",
    "        plot2 = ax2.imshow(absdiff, extent=(minlon, maxlon, minlat, maxlat), \n",
    "                           cmap='RdBu_r',vmin=vmin, vmax=vmax)\n",
    "        ax2.set_title('Difference\\n(Dev - Ref)') \n",
    "        cb = plt.colorbar(plot2, ax=ax2, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.1 or (vmax-vmin) > 100:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "        \n",
    "        # Subplot (1,1): Fractional Difference (restrict to +/-2)\n",
    "        fracdiff = (ds_dev - ds_ref) / ds_ref\n",
    "        if full_ratio_range: \n",
    "            [vmin, vmax] = [None, None]\n",
    "        else: \n",
    "            [vmin, vmax] = [-2, 2]\n",
    "        ax3.coastlines()\n",
    "        plot3 = ax3.imshow(fracdiff, extent=(minlon, maxlon, minlat, maxlat), \n",
    "                           cmap='RdBu_r',vmin=vmin, vmax=vmax)\n",
    "        ax3.set_title('Fractional Difference\\n(Dev - Ref)/Ref') \n",
    "        cb = plt.colorbar(plot3, ax=ax3, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.1 or (vmax-vmin) > 100:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)  \n",
    "            \n",
    "        if savepdf:    \n",
    "            pdf.savefig(figs)\n",
    "            plt.close(figs)\n",
    "            \n",
    "    if savepdf: pdf.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to plot and compare zonal mean\n",
    "\n",
    "Many of the same features available for plotting a single level above are also available for this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def compare_zonal_mean(refdata, refstr, devdata, devstr, llres, varlist=None, \n",
    "                       weightsdir='.', itime=0, savepdf=False, \n",
    "                       pdfname='gchp_vs_gchp_map.pdf', match_cbar=True, full_ratio_range=False, \n",
    "                       normalize_by_area=False ):\n",
    "\n",
    "    # If no varlist is passed, plot all 3D variables in the dataset\n",
    "    if varlist == None:\n",
    "        [commonvars, commonvars2D, varlist] = compare_varnames(refdata, devdata)\n",
    "        print('Plotting all 3D variables')\n",
    "    n_var = len(varlist)\n",
    "    \n",
    "    # Get lat-lon grid\n",
    "    llgrid = make_grid_LL(llres)\n",
    "    \n",
    "    # Universal plot setup\n",
    "    xtick_positions = np.arange(-90,91,30)\n",
    "    xticklabels = ['{}$\\degree$'.format(x) for x in xtick_positions]\n",
    "    ytick_positions = np.arange(0,61,20)\n",
    "    yticklabels = [str(y) for y in ytick_positions]\n",
    "    \n",
    "    # Create pdf (if saving)\n",
    "    if savepdf:\n",
    "        print('\\nCreating {} for {} variables'.format(pdfname, n_var))\n",
    "        pdf = PdfPages(pdfname)\n",
    "\n",
    "    # Loop over variables\n",
    "    for ivar in range(n_var):\n",
    "        if savepdf: print('{} '.format(ivar), end='')\n",
    "        varname = varlist[ivar]\n",
    "        \n",
    "        # Do some checks: dimensions and units\n",
    "        varndim_ref = refdata[varname].ndim\n",
    "        varndim_dev = devdata[varname].ndim\n",
    "        nlev = 72\n",
    "        assert varndim_ref == varndim_dev, 'Dimensions do not agree for {}!'.format(varname)\n",
    "        units_ref = refdata[varname].units\n",
    "        units_dev = devdata[varname].units\n",
    "        assert units_ref == units_dev, 'Units do not match for {}!'.format(varname)\n",
    "        \n",
    "        # Set plot extent\n",
    "        extent=(-90,90,0,nlev)\n",
    "        \n",
    "        # if normalizing by area, adjust units to be per m2, and adjust title string\n",
    "        units = units_ref\n",
    "        varndim = varndim_ref\n",
    "        subtitle_extra = ''\n",
    "        \n",
    "        # Slice the data\n",
    "        ds_ref = refdata[varname].isel(time=itime)\n",
    "        ds_dev = devdata[varname].isel(time=itime)\n",
    "\n",
    "        # if normalizing by area, transform on the native grid and adjust units and subtitle string\n",
    "        exclude_list = ['WetLossConvFrac','Prod_','Loss_']\n",
    "        if normalize_by_area and not any(s in varname for s in exclude_list):\n",
    "            ds_ref.values = ds_ref.values / refdata['AREAM2'].values[np.newaxis,:,:]\n",
    "            ds_dev.values = ds_dev.values / devdata['AREAM2'].values[np.newaxis,:,:]\n",
    "            units = '{} m-2'.format(units)\n",
    "            subtitle_extra = ', Normalized by Area'\n",
    "            \n",
    "        # Calculate zonal mean of the regridded data\n",
    "        zm_ref = ds_ref.mean(axis=2)\n",
    "        zm_dev = ds_dev.mean(axis=2)\n",
    "            \n",
    "        # Get min and max for colorbar limits\n",
    "        [vmin_ref, vmax_ref] = [zm_ref.min(), zm_ref.max()]\n",
    "        [vmin_dev, vmax_dev] = [zm_dev.min(), zm_dev.max()]\n",
    "        vmin_cmn = np.min([vmin_ref, vmin_dev])\n",
    "        vmax_cmn = np.max([vmax_ref, vmax_dev])\n",
    "        if match_cbar: [vmin, vmax] = [vmin_cmn, vmax_cmn]\n",
    "        \n",
    "        # Create 2x2 figure\n",
    "        figs, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2, 2, figsize=[12,12], \n",
    "                                                      subplot_kw={'projection': crs.PlateCarree()})\n",
    "        # Give the page a title\n",
    "        offset = 0.96\n",
    "        fontsize=25\n",
    "        figs.suptitle('{}, Zonal Mean'.format(varname), fontsize=fontsize, y=offset)\n",
    "\n",
    "        # Subplot 0: Ref\n",
    "        if not match_cbar: [vmin, vmax] = [vmin_ref, vmax_ref]\n",
    "        plot0 = ax0.imshow(zm_ref, cmap=WhGrYlRd, extent=extent, vmin=vmin, vmax=vmax)\n",
    "        ax0.set_title('{} (Ref){}\\n{}'.format(refstr, subtitle_extra, llres ))\n",
    "        ax0.set_aspect('auto')\n",
    "        ax0.set_xticks(xtick_positions)\n",
    "        ax0.set_xticklabels(xticklabels)\n",
    "        ax0.set_yticks(ytick_positions)\n",
    "        ax0.set_yticklabels(yticklabels)\n",
    "        cb = plt.colorbar(plot0, ax=ax0, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.001 or (vmax-vmin) > 1000:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "        \n",
    "        # Subplot 1: Dev\n",
    "        if not match_cbar: [vmin, vmax] = [vmin_dev, vmax_dev]\n",
    "        plot1 = ax1.imshow(zm_dev, cmap=WhGrYlRd, extent=extent, vmin=vmin, vmax=vmax)\n",
    "        ax1.set_title('{} (Dev){}\\n{}'.format(devstr, subtitle_extra, llres))\n",
    "        ax1.set_aspect('auto')\n",
    "        ax1.set_xticks(xtick_positions)\n",
    "        ax1.set_xticklabels(xticklabels)\n",
    "        ax1.set_yticks(ytick_positions)\n",
    "        ax1.set_yticklabels(yticklabels)\n",
    "        cb = plt.colorbar(plot1, ax=ax1, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.001 or (vmax-vmin) > 1000:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "            \n",
    "        # Subplot 2: Difference\n",
    "        zm_absdiff = zm_dev - zm_ref\n",
    "        diffabsmax = max([np.abs(zm_absdiff.min()), np.abs(zm_absdiff.max())])\n",
    "        [vmin, vmax] = [-diffabsmax, diffabsmax]\n",
    "        plot2 = ax2.imshow(zm_absdiff, cmap='RdBu_r', extent=extent, vmin=vmin, vmax=vmax)\n",
    "        ax2.set_title('Difference\\n(Dev - Ref)')\n",
    "        ax2.set_aspect('auto')\n",
    "        ax2.set_xticks(xtick_positions)\n",
    "        ax2.set_xticklabels(xticklabels)\n",
    "        ax2.set_yticks(ytick_positions)\n",
    "        ax2.set_yticklabels(yticklabels)\n",
    "        cb = plt.colorbar(plot2, ax=ax2, orientation='horizontal', pad=0.10)\n",
    "        if (vmax-vmin) < 0.001 or (vmax-vmin) > 1000:\n",
    "            cb.locator = ticker.MaxNLocator(nbins=4)\n",
    "            cb.update_ticks()\n",
    "        cb.set_label(units)\n",
    "        \n",
    "        # Subplot 3: Fractional Difference (restrict to +/-2)\n",
    "        zm_fracdiff = (zm_dev - zm_ref) / zm_ref\n",
    "        if full_ratio_range: [vmin, vmax] = [None, None]\n",
    "        else: [vmin, vmax] = [-2, 2]\n",
    "        plot3 = ax3.imshow(zm_fracdiff, vmin=vmin, vmax=vmax, cmap='RdBu_r', extent=extent)\n",
    "        ax3.set_title('Fractional Difference\\n(Dev - Ref)/Ref')\n",
    "        ax3.set_aspect('auto')\n",
    "        ax3.set_xticks(xtick_positions)\n",
    "        ax3.set_xticklabels(xticklabels)\n",
    "        ax3.set_yticks(ytick_positions)\n",
    "        ax3.set_yticklabels(yticklabels)\n",
    "        cb = plt.colorbar(plot3, ax=ax3, orientation='horizontal', pad=0.10)\n",
    "        cb.set_clim(vmin=vmin, vmax=vmax)\n",
    "        cb.set_label('unitless')      \n",
    "            \n",
    "        if savepdf:    \n",
    "            pdf.savefig(figs)\n",
    "            plt.close(figs)\n",
    "            \n",
    "    if savepdf: pdf.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define directories and global information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Path 1 exists: /n/home08/elundgren/GC/testruns/12.0.3/geosfp_4x5_standard\n",
      "Path 2 exists: /n/home08/elundgren/GC/testruns/12.0.3/geosfp_4x5_standard\n"
     ]
    }
   ],
   "source": [
    "# Shared high-level directory\n",
    "testdir = '/n/home08/elundgren/GC/testruns/12.0.3'\n",
    "\n",
    "# Ref and dev run output directories\n",
    "refdir = os.path.join(testdir,'geosfp_4x5_standard')\n",
    "devdir = os.path.join(testdir,'geosfp_4x5_standard')\n",
    "\n",
    "# Ref and dev strings (for use in plots and messages)\n",
    "refstr='12.0.3'\n",
    "devstr='12.0.3'\n",
    "\n",
    "# Set directory to store plots\n",
    "plotsdir = os.path.join(testdir,'plots')\n",
    "\n",
    "# Define cubed sphere resolutions (used for restart file comparison)\n",
    "llres = '4x5'\n",
    "\n",
    "# Check that paths exist\n",
    "check_paths(refdir, devdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Print Ref netcdf filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GEOSChem.StateMet.20160701_0000z.nc4\n",
      "HEMCO_diagnostics.201607010000.nc\n",
      "GEOSChem.DryDep.20160701_0000z.nc4\n",
      "GEOSChem.ConcAfterChem.20160701_0000z.nc4\n",
      "GEOSChem.SpeciesConc.20160701_0000z.nc4\n",
      "GEOSChem.AerosolMass.20160701_0000z.nc4\n",
      "GEOSChem.JValues.20160701_0000z.nc4\n",
      "GEOSChem.Aerosols.20160701_0000z.nc4\n",
      "GEOSChem.CloudConvFlux.20160701_0000z.nc4\n",
      "GEOSChem_restart.201608010000.nc\n",
      "GEOSChem.WetLossConv.20160701_0000z.nc4\n",
      "GEOSChem.WetLossLS.20160701_0000z.nc4\n",
      "GEOSChem.StateChm.20160701_0000z.nc4\n",
      "GEOSChem.ProdLoss.20160701_0000z.nc4\n",
      "HEMCO_restart.201608010000.nc\n",
      "GEOSChem.JValuesLocalNoon.20160701_0000z.nc4\n",
      "GEOSChem.LevelEdgeDiags.20160701_0000z.nc4\n",
      "GEOSChem_restart.201607010000.nc\n"
     ]
    }
   ],
   "source": [
    "reffiles = [k for k in os.listdir(refdir) if '.nc' in k]\n",
    "for k in reffiles:\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Print Dev netcdf filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GEOSChem.StateMet.20160701_0000z.nc4\n",
      "HEMCO_diagnostics.201607010000.nc\n",
      "GEOSChem.DryDep.20160701_0000z.nc4\n",
      "GEOSChem.ConcAfterChem.20160701_0000z.nc4\n",
      "GEOSChem.SpeciesConc.20160701_0000z.nc4\n",
      "GEOSChem.AerosolMass.20160701_0000z.nc4\n",
      "GEOSChem.JValues.20160701_0000z.nc4\n",
      "GEOSChem.Aerosols.20160701_0000z.nc4\n",
      "GEOSChem.CloudConvFlux.20160701_0000z.nc4\n",
      "GEOSChem_restart.201608010000.nc\n",
      "GEOSChem.WetLossConv.20160701_0000z.nc4\n",
      "GEOSChem.WetLossLS.20160701_0000z.nc4\n",
      "GEOSChem.StateChm.20160701_0000z.nc4\n",
      "GEOSChem.ProdLoss.20160701_0000z.nc4\n",
      "HEMCO_restart.201608010000.nc\n",
      "GEOSChem.JValuesLocalNoon.20160701_0000z.nc4\n",
      "GEOSChem.LevelEdgeDiags.20160701_0000z.nc4\n",
      "GEOSChem_restart.201607010000.nc\n"
     ]
    }
   ],
   "source": [
    "devfiles = [k for k in os.listdir(devdir) if '.nc' in k]\n",
    "for k in devfiles:\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Restart Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#reffile_rst = os.path.join(refdir,'GEOSChem_restart.201607010000.nc')\n",
    "#devfile_rst = os.path.join(devdir,'GEOSChem_restart.201607010000.nc')\n",
    "#check_paths(reffile_rst, devfile_rst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#refdata_rst = xr.open_dataset(reffile_rst)\n",
    "#devdata_rst = xr.open_dataset(devfile_rst)\n",
    "#[commonvars, commonvars2D, commonvars3D] = compare_varnames(refdata_rst, refstr, devdata_rst, devstr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get_stats(refdata_rst, refstr, devdata_rst, devstr, 'SPC_CFC11')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#compare_zonal_mean( refdata_rst, refstr, devdata_rst, devstr, llres, varlist=['SPC_O3'] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#compare_single_level( refdata_rst, refstr, devdata_rst, devstr, llres, varlist=['SPC_O3'] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect time-averaged species concentration diagnostic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Path 1 exists: /n/home08/elundgren/GC/testruns/12.0.3/geosfp_4x5_standard/GEOSChem.SpeciesConc.20160701_0000z.nc4\n",
      "Path 2 exists: /n/home08/elundgren/GC/testruns/12.0.3/geosfp_4x5_standard/GEOSChem.SpeciesConc.20160701_0000z.nc4\n"
     ]
    }
   ],
   "source": [
    "day = '20160701'\n",
    "time = '0000'\n",
    "collection = 'SpeciesConc'\n",
    "[refdata, devdata] = get_collection_data(refdir, devdir, collection, day, time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168 common variables (0 are 2-dim and 162 are 3-dim)\n"
     ]
    }
   ],
   "source": [
    "[commonvars, commonvars2D, commonvars3D] = compare_varnames(refdata, refstr, devdata, devstr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['SpeciesConc_CFC11', 'SpeciesConc_CFC113', 'SpeciesConc_CFC114', 'SpeciesConc_CFC115', 'SpeciesConc_CFC12', 'SpeciesConc_HCFC123', 'SpeciesConc_HCFC141b', 'SpeciesConc_HCFC142b', 'SpeciesConc_HCFC22']\n"
     ]
    }
   ],
   "source": [
    "CFCs=[k for k in commonvars if 'CFC' in k]\n",
    "print(CFCs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data units: mol mol-1 dry\n",
      "Array sizes:\n",
      "    12.0.3:  (1, 72, 46, 72)\n",
      "    12.0.3:  (1, 72, 46, 72)\n",
      "Global stats:\n",
      "  Mean:\n",
      "    12.0.3:  1.1334942012775073e-10\n",
      "    12.0.3:  1.1334942012775073e-10\n",
      "  Min:\n",
      "    12.0.3:  0.0\n",
      "    12.0.3:  0.0\n",
      "  Max:\n",
      "    12.0.3:  2.193199927225109e-10\n",
      "    12.0.3:  2.193199927225109e-10\n",
      "  Sum:\n",
      "    12.0.3:  2.7029756893170998e-05\n",
      "    12.0.3:  2.7029756893170998e-05\n"
     ]
    }
   ],
   "source": [
    "varname='SpeciesConc_CFC11'\n",
    "get_stats(refdata, refstr, devdata, devstr, varname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#compare_single_level( refdata, refstr, devdata, devstr, llres, varlist=[varname] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#compare_zonal_mean( refdata, refstr, devdata, devstr, llres, varlist=[varname] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create PDF of plots for list of species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define subset of variables in files\n",
    "#varlist = commonvars3D\n",
    "varlist = [varname]\n",
    "#varlist = [k for k in commonvars3D if 'CFC' in k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Creating /n/home08/elundgren/GC/testruns/12.0.3/plots/SpeciesConc_CFC11_Surface.pdf for 1 variables\n",
      "0 "
     ]
    }
   ],
   "source": [
    "# Surface\n",
    "#pdfname = os.path.join(plotsdir,'{}_Surface.pdf'.format(ref_collection))\n",
    "pdfname = os.path.join(plotsdir,'{}_Surface.pdf'.format(varname))\n",
    "#pdfname = os.path.join(plotsdir,'{}_Surface.pdf'.format('CFCs'))\n",
    "compare_single_level( refdata, refstr, devdata, devstr, llres, varlist=varlist, ilev=0,\n",
    "                     savepdf=True, pdfname=pdfname )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Creating /n/home08/elundgren/GC/testruns/12.0.3/plots/SpeciesConc_CFC11_ZonalMean.pdf for 1 variables\n",
      "0 "
     ]
    }
   ],
   "source": [
    "# Zonal mean\n",
    "#pdfname = os.path.join(plotsdir,'{}_ZonalMean.pdf'.format(ref_collection))\n",
    "pdfname = os.path.join(plotsdir,'{}_ZonalMean.pdf'.format(varname))\n",
    "#pdfname = os.path.join(plotsdir,'{}_ZonalMean.pdf'.format('CFCs'))\n",
    "compare_zonal_mean(refdata, refstr, devdata, devstr, llres, varlist=varlist, \n",
    "                   savepdf=True, pdfname=pdfname )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other diagnostic collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['GEOSChem.StateMet.20160701_0000z.nc4',\n",
       " 'HEMCO_diagnostics.201607010000.nc',\n",
       " 'GEOSChem.DryDep.20160701_0000z.nc4',\n",
       " 'GEOSChem.ConcAfterChem.20160701_0000z.nc4',\n",
       " 'GEOSChem.SpeciesConc.20160701_0000z.nc4',\n",
       " 'GEOSChem.AerosolMass.20160701_0000z.nc4',\n",
       " 'GEOSChem.JValues.20160701_0000z.nc4',\n",
       " 'GEOSChem.Aerosols.20160701_0000z.nc4',\n",
       " 'GEOSChem.CloudConvFlux.20160701_0000z.nc4',\n",
       " 'GEOSChem_restart.201608010000.nc',\n",
       " 'GEOSChem.WetLossConv.20160701_0000z.nc4',\n",
       " 'GEOSChem.WetLossLS.20160701_0000z.nc4',\n",
       " 'GEOSChem.StateChm.20160701_0000z.nc4',\n",
       " 'GEOSChem.ProdLoss.20160701_0000z.nc4',\n",
       " 'HEMCO_restart.201608010000.nc',\n",
       " 'GEOSChem.JValuesLocalNoon.20160701_0000z.nc4',\n",
       " 'GEOSChem.LevelEdgeDiags.20160701_0000z.nc4',\n",
       " 'GEOSChem_restart.201607010000.nc']"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inspect the available files again, but only show those that match\n",
    "commonfiles = [k for k in reffiles if k in devfiles]\n",
    "commonfiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Template for inspecting any collection:\n",
    "\n",
    "# To specify file:\n",
    "#day = {edit}\n",
    "#time = {edit}\n",
    "#collection = {edit}\n",
    "\n",
    "# For retrieving the data:\n",
    "#[refdata, devdata] = get_collection_data(refdir, devdir, collection, day, time)\n",
    "\n",
    "# For inspecting interactively:\n",
    "#refdata\n",
    "#devdata\n",
    "#[commonvars, commonvars2D, commonvars3D] = compare_varnames(refdata, refstr, devdata, devstr)\n",
    "# varname = {edit}\n",
    "#get_stats(refdata, refstr, devdata, devstr, varname)\n",
    "# varlist = {edit}\n",
    "#compare_single_level(refdata, refstr, devdata, devstr, llres, varlist=[varname] )\n",
    "#compare_zonal_mean(refdata, refstr, devdata, devstr, llres, varlist=[varname] )\n",
    "\n",
    "# For creating PDFs:\n",
    "#prefix = {edit}\n",
    "#pdfname = os.path.join(plotsdir,'{}_Surface.pdf'.format(prefix))\n",
    "#compare_single_level(refdata, refstr, devdata, devstr, llres, varlist=varlist, ilev=0,\n",
    "#                     savepdf=True, pdfname=pdfname )\n",
    "#pdfname = os.path.join(plotsdir,'{}_ZonalMean.pdf'.format(prefix))\n",
    "#compare_zonal_mean(refdata, refstr, devdata, devstr, varlist=varlist, llres, \n",
    "#                   savepdf=True, pdfname=pdfname )\n",
    "\n",
    "# For other function options use shift-tab while hovering over the function name in a cell that calls it"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}