diff --git a/Bottom_Stratification.ipynb b/Bottom_Stratification.ipynb new file mode 100644 index 0000000..b6018e1 --- /dev/null +++ b/Bottom_Stratification.ipynb @@ -0,0 +1,475 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8bfb94a9-e573-406d-b140-86af797d9523", + "metadata": {}, + "source": [ + "# Prototype for combined model-observation bias and model-intercomparison notebook" + ] + }, + { + "cell_type": "markdown", + "id": "bb7a4fa0-db98-49a6-85c4-94096b405a03", + "metadata": {}, + "source": [ + "This notebook plots the bottom stratification bias for a series of simulations vs. climataology derived from the World Ocean Atlas (WOA) 2018.\n", + "\n", + "It also plots the differences between the simulations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "68078c46-4b7b-4b2f-b8c1-2d2538f46ebc", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import datetime" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50093a3b-2193-4041-8409-de3dc594cf7d", + "metadata": {}, + "outputs": [], + "source": [ + "config = {\n", + " \"startyr\": \"1980\",\n", + " \"endyr\": \"2017\",\n", + " \"dora_id\": \"odiv-384, odiv-391, odiv-395, odiv-393, odiv-396\", # try with just \"odiv-319\" as well.\n", + " \"pathPP\": None,\n", + "}" + ] + }, + { + "cell_type": "markdown", + "id": "b3f810d8-4fac-4b1b-afc3-14077762b34a", + "metadata": {}, + "source": [ + "MAR will pass for environment variables to the script when running via the web engine:\n", + "\n", + "* `MAR_STARTYR`: A `str` of the beginning year of analysis from model\n", + "* `MAR_ENDYR`: A `str` of the ending year of analysis from model\n", + "* `MAR_DORA_ID`: A `str` of the experiment ID in the database for a single experiment, e.g. `\"odiv-1\"`, or comma-separated experiments, e.g. `\"odiv-1, odiv-2\"`\n", + "* `MAR_PATHPP`: A `str` of the top-level path to the post-processing experiment directory of the experiment\n", + "\n", + "The block below will use values passed in by Dora but default to the values defined above in `config`. This is useful for interactive usage and debugging.\n", + "\n", + "If executed from Dora, there will also be a `DORA_EXECUTE` variable that is set." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3ed173d4-cc6d-4693-934a-ff02dd96e632", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in config.items():\n", + " if f\"MAR_{k.upper()}\" in os.environ.keys():\n", + " c = os.environ[f\"MAR_{k.upper()}\"]\n", + " else:\n", + " c = v\n", + " if k==\"dora_id\" and (type(v) is not list): \n", + " c = [e.strip(\" \") for e in c.split(\",\")]\n", + " config[k] = c" + ] + }, + { + "cell_type": "markdown", + "id": "74366b78-6b75-4210-9ed3-24cef6af4550", + "metadata": {}, + "source": [ + "### Import Python Modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dfcdc828-ad45-4ef0-92bc-05df9c036270", + "metadata": {}, + "outputs": [], + "source": [ + "import doralite\n", + "import glob\n", + "import momlevel\n", + "import subprocess \n", + "\n", + "import numpy as np\n", + "import xarray as xr\n", + "import matplotlib.pyplot as plt\n", + "import cartopy.crs as ccrs\n", + "\n", + "from matplotlib.colors import ListedColormap, BoundaryNorm" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "806241bd-ebea-4640-ab69-cea6edba62de", + "metadata": {}, + "outputs": [], + "source": [ + "# momgrid will use a directory of pre-computed weights that is used for calculating basic area-weighted statistics later\n", + "import momgrid\n", + "os.environ[\"MOMGRID_WEIGHTS_DIR\"] = \"/nbhome/John.Krasting/grid_weights\"" + ] + }, + { + "cell_type": "markdown", + "id": "d3ddbddd-eed1-4ac4-9123-6c7df5fcab03", + "metadata": {}, + "source": [ + "### Define Local Parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "04229c22-3781-400f-a390-8c5eaaef5189", + "metadata": {}, + "outputs": [], + "source": [ + "# Define some local variables. These are taken from the doralite object\n", + "# or they can be defined locally\n", + "\n", + "def get_local_params(config):\n", + " experiments = {dora_id: doralite.dora_metadata(dora_id) for dora_id in config[\"dora_id\"]}\n", + " params = {\n", + " dora_id: {\n", + " \"experiment\": experiment,\n", + " \"pathPP\": experiment[\"pathPP\"],\n", + " 'expName': experiment['expName']\n", + " }\n", + " for dora_id, experiment in experiments.items()\n", + " }\n", + " return params\n", + "\n", + "start = int(config[\"startyr\"])\n", + "end = int(config[\"endyr\"])\n", + "params = get_local_params(config)" + ] + }, + { + "cell_type": "markdown", + "id": "a8642873-93e4-4e1f-9486-435d6070fc43", + "metadata": {}, + "source": [ + "### Determine What Files to Load" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61b3b4a4-8b70-4a01-bae7-1c0496347c00", + "metadata": {}, + "outputs": [], + "source": [ + "# Determine what files are needed (leave this up to the developer for flexibility)\n", + "# This is an example of what someone might do:\n", + "\n", + "component = \"ocean_annual\"\n", + "static = f\"{component}/{component}.static.nc\"\n", + "varname = \"N2_b\"\n", + "\n", + "chunk = \"5yr\"\n", + "\n", + "filelist_dict = {odiv:\n", + " glob.glob(f\"{p['pathPP']}{component}/ts/**/{chunk}/{component}.*.{varname}.nc\", recursive=True)\n", + " for odiv, p in params.items() \n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "382835f3-a3d7-419f-a92f-e65808230c59", + "metadata": {}, + "outputs": [], + "source": [ + "def is_in_range(file,start,end):\n", + " start = int(start)\n", + " end = int(end)\n", + " target = set(list(range(start,end+1)))\n", + " fname = os.path.basename(file)\n", + " times = fname.split(\".\")[1]\n", + " times = times.split(\"-\")\n", + " times = [int(x[0:4]) for x in times]\n", + " candidate = set(list(range(times[0],times[1]+1)))\n", + " return len(candidate.intersection(target)) > 0\n", + "\n", + "filelist_dict = {odiv: [x for x in filelist if is_in_range(x,start,end)] for odiv, filelist in filelist_dict.items()}\n", + "\n", + "import itertools\n", + "list_of_filelists = [filelist for filelist in filelist_dict.values() if len(filelist) > 1]\n", + "filelist = sorted(list(itertools.chain.from_iterable(list_of_filelists)))\n", + "\n", + "staticfiles = [f\"{p['pathPP']}/{static}\" for p in params.values()]\n", + "\n", + "_ = [print(x) for x in filelist]" + ] + }, + { + "cell_type": "markdown", + "id": "4bfd8545-cb3b-43ee-b1fb-e71a99fe3262", + "metadata": {}, + "source": [ + "### DMgetting Files" + ] + }, + { + "cell_type": "markdown", + "id": "736dbdaa-f8ed-4ad8-a641-eee616041458", + "metadata": {}, + "source": [ + "Dora cannot issue calls to dmget" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7096dc5b-559b-4221-a189-0836d77afd12", + "metadata": {}, + "outputs": [], + "source": [ + "if not \"DORA_EXECUTE\" in os.environ.keys():\n", + " print(\"Calling dmget on files ...\")\n", + " cmd = [\"dmget\"]+filelist+staticfiles\n", + " _ = subprocess.check_output(cmd)" + ] + }, + { + "cell_type": "markdown", + "id": "b32a4d6d-e5a7-4477-af8f-9e5468e3babc", + "metadata": {}, + "source": [ + "### Load model data and grid" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70b1b444-20e4-4e63-b8a6-4d076d977eb2", + "metadata": {}, + "outputs": [], + "source": [ + "models = {}\n", + "for (odiv, filelist), staticfile in zip(filelist_dict.items(), staticfiles):\n", + " if len(filelist) == 0:\n", + " print(f\"The required diagnostics are not available for {odiv}. Skipping.\")\n", + " continue\n", + " \n", + " gds = momgrid.Gridset(filelist, force_symmetric=True, return_corners=True)\n", + " gds.data = gds.data.sel(time=slice(f\"{str(start).zfill(4)}-01-01\",f\"{str(end).zfill(4)}-12-31\"))\n", + "\n", + " og = xr.open_dataset(staticfile)\n", + " og = og.assign_coords({\"xh\":gds.data.xh, \"yh\":gds.data.yh})\n", + "\n", + " models[odiv] = {\"gds\":gds, \"og\":og}" + ] + }, + { + "cell_type": "markdown", + "id": "20efeacb-e03c-4ddb-8f12-1cb67dee8a4e", + "metadata": {}, + "source": [ + "### Define some helper function for the plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85c98f07-28d4-4b44-9906-804c50e5cd05", + "metadata": {}, + "outputs": [], + "source": [ + "def set_annotaions(ax, odiv):\n", + " _ = ax.set_xticks([])\n", + " _ = ax.set_yticks([])\n", + " _ = ax.text(0.0,1.06, f\"Bottom stratification ({start}-{end})\", weight=\"bold\", fontsize=12, transform=ax.transAxes)\n", + " _ = ax.text(0.0,1.02, f\"{odiv} [{params[odiv]['expName']}]\", style=\"italic\", fontsize=10, transform=ax.transAxes)\n", + " _ = ax.text(1.0,1.05, str(starttime.values), ha=\"right\", fontsize=8, transform=ax.transAxes)\n", + " _ = ax.text(1.0,1.02, str(endtime.values), ha=\"right\", fontsize=8, transform=ax.transAxes)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aadf7139-8861-4a4a-8a2a-1a6c6244cc37", + "metadata": {}, + "outputs": [], + "source": [ + "def add_colorbar(fig, cb):\n", + " cbar_ax = fig.add_axes([0.16, 0.06, 0.7, 0.03])\n", + " fig.colorbar(\n", + " cb, cax=cbar_ax, orientation=\"horizontal\", extend=\"both\", label=r\"[s$^{-2}$]\"\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "b2986841-0e5b-453a-b062-3fc47f652f89", + "metadata": {}, + "source": [ + "# Plotting model values" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4dd47f99-ac67-49b1-9552-4fc81da73a75", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.colors as colors" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c52c87b8-a5b7-4742-a61b-6bb51be82e96", + "metadata": {}, + "outputs": [], + "source": [ + "for (odiv, model_dict) in models.items():\n", + " gds = model_dict[\"gds\"]\n", + " \n", + " # Time-average the model data\n", + " model = gds.data.N2_b\n", + " starttime = model.time[0]\n", + " endtime = model.time[-1]\n", + " model = model.mean(\"time\", keep_attrs=True).load()\n", + " \n", + " # Setup plot\n", + " fig = plt.figure(figsize=(10,6))\n", + " ax = plt.subplot(1,1,1, facecolor=\"lightgray\")\n", + " \n", + " # Definie geolon and geolat for plotting (use corners!)\n", + " x = gds.data.geolon_c\n", + " y = gds.data.geolat_c\n", + " \n", + " # Run pcolormesh\n", + " cb = plt.pcolormesh(x,y,model, cmap=\"cividis\", norm=colors.LogNorm(vmin=1e-8, vmax=1e-4))\n", + " \n", + " # Clean up figure and add labels\n", + " set_annotaions(ax, odiv)\n", + " \n", + " # Add colorbar\n", + " add_colorbar(fig, cb)" + ] + }, + { + "cell_type": "markdown", + "id": "1f47518f-00ea-4d16-8e9d-554edc14cb2f", + "metadata": {}, + "source": [ + "# Plotting inter-model ratios" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "294f6a30-47b8-4938-8fa4-f6390e841234", + "metadata": {}, + "outputs": [], + "source": [ + "def set_annotaions_ratio(ax, odiv, odiv_ref):\n", + " _ = ax.set_xticks([])\n", + " _ = ax.set_yticks([])\n", + " _ = ax.text(0.0,1.06, f\"percent change in stratification ({start}-{end})\", weight=\"bold\", fontsize=12, transform=ax.transAxes)\n", + " _ = ax.text(0.0,1.02, rf\"relative change = 100% (({odiv} [{params[odiv]['expName']}] - {odiv_ref} [{params[odiv_ref]['expName']}]) $/$ {odiv_ref} [{params[odiv_ref]['expName']}]))\", style=\"italic\", fontsize=10, transform=ax.transAxes)\n", + " _ = ax.text(1.0,1.05, str(starttime.values), ha=\"right\", fontsize=8, transform=ax.transAxes)\n", + " _ = ax.text(1.0,1.02, str(endtime.values), ha=\"right\", fontsize=8, transform=ax.transAxes)\n", + "\n", + "def add_colorbar_relative_change(fig, cb):\n", + " cbar_ax = fig.add_axes([0.16, 0.06, 0.7, 0.03])\n", + " fig.colorbar(\n", + " cb, cax=cbar_ax, orientation=\"horizontal\", extend=\"both\", label=r\"% change, relative to reference simulation\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "93f368ae-718c-44c2-a0c4-a269771f9b37", + "metadata": {}, + "outputs": [], + "source": [ + "first = True\n", + "\n", + "if len(models)==1: \n", + " print(\"No difference plot because only one experiment ID was provided!\")\n", + "\n", + "else:\n", + " for (odiv, model_dict) in models.items():\n", + " if first:\n", + " first = False\n", + " gds_ref = model_dict[\"gds\"].regrid(resolution=0.25)\n", + " odiv_ref = odiv\n", + " continue\n", + " \n", + " gds = model_dict[\"gds\"].regrid(resolution=0.25)\n", + " \n", + " # Time-average the model data\n", + " model = gds.N2_b\n", + " starttime = model.time[0]\n", + " endtime = model.time[-1]\n", + " model = model.mean(\"time\", keep_attrs=True).load()\n", + " \n", + " model_ref = gds_ref.N2_b\n", + " starttime = model_ref.time[0]\n", + " endtime = model_ref.time[-1]\n", + " model_ref = model_ref.mean(\"time\", keep_attrs=True).load()\n", + " \n", + " # Setup plot\n", + " fig = plt.figure(figsize=(12,6))\n", + " ax = plt.subplot(1,1,1, facecolor=\"lightgray\")\n", + " \n", + " # Definie geolon and geolat for plotting (use corners!)\n", + " x = gds.lon_b\n", + " y = gds.lat_b\n", + " \n", + " # Run pcolormesh\n", + " cb = plt.pcolormesh(x,y,100*(model-model_ref)/model_ref, cmap=\"RdBu_r\", vmin=-100, vmax=100)\n", + " \n", + " # Clean up figure and add labels\n", + " set_annotaions_ratio(ax, odiv, odiv_ref)\n", + " \n", + " # Add colorbar\n", + " add_colorbar_relative_change(fig, cb)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b68a042-87f9-4241-a1a1-0d8ea6816882", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}