{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f30aa9dc",
   "metadata": {
    "papermill": {},
    "tags": []
   },
   "source": [
    "# ocean: fx\n",
    "\n",
    "fx ocean grid 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",
    "* ti-u-hxy-sea\n",
    "* ti-ol-hxy-sea\n",
    "    - ti: time independent\n",
    "    - u: ocean surface\n",
    "    - ol: ocean levels\n",
    "    - hxy: horizontal curvlinear grids\n",
    "    - sea: ocean domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f107ce",
   "metadata": {
    "papermill": {},
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "## Import libraries\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.colors import ListedColormap, BoundaryNorm\n",
    "import cartopy.feature as cfeature\n",
    "import cartopy.crs as ccrs\n",
    "import cmocean\n",
    "import sys\n",
    "import os\n",
    "import glob\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "sys.path.append(os.getcwd())\n",
    "from utils import load_grid_vertex, read_variables, read_compound_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1227985d",
   "metadata": {
    "papermill": {},
    "tags": [
     "parameters",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# parameters for the cmorized data\n",
    "cmorout=''                  # root for cmorized data, e.g., '/scratch/$USER/cmorout'\n",
    "source_id      = ''         # model name, e.g., 'NorESM3-LM'\n",
    "experiment_id  = ''         # experiment name, e.g., 'historical', 'ssp585', 'piControl'\n",
    "variant_label  = ''         # variant label, e.g., 'r1i1p1f1'\n",
    "grid_label     = ''         # grid label, e.g., 'gn', 'gr', 'g999'\n",
    "version        = ''         # version, e.g., 'v20260601'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ac8e69",
   "metadata": {
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "cmorout = \"/scratch/yanchun/cmorout\"\n",
    "source_id = \"UKESM1-3-LL\"\n",
    "experiment_id = \"piControl\"\n",
    "variant_label = \"r1i1p1f1\"\n",
    "grid_label = \"g143\"\n",
    "version = \"v20260821\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8314513c",
   "metadata": {
    "papermill": {},
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# data path\n",
    "data_path = os.path.join(cmorout, source_id, experiment_id, version)\n",
    "\n",
    "# load grid\n",
    "grid_file = 'data/grid.nc'\n",
    "lat, lon, clat, clon = load_grid_vertex(grid_file)\n",
    "with xr.open_dataset(grid_file) as ds:\n",
    "    pmask = ds['pmask']\n",
    "    parea = ds['parea']\n",
    "\n",
    "parea = parea.rename({'y': 'j', 'x': 'i'})\n",
    "pmask = pmask.rename({'y': 'j', 'x': 'i'})\n",
    "# load methods for plotting and set defaults\n",
    "methods =read_variables('data/methods.txt')\n",
    "#print(methods.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f06c51b7",
   "metadata": {
    "papermill": {},
    "tags": []
   },
   "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",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a5a359",
   "metadata": {
    "papermill": {},
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# load compound names\n",
    "\n",
    "coords = ('ti-u-hxy-sea', 'ti-u-hxy-u', 'ti-ol-hxy-sea')\n",
    "cnames = read_compound_names('data/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",
    "        data_file = f\"{var}_{coord}_{freq}_{region}_{grid_label}_{source_id}_{experiment_id}_{variant_label}.nc\"\n",
    "        if not glob.glob(os.path.join(data_path, data_file)):\n",
    "            print(cname)\n",
    "            continue\n",
    "        else:\n",
    "            display(HTML(f'<a href=\"#{cname}\">{cname}</a>'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2cb0804",
   "metadata": {
    "papermill": {},
    "tags": []
   },
   "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": {
    "papermill": {},
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "# 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.keys():\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].keys():\n",
    "                mth_vert = methods[cname]['vertical']\n",
    "\n",
    "            if 'timeseries' in methods[cname].keys():\n",
    "                mth_ts = methods[cname]['timeseries']\n",
    "\n",
    "            if 'cmap' in methods[cname].keys():\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",
    "    data_file = f\"{var}_{coord}_{freq}_{region}_{grid_label}_{source_id}_{experiment_id}_{variant_label}.nc\"\n",
    "\n",
    "    if not glob.glob(os.path.join(data_path, data_file)):\n",
    "        continue\n",
    "\n",
    "    #with xr.open_mfdataset(os.path.join(data_path, data_file)) as ds:\n",
    "    data_file = glob.glob(os.path.join(data_path, data_file))[0]\n",
    "    with xr.open_dataset(os.path.join(data_path, data_file)) as ds:\n",
    "        if var in ds:\n",
    "            data = ds[var]\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",
    "    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'timeseries aggregation method for 3D->1D plot: {mth_ts}')\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",
    "    if coord == 'ti-ol-hxy-sea':\n",
    "        if mth_vert == 'mean':\n",
    "            data2d = data.mean(dim='lev',keep_attrs=True).where(pmask == 1)\n",
    "        elif mth_vert == 'sum':\n",
    "            data2d = data.sum(dim='lev',keep_attrs=True).where(pmask == 1)\n",
    "        else:\n",
    "            raise ValueError(f'Unsupported vertical aggregation method: {mth_vert}')\n",
    "        print(f'vertical aggregation method for 3D->2D plot: {mth_vert}')\n",
    "    else:\n",
    "        data2d = data \n",
    "\n",
    "\n",
    "    #ax1 = fig.add_subplot(1, 2, 1, projection=proj)\n",
    "    #fig, (ax1, ax2) = plt.subplots( nrows=1, ncols=2, figsize=(12, 5), gridspec_kw={'width_ratios': [2, 1]}, subplot_kw={'projection': proj})\n",
    "\n",
    "    proj = ccrs.PlateCarree(central_longitude=90.0)\n",
    "    fig, ax = plt.subplots(1, figsize=(11.7, 4), dpi=96, subplot_kw={\"projection\": proj})\n",
    "\n",
    "    # plot 2D map\n",
    "    #ax1 = fig.add_subplot(gs[0], projection=proj)\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",
    "    # Plot the mesh\n",
    "    if vmin == vmax:\n",
    "        pm = ax.pcolormesh(lon, lat, data2d, cmap=cmap,\n",
    "                            transform=ccrs.PlateCarree(), shading='auto', rasterized=True)\n",
    "    else:\n",
    "        pm = ax.pcolormesh(lon, lat, data2d, cmap=cmap, norm=norm,\n",
    "                            transform=ccrs.PlateCarree(), shading='auto', rasterized=True)\n",
    "\n",
    "    # Add map features\n",
    "    #ax.stock_img()\n",
    "    ax.add_feature(cfeature.LAND, facecolor='lightgray')\n",
    "\n",
    "    gl = ax.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=ax, fraction=0.4, shrink=0.8, label='[yr]')\n",
    "    cb.set_label(label=data.units, size=14)\n",
    "    cb.ax.tick_params(labelsize=12)\n",
    "    plt.tight_layout()\n",
    "        \n",
    "    #fig, ax = plot_map2d(lon, lat, data2d.mean(dim='time'), proj='PlateCarree', norm=None, cmap=eval(mth_cmap))\n",
    "    ax.coastlines(resolution='110m')\n",
    "    ax.add_feature(cfeature.LAND, facecolor='lightgray')\n",
    "    plt.title(data2d.long_name)\n",
    "\n",
    "\n",
    "    plt.title(data2d.long_name)\n",
    "    plt.show()\n",
    "\n",
    "    del data, data2d\n",
    "    del ax, pm, cb, gl, fig"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cmip7validate-env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.3"
  },
  "papermill": {
   "default_parameters": {},
   "environment_variables": {},
   "input_path": "notebooks/ocean_fx.ipynb",
   "output_path": "ocean_fx.ipynb",
   "parameters": {
    "cmorout": "/scratch/yanchun/cmorout",
    "experiment_id": "piControl",
    "grid_label": "g143",
    "source_id": "UKESM1-3-LL",
    "variant_label": "r1i1p1f1",
    "version": "v20260821"
   },
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}