{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f30aa9dc",
   "metadata": {},
   "source": [
    "# ocean: 3D\n",
    "\n",
    "3D 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-hxy-sea\n",
    "    - tavg: time average\n",
    "    - ol: ocean levels\n",
    "    - hxy: horizontal curvlinear grids\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",
    "\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",
    "from pathlib import Path\n",
    "from pyprojroot import here\n",
    "from xarray.coders import CFDatetimeCoder\n",
    "\n",
    "from cmip7validate.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_3D\")\n",
    "log.info(\"Starting ocean_3D.ipynb\")\n",
    "\n",
    "# Get project root\n",
    "root_path = Path(here())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a479373d",
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# load parameters\n",
    "import yaml\n",
    "\n",
    "with open(root_path.joinpath('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'\n",
    "\n",
    "# data path\n",
    "data_path = Path(cmorout).joinpath(source_id, experiment_id, version)\n",
    "\n",
    "# load grid\n",
    "grid_file = data_path.joinpath(f\"areacello_ti-u-hxy-u_fx_glb_{grid_label}_{source_id}_{experiment_id}_{variant_label}.nc\")\n",
    "if grid_file.exists():\n",
    "    with xr.open_dataset(grid_file) as ds:\n",
    "        parea = ds['areacello']\n",
    "else:\n",
    "    grid_file = root_path.joinpath('assets/grid.nc')\n",
    "    with xr.open_dataset(grid_file) as ds:\n",
    "        parea = ds['parea']\n",
    "        parea = parea.rename({'y': 'j', 'x': 'i'})\n",
    "\n",
    "grid_file = data_path.joinpath(f\"sftof_ti-u-hxy-u_fx_glb_{grid_label}_{source_id}_{experiment_id}_{variant_label}.nc\")\n",
    "if grid_file.exists():\n",
    "    with xr.open_dataset(grid_file) as ds:\n",
    "        pmask = ds['sftof']\n",
    "else:\n",
    "    grid_file = root_path.joinpath('assets/grid.nc')\n",
    "    with xr.open_dataset(grid_file) as ds:\n",
    "        pmask = ds['pmask']\n",
    "        pmask = pmask.rename({'y': 'j', 'x': 'i'})\n",
    "\n",
    "# load methods for plotting and set defaults\n",
    "methods =read_variables(root_path.joinpath('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",
    "grid_label = 'g???'\n",
    "coords = ('tavg-ol-hxy-sea')\n",
    "cnames = read_compound_names(root_path.joinpath('assets/variables.nml'))\n",
    "# examples of compound names:\n",
    "# cnames = ['ocean.thetao.tavg-ol-hxy-sea.mon.glb', 'ocean.masscello.tavg-ol-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 sorted(data_path.glob(file_pattern)):\n",
    "            print(cname)\n",
    "            continue\n",
    "        else:\n",
    "            display(HTML(f'<a href=\"#{cname}\">{cname}</a>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c549e24",
   "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": [
    "import time\n",
    "\n",
    "# loop through compound names and plot\n",
    "for cname in cnames:\n",
    "    mth_vert = 'mean'\n",
    "    mth_ts = 'mean'\n",
    "    mth_cmap = 'mpl.colormaps[\"viridis\"]'\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:\n",
    "            if 'vertical' in methods[cname]:\n",
    "                mth_vert = methods[cname]['vertical']\n",
    "\n",
    "            if 'timeseries' in methods[cname]:\n",
    "                mth_ts = methods[cname]['timeseries']\n",
    "\n",
    "            if '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(data_path.glob(file_pattern))\n",
    "    if data_files:\n",
    "        log.info(f\"{cname}\")\n",
    "        #start = time.time()\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",
    "                lat = ds['latitude']\n",
    "                lon = ds['longitude']\n",
    "                lev_bnds = ds.isel(time=0,drop=True)['lev_bnds']\n",
    "                thk = lev_bnds.diff(dim='bnds').squeeze()\n",
    "                thk = thk/thk.sum(dim='lev')\n",
    "            else:\n",
    "                continue\n",
    "    else:\n",
    "        continue\n",
    "\n",
    "    # remove the las row of grid info if the data has not the last row\n",
    "    if data.sizes['j'] == 384:\n",
    "        parea = parea.isel(j=slice(0, 384))\n",
    "        pmask = pmask.isel(j=slice(0, 384))\n",
    "\n",
    "    if 'lev' not in data.dims:\n",
    "        print(f\"{cname} does not have vertical dimension 'lev' , skipping vertical aggregation.\")\n",
    "        continue\n",
    "\n",
    "    # time mean\n",
    "    datatmean = data.mean(dim='time', keep_attrs=True).compute()\n",
    "    # vertical mean\n",
    "    if mth_vert == 'mean':\n",
    "        data2d = datatmean.weighted(thk).mean(dim='lev',keep_attrs=True).where(pmask == 1)\n",
    "    elif mth_vert == 'sum':\n",
    "        data2d = datatmean.weighted(thk).sum(dim='lev',keep_attrs=True).where(pmask == 1)\n",
    "    else:\n",
    "        raise ValueError(f'Unsupported vertical aggregation method: {mth_vert}')\n",
    "\n",
    "    #print(f\"l75:{time.time()-start}\")\n",
    "    # remove the las row of grid info if the data has not the last row\n",
    "    if data2d.sizes['j'] == 384:\n",
    "        parea = parea.isel(j=slice(0, 384))\n",
    "        pmask = pmask.isel(j=slice(0, 384))\n",
    "\n",
    "    display(HTML(f'<div id=\"{cname}\"></div>'))\n",
    "    print(f'\\033[1m{cname}\\033[0m')\n",
    "    print(f'long name: {data2d.long_name} ({data2d.units})')\n",
    "    print(f'vertical aggregation method for 3D->2D plot: {mth_vert}')\n",
    "    print(f'timeseries aggregation method for 3D->1D plot: {mth_ts}')\n",
    "    print(f'original_name: {data2d.attrs[\"original_name\"]} -> {var}')\n",
    "    if 'history' in data2d.attrs:\n",
    "        print(f'history: {data2d.attrs[\"history\"]}')\n",
    "    if 'comment' in data2d.attrs:\n",
    "        print(f'comment: {data2d.attrs[\"comment\"]}')\n",
    "\n",
    "    # #Create figure\n",
    "    fig = plt.figure(figsize=(16, 3), dpi=96)\n",
    "    gs = fig.add_gridspec(nrows=1, ncols=2, width_ratios=[2, 1]) \n",
    "    proj = ccrs.PlateCarree(central_longitude=90.0)\n",
    "    # rescale\n",
    "    vmin, vmax = data2d.quantile([0.001, 0.999], dim=None).values\n",
    "    if mth_cmap == 'cmocean.cm.balance':\n",
    "        vmax = max(abs(vmin),abs(vmax))\n",
    "        vmin = -vmax\n",
    "    # Define the discrete levels (boundaries) and the corresponding colors\n",
    "    levels = np.linspace(vmin, vmax, 21)\n",
    "    n_levels = 20\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",
    "    ax1 = fig.add_subplot(gs[0], projection=proj)\n",
    "    if vmin == vmax:\n",
    "        pm = ax1.pcolormesh(lon, lat, data2d, cmap=cmap,\n",
    "                            transform=ccrs.PlateCarree(), shading='auto', rasterized=True)\n",
    "    else:\n",
    "        pm = ax1.pcolormesh(lon, lat, data2d, cmap=cmap, norm=norm,\n",
    "                            transform=ccrs.PlateCarree(), shading='auto', rasterized=True)\n",
    "\n",
    "    #print(f\"l114:{time.time()-start}\")\n",
    "    gl = ax1.gridlines(ylocs=range(-90, 90, 30), draw_labels=True)\n",
    "    gl.ylocator = mpl.ticker.FixedLocator(range(-90,90,30))\n",
    "\n",
    "    # Add colorbar\n",
    "    cb = plt.colorbar(pm, ax=ax1, fraction=0.4, shrink=0.8, label='[yr]', extend='both')\n",
    "    cb.set_label(label=data.units, size=14)\n",
    "    cb.ax.tick_params(labelsize=12)\n",
    "    plt.tight_layout()\n",
    "        \n",
    "    # Add map features\n",
    "    ax1.coastlines(resolution='110m')\n",
    "    ax1.add_feature(cfeature.LAND, facecolor='lightgray')\n",
    "    plt.suptitle(data2d.long_name, y=1.1)\n",
    "\n",
    "    # plot 1D timeseries\n",
    "    if mth_ts == 'mean':\n",
    "        data1d = data.weighted(parea*thk).mean(dim=['lev','j', 'i'], keep_attrs=True).compute()\n",
    "    elif mth_ts == 'sum':\n",
    "        data1d = data.weighted(parea*thk).sum(dim=['lev','j', 'i'], keep_attrs=True).compute()\n",
    "    else:\n",
    "        raise ValueError(f'Unsupported timeseries method: {mth_ts}')\n",
    "   \n",
    "    data1d.attrs = data2d.attrs.copy()\n",
    "\n",
    "    ax2 = fig.add_subplot(gs[1])\n",
    "    if data1d.sizes[\"time\"] < 2:\n",
    "        ax2.text(0.1, 0.8, (f\"global {mth_ts} value: {data1d.values[0]} {data1d.units}\"))\n",
    "        ax2.text(0.1, 0.6, (\"less than 2 time steps.\"))\n",
    "        ax2.text(0.1, 0.4, (\"skipping timeseries plot.\"))\n",
    "    else:\n",
    "        data1d.plot(ax=ax2)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "    del data, data2d, data1d\n",
    "    del fig, gs, ax1, ax2, pm, cb, gl"
   ]
  }
 ],
 "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
}