xgcm demonstration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from xgcm import Grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Open dataset with u, v and h data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:     (nv: 2, rho_i: 38, rho_l: 37, xh: 160, xq: 161, yh: 800, yq: 801)\n",
       "Coordinates:\n",
       "  * xh          (xh) float64 0.125 0.375 0.625 0.875 1.125 1.375 1.625 1.875 ...\n",
       "  * yh          (yh) float64 -70.15 -70.07 -69.98 -69.89 -69.81 -69.72 ...\n",
       "  * rho_l       (rho_l) float64 1.03e+03 1.031e+03 1.032e+03 1.033e+03 ...\n",
       "  * rho_i       (rho_i) float64 1.029e+03 1.03e+03 1.032e+03 1.033e+03 ...\n",
       "    Time        float64 2.047e+04\n",
       "  * nv          (nv) float64 1.0 2.0\n",
       "  * xq          (xq) float64 0.0 0.25 0.5 0.75 1.0 1.25 1.5 1.75 2.0 2.25 ...\n",
       "  * yq          (yq) float64 -70.19 -70.11 -70.02 -69.94 -69.85 -69.77 ...\n",
       "Data variables:\n",
       "    h_rho       (rho_l, yh, xh) float32 ...\n",
       "    uh_rho      (rho_l, yh, xq) float32 ...\n",
       "    vh_rho      (rho_l, yq, xh) float32 ...\n",
       "    average_T1  float64 ...\n",
       "    average_T2  float64 ...\n",
       "    average_DT  float64 ...\n",
       "    Time_bnds   (nv) float64 ...\n",
       "Attributes:\n",
       "    filename:          ave_cont_rho.nc\n",
       "    title:             Global ALE Experiment\n",
       "    associated_files:  area_t: ocean_static.nc\n",
       "    grid_type:         regular\n",
       "    grid_tile:         N/A"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = xr.open_dataset('/g/data/v45/ahg157/mom6-archive/tub-spinup/output028/ave_cont_rho.nc',\n",
    "                     decode_times=False).isel(Time=0)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method 1: Define a grid with explicit coordinates (in documentation, but doesn't work)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "__init__() got an unexpected keyword argument 'coords'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-f81961d5ddb3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0;34m'X'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m'center'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'xh'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'outer'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'xq'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0;34m'Y'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m'center'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'yh'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'outer'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'yq'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m     \u001b[0;34m'Z'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m'center'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'rho_l'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'outer'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'rho_i'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m })\n",
      "\u001b[0;31mTypeError\u001b[0m: __init__() got an unexpected keyword argument 'coords'"
     ]
    }
   ],
   "source": [
    "g = Grid(ds, coords={\n",
    "    'X': {'center': 'xh', 'outer': 'xq'},\n",
    "    'Y': {'center': 'yh', 'outer': 'yq'},\n",
    "    'Z': {'center': 'rho_l', 'outer': 'rho_i'}\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method 2: Set grid attributes on dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.coords['xh'].attrs.update(axis='X')\n",
    "ds.coords['xq'].attrs.update(axis='X', c_grid_axis_shift=0.5)\n",
    "ds.coords['yh'].attrs.update(axis='Y')\n",
    "ds.coords['yq'].attrs.update(axis='Y', c_grid_axis_shift=0.5)\n",
    "ds.coords['rho_l'].attrs.update(axis='Z')\n",
    "ds.coords['rho_i'].attrs.update(axis='Z', c_grid_axis_shift=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xgcm.Grid>\n",
       "X Axis (periodic):\n",
       "  * center   xh (160) --> outer\n",
       "  * outer    xq (161) --> center\n",
       "Y Axis (not periodic):\n",
       "  * center   yh (800) --> outer\n",
       "  * outer    yq (801) --> center\n",
       "Z Axis (not periodic):\n",
       "  * center   rho_l (37) --> outer\n",
       "  * outer    rho_i (38) --> center"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = Grid(ds, periodic=['X'])\n",
    "# if periodic isn't specified, it assumes all dimensions are...\n",
    "# there's probably another attribute that modifies that behaviour\n",
    "g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (rho_l: 37, yh: 800, xh: 160)>\n",
       "array([[[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]]], dtype=float32)\n",
       "Coordinates:\n",
       "  * rho_l    (rho_l) float64 1.03e+03 1.031e+03 1.032e+03 1.033e+03 ...\n",
       "  * yh       (yh) float64 -70.15 -70.07 -69.98 -69.89 -69.81 -69.72 -69.64 ...\n",
       "  * xh       (xh) float64 0.125 0.375 0.625 0.875 1.125 1.375 1.625 1.875 ..."
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.interp(ds.uh_rho, 'X')**2 + g.interp(ds.vh_rho, 'Y')**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differencing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray (rho_i: 38, yh: 800, xh: 160)>\n",
       "array([[[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]],\n",
       "\n",
       "       [[nan, nan, ..., nan, nan],\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        ...,\n",
       "        [nan, nan, ...,  0., nan],\n",
       "        [nan, nan, ..., nan, nan]]], dtype=float32)\n",
       "Coordinates:\n",
       "  * rho_i    (rho_i) float64 1.029e+03 1.03e+03 1.032e+03 1.033e+03 ...\n",
       "  * yh       (yh) float64 -70.15 -70.07 -69.98 -69.89 -69.81 -69.72 -69.64 ...\n",
       "  * xh       (xh) float64 0.125 0.375 0.625 0.875 1.125 1.375 1.625 1.875 ..."
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g.diff(ds.h_rho, 'Z', boundary='extend')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-18.10]",
   "language": "python",
   "name": "conda-env-analysis3-18.10-py"
  },
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}