{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Analysis of Surface Fields\n",
    "\n",
    "`mom6_tools.MOM6grid` returns an object with MOM6 grid data.\n",
    "\n",
    "`mom6_tools.latlon_analysis` has a collection of tools used to perform spatial analysis  (e.g., time averages and spatial mean).\n",
    "\n",
    "The goal of this notebook is the following:\n",
    "\n",
    "1) server as an example of how to post-process CESM/MOM6 output;\n",
    "\n",
    "2) create time averages of surface fields;\n",
    "\n",
    "3) create time-series of globally-averaged surface fields;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import os\n",
    "import warnings\n",
    "\n",
    "from dask.distributed import Client, LocalCluster\n",
    "import intake\n",
    "from mom6_tools.MOM6grid import MOM6grid\n",
    "from mom6_tools.surface import get_MLD, get_BLD\n",
    "import xarray as xr\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "CESM_output_dir = \"/glade/campaign/cesm/development/cross-wg/diagnostic_framework/CESM_output_for_testing\"\n",
    "serial = False  # use dask LocalCluster\n",
    "Case = \"b.e23_alpha16b.BLT1850.ne30_t232.054\"\n",
    "savefigs = False\n",
    "mom6_tools_config = {}\n",
    "lc_kwargs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTDIR = f\"{CESM_output_dir}/{Case}/ocn/hist/\"\n",
    "print(\"Output directory is:\", OUTDIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# The following parameters must be set accordingly\n",
    "######################################################\n",
    "\n",
    "\n",
    "# create an empty class object\n",
    "class args:\n",
    "    pass\n",
    "\n",
    "\n",
    "args.start_date = mom6_tools_config[\"start_date\"]\n",
    "args.end_date = mom6_tools_config[\"end_date\"]\n",
    "args.casename = Case\n",
    "args.native = f\"{Case}.{mom6_tools_config['Fnames']['native']}\"\n",
    "args.static = f\"{Case}.{mom6_tools_config['Fnames']['static']}\"\n",
    "args.mld_obs = \"mld-deboyer-tx2_3v2\"\n",
    "args.savefigs = savefigs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.isdir(\"PNG/BLD\"):\n",
    "    print(\"Creating a directory to place figures (PNG/BLD)... \\n\")\n",
    "    os.system(\"mkdir -p PNG/BLD\")\n",
    "if not os.path.isdir(\"PNG/MLD\"):\n",
    "    print(\"Creating a directory to place figures (PNG/MLD)... \\n\")\n",
    "    os.system(\"mkdir -p PNG/MLD\")\n",
    "if not os.path.isdir(\"ncfiles\"):\n",
    "    print(\"Creating a directory to place netcdf files (ncfiles)... \\n\")\n",
    "    os.system(\"mkdir ncfiles\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spin up cluster (if running in parallel)\n",
    "client = None\n",
    "if not serial:\n",
    "    cluster = LocalCluster(**lc_kwargs)\n",
    "    client = Client(cluster)\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load mom6 grid\n",
    "grd = MOM6grid(OUTDIR + args.static)\n",
    "grd_xr = MOM6grid(OUTDIR + args.static, xrformat=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Reading native dataset...\")\n",
    "startTime = datetime.now()\n",
    "\n",
    "\n",
    "def preprocess(ds):\n",
    "    \"\"\"Compute montly averages and return the dataset with variables\"\"\"\n",
    "    variables = [\"oml\", \"mlotst\", \"tos\", \"SSH\", \"SSU\", \"SSV\", \"speed\", \"time_bnds\"]\n",
    "    for v in variables:\n",
    "        if v not in ds.variables:\n",
    "            ds[v] = xr.zeros_like(ds.SSH)\n",
    "    return ds[variables]\n",
    "\n",
    "\n",
    "ds1 = xr.open_mfdataset(OUTDIR + args.native, parallel=False)\n",
    "ds = preprocess(ds1)\n",
    "\n",
    "print(\"Time elasped: \", datetime.now() - startTime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Selecting data between {} and {}...\".format(args.start_date, args.end_date))\n",
    "ds_sel = ds.sel(time=slice(args.start_date, args.end_date))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catalog = intake.open_catalog(mom6_tools_config[\"oce_cat\"])\n",
    "mld_obs = catalog[args.mld_obs].to_dask()\n",
    "# uncomment to list all datasets available\n",
    "# list(catalog)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mixed layer depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "# MLD\n",
    "get_MLD(ds, \"mlotst\", mld_obs, grd, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boundary layer depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_BLD(ds, \"oml\", grd, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SSH (not working)\n",
    "# get_SSH(ds, 'SSH', grd, args)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NPL 2022b",
   "language": "python",
   "name": "npl-2022b"
  },
  "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.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}