{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 03: Loading and visualizing groundwater models\n",
    "\n",
    "This exercise, we will load an existing model into Flopy, run the model and then use [pandas](https://pandas.pydata.org/), [matplotlib](https://matplotlib.org/) and [numpy](https://www.numpy.org/) to look at the results and compare them to observed data. We will also export model input and output to shapefiles and rasters.\n",
    "\n",
    "#### Required executables\n",
    "* MODFLOW-6; available here: https://github.com/MODFLOW-USGS/executables\n",
    "\n",
    "#### Operations\n",
    "* reading tabular data from a file or url using the powerful [pandas.read_csv](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html) method\n",
    "* getting `pandas.DataFrame`s of Hydmod, SFR, and global mass balance output\n",
    "* converting model times to real date-times to allow plotting against other temporally-referenced data\n",
    "* quickly subsetting data by category, attribute values, times, index position, etc.\n",
    "* computing quantiles and other basic statistics\n",
    "* making plots using `matplotlib` and the built-in hooks to it in `pandas`\n",
    "\n",
    "#### The Pleasant Lake example\n",
    "The example model is a simplified version of the MODFLOW-6 model published by Fienen et al (2022, 2021; Figure 1), who used a multi-scale modeling approach to evaluate the effects of agricultural groundwater abstraction on the ecology of 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; 200 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",
    "* Recharge 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 too 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": [
    "m = 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": [
    "m.modelgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.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 ``epsg`` attribute (in this case 3070 for Wisconsin Transverse Mercator)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.modelgrid.crs = 3070"
   ]
  },
  {
   "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(m, ax=ax)\n",
    "lc = pmv.plot_grid()\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'{m.modelgrid.units.capitalize()} easting, {m.modelgrid.crs.name}')\n",
    "ax.set_ylabel(f'{m.modelgrid.units.capitalize()} northing, {m.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(m, ax=ax)\n",
    "lc = pmv.plot_grid()\n",
    "top = pmv.plot_array(m.dis.top.array)\n",
    "ax.set_xlabel(f'{m.modelgrid.units.capitalize()} easting, {m.modelgrid.crs.name}')\n",
    "ax.set_ylabel(f'{m.modelgrid.units.capitalize()} northing, {m.modelgrid.crs.name}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting the model grid to a shapefile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.modelgrid.write_shapefile(str(output_folder / 'pleasant_grid.shp'))"
   ]
  },
  {
   "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=(10, 3))\n",
    "xs = flopy.plot.PlotCrossSection(model=m, 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 {m.modelgrid.units.capitalize()}')\n",
    "ax.set_ylabel(f'Elevation, in {m.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=(10, 3))\n",
    "xs_line = [(552400, 393000), (552400 + 5000, 393000 - 4000)]\n",
    "xs = flopy.plot.PlotCrossSection(model=m, \n",
    "                                 line={\"line\": xs_line}, ax=ax,\n",
    "                                 geographic_coords=True)\n",
    "lc = xs.plot_grid(zorder=4)\n",
    "\n",
    "pc = xs.plot_array(m.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 = m.modelgrid.shape[1] # set frames to number of rows\n",
    "\n",
    "def update(i):\n",
    "    ax.cla()\n",
    "    xs = flopy.plot.PlotCrossSection(model=m, 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 {m.modelgrid.units.capitalize()}')\n",
    "    ax.set_ylabel(f'Elevation, in {m.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",
    "m.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 = m.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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.rcha.recharge.array[:, 0, :, :].sum(axis=(1, 2)) * 100**2"
   ]
  },
  {
   "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 m.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 get the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.run_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the output\n",
    "With MODFLOW 6 models, we can get the output from the model object, without having to reference additional files. Sometimes though, it may be easier to read the file directly.\n",
    "\n",
    "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",
    "from mfexport.utils import get_water_table\n",
    "\n",
    "hds = m.output.head().get_data(kstpkper=(0, 0))\n",
    "wt = get_water_table(hds, nodata=-1e30)\n",
    "\n",
    "# add the lake stage to the water table\n",
    "lak_output = pd.read_csv(sim_ws / 'lake1.obs.csv')\n",
    "stage = lak_output['STAGE'][0]\n",
    "cnd = pd.DataFrame(m.lak.connectiondata.array)\n",
    "k, i, j = zip(*cnd['cellid'])\n",
    "wt[i, j] = stage\n",
    "# add the SFR stage as well\n",
    "sfr_stage = m.sfr.output.stage().get_data()[0, 0, :]\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(*m.sfr.packagedata.array['cellid'])\n",
    "wt[sfr_i, sfr_j] = sfr_stage\n",
    "\n",
    "cbc = m.output.budget()\n",
    "lak = cbc.get_data(text='lak', full3D=True)[0].sum(axis=0)\n",
    "sfr = cbc.get_data(text='sfr', full3D=True)[0]"
   ]
  },
  {
   "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(m, ax=ax)\n",
    "ctr = pmv.contour_array(wt, 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(lak, cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "im = pmv.plot_array(sfr.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",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 6))\n",
    "pmv = flopy.plot.PlotMapView(m, ax=ax, extent=(554500, 557500, 388500, 392000))\n",
    "ctr = pmv.contour_array(wt, 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(lak, cmap='coolwarm', vmin=vmin, vmax=vmax)\n",
    "im = pmv.plot_array(sfr.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(m.modelgrid, str(output_folder / 'water_table.tif'), wt)"
   ]
  },
  {
   "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 = wt - m.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(m.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=m, 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=m, \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(hds, head=hds, alpha=0.5, masked_values=[1e30])\n",
    "ctr = xs.contour_array(hds, head=hds, levels=levels, colors=\"b\", masked_values=[1e30])\n",
    "surf = xs.plot_surface(wt, masked_values=[1e30], color=\"blue\", lw=2)\n",
    "\n",
    "labels = pmv.ax.clabel(\n",
    "    ctr, inline=True, \n",
    "    fontsize=8, inline_spacing=5)"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "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": [
    "m.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": [
    "m.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(m.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 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": [
    "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}