docs/notebooks/sdba-advanced.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Statistical Downscaling and Bias-Adjustment - Advanced tools\n",
    "\n",
    "The previous notebook covered the most common utilities of `xclim.sdba` for conventionnal cases. Here we explore more advanced usage of `xclim.sdba` tools.\n",
    "\n",
    "## Optimization with dask\n",
    "\n",
    "Adjustment processes can be very heavy when we need to compute them over large regions and long timeseries. Using small groupings (like `time.dayofyear` ) adds precision and robustness, but also decuplates the load and computing complexity. Fortunately, unlike the heroic pioneers of scientific computing who managed to write parallelized Fortran, we now have [dask](https://dask.org/). With only a few parameters, we can magically distribute the computing load to multiple workers and threads.\n",
    "\n",
    "A good first read on the use of dask within xarray are the latter's [Optimization tips](http://xarray.pydata.org/en/stable/dask.html#optimization-tips).\n",
    "\n",
    "Some xclim.sdba-specific tips:\n",
    "\n",
    "* Most adjustment method will need to perform operation on the whole `time` coordinate, so it is best to optimize chunking along the other dimensions. This is often different from how public data is shared, where more universal 3D chunks are used. \n",
    "\n",
    "    Chunking of outputs can be controlled in xarray's [to_netcdf](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_netcdf.html?highlight=to_netcdf#xarray.Dataset.to_netcdf). We also suggest using [Zarr](https://zarr.readthedocs.io/en/stable/) files. According to [its creators](https://ui.adsabs.harvard.edu/abs/2018AGUFMIN33A..06A/abstract), `zarr` stores should give better performances, especially because of their better ability for parallel I/O. See [Dataset.to_zarr](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.to_zarr.html?highlight=to_zarr#xarray.Dataset.to_zarr) and this useful [rechunking package](https://rechunker.readthedocs.io).\n",
    "\n",
    "\n",
    "* One of the main bottleneck for adjustments with small groups is that dask needs to build and optimize an enormous task graph. This issue has been greatly reduced with xclim 0.27 and the use of `map_blocks` in the adjustment methods. However, not all adjustment methods use this optimized syntax.\n",
    "\n",
    "    In order to help dask, one can split the processing in parts. For splitting traning and adjustment, see [the section below](#Initializing-an-Adjustment-object-from-a-training-dataset).\n",
    "\n",
    "\n",
    "* Another massive bottleneck of parallelization of xarray is the thread-locking behaviour of some methods. It is quite difficult to isolate and avoid those locking, so one of the best workaround is to use Dask configurations with many _processes_ and few _threads_. The former do not share memory and thus are not impacted when the other is locked. However, this adds many memory transfer operations and, by experience, reduces dask's ability to parallelize some pipelines. \n",
    " \n",
    "   This is done usually by creating a dask client with a large `n_workers` and a small `threads_per_worker`. The last section of this notebook goes into a bit more details by suggesting a template for large-scale bias-adjustment pipelines.\n",
    "\n",
    "* Sometimes, datasets have auxiliary coordinates (for example : lat / lon in a rotated pole dataset). Xarray handles these variables as data variables and will **not** load them if dask is used. However, in some operations, xclim or xarray will trigger an access to those variables, triggering computations each time, since they are dask-backed. To avoid this behaviour, one can load the coordinates, or simply remove them from the inputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOESS smoothing and detrending\n",
    "\n",
    "As described in Cleveland (1979), locally weighted linear regressions are multiple regression methods using a nearest-neighbor approach. Instead of using all data points to compute a linear or polynomial regression, LOESS algorithms compute a local regression for each point in the dataset, using only the k-nearest neighbors as selected by a weighting function. This weighting function must fulfill some strict requirements, see the doc of `xclim.sdba.loess.loess_smoothing` for more details.\n",
    "\n",
    "In xclim's implementation, the user can choose between local _constancy_ ($d=0$, local estimates are weighted averages) and local _linearity_ ($d=1$, local estimates are taken from linear regressions). Two weighting functions are currently implemented : \"tricube\" ($w(x) = (1 - x^3)^3$) and \"gaussian\" ($w(x) = e^{-x^2 / 2\\sigma^2}$). Finally, the number of Cleveland's _robustifying iterations_ is controllable through `niter`. After computing an estimate of $y(x)$, the weights are modulated by a function of the distance between the estimate and the points and the procedure is started over. These iterations are made to weaken the effect of outliers on the estimate.\n",
    "\n",
    "The next example shows the application of the LOESS to daily temperature data. The black line and dot are the estimated $y$, outputs of the `sdba.loess.loess_smoothing` function, using local linear regression (passing $d = 1$), a window spanning 20% ($f = 0.2$) of the domain, the \"tricube\" weighting function and only one iteration. The red curve illustrates the weighting function on January 1st 2014, where the red circles are the nearest-neighbors used in the estimation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from xclim.sdba import loess\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Daily temperature data from xarray's tutorials\n",
    "ds = xr.tutorial.open_dataset('air_temperature').resample(time='D').mean()\n",
    "tas = ds.isel(lat=0, lon=0).air\n",
    "\n",
    "# Compute the smoothed series\n",
    "f = 0.2\n",
    "ys = loess.loess_smoothing(tas, d=1, weights='tricube', f=f, niter=1)\n",
    "\n",
    "# Plot data points and smoothed series\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(tas.time, tas, 'o', fillstyle='none')\n",
    "ax.plot(tas.time, ys, 'k')\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_ylabel('Temperature [K]')\n",
    "\n",
    "## The code below calls internal functions to demonstrate how the weights are computed. \n",
    "\n",
    "# LOESS algorithms as implemented here use scaled coordinates.\n",
    "x = tas.time\n",
    "x = (x - x[0]) / (x[-1] - x[0])\n",
    "xi = x[366]\n",
    "ti = tas.time[366]\n",
    "\n",
    "# Weighting function take the distance with all neighbors scaled by the r parameter as input\n",
    "r = int(f * tas.time.size)\n",
    "h = np.sort(np.abs(x - xi))[r]\n",
    "weights = loess._tricube_weighting(np.abs(x - xi).values / h)\n",
    "\n",
    "# Plot nearest neighbors and weighing function\n",
    "wax = ax.twinx()\n",
    "wax.plot(tas.time, weights, color='indianred')\n",
    "ax.plot(tas.time, tas.where(tas * weights > 0), 'o', color='lightcoral', fillstyle='none')\n",
    "\n",
    "ax.plot(ti, ys[366], 'ko')\n",
    "wax.set_ylabel('Weights')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LOESS smoothing can suffer from heavy boundary effects. On the previous graph, we can associate the strange bend on the left end of the line to them. The next example shows a stronger case. Usually, $\\frac{f}{2}N$ points on each side should be discarded. On the other hand, LOESS has the advantage of always staying within the bounds of the data.\n",
    "\n",
    "\n",
    "### LOESS Detrending\n",
    "\n",
    "In climate science, it can be used in the detrending process. `xclim` provides `sdba.detrending.LoessDetrend` in order to compute trend with the LOESS smoothing and remove them from timeseries.\n",
    "\n",
    "First we create some toy data with a sinusoidal annual cycle, random noise and a linear temperature increase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time = xr.cftime_range('1990-01-01', '2049-12-31', calendar='noleap')\n",
    "tas = xr.DataArray(\n",
    "   (10 * np.sin(time.dayofyear * 2 * np.pi / 365) +  # Annual variability\n",
    "    5 * (np.random.random_sample(time.size) - 0.5) +  # Random noise\n",
    "    np.linspace(0, 1.5, num=time.size)),  # 1.5 degC increase in 60 years\n",
    "    dims=('time',), coords={'time': time},\n",
    "    attrs={'units': 'degC'}, name='temperature',\n",
    ")\n",
    "tas.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we compute the trend on the data. Here, we compute on the whole timeseries (`group='time'`) with the parameters suggested above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from xclim.sdba.detrending import LoessDetrend\n",
    "\n",
    "# Create the detrending object\n",
    "det = LoessDetrend(group='time', d=0, niter=2, f=0.2)\n",
    "# Fitting returns a new object and computes the trend.\n",
    "fit = det.fit(tas)\n",
    "# Get the detrended series\n",
    "tas_det = fit.detrend(tas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "fit.ds.trend.plot(ax=ax, label='Computed trend')\n",
    "ax.plot(time, np.linspace(0, 1.5, num=time.size), label='Expected tred')\n",
    "ax.plot([time[0], time[int(0.1 * time.size)]], [0.4, 0.4], linewidth=6, color=\"gray\")\n",
    "ax.plot([time[-int(0.1 * time.size)], time[-1]], [1.1, 1.1], linewidth=6, color=\"gray\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As said earlier, this example shows how the Loess has strong boundary effects. It is recommended to remove the $\\frac{f}{2}\\cdot N$ outermost points on each side, as shown by the  gray bars in the graph above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initializing an Adjustment object from a training dataset\n",
    "\n",
    "For large scale uses, when the training step deserves its own computation and write to disk, or simply when there are multiples `sim` to be adjusted with the same training, it is helpful to be able to instantiate the Adjustment objects from the training dataset itself. This trick relies on a global attribute \"adj_params\" set on the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "# Create toy data for the example, here fake temperature timeseries\n",
    "t = xr.cftime_range('2000-01-01', '2030-12-31', freq='D', calendar='noleap')\n",
    "ref = xr.DataArray((-20 * np.cos(2 * np.pi * t.dayofyear / 365) + 2 * np.random.random_sample((t.size,)) + 273.15\n",
    "                    + 0.1 * (t - t[0]).days / 365),  # \"warming\" of 1K per decade,\n",
    "                   dims=('time',), coords={'time': t}, attrs={'units': 'K'})\n",
    "sim = xr.DataArray((-18 * np.cos(2 * np.pi * t.dayofyear / 365) + 2 * np.random.random_sample((t.size,)) + 273.15\n",
    "                    + 0.11 * (t - t[0]).days / 365),  # \"warming\" of 1.1K per decade\n",
    "                   dims=('time',), coords={'time': t}, attrs={'units': 'K'})\n",
    "\n",
    "ref = ref.sel(time=slice(None, '2015-01-01'))\n",
    "hist = sim.sel(time=slice(None, '2015-01-01'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from xclim.sdba.adjustment import QuantileDeltaMapping\n",
    "\n",
    "QDM = QuantileDeltaMapping(nquantiles=15, kind='+', group='time.dayofyear')\n",
    "QDM.train(ref, hist)\n",
    "QDM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trained `QDM` exposes the training data in the `ds` attribute, Here, we will write it to disk, read it back and initialize an new object from it. Notice the `adj_params` in the dataset, that has the same value as the repr string printed just above. Also, notice the `_xclim_adjustment` attribute that contains a json string so we can rebuild the adjustment object later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QDM.ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QDM.ds.to_netcdf('QDM_training.nc')\n",
    "ds = xr.open_dataset('QDM_training.nc')\n",
    "QDM2 = QuantileDeltaMapping.from_dataset(ds)\n",
    "QDM2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case above, creating a full object from the dataset doesn't make the most sense since we are in the same python session, with the \"old\" object still available. This method effective when we reload the training data in a different python session, say on another computer. **However, take note that there is no retrocompatiblity insurance.** If the QuantileDeltaMapping class was to change in a new xclim version, one would not be able to create the new object from a dataset saved with the old one.\n",
    "\n",
    "For the case where we stay in the same python session, it is still useful to trigger the dask computations. For small datasets, that could mean a simple `QDM.ds.load()`, but sometimes even the training data is too large to be full loaded in memory. In that case, we could also do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QDM.ds.to_netcdf('QDM_training2.nc')\n",
    "ds = xr.open_dataset('QDM_training2.nc')\n",
    "ds.attrs['title'] = 'This is the dataset, but read from disk.'\n",
    "QDM.set_dataset(ds)\n",
    "QDM.ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QDM2.adjust(sim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieving extra output diagnostics\n",
    "\n",
    "To fully understand what is happening during the bias-adjustment process, sdba can output _diagnostic_ variables, giving more visibility to what the adjustment is doing behind the scene. This behaviour, a `verbose` option, is controlled by the `sdba_extra_output` option, set with `xclim.set_options`. When `True`, `train` calls are instructed to include additional variables to the training datasets. In addition, the `adjust` calls will always output a dataset, with `scen` and, depending on the algorithm, other diagnostics variables. See the documentation of each `Adjustment` objects to see what extra variables are available.\n",
    "\n",
    "For the moment, this feature is still in construction and only a few `Adjustment` actually provide extra outputs. Please open issues on the Github repo if you have needs or ideas of interesting diagnostic variables.\n",
    "\n",
    "For example, `QDM.adjust` adds `sim_q`, which gives the quantile of each element of `sim` within its group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from xclim import set_options\n",
    "\n",
    "with set_options(sdba_extra_output=True):\n",
    "    QDM = QuantileDeltaMapping(nquantiles=15, kind='+', group='time.dayofyear')\n",
    "    QDM.train(ref, hist)\n",
    "    out = QDM.adjust(sim)\n",
    "\n",
    "out.sim_q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moving window for adjustments\n",
    "\n",
    "Some Adjustment methods require that the adjusted data (`sim`) be of the same length (same number of points) than the training data (`ref` and `hist`). This requirements often ensure conservation of statistical properties and a better representation of the climate change signal over the long adjusted timeseries.\n",
    "\n",
    "In opposition to a conventionnal \"rolling window\", here it is the _years_ that are the base units of the window, not the elements themselves. xclim implements `sdba.construct_moving_yearly_window` and `sdba.unpack_moving_yearly_window` to manipulate data in that goal. The \"construct\" function cuts the data in overlapping windows of a certain length (in years) and stacks them along a new `\"movingdim\"` dimension, alike to xarray's `da.rolling(time=win).construct('movingdim')`, but with yearly steps. The step between each window can also be controlled. This argument is an indicator of how many years overlap between each window. With a value of 1 (the default), a window will have `window - 1` years overlapping with the previous one. `step = window` will result in no overlap at all.\n",
    "\n",
    "By default, the result is chunked along this `'movingdim'` dimension. For this reason, the method is expected to be more computationally efficient (when using `dask`) than looping over the windows.\n",
    "\n",
    "Note that this results in two restrictions:\n",
    "\n",
    "1. The constructed array has the same \"time\" axis for all windows. This is a problem if the actual _year_ is of importance for the adjustment, but this is not the case for any of xclim's current adjustment methods.\n",
    "2. The input timeseries must be in a calendar with uniform year lengths. For daily data, this means only the \"360_day\", \"noleap\" and \"all_leap\" calendars are supported.\n",
    "\n",
    "The \"unpack\" function does the opposite : it concatenates the windows together to recreate the original timeseries. Only the central `step` years are kept from each window. Which means the final timeseries has  `(window - step) / 2` years missing on either side, with the extra year missing on the right in case of an odd `(window - step)`.\n",
    "\n",
    "Here, as`ref` and `hist` cover 15 years, we will use a window of 15 on sim. With a step of 2, this means the first window goes from 2000 to 2014 (inclusive). The last window goes from 2016 to 2030. `window - step = 13`, so 6 years will be missing at the beginning of the final `scen` and 7 years at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QDM = QuantileDeltaMapping(nquantiles=15, kind='+', group='time.dayofyear')\n",
    "QDM.train(ref, hist)\n",
    "\n",
    "scen_nowin = QDM.adjust(sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from xclim.sdba import construct_moving_yearly_window, unpack_moving_yearly_window\n",
    "\n",
    "sim_win = construct_moving_yearly_window(sim, window=15, step=2)\n",
    "sim_win"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scen_win = unpack_moving_yearly_window(\n",
    "    QDM.adjust(sim_win)\n",
    ")\n",
    "scen_win"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another short example, with an uneven number of years. Here `sim` goes from 2000 to 2029 (30 years instead of 31). With a step of 2 and a window of 15, the first window goes again from 2000 to 2014, but the last one is now from 2014 to 2028. The next window would be 2016-2030, but that last year doesn't exist. The final timeseries is thus from 2006 to 2021, 6 years missing at the beginning, like last time and **8** years missing at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_win = construct_moving_yearly_window(sim.sel(time=slice('2000', '2029')), window=15, step=2)\n",
    "sim_win"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim2 = unpack_moving_yearly_window(sim_win)\n",
    "sim2"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  },
  "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
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}