{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1ca2a6f5-bd5a-4520-8380-baf5cb536d41",
   "metadata": {},
   "source": [
    "## ADF Diagnostics In Jupyter\n",
    "This notebook will run the Atmospheric Diagnostic Framework using the settings in a config.yaml file in your ADF directory. \n",
    "\n",
    "Note that it was developed to run on Cheyenne/Caspar *with the NPL (conda) kernel*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eea8d1a-d6a7-4464-bd6d-8194aa2368d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Setup\n",
    "#### Required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d0264e2-c116-4e18-af33-0ce3adabcd60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path\n",
    "from pathlib import Path\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6566766f-9d8f-4428-bf01-3c36a1ffeb63",
   "metadata": {},
   "source": [
    "#### Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e627890-a0a9-42d9-bb23-302abfedd944",
   "metadata": {},
   "outputs": [],
   "source": [
    "### default parameters\n",
    "# adf_path = \"../../externals/ADF\"\n",
    "# config_path = \".\"\n",
    "# config_fil_str = \"config_f.cam6_3_119.FLTHIST_ne30.r328_gamma0.33_soae_numcin3.001_vs_f.cam6_3_119.FLTHIST_ne30.r328_gamma0.33_soae.001.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec97fae-f7f3-40cc-adf0-3f6726753215",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# cell to insert parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88e305c-5473-4de6-9aa5-bf674ecd899f",
   "metadata": {},
   "source": [
    "#### All path processing:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66a55cfa-2d81-4298-8623-4ceaf3c0fb1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine ADF directory path\n",
    "# If it is in your cwd, set adf_path = local_path,\n",
    "# otherwise set adf_path appropriately\n",
    "\n",
    "local_path = os.path.abspath(\"\")\n",
    "\n",
    "# Set up the ADF for your location/user name\n",
    "# user = \"richling\"\n",
    "# adf_path = f\"/glade/work/{user}/ADF/\"\n",
    "\n",
    "print(f\"current working directory = {local_path}\")\n",
    "print(f\"ADF path                  = {adf_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1689b448-daf6-44ca-9c9e-b1345490b552",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set path to ADF lib\n",
    "lib_path = os.path.join(adf_path, \"lib\")\n",
    "print(f\"The lib scripts live here, right? {lib_path}\")\n",
    "\n",
    "# Add paths to python path:\n",
    "sys.path.append(lib_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f927cc3e-7fa9-4347-b44e-956c7ebdc742",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set path to ADF plotting scripts directory\n",
    "plotting_scripts_path = os.path.join(adf_path, \"scripts\", \"plotting\")\n",
    "print(f\"The plotting scripts live here, right? {plotting_scripts_path}\")\n",
    "\n",
    "# Add paths to python path:\n",
    "sys.path.append(plotting_scripts_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159be7f2-8593-483a-b200-daf8c3c992c8",
   "metadata": {},
   "source": [
    "#### Import config file into ADF object\n",
    "If there are errors, here, it is likely due to path errors above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd4918d2-dcc0-417b-8867-d368c4cb1d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = os.path.join(config_path, config_fil_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f553607f-cd31-4aa3-8142-4959b3569a1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import ADF diagnostics object\n",
    "from adf_diag import AdfDiag\n",
    "\n",
    "# If this fails, check your paths output in the cells above,\n",
    "# and that you are running the NPL (conda) Kernel\n",
    "# You can see all the paths being examined by un-commenting the following:\n",
    "# sys.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ac46243-f980-41e1-9f8c-8d8ad361bb89",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize ADF object\n",
    "adf = AdfDiag(config_file)\n",
    "adf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7d266a8-f006-468d-96d5-51d9a93dd6dd",
   "metadata": {
    "tags": []
   },
   "source": [
    "### ADF Standard Work Flow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b1d075b-232d-4466-a6be-192c457ae5a9",
   "metadata": {},
   "source": [
    "#### Calculate the Time Series files\n",
    "\n",
    "* NOTE: If <strong>not</strong> comparing against observations, you must run <strong>_create_time_series()_</strong> again with baseline flag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f056e6d-5828-434d-b70b-f60627f30ba1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create model time series.\n",
    "adf.create_time_series()\n",
    "\n",
    "# Create model baseline time series (if needed):\n",
    "if not adf.compare_obs:\n",
    "    adf.create_time_series(baseline=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec72ce9-8167-435c-8334-728e2df24a01",
   "metadata": {},
   "source": [
    "#### Calculate the Climo files\n",
    "\n",
    "* NOTE: <strong>Do not</strong> need to specify or repeat for baseline case unlike time series generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a611c7-a90b-4d24-a518-046407564e35",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create model climatology (climo) files.\n",
    "adf.create_climo()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dc28ffb-f501-4a42-8e21-cd8b993163ea",
   "metadata": {},
   "source": [
    "#### Regrid the Climo files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9673eada-6eeb-4c68-b59a-bd7baf62aca1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Regrid model climatology files to match either\n",
    "# observations or CAM baseline climatologies.\n",
    "# This call uses the \"regridding_scripts\" specified in the config file:\n",
    "adf.regrid_climo()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56434dfe-187d-43a8-967c-5cf027e1b05b",
   "metadata": {},
   "source": [
    "#### Run statistics on Time Series files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40c90be-b041-46b6-b268-3f3f02ea855a",
   "metadata": {
    "tags": []
   },
   "source": [
    "#Perform analyses on the simulation(s).\n",
    "#This call uses the \"analysis_scripts\" specified in the config file:\n",
    "adf.perform_analyses()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e12c5b56-2a54-49c9-816c-9d447d511a9d",
   "metadata": {},
   "source": [
    "#### Create Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31c8782b-1492-4260-bfea-48a69169e30a",
   "metadata": {
    "tags": []
   },
   "source": [
    "#Create plots.\n",
    "#This call uses the \"plotting_scripts\" specified in the config file:\n",
    "adf.create_plots()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "898b1090-aad9-45b2-a0f2-2882ae8f9834",
   "metadata": {},
   "source": [
    "#### Generate HTML files\n",
    "\n",
    "* This will create html files that you can view via webbrower either:\n",
    "    - in Casper/Cheyenne\n",
    "    - pushing it to CGD projects webpage through Tungsten"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f391772b-8f17-40bc-b87d-d9b27b0a8bba",
   "metadata": {
    "tags": []
   },
   "source": [
    "#Create website.\n",
    "if adf.create_html:\n",
    "    adf.create_website()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11d4e7cd-49f5-4a68-85a1-ad18fa2a6706",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29b2605d-519c-4305-ab2b-47c5c053a00c",
   "metadata": {},
   "source": [
    "## ADF Helpful Methods and Structures \n",
    "\n",
    "Demonstration of a few methods to get information from the ADF object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdd8063d-195b-42bc-99a4-37ea0f9a62a0",
   "metadata": {},
   "source": [
    "#### List all adf object related methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d9703e-5053-4a3f-9745-2a6cddfa4f6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dir(adf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f3971d-14de-4d7d-996b-39e4bda99d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "adf.num_procs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4583376e-6ac3-4f0f-81d0-f4e92ce13124",
   "metadata": {},
   "source": [
    "### Get information from the subsections of the config yaml file\n",
    "\n",
    "Remember the different sub-sections?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad924a19-81ba-49dc-8e3b-2bb44f2b5114",
   "metadata": {},
   "source": [
    "#### Basic Info Section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3904cb9-a618-43af-9605-e6fcd4e7fd59",
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_info_dict = adf.read_config_var(\"diag_basic_info\")\n",
    "\n",
    "for key, val in basic_info_dict.items():\n",
    "    print(f\"{key}: {val}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8b1de2f-2a36-472f-b0ae-517845245473",
   "metadata": {},
   "source": [
    "#### Test Case Info Section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cccdb0b2-ba50-45d5-8ce5-af6a9161c762",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dict = adf.read_config_var(\"diag_cam_climo\")\n",
    "\n",
    "for key, val in test_dict.items():\n",
    "    print(f\"{key}: {val}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33a7da55-13a0-49a6-988b-ab7451b49dac",
   "metadata": {},
   "source": [
    "#### Baseline Case Info Section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf8b802-300f-41f6-b6ad-32c4f08c463f",
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_dict = adf.read_config_var(\"diag_cam_baseline_climo\")\n",
    "\n",
    "for key, val in baseline_dict.items():\n",
    "    print(f\"{key}: {val}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b878fc07-a39a-4ef6-a04c-499280d3d6cb",
   "metadata": {},
   "source": [
    "### Get information <strong>not directly</strong> from the subsections of the config yaml file\n",
    "\n",
    "This just represents a different way to get some ADF info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "721ba2f9-29c8-40ee-a813-181194f46b40",
   "metadata": {},
   "source": [
    "#### Get Case/Baseline Names\n",
    "\n",
    "This is a different wat to get case names than from the *adf.read_config_var()* method that read in data from sub-sections above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b72978-0af4-4d04-87d6-0c56de58400b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of case names (list by default)\n",
    "case_names = adf.get_cam_info(\"cam_case_name\", required=True)\n",
    "print(case_names)\n",
    "\n",
    "base_name = adf.get_baseline_info(\"cam_case_name\")\n",
    "print(base_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "898f441c-b001-4165-a614-766183b638ed",
   "metadata": {},
   "source": [
    "#### Get Case/Baseline Climo file locations\n",
    "\n",
    "Here we are calling directly from the config file, no subsection "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ba4bcb-70d7-47cd-a7c5-0b1827baaaf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_climo_loc = adf.get_cam_info(\"cam_climo_loc\", required=True)\n",
    "base_climo_loc = adf.get_baseline_info(\"cam_climo_loc\")\n",
    "case_climo_loc, base_climo_loc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "864e9b1f-fef8-4d45-94b6-5e4cc8e099c9",
   "metadata": {},
   "source": [
    "#### Get Desired Variable Names\n",
    "\n",
    "Here we are calling directly from the config file, no subsection "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff7a34b1-6795-46f3-a3b5-a40507b6aa52",
   "metadata": {},
   "outputs": [],
   "source": [
    "var_list = adf.diag_var_list\n",
    "print(var_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8504642-fcfc-482c-b5a8-7494f9fde680",
   "metadata": {},
   "source": [
    "#### Get variable defaults from adf_variable_defaults.yaml\n",
    "\n",
    "Take a look at what defaults are for <strong>TS</strong>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3a780fd-6234-4574-8535-ebf6e7bba3fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "adf.variable_defaults[\"TS\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2701110-87da-44be-8cc5-a67fe759870e",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e847946-aed1-4226-be42-46b46fc945ea",
   "metadata": {},
   "source": [
    "## Exploration of the Output Data\n",
    "\n",
    "Now that the ADF has created all the necessary timeseries/climo/regridded data let's run a quick set of functions to plot time series for RESTOM, TS, and ICEFRAC"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eddc165-065b-4874-a53d-0a53486c1817",
   "metadata": {},
   "source": [
    "##### Let's grab the case names, time series locations, variable defaults dictionary and climo years"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001f31de-2897-4cfd-a7a2-a63a47b63827",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_names = adf.get_cam_info(\"cam_case_name\", required=True)\n",
    "case_names_len = len(case_names)\n",
    "data_name = adf.get_baseline_info(\"cam_case_name\", required=False)\n",
    "\n",
    "case_ts_locs = adf.get_cam_info(\"cam_ts_loc\", required=True)\n",
    "data_ts_loc = adf.get_baseline_info(\"cam_ts_loc\", required=False)\n",
    "\n",
    "res = adf.variable_defaults  # dict of variable-specific plot preferences\n",
    "# or an empty dictionary if use_defaults was not specified in YAML.\n",
    "\n",
    "start_year = adf.climo_yrs[\"syears\"]\n",
    "end_year = adf.climo_yrs[\"eyears\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e35c52-2044-4677-8b7c-d8266bd19b20",
   "metadata": {},
   "source": [
    "### Time Series Plotting Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a123c21a-e27b-4b2b-b674-c77b7912d996",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _load_dataset(fils):\n",
    "    if len(fils) == 0:\n",
    "        print(\"Input file list is empty.\")\n",
    "        return None\n",
    "    elif len(fils) > 1:\n",
    "        return xr.open_mfdataset(fils, combine=\"by_coords\")\n",
    "    else:\n",
    "        sfil = str(fils[0])\n",
    "        return xr.open_dataset(sfil)\n",
    "    # End if\n",
    "\n",
    "\n",
    "# End def"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff599aa-4ad7-43b0-8381-a94177e40c6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _data_calcs(ts_loc, var, subset=None):\n",
    "    \"\"\"\n",
    "    args\n",
    "    ----\n",
    "     - ts_loc: Path\n",
    "         path to time series file\n",
    "\n",
    "     - var: str\n",
    "         name of variable\n",
    "\n",
    "     - subset (optional): dict\n",
    "         lat/lon extents (south, north, east, west)\n",
    "    \"\"\"\n",
    "    fils = sorted(list(Path(ts_loc).glob(f\"*{var}*.nc\")))\n",
    "\n",
    "    ts_ds = _load_dataset(fils)\n",
    "\n",
    "    time = ts_ds[\"time\"]\n",
    "    time = xr.DataArray(\n",
    "        ts_ds[\"time_bnds\"].load().mean(dim=\"nbnd\").values,\n",
    "        dims=time.dims,\n",
    "        attrs=time.attrs,\n",
    "    )\n",
    "    ts_ds[\"time\"] = time\n",
    "    ts_ds.assign_coords(time=time)\n",
    "    ts_ds = xr.decode_cf(ts_ds)\n",
    "\n",
    "    if subset != None:\n",
    "        ts_ds = ts_ds.sel(\n",
    "            lat=slice(subset[\"s\"], subset[\"n\"]), lon=slice(subset[\"w\"], subset[\"e\"])\n",
    "        )\n",
    "\n",
    "    data = ts_ds[var].squeeze()\n",
    "    unit = data.units\n",
    "\n",
    "    # global weighting\n",
    "    w = np.cos(np.radians(data.lat))\n",
    "    avg = data.weighted(w).mean(dim=(\"lat\", \"lon\"))\n",
    "\n",
    "    yrs = np.unique(\n",
    "        [\n",
    "            str(val.item().timetuple().tm_year).zfill(4)\n",
    "            for _, val in enumerate(ts_ds[\"time\"])\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return avg, yrs, unit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8141a995-7a66-4562-9ea6-472cde4e21f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_plot(ax, name, vals, yrs, unit, color_dict, linewidth=None, zorder=1):\n",
    "    \"\"\"\n",
    "    args\n",
    "    ----\n",
    "     - color_dict: dict\n",
    "         color and marker style for variable\n",
    "    \"\"\"\n",
    "\n",
    "    ax.plot(\n",
    "        yrs,\n",
    "        vals,\n",
    "        color_dict[\"marker\"],\n",
    "        c=color_dict[\"color\"],\n",
    "        label=name,\n",
    "        linewidth=linewidth,\n",
    "        zorder=zorder,\n",
    "    )\n",
    "\n",
    "    ax.set_xlabel(\"Years\", fontsize=15, labelpad=20)\n",
    "    ax.set_ylabel(unit, fontsize=15, labelpad=20)\n",
    "\n",
    "    # For the minor ticks, use no labels; default NullFormatter.\n",
    "    ax.tick_params(which=\"major\", length=7)\n",
    "    ax.tick_params(which=\"minor\", length=5)\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7aaeed2-b21f-4610-8b2e-0530d5556b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_var_details(ax, var, vals_cases, vals_base):\n",
    "    mins = []\n",
    "    maxs = []\n",
    "    for i, val in enumerate(vals_cases):\n",
    "        mins.append(np.nanmin(vals_cases[i]))\n",
    "        maxs.append(np.nanmax(vals_cases[i]))\n",
    "\n",
    "    mins.append(np.nanmin(vals_base))\n",
    "    maxs.append(np.nanmax(vals_base))\n",
    "\n",
    "    if var == \"SST\":\n",
    "        ax.set_ylabel(\"K\", fontsize=20, labelpad=12)\n",
    "        tick_spacing = 0.5\n",
    "        ax.yaxis.set_major_locator(MultipleLocator(1))\n",
    "        ax.set_title(f\"Time Series Global: {var} - ANN\", loc=\"left\", fontsize=22)\n",
    "\n",
    "    if var == \"TS\":\n",
    "        ax.set_ylabel(\"K\", fontsize=20, labelpad=12)\n",
    "        tick_spacing = 0.5\n",
    "        ax.yaxis.set_minor_locator(MultipleLocator(0.5))\n",
    "        ax.set_title(f\"Time Series Global: {var} - ANN\", loc=\"left\", fontsize=22)\n",
    "\n",
    "    if var == \"ICEFRAC\":\n",
    "        ax.set_ylabel(\"frac\", fontsize=20, labelpad=12)\n",
    "        tick_spacing = 0.1\n",
    "        ax.set_ylim(np.floor(min(mins)), np.ceil(max(maxs)))\n",
    "        ax.set_title(f\"Time Series LabSea: {var} - ANN\", loc=\"left\", fontsize=22)\n",
    "\n",
    "    if var == \"RESTOM\":\n",
    "        ax.set_ylabel(\"W/m2\", fontsize=20, labelpad=12)\n",
    "        tick_spacing = 0.5\n",
    "        ax.yaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "        ax.set_title(f\"Time Series Global: {var} - ANN\", loc=\"left\", fontsize=22)\n",
    "\n",
    "        # Set label to show if RESTOM is 1 or 5-yr avg\n",
    "        line_1yr = Line2D(\n",
    "            [],\n",
    "            [],\n",
    "            label=\"1-yr avg\",\n",
    "            color=\"k\",\n",
    "            linewidth=1,\n",
    "            marker=\"*\",\n",
    "        )\n",
    "        line_5yr = Line2D(\n",
    "            [],\n",
    "            [],\n",
    "            label=\"5-yr avg\",\n",
    "            color=\"k\",\n",
    "            linewidth=1,\n",
    "        )\n",
    "        ax.legend(handles=[line_1yr, line_5yr], bbox_to_anchor=(0.99, 0.99))\n",
    "\n",
    "    # Add extra space on the y-axis, except for ICEFRAC\n",
    "    if var != \"ICEFRAC\":\n",
    "        ax.set_ylim(np.floor(min(mins)), np.ceil(max(maxs)) + tick_spacing)\n",
    "\n",
    "    ax.yaxis.set_major_locator(MultipleLocator(tick_spacing))\n",
    "\n",
    "    ax.tick_params(axis=\"y\", which=\"major\", labelsize=16)\n",
    "    ax.tick_params(axis=\"y\", which=\"minor\", labelsize=16)\n",
    "\n",
    "    ax.tick_params(axis=\"x\", which=\"major\", labelsize=14)\n",
    "    ax.tick_params(axis=\"x\", which=\"minor\", labelsize=14)\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26eb6dfb-eec3-4c89-9964-6495891ca572",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1fa7a0d8-6d8f-45e1-b606-9507a11764dc",
   "metadata": {},
   "source": [
    "### Plot the time series!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67e22347-a3b8-4b3c-a68f-38a84ff58454",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_var_list = [\"RESTOM\", \"TS\", \"ICEFRAC\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90848bf6-403d-4a3b-a69a-b4ff12f3891b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "from matplotlib.ticker import MultipleLocator\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "fig = plt.figure(figsize=(30, 15))\n",
    "\n",
    "# Change the layout/number of subplots based off number of variables desired\n",
    "rows = 2\n",
    "cols = 3\n",
    "gs = fig.add_gridspec(rows, cols, hspace=0.3, wspace=0.2)\n",
    "\n",
    "# Rough subset for Lab Sea\n",
    "w = -63.5 + 360\n",
    "e = -47.5 + 360\n",
    "s = 53.5\n",
    "n = 65.5\n",
    "subset = {\"s\": s, \"n\": n, \"e\": e, \"w\": w}\n",
    "\n",
    "# Add more colors as needed for number of test cases\n",
    "# ** Baseline is already added as green dashed line in plotting function **\n",
    "# matplotlib colors here: https://matplotlib.org/stable/gallery/color/named_colors.html\n",
    "colors = [\"k\", \"aqua\", \"orange\", \"b\", \"magenta\", \"goldenrod\", \"slategrey\", \"rosybrown\"]\n",
    "\n",
    "# Setup plotting\n",
    "# ---------------\n",
    "\n",
    "# Loop over variables:\n",
    "for i, var in enumerate(ts_var_list):\n",
    "    print(\"Plotting variable:\", var)\n",
    "\n",
    "    if var == \"RESTOM\":\n",
    "        ax = plt.subplot(gs[0, 0])\n",
    "    if var == \"TS\":\n",
    "        ax = plt.subplot(gs[0, 1])\n",
    "    if var == \"ICEFRAC\":\n",
    "        ax = plt.subplot(gs[0, 2])\n",
    "\n",
    "    # Grab baseline case:\n",
    "    # --------------------\n",
    "\n",
    "    if var == \"RESTOM\":\n",
    "        avg_base_FSNT, yrs_base, unit = _data_calcs(data_ts_loc, \"FSNT\")\n",
    "        avg_base_FLNT, _, _ = _data_calcs(data_ts_loc, \"FLNT\")\n",
    "        if len(yrs_base) < 5:\n",
    "            print(\n",
    "                f\"Not a lot of climo years for {data_name}, only doing 1-yr avg for RESTOM...\"\n",
    "            )\n",
    "            FSNT_base = avg_base_FSNT\n",
    "            FLNT_base = avg_base_FLNT\n",
    "        else:\n",
    "            FSNT_base = avg_base_FSNT.rolling(time=60, center=True).mean()\n",
    "            FLNT_base = avg_base_FLNT.rolling(time=60, center=True).mean()\n",
    "\n",
    "        avg_base = FSNT_base - FLNT_base\n",
    "\n",
    "    if var == \"TS\" or var == \"SST\":\n",
    "        avg_base, yrs_base, unit = _data_calcs(data_ts_loc, var)\n",
    "\n",
    "    if var == \"ICEFRAC\":\n",
    "        avg_base, yrs_base, unit = _data_calcs(data_ts_loc, var, subset)\n",
    "\n",
    "    # Get int of years for plotting on x-axis\n",
    "    yrs_base_int = yrs_base.astype(int)\n",
    "\n",
    "    # Create yearly averages\n",
    "    vals_base = [avg_base.sel(time=i).mean() for i in yrs_base]\n",
    "\n",
    "    # Plot baseline data\n",
    "    color_dict = {\"color\": \"g\", \"marker\": \"--\"}\n",
    "    ax = ts_plot(ax, data_name, vals_base, yrs_base_int, unit, color_dict)\n",
    "\n",
    "    # Loop over test cases:\n",
    "    # ----------------------\n",
    "    # Create lists to hold all sets of years (for each case) and\n",
    "    # sets of var data (for each case)\n",
    "    vals_cases = []\n",
    "    yrs_cases = []\n",
    "    for case_idx, case_name in enumerate(case_names):\n",
    "        if var == \"RESTOM\":\n",
    "            avg_case_FSNT, yrs_case, unit = _data_calcs(case_ts_locs[case_idx], \"FSNT\")\n",
    "            avg_case_FLNT, _, _ = _data_calcs(case_ts_locs[case_idx], \"FLNT\")\n",
    "            if len(yrs_case) < 5:\n",
    "                print(\n",
    "                    f\"Not a lot of climo years for {case_name}, only doing 1-yr avg for RESTOM...\"\n",
    "                )\n",
    "                FSNT_case = avg_case_FSNT\n",
    "                FLNT_case = avg_case_FLNT\n",
    "                color_dict = {\"color\": colors[case_idx], \"marker\": \"-*\"}\n",
    "            else:\n",
    "                FSNT_case = avg_case_FSNT.rolling(time=60, center=True).mean()\n",
    "                FLNT_case = avg_case_FLNT.rolling(time=60, center=True).mean()\n",
    "                color_dict = {\"color\": colors[case_idx], \"marker\": \"-\"}\n",
    "\n",
    "            avg_case = FSNT_case - FLNT_case\n",
    "\n",
    "        if var == \"TS\":\n",
    "            avg_case, yrs_case, unit = _data_calcs(case_ts_locs[case_idx], var)\n",
    "            color_dict = {\"color\": colors[case_idx], \"marker\": \"-\"}\n",
    "\n",
    "        if var == \"ICEFRAC\":\n",
    "            avg_case, yrs_case, unit = _data_calcs(case_ts_locs[case_idx], var, subset)\n",
    "            color_dict = {\"color\": colors[case_idx], \"marker\": \"-\"}\n",
    "\n",
    "        # Get yearly averages for all available years\n",
    "        vals_case = [avg_case.sel(time=i).mean() for i in yrs_case]\n",
    "        vals_cases.append(vals_case)\n",
    "\n",
    "        # Get int of years for plotting on x-axis\n",
    "        yrs_case_int = yrs_case.astype(int)\n",
    "        yrs_cases.append(yrs_case_int)\n",
    "\n",
    "        # Add case to plot (ax)\n",
    "        ax = ts_plot(ax, case_name, vals_case, yrs_case_int, unit, color_dict)\n",
    "\n",
    "    # End for (case names)\n",
    "\n",
    "    # Get variable details\n",
    "    ax = plot_var_details(ax, var, vals_cases, vals_base)\n",
    "\n",
    "    # Grab all unique years and find min/max years\n",
    "    uniq_yrs = sorted(x for v in yrs_cases for x in v)\n",
    "    max_year = int(max(uniq_yrs))\n",
    "    min_year = int(min(uniq_yrs))\n",
    "\n",
    "    last_year = max_year - max_year % 5\n",
    "    if (max_year > 5) and (last_year < max_year):\n",
    "        last_year += 5\n",
    "\n",
    "    first_year = min_year - min_year % 5\n",
    "    if min_year < 5:\n",
    "        first_year = 0\n",
    "\n",
    "    ax.set_xlim(first_year, last_year)\n",
    "    ax.set_xlabel(\"Years\", fontsize=15, labelpad=20)\n",
    "    # Set the x-axis plot limits\n",
    "    # to guarantee data from all cases (including baseline) are on plot\n",
    "    ax.set_xlim(min_year, max_year + 1)\n",
    "\n",
    "    # x-axis ticks and numbers\n",
    "    if max_year - min_year > 120:\n",
    "        ax.xaxis.set_major_locator(MultipleLocator(20))\n",
    "        ax.xaxis.set_minor_locator(MultipleLocator(10))\n",
    "    if 10 <= max_year - min_year <= 120:\n",
    "        ax.xaxis.set_major_locator(MultipleLocator(5))\n",
    "        ax.xaxis.set_minor_locator(MultipleLocator(1))\n",
    "    if 0 < max_year - min_year < 10:\n",
    "        ax.xaxis.set_major_locator(MultipleLocator(1))\n",
    "        ax.xaxis.set_minor_locator(MultipleLocator(1))\n",
    "\n",
    "    # End for (case loop)\n",
    "# End for (variables loop)\n",
    "\n",
    "# Set up legend\n",
    "# Gather labels based on case names and plotted line format (color, style, etc)\n",
    "lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]\n",
    "lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]\n",
    "fig.legend(\n",
    "    lines[: case_names_len + 1],\n",
    "    labels[: case_names_len + 1],\n",
    "    loc=\"center left\",\n",
    "    fontsize=18,\n",
    "    bbox_to_anchor=(0.365, 0.4, 0.02, 0.05),\n",
    ")  # bbox_to_anchor(x0, y0, width, height)\n",
    "\n",
    "fig.show()\n",
    "\n",
    "# plt.savefig(\"TimeSeries_ANN.png\", facecolor='w',bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b037d2f-85c0-4849-a8f1-f4a18d3ddc31",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NPL 2023a",
   "language": "python",
   "name": "npl-2023a"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}