Skip to content

Instantly share code, notes, and snippets.

@angus-g

angus-g/xgcm demonstration.ipynb

Created November 16, 2018 00:22
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save angus-g/d4556cfbb89e39446737926ed8ef3f7b to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save angus-g/d4556cfbb89e39446737926ed8ef3f7b to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment