{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95e17c3b",
   "metadata": {},
   "source": [
    "# Working with grid rotation\n",
    "\n",
    "Rotated grids are supported in nlmod. It is implemented in the following manner:\n",
    "\n",
    "- `angrot`, `xorigin` and `yorigin` (naming identical to modflow 6) are added to the attributes of the model Dataset.\n",
    "- `angrot` is the counter-clockwise rotation angle (in degrees) of the model grid coordinate system relative to a real-world coordinate system (identical to definition in modflow 6)\n",
    "- when a grid is rotated:\n",
    "    - `x` and `y` (and `xv` and `yv` for a vertex grid) are in model coordinates, instead of real-world coordinates.\n",
    "    - `xc` and `yc` are added to the Dataset and represent the cell centers in real-world coordinates (naming identical to rioxarray rotated grids)\n",
    "    - the plot-methods in nlmod plot the grid in model coordinates by default (can be overridden by the setting the parameter `rotated=True`)\n",
    "    - before intersecting with the grid, GeoDataFrames are automatically transformed to model coordinates.\n",
    "\n",
    "When grids are not rotated, the model Dataset does not contain an attribute named `angrot` (or it is 0). The x- and y-coordinates of the model then respresent real-world coordinates.\n",
    "\n",
    "In this notebook we generate a model of 1 by 1 km, with a grid that is rotated 10 degrees relative to the real-world coordinates system (EPSG:28992: RD-coordinates)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d30c1fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import matplotlib\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe4ef601",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67dffb64",
   "metadata": {},
   "source": [
    "## Generate a model Dataset\n",
    "We generate a model dataset with a rotation of 10 degrees counterclockwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df7f488",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.get_ds(\n",
    "    [0, 1000, 0, 1000],\n",
    "    angrot=10.0,\n",
    "    xorigin=200_000,\n",
    "    yorigin=500_000,\n",
    "    delr=10.0,\n",
    "    model_name=\"nlmod\",\n",
    "    model_ws=\"02_grid_rotation\",\n",
    ")\n",
    "\n",
    "ds = nlmod.time.set_ds_time(ds, time=\"2023-01-01\", start=\"2013-01-01\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee27ce59",
   "metadata": {},
   "source": [
    "## Use a disv-grid\n",
    "We call the refine method to generate a vertex grid (with the option of grid-refinement), instead of a structured grid. We can comment the next line to model a structured grid, and the rest of the notebook will run without problems as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8dbc46c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.grid.refine(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "505b37e5",
   "metadata": {},
   "source": [
    "## Add AHN\n",
    "Download the ahn, resample to the new grid (using the method 'average') and compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fba7c32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download AHN\n",
    "extent = nlmod.grid.get_extent(ds)\n",
    "ahn = nlmod.read.ahn.download_ahn3(extent)\n",
    "\n",
    "# Resample to the grid\n",
    "ds[\"ahn\"] = nlmod.resample.structured_da_to_ds(ahn, ds, method=\"average\")\n",
    "\n",
    "# Compare original ahn to the resampled one\n",
    "f, axes = nlmod.plot.get_map(extent, ncols=2)\n",
    "norm = matplotlib.colors.Normalize()\n",
    "pc = nlmod.plot.data_array(ahn, ax=axes[0], norm=norm)\n",
    "nlmod.plot.colorbar_inside(pc, ax=axes[0])\n",
    "pc = nlmod.plot.data_array(\n",
    "    ds[\"ahn\"], ds=ds, ax=axes[1], rotated=True, norm=norm, edgecolor=\"face\"\n",
    ")\n",
    "nlmod.plot.colorbar_inside(pc, ax=axes[1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8441d8b2",
   "metadata": {},
   "source": [
    "## Download surface water\n",
    "Download BGT-polygon data, add stage information from the waterboard, and grid the polygons. Because we use a rotated grid, the bgt-polygons are in model coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e7cc6b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "bgt = nlmod.read.bgt.download_bgt(extent)\n",
    "bgt = nlmod.gwf.surface_water.add_stages_from_waterboards(bgt, extent=extent)\n",
    "bgt = nlmod.grid.gdf_to_grid(bgt, ds).set_index(\"cellid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111ac670",
   "metadata": {},
   "source": [
    "## Download knmi-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc70036e",
   "metadata": {},
   "outputs": [],
   "source": [
    "knmi_ds = nlmod.read.knmi.get_recharge(ds)\n",
    "ds.update(knmi_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883fae7d",
   "metadata": {},
   "source": [
    "## Generate flopy-model\n",
    "We generate a simulation and a groundwater flow model, with some standard packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "939f3a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%\n",
    "# create simulation\n",
    "sim = nlmod.sim.sim(ds)\n",
    "\n",
    "# create time discretisation\n",
    "tdis = nlmod.sim.tdis(ds, sim)\n",
    "\n",
    "# create ims\n",
    "ims = nlmod.sim.ims(sim, complexity=\"complex\")\n",
    "\n",
    "# create groundwater flow model\n",
    "gwf = nlmod.gwf.gwf(ds, sim)\n",
    "\n",
    "# Create discretization\n",
    "dis = nlmod.gwf.dis(ds, gwf)\n",
    "\n",
    "# create node property flow\n",
    "npf = nlmod.gwf.npf(ds, gwf)\n",
    "\n",
    "# Create the initial conditions package\n",
    "ic = nlmod.gwf.ic(ds, gwf, starting_head=0.0)\n",
    "\n",
    "# Create the output control package\n",
    "oc = nlmod.gwf.oc(ds, gwf)\n",
    "\n",
    "# create recharge package\n",
    "rch = nlmod.gwf.rch(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703913be",
   "metadata": {},
   "source": [
    "## Add surface water\n",
    "To the groundwater flow model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1161cf59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add surface water with a winter and a summer stage\n",
    "# (which are both added with about half their conductance in a steady state simulation)\n",
    "drn = nlmod.gwf.surface_water.gdf_to_seasonal_pkg(bgt, gwf, ds, print_input=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d5fffa",
   "metadata": {},
   "source": [
    "## Run the model and read the heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e91eed52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run the model\n",
    "nlmod.sim.write_and_run(sim, ds)\n",
    "\n",
    "# read the heads\n",
    "head = nlmod.gwf.get_heads_da(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a2886ae",
   "metadata": {},
   "source": [
    "## Plot the heads in layer 1\n",
    "When grid rotation is used, nlmod.plot.data_array() plots a DataArray in model coordinates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba48f036",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(ds.extent)\n",
    "pc = nlmod.plot.data_array(head.sel(layer=1).mean(\"time\"), ds=ds, edgecolor=\"k\")\n",
    "cbar = nlmod.plot.colorbar_inside(pc)\n",
    "bgt.plot(ax=ax, edgecolor=\"k\", facecolor=\"none\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf81227",
   "metadata": {},
   "source": [
    "If we want to plot in realworld coordinates, we set the optional parameter `rotated=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24dd383d",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(extent)\n",
    "pc = nlmod.plot.data_array(\n",
    "    head.sel(layer=1).mean(\"time\"), ds=ds, edgecolor=\"k\", rotated=True\n",
    ")\n",
    "cbar = nlmod.plot.colorbar_inside(pc)\n",
    "# as the surface water shapes are in model coordinates, we need to transform them\n",
    "# to real-world coordinates before plotting\n",
    "affine = nlmod.grid.get_affine_mod_to_world(ds)\n",
    "bgt_rw = nlmod.grid.affine_transform_gdf(bgt, affine)\n",
    "bgt_rw.plot(ax=ax, edgecolor=\"k\", facecolor=\"none\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91472ed9",
   "metadata": {},
   "source": [
    "Export the model dataset to a netcdf-file, which you can open in qgis using 'Add mesh layer'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba33425e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = os.path.join(ds.model_ws, \"ugrid_ds.nc\")\n",
    "nlmod.gis.ds_to_ugrid_nc_file(ds, fname)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}