notebooks/cesm_bias.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc65f494-4fbe-4af0-85ca-e3a424b56454",
   "metadata": {},
   "source": [
    "# Bias-correct CESM2 LENS temperature data using ERA5 reanalysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ca3999-b076-4dc9-8172-6d14ba082229",
   "metadata": {},
   "source": [
    "### Input Data Access\n",
    "- This notebook illustrates how to bias-correct daily temperature data from CESM2 Large Ensemble Dataset (https://www.cesm.ucar.edu/community-projects/lens2) hosted on AWS.\n",
    "- This data is open access (https://aws.amazon.com/marketplace/pp/prodview-xilranwbl2ep2#resources)\n",
    "- We will access the data using OSDF's AWS open data origin.\n",
    "- The OSDF zarr paths are obtained from an intake catalog which lives on NCAR's Research Data Archive and is publicly accessible via https.\n",
    "- We will use NCAR's origin to access the publicly available ERA5 reanalysis.\n",
    "- In summary, this notebook illustrates how we can access data from two different OSDF origins using PelicanFS and perform an interesting computation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75379e41-148d-4236-be98-9e25e9e47cf0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### Computational resources and Output data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75c0b21-370e-46ff-917a-18aee9f096bd",
   "metadata": {},
   "source": [
    "- If you don't have access to NCAR's HPC system, please select the appropriate cluster\n",
    "- All the intermediate results and the final result will be written to NCAR's GLADE storage system, which doesn't have public write access.\n",
    "- You are welcome to modify this to suit your needs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b9a4b04-ef0a-49bf-b2f2-17c5ff869737",
   "metadata": {},
   "source": [
    "#### Import package, define parameters and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "20444ba4-275a-45a7-bda5-c22b5cb5cee4",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'ESMF'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "File \u001b[0;32m~/venvs/osdf/lib/python3.10/site-packages/xesmf/util.py:8\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m----> 8\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mesmpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mESMF\u001b[39;00m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m:\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'esmpy'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 15\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mre\u001b[39;00m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# import nest_asyncio\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;66;03m# nest_asyncio.apply()\u001b[39;00m\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mxesmf\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mxe\u001b[39;00m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpyplot\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mplt\u001b[39;00m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# import cartopy as cart\u001b[39;00m\n",
      "File \u001b[0;32m~/venvs/osdf/lib/python3.10/site-packages/xesmf/__init__.py:3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m----> 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data, util\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfrontend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Regridder, SpatialAverager\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "File \u001b[0;32m~/venvs/osdf/lib/python3.10/site-packages/xesmf/util.py:10\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mesmpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mESMF\u001b[39;00m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m:\n\u001b[0;32m---> 10\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mESMF\u001b[39;00m\n\u001b[1;32m     12\u001b[0m LON_CF_ATTRS \u001b[38;5;241m=\u001b[39m {\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mstandard_name\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlongitude\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124munits\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdegrees_east\u001b[39m\u001b[38;5;124m'\u001b[39m}\n\u001b[1;32m     13\u001b[0m LAT_CF_ATTRS \u001b[38;5;241m=\u001b[39m {\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mstandard_name\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlatitude\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124munits\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdegrees_north\u001b[39m\u001b[38;5;124m'\u001b[39m}\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'ESMF'"
     ]
    }
   ],
   "source": [
    "# Display output of plots directly in Notebook\n",
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import intake\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "# import s3fs\n",
    "import seaborn as sns\n",
    "import re\n",
    "# import nest_asyncio\n",
    "# nest_asyncio.apply()\n",
    "import xesmf as xe\n",
    "import matplotlib.pyplot as plt\n",
    "# import cartopy as cart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e53607-bf0a-41a4-9f72-5769197e4570",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fsspec.implementations.http as fshttp\n",
    "from pelicanfs.core import PelicanFileSystem, PelicanMap, OSDFFileSystem "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9928a424-cf23-4a2f-bf16-afe7137d296f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask \n",
    "from dask_jobqueue import PBSCluster\n",
    "from dask.distributed import Client\n",
    "from dask.distributed import performance_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e969a1-f9fc-467c-9732-fd2f9bb44361",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This overwrites the default scheduler with a single-threaded scheduler\n",
    "dask.config.set(scheduler='synchronous')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961a7413-7c37-42a3-9891-14935fb12bd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_year0  = '1991'\n",
    "init_year1  = '2020'\n",
    "final_year0 = '2071'\n",
    "final_year1 = '2100'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9b6c48-22f6-476b-8261-7427e6186b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_daily(ds):\n",
    "    year = ds.time.dt.year\n",
    "    day = ds.time.dt.dayofyear\n",
    "\n",
    "    # assign new coords\n",
    "    ds = ds.assign_coords(year=(\"time\", year.data), day=(\"time\", day.data))\n",
    "\n",
    "    # reshape the array to (..., \"day\", \"year\")\n",
    "    return ds.set_index(time=(\"year\", \"day\")).unstack(\"time\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9355d14-86f5-4236-a986-44af18453483",
   "metadata": {},
   "outputs": [],
   "source": [
    "rda_scratch = '/glade/campaign/collections/rda/scratch/harshah'\n",
    "zarr_path   = rda_scratch + \"/tas_zarr/\"\n",
    "mean_path   = zarr_path + \"/means/\"\n",
    "stdev_path  = zarr_path + \"/stdevs/\"\n",
    "#\n",
    "rda_url        =  'https://data.rda.ucar.edu/'\n",
    "intake_url     = rda_url + 'harshah/intake_catalogs/osdf/cesm2-lens-osdf/aws-cesm2-le.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67acfb7e-79c0-4484-b0f8-2097846d7bb1",
   "metadata": {},
   "source": [
    "### Create a Dask cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c1a9cf-887d-424d-b863-2f30a48e4211",
   "metadata": {},
   "source": [
    "#### Dask Introduction\n",
    "\n",
    "Dask is a solution that enables the scaling of Python libraries. It mimics popular scientific libraries such as numpy, pandas, and xarray that enables an easier path to parallel processing without having to refactor code.\n",
    "\n",
    "There are 3 components to parallel processing with Dask: the client, the scheduler, and the workers.\n",
    "\n",
    "The Client is best envisioned as the application that sends information to the Dask cluster. In Python applications this is handled when the client is defined with client = Client(CLUSTER_TYPE). A Dask cluster comprises of a single scheduler that manages the execution of tasks on workers. The CLUSTER_TYPE can be defined in a number of different ways.\n",
    "\n",
    "- There is LocalCluster, a cluster running on the same hardware as the application and sharing the available resources, directly in Python with dask.distributed.\n",
    "\n",
    "- In certain JupyterHubs Dask Gateway may be available and a dedicated dask cluster with its own resources can be created dynamically with dask.gateway.\n",
    "\n",
    "- On HPC systems dask_jobqueue is used to connect to the HPC Slurm, PBS or HTCondor job schedulers to provision resources.\n",
    "\n",
    "The dask.distributed client python module can also be used to connect to existing clusters. A Dask Scheduler and Workers can be deployed in containers, or on Kubernetes, without using a Python function to create a dask cluster. The dask.distributed Client is configured to connect to the scheduler either by container name, or by the Kubernetes service name."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65d84b6a-cb23-42c4-be4f-d0c74eac2814",
   "metadata": {},
   "source": [
    "#### Select the Dask cluster type\n",
    "The default will be LocalCluster as that can run on any system.\n",
    "\n",
    "If running on a HPC computer with a PBS Scheduler, set to True. Otherwise, set to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcac81b8-5fdd-4e5d-b4b7-39796850545c",
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_PBS_SCHEDULER = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52991f0b-f035-46fe-9638-5a7011f2feb6",
   "metadata": {},
   "source": [
    "If running on Jupyter server with Dask Gateway configured, set to True. Otherwise, set to False."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa6ea30-bde0-423d-b0e7-0a66c16b450c",
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_DASK_GATEWAY = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbd16b1e-05ff-4d3c-a461-2aeb15997fa9",
   "metadata": {},
   "source": [
    "#### Python function for a PBS cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64d70f4e-ea6e-42aa-9430-0630d3566ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a PBS cluster object\n",
    "def get_pbs_cluster():\n",
    "    \"\"\" Create cluster through dask_jobqueue.   \n",
    "    \"\"\"\n",
    "    from dask_jobqueue import PBSCluster\n",
    "    cluster = PBSCluster(\n",
    "        job_name = 'dask-osdf-24',\n",
    "        cores = 1,\n",
    "        memory = '4GiB',\n",
    "        processes = 1,\n",
    "        local_directory = rda_scratch + '/dask/spill',\n",
    "        log_directory = rda_scratch + '/dask/logs/',\n",
    "        resource_spec = 'select=1:ncpus=1:mem=4GB',\n",
    "        queue = 'casper',\n",
    "        walltime = '3:00:00',\n",
    "        #interface = 'ib0'\n",
    "        interface = 'ext'\n",
    "    )\n",
    "    return cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50dfec14-81f7-4918-9f2a-f18493d961f8",
   "metadata": {},
   "source": [
    "#### Python function for a Gateway Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9da16d2f-4fdf-413e-89bf-e4a7e077d475",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gateway_cluster():\n",
    "    \"\"\" Create cluster through dask_gateway\n",
    "    \"\"\"\n",
    "    from dask_gateway import Gateway\n",
    "\n",
    "    gateway = Gateway()\n",
    "    cluster = gateway.new_cluster()\n",
    "    cluster.adapt(minimum=2, maximum=4)\n",
    "    return cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "731d0188-d11c-4776-b7b2-cc51ed7bfef6",
   "metadata": {},
   "source": [
    "#### Python function for a Local Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d95ce7e7-23ec-452d-9e4c-77b63b4b3dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_local_cluster():\n",
    "    \"\"\" Create cluster using the Jupyter server's resources\n",
    "    \"\"\"\n",
    "    from distributed import LocalCluster, performance_report\n",
    "    cluster = LocalCluster()    \n",
    "\n",
    "    cluster.scale(4)\n",
    "    return cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd47ae54-0432-48f7-801c-0843d49b1c15",
   "metadata": {},
   "source": [
    "#### Python logic to select the Dask Cluster type\n",
    "- This uses True/False boolean logic based on the variables set in the previous cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b72fef-ff00-4bd3-97de-4cf93e91e15f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain dask cluster in one of three ways\n",
    "\n",
    "if USE_PBS_SCHEDULER:\n",
    "    cluster = get_pbs_cluster()\n",
    "elif USE_DASK_GATEWAY:\n",
    "    cluster = get_gateway_cluster()\n",
    "else:\n",
    "    cluster = get_local_cluster()\n",
    "\n",
    "# Connect to cluster\n",
    "from distributed import Client\n",
    "client = Client(cluster)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1ab50dc-0496-455e-8b02-acf4ef314b37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scale the cluster and display cluster dashboard URL\n",
    "cluster.scale(4)\n",
    "\n",
    "cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e971bb08-6952-4e5a-8c34-5b03fe189562",
   "metadata": {},
   "source": [
    "## Load CESM LENS2 temperature data from AWS using an intake catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be45a360-624a-403b-be3f-6879ec0a5c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "cesm_cat = intake.open_esm_datastore(intake_url)\n",
    "cesm_cat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2dc1598-96bc-491a-8bae-80a90b8dbd3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cesm_temp = cesm_cat.search(variable ='TREFHTMX', frequency ='daily')\n",
    "cesm_temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a192979f-0901-4a6d-8399-f03c81023e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "cesm_temp.df['path'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc63c711-865e-4f21-86dc-03947dac7a04",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "dsets_cesm = cesm_temp.to_dataset_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ba6df1-c054-4893-a87f-9f87a6b412d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "historical_smbb  = dsets_cesm['atm.historical.daily.smbb']\n",
    "future_smbb      = dsets_cesm['atm.ssp370.daily.smbb']\n",
    "\n",
    "historical_cmip6 = dsets_cesm['atm.historical.daily.cmip6']\n",
    "future_cmip6     = dsets_cesm['atm.ssp370.daily.cmip6']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1944d35a-cf70-47b0-94a2-2861cb6c8a9b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "historical_smbb_init = historical_smbb.TREFHTMX.sel(time=slice(init_year0, init_year1))\n",
    "historical_smbb_init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f36a8b3-fdb6-4d16-8fa7-1e42da3eebc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Plot sample data\n",
    "historical_smbb.TREFHTMX.isel(member_id=0,time=0).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43bbedd-909b-4bba-bbe9-0043b6cd5fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# merge_ds_smbb = xr.concat([historical_smbb, future_smbb], dim='time')\n",
    "# merge_ds_smbb = merge_ds_smbb.dropna(dim='member_id')\n",
    "\n",
    "# merge_ds_cmip6= xr.concat([historical_cmip6, future_cmip6], dim='time')\n",
    "# merge_ds_cmip6 = merge_ds_cmip6.dropna(dim='member_id')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343c5b75-1706-455c-99a1-984c666f762f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# t_smbb      = merge_ds_smbb.TREFHTMX\n",
    "# t_cmip6     = merge_ds_cmip6.TREFHTMX\n",
    "# t_init_cmip6 = t_cmip6.sel(time=slice(init_year0, init_year1))\n",
    "# t_init_smbb  = t_smbb.sel(time=slice(init_year0, init_year1))\n",
    "# t_init       = xr.concat([t_init_cmip6,t_init_smbb],dim='member_id')\n",
    "# t_init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc673ff-036a-49d6-97c6-b48835a95906",
   "metadata": {},
   "outputs": [],
   "source": [
    "# t_init_day = to_daily(t_init)\n",
    "# #t_init_day"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3d5ac5-957b-4fa8-81b9-f6eda0601a35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# t_fut_cmip6 = t_cmip6.sel(time=slice(final_year0, final_year1))\n",
    "# t_fut_smbb  = t_smbb.sel(time=slice(final_year0, final_year1))\n",
    "# t_fut       = xr.concat([t_fut_cmip6,t_fut_smbb],dim='member_id')\n",
    "# t_fut_day   = to_daily(t_fut)\n",
    "# t_fut_day"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec85a550-5e99-4a66-98ab-25b984cd481a",
   "metadata": {},
   "source": [
    "### Save means and standard deviations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92dc4f5c-940c-4c86-82b4-43b2cd72f631",
   "metadata": {},
   "outputs": [],
   "source": [
    "# init_means   = t_init_day.mean({'year','member_id'})\n",
    "# init_stdevs  = t_init_day.std({'year','member_id'})\n",
    "# final_means  = t_fut_day.mean({'year','member_id'})\n",
    "# final_stdevs = t_fut_day.std({'year','member_id'})\n",
    "#\n",
    "# init_ensemble_means  = t_init_day.mean({'member_id'})\n",
    "# final_ensemble_means = t_fut_day.mean({'member_id'})\n",
    "# init_ensemble_means  = init_ensemble_means.chunk({'lat':192,'lon':288,'year':2,'day':365})\n",
    "# final_ensemble_means = final_ensemble_means.chunk({'lat':192,'lon':288,'year':2,'day':365})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ccc8ce3-a11e-4cfb-b563-2be618d4cfd2",
   "metadata": {},
   "source": [
    "- Save the overall means, standard devaitions and the ensemble means\n",
    "- We will regrid the 'final/EOC' ensemble means onto the ERA5 grid.\n",
    "- We will then compare it with the bias-corrected future predictions obtained from ERA5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410f94d7-2754-4e5b-82b9-398288dd4842",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# init_means.to_dataset().to_zarr(mean_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_means.zarr',mode='w')\n",
    "# init_stdevs.to_dataset().to_zarr(stdev_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_stdevs.zarr',mode='w') \n",
    "# final_means.to_dataset().to_zarr(mean_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_means.zarr',mode='w')\n",
    "# final_stdevs.to_dataset().to_zarr(stdev_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_stdevs.zarr',mode='w') \n",
    "# init_ensemble_means.to_dataset().to_zarr(mean_path + 'cesm2_'+ init_year0 + '_' + init_year1 \\\n",
    "#                                          + '_ensemble_means.zarr',mode='w')\n",
    "# final_ensemble_means.to_dataset().to_zarr(mean_path + 'cesm2_'+ final_year0 + '_' + final_year1 \\\n",
    "#                                           + '_ensemble_means.zarr',mode='w')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e248333-8828-46db-a491-5697a103b485",
   "metadata": {},
   "source": [
    "## Access ERA5 data and regrid CESM2 LENS data on the finer, ERA5 grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc3142db-6ed0-4331-b994-3b78db3d7d29",
   "metadata": {},
   "source": [
    "- In this section, we will load pre-processed ERA5 data from the NCAR origin\n",
    "- We will not use an intake catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdbb0965-b253-4be3-ae6f-05e38fe431cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create OSDF file path for the ERA5 zarr store\n",
    "namespace        = 'ncar/'   #NCAR's Research Data Archive\n",
    "# osdf_director    = 'https://osdf-director.osg-htc.org/'\n",
    "osdf_director    = 'osdf:///'\n",
    "era5_zarr_path   = namespace + 'rda/harshah/era5_tas/zarr/e5_tas2m_daily_1940_2023.zarr'\n",
    "#\n",
    "# pelfs = PelicanFileSystem(osdf_director)\n",
    "pelfs = OSDFFileSystem()\n",
    "#\n",
    "# pel_zarr         = PelicanMap(era5_zarr_path, pelfs)\n",
    "print(era5_zarr_path)\n",
    "osdf_protocol_era5path = osdf_director + era5_zarr_path\n",
    "print(osdf_protocol_era5path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1153f99-700e-49a8-ac1e-79c5f28fcfd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "tas_obs_daily = xr.open_zarr(osdf_protocol_era5path).VAR_2T\n",
    "tas_obs_init = tas_obs_daily.sel(time=slice(init_year0, init_year1))\n",
    "tas_obs_init"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a3b6339-fe60-49a7-8ed5-3fe268cd80d8",
   "metadata": {},
   "source": [
    "### Perform Bias Correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a0e4932-7f2f-44a8-a934-39853a1667ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_means_ds = xr.open_zarr(mean_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_means.zarr')\n",
    "init_means    = init_means_ds.TREFHTMX\n",
    "init_means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "686d452c-c1ea-4b67-8f0a-0301ac92763a",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_means  = xr.open_zarr(mean_path  + 'cesm2_'+ final_year0 + '_' + final_year1+ '_means.zarr').TREFHTMX\n",
    "init_stdevs  = xr.open_zarr(stdev_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_stdevs.zarr').TREFHTMX\n",
    "final_stdevs = xr.open_zarr(stdev_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_stdevs.zarr').TREFHTMX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73737f3e-ee1d-4bf8-88cc-a4773158b2a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create output grid\n",
    "ds_out = xr.Dataset(\n",
    "    coords={\n",
    "        'latitude': tas_obs_init.coords['latitude'],\n",
    "        'longitude': tas_obs_init.coords['longitude']\n",
    "    }\n",
    ")\n",
    "ds_out = ds_out.rename({'latitude':'lat','longitude':'lon'})\n",
    "ds_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ee7e7fe-fca4-4bb6-9846-0060e1744b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tas_obs_initial    = to_daily(tas_obs_init)\n",
    "tas_obs_initial    = tas_obs_initial.rename({'latitude':'lat','longitude':'lon'})\n",
    "tas_obs_initial    = tas_obs_initial.chunk({'lat':139,'lon':544,'year':3,'day':90})\n",
    "tas_obs_initial "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e76d693-5ec0-4aa6-8c7c-a9dbb935a92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# tas_obs_initial.to_dataset().to_zarr(zarr_path + \"e5_tas2m_initial_1991_2020.zarr\",mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bd94a5f-d6cb-4f24-a71b-d7940e66bdaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "tas_obs_initial = xr.open_zarr(zarr_path + \"e5_tas2m_initial_1991_2020.zarr\").VAR_2T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5667476c-d546-4721-a221-c3febf878ac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "regridder = xe.Regridder(init_means_ds, ds_out, \"bilinear\")\n",
    "regridder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca98aefa-61b8-4e30-b34c-a3c716ef2334",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_means_regrid = regridder(init_means, keep_attrs=True)\n",
    "init_means_regrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584bdf59-6f0f-403d-b8fa-9cc08c4d925a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Regrid other variables\n",
    "init_stdevs_regrid  = regridder(init_stdevs, keep_attrs=True)\n",
    "final_means_regrid  = regridder(final_means, keep_attrs=True)\n",
    "final_stdevs_regrid = regridder(final_stdevs, keep_attrs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "665b1cd5-0d89-48ae-b5af-9091997d43c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# #Save regridded data\n",
    "# init_means_regrid.to_dataset().to_zarr(mean_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_means_regridded.zarr',mode='w')\n",
    "# init_stdevs_regrid.to_dataset().to_zarr(stdev_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_stdevs_regridded.zarr',mode='w') \n",
    "# final_means_regrid.to_dataset().to_zarr(mean_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_means_regridded.zarr',mode='w')\n",
    "# final_stdevs_regrid.to_dataset().to_zarr(stdev_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_stdevs_regridded.zarr',mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d1d097d-32bd-4165-9e62-10c006513435",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Open regridded data\n",
    "init_means_regrid  = xr.open_zarr(mean_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_means_regridded.zarr').TREFHTMX\n",
    "init_stdevs_regrid = xr.open_zarr(stdev_path + 'cesm2_'+ init_year0 + '_' + init_year1+ '_stdevs_regridded.zarr').TREFHTMX\n",
    "final_means_regrid  = xr.open_zarr(mean_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_means_regridded.zarr').TREFHTMX\n",
    "final_stdevs_regrid = xr.open_zarr(stdev_path + 'cesm2_'+ final_year0 + '_' + final_year1+ '_stdevs_regridded.zarr').TREFHTMX"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ca7a7ca-30b1-4048-acf4-161d9cdc91ea",
   "metadata": {},
   "source": [
    "## Now, perform bias correction by only adjusting the first moment, i.e, mean and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e94183-bcf4-4803-84c0-aad9c2cac2d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tas_bc = tas_obs_initial  + (final_means_regrid - init_means_regrid)\n",
    "tas_bc = tas_bc.chunk({'lat':139,'lon':544,'year':3,'day':90})\n",
    "tas_bc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad63b9e2-6e4b-465d-bc60-684429a46e0c",
   "metadata": {},
   "source": [
    "### Plot bias corrected temperature and CESM model's predictions for the End of the 21st century (2100)\n",
    "- Since tas_bc are predictions for the years 2070-2100, we need to change the year coordinated\n",
    "- We will then save the bias-corrected surface air temperatures (tas) to a zarr store.\n",
    "- Finally, we will read from this zarr store and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d77ede3c-1674-4d32-b34b-5474de11411f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the year coordinate\n",
    "tas_bc['year'] = tas_bc['year'] + 80\n",
    "tas_bc         = tas_bc.rename('bias_corrected_tas')\n",
    "tas_bc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bed64597-bd12-4ba8-b53c-d91987377c5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# tas_bc.to_dataset().to_zarr(zarr_path + 'bias_corrected_tas_1991_2020.zarr',mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c987333-de8b-45f6-b690-7c7f3c42ff23",
   "metadata": {},
   "outputs": [],
   "source": [
    "tas_bc = xr.open_zarr(zarr_path + 'bias_corrected_tas_1991_2020.zarr').bias_corrected_tas\n",
    "tas_bc = tas_bc.sortby('lat',ascending=True)\n",
    "tas_bc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e697119b-d081-4581-9452-71833ac59fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_ensemble_means = xr.open_zarr(mean_path + 'cesm2_'+ final_year0 + '_' + final_year1 + '_ensemble_means.zarr').TREFHTMX\n",
    "final_ensemble_means = final_ensemble_means.sortby('lat',ascending=True)\n",
    "final_ensemble_means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93349d05-ec38-48d9-95a5-55715c435aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "tas_bc.sel(year = 2100, day = 211).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa03e7c9-537d-4393-9d8b-218d6ba8e73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_ensemble_means.sel(year = 2100, day = 211).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abed4842-ead0-4a97-b04e-f3069f4bf1e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "fig, axs = plt.subplots(1, 2, figsize=(10, 5))\n",
    "\n",
    "# Plot the first array\n",
    "tas_bc.sel(year = 2100, day = 211).plot(ax = axs[0],cmap='RdBu_r',add_colorbar=False)\n",
    "axs[0].set_title('Bias-corrected projections')\n",
    "#axs[0].coastlines(color=\"black\")\n",
    "\n",
    "# Plot the second array\n",
    "final_ensemble_means.sel(year = 2100, day = 211).plot(ax = axs[1],cmap='RdBu_r')\n",
    "axs[1].set_title('EOC model ensemble means')\n",
    "\n",
    "# Display the plots\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python:osdf",
   "language": "python",
   "name": "osdf"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}