{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f30aa9dc",
   "metadata": {},
   "source": [
    "# ocean: zonal\n",
    "\n",
    "zonal ocean data\n",
    "\n",
    "```{dropdown}  **Branding suffix of datasets**\n",
    "Datasets with the folloing branding suffix are shown in this page: \\\n",
    "[time]-[vertical]-[horizontal]-[domain]\n",
    "```\n",
    "* tavg-ol-hyb-sea\n",
    "    - tavg: time average\n",
    "    - ol: ocean levels\n",
    "    - hyb: zonal mean\n",
    "    - sea: ocean domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f107ce",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "## Import libraries\n",
    "import glob\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import cmocean\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from IPython.display import HTML, display\n",
    "from matplotlib.colors import BoundaryNorm, ListedColormap\n",
    "\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "from utils import read_compound_names, read_variables\n",
    "\n",
    "logging.basicConfig(\n",
    "    filename=\"execute.log\",\n",
    "    level=logging.INFO,\n",
    "    format=\"%(asctime)s %(message)s\",\n",
    "    force=True,\n",
    ")\n",
    "log = logging.getLogger(\"ocean_zonal\")\n",
    "log.info(\"Starting ocean_zonal.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a6ede6",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# load parameters\n",
    "import yaml\n",
    "\n",
    "with open('params.yml', 'r') as f:\n",
    "    param = yaml.safe_load(f)\n",
    "    cmorout = param['cmorout']\n",
    "    source_id = param['source_id']\n",
    "    experiment_id = param['experiment_id']\n",
    "    variant_label = param['variant_label']\n",
    "    version = param['version']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8314513c",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# grid label\n",
    "grid_label = 'g202'         # grid label, e.g., 'gn', 'gr', 'g999'\n",
    "\n",
    "# data path\n",
    "data_path = os.path.join(cmorout, source_id, experiment_id, version)\n",
    "\n",
    "# load methods for plotting and set defaults\n",
    "methods =read_variables('assets/methods.txt')\n",
    "#print(methods.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f06c51b7",
   "metadata": {},
   "source": [
    "---\n",
    "```{dropdown}  **List of datasets:** \n",
    "List of compound names of datasets that needs to be cmorized. \\\n",
    "Those datasets which are not cmorized have no link.\n",
    "``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a5a359",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# load compound names\n",
    "\n",
    "coords = ('tavg-ol-hyb-sea','tavg-rho-hyb-sea','tavg-u-hyb-sea')\n",
    "cnames = read_compound_names('assets/variables.nml')\n",
    "# examples of compound names:\n",
    "# cnames = ['ocean.tos.tavg-u-hxy-sea.mon.glb', 'ocean.ficeberg.tavg-u-hxy-sea.mon.glb']\n",
    "for cname in cnames:\n",
    "    realm = cname.split('.')[0]\n",
    "    var = cname.split('.')[1]\n",
    "    coord = cname.split('.')[2]\n",
    "    freq = cname.split('.')[3]\n",
    "    region = cname.split('.')[4]\n",
    "\n",
    "    if realm != 'ocean' or coord not in coords:\n",
    "        continue\n",
    "    else:\n",
    "        file_pattern = f\"{var}_{coord}_{freq}_{region}_{grid_label}_{source_id}_{experiment_id}_{variant_label}_*.nc\"\n",
    "        if not glob.glob(os.path.join(data_path, file_pattern)):\n",
    "            print(cname)\n",
    "            continue\n",
    "        else:\n",
    "            display(HTML(f'<a href=\"#{cname}\">{cname}</a>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f2b50e9",
   "metadata": {},
   "source": [
    "---\n",
    "```{dropdown}  **Dataset validated:**\n",
    "Each validated dataset has some text information and some plots (2D and/or 1D)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25072041",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "from xarray.coders import CFDatetimeCoder\n",
    "\n",
    " # loop through compound names and plot\n",
    "for cname in cnames:\n",
    "    mth_cmap = 'mpl.colormaps[\"seismic\"]'\n",
    "    if cname not in methods:\n",
    "        print(f\"{cname} not found in methods.txt, using default methods for plotting.\")\n",
    "    else:\n",
    "        if methods[cname] is not None and 'cmap' in methods[cname]:\n",
    "            mth_cmap = methods[cname]['cmap']\n",
    "\n",
    "    realm = cname.split('.')[0]\n",
    "    var = cname.split('.')[1]\n",
    "    coord = cname.split('.')[2]\n",
    "    freq = cname.split('.')[3]\n",
    "    region = cname.split('.')[4]\n",
    "\n",
    "    if realm != 'ocean' or coord not in coords:\n",
    "        continue\n",
    "\n",
    "    file_pattern = f\"{var}_{coord}_{freq}_{region}_{grid_label}_{source_id}_{experiment_id}_{variant_label}_*.nc\"\n",
    "    data_files = sorted(glob.glob(os.path.join(data_path, file_pattern)))\n",
    "    if data_files:\n",
    "        log.info(f\"{cname}\")\n",
    "        with xr.open_mfdataset(data_files, decode_times=CFDatetimeCoder(use_cftime=True)) as ds:\n",
    "            if var in ds:\n",
    "                data = ds[var]\n",
    "                sector = ds['sector']\n",
    "            else:\n",
    "                continue\n",
    "    else:\n",
    "        continue\n",
    "\n",
    "\n",
    "    display(HTML(f'<div id=\"{cname}\"></div>'))\n",
    "    print(f'\\033[1m{cname}\\033[0m')\n",
    "    print(f'long name: {data.long_name} ({data.units})')\n",
    "    print(f'original_name: {data.attrs[\"original_name\"]} -> {var}')\n",
    "    if 'history' in data.attrs:\n",
    "        print(f'history: {data.attrs[\"history\"]}')\n",
    "    if 'comment' in data.attrs:\n",
    "        print(f'comment: {data.attrs[\"comment\"]}')\n",
    "\n",
    "    nlines = data.basin.size\n",
    "    fig, axs = plt.subplots((nlines+1)//2, 2, figsize=(11.69, 8.27), dpi=96, sharex=True, sharey=False)\n",
    "    ax = axs.flatten()\n",
    "    for n in range(nlines):\n",
    "        if any(dim in data.dims for dim in ['lev', 'rho']):\n",
    "            data2d = data.isel(basin=n).mean(dim='time')\n",
    "            vmin, vmax = data2d.quantile([0.001, 0.999], dim=None).values\n",
    "            vmax = max(abs(vmin),abs(vmax))\n",
    "            vmin = -vmax\n",
    "            # Define the discrete levels (boundaries) and the corresponding colors\n",
    "            n_levels = 20\n",
    "            levels = np.linspace(vmin, vmax, n_levels+1)\n",
    "            colors = eval(mth_cmap)(np.linspace(0,1,n_levels))\n",
    "            cmap = ListedColormap(colors)\n",
    "\n",
    "            # Create a BoundaryNorm to map data to discrete color indices\n",
    "            norm = BoundaryNorm(levels, ncolors=len(cmap.colors), clip=False)\n",
    "\n",
    "            if vmin == vmax:\n",
    "                data2d.plot(ax=ax[n], cmap=cmap, bar_kwargs={\"extend\":\"both\", \"label\": f\"{data.name} ({data.units})\"})\n",
    "            else:\n",
    "                data2d.plot(ax=ax[n], cmap=cmap, norm=norm, cbar_kwargs={\"extend\":\"both\", \"label\": f\"{data.name} ({data.units})\"})\n",
    "            ax[n].invert_yaxis()\n",
    "        else:\n",
    "            data.isel(basin=n).mean(dim='time').plot(ax=ax[n])\n",
    "            ax[n].set_ylabel(data.units)\n",
    "        \n",
    "        ax[n].set_title(sector[n].compute().astype(str).item(0).strip())\n",
    "        ax[n].set_xlabel('')\n",
    "        ax[n].grid()\n",
    "\n",
    "    # Turn off x‑axis ticks for all subplots\n",
    "    for ax in axs.flatten():\n",
    "        ax.xaxis.set_tick_params(labelbottom=False, bottom=False)\n",
    "\n",
    "    # Turn on x‑axis ticks only for the bottom row (last two axes)\n",
    "    for ax in axs[-1,:]:\n",
    "        ax.xaxis.set_tick_params(labelbottom=True, bottom=True)\n",
    "        ax.set_xlabel(data.lat.long_name)\n",
    "\n",
    "    fig.suptitle(f\"{data.long_name}\", fontsize=16)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    del data\n",
    "    del ax, axs, fig "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cmip7validate-env",
   "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}