{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Utilities for observations handling\n",
    "\n",
    "In previous notebooks, the `xs.spatial.subset` and `xs.spatial_mean` methods have already been shown, but `xscen` ships with a few more spatial utilities and other functions that come handy when handling observations (point timeseries)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xclim as xc\n",
    "import xscen as xs\n",
    "\n",
    "import string\n",
    "import random\n",
    "\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeat\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Let's first create synthetic data for the example. We have here one year of daily `tas` data on a 50x50 grid with a Rotated Pole projection. We also have the same year of tas data at 102 \"stations\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = xs.testing.datablock_3d(np.random.random((366, 50, 50)), \"tas\", \"rlon\", -25, \"rlat\", -25, 1, 1, \"2000-01-01\", as_dataset=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 100\n",
    "obs = xc.testing.helpers.test_timeseries(np.random.random(366,), 'tas', start='2000-01-01', as_dataset=True)\n",
    "stations = ['00', '01'] + [''.join(random.choices(string.ascii_letters, k=4)) for i in range(N)]\n",
    "lon = [-100, -100] + list(np.random.randint(-120, -60, size=N))\n",
    "lat = [45, 45] + list(np.random.randint(30, 60, size=N))\n",
    "obs = obs.expand_dims(station=stations).assign_coords(lon=(('station',), lon), lat=(('station',), lat))\n",
    "obs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Merge colocated stations\n",
    "\n",
    "You might have noticed that the \"observation\" data has the same coordinates for the first two stations : 00 and 01 are colocated. This can happen in real life when we merge observational datasets from multiple non-independent sources (ex: GHCNd and AHCCD), in which case the same values are repeated, or when a station changes name or has multiple instruments measuring the same variable (ex: Info-Climat), in which case one station entry as values up to a certain time before switching over to the other entry. In some context, when we consider all there sources equal, we'd want to merge them and have only one timeseries per location.\n",
    "\n",
    "This can be done with `xs.spatial.merge_duplicated_stations`. It compares the `lat` and `lon` coordinates and merge all co-located elements by taking the first non-nan value along the station dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "obs = xs.spatial.merge_duplicated_stations(obs)\n",
    "obs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Voronoi weights\n",
    "When comparing observations to gridded data, one often averages some statistics over multiple stations. When done without weighting, this average treats all stations as equal, which implictely gives a stronger weight to regions with many stations. This is not always wanted. Voronoi weights are a way to enhance the geographical representativity of the average by giving less weight to stations in densely observed areas. The process divides an area with small polygons in such a way that each polygon includes exactly one station. The weights are then the fraction of the total area covered by this polygon. Further info is available on [Wikipedia](https://en.wikipedia.org/wiki/Voronoi_diagram).\n",
    "\n",
    "The xscen function does not store the polygons, it only returns the weight.\n",
    "\n",
    "Here, we limit the voronoi polygon to the total extent of the dataset (using another handy xscen function). We could also pass a `GeoDataFrame` with a collection of polygons so that the weights are computed for each region. One can see in the figure that stations outside the region were attributed a weight of 0. We also limit the maximum weight to 5% (0.05) in order to lessen the impact of outlier stations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "extent = xs.spatial.dataset_extent(ds, method='shape')\n",
    "vw = xs.spatial.voronoi_weights(obs, region=extent['shape'], maxfrac=0.05)\n",
    "obs = obs.assign_coords(weight=vw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(subplot_kw={'projection': ccrs.PlateCarree()})\n",
    "sc = ax.scatter(x=obs.lon, y=obs.lat, c=obs.weight, transform=ccrs.PlateCarree())\n",
    "ax.add_geometries([extent['shape']], facecolor='none', edgecolor='red', crs=ccrs.PlateCarree())\n",
    "ax.coastlines()\n",
    "ax.set_facecolor('grey')\n",
    "bs = extent['shape'].bounds\n",
    "ax.set_extent([bs[0], bs[2], bs[1], bs[3]], crs=ccrs.PlateCarree())\n",
    "fig.colorbar(sc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {
    "tags": []
   },
   "source": [
    "If we computed some RSME, for example, we could do the average with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds_at_obs = xs.spatial.subset(ds, 'gridpoint', lon=obs.lon, lat=obs.lat)\n",
    "rmse = ((ds_at_obs.tas - obs.tas)**2).mean('time')**0.5\n",
    "rmse.weighted(obs.weight).mean()"
   ]
  }
 ],
 "metadata": {
  "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}