03_MOM6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting started with xgcm for MOM6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* MOM6 variables are staggered according to the Arakawa C-grid \n",
    "* It uses a north-east index convention\n",
    "* center points are labelled (xh, yh) and corner points are labelled (xq, yq)\n",
    "* important: variables xh/yh, xq/yq that are named \"nominal\" longitude/latitude **are not** the true geographical coordinates and are not suitable for plotting (more later)\n",
    "\n",
    "See [indexing](https://mom6.readthedocs.io/en/dev-gfdl/api/generated/pages/Horizontal_indexing.html) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from xgcm import Grid\n",
    "import warnings\n",
    "import matplotlib.pylab as plt\n",
    "from cartopy import crs as ccrs\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "_ = xr.set_options(display_style='text')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial, we use sample data for the $\\frac{1}{2}^{\\circ}$ global model OM4p05 hosted on a GFDL thredds server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataurl = 'http://35.188.34.63:8080/thredds/dodsC/OM4p5/'\n",
    "\n",
    "ds = xr.open_dataset(f'{dataurl}/ocean_monthly_z.200301-200712.nc4',\n",
    "                     chunks={'time':1, 'z_l': 1}, drop_variables=['average_DT',\n",
    "                                                                  'average_T1',\n",
    "                                                                  'average_T2'],\n",
    "                     engine='pydap')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:       (nv: 2, time: 60, xh: 720, xq: 720, yh: 576, yq: 576, z_i: 36, z_l: 35)\n",
       "Coordinates:\n",
       "  * nv            (nv) float64 1.0 2.0\n",
       "  * xh            (xh) float64 -299.8 -299.2 -298.8 -298.2 ... 58.75 59.25 59.75\n",
       "  * xq            (xq) float64 -299.5 -299.0 -298.5 -298.0 ... 59.0 59.5 60.0\n",
       "  * yh            (yh) float64 -77.91 -77.72 -77.54 -77.36 ... 89.47 89.68 89.89\n",
       "  * yq            (yq) float64 -77.82 -77.63 -77.45 -77.26 ... 89.58 89.79 90.0\n",
       "  * z_i           (z_i) float64 0.0 5.0 15.0 25.0 ... 5.75e+03 6.25e+03 6.75e+03\n",
       "  * z_l           (z_l) float64 2.5 10.0 20.0 32.5 ... 5.5e+03 6e+03 6.5e+03\n",
       "  * time          (time) object 2003-01-16 12:00:00 ... 2007-12-16 12:00:00\n",
       "Data variables:\n",
       "    Coriolis      (yq, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    areacello     (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    areacello_bu  (yq, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    areacello_cu  (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    areacello_cv  (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    deptho        (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dxCu          (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dxCv          (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dxt           (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dyCu          (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dyCv          (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    dyt           (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolat        (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolat_c      (yq, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolat_u      (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolat_v      (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolon        (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolon_c      (yq, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolon_u      (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    geolon_v      (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    hfgeou        (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    sftof         (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    thkcello      (z_l, yh, xh) float32 dask.array&lt;chunksize=(1, 576, 720), meta=np.ndarray&gt;\n",
       "    wet           (yh, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    wet_c         (yq, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    wet_u         (yh, xq) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    wet_v         (yq, xh) float32 dask.array&lt;chunksize=(576, 720), meta=np.ndarray&gt;\n",
       "    so            (time, z_l, yh, xh) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    time_bnds     (time, nv) timedelta64[ns] dask.array&lt;chunksize=(1, 2), meta=np.ndarray&gt;\n",
       "    thetao        (time, z_l, yh, xh) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    umo           (time, z_l, yh, xq) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    uo            (time, z_l, yh, xq) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    vmo           (time, z_l, yq, xh) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    vo            (time, z_l, yq, xh) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    volcello      (time, z_l, yh, xh) float32 dask.array&lt;chunksize=(1, 1, 576, 720), meta=np.ndarray&gt;\n",
       "    zos           (time, yh, xh) float32 dask.array&lt;chunksize=(1, 576, 720), meta=np.ndarray&gt;\n",
       "Attributes:\n",
       "    filename:                        ocean_monthly.200301-200712.zos.nc\n",
       "    title:                           OM4p5_IAF_BLING_CFC_abio_csf_mle200\n",
       "    associated_files:                areacello: 20030101.ocean_static.nc\n",
       "    grid_type:                       regular\n",
       "    grid_tile:                       N/A\n",
       "    external_variables:              areacello\n",
       "    DODS_EXTRA.Unlimited_Dimension:  time</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:       (nv: 2, time: 60, xh: 720, xq: 720, yh: 576, yq: 576, z_i: 36, z_l: 35)\n",
       "Coordinates:\n",
       "  * nv            (nv) float64 1.0 2.0\n",
       "  * xh            (xh) float64 -299.8 -299.2 -298.8 -298.2 ... 58.75 59.25 59.75\n",
       "  * xq            (xq) float64 -299.5 -299.0 -298.5 -298.0 ... 59.0 59.5 60.0\n",
       "  * yh            (yh) float64 -77.91 -77.72 -77.54 -77.36 ... 89.47 89.68 89.89\n",
       "  * yq            (yq) float64 -77.82 -77.63 -77.45 -77.26 ... 89.58 89.79 90.0\n",
       "  * z_i           (z_i) float64 0.0 5.0 15.0 25.0 ... 5.75e+03 6.25e+03 6.75e+03\n",
       "  * z_l           (z_l) float64 2.5 10.0 20.0 32.5 ... 5.5e+03 6e+03 6.5e+03\n",
       "  * time          (time) object 2003-01-16 12:00:00 ... 2007-12-16 12:00:00\n",
       "Data variables:\n",
       "    Coriolis      (yq, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    areacello     (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    areacello_bu  (yq, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    areacello_cu  (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    areacello_cv  (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    deptho        (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dxCu          (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dxCv          (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dxt           (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dyCu          (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dyCv          (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    dyt           (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolat        (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolat_c      (yq, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolat_u      (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolat_v      (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolon        (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolon_c      (yq, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolon_u      (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    geolon_v      (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    hfgeou        (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    sftof         (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    thkcello      (z_l, yh, xh) float32 dask.array<chunksize=(1, 576, 720), meta=np.ndarray>\n",
       "    wet           (yh, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    wet_c         (yq, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    wet_u         (yh, xq) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    wet_v         (yq, xh) float32 dask.array<chunksize=(576, 720), meta=np.ndarray>\n",
       "    so            (time, z_l, yh, xh) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    time_bnds     (time, nv) timedelta64[ns] dask.array<chunksize=(1, 2), meta=np.ndarray>\n",
       "    thetao        (time, z_l, yh, xh) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    umo           (time, z_l, yh, xq) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    uo            (time, z_l, yh, xq) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    vmo           (time, z_l, yq, xh) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    vo            (time, z_l, yq, xh) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    volcello      (time, z_l, yh, xh) float32 dask.array<chunksize=(1, 1, 576, 720), meta=np.ndarray>\n",
       "    zos           (time, yh, xh) float32 dask.array<chunksize=(1, 576, 720), meta=np.ndarray>\n",
       "Attributes:\n",
       "    filename:                        ocean_monthly.200301-200712.zos.nc\n",
       "    title:                           OM4p5_IAF_BLING_CFC_abio_csf_mle200\n",
       "    associated_files:                areacello: 20030101.ocean_static.nc\n",
       "    grid_type:                       regular\n",
       "    grid_tile:                       N/A\n",
       "    external_variables:              areacello\n",
       "    DODS_EXTRA.Unlimited_Dimension:  time"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## xgcm grid definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The horizontal dimensions are a combination of (xh or xq) and (yh or yq) corresponding to the staggered point. In the vertical z_l refers to the depth of the center of the layer and z_i to the position of the interfaces, such as len(z_i) = len(z_l) +1.\n",
    "\n",
    "* the geolon/geolat family are the **TRUE** geographical coordinates and are the longitude/latitude you want to use to plot results. The subscript correspond to the staggered point (c: corner, u: U-point, v: V-point, no subscript: center)\n",
    "\n",
    "* the areacello family is the area of the ocean cell at various points with a slightly naming convention (bu: corner, cu: U-point, cv: V-point, no subscript: center). Warning, because of the curvilinear grid:\n",
    "\n",
    "   $$areacello \\neq dxt * dyt$$\n",
    "\n",
    "* the dx/dy family has the following naming convention: dx(Cu: U-point, Cv: V-point, no suffix: center)\n",
    "\n",
    "* thkcello is the layer thickness for each cell (variable). volcello is the volume of the cell, such as:\n",
    "\n",
    "  $$volcello = areacello * thkcello$$\n",
    "\n",
    "\n",
    "The MOM6 output can be written in Symetric (**len(Xq) = len(Xh) + 1**) or Non-symetric mode (**len(Xq) = len(Xh)**), where X is a notation for both x and y. In Symetric mode, one would define the grid for the global as:\n",
    "\n",
    "```python\n",
    "grid = Grid(ds, coords={'X': {'center': 'xh', 'outer': 'xq'},\n",
    "                        'Y': {'center': 'yh', 'outer': 'yq'},\n",
    "                        'Z': {'inner': 'z_l', 'outer': 'z_i'} }, periodic=['X'])\n",
    "```\n",
    "\n",
    "and in Non-symetric mode:\n",
    "\n",
    "```python\n",
    "\n",
    "grid = Grid(ds, coords={'X': {'center': 'xh', 'right': 'xq'},\n",
    "                        'Y': {'center': 'yh', 'right': 'yq'},\n",
    "                        'Z': {'inner': 'z_l', 'outer': 'z_i'} }, periodic=['X'])\n",
    "```\n",
    "\n",
    "Of course, **don't forget to drop the periodic option if you're running a regional model**. Our data is written in Non-symetric mode hence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = Grid(ds, coords={'X': {'center': 'xh', 'right': 'xq'},\n",
    "                        'Y': {'center': 'yh', 'right': 'yq'},\n",
    "                        'Z': {'inner': 'z_l', 'outer': 'z_i'} }, periodic=['X'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A note on geographical coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MOM6 uses land processor elimination, which creates blank holes in the produced geolon/geolat fields. This can result in problems while plotting. It is recommended to overwrite them by the full arrays that are produced by running the model for a few steps without land processor elimination. Here we copy one of these files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 6512k  100 6512k    0     0  13.2M      0 --:--:-- --:--:-- --:--:-- 13.2M\n"
     ]
    }
   ],
   "source": [
    "!curl -O https://raw.githubusercontent.com/raphaeldussin/MOM6-AnalysisCookbook/master/docs/notebooks/data/ocean_grid_sym_OM4_05.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "netcdf ocean_grid_sym_OM4_05 {\n",
      "dimensions:\n",
      "\tyh = 576 ;\n",
      "\txh = 720 ;\n",
      "\tyq = 577 ;\n",
      "\txq = 721 ;\n",
      "variables:\n",
      "\tfloat geolat(yh, xh) ;\n",
      "\t\tgeolat:long_name = \"Latitude of tracer (T) points\" ;\n",
      "\t\tgeolat:units = \"degrees_north\" ;\n",
      "\t\tgeolat:missing_value = 1.e+20f ;\n",
      "\t\tgeolat:_FillValue = 1.e+20f ;\n",
      "\t\tgeolat:cell_methods = \"time: point\" ;\n",
      "\tfloat geolat_c(yq, xq) ;\n",
      "\t\tgeolat_c:long_name = \"Latitude of corner (Bu) points\" ;\n",
      "\t\tgeolat_c:units = \"degrees_north\" ;\n",
      "\t\tgeolat_c:missing_value = 1.e+20f ;\n",
      "\t\tgeolat_c:_FillValue = 1.e+20f ;\n",
      "\t\tgeolat_c:cell_methods = \"time: point\" ;\n",
      "\t\tgeolat_c:interp_method = \"none\" ;\n",
      "\tfloat geolon(yh, xh) ;\n",
      "\t\tgeolon:long_name = \"Longitude of tracer (T) points\" ;\n",
      "\t\tgeolon:units = \"degrees_east\" ;\n",
      "\t\tgeolon:missing_value = 1.e+20f ;\n",
      "\t\tgeolon:_FillValue = 1.e+20f ;\n",
      "\t\tgeolon:cell_methods = \"time: point\" ;\n",
      "\tfloat geolon_c(yq, xq) ;\n",
      "\t\tgeolon_c:long_name = \"Longitude of corner (Bu) points\" ;\n",
      "\t\tgeolon_c:units = \"degrees_east\" ;\n",
      "\t\tgeolon_c:missing_value = 1.e+20f ;\n",
      "\t\tgeolon_c:_FillValue = 1.e+20f ;\n",
      "\t\tgeolon_c:cell_methods = \"time: point\" ;\n",
      "\t\tgeolon_c:interp_method = \"none\" ;\n",
      "\tdouble xh(xh) ;\n",
      "\t\txh:long_name = \"h point nominal longitude\" ;\n",
      "\t\txh:units = \"degrees_east\" ;\n",
      "\t\txh:cartesian_axis = \"X\" ;\n",
      "\tdouble xq(xq) ;\n",
      "\t\txq:long_name = \"q point nominal longitude\" ;\n",
      "\t\txq:units = \"degrees_east\" ;\n",
      "\t\txq:cartesian_axis = \"X\" ;\n",
      "\tdouble yh(yh) ;\n",
      "\t\tyh:long_name = \"h point nominal latitude\" ;\n",
      "\t\tyh:units = \"degrees_north\" ;\n",
      "\t\tyh:cartesian_axis = \"Y\" ;\n",
      "\tdouble yq(yq) ;\n",
      "\t\tyq:long_name = \"q point nominal latitude\" ;\n",
      "\t\tyq:units = \"degrees_north\" ;\n",
      "\t\tyq:cartesian_axis = \"Y\" ;\n",
      "\n",
      "// global attributes:\n",
      "\t\t:filename = \"19000101.ocean_static.nc\" ;\n",
      "\t\t:title = \"OM4_SIS2_cgrid_05\" ;\n",
      "\t\t:grid_type = \"regular\" ;\n",
      "\t\t:grid_tile = \"N/A\" ;\n",
      "\t\t:history = \"Tue Mar  3 13:41:58 2020: ncks -v geolon,geolon_c,geolat,geolat_c /archive/gold/datasets/OM4_05/mosaic_ocean.v20180227.unpacked/ocean_static_sym.nc -o ocean_grid_sym_OM5_05.nc\" ;\n",
      "\t\t:NCO = \"4.0.3\" ;\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "!ncdump -h ocean_grid_sym_OM4_05.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocean_grid_sym = xr.open_dataset('ocean_grid_sym_OM4_05.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:   (xh: 720, xq: 721, yh: 576, yq: 577)\n",
       "Coordinates:\n",
       "  * xh        (xh) float64 -299.8 -299.2 -298.8 -298.2 ... 58.75 59.25 59.75\n",
       "  * xq        (xq) float64 -300.0 -299.5 -299.0 -298.5 ... 58.5 59.0 59.5 60.0\n",
       "  * yh        (yh) float64 -77.91 -77.72 -77.54 -77.36 ... 89.47 89.68 89.89\n",
       "  * yq        (yq) float64 -78.0 -77.82 -77.63 -77.45 ... 89.37 89.58 89.79 90.0\n",
       "Data variables:\n",
       "    geolat    (yh, xh) float32 ...\n",
       "    geolat_c  (yq, xq) float32 ...\n",
       "    geolon    (yh, xh) float32 ...\n",
       "    geolon_c  (yq, xq) float32 ...\n",
       "Attributes:\n",
       "    filename:   19000101.ocean_static.nc\n",
       "    title:      OM4_SIS2_cgrid_05\n",
       "    grid_type:  regular\n",
       "    grid_tile:  N/A\n",
       "    history:    Tue Mar  3 13:41:58 2020: ncks -v geolon,geolon_c,geolat,geol...\n",
       "    NCO:        4.0.3</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:   (xh: 720, xq: 721, yh: 576, yq: 577)\n",
       "Coordinates:\n",
       "  * xh        (xh) float64 -299.8 -299.2 -298.8 -298.2 ... 58.75 59.25 59.75\n",
       "  * xq        (xq) float64 -300.0 -299.5 -299.0 -298.5 ... 58.5 59.0 59.5 60.0\n",
       "  * yh        (yh) float64 -77.91 -77.72 -77.54 -77.36 ... 89.47 89.68 89.89\n",
       "  * yq        (yq) float64 -78.0 -77.82 -77.63 -77.45 ... 89.37 89.58 89.79 90.0\n",
       "Data variables:\n",
       "    geolat    (yh, xh) float32 ...\n",
       "    geolat_c  (yq, xq) float32 ...\n",
       "    geolon    (yh, xh) float32 ...\n",
       "    geolon_c  (yq, xq) float32 ...\n",
       "Attributes:\n",
       "    filename:   19000101.ocean_static.nc\n",
       "    title:      OM4_SIS2_cgrid_05\n",
       "    grid_type:  regular\n",
       "    grid_tile:  N/A\n",
       "    history:    Tue Mar  3 13:41:58 2020: ncks -v geolon,geolon_c,geolat,geol...\n",
       "    NCO:        4.0.3"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ocean_grid_sym"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I used here a symetric grid to highlight the differences with the non-symetric. Since MOM6 uses the north-east convention, we can obtain the non-symetric grid from the symetric by removing the first row and column in our arrays.\n",
    "This overwrites our \"gruyere\" coordinates in our Non-symetric dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds['geolon_c'] = xr.DataArray(data=ocean_grid_sym['geolon_c'][1:,1:], dims=('yq', 'xq'))\n",
    "ds['geolat_c'] = xr.DataArray(data=ocean_grid_sym['geolat_c'][1:,1:], dims=('yq', 'xq'))\n",
    "\n",
    "ds['geolon'] = xr.DataArray(data=ocean_grid_sym['geolon'], dims=('yh', 'xh'))\n",
    "ds['geolat'] = xr.DataArray(data=ocean_grid_sym['geolat'], dims=('yh', 'xh'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vorticity computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can borrow the expression for vorticity from the MITgcm example and adapt it for MOM6:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "vorticity = ( - grid.diff(ds.uo * ds.dxCu, 'Y', boundary='fill')\n",
    "              + grid.diff(ds.vo * ds.dyCv, 'X', boundary='fill') ) / ds.areacello_bu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.DataArray (time: 60, z_l: 35, yq: 576, xq: 720)&gt;\n",
       "dask.array&lt;truediv, shape=(60, 35, 576, 720), dtype=float32, chunksize=(1, 1, 575, 719), chunktype=numpy.ndarray&gt;\n",
       "Coordinates:\n",
       "  * time     (time) object 2003-01-16 12:00:00 ... 2007-12-16 12:00:00\n",
       "  * z_l      (z_l) float64 2.5 10.0 20.0 32.5 ... 5e+03 5.5e+03 6e+03 6.5e+03\n",
       "  * yq       (yq) float64 -77.82 -77.63 -77.45 -77.26 ... 89.37 89.58 89.79 90.0\n",
       "  * xq       (xq) float64 -299.5 -299.0 -298.5 -298.0 ... 58.5 59.0 59.5 60.0</pre>"
      ],
      "text/plain": [
       "<xarray.DataArray (time: 60, z_l: 35, yq: 576, xq: 720)>\n",
       "dask.array<truediv, shape=(60, 35, 576, 720), dtype=float32, chunksize=(1, 1, 575, 719), chunktype=numpy.ndarray>\n",
       "Coordinates:\n",
       "  * time     (time) object 2003-01-16 12:00:00 ... 2007-12-16 12:00:00\n",
       "  * z_l      (z_l) float64 2.5 10.0 20.0 32.5 ... 5e+03 5.5e+03 6e+03 6.5e+03\n",
       "  * yq       (yq) float64 -77.82 -77.63 -77.45 -77.26 ... 89.37 89.58 89.79 90.0\n",
       "  * xq       (xq) float64 -299.5 -299.0 -298.5 -298.0 ... 58.5 59.0 59.5 60.0"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vorticity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take the surface relative vorticity at the first time record:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_plot = 1e5 * vorticity.isel(time=0, z_l=0)\n",
    "_ = data_plot.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we want to be careful and make sure we use the right set of coordinates (geolon_c/geolat_c). Since they are not present in the DataArray, we can add them easily with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_plot = data_plot.assign_coords({'geolon_c': ds['geolon_c'],\n",
    "                                     'geolat_c': ds['geolat_c']})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing worth noting is that geolon_c is not monotonic in the uppermost row. Hence this row needs to be removed for cartopy to properly plot. Another option is to subsample `x` in the MOM6 supergrid, usually named `ocean_hgrid.nc`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_plot = data_plot.isel(xq=slice(0,-1), yq=slice(0,-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define a function that will produce a publication-quality plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subplot_kws=dict(projection=ccrs.PlateCarree(),\n",
    "                 facecolor='grey')\n",
    "\n",
    "plt.figure(figsize=[12,8])\n",
    "p = data_plot.plot(x='geolon_c', y='geolat_c',\n",
    "                   vmin=-1, vmax=1,\n",
    "                   cmap='bwr',\n",
    "                   subplot_kws=subplot_kws,\n",
    "                   transform=ccrs.PlateCarree(),\n",
    "                   add_labels=False,\n",
    "                   add_colorbar=False)\n",
    "\n",
    "# add separate colorbar\n",
    "cb = plt.colorbar(p, ticks=[-1,-0.5,0,0.5,1], shrink=0.6)\n",
    "cb.ax.tick_params(labelsize=18)\n",
    "\n",
    "# optional add grid lines\n",
    "p.axes.gridlines(color='black', alpha=0.5, linestyle='--')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please look at the MITgcm examples for more about what xgcm can do. Also for MOM6 analysis examples using xarray and its companion software, please visit the [MOM6 Analysis Cookbook](https://mom6-analysiscookbook.readthedocs.io/en/latest/index.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}