{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Interpolate a meridional section of CROCO outputs to regular depths\n\nThis notebook, we show:\n\n* how to compute the depths from s-coordinates,\n* how to easily find the name of variables and coordinates,\n* how to interpolate a 3D field with varying depths to regular depths.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialisations\n\nImport needed modules.\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 cmocean\nimport xoa\nfrom xoa.regrid import regrid1d\n\nxr.set_options(display_style=\"text\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Register the :meth:`xarray.Dataset.decode_sigma` callable accessor.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "xoa.register_accessors(decode_sigma=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `xoa <accessors>` accessor is also registered by default, and give access\nto most of the fonctionalities of the other accessors.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Read the model\nThis sample is a meridional extraction of a full 3D CROCO output.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds = xoa.open_data_sample(\"croco.south-africa.meridional.nc\")\nprint(ds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute depths from s-coordinates\n\nDecode the dataset according to the CF conventions:\n\n1. Find sigma terms\n2. Compute depths\n3. Assign depths as coordinates\n\nNote that the :meth:`xarray.Dataset.decode_sigma` callable accessor\ncalls the :func:`xoa.sigma.decode_cf_sigma` function.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds = ds.decode_sigma()\nprint(ds.depth)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Find coordinate names from CF conventions\n\nThe `depth` was assigned as coordinates at the previous stage.\nWe use the `xoa <accessors.dataset>` accessor to easily access the temperature, latitude and depth arrays.\nThe default configuration exposes shortcuts for some variables and coordinates\nas shown in :cfsec:`accessors`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "temp = ds.xoa.temp.squeeze()\nlat_name = temp.xoa.lat.name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Interpolate at regular depths\n\nWe interpolate the temperature array from irregular to regular depths.\n\nLet's create the output depths.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "depth = xr.DataArray(np.linspace(ds.depth.min(), ds.depth.max(), 100),\n                     name=\"depth\", dims=\"depth\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's interpolate the temperature.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tempz = regrid1d(temp, depth)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plots\n\nMake a basic comparison plots.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(ncols=2, sharex=True, sharey=True, figsize=(10, 4))\nkw = dict(levels=np.arange(0, 23))\ntemp.plot.contourf(lat_name, \"depth\", cmap=\"cmo.thermal\", ax=axs[0], **kw)\ntemp.plot.contour(lat_name, \"depth\", colors='w', linewidths=.3, ax=axs[0], **kw)\ntempz.plot.contourf(lat_name, \"depth\", cmap=\"cmo.thermal\", ax=axs[1], **kw)\ntempz.plot.contour(lat_name, \"depth\", colors='w', linewidths=.3, ax=axs[1], **kw);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Et voil\u00e0!\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.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}