examples/notebooks/Head_Monitor_Example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ee31d799",
   "metadata": {},
   "source": [
    "# MODFLOW-API Head monitor example\n",
    "\n",
    "In this example the modflow-api is used in a more complex callback function to create a Head Monitor that updates at the timestep level. This example reverses `CHD` boundary conditions each stress period on a simple 10 x 10 model and displays the head results for each timestep.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4526a124",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output, display  # remove this import if adapted to python script\n",
    "\n",
    "from modflowapi import run_simulation, Callbacks\n",
    "from flopy.discretization import StructuredGrid\n",
    "from flopy.plot import PlotMapView, styles\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f7ac8a2",
   "metadata": {},
   "source": [
    "### Create a class that includes a callback function\n",
    "\n",
    "This class handles changing the `CHD` boundary condition as well as updating the matplotlib plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd38b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "class StructuredHeadMonitor:\n",
    "    \"\"\"\n",
    "    An example class that reverses the model gradient by\n",
    "    swapping CHD boundary conditions each stress period, \n",
    "    and monitors the head at each timestep by updating\n",
    "    a matplotlib plot. This class could be adapted to \n",
    "    be used as a head monitor to observe other changes\n",
    "    in the model by modifying the callback class.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    layer: int\n",
    "        zero based model layer to plot\n",
    "    vmin : float\n",
    "        minimum head value for color scaling on the plot\n",
    "    vmax : float\n",
    "        maximum head value for color scaling on the plot\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, layer, vmin, vmax):\n",
    "        self.modelgrid = None\n",
    "        self.ax = None\n",
    "        self.pmv = None\n",
    "        self.pc = None\n",
    "        self.ax = None\n",
    "        self.layer = layer\n",
    "        self.vmin = vmin\n",
    "        self.vmax = vmax\n",
    "        self.kperold = None\n",
    "\n",
    "    def build_modelgrid(self, ml):\n",
    "        \"\"\"\n",
    "        Method to update the matplotlib plot\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        ml : ApiModel\n",
    "            modflow-api ApiModel object\n",
    "        \"\"\"\n",
    "        delc = ml.dis.get_advanced_var(\"delc\")\n",
    "        delr = ml.dis.get_advanced_var(\"delr\")\n",
    "        top = ml.dis.top.values[0]\n",
    "        botm = ml.dis.bot.values\n",
    "        idomain = ml.dis.idomain.values\n",
    "        self.modelgrid = StructuredGrid(\n",
    "            delc=delc,\n",
    "            delr=delr,\n",
    "            top=top,\n",
    "            botm=botm,\n",
    "            idomain=idomain\n",
    "        )\n",
    "\n",
    "    def initialize_plot(self):\n",
    "        \"\"\"\n",
    "        Method to initalize a matplotlib plot using flopy\n",
    "        \"\"\"\n",
    "        fig, ax = plt.subplots(figsize=(8, 8))\n",
    "        self.fig = fig\n",
    "        self.ax = ax\n",
    "        self.pmv = PlotMapView(modelgrid=self.modelgrid, ax=ax, layer=self.layer)\n",
    "        grid = self.pmv.plot_grid()\n",
    "        idm = self.pmv.plot_inactive()\n",
    "        initial = np.full(self.modelgrid.shape, np.nan)\n",
    "        self.pc = self.pmv.plot_array(initial, vmin=self.vmin, vmax=self.vmax)\n",
    "        plt.colorbar(self.pc)\n",
    "\n",
    "    def update_plot(self, ml):\n",
    "        \"\"\"\n",
    "        Method to update the matplotlib plot\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        ml : ApiModel\n",
    "            modflow-api ApiModel object\n",
    "        \"\"\"\n",
    "        heads = ml.X\n",
    "        self.ax.cla()\n",
    "        grid = self.pmv.plot_grid()\n",
    "        idm = self.pmv.plot_inactive()\n",
    "        self.pc = self.pmv.plot_array(heads, vmin=self.vmin, vmax=self.vmax)\n",
    "        \n",
    "        # only applicable to jupyter notebooks, remove these two lines in python scipt\n",
    "        display(self.fig)  \n",
    "        if ml.kper == (ml.nper - 1) and ml.kstp == (ml.nstp - 1):\n",
    "            pass\n",
    "        else:\n",
    "            clear_output(wait = True)  \n",
    "        \n",
    "        # the pause time can be reduced if adapted in python script \n",
    "        plt.pause(0.1)  \n",
    "\n",
    "    def callback(self, sim, callback_step):\n",
    "        \"\"\"\n",
    "        A demonstration function that dynamically adjusts the CHD\n",
    "        boundary conditions each stress period in a modflow-6 model\n",
    "        through the MODFLOW-API and then updates heads on a matplotlib\n",
    "        plot for each timestep.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        sim : modflowapi.Simulation\n",
    "            A simulation object for the solution group that is \n",
    "            currently being solved\n",
    "        callback_step : enumeration\n",
    "            modflowapi.Callbacks enumeration object that indicates\n",
    "            the part of the solution modflow is currently in.\n",
    "        \"\"\"\n",
    "        if callback_step == Callbacks.initialize:\n",
    "            ml = sim.get_model()\n",
    "            self.build_modelgrid(ml)\n",
    "            self.initialize_plot()\n",
    "\n",
    "        if callback_step == Callbacks.timestep_start:\n",
    "            ml = sim.get_model()\n",
    "            if ml.kper == 0:\n",
    "                self.kperold = ml.kper\n",
    "                head = ml.chd.stress_period_data.dataframe[\"head\"].values\n",
    "                self.head = head\n",
    "            else:\n",
    "                df = ml.chd.stress_period_data.dataframe\n",
    "                if self.kperold != ml.kper:\n",
    "                    self.kperold = ml.kper\n",
    "                    self.head = self.head[::-1]\n",
    "\n",
    "                df[\"head\"] = self.head\n",
    "                ml.chd.stress_period_data.dataframe = df\n",
    "\n",
    "        if callback_step == Callbacks.timestep_end:\n",
    "            ml = sim.get_model()\n",
    "            self.update_plot(ml)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5cf1d9",
   "metadata": {},
   "source": [
    "Run the model using the and supply the `StructuredHeadMonitor`'s `callback` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2902aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "hdmon = StructuredHeadMonitor(layer=0, vmin=70, vmax=95)\n",
    "dll = \"libmf6\"\n",
    "sim_ws = Path(\"../data/dis_model\")\n",
    "run_simulation(dll, sim_ws, hdmon.callback, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba6ce60",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}