{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a3f62867",
   "metadata": {},
   "source": [
    "# Plot methods in nlmod\n",
    "\n",
    "This notebook shows different methods of plotting data with nlmod. \n",
    "\n",
    "There are many ways to plot data and it depends on the type of data and plot which of\n",
    "these method is the most convenient:\n",
    "\n",
    "- using `nlmod.plot` utilities\n",
    "- using `flopy` plot methods\n",
    "- using `xarray` plot methods\n",
    "\n",
    "The default plot methods in nlmod use a model Dataset as input (this is an xarray\n",
    "Dataset with some required variables and attributes). These plotting methods are\n",
    "accessible through `nlmod.plot`.\n",
    "\n",
    "Flopy contains its own plotting utilities and nlmod contains some wrapper functions that\n",
    "use flopy's plotting utilities under the hood. These require a flopy modelgrid or model\n",
    "object. These plotting methods are accessible through `nlmod.plot.flopy`.\n",
    "\n",
    "Finally, xarray also allows plotting of data with `.plot()`. This is used in a few\n",
    "cases in this notebook but for more detailed information, refer to the \n",
    "[xarray documentation](https://xarray.pydata.org/en/v2023.08.0/gallery.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9ae5a49",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import flopy\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "from shapely.geometry import LineString\n",
    "\n",
    "import nlmod\n",
    "from nlmod.plot import DatasetCrossSection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5a389c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ca2bf5",
   "metadata": {},
   "source": [
    "First we read a fully run model, from the notebook 09_schoonhoven.ipynb. Please run that notebook first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca958d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"Schoonhoven\"\n",
    "model_ws = os.path.join(\"..\", \"examples\", \"09_schoonhoven\")\n",
    "fname_nc = os.path.join(model_ws, f\"{model_name}.nc\")\n",
    "if not os.path.isfile(fname_nc):\n",
    "    raise (\n",
    "        Exception(\n",
    "            f\"{fname_nc} not found. Please run notebook 09_schoonhoven.ipynb in the 'examples' directory first\"\n",
    "        )\n",
    "    )\n",
    "ds = xr.open_dataset(fname_nc)\n",
    "ds.attrs[\"model_ws\"] = model_ws\n",
    "ds[\"head\"] = nlmod.gwf.output.get_heads_da(ds)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3909cdd5",
   "metadata": {},
   "source": [
    "For the flopy plot-methods we need a modelgrid object. We generate this from the model Dataset using the method. nlmod.grid.modelgrid_from_ds()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36f98ba6",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid = nlmod.grid.modelgrid_from_ds(ds)\n",
    "modelgrid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f594a18d",
   "metadata": {},
   "source": [
    "## Maps\n",
    "We can plot variables on a map using nlmod.plot.data_array(). We can also use the PlotMapView-class from flopy, and plot an array using the plot_array method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909de9e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(ds.extent, ncols=2)\n",
    "\n",
    "# plot using nlmod\n",
    "pc = nlmod.plot.data_array(ds[\"top\"], ds=ds, ax=ax[0])\n",
    "\n",
    "# plot using flopy\n",
    "pmv = flopy.plot.PlotMapView(modelgrid=modelgrid, ax=ax[1])\n",
    "pmv.plot_array(ds[\"top\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98368c1d",
   "metadata": {},
   "source": [
    "## Cross-sections\n",
    "We can also plot cross-sections, either with DatasetCrossSection in nlmod, or using the PlotCrossSection class of flopy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32a89356",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = (ds.extent[2] + ds.extent[3]) / 2 + 0.1\n",
    "line = [(ds.extent[0], y), (ds.extent[1], y)]\n",
    "zmin = -100.0\n",
    "zmax = 10.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "044aa8f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(figsize=(10, 5), nrows=2)\n",
    "\n",
    "# plot using nlmod\n",
    "dcs = DatasetCrossSection(ds, line=line, zmin=zmin, zmax=zmax, ax=ax[0])\n",
    "dcs.plot_array(ds[\"kh\"])\n",
    "\n",
    "# plot using flopy\n",
    "pcs = flopy.plot.PlotCrossSection(modelgrid=modelgrid, line={\"line\": line}, ax=ax[1])\n",
    "pcs.plot_array(ds[\"kh\"])\n",
    "pcs.ax.set_ylim((zmin, zmax));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fba12db0",
   "metadata": {},
   "source": [
    "With the DatasetCrossSection in `nlmod` it is also possible to plot the layers according to the official colors of REGIS, to plot the layer names on the plot, or to plot the model grid in the cross-section. An example is shown in the plot below.\n",
    "\n",
    "The location of the cross-section and the cross-section labels can be added using `nlmod.plot.inset_map()` and `nlmod.plot.add_xsec_line_and_labels()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45bc5231",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(figsize=(10, 5))\n",
    "dcs = DatasetCrossSection(ds, line=line, zmin=-200, zmax=10, ax=ax)\n",
    "colors = nlmod.read.regis.get_legend()\n",
    "dcs.plot_layers(colors=colors, min_label_area=1000)\n",
    "dcs.plot_grid(vertical=False, linewidth=0.5)\n",
    "mapax = nlmod.plot.inset_map(ax, ds.extent)\n",
    "nlmod.plot.add_xsec_line_and_labels(line, ax, mapax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5f4fed5",
   "metadata": {},
   "source": [
    "Another feature of the `DatasetCrossSection` is the `plot_wells` method to plot the screen depths in a cross section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a85292d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data frame with well data (random)\n",
    "nwells = 11\n",
    "max_dist = 1000\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"screen_top\": np.random.randint(-80, -40, nwells),\n",
    "        \"screen_bottom\": np.random.randint(-125, -81, nwells),\n",
    "        \"x\": np.random.randint(ds.extent[0], ds.extent[1], nwells),\n",
    "        \"y\": np.random.randint(ds.extent[2], ds.extent[3], nwells),\n",
    "    },\n",
    "    index=[f\"well{i}\" for i in range(nwells)],\n",
    ")\n",
    "cmap = plt.get_cmap(\"tab20\")\n",
    "df[\"filtercolor_face\"] = [\n",
    "    cmap(i)\n",
    "    for i in np.linspace(0, int((len(df) - 1) * cmap.N / len(df)), len(df), dtype=int)\n",
    "]\n",
    "df['tubecolor'] = 'white'\n",
    "gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.x, df.y))\n",
    "\n",
    "# plot cross section\n",
    "f, ax = plt.subplots(figsize=(10, 5))\n",
    "dcs = DatasetCrossSection(ds, line=line, zmin=-200, zmax=10, ax=ax)\n",
    "colors = nlmod.read.regis.get_legend()\n",
    "dcs.plot_layers(colors=colors)\n",
    "\n",
    "# plot wells in cross section\n",
    "dcs.plot_wells(\n",
    "    df,\n",
    "    legend=True,\n",
    "    max_dist=max_dist,\n",
    "    filter_patch_kwargs={'alpha':0.5, 'zorder':10},\n",
    "    filter_rect_kwargs={'alpha':0.5, 'zorder':10},\n",
    "    tubeline_kwargs={'alpha':0.5},\n",
    "    legend_kwds={\"loc\": (0, 1), \"frameon\": False, \"ncol\": 7},\n",
    ")\n",
    "\n",
    "# inset map\n",
    "mapax = nlmod.plot.inset_map(ax, ds.extent)\n",
    "\n",
    "# plot all wells in grey, wells within x meter of cross section in color\n",
    "gdf.plot(ax=mapax, color=\"grey\", edgecolor=\"grey\", markersize=20)\n",
    "gdf_plotted = gdf[gdf.geometry.distance(LineString(line)) <= max_dist]\n",
    "gdf_plotted.plot(\n",
    "    ax=mapax, color=gdf_plotted[\"filtercolor_face\"], edgecolor=\"k\", markersize=20\n",
    ")\n",
    "\n",
    "nlmod.plot.add_xsec_line_and_labels(line, ax, mapax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "336b85c9",
   "metadata": {},
   "source": [
    "## Animation\n",
    "\n",
    "There is also an option to make an animation of a cross section with variations in heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45382984",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make animation of the heads over time\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "\n",
    "f, ax = plt.subplots(figsize=(15, 5))\n",
    "dcs = DatasetCrossSection(ds, line=line, zmin=-3.0, zmax=0.0, ax=ax)\n",
    "\n",
    "dcs.plot_grid(vertical=False, linewidth=0.5)\n",
    "\n",
    "vmin = np.nanmin(ds[\"head\"])\n",
    "vmax = np.nanmax(ds[\"head\"])\n",
    "norm = mpl.colors.Normalize(vmin, vmax)\n",
    "\n",
    "mapax = nlmod.plot.inset_map(ax, ds.extent)\n",
    "nlmod.plot.add_xsec_line_and_labels(line, ax, mapax)\n",
    "\n",
    "plt.close(f)\n",
    "from IPython.display import HTML\n",
    "\n",
    "HTML(dcs.animate(ds[\"head\"], head=ds[\"head\"], norm=norm).to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f49969",
   "metadata": {},
   "source": [
    "## Time series\n",
    "For time series we use the functionality of xarray, as we have read the heads in a xarray DataArray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee45ab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 118228\n",
    "y = 439870\n",
    "head_point = nlmod.gwf.get_head_at_point(ds[\"head\"], x=x, y=y, ds=ds)\n",
    "head_point.plot.line(hue=\"layer\", size=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d56cc81c",
   "metadata": {},
   "source": [
    "We can also use pandas to plot the heads. First transform the data to a Pandas DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b129675",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = head_point.to_pandas()\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02112ab3",
   "metadata": {},
   "source": [
    "And then plot this DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd9d1f42",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.plot(figsize=(10, 10));"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}