{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cordex domains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The domain module should give some tools to work with preconfigured or user defined domains. Domains are defined as xarray datasets that will contain dimensions and coodinates according to CF-conventions.\n",
    "\n",
    "**NOTE**: The domain module mostly focuses on working with rotated cordex domains and how they are defined in the [cordex archive specifications](https://is-enes-data.github.io/cordex_archive_specifications.pdf). However, there are some regional models that use different mappings instead of `rotated_pole` or `rotated_latitude_longitude` which we focus on. Any expertise working with those different mappings is highly welcome!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with domain information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cordex as cx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The domain module contains some useful functions to work with cordex meta data, e.g., you can get some domain grid information using"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.domain_info(\"EUR-11\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The domain information is stored in a number of csv tables. The module contains a tables dictionary that sorts the tables by resolution or project, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.domains.tables.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All available cordex domains are in those tables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.domains.table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `EUR-11` example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The heart of the module are some functions that create a dataset from the grid information, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eur11 = cx.cordex_domain(\"EUR-11\", dummy=\"topo\")\n",
    "eur11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `dummy='topo'` argument means, we want a dummy variable in the dataset to see how the domain looks like. For the dummy topography, we use the `cdo topo` operator in the background. So maybe you have to install `python-cdo`, e.g., `conda install -c conda-forge python-cdo`. Working with xarray datasets means, that we can use all the nice functions of xarray including plotting. The latest versions of `py-cordex` also come with an xarray accessor that allows you, e.g., to get a quick overview of the domain, using, e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from matplotlib import pyplot as plt\n",
    "\n",
    "# plt.figure(figsize=(8,8))\n",
    "eur11.cx.map()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the dummy topography to plot an overview:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eur11.topo.plot(cmap=\"terrain\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eur11.topo.plot(x=\"lon\", y=\"lat\", cmap=\"terrain\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a slightly more sophisticated plotting function that uses cartopy for the right [projection](https://scitools.org.uk/cartopy/docs/latest/tutorials/understanding_transform.html) with a rotated pole:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(da, pole, vmin=None, vmax=None, borders=True, title=None):\n",
    "    \"\"\"plot a domain using the right projection with cartopy\"\"\"\n",
    "    import cartopy.crs as ccrs\n",
    "    import cartopy.feature as cf\n",
    "    import matplotlib.pyplot as plt\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    projection = ccrs.PlateCarree()\n",
    "    transform = ccrs.RotatedPole(pole_latitude=pole[1], pole_longitude=pole[0])\n",
    "    # ax = plt.axes(projection=projection)\n",
    "    ax = plt.axes(projection=transform)\n",
    "    # ax.set_extent([ds_sub.rlon.min(), ds_sub.rlon.max(), ds_sub.rlat.min(), ds_sub.rlat.max()], crs=transform)\n",
    "    ax.gridlines(\n",
    "        draw_labels=True,\n",
    "        linewidth=0.5,\n",
    "        color=\"gray\",\n",
    "        xlocs=range(-180, 180, 10),\n",
    "        ylocs=range(-90, 90, 5),\n",
    "    )\n",
    "    da.plot(\n",
    "        ax=ax,\n",
    "        cmap=\"terrain\",\n",
    "        transform=transform,\n",
    "        vmin=vmin,\n",
    "        vmax=vmax,\n",
    "        x=\"rlon\",\n",
    "        y=\"rlat\",\n",
    "    )\n",
    "    ax.coastlines(resolution=\"50m\", color=\"black\", linewidth=1)\n",
    "    if borders:\n",
    "        ax.add_feature(cf.BORDERS)\n",
    "    if title is not None:\n",
    "        ax.set_title(title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pole = (\n",
    "    eur11.rotated_latitude_longitude.grid_north_pole_longitude,\n",
    "    eur11.rotated_latitude_longitude.grid_north_pole_latitude,\n",
    ")\n",
    "pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(eur11.topo, pole)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User defined domain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The domains are actually created from [a csv table](https://github.com/WCRP-CORDEX/domain-tables/blob/main/rotated-latitude-longitude.csv). The domains are created using the `create_dataset` function, in which the data from the table is fed, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.create_dataset?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create the EUR-11 domain manually from the numbers in the table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.domains.table.loc[\"EUR-11\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eur11_user = cx.create_dataset(\n",
    "    nlon=424,\n",
    "    nlat=412,\n",
    "    dlon=0.11,\n",
    "    dlat=0.11,\n",
    "    ll_lon=-28.375,\n",
    "    ll_lat=-23.375,\n",
    "    pollon=-162.00,\n",
    "    pollat=39.25,\n",
    "    dummy=\"topo\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check that this gives the same result as our preconfigured domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eur11_user.equals(eur11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now use the `create_dataset` function to create any domain as an xarray dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot all cordex-core domains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need a slightly modified plotting routine for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plots(dsets, vmin=None, vmax=None, borders=True, title=None):\n",
    "    \"\"\"plot a domain using the right projection with cartopy\"\"\"\n",
    "    import cartopy.crs as ccrs\n",
    "    import cartopy.feature as cf\n",
    "    import matplotlib.patheffects as pe\n",
    "    import matplotlib.pyplot as plt\n",
    "\n",
    "    plt.figure(figsize=(20, 10))\n",
    "    projection = ccrs.PlateCarree()\n",
    "    # transform = ccrs.RotatedPole(pole_latitude=pole[1], pole_longitude=pole[0])\n",
    "    # ax = plt.axes(projection=projection)\n",
    "    ax = plt.axes(projection=projection)\n",
    "    # ax.set_extent([ds_sub.rlon.min(), ds_sub.rlon.max(), ds_sub.rlat.min(), ds_sub.rlat.max()], crs=transform)\n",
    "    ax.gridlines(\n",
    "        draw_labels=True,\n",
    "        linewidth=0.5,\n",
    "        color=\"gray\",\n",
    "        xlocs=range(-180, 180, 15),\n",
    "        ylocs=range(-90, 90, 10),\n",
    "    )\n",
    "    # path_effects = [pe.Stroke(linewidth=50, foreground='g'), pe.Normal()]\n",
    "    for ds in dsets:\n",
    "        pole = (\n",
    "            ds.rotated_latitude_longitude.grid_north_pole_longitude,\n",
    "            ds.rotated_latitude_longitude.grid_north_pole_latitude,\n",
    "        )\n",
    "        transform = ccrs.RotatedPole(pole_latitude=pole[1], pole_longitude=pole[0])\n",
    "        ds.topo.plot(\n",
    "            ax=ax,\n",
    "            cmap=\"terrain\",\n",
    "            transform=transform,\n",
    "            vmin=vmin,\n",
    "            vmax=vmax,\n",
    "            x=\"rlon\",\n",
    "            y=\"rlat\",\n",
    "            add_colorbar=False,\n",
    "        )\n",
    "    ax.coastlines(resolution=\"50m\", color=\"black\", linewidth=1)\n",
    "    if borders:\n",
    "        ax.add_feature(cf.BORDERS)\n",
    "    if title is not None:\n",
    "        ax.set_title(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's plot all cordex core domains into one overview. We will identify them by their resolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table = cx.domains.table\n",
    "\n",
    "plots(\n",
    "    [\n",
    "        cx.cordex_domain(name, dummy=\"topo\")\n",
    "        for name in table.index\n",
    "        if table.loc[name].dlon == 0.22\n",
    "    ],\n",
    "    borders=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the grid information for interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The gridded Cordex datasets are particular usefule for regridding either with `CDO` or other interpolation packages."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use some CMIP6 model data from the ESGF to show how we can do the regridding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/remo-rcm/pyremo-data/main/orog_fx_MPI-ESM1-2-HR_historical_r1i1p1f1_gn.nc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "\n",
    "# https://raw.githubusercontent.com/remo-rcm/pyremo-data/main/orog_fx_MPI-ESM1-2-HR_historical_r1i1p1f1_gn.nc\n",
    "ds = xr.open_dataset(\"orog_fx_MPI-ESM1-2-HR_historical_r1i1p1f1_gn.nc\")\n",
    "# ds = xr.open_dataset(\n",
    "#    \"https://esgf3.dkrz.de/thredds/dodsC/cmip6/CMIP/MPI-M/MPI-ESM1-2-HR/historical/r1i1p1f1/Amon/tas/gn/v20190710/tas_Amon_MPI-ESM1-2-HR_historical_r1i1p1f1_gn_199501-199912.nc\"\n",
    "# )\n",
    "# ds = xr.open_dataset(\n",
    "#    \"http://esgf-data3.ceda.ac.uk/thredds/dodsC/esg_cmip6/CMIP6/CMIP/MOHC/UKESM1-0-LL/historical/r1i1p1f2/Amon/tas/gn/v20190406/tas_Amon_UKESM1-0-LL_historical_r1i1p1f2_gn_185001-194912.nc\"\n",
    "# )\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolation using `CDO`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use `CDO`'s python bindings to control the cdo operators. Please note, that the python bindings in the background actually execute the cdo commands on the command line. `CDO` does have a huge IO overhead since it will always write a file between each operator step and will always need data from a file on the filesystem as input. If you give an xarray dataset as input (see below), the python binding will actually trigger a `to_netcdf` call to write the input as a temporary file to disk. You should be aware of this if you use huge xarray datasets as input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will first write the EUR-11 grid into a file on the disk so that we can use it as input to `CDO`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cdo import Cdo\n",
    "\n",
    "eur11.to_netcdf(\"EUR-11_grid.nc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will remap the first timestep of the CMIP6 modeldata to the EUR-11 grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remap_cdo = Cdo().remapbil(\"EUR-11_grid.nc\", input=ds, returnXArray=\"orog\")\n",
    "remap_cdo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remap_cdo.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternative interpolation methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A nice alternative is [xesmf](https://pangeo-xesmf.readthedocs.io) since it is based on xarray and will also very nicely work with dask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xesmf as xe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regridder = xe.Regridder(ds, eur11, \"bilinear\", periodic=True)\n",
    "regridder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remap_xe = regridder(ds.orog)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remap_xe.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can easily compare both approaches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(remap_cdo - remap_xe).plot(x=\"lon\", y=\"lat\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The nice thing about `xesmf` is that it works together with xarray and will keep all meta information. Another consequence is that `xesmf` [works well with dask](https://pangeo-xesmf.readthedocs.io/en/latest/notebooks/Dask.html) and it's vectorization. That means, if we have a long time axis along we want to regrid, this can easily be parallelized using, e.g., `dask.distributed` and will also work nicely with large datasets that don't fit into memory. The `xesmf` regridder can also store and reuse regridding weights for faster interpolation."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}