From 89f4619593d1898e5038196e9e093bf15d755e41 Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Sat, 2 Sep 2023 09:02:27 -0400 Subject: [PATCH] Update xgcm.Grid API - See related comments in https://github.com/hdrake/sectionate/commit/39c30090cde6777956026c1bcfeb4c762ee70bbd - Temporary bug-fix for xhistogram call (see https://github.com/xgcm/xhistogram/issues/16) - Re-ran all notebooks --- examples/benchmarking_transformations.ipynb | 131 +- examples/full3D_tutorial.ipynb | 8540 +------------------ examples/surface2D_tutorial.ipynb | 83 +- xwmt/wm.py | 8 +- xwmt/wmt.py | 2 + 5 files changed, 394 insertions(+), 8370 deletions(-) diff --git a/examples/benchmarking_transformations.ipynb b/examples/benchmarking_transformations.ipynb index 5d0cccd..2016635 100644 --- a/examples/benchmarking_transformations.ipynb +++ b/examples/benchmarking_transformations.ipynb @@ -10,10 +10,19 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 26, "id": "46326d2c-cdd6-47ca-bd4d-f5ee769e1403", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The autoreload extension is already loaded. To reload it, use:\n", + " %reload_ext autoreload\n" + ] + } + ], "source": [ "%load_ext autoreload\n", "%autoreload 2" @@ -21,7 +30,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 27, "id": "d6983c65-69f4-42a5-ba30-7c06813ac830", "metadata": {}, "outputs": [], @@ -44,7 +53,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 28, "id": "74a1e530-2ba7-4fcc-8b2c-c782cfe7c5c4", "metadata": {}, "outputs": [ @@ -56,7 +65,7 @@ "pandas version 1.5.2 \n", "xarray version 2022.12.0 \n", "xhistogram version 0.3.1 \n", - "xgcm version 0.8.2.dev15+g7492277 \n", + "xgcm version 0.1.dev565+g7492277 \n", "xwmt version 0.0.3 \n", "xbudget version 0.0.1\n" ] @@ -87,7 +96,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 29, "id": "129cadec-3814-4bbb-8a3a-0952bf6e3eb8", "metadata": {}, "outputs": [], @@ -109,7 +118,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 30, "id": "39bb6ecb-5629-44a0-bd14-63651ccdf561", "metadata": {}, "outputs": [ @@ -143,7 +152,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 31, "id": "e3151cd7-264c-42bc-b634-ec221287daad", "metadata": {}, "outputs": [ @@ -172,7 +181,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 32, "id": "3f8c5969-8f26-4789-ba20-f7b956117d3e", "metadata": {}, "outputs": [], @@ -185,7 +194,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 33, "id": "29989b0a-b988-4395-a7f1-bbf14c97d2f3", "metadata": {}, "outputs": [], @@ -205,22 +214,22 @@ "metrics = {\n", " ('X','Y'): \"areacello\", # Required for area-integration\n", "}\n", - "grid = xgcm.Grid(ds, coords=coords, metrics=metrics, periodic=False, autoparse_metadata=False)" + "grid = xgcm.Grid(ds, coords=coords, metrics=metrics, boundary={'X':'periodic', 'Y':'periodic'}, autoparse_metadata=False)" ] }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 34, "id": "952682fd-ef28-43c0-87ab-16b75ce11593", "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "" + "" ] }, - "execution_count": 9, + "execution_count": 34, "metadata": {}, "output_type": "execute_result" }, @@ -241,7 +250,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 35, "id": "a75d81f7-e5de-4c66-a838-430f50dacb2a", "metadata": {}, "outputs": [], @@ -255,7 +264,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 36, "id": "283444ba-f497-43ed-85ac-5544447fdff8", "metadata": {}, "outputs": [], @@ -265,7 +274,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 37, "id": "643eb290-cc54-4006-8a66-b9f486fb30c8", "metadata": {}, "outputs": [], @@ -283,7 +292,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 38, "id": "e59e9466-1c3c-46fd-9c03-cc863b0e7c8a", "metadata": {}, "outputs": [ @@ -291,8 +300,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 7.15 s, sys: 1.79 s, total: 8.93 s\n", - "Wall time: 3.92 s\n" + "CPU times: user 6.85 s, sys: 1.16 s, total: 8.01 s\n", + "Wall time: 3.83 s\n" ] }, { @@ -301,7 +310,7 @@ "''" ] }, - "execution_count": 13, + "execution_count": 38, "metadata": {}, "output_type": "execute_result" } @@ -320,7 +329,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 39, "id": "e3fe5d80-5d6a-4d58-8875-48276fac3e49", "metadata": {}, "outputs": [ @@ -328,7 +337,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/work/hfd/.conda/envs/generic-xstack/lib/python3.10/site-packages/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", + "/work2/hfd/codedev/xgcm/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", " out = xr.apply_ufunc(\n" ] }, @@ -336,8 +345,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 58.4 s, sys: 1min 26s, total: 2min 24s\n", - "Wall time: 37.5 s\n" + "CPU times: user 57.2 s, sys: 2min 17s, total: 3min 14s\n", + "Wall time: 38.7 s\n" ] }, { @@ -346,7 +355,7 @@ "''" ] }, - "execution_count": 14, + "execution_count": 39, "metadata": {}, "output_type": "execute_result" } @@ -367,7 +376,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 40, "id": "c46600f5-cc36-448b-8553-63cb2ee8a23c", "metadata": {}, "outputs": [ @@ -375,8 +384,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 1min 22s, sys: 3min 36s, total: 4min 59s\n", - "Wall time: 1min 5s\n" + "CPU times: user 1min 25s, sys: 5min 50s, total: 7min 15s\n", + "Wall time: 1min 18s\n" ] }, { @@ -385,7 +394,7 @@ "''" ] }, - "execution_count": 15, + "execution_count": 40, "metadata": {}, "output_type": "execute_result" } @@ -404,7 +413,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 41, "id": "4a3ed118-ae66-454a-8e7c-76bf95f344ff", "metadata": {}, "outputs": [ @@ -466,7 +475,7 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 42, "id": "dbd7df54-707f-4efc-914c-bbc256fd97b1", "metadata": {}, "outputs": [], @@ -476,7 +485,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 43, "id": "872104a6-a8fb-4417-af96-a3c2828ee92b", "metadata": {}, "outputs": [ @@ -484,7 +493,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/work/hfd/.conda/envs/generic-xstack/lib/python3.10/site-packages/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", + "/work2/hfd/codedev/xgcm/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", " out = xr.apply_ufunc(\n" ] }, @@ -492,8 +501,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 25.4 s, sys: 38.4 s, total: 1min 3s\n", - "Wall time: 14.5 s\n" + "CPU times: user 25.5 s, sys: 41.7 s, total: 1min 7s\n", + "Wall time: 14.8 s\n" ] }, { @@ -502,7 +511,7 @@ "''" ] }, - "execution_count": 18, + "execution_count": 43, "metadata": {}, "output_type": "execute_result" } @@ -522,7 +531,7 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 44, "id": "6fafe3fd-9662-49c0-824e-500dd5f9a8de", "metadata": {}, "outputs": [ @@ -530,8 +539,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 1min 7s, sys: 3min 10s, total: 4min 17s\n", - "Wall time: 55.9 s\n" + "CPU times: user 1min 14s, sys: 3min 40s, total: 4min 55s\n", + "Wall time: 43.9 s\n" ] }, { @@ -540,7 +549,7 @@ "''" ] }, - "execution_count": 19, + "execution_count": 44, "metadata": {}, "output_type": "execute_result" } @@ -559,17 +568,17 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 45, "id": "58d2bf0f-e3d7-49cf-99d8-88da04f6cd7a", "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "" + "" ] }, - "execution_count": 20, + "execution_count": 45, "metadata": {}, "output_type": "execute_result" }, @@ -604,7 +613,7 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 46, "id": "caa304d3-370d-4317-995d-2498b2578c91", "metadata": {}, "outputs": [ @@ -612,7 +621,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/work/hfd/.conda/envs/generic-xstack/lib/python3.10/site-packages/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", + "/work2/hfd/codedev/xgcm/xgcm/transform.py:245: FutureWarning: ``output_sizes`` should be given in the ``dask_gufunc_kwargs`` parameter. It will be removed as direct parameter in a future version.\n", " out = xr.apply_ufunc(\n" ] }, @@ -620,8 +629,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 1min 57s, sys: 2min 57s, total: 4min 54s\n", - "Wall time: 1min 32s\n" + "CPU times: user 1min 57s, sys: 3min 20s, total: 5min 17s\n", + "Wall time: 1min 17s\n" ] }, { @@ -630,7 +639,7 @@ "''" ] }, - "execution_count": 21, + "execution_count": 46, "metadata": {}, "output_type": "execute_result" } @@ -650,7 +659,7 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 47, "id": "4d1ec453-f114-4f39-88db-aeb151a77b11", "metadata": {}, "outputs": [ @@ -658,8 +667,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "CPU times: user 2min 30s, sys: 5min 41s, total: 8min 12s\n", - "Wall time: 2min 7s\n" + "CPU times: user 2min 29s, sys: 6min 3s, total: 8min 32s\n", + "Wall time: 1min 42s\n" ] }, { @@ -668,7 +677,7 @@ "''" ] }, - "execution_count": 22, + "execution_count": 47, "metadata": {}, "output_type": "execute_result" } @@ -695,7 +704,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": 48, "id": "913ff315-3b16-48a8-847e-31eb76303946", "metadata": {}, "outputs": [], @@ -705,7 +714,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": 49, "id": "87715a72-6e98-49a9-9b7f-0c8d622f0096", "metadata": {}, "outputs": [ @@ -1080,18 +1089,18 @@ " 1, 1, 2, 3, 3, 4, 3, 4, 3, 8, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n", " 0, 0, 0, 0, 0, 0, 0, 0, 0])\n", "Coordinates:\n", - " * tos_bin (tos_bin) float64 -3.5 -2.5 -1.5 -0.5 0.5 ... 45.5 46.5 47.5 48.5
    • " ], "text/plain": [ "\n", @@ -1102,7 +1111,7 @@ " * tos_bin (tos_bin) float64 -3.5 -2.5 -1.5 -0.5 0.5 ... 45.5 46.5 47.5 48.5" ] }, - "execution_count": 24, + "execution_count": 49, "metadata": {}, "output_type": "execute_result" } @@ -1113,7 +1122,7 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": 50, "id": "66dac032-19e7-48a7-8417-c328a5dac8f7", "metadata": {}, "outputs": [ diff --git a/examples/full3D_tutorial.ipynb b/examples/full3D_tutorial.ipynb index 613b94e..39e6ade 100644 --- a/examples/full3D_tutorial.ipynb +++ b/examples/full3D_tutorial.ipynb @@ -2,8184 +2,172 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, - "id": "182a8d3a-09f8-448e-a029-d03db69bbe44", - "metadata": {}, - "outputs": [], - "source": [ - "%load_ext autoreload\n", - "%autoreload 2" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "99c3a046-9435-46a6-b9c2-81becd029d93", - "metadata": {}, - "outputs": [], - "source": [ - "import warnings\n", - "import numpy as np\n", - "from dask.diagnostics import ProgressBar\n", - "import xarray as xr\n", - "import xhistogram\n", - "import xgcm\n", - "from xgcm import Grid\n", - "import matplotlib.pyplot as plt" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "74b06942-f576-4742-8339-0a4032cdb7ce", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "numpy version 1.22.4 \n", - "xarray version 2022.12.0 \n", - "xhistogram version 0.3.1 \n", - "xgcm version 0.8.2.dev15+g7492277\n" - ] - } - ], - "source": [ - "print('numpy version', np.__version__, '\\nxarray version', xr.__version__, '\\nxhistogram version', xhistogram.__version__, '\\nxgcm version', xgcm.__version__)" - ] - }, - { - "cell_type": "markdown", - "id": "5fa6652e-9d94-4870-8f33-037ae4d593e5", - "metadata": {}, - "source": [ - "# Tutorial: evaluating water mass transformations from finite-volume tracer budget terms" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "0b71a055-e428-4a25-9edb-45ea25ce9b95", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "xwmt version 0.0.3\n" - ] - } - ], - "source": [ - "import xwmt\n", - "print('xwmt version', xwmt.__version__)" - ] - }, - { - "cell_type": "markdown", - "id": "2f73cabd-5dab-44f4-91de-9d8c6a0d6f92", - "metadata": {}, - "source": [ - "### Load grids and data" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "aa99ea92-986b-47d4-a387-1ab5dd951056", - "metadata": {}, - "outputs": [], - "source": [ - "dt = 'daily' # choose from ['monthly', 'daily', 'hourly']\n", - "gridname = 'natv' # choose from ['zstr', 'natv', 'rho2']\n", - "# Caution: non-natv diagnostics are currently broken--see https://github.com/NCAR/MOM6/pull/215\n", - "\n", - "# Load data\n", - "sim = \"wmt_incsurffluxes.natv_rho2_zstr.monthly_daily_hourly.13months\"\n", - "rootdir = f\"/archive/Graeme.Macgilchrist/MOM6-examples/ice_ocean_SIS2/Baltic_OM4_025/{sim}/Baltic_OM4_025/\"\n", - "prefix = '19000101.ocean_'+dt+'_'\n", - "time = \"190*\"\n", - "\n", - "# Merge diagnostics that were saved into different files into one main dataset\n", - "suffixs = ['surf','thck','heat','salt','xtra']\n", - "Zprefixes = {'rho2':'rho2_', 'zstr':'z_', 'natv':'z'}\n", - "Zprefix = Zprefixes[gridname]\n", - "ds = xr.Dataset()\n", - "for suffix in suffixs:\n", - " if suffix == \"surf\":\n", - " filename = prefix+suffix+'_'+time+'.nc'\n", - " else:\n", - " filename = prefix+gridname+'_'+suffix+'_'+time+'.nc'\n", - " dsnow = xr.open_mfdataset(rootdir+filename, chunks={'time':1})\n", - " ds = xr.merge([ds,dsnow])\n", - "\n", - "# Select year-long subset\n", - "ds = ds.sel(time=slice('1900-02-01 00', '1901-02-01 00'))\n", - "\n", - "# Load grid\n", - "oceangridname = '19000101.ocean_static.nc'\n", - "ocean_grid = xr.open_dataset(rootdir+oceangridname).squeeze()\n", - "\n", - "# Some renaming of variables to match conventions\n", - "ocean_grid = ocean_grid.rename({'depth_ocean':'deptho'})\n", - "ds = ds.rename({'temp':'thetao', 'salt':'so'})\n", - "ds['sigma2'] = ds['rhopot2']-1000." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "d23a62b4-ed21-4116-8f91-ae9729b57937", - "metadata": {}, - "outputs": [], - "source": [ - "# Add core coordinates of ocean_grid to ds\n", - "ds = ds.assign_coords({\n", - " \"areacello\": xr.DataArray(ocean_grid[\"areacello\"].values, dims=('yh', 'xh',)), # Required for area-integration\n", - " \"lon\": xr.DataArray(ocean_grid[\"geolon\"].values, dims=('yh', 'xh',)), # Required for calculating density if not already provided!\n", - " \"lat\": xr.DataArray(ocean_grid[\"geolat\"].values, dims=('yh', 'xh',)), # Required for calculating density if not already provided!\n", - "})\n", - "\n", - "# xgcm grid for dataset\n", - "coords = {\n", - " 'X': {'center': 'xh', 'right': 'xq'},\n", - " 'Y': {'center': 'yh', 'right': 'yq'},\n", - " 'Z': {'center': f'{Zprefix}l', 'outer': f'{Zprefix}i'}\n", - "}\n", - "metrics = {\n", - " ('X','Y'): \"areacello\", # Required for area-integration\n", - "}\n", - "grid = Grid(ds.copy(), coords=coords, metrics=metrics, periodic=None, autoparse_metadata=False)" - ] - }, - { - "cell_type": "markdown", - "id": "e3d2b510-a73e-482e-a335-751ffc8291c1", - "metadata": {}, - "source": [ - "**Note:** In addition to `grid` containing information about the static model grid, its attribute `grid._ds` points to the full dataset containing model diagnostics!" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "f176fd1e-c8d1-4722-a1ea-8b0e12d099c7", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
    \n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "
    <xarray.Dataset>\n",
-       "Dimensions:                         (xh: 97, yh: 105, time: 365, nv: 2, zl: 75,\n",
-       "                                     zi: 76, xq: 98, yq: 106)\n",
-       "Coordinates:\n",
-       "  * xh                              (xh) float64 6.125 6.375 ... 29.88 30.12\n",
-       "  * yh                              (yh) float64 53.01 53.16 ... 65.48 65.56\n",
-       "  * time                            (time) object 1900-02-01 12:00:00 ... 190...\n",
-       "  * nv                              (nv) float64 1.0 2.0\n",
-       "  * zl                              (zl) float64 1.0 3.0 ... 3.591e+03 5.111e+03\n",
-       "  * zi                              (zi) float64 0.0 2.0 ... 3.722e+03 6.5e+03\n",
-       "  * xq                              (xq) float64 6.0 6.25 6.5 ... 30.0 30.25\n",
-       "  * yq                              (yq) float64 52.94 53.09 ... 65.52 65.6\n",
-       "    areacello                       (yh, xh) float32 2.806e+08 ... 7.306e+07\n",
-       "    lon                             (yh, xh) float32 6.125 6.375 ... 29.96 30.21\n",
-       "    lat                             (yh, xh) float32 53.01 53.01 ... 65.08 65.07\n",
-       "Data variables: (12/61)\n",
-       "    tos                             (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    sos                             (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    hfds                            (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    hflso                           (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    hfsso                           (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    rlntds                          (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n",
-       "    ...                              ...\n",
-       "    drho_dT                         (time, zl, yh, xh) float32 dask.array<chunksize=(1, 75, 105, 97), meta=np.ndarray>\n",
-       "    drho_dS                         (time, zl, yh, xh) float32 dask.array<chunksize=(1, 75, 105, 97), meta=np.ndarray>\n",
-       "    umo                             (time, zl, yh, xq) float32 dask.array<chunksize=(1, 75, 105, 98), meta=np.ndarray>\n",
-       "    vmo                             (time, zl, yq, xh) float32 dask.array<chunksize=(1, 75, 106, 97), meta=np.ndarray>\n",
-       "    e                               (time, zi, yh, xh) float32 dask.array<chunksize=(1, 76, 105, 97), meta=np.ndarray>\n",
-       "    sigma2                          (time, zl, yh, xh) float32 dask.array<chunksize=(1, 75, 105, 97), meta=np.ndarray>
    " - ], - "text/plain": [ - "\n", - "Dimensions: (xh: 97, yh: 105, time: 365, nv: 2, zl: 75,\n", - " zi: 76, xq: 98, yq: 106)\n", - "Coordinates:\n", - " * xh (xh) float64 6.125 6.375 ... 29.88 30.12\n", - " * yh (yh) float64 53.01 53.16 ... 65.48 65.56\n", - " * time (time) object 1900-02-01 12:00:00 ... 190...\n", - " * nv (nv) float64 1.0 2.0\n", - " * zl (zl) float64 1.0 3.0 ... 3.591e+03 5.111e+03\n", - " * zi (zi) float64 0.0 2.0 ... 3.722e+03 6.5e+03\n", - " * xq (xq) float64 6.0 6.25 6.5 ... 30.0 30.25\n", - " * yq (yq) float64 52.94 53.09 ... 65.52 65.6\n", - " areacello (yh, xh) float32 2.806e+08 ... 7.306e+07\n", - " lon (yh, xh) float32 6.125 6.375 ... 29.96 30.21\n", - " lat (yh, xh) float32 53.01 53.01 ... 65.08 65.07\n", - "Data variables: (12/61)\n", - " tos (time, yh, xh) float32 dask.array\n", - " sos (time, yh, xh) float32 dask.array\n", - " hfds (time, yh, xh) float32 dask.array\n", - " hflso (time, yh, xh) float32 dask.array\n", - " hfsso (time, yh, xh) float32 dask.array\n", - " rlntds (time, yh, xh) float32 dask.array\n", - " ... ...\n", - " drho_dT (time, zl, yh, xh) float32 dask.array\n", - " drho_dS (time, zl, yh, xh) float32 dask.array\n", - " umo (time, zl, yh, xq) float32 dask.array\n", - " vmo (time, zl, yq, xh) float32 dask.array\n", - " e (time, zi, yh, xh) float32 dask.array\n", - " sigma2 (time, zl, yh, xh) float32 dask.array" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" + "execution_count": 29, + "id": "182a8d3a-09f8-448e-a029-d03db69bbe44", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The autoreload extension is already loaded. To reload it, use:\n", + " %reload_ext autoreload\n" + ] + } + ], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "id": "99c3a046-9435-46a6-b9c2-81becd029d93", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "import numpy as np\n", + "from dask.diagnostics import ProgressBar\n", + "import xarray as xr\n", + "import xhistogram\n", + "import xgcm\n", + "from xgcm import Grid\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "74b06942-f576-4742-8339-0a4032cdb7ce", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "numpy version 1.22.4 \n", + "xarray version 2022.12.0 \n", + "xhistogram version 0.3.1 \n", + "xgcm version 0.1.dev565+g7492277\n" + ] + } + ], + "source": [ + "print('numpy version', np.__version__, '\\nxarray version', xr.__version__, '\\nxhistogram version', xhistogram.__version__, '\\nxgcm version', xgcm.__version__)" + ] + }, + { + "cell_type": "markdown", + "id": "5fa6652e-9d94-4870-8f33-037ae4d593e5", + "metadata": {}, + "source": [ + "# Tutorial: evaluating water mass transformations from finite-volume tracer budget terms" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "0b71a055-e428-4a25-9edb-45ea25ce9b95", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "xwmt version 0.0.3\n" + ] } ], "source": [ - "grid._ds" + "import xwmt\n", + "print('xwmt version', xwmt.__version__)" + ] + }, + { + "cell_type": "markdown", + "id": "2f73cabd-5dab-44f4-91de-9d8c6a0d6f92", + "metadata": {}, + "source": [ + "### Load grids and data" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "aa99ea92-986b-47d4-a387-1ab5dd951056", + "metadata": {}, + "outputs": [], + "source": [ + "dt = 'daily' # choose from ['monthly', 'daily', 'hourly']\n", + "gridname = 'natv' # choose from ['zstr', 'natv', 'rho2']\n", + "# Caution: non-natv diagnostics are currently broken--see https://github.com/NCAR/MOM6/pull/215\n", + "\n", + "# Load data\n", + "sim = \"wmt_incsurffluxes.natv_rho2_zstr.monthly_daily_hourly.13months\"\n", + "rootdir = f\"/archive/Graeme.Macgilchrist/MOM6-examples/ice_ocean_SIS2/Baltic_OM4_025/{sim}/Baltic_OM4_025/\"\n", + "prefix = '19000101.ocean_'+dt+'_'\n", + "time = \"190*\"\n", + "\n", + "# Merge diagnostics that were saved into different files into one main dataset\n", + "suffixs = ['surf','thck','heat','salt','xtra']\n", + "Zprefixes = {'rho2':'rho2_', 'zstr':'z_', 'natv':'z'}\n", + "Zprefix = Zprefixes[gridname]\n", + "ds = xr.Dataset()\n", + "for suffix in suffixs:\n", + " if suffix == \"surf\":\n", + " filename = prefix+suffix+'_'+time+'.nc'\n", + " else:\n", + " filename = prefix+gridname+'_'+suffix+'_'+time+'.nc'\n", + " dsnow = xr.open_mfdataset(rootdir+filename, chunks={'time':1})\n", + " ds = xr.merge([ds,dsnow])\n", + "\n", + "# Select year-long subset\n", + "ds = ds.sel(time=slice('1900-02-01 00', '1901-02-01 00'))\n", + "\n", + "# Load grid\n", + "oceangridname = '19000101.ocean_static.nc'\n", + "ocean_grid = xr.open_dataset(rootdir+oceangridname).squeeze()\n", + "\n", + "# Some renaming of variables to match conventions\n", + "ocean_grid = ocean_grid.rename({'depth_ocean':'deptho'})\n", + "ds = ds.rename({'temp':'thetao', 'salt':'so'})\n", + "ds['sigma2'] = ds['rhopot2']-1000." + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "d23a62b4-ed21-4116-8f91-ae9729b57937", + "metadata": {}, + "outputs": [], + "source": [ + "# Add core coordinates of ocean_grid to ds\n", + "ds = ds.assign_coords({\n", + " \"areacello\": xr.DataArray(ocean_grid[\"areacello\"].values, dims=('yh', 'xh',)), # Required for area-integration\n", + " \"lon\": xr.DataArray(ocean_grid[\"geolon\"].values, dims=('yh', 'xh',)), # Required for calculating density if not already provided!\n", + " \"lat\": xr.DataArray(ocean_grid[\"geolat\"].values, dims=('yh', 'xh',)), # Required for calculating density if not already provided!\n", + "})\n", + "\n", + "# xgcm grid for dataset\n", + "coords = {\n", + " 'X': {'center': 'xh', 'right': 'xq'},\n", + " 'Y': {'center': 'yh', 'right': 'yq'},\n", + " 'Z': {'center': f'{Zprefix}l', 'outer': f'{Zprefix}i'}\n", + "}\n", + "metrics = {\n", + " ('X','Y'): \"areacello\", # Required for area-integration\n", + "}\n", + "grid = Grid(ds.copy(), coords=coords, metrics=metrics, boundary={'X':'extend', 'Y':'extend', 'Z':'extend'}, autoparse_metadata=False)" + ] + }, + { + "cell_type": "markdown", + "id": "e3d2b510-a73e-482e-a335-751ffc8291c1", + "metadata": {}, + "source": [ + "**Note:** In addition to `grid` containing information about the static model grid, its attribute `grid._ds` points to the full dataset containing model diagnostics!" ] }, { @@ -8202,7 +190,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 35, "id": "3f01ee64-8dba-4039-b8e5-a1100e4cdeac", "metadata": {}, "outputs": [], @@ -8238,7 +226,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 36, "id": "b392f49c-1562-4348-99e4-203867db5daa", "metadata": {}, "outputs": [], @@ -8256,7 +244,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 37, "id": "5772b7f6-4d44-4da5-9ce7-5acd3a3af0a4", "metadata": {}, "outputs": [], @@ -8274,7 +262,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 38, "id": "3b25f4ee-a600-4c14-87f6-eaef0ca1c981", "metadata": {}, "outputs": [ @@ -8284,7 +272,7 @@ "['diffusion']" ] }, - "execution_count": 11, + "execution_count": 38, "metadata": {}, "output_type": "execute_result" } @@ -8303,7 +291,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 39, "id": "ebe8bd8b-1d8c-4ad2-a7a6-fcac8ea7768d", "metadata": {}, "outputs": [ @@ -8311,7 +299,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "[########################################] | 100% Completed | 8.56 ss\n" + "[########################################] | 100% Completed | 53.20 s\n" ] } ], @@ -8322,17 +310,17 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 40, "id": "4d44a979-4e85-4ae7-a6b2-542cd32938df", "metadata": {}, "outputs": [ { "data": { "text/plain": [ - "" + "" ] }, - "execution_count": 13, + "execution_count": 40, "metadata": {}, "output_type": "execute_result" }, @@ -8379,7 +367,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 41, "id": "ccddff5a-369e-4b1b-93e5-8766cb0a6a31", "metadata": {}, "outputs": [ @@ -8536,7 +524,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 42, "id": "d38d0034-75d0-4fe3-8d70-7539bb356dae", "metadata": {}, "outputs": [ @@ -8919,7 +907,7 @@ " * zl (zl) float64 1.0 3.0 5.0 7.0 ... 3.338e+03 3.591e+03 5.111e+03\n", " areacello (yh, xh) float32 2.806e+08 2.806e+08 ... 7.357e+07 7.306e+07\n", " lon (yh, xh) float32 6.125 6.375 6.625 6.875 ... 29.71 29.96 30.21\n", - " lat (yh, xh) float32 53.01 53.01 53.01 53.01 ... 65.09 65.08 65.07" + " dtype='float64', name='zl'))
    • " ], "text/plain": [ "" ] @@ -9464,7 +1452,7 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 48, "id": "e0a5252b-412a-4273-87dc-5f58f843f07b", "metadata": {}, "outputs": [ @@ -9472,7 +1460,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "[########################################] | 100% Completed | 205.87 s\n" + "[########################################] | 100% Completed | 181.03 s\n" ] } ], @@ -9494,7 +1482,7 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 49, "id": "d07777e2-b1f1-4841-8d36-393de2551c50", "metadata": {}, "outputs": [ @@ -9539,7 +1527,7 @@ " 'surface_ocean_flux_advective_negative_rhs': 'salt_rhs_surface_ocean_flux_advective_negative_rhs_product'}}}" ] }, - "execution_count": 22, + "execution_count": 49, "metadata": {}, "output_type": "execute_result" } @@ -9550,7 +1538,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": 50, "id": "95aacfa7-1b14-4ea0-b8b9-1753ff3a4f0e", "metadata": {}, "outputs": [ @@ -9606,7 +1594,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": 51, "id": "8e676faf-1fb7-4799-b31c-b68b3aa1ce86", "metadata": {}, "outputs": [ @@ -10005,7 +1993,7 @@ " salt_rhs_surface_exchange_flux_nonadvective_sum (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n", " salt_rhs_surface_exchange_flux_advective_product (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n", " salt_rhs_surface_exchange_flux_sum (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>\n", - " salt_rhs_surface_ocean_flux_advective_negative_rhs_product (time, yh, xh) float32 dask.array<chunksize=(1, 105, 97), meta=np.ndarray>