{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "from matplotlib import rcParams\n",
    "\n",
    "rcParams[\"figure.dpi\"] = 300\n",
    "rcParams[\"font.size\"] = 8\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{{ prolog }}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create 3D mask of approximate fractional overlap\n",
    "\n",
    "Working with fractional overlap - indicating how much of the grid cell is covered by the region - can help to create more exact regional means, allow to exclude gridpoints from regions if the overlap is too small etc.. Since v0.12.0 regionmask can create a 3D masks with the approximate fractional overlap of a set of regions for equally-spaced latitude and longitude grids. The mask is a float `xarray.DataArray` with shape `region x lat x lon`, with the overlap given as fraction (i.e. between 0 and 1). Many concepts were already introduced in {doc}`3D masks<mask_3D>`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{attention}\n",
    "\n",
    "   There are three caveats when creating fractional overlaps:\n",
    "\n",
    "   1. The passed longitude and latitude coordinates must be equally spaced - otherwise an ``InvalidCoordsError`` is raised.\n",
    "   2. Calculating the fractional overlap can be memory intensive (especially when passing many coordinates).\n",
    "   3. The resulting mask is correct to about 0.05 (i.e. 5%).\n",
    "\n",
    "Despite these restrictions using an approximation has advantages over calculating the exact overlap - it is considerably faster. \n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import libraries and check the regionmask version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "import regionmask\n",
    "\n",
    "# don't expand data\n",
    "xr.set_options(display_style=\"text\", display_expand_data=False, display_width=60)\n",
    "\n",
    "regionmask.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a mask\n",
    "\n",
    "Define a lon/ lat grid with a 2° grid spacing, where the points define the center of the grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lon = np.arange(0, 360, 2)\n",
    "lat = np.arange(90, -91, -2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `mask_3D_frac_approx` calculates the fractional overlap of each gridpoint with each region. Here using the AR6 land regions (Iturbide et al., [2020](https://doi.org/10.5194/essd-12-2959-2020)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = regionmask.defined_regions.ar6.land.mask_3D_frac_approx(lon, lat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Illustration\n",
    "\n",
    "As mentioned, `mask` is a float `xarray.DataArray` with shape `region x lat x lon`. It contains `region` (=`numbers`) as dimension coordinate as well as `abbrevs` and `names` as non-dimension coordinates (see the xarray docs for the details on the [terminology](http://xarray.pydata.org/en/stable/terminology.html)). \n",
    "\n",
    "The regions northwestern North-America, and northeastern North-America (regions 1, and 2) look as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cartopy.crs as ccrs\n",
    "\n",
    "fg = mask.sel(region=slice(1, 2)).plot(\n",
    "    subplot_kws=dict(projection=ccrs.PlateCarree()),\n",
    "    col=\"region\",\n",
    "    col_wrap=2,\n",
    "    transform=ccrs.PlateCarree(),\n",
    "    add_colorbar=True,\n",
    "    aspect=1.5,\n",
    "    cmap=\"Blues\",\n",
    "    cbar_kwargs={\"pad\": 0.01, \"shrink\": 0.65},\n",
    ")\n",
    "\n",
    "fg.cbar.set_label(\"Fractional overlap\")\n",
    "\n",
    "for ax in fg.axs.flatten():\n",
    "    regionmask.defined_regions.ar6.land.plot(\n",
    "        ax=ax, add_label=False, line_kws={\"lw\": 0.5}\n",
    "    )\n",
    "    ax.set_extent([-172, -47.5, 35, 90], ccrs.PlateCarree())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with a 3D mask\n",
    "\n",
    "masks can be used to select data in a certain region and to calculate regional averages. Let's illustrate this with a 'real' dataset - the example data is a temperature field over North America."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "airtemps = xr.tutorial.load_dataset(\"air_temperature\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An xarray object can be passed to the `mask_3D_frac_approx` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_3D_frac_approx = regionmask.defined_regions.ar6.land.mask_3D_frac_approx(airtemps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As `airtemps` has another grid than the example above, the resulting mask looks different:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fg = mask_3D_frac_approx.sel(region=slice(1, 2)).plot(\n",
    "    subplot_kws=dict(projection=ccrs.PlateCarree()),\n",
    "    col=\"region\",\n",
    "    col_wrap=2,\n",
    "    transform=ccrs.PlateCarree(),\n",
    "    add_colorbar=True,\n",
    "    aspect=1.5,\n",
    "    cmap=\"Blues\",\n",
    "    cbar_kwargs={\"pad\": 0.01, \"shrink\": 0.65},\n",
    ")\n",
    "\n",
    "fg.cbar.set_label(\"Fractional overlap\")\n",
    "\n",
    "for ax in fg.axs.flatten():\n",
    "    regionmask.defined_regions.ar6.land.plot(\n",
    "        ax=ax, add_label=False, line_kws={\"lw\": 0.5}\n",
    "    )\n",
    "    ax.set_extent([-172, -47.5, 35, 90], ccrs.PlateCarree())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use an overlap threshold\n",
    "\n",
    "To restrict the region to gridcells that overlap more to a certain threshold, grid points can be masked out using `where`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 0.5\n",
    "\n",
    "mask_3D_ge050 = mask_3D_frac_approx.where(mask_3D_frac_approx >= threshold, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This sets all grid points with an overlap of less than 50% to 0. The second options is to convert the fractional mask to a boolean one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_3D_bool = mask_3D_frac_approx >= threshold"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate weighted regional averages\n",
    "\n",
    "As for the boolan 3D mask, we can calculate the regional averages using fractional mask. In this case each grid point contributes according to its overlap and area. As proxy of the grid cell area we use `cos(lat)`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "It is better to use a model's original grid cell area (e.g. ``areacella``). ``cos(lat)`` works reasonably well for regular lat/ lon grids. For \n",
    "irregular grids (regional models, ocean models, ...) it is not appropriate.\n",
    "\n",
    ":::\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = np.cos(np.deg2rad(airtemps.lat))\n",
    "\n",
    "ts_airtemps_regional = airtemps.weighted(mask_3D_frac_approx * weights).mean(\n",
    "    dim=(\"lat\", \"lon\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is almost the same as for the boolean 3D mask: by multiplying `mask_3D * weights` we get a DataArray where the fractional overlap is scaled by the grid cell area. `airtemps.weighted(mask_3D * weights).mean([\"lat\", \"lon\"])` calculates the weighted `mean` over the lat and lon dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_airtemps_regional"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The regionally-averaged time series can be plotted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_airtemps_regional.air.sel(region=slice(0, 2)).plot(col=\"region\", col_wrap=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Iturbide, M., Gutiérrez, J. M., Alves, L. M., Bedia, J., Cerezo-Mota, R., Cimadevilla, E., Cofiño, A. S., Di Luca, A., Faria, S. H., Gorodetskaya, I. V., Hauser, M., Herrera, S., Hennessy, K., Hewitt, H. T., Jones, R. G., Krakovska, S., Manzanas, R., Martínez-Castro, D., Narisma, G. T., Nurhati, I. S., Pinto, I., Seneviratne, S. I., van den Hurk, B., and Vera, C. S.: An update of IPCC climate reference regions for subcontinental analysis of climate model data: definition and aggregated datasets, Earth Syst. Sci. Data, 12, 2959–2970, https://doi.org/10.5194/essd-12-2959-2020, 2020. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-regionmask-docs]",
   "language": "python",
   "name": "conda-env-.conda-regionmask-docs-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}