{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resampling raster data\n",
    "Resampling data is a very common operation when building a Modflow model. Usually it is used to project data from one grid onto the other. There are many different ways to do this. This notebook shows some examples of resampling methods that are incorporated in the `nlmod` package. These methods rely heavily on resampling methods in packages such as `rioxarray` and `scipy.interpolate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import geopandas as gpd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from matplotlib.colors import Normalize\n",
    "from scipy.interpolate import RectBivariateSpline\n",
    "from shapely.geometry import LineString, Point\n",
    "\n",
    "import nlmod\n",
    "from nlmod import resample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grid types\n",
    "\n",
    "Two different gridtypes are supported in `nlmod`:\n",
    "\n",
    "- structured grids where the cellsize is fixed for all cells\n",
    "- vertex grids where the cellsize differs locally. These grids are usually created using local grid refinement algorithms.\n",
    "\n",
    "In this notebook we define a few xarray DataArrays of structured and vertex grids. We use these grids in the next section to show the resampling functions in `nlmod`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### structured grid\n",
    "This structured grid consits of 3 x 3 cells, and has random numbers between 0 and 9."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.get_ds([950, 1250, 20050, 20350], delr=100)\n",
    "rng = np.random.default_rng(123)\n",
    "ds[\"data\"] = (\"y\", \"x\"), rng.random((len(ds.y), len(ds.x))) * 10\n",
    "\n",
    "norm=Normalize(0, 10)\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect(\"equal\")\n",
    "ds[\"data\"].plot(ax=ax, lw=0.1, edgecolor=\"k\", norm=norm);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### structured grid with nan value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"data_nan\"] = ds[\"data\"].copy()\n",
    "ds[\"data_nan\"].data[0, 1] = np.nan\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect(\"equal\")\n",
    "ds[\"data_nan\"].plot(ax=ax, lw=0.1, edgecolor=\"k\", norm=norm);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### vertex grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsv = nlmod.grid.refine(\n",
    "    ds, refinement_features=[([Point(1200, 20200)], \"point\", 1)], model_ws=\"07_resampling\"\n",
    ")\n",
    "dsv[\"data\"] = \"icell2d\", rng.random(len(dsv.data)) * 10\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect(\"equal\")\n",
    "pc = nlmod.plot.data_array(dsv[\"data\"], ds=dsv, edgecolor=\"k\", norm=norm)\n",
    "plt.colorbar(pc);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### vertex grid with nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsv[\"data_nan\"] = dsv[\"data\"].copy()\n",
    "dsv[\"data_nan\"][7] = np.nan\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect(\"equal\")\n",
    "nlmod.plot.data_array(dsv[\"data_nan\"], ds=dsv, edgecolor=\"k\", norm=norm);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured grid to fine structured grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate a finer model dataset\n",
    "ds_fine = nlmod.get_ds(extent=[950.0, 1250.0, 20050.0, 20350.0], delr=50.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_structured_data_arrays(da1, da2, method, edgecolor=\"k\"):\n",
    "    fig, axes = plt.subplots(ncols=2, figsize=(12, 6))\n",
    "    da1.plot(ax=axes[0], edgecolor=edgecolor, norm=norm)\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    axes[0].set_title(\"original grid\")\n",
    "    da2.plot(ax=axes[1], edgecolor=edgecolor, norm=norm)\n",
    "    axes[1].set_aspect(\"equal\")\n",
    "    axes[1].set_title(f\"resampled grid, method {method}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\", \"cubic\", \"average\", \"min\"]:\n",
    "    struc_out = resample.structured_da_to_ds(ds[\"data\"], ds_fine, method=method)\n",
    "    compare_structured_data_arrays(ds[\"data\"], struc_out, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\", \"cubic\", \"average\", \"mode\"]:\n",
    "    struc_out = resample.structured_da_to_ds(ds[\"data_nan\"], ds_fine, method=method)\n",
    "    compare_structured_data_arrays(ds[\"data_nan\"], struc_out, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rectangular Bivariate Spline\n",
    "\n",
    "*Note: not yet included as a method in nlmod*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interp_spline = RectBivariateSpline(\n",
    "    ds.x.values,\n",
    "    ds.y.values[::-1],\n",
    "    ds[\"data\"].values[::-1],\n",
    "    ky=min(3, len(ds.y) - 1),\n",
    "    kx=min(3, len(ds.x) - 1),\n",
    ")\n",
    "arr_out = interp_spline(ds_fine.x, ds_fine.y[::-1], grid=True)[::-1]\n",
    "struc_out = xr.DataArray(\n",
    "    arr_out, dims=(\"y\", \"x\"), coords={\"x\": ds_fine.x, \"y\": ds_fine.y}\n",
    ")\n",
    "compare_structured_data_arrays(ds[\"data\"], struc_out, \"Rectangular Bivariate Spline\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rectangular Bivariate Spline with nans\n",
    "\n",
    "*Note: not yet included as a method in nlmod*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interp_spline = RectBivariateSpline(\n",
    "    ds.x.values,\n",
    "    ds.y.values[::-1],\n",
    "    ds[\"data_nan\"].values[::-1],\n",
    "    ky=min(3, len(ds.y) - 1),\n",
    "    kx=min(3, len(ds.x) - 1),\n",
    ")\n",
    "arr_out = interp_spline(ds_fine.x, ds_fine.y[::-1], grid=True)[::-1]\n",
    "struc_out = xr.DataArray(\n",
    "    arr_out, dims=(\"y\", \"x\"), coords={\"x\": ds_fine.x, \"y\": ds_fine.y}\n",
    ")\n",
    "compare_structured_data_arrays(\n",
    "    ds[\"data_nan\"], struc_out, \"Rectangular Bivariate Spline\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured grid to locally refined grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_struct_to_vertex(struc2d, res_vertex2d_n, dsv, method):\n",
    "    fig, axes = plt.subplots(ncols=2, figsize=(12, 6))\n",
    "    struc2d.plot(ax=axes[0], edgecolor=\"k\", norm=norm)\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    axes[0].set_title(\"structured grid\")\n",
    "\n",
    "    pc = nlmod.plot.data_array(\n",
    "        res_vertex2d_n, ds=dsv, ax=axes[1], edgecolor=\"k\", norm=norm\n",
    "    )\n",
    "    plt.colorbar(pc)\n",
    "    axes[1].set_aspect(\"equal\")\n",
    "    axes[1].set_title(f\"locally refined grid, method {method}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### WIthout NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\", \"cubic\"]:\n",
    "    res_vertex2d_n = resample.structured_da_to_ds(ds[\"data\"], dsv, method=method)\n",
    "    compare_struct_to_vertex(ds[\"data\"], res_vertex2d_n, dsv, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Locally refined grid to structured grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_vertex_to_struct(vertex1, dsv, struc_out_n, method):\n",
    "    fig, axes = plt.subplots(ncols=2, figsize=(12, 6))\n",
    "    pc = nlmod.plot.data_array(vertex1, ds=dsv, ax=axes[0], edgecolor=\"k\", norm=norm)\n",
    "    plt.colorbar(pc)\n",
    "    axes[0].set_title(\"original\")\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    struc_out_n.plot(ax=axes[1], edgecolor=\"k\", norm=norm)\n",
    "    axes[1].set_title(f\"resampled, method {method}\")\n",
    "    axes[1].set_aspect(\"equal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\", \"cubic\"]:\n",
    "    struc_out_n = resample.vertex_da_to_ds(dsv[\"data\"], ds=ds, method=method)\n",
    "    compare_vertex_to_struct(dsv[\"data\"], dsv, struc_out_n, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\", \"cubic\"]:\n",
    "    struc_out_n = resample.vertex_da_to_ds(dsv[\"data_nan\"], ds=ds, method=method)\n",
    "    compare_vertex_to_struct(dsv[\"data_nan\"], dsv, struc_out_n, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fill nan values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Structured grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\"]:\n",
    "    struc2d_nan_filled = resample.fillnan_da_structured_grid(\n",
    "        ds[\"data_nan\"], method=method\n",
    "    )\n",
    "    compare_structured_data_arrays(ds[\"data_nan\"], struc2d_nan_filled, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### vertex grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_vertex_arrays(vertex1, vertex2, dsv, method):\n",
    "    fig, axes = plt.subplots(ncols=2, figsize=(12, 6))\n",
    "    pc = nlmod.plot.data_array(vertex1, ds=dsv, ax=axes[0], edgecolor=\"k\", norm=norm)\n",
    "    plt.colorbar(pc)\n",
    "    axes[0].set_title(\"original\")\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    pc = nlmod.plot.data_array(vertex2, ds=dsv, ax=axes[1], edgecolor=\"k\", norm=norm)\n",
    "    plt.colorbar(pc)\n",
    "    axes[1].set_title(f\"resampled, method {method}\")\n",
    "    axes[1].set_aspect(\"equal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"linear\"]:\n",
    "    vertex1_nan_filled = resample.fillnan_da_vertex_grid(\n",
    "        dsv[\"data_nan\"], ds=dsv, method=method\n",
    "    )\n",
    "    compare_vertex_arrays(dsv[\"data_nan\"], vertex1_nan_filled, dsv, method)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real world example\n",
    "In this example we will resample the values of the dutch Digital Terrain Model (DTM) from AHN4 to a structured grid and a vertex grid, using several methods. First we will download the AHN-information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [133000, 134000, 402000, 403000]\n",
    "ahn = nlmod.read.ahn.download_ahn4(extent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transform ahn data to structured grid\n",
    "We create a dummy dataset with a structured grid, to which we will resample the AHN-data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create an empty model dataset\n",
    "ds_ahn = nlmod.get_ds(extent, delr=100.0, layer=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm = Normalize(ahn.min(), ahn.max())\n",
    "for method in [\"nearest\", \"linear\", \"average\", \"min\", \"max\"]:\n",
    "    ahn_res = nlmod.resample.structured_da_to_ds(ahn, ds_ahn, method=method)\n",
    "\n",
    "    fig, axes = nlmod.plot.get_map(extent, ncols=2, figsize=(12, 6))\n",
    "    pc = nlmod.plot.data_array(ahn, ax=axes[0], norm=norm)\n",
    "    nlmod.plot.colorbar_inside(pc, ax=axes[0])\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    axes[0].set_title(\"original grid\")\n",
    "    pc = nlmod.plot.data_array(ahn_res, ds_ahn, ax=axes[1], edgecolor=\"k\", norm=norm)\n",
    "    nlmod.plot.colorbar_inside(pc, ax=axes[1])\n",
    "    axes[1].set_aspect(\"equal\")\n",
    "    axes[1].set_title(f\"resampled grid, method {method}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transform ahn data to vertex grid\n",
    "We create a vertex grid by refining the cells along a line from the southwest to the northeast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gdf = gpd.GeoDataFrame(\n",
    "    geometry=[LineString([(extent[0], extent[2]), (extent[1], extent[3])]).buffer(10.0)]\n",
    ")\n",
    "dsv = nlmod.grid.refine(ds_ahn, model_ws=\"07_resampling\", refinement_features=[(gdf, 1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for method in [\"nearest\", \"average\", \"min\", \"max\"]:\n",
    "    ahn_res = nlmod.resample.structured_da_to_ds(ahn, dsv, method=method)\n",
    "\n",
    "    fig, axes = nlmod.plot.get_map(extent, ncols=2, figsize=(12, 6))\n",
    "    pc = nlmod.plot.data_array(ahn, ax=axes[0], norm=norm)\n",
    "    nlmod.plot.colorbar_inside(pc, ax=axes[0])\n",
    "    axes[0].set_aspect(\"equal\")\n",
    "    axes[0].set_title(\"original grid\")\n",
    "    pc = nlmod.plot.data_array(ahn_res, dsv, ax=axes[1], edgecolor=\"k\", norm=norm)\n",
    "    nlmod.plot.colorbar_inside(pc, ax=axes[1])\n",
    "    axes[1].set_aspect(\"equal\")\n",
    "    axes[1].set_title(f\"resampled grid, method {method}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}