{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# NCL_traj_1.py\nThis script illustrates the following concepts:\n   - Plotting a simple trajectory plot\n   - Plotting multiple trajectories in different colors\n   - Plotting every nth time step in a trajectory\n\nSee following URLs to see the reproduced NCL plot & script:\n    - Original NCL script: https://www.ncl.ucar.edu/Applications/Scripts/traj_1.ncl\n    - Original NCL plot: https://www.ncl.ucar.edu/Applications/Images/traj_1_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 geocat.datafiles as gdf\nimport cartopy.crs as ccrs\nimport cartopy.feature as cfeature\nimport matplotlib.pyplot as plt\nimport matplotlib.ticker as ticker\nimport numpy as np\nfrom geocat.viz import 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 load the data into xarrays\nds = xr.open_dataset(gdf.get('netcdf_files/traj_data.nc'))\nsdata = ds.get('sdata')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper function to plot every n-th timestep:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plot_nth_timestep(nparrayy, nparrayx, n):\n\n    for x in range(0, len(nparrayx)):\n\n        # Plot green starting point of each trajectory\n        if x == 0:\n            y, x = nparrayy[x], nparrayx[x]\n            plt.scatter(x, y, color='green', s=1, zorder=2.5)\n\n        # Plot every n-th timestamp\n        if x % n == 0:\n            y, x = nparrayy[x], nparrayx[x]\n            plt.scatter(x, y, color='black', s=1, zorder=2.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Initialize axes\nax = plt.axes(projection=ccrs.PlateCarree())\nax.set_extent([-75, -25, -60, -20], crs=None)\n\n# Set title and subtitle\nplt.suptitle('Example of a Trajectory Plot')\nplt.title('markers every 4th timestep', fontsize=10, pad=10)\n\n# Set land feature and change color to 'lightgrey'\n# See link for extensive list of colors:\n# https://matplotlib.org/3.1.0/gallery/color/named_colors.html\nax.add_feature(cfeature.LAND, color='lightgrey')\n\n# Use geocat.viz.util convenience function to make plots look like NCL plots by using latitude, longitude tick labels\ngvutil.add_lat_lon_ticklabels(ax)\n\n# Use geocat.viz.util convenience function to set axes tick values\ngvutil.set_axes_limits_and_ticks(ax,\n                                 xticks=np.linspace(-70, -30, 5),\n                                 yticks=np.linspace(-60, -20, 5))\n\n# Select trajectories to plot\ntraj = [1, 10, 53, 67, 80]\n\n# Set colors of each trajectory line\ntrajlinecolors = [\"red\", \"blue\", \"green\", \"grey\", \"magenta\"]\n\n# Plot each trajectory\nfor i in range(len(traj)):\n\n    # Extract latitude\n    ypt = (np.array(sdata[1, :, traj[i]]) - 360)\n\n    # Extract longitude\n    xpt = np.array(sdata[2, :, traj[i]])\n\n    plt.plot(ypt, xpt, color=trajlinecolors[i], linewidth=0.4)\n\n    plot_nth_timestep(xpt, ypt, n=4)\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
}