{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_sat_2.py\nThis script illustrates the following concepts:\n   - Converting float data into short data\n   - Drawing filled contours over a satellite map\n   - Explicitly setting contour fill colors\n   - Finding local high pressure values\n\nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/sat_2.ncl\n    - Original NCL plot: https://www.ncl.ucar.edu/Applications/Images/sat_2_lg.png\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import packages:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import xarray as xr\nimport cartopy.crs as ccrs\nimport cartopy.feature as cfeature\nimport numpy as np\nfrom sklearn.cluster import DBSCAN\nimport matplotlib.pyplot as plt\nfrom matplotlib import colors\nimport matplotlib.ticker as mticker\nimport warnings\n\nimport geocat.datafiles as gdf\nimport geocat.viz.util as gvutil"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read in data:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Open a netCDF data file using xarray default engine and\n# load the data into xarrays\nds = xr.open_dataset(gdf.get(\"netcdf_files/slp.1963.nc\"), decode_times=False)\n\n# Get data from the 21st timestep\npressure = ds.slp[21, :, :]\n\n# Translate float values to short values\npressure = pressure.astype('float32')\n\n# Convert Pa to hPa data\npressure = pressure * 0.01\n\n# Fix the artifact of not-shown-data around 0 and 360-degree longitudes\nwrap_pressure = gvutil.xr_add_cyclic_longitudes(pressure, \"lon\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def findLocalExtrema(da, highVal=0, lowVal=1000, eType='Low'):\n    \"\"\"\n    Utility function to find local low/high field variable coordinates on a contour map. To classify as a local high, the data\n    point must be greater than highVal, and to classify as a local low, the data point must be less than lowVal.\n    Args:\n        da: (:class:`xarray.DataArray`):\n            Xarray data array containing the lat, lon, and field variable (ex. pressure) data values\n        highVal (:class:`int`):\n            Data value that the local high must be greater than to qualify as a \"local high\" location.\n            Default highVal is 0.\n        lowVal (:class:`int`):\n            Data value that the local low must be less than to qualify as a \"local low\" location.\n            Default lowVal is 1000.\n        eType (:class:`str`):\n            'Low' or 'High'\n            Determines which extrema are being found- minimum or maximum, respectively.\n            Default eType is 'Low'.\n    Returns:\n        clusterExtremas (:class:`list`):\n            List of coordinate tuples in GPS form (lon in degrees, lat in degrees)\n            that specify local low/high locations\n    \"\"\"\n\n    # Create a 2D array of coordinates in the same shape as the field variable data\n    # so each coordinate is easily mappable to a data value\n    # ex:\n    # (1, 1), (2, 1), (3, 1)\n    # (1, 2)................\n    # (1, 3)................\n    lons, lats = np.meshgrid(np.array(da.lon), np.array(da.lat))\n    coordarr = np.dstack((lons, lats))\n\n    # Find all zeroes that also qualify as low or high values\n    extremacoords = []\n\n    if eType == 'Low':\n        coordlist = np.argwhere(da.data < lowVal)\n        extremacoords = [tuple(coordarr[x[0]][x[1]]) for x in coordlist]\n    if eType == 'High':\n        coordlist = np.argwhere(da.data > highVal)\n        extremacoords = [tuple(coordarr[x[0]][x[1]]) for x in coordlist]\n\n    if extremacoords == []:\n        if eType == 'Low':\n            warnings.warn(\n                'No local extrema with data value less than given lowVal')\n            return []\n        if eType == 'High':\n            warnings.warn(\n                'No local extrema with data value greater than given highVal')\n            return []\n\n    # Clean up noisy data to find actual extrema\n\n    # Use Density-based spatial clustering of applications with noise\n    # to cluster and label coordinates\n    db = DBSCAN(eps=10, min_samples=1)\n    new = db.fit(extremacoords)\n    labels = new.labels_\n\n    # Create an dictionary of values with key being coordinate\n    # and value being cluster label.\n    coordsAndLabels = {label: [] for label in labels}\n    for label, coord in zip(labels, extremacoords):\n        coordsAndLabels[label].append(coord)\n\n    # Initialize array of coordinates to be returned\n    clusterExtremas = []\n\n    # Iterate through the coordinates in each cluster\n    for key in coordsAndLabels:\n\n        # Create array to hold all the field variable values for that cluster\n        datavals = []\n        for coord in coordsAndLabels[key]:\n\n            # Find pressure data at that coordinate\n            cond = np.logical_and(coordarr[:, :, 0] == coord[0],\n                                  coordarr[:, :, 1] == coord[1])\n            x, y = np.where(cond)\n            datavals.append(da.data[x[0]][y[0]])\n\n        # Find the index of the smallest/greatest field variable value of each cluster\n        if eType == 'Low':\n            index = np.argmin(np.array(datavals))\n        if eType == 'High':\n            index = np.argmax(np.array(datavals))\n\n        # Append the coordinate corresponding to that index to the array to be returned\n        clusterExtremas.append(\n            (coordsAndLabels[key][index][0], coordsAndLabels[key][index][1]))\n\n    return clusterExtremas"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plotCLabels(ax,\n                contours,\n                transform,\n                proj,\n                clabel_locations=[],\n                fontsize=12,\n                whitebbox=False,\n                horizontal=False):\n    \"\"\"\n    Utility function to plot contour labels by passing in a coordinate to the clabel function.\n    This allows the user to specify the exact locations of the labels, rather than having matplotlib\n    plot them automatically.\n    This function is exemplified in the python version of https://www.ncl.ucar.edu/Applications/Images/sat_1_lg.png\n    Args:\n        ax (:class:`matplotlib.pyplot.axis`):\n            Axis containing the contour set.\n        contours (:class:`cartopy.mpl.contour.GeoContourSet`):\n            Contour set that is being labeled.\n        transform (:class:`cartopy._crs`):\n            Instance of CRS that represents the source coordinate system of coordinates.\n            (ex. ccrs.Geodetic()).\n        proj (:class:`cartopy.crs`):\n            Projection 'ax' is defined by.\n            This is the instance of CRS that the coordinates will be transformed to.\n        clabel_locations (:class:`list`):\n            List of coordinate tuples in GPS form (lon in degrees, lat in degrees)\n            that specify where the contours with regular field variable values should be plotted.\n        fontsize (:class:`int`):\n            Font size of contour labels.\n        whitebbox (:class:`bool`):\n            Setting this to \"True\" will cause all labels to be plotted with white backgrounds\n        horizontal (:class:`bool`):\n            Setting this to \"True\" will cause the contour labels to be horizontal.\n    Returns:\n        cLabels (:class:`list`):\n            List of text instances of all contour labels\n    \"\"\"\n\n    # Initialize empty array that will be filled with contour label text objects and returned\n    cLabels = []\n\n    # Plot any regular contour levels\n    if clabel_locations != []:\n        clevelpoints = proj.transform_points(\n            transform, np.array([x[0] for x in clabel_locations]),\n            np.array([x[1] for x in clabel_locations]))\n        transformed_locations = [(x[0], x[1]) for x in clevelpoints]\n        ax.clabel(contours,\n                  manual=transformed_locations,\n                  inline=True,\n                  fontsize=fontsize,\n                  colors='black',\n                  fmt=\"%.0f\")\n        [cLabels.append(txt) for txt in contours.labelTexts]\n\n        if horizontal is True:\n            [txt.set_rotation('horizontal') for txt in contours.labelTexts]\n\n    if whitebbox is True:\n        [\n            txt.set_bbox(dict(facecolor='white', edgecolor='none', pad=2))\n            for txt in cLabels\n        ]\n\n    return cLabels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plotELabels(transform,\n                proj,\n                da,\n                clabel_locations=[],\n                label='L',\n                fontsize=22,\n                whitebbox=False,\n                horizontal=True):\n    \"\"\"\n    Utility function to plot contour labels. High/Low contour labels will be plotted using text boxes for more accurate label values\n    and placement.\n    This function is exemplified in the python version of https://www.ncl.ucar.edu/Applications/Images/sat_1_lg.png\n    Args:\n        da: (:class:`xarray.DataArray`):\n            Xarray data array containing the lat, lon, and field variable data values.\n        transform (:class:`cartopy._crs`):\n            Instance of CRS that represents the source coordinate system of coordinates.\n            (ex. ccrs.Geodetic()).\n        proj (:class:`cartopy.crs`):\n            Projection 'ax' is defined by.\n            This is the instance of CRS that the coordinates will be transformed to.\n        clabel_locations (:class:`list`):\n            List of coordinate tuples in GPS form (lon in degrees, lat in degrees)\n            that specify where the contour labels should be plotted.\n        label (:class:`str`):\n            ex. 'L' or 'H'\n            The data value will be plotted as a subscript of this label.\n        fontsize (:class:`int`):\n            Font size of regular contour labels.\n        horizontal (:class:`bool`):\n            Setting this to \"True\" will cause the contour labels to be horizontal.\n        whitebbox (:class:`bool`):\n            Setting this to \"True\" will cause all labels to be plotted with white backgrounds\n    Returns:\n        extremaLabels (:class:`list`):\n            List of text instances of all contour labels\n    \"\"\"\n\n    # Create array of coordinates in the same shape as field variable data\n    # so each coordinate can be easily mapped to its data value.\n    # ex:\n    # (1, 1), (2, 1), (3, 1)\n    # (1, 2)................\n    # (1, 3)................\n    lons, lats = np.meshgrid(np.array(da.lon), np.array(da.lat))\n    coordarr = np.dstack((lons, lats))\n\n    # Initialize empty array that will be filled with contour label text objects and returned\n    extremaLabels = []\n\n    # Plot any low contour levels\n    clabel_points = proj.transform_points(\n        transform, np.array([x[0] for x in clabel_locations]),\n        np.array([x[1] for x in clabel_locations]))\n    transformed_locations = [(x[0], x[1]) for x in clabel_points]\n\n    for x in range(len(transformed_locations)):\n\n        try:\n            # Find field variable data at that coordinate\n            coord = clabel_locations[x]\n            cond = np.logical_and(coordarr[:, :, 0] == coord[0],\n                                  coordarr[:, :, 1] == coord[1])\n            z, y = np.where(cond)\n            p = int(round(da.data[z[0]][y[0]]))\n\n            lab = plt.text(transformed_locations[x][0],\n                           transformed_locations[x][1],\n                           label + \"$_{\" + str(p) + \"}$\",\n                           fontsize=fontsize,\n                           horizontalalignment='center',\n                           verticalalignment='center')\n\n            if horizontal is True:\n                lab.set_rotation('horizontal')\n\n            extremaLabels.append(lab)\n\n        except Exception as E:\n            print(E)\n            continue\n\n    if whitebbox is True:\n        [\n            txt.set_bbox(dict(facecolor='white', edgecolor='none', pad=2))\n            for txt in extremaLabels\n        ]\n\n    return extremaLabels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Create plot\n\n# Set figure size\nfig = plt.figure(figsize=(8, 8))\n\n# Set global axes with an orthographic projection\nproj = ccrs.Orthographic(central_longitude=270, central_latitude=45)\nax = plt.axes(projection=proj)\nax.set_global()\n\n# Add land, coastlines, and ocean features\nax.add_feature(cfeature.LAND, facecolor='lightgray', zorder=1)\nax.add_feature(cfeature.COASTLINE, linewidth=.3, zorder=2)\nax.add_feature(cfeature.OCEAN, facecolor='white')\nax.add_feature(cfeature.BORDERS, linewidth=.3)\nax.add_feature(cfeature.LAKES, facecolor='white', edgecolor='black', linewidth=.3)\n\n# Create color map\ncolorvalues = [1020, 1036, 1500]\ncmap = colors.ListedColormap(['None', 'lightgray', 'dimgrey'])\nnorm = colors.BoundaryNorm(colorvalues, 2)\n\n# Plot contour data\np = wrap_pressure.plot.contourf(ax=ax,\n                                zorder=2,\n                                transform=ccrs.PlateCarree(),\n                                levels=30,\n                                cmap=cmap,\n                                norm=norm,\n                                add_labels=False,\n                                add_colorbar=False)\n\np = wrap_pressure.plot.contour(ax=ax,\n                               transform=ccrs.PlateCarree(),\n                               linewidths=0.3,\n                               levels=30,\n                               cmap='black',\n                               add_labels=False)\n\n# low pressure contour levels- these will be plotted\n# as a subscript to an 'L' symbol.\nlowClevels = findLocalExtrema(pressure, lowVal=995, eType='Low')\nhighClevels = findLocalExtrema(pressure, highVal=1042, eType='High')\n\n# Label regular contours with automatic matplotlib labeling\n# Specify the levels to label every other contour level\nax.clabel(p,\n          levels=np.arange(956, 1064, 8),\n          inline=True,\n          fontsize=12,\n          colors='black',\n          fmt=\"%.0f\")\n\n# Label low and high contours\nplotELabels(ccrs.Geodetic(),\n            proj,\n            wrap_pressure,\n            clabel_locations=lowClevels,\n            label='L')\nplotELabels(ccrs.Geodetic(),\n            proj,\n            wrap_pressure,\n            clabel_locations=highClevels,\n            label='H')\n\n# Use gvutil function to set title and subtitles\ngvutil.set_titles_and_labels(ax,\n                             maintitle=r\"$\\bf{SLP}$\" + \" \" + r\"$\\bf{1963,}$\" +\n                             \" \" + r\"$\\bf{January}$\" + \" \" + r\"$\\bf{24th}$\",\n                             maintitlefontsize=20,\n                             lefttitle=\"mean Daily Sea Level Pressure\",\n                             lefttitlefontsize=16,\n                             righttitle=\"hPa\",\n                             righttitlefontsize=16)\n\n# Set characteristics of text box\nprops = dict(facecolor='white', edgecolor='black', alpha=0.5)\n\n# Place text box\nax.text(0.40,\n        -0.1,\n        'CONTOUR FROM 948 TO 1064 BY 4',\n        transform=ax.transAxes,\n        fontsize=16,\n        bbox=props)\n\n# Add gridlines to axis\ngl = ax.gridlines(color='gray', linestyle='--')\ngl.xlocator = mticker.FixedLocator(np.arange(-180, 180, 20))\ngl.ylocator = mticker.FixedLocator(np.arange(-90, 90, 20))\n\n# Make layout tight\nplt.tight_layout()\n\nplt.show()"
      ]
    }
  ],
  "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.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}