{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Compare Hycom3d with a GDP drifter\n\nIn this notebook, we show :\n\n* how to decode dataset so that it is easy to access generic coordinates and variables,\n* how to compute depths from layer thicknesses,\n* how to interpolate currents from U and V positions to T position on an arakawa C grid,\n* how to perform a 4D interpolation with a variable depth coordinate to random positions,\n* how to make an horizontal slice of a 4D variable with a variable depth coordinate.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialisations\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport xarray as xr\nimport matplotlib.pyplot as plt\nimport cartopy.crs as ccrs\nimport xoa\nfrom xoa.grid import dz2depth, shift\nfrom xoa.regrid import grid2loc, regrid1d\nimport xoa.cf as xcf\nimport xoa.geo as xgeo\nfrom xoa.plot import plot_flow\n\nxr.set_options(display_style=\"text\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Register the `xoa <accessors>` accessors :\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "xoa.register_accessors()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Register the Hycom naming specifications\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "xcf.register_cf_specs(xoa.get_data_sample(\"hycom.cfg\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here is what these CF specifications contain\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with open(xoa.get_data_sample(\"hycom.cfg\")) as f:\n    print(f.read())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that most of these specifications are not needed to decode the dataset, i.e to find\nknown variables and coordinates, since the default specifications help for that.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Read and decode the datasets\n\n### Hycom velocities\n\nU and V are stored in separate files at U and V locations on the Hycom Arakawa C-Grid.\nWe choose to rename the dataset contents so that we can use generic names to access it.\nWe call for that the :meth:`~xarray.Dataset.xoa.decode` accessor method.\nNote that the staggered grid location is automatically appended to horizontal\ncoordinates and dimensions of the U and V velocity components since they are\nplaced at U and V locations. This prevent any conflict with the bathymetry and\nlayer thickness horizontal coordinates and dimensions, which point to the T location.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "U velocity component\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "u_hycom = xoa.open_data_sample(\"hycom.gdp.u.nc\").xoa.decode().u.squeeze(drop=True)\nprint(u_hycom)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "V velocity component\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "v_hycom = xoa.open_data_sample(\"hycom.gdp.v.nc\").xoa.decode().v.squeeze(drop=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Layer thicknesses dataset\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "hycom = xoa.open_data_sample(\"hycom.gdp.h.nc\").xoa.decode().squeeze(drop=True)\nprint(hycom)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now compute the depths from the layer thicknesses at the T location\nthanks to the :func:`xoa.grid.dz2depth` function:\nwe integrate from a null SSH and interpolate the depth from W locations to T locations.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "hycom.coords[\"depth\"] = dz2depth(hycom.dz, centered=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then we interpolate the velocity components to the T location with\nthe :func:`xoa.grid.shift` function.\nThe call to :meth:`reloc <xoa.cf.CFSpecs.reloc>` helps removing the\nstaggered grid location prefixes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ut_hycom = shift(u_hycom, {\"x\": \"left\", \"y\": \"left\"}).xoa.reloc(u=False)\nvt_hycom = shift(v_hycom, {\"x\": \"left\", \"y\": \"left\"}).xoa.reloc(v=False)\nhycom[\"u\"] = ut_hycom.assign_coords(**hycom.coords)\nhycom[\"v\"] = vt_hycom.assign_coords(**hycom.coords)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "So, finally we obtain:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(hycom)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### GDP drifter\n\nThe drifter comes as a `csv` file and we read it as :class:`pandas.DataFrame` instance.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter = xoa.open_data_sample(\"gdp-6203641.csv\", header=0, skiprows=[1], parse_dates=[2], index_col=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Since the sampling is not that nice, we resample it to 3-hour intervals.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter = drifter.resample(\"3H\").mean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We convert it to and :class:`xarray.Dataset`, fix the time coordinate and decode it.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter = drifter.to_xarray().assign_coords(time=drifter.index.values).xoa.decode()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We drop missing values.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter = drifter.where(~drifter.lon.isnull() & ~drifter.lat.isnull(), drop=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We add a constant depth of 15 m.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter.coords[\"depth\"] = drifter.lon*0 + 15"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here is what we obtain.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(drifter)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute and interpolate velocities\n\nWe compute the drifter velocity components\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "drifter[\"u\"] = drifter.lon.differentiate(\"time\", datetime_unit=\"s\") * xgeo.EARTH_RADIUS*np.pi/180\ndrifter[\"u\"] *= np.cos(np.radians(drifter.lat.values))\ndrifter[\"v\"] = drifter.lat.differentiate(\"time\", datetime_unit=\"s\") * xgeo.EARTH_RADIUS*np.pi/180"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We make a 4D linear interpolation the Hycom velocity to the drifter positions with :func:`xoa.regrid.grid2loc`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "uloc = grid2loc(hycom[\"u\"], drifter)\nvloc = grid2loc(hycom[\"v\"], drifter)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Instead of just showing the velocities along the drifter positions,\nwe can plot the mean model velocities over the same period as a background.\nSo we interpolate them at 15 m with :func:`xoa.regrid.regrid1d` and compute a time average.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "d15 = xr.DataArray([15.], dims=\"depth\", name=\"depth\")\nuh15 = regrid1d(hycom[\"u\"], d15).squeeze(drop=True).mean(dim=\"time\")\nvh15 = regrid1d(hycom[\"v\"], d15).squeeze(drop=True).mean(dim=\"time\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot\n\nThe geographic extent is easily computed with the :func:`xoa.geo.get_extent` function.\nThe background currents are plotted with the :func:`xoa.plot.plot_flow` function.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pmerc = ccrs.Mercator()\npcarr = ccrs.PlateCarree()\nfig, ax = plt.subplots(figsize=(7, 7), subplot_kw={\"facecolor\": \"teal\", \"projection\": pmerc})\nax.gridlines(draw_labels=True, dms=True)\nax.set_extent(xgeo.get_extent(uh15))\nkwqv = dict(scale_units=\"dots\", scale=0.1/20, units=\"dots\", transform=pcarr)\nplot_flow(uh15, vh15, axes=ax, transform=pcarr, color=\"w\", alpha=(.3, 1), linewidth=.6)\nqv = ax.quiver(uloc.lon.values, uloc.lat.values, uloc.values, vloc.values,\n    color=\"w\", width=2, label=\"Model\", **kwqv)\nax.plot(drifter.lon.values, drifter.lat.values, '-', color=\"C1\", transform=pcarr, lw=.5)\nax.quiver(drifter.lon.values, drifter.lat.values, drifter.u.values, drifter.v.values,\n    color=\"C1\", label=\"Drifter\", width=2, **kwqv)\nplt.quiverkey(qv, 0.1, 1.06, 0.1, r\"0.1 $m\\,s^{-1}$\", color=\"k\", alpha=1, labelpos=\"E\")\nplt.legend();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The discrepancies between the lagrangian and mean eulerian currents highlight\nthe variability due to the mesocale activity.\nThe differences between the observed and modeled currents are partly due to\nthe lack of data assimilation in the model.\n\n"
      ]
    }
  ],
  "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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}