{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c59d0ba",
   "metadata": {},
   "source": [
    "# Comparing modelled heads to observations\n",
    "\n",
    "This notebook showcases some methods to compare modelled heads to head observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28629a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import hydropandas as hpd\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e70ae329",
   "metadata": {},
   "source": [
    "First read in model results from the IJmuiden model (03_local_grid_refinement.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed82278c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ws = Path(\"../examples/03_local_grid_refinement\")\n",
    "model_name = \"IJm_planeten\"\n",
    "\n",
    "# read model dataset\n",
    "ds = xr.open_dataset(model_ws / f\"{model_name}.nc\")\n",
    "ds.attrs[\"model_ws\"] = model_ws\n",
    "\n",
    "# read heads\n",
    "head = nlmod.gwf.output.get_heads_da(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "283953be",
   "metadata": {},
   "source": [
    "Compute the groundwater level and plot the result for the first timestep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b9fd1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the groundwater level in each time step\n",
    "gwl = nlmod.gwf.output.get_gwl_from_wet_cells(head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d9d08a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the heads in the first aquifer\n",
    "ax = nlmod.plot.map_array(\n",
    "    gwl.isel(time=0), ds=ds, cmap=\"RdBu_r\", colorbar_label=\"head [m+NAP]\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "787b9031",
   "metadata": {},
   "source": [
    "Load the measurements and plot the locations of the observation wells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b86ef70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# df = pd.read_pickle(\"./data/20250428_bro_ijmuiden_np1_26_4.pklz\", compression=\"zip\")\n",
    "df = hpd.read_excel(\"../examples/data/20250428_bro_ijmuiden.xlsx\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ee346e",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(nlmod.grid.get_extent(ds), background=True)\n",
    "ax.plot(df.x, df.y, \"ko\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "253d650b",
   "metadata": {},
   "source": [
    "## Get the modeled heads\n",
    "\n",
    "Get the heads from the cells in which the observation wells are located. \n",
    "\n",
    "For this we use the `nlmod.layers.get_modellayers_indexer()` method which takes\n",
    "a model dataset (defining the model grid) and a dataframe (with the observation\n",
    "well metadata) as input. \n",
    "\n",
    "The dataframe must define the x,y locations of the\n",
    "observation wells, and the top and bottom of the screens. By default it is\n",
    "assumed these column names follow the hydropandas standard: `x`, `y`\n",
    "`screen_top` and `screen_bottom`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1893e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = nlmod.layers.get_modellayers_indexer(ds, df)\n",
    "idx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd812ae",
   "metadata": {},
   "source": [
    "This indexer can be used directly (if no warnings were raised or if `drop_nan_layers=True`) to obtain the heads in the cells with observation wells.\n",
    "\n",
    "````{note}\n",
    "If warnings were raised, this means there are observation wells for which the corresponding model layer could not be determined (these probably lie above or below the model). In this case the modellayer is returned as a float array and contains NaNs.\n",
    "\n",
    "Some post-processing will be necessary to be able to use the indexer e.g. dropping the NaN values:\n",
    "\n",
    "```python\n",
    "idx.dropna(\"name\", subset=[\"layer\"])\n",
    "```\n",
    "\n",
    "Additionally, the layer might also have to be renamed to get the layer names\n",
    "corresponding to the layer indices:\n",
    "\n",
    "```python\n",
    "idx[\"layer\"].values = ds[\"layer\"].values[idx[\"layer\"].astype(int)]\n",
    "```\n",
    "````\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89ee1cfe",
   "metadata": {},
   "source": [
    "Try using the indexer to get the modelled heads for each observation well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e9372c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "hsim = head.sel(**idx)\n",
    "hsim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b58bc4e",
   "metadata": {},
   "source": [
    "Get and plot the result for the a random observation well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad3c98c",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 20\n",
    "hsim.isel(name=i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed0afb0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "hsim.isel(name=i).plot(marker=\"o\", figsize=(10, 3))\n",
    "# plot observations\n",
    "df.obs.loc[hsim[\"name\"].values[i]].loc[\n",
    "    pd.Timestamp(hsim.time[0].item()) : pd.Timestamp(hsim.time[-1].item())\n",
    "].plot(y=\"value\", ax=plt.gca(), marker=\"x\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf0083ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "hmean = hsim.mean(\"time\")\n",
    "\n",
    "f, ax = nlmod.plot.get_map(nlmod.grid.get_extent(ds), background=True)\n",
    "cm = ax.scatter(\n",
    "    hmean.x,\n",
    "    hmean.y,\n",
    "    s=75,\n",
    "    c=hmean.values,\n",
    "    cmap=\"RdBu_r\",\n",
    "    edgecolors=\"k\",\n",
    "    linewidths=0.75,\n",
    ")\n",
    "cbar = f.colorbar(cm, ax=ax, label=\"head [m+NAP]\", shrink=0.85)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec5051ed",
   "metadata": {},
   "source": [
    "## Interpolating heads\n",
    "It is also possible to interpolate the heads at the locations of the observation wells. For this we need the original x, y coordinates of the observation wells as well as the layer each well is measuring in. \n",
    "\n",
    "To get all this information it can be useful to use `full_output=True` in\n",
    "`nlmod.layers.get_modellayers_indexer()`. This returns every variable that is\n",
    "necessary to compute the layer for each observation well. \n",
    "\n",
    "```{note}\n",
    "This also returns the model layer for the `screen_top` and `screen_bottom` separately, allowing you to identify observation wells spanning multiple layers.\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83c06441",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx_full = nlmod.layers.get_modellayers_indexer(ds, df, full_output=True)\n",
    "idx_full"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb91e6e",
   "metadata": {},
   "source": [
    "Now we can use `nlmod.observations.interpolate_points_ds()` to compute the interpolated heads\n",
    "at each observation well. The first argument is the data array we want to\n",
    "interpolate. The second argument is a dataset containing information about the\n",
    "location and layer for each observation well.\n",
    "\n",
    "We need to pass the correct names for each variable:\n",
    "\n",
    " - `x`, `y`: the coordinate names for the locations of the computed heads, the default is `\"x\"` and `\"y\"`\n",
    " - `xi`, `yi`: the coordinate names of the observation wells in `idx_full`, the default is `\"x\"` and `\"y\"`\n",
    " - `layer`: the layer dimension, the default is \"layer\"\n",
    "\n",
    "Our data matches the default so we don't need to adjust anything.\n",
    "\n",
    "```{note}\n",
    "For structured grids the returned x and y-coordinates in `nlmod.layers.get_modellayers_indexer()` are the coordinates corresponding to the cell centers. This way the result can be directly used for indexing a data array. The original locations of the observation wells are stored under `x_obs` `y_obs`. When using structured grids make sure to pass the correct coordinate names for `xi` and `yi` to the interpolate function.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87168985",
   "metadata": {},
   "outputs": [],
   "source": [
    "hsim_i = nlmod.observations.interpolate_to_points(head, idx_full)\n",
    "hsim_i"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ccfc95a",
   "metadata": {},
   "source": [
    "Compare the interpolated result to the earlier result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f04e8541",
   "metadata": {},
   "outputs": [],
   "source": [
    "hsim.isel(name=i).plot(marker=\"o\", figsize=(10, 3))\n",
    "hsim_i.isel(name=i).plot(marker=\"o\", ax=plt.gca());"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "049430d2",
   "metadata": {},
   "source": [
    "Plot the location of the observation well in the grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "512bc669",
   "metadata": {},
   "outputs": [],
   "source": [
    "obswell = idx_full.isel(name=i)\n",
    "extent = [obswell.x - 200, obswell.x + 200, obswell.y - 200, obswell.y + 200]\n",
    "f, ax = nlmod.plot.get_map(extent, background=True, figsize=6)\n",
    "nlmod.plot.modelgrid(ds, ax=ax)\n",
    "ax.plot(head.x, head.y, \"k.\", label=\"cell centers\")\n",
    "ax.plot(obswell.x, obswell.y, \"ro\", markersize=10, label=obswell.name.item())\n",
    "ax.legend(loc=(0, 1), frameon=False, ncol=2, fontsize=\"small\");"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}