{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 03: Loading and visualizing groundwater models\n",
    "\n",
    "This exercise, we will load an existing model using Flopy, look at some of the inputs, run the model, and load and plot some of the results.\n",
    "\n",
    "#### The Pleasant Lake example\n",
    "The example model is a simplified version of the MODFLOW-6 model published by Fienen and others (2022, 2021), of the area around Pleasant Lake in central Wisconsin, USA. The original report and model files are available at the links below.\n",
    "\n",
    "##### Example model details:\n",
    "\n",
    "* Transient MODFLOW-6 simulation with monthly stress periods for calendar year 2012\n",
    "* units of meters and days\n",
    "* 4 layers; 100 meter uniform grid spacing\n",
    "    * layers 1-3 represent surficial deposits\n",
    "    * layer 4 represents Paleozoic bedrock (sandstone)\n",
    "* Transient specified head perimeter boundary (CHD package) from a regional model solution\n",
    "* Transient recharge from a soil water balance model, specified with RCHA\n",
    "* Streams specified with SFR\n",
    "* Pleasant Lake simulated with the Lake Package\n",
    "* Head observations specified with the OBS utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import flopy\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make an output folder\n",
    "output_folder = Path('03-output')\n",
    "output_folder.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load a preexisting MODFLOW 6 model\n",
    "Because this is MODFLOW 6, we need to load the simulation first, and then get the model.\n",
    "\n",
    "**Note:** To avoid loading certain packages (that may be slow) use the ``load_only`` argument to specify the packages that should be loaded.   \n",
    "e.g. ``load_only=['dis']``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "sim_ws = Path('../data/pleasant-lake/')\n",
    "sim = flopy.mf6.MFSimulation.load('pleasant', sim_ws=str(sim_ws), exe_name='mf6',\n",
    "                                  #load_only=['dis']\n",
    "                         )\n",
    "sim.model_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = sim.get_model('pleasant')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the model\n",
    "\n",
    "First let's check that the model grid is correctly located. It is, in this case, because the model has the origin and rotation specified in the DIS package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.modelgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_package_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, in order to write shapefiles with a ``.prj`` file that specifies the coordinate references system (CRS), we need to assign one to the grid (there currently is no CRS input for MODFLOW 6). We can do this by simply specifying an [EPSG code](https://epsg.io/) to the ``crs`` attribute (in this case 3070 for Wisconsin Transverse Mercator)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.modelgrid.crs = 3070\n",
    "model.modelgrid.units"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also assign the modelgrid to a variable so we don't have to type `model.modelgrid` every time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid = model.modelgrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### On a map\n",
    "We can plot the model in the CRS using the ``PlotMapView`` object. More examples in the Flopy demo here (for unstructured grids too!): https://github.com/modflowpy/flopy/blob/develop/examples/Notebooks/flopy3.3_PlotMapView.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "pmv = flopy.plot.PlotMapView(model, ax=ax)\n",
    "lc = pmv.plot_grid(lw=0.5)\n",
    "pmv.plot_bc(\"WEL\", plotAll=True)\n",
    "pmv.plot_bc(\"LAK\", plotAll=True)\n",
    "pmv.plot_bc(\"SFR\", plotAll=True)\n",
    "pmv.plot_bc(\"CHD\", plotAll=True)\n",
    "ax.set_xlabel(f'{modelgrid.units.capitalize()} easting, {modelgrid.crs.name}')\n",
    "ax.set_ylabel(f'{modelgrid.units.capitalize()} northing, {modelgrid.crs.name}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "pmv = flopy.plot.PlotMapView(model, ax=ax)\n",
    "lc = pmv.plot_grid(lw=0.5)\n",
    "top = pmv.plot_array(model.dis.top.array)\n",
    "ax.set_xlabel(f'{modelgrid.units.capitalize()} easting, {modelgrid.crs.name}')\n",
    "ax.set_ylabel(f'{modelgrid.units.capitalize()} northing, {modelgrid.crs.name}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting the model grid to a shapefile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid.write_shapefile(str(output_folder / 'pleasant_grid.shp'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the model grid as a GeoPandas `GeoDataFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid_gdf = modelgrid.geo_dataframe\n",
    "modelgrid_gdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add row, column information to model grid `GeoDataFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelgrid_gdf['i'] = modelgrid_gdf.index // modelgrid.ncol\n",
    "modelgrid_gdf['j'] = modelgrid_gdf.index % modelgrid.ncol\n",
    "modelgrid_gdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making a cross section through the model\n",
    "\n",
    "more examples in the Flopy demo here: https://github.com/modflowpy/flopy/blob/develop/examples/Notebooks/flopy3.3_PlotCrossSection.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### By row or column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 5))\n",
    "xs = flopy.plot.PlotCrossSection(model=model, line={\"row\": 30}, ax=ax)\n",
    "lc = xs.plot_grid()\n",
    "xs.plot_bc(\"LAK\")\n",
    "xs.plot_bc(\"SFR\")\n",
    "ax.set_xlabel(f'Distance, in {modelgrid.units.capitalize()}')\n",
    "ax.set_ylabel(f'Elevation, in {modelgrid.units.capitalize()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Along an arbitrary line\n",
    "(and in Geographic Coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 5))\n",
    "xs_line = [(552400, 393000), (552400 + 5000, 393000 - 4000)]\n",
    "xs = flopy.plot.PlotCrossSection(model=model, \n",
    "                                 line={\"line\": xs_line}, ax=ax,\n",
    "                                 geographic_coords=True)\n",
    "lc = xs.plot_grid(zorder=4)\n",
    "\n",
    "pc = xs.plot_array(model.npf.k.array)\n",
    "fig.colorbar(pc, label='Hydraulic Conductivity, in m/day');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### What if we want to look at cross sections for each row or column?\n",
    "This code allows for every row or column to be visualized in cross section within the Jupyter Notebook session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 5))\n",
    "frames = modelgrid.shape[1] # set frames to number of rows\n",
    "\n",
    "def update(i):\n",
    "    ax.cla()\n",
    "    xs = flopy.plot.PlotCrossSection(model=model, line={\"row\": i}, ax=ax)\n",
    "    lc = xs.plot_grid()\n",
    "    xs.plot_bc(\"LAK\")\n",
    "    xs.plot_bc(\"SFR\")\n",
    "    ax.set_title(f\"row: {i}\")\n",
    "    ax.set_xlabel(f'Distance, in {modelgrid.units.capitalize()}')\n",
    "    ax.set_ylabel(f'Elevation, in {modelgrid.units.capitalize()}')\n",
    "    return\n",
    "\n",
    "import matplotlib.animation as animation\n",
    "ani = animation.FuncAnimation(fig=fig, func=update, frames=frames)\n",
    "plt.close()\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An aside on working with model input as numpy arrays\n",
    "Every input to MODFOW is attached to a Flopy object (with the attribute name of the variable) as a numpy ``ndarray`` (for ndarray-type data) or a ``recarray`` for tabular or list-type data. For example, we can access the recharge array (4D-- nper x nlay x nrow x ncol) with:\n",
    "\n",
    "```\n",
    "model.rcha.recharge.array\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ``ndarray`` example: plot spatial average recharge by stress period\n",
    "To minimize extra typing, it often makes sense to reassign the numpy array to a new variable to work with it further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rch_inches = model.rcha.recharge.array[:, 0, :, :].mean(axis=(1, 2)) * 12 * 30.4 / .3048 \n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(rch_inches)\n",
    "ax.axhline(rch_inches.mean(), c='C1')\n",
    "ax.set_ylabel(f\"Average recharge, in monthly inches\")\n",
    "ax.set_xlabel(\"Model stress period\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tabular data example: plot pumping by stress period\n",
    "Most tabular input for the 'basic' stress packages (Constant Head, Drain, General Head, RIV, WEL, etc) are accessible via a ``stress_period_data`` attribute.   \n",
    "* To access the data, we have to call another ``.data`` attribute, which gives us a dictionary of ``recarray``s by stress period.  \n",
    "* Any one of these can be converted to a ``pandas.DataFrame`` individually, or we can make a dataframe of all of them with a simple loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = []\n",
    "for kper, df in model.wel.stress_period_data.get_dataframe().items():\n",
    "    df['per'] = kper\n",
    "    dfs.append(df)\n",
    "df = pd.concat(dfs)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can sum by stress period, or plot individual wells across stress periods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.groupby('per').sum()['q'].plot()\n",
    "plt.title('Total pumpage by Stress period')\n",
    "plt.ylabel('$m^3$/day')\n",
    "plt.xlabel('Model stress period')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: plot pumping for the well at cellid: 2, 24, 2\n",
    "About many gallons did this well pump in stress periods 1 through 12?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = df.groupby('boundname').get_group('pleasant_2-13-2').plot(x='per')\n",
    "ax.set_ylabel('$m^3$/day')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution:\n",
    "\n",
    "1) get the pumping rates by stress period"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rates = df.groupby('boundname').get_group('pleasant_2-13-2')[1:]\n",
    "rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2) Google \"python get days in month\" or similar (simply using 30.4 would be fine too for an approximate answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import calendar\n",
    "rates['days'] = [calendar.monthrange(2012,m)[1] for m in rates['per']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3) multiply the days by the daily pumping rate to get the totals; convert units and sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rates['gallons'] = rates['q'] * rates['days'] * 264.172\n",
    "print(f\"{rates['gallons'].sum():,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing model output\n",
    "\n",
    "### Run the model first to generate the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two ways to read MODFLOW output\n",
    "\n",
    "#### Getting MODFLOW output from the model object\n",
    "\n",
    "With MODFLOW 6 models, we can get the output from the model object, without having to reference the output files.\n",
    "\n",
    "In Flopy, MODFLOW binary output is read using objects that represent the various file types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "headfile_object = model.output.head()\n",
    "headfile_object"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output file objects include various methods and attributes with information about the output. For example, the elapsed simulation times with output recrods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "headfile_object.get_times()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Including a `get_data()` method to get the output values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "heads = headfile_object.get_data(kstpkper=(0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sfr_output_object = model.sfr.output\n",
    "sfr_output_object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sfr_stage = sfr_output_object.stage().get_data()[0, 0, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cell Budget output works the same way.\n",
    "\n",
    "In this case, we specified `save_flows` as an option in the various packages to save flux output for those packages, include fluxes to and from the Lake and SFR Packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cbc_file_object = model.output.budget()\n",
    "cbc_file_object.get_unique_record_names()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cell Budget results are returned in a list that allows for auxilliary information besides the output to be included in subsequent items.\n",
    "\n",
    "The values are the first item."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lake_fluxes = cbc_file_object.get_data(text='lak', full3D=True)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Especially for structured (layer, row, column) grids, we usually want to specify `full3D=True`. This returns the results in a 3D Numpy array. \n",
    "\n",
    "Otherwise, results are returned in a Numpy `recarray`, referenced by unique node numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cbc_file_object.get_data(text='lak')[0][:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cbc_file_object = model.output.budget()\n",
    "lake_fluxes = cbc_file_object.get_data(text='lak', full3D=True)[0].sum(axis=0)\n",
    "sfr_fluxes = cbc_file_object.get_data(text='sfr', full3D=True)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Alternatively, read MODFLOW binary output from the files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flopy.utils import binaryfile as bf\n",
    "\n",
    "headfile_object = bf.HeadFile(sim_ws / f\"{model.name}.hds\")\n",
    "heads_from_the_file = headfile_object.get_data(kstpkper=(0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cbc_file_object = bf.CellBudgetFile(sim_ws / f\"{model.name}.cbc\")\n",
    "cbc_fluxes_from_the_file = cbc_file_object.get_data(kstpkper=(0, 0), full3D=True)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MODFLOW 6 CSV observation output can also be easily read using `pandas`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lak_output = pd.read_csv(sim_ws / 'lake1.obs.csv')\n",
    "lak_output.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "head_obs_output = pd.read_csv(sim_ws / f\"{model.name}.head.obs\")\n",
    "head_obs_output.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The head solution is reported for each layer. \n",
    "* We can use the ``get_water_table`` utility to get a 2D surface of the water table position in each cell. \n",
    "* To accurately portray the water table around the lake, we can read the lake stage from the observation file and assign it to the relevant cells in the water table array. \n",
    "* Otherwise, depending on how the lake is constructed, the lake area would be shown as a nodata/no-flow area, or as the heads in the groundwater system below the lakebed.\n",
    "* In this case, we are getting the solution from the initial steady-state period"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flopy.utils.postprocessing import get_water_table\n",
    "\n",
    "water_table = get_water_table(heads)\n",
    "\n",
    "# add the lake stage to the water table\n",
    "stage = lak_output['STAGE'][0]\n",
    "cnd = pd.DataFrame(model.lak.connectiondata.array)\n",
    "k, i, j = zip(*cnd['cellid'])\n",
    "water_table[i, j] = stage\n",
    "# add the SFR stage as well\n",
    "# get the SFR cell i, j locations\n",
    "# by unpacking the cellid tuples in the packagedata\n",
    "sfr_k, sfr_i, sfr_j = zip(*model.sfr.packagedata.array['cellid'])\n",
    "water_table[sfr_i, sfr_j] = sfr_stage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot head and surface water flux results\n",
    "We can add output to a PlotMapView instance as arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels=np.arange(280, 315, 2)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "pmv = flopy.plot.PlotMapView(model, ax=ax)\n",
    "ctr = pmv.contour_array(water_table, levels=levels, \n",
    "                        linewidths=1, colors='b')\n",
    "labels = pmv.ax.clabel(ctr, inline=True, fontsize=8, inline_spacing=1)\n",
    "vmin, vmax = -100, 100\n",
    "im = pmv.plot_array(lake_fluxes, cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "im = pmv.plot_array(sfr_fluxes.sum(axis=0), cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "cb = fig.colorbar(im, shrink=0.7, label='Leakage, in m$^3$/day')\n",
    "ax.set_ylabel(\"Northing, WTM meters\")\n",
    "ax.set_xlabel(\"Easting, WTM meters\")\n",
    "ax.set_aspect(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Zoom in on the lake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels=np.arange(280, 315, 1)\n",
    "new_extent = (554500, 557500, 388500, 392000)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "pmv = flopy.plot.PlotMapView(model, ax=ax, extent=new_extent)\n",
    "ctr = pmv.contour_array(water_table, levels=levels, \n",
    "                        linewidths=1, colors='b')\n",
    "labels = pmv.ax.clabel(ctr, inline=True, fontsize=8, inline_spacing=1)\n",
    "vmin, vmax = -100, 100\n",
    "im = pmv.plot_array(lake_fluxes, cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "im = pmv.plot_array(sfr_fluxes.sum(axis=0), cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "cb = fig.colorbar(im, shrink=0.7, label='Leakage, in m$^3$/day')\n",
    "ax.set_ylabel(\"Northing, WTM meters\")\n",
    "ax.set_xlabel(\"Easting, WTM meters\")\n",
    "ax.set_aspect(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting Rasters\n",
    "We can use the ``export_array`` utility to make a GeoTIFF of any 2D array on a structured grid. For example, make a raster of the simulated water table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flopy.export.utils import export_array\n",
    "\n",
    "export_array(model.modelgrid, str(output_folder / 'water_table.tif'), water_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise: evaluating overpressurization\n",
    "A common issue with groundwater flow models is overpressurization- where heads above the land surface are simulated. Sometimes, these indicate natural wetlands that aren't explicitly simulated in the model, but other times they are a sign of unrealistic parameters. Use the information in this lesson to answer the following questions:\n",
    "\n",
    "1) Does this model solution have any overpressiuzation? If so, where? Is it appropriate?\n",
    "\n",
    "2) What is the maximum value of overpressurization?\n",
    "\n",
    "3) What is the maximum depth to water simulated? Where are the greatest depths to water? Do they look appropriate?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution\n",
    "\n",
    "1) Make a numpy array of overpressurization and get the max and min values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "op = water_table - model.dis.top.array\n",
    "op.max(), op.min()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2) Plot it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(op); plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3) Export a raster of overpressurization so we can compare it against air photos (or mapped wetlands if we have them!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_array(model.modelgrid, str(output_folder / 'op.tif'), op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The highest levels of overpressurization correspond to Pleasant Lake, where the model top represents the lake bottom. Other areas of *OP* appear to correspond to lakes or wetlands, especially the spring complex south of Pleasant Lake, where Tagatz Creek originates.\n",
    "\n",
    "The greatest depths to water correspond to a topographic high in the southwest part of the model domain. A cross section through the area confirms that it is a bedrock high that rises more than 50 meters above the surrounding topography, so a depth to water of 76 meters in this area seems reasonable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 3))\n",
    "xs = flopy.plot.PlotCrossSection(model=model, line={\"row\": 62}, ax=ax)\n",
    "lc = xs.plot_grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot a cross section of the head solution with the water table\n",
    "We can also view output in cross section. In this case, ``PlotMatView`` plots the head solution where the model is saturated. We can add the water table we created above that includes the lake surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "xs_line = [(552400, 393000), (552400 + 5000, 393000 - 4000)]\n",
    "xs = flopy.plot.PlotCrossSection(model=model, \n",
    "                                 line={\"line\": xs_line}, \n",
    "                                 #line={\"row\": 32}, \n",
    "                                 ax=ax,\n",
    "                                 geographic_coords=True)\n",
    "lc = xs.plot_grid()\n",
    "pc = xs.plot_array(heads, head=heads, alpha=0.5, masked_values=[1e30])\n",
    "ctr = xs.contour_array(heads, head=heads, levels=levels, colors=\"b\", masked_values=[1e30])\n",
    "surf = xs.plot_surface(water_table, masked_values=[1e30], color=\"blue\", lw=2)\n",
    "\n",
    "labels = pmv.ax.clabel(\n",
    "    ctr, inline=True, \n",
    "    fontsize=8, inline_spacing=5)\n",
    "\n",
    "ax.set_xlabel(f'Distance, in {modelgrid.units.capitalize()}')\n",
    "ax.set_ylabel(f'Elevation, in {modelgrid.units.capitalize()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting observation output two ways\n",
    "In this model, head \"observations\" were specified at various locations using the MODFLOW 6 Observation Utility. MODFLOW reports simulated values of head at these locations, which can then be compared with equivalent field observations for model parameter estimation.\n",
    "\n",
    "Earlier we obtained a DataFrame of Lake Package observation output for pleasant lake by reading in 'lake1.obs.csv' with pandas.\n",
    "We can read the head observation output with pandas too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "headobs = pd.read_csv(sim_ws / 'pleasant.head.obs')\n",
    "headobs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Head observations can also be accessed via the ``.output`` attribute for their respective package. First we have to find the name associated with that package though. We can get this by calling ``get_package_list()``. Looks like ``\"OBS_3\"`` is it (since the only `OBS` packages in the model is for heads)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_package_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next let's query the available output methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.obs_3.output.methods()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now get the output. It comes back as a Numpy recarray be default, but we can easily cast it to a DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(model.obs_3.output.obs().get_data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using boundnames to define observations\n",
    "In MODFLOW 6, we can use boundnames to create observations for groups of cells. For example, in this model, each head value specified in the Constant Head (CHD) Package has a ``boundname`` of east, west, north or south, to indicate the side of the model perimeter it's on.\n",
    "\n",
    "Example of boundnames specified in an external input file for the CHD Package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.read_csv(sim_ws / 'external/chd_001.dat', delim_whitespace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Constant Head Observation Utility input is then set up like so:\n",
    "```\n",
    "BEGIN options\n",
    "END options\n",
    "\n",
    "BEGIN continuous  FILEOUT  pleasant.chd.obs.output.csv\n",
    "# obsname obstype ID\n",
    "  east  chd  east\n",
    "  west  chd  west\n",
    "  north  chd  north\n",
    "  south  chd  south\n",
    "END continuous  FILEOUT  pleasant.chd.obs.output.csv\n",
    "```\n",
    "\n",
    "The resulting observation output (net flow across the boundary faces in model units of cubic meters per day) can be found in ``pleasant.chd.obs.output.csv``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(sim_ws / 'pleasant.chd.obs.output.csv')\n",
    "df.index = df['time']\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting global mass balance from the listing file\n",
    "The ``Mf6ListBudget`` and ``MfListBudget`` (for earlier MODFLOW versions) utilities can assemble the global mass balance output (printed in the Listing file) into a DataFrame. A ``start_datetime`` can be added to convert the MODFLOW time to actual dates.\n",
    "\n",
    "**Note:** The ``start_datetime`` functionality is unaware of steady-state periods, so if we put in the actual model start date of 2012-01-01, the 1-day initial steady-state will be included, resulting in the stress periods being offset by one day. Also note that the dates here represent the *end* of each stress period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flopy.utils import Mf6ListBudget "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mfl = Mf6ListBudget(sim_ws / 'pleasant.list')\n",
    "flux, vol = mfl.get_dataframes(start_datetime='2011-12-30')\n",
    "flux.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flux.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the model mass balance error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = flux['PERCENT_DISCREPANCY'].plot()\n",
    "ax.set_ylabel('Percent mass balance error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make a stacked bar plot of the global mass balance\n",
    "\n",
    "Note: This works best if the in and out columns are aligned, such that ``STO-SY_IN`` and ``STO-SY_OUT`` are both colored orange, for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "in_cols = ['STO-SS_IN', 'STO-SY_IN', 'WEL_IN', 'RCHA_IN', 'CHD_IN', 'SFR_IN', 'LAK_IN']\n",
    "out_cols = [c.replace('_IN', '_OUT') for c in in_cols]\n",
    "flux[in_cols].plot.bar(stacked=True, ax=ax)\n",
    "(-flux[out_cols]).plot.bar(stacked=True, ax=ax)\n",
    "ax.legend(loc='lower left', bbox_to_anchor=(1, 0))\n",
    "ax.axhline(0, lw=0.5, c='k')\n",
    "ax.set_ylabel('Simulated Flux, in $m^3/d$')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "Fienen, M. N., Haserodt, M. J., Leaf, A. T., and Westenbroek, S. M. (2022). Simulation of regional groundwater flow and groundwater/lake interactions in the central Sands, Wisconsin. U.S. Geological Survey Scientific Investigations Report 2022-5046. doi:10.3133/sir20225046\n",
    "\n",
    "Fienen, M. N., Haserodt, M. J., and Leaf, A. T. (2021). MODFLOW models used to simulate groundwater flow in the Wisconsin Central Sands Study Area, 2012-2018. New York: U.S. Geological Survey Data Release. doi:10.5066/P9BVFSGJ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyclass",
   "language": "python",
   "name": "python3"
  },
  "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}