{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "afad5282",
   "metadata": {},
   "source": [
    "# 09: Demonstration of MODFLOW 6 Groundwater Transport with a Voronoi Grid\n",
    "\n",
    "MODFLOW 6 includes a Groundwater Transport (GWT) Model for simulation of solute transport through the subsurface.  The GWT Model can be used with structured or unstructured model grids.  The purpose of this example is to demonstrate the construction, running, and post-processing of a simple solute transport model.\n",
    "\n",
    "The solute transport model is based on an existing flow model of the Freyberg example.  The flow model usesd a voronoi model grid to simulate steady-state conditions.  In this notebook, we create a transient solute transport model using the same voronoi grid that was used for flow.\n",
    "\n",
    "The following steps are used in this notebook.\n",
    "* Load the existing flow model into FloPy\n",
    "* Plot the model grid\n",
    "* Create the solute transport model\n",
    "* Run the solute transport model with MODFLOW 6\n",
    "* Create animations of the model results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88945275-63d6-4c07-8189-45e2911bc5cc",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc15635f-e887-417c-a9b6-583f1d0c758e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as colors\n",
    "\n",
    "import flopy\n",
    "from flopy.utils.triangle import Triangle as Triangle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01909c79-e3d5-45c9-b058-ffb1d8216f51",
   "metadata": {},
   "source": [
    "### Load and Plot the Existing Flow Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55b0c73-ec82-4654-ac95-d840845c6a80",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ws_load = \"./data/voronoi/\"\n",
    "model_ws = \"./temp/voronoi-gwt/\"\n",
    "name = \"voronoi\"\n",
    "name_load = \"project\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "945eddc2-b144-4305-b51e-d659feb6cf50",
   "metadata": {},
   "source": [
    "Load a few shapefiles with geopandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5e72a3-edfd-4595-ba6d-420f210a5d84",
   "metadata": {},
   "outputs": [],
   "source": [
    "river_shp = \"data_project/Flowline_river.shp\"\n",
    "wells_shp = \"data_project/pumping_well_locations.shp\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30df1ad3-a236-4e1a-9dc2-6093c2d1c372",
   "metadata": {
    "tags": []
   },
   "source": [
    "Load the existing voronoi groundwater flow model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f556d4-fc16-4114-a9d0-5c14181d8ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "sim = flopy.mf6.MFSimulation.load(sim_ws=model_ws_load, sim_name=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3828269-85ab-4c14-879e-b6c28450a465",
   "metadata": {},
   "source": [
    "Get the gwf model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48b8a191-3230-4e41-b0df-6e00d43a5b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwf = sim.get_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24a69ac7-9fd0-4dbb-940e-138f7a7493e5",
   "metadata": {},
   "source": [
    "Assign a constant concentration condition to the x, y location of 550, 7900.  Use the `modelgrid.intersect` method to determine the cell number for the constant concentration condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd922664-30cf-4be1-941f-e0d8ce186c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwf.modelgrid.intersect(550, 7900)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dedb537b-3a25-4183-b573-7f2f72c54c62",
   "metadata": {},
   "source": [
    "Plot the grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d5d5361-3361-4568-b711-3cfa482f2993",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwf.modelgrid.plot()\n",
    "ax = plt.gca()\n",
    "ax.plot(550, 7900, marker=\"o\", lw=0, color=\"red\", )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e664ff92-50bf-444b-b78d-86d85e2497e8",
   "metadata": {},
   "source": [
    "### Create the Groundwater Transport Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e64ee71d-342a-4e5f-aad7-adcadf6ce700",
   "metadata": {},
   "source": [
    "Get data from the GWF DISV package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de99b252-29f7-4da8-b0f4-207ed36a26fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlay, ncpl = gwf.disv.nlay.array, gwf.disv.ncpl.array\n",
    "nlay, ncpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "806af036-109a-40dd-bf47-c8e5c9e92dad",
   "metadata": {},
   "outputs": [],
   "source": [
    "top, botm = gwf.disv.top.array, gwf.disv.botm.array\n",
    "top.shape, botm.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8581e8ec-c521-4758-a04b-438c4ef6fd1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "nverts = gwf.disv.nvert.array\n",
    "nverts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab044fde-018d-43f8-a743-1637358d4b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertices, cell2d = gwf.disv.vertices.array, gwf.disv.cell2d.array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b8fff6-39db-4277-a6c5-839de61e3ff6",
   "metadata": {},
   "source": [
    "Create the GWT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f03b0549-d677-4f23-b01d-b28a164f953b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_gwt = flopy.mf6.MFSimulation(sim_name=name, sim_ws=model_ws)\n",
    "tdis = flopy.mf6.ModflowTdis(sim_gwt, \n",
    "                             time_units=\"days\",\n",
    "                             perioddata=((10000.0, 100, 1.0),),\n",
    "                            )\n",
    "ims = flopy.mf6.ModflowIms(\n",
    "    sim_gwt,\n",
    "    linear_acceleration=\"bicgstab\",\n",
    "    outer_maximum=200,\n",
    "    inner_maximum=100,\n",
    "    print_option=\"all\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845f8a67-c82b-48eb-b151-adc0f8e58b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwt = flopy.mf6.ModflowGwt(sim_gwt, modelname=name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c690cae9-5a17-45e8-b7ff-f9f99d7a1cdc",
   "metadata": {},
   "source": [
    "Create the GWT packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c5fc188-4f1a-4a43-a997-85598d15e9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dis = flopy.mf6.ModflowGwtdisv(\n",
    "    gwt,\n",
    "    length_units=\"feet\",\n",
    "    nlay=nlay,\n",
    "    ncpl=ncpl,\n",
    "    nvert=nverts,\n",
    "    top=top,\n",
    "    botm=botm,\n",
    "    vertices=vertices,\n",
    "    cell2d=cell2d,\n",
    ")\n",
    "ic = flopy.mf6.ModflowGwtic(gwt, strt=0.0)\n",
    "adv = flopy.mf6.ModflowGwtadv(\n",
    "    gwt, \n",
    "    scheme=\"tvd\",\n",
    ")\n",
    "dsp = flopy.mf6.ModflowGwtdsp(gwt, alh=50.0, ath1=5)\n",
    "mst = flopy.mf6.ModflowGwtmst(gwt, porosity=0.2)\n",
    "pd = [\n",
    "    (\"GWFHEAD\", f\"../../{model_ws_load}{name_load}.hds\", None),\n",
    "    (\"GWFBUDGET\", f\"../../{model_ws_load}{name_load}.cbc\", None),\n",
    "]\n",
    "fmi = flopy.mf6.ModflowGwtfmi(gwt, packagedata=pd)\n",
    "ssm = flopy.mf6.ModflowGwtssm(gwt)\n",
    "cnc = flopy.mf6.ModflowGwtcnc(gwt, stress_period_data=[(0, 2027, 100.)])\n",
    "\n",
    "oc = flopy.mf6.ModflowGwtoc(\n",
    "    gwt,\n",
    "    concentration_filerecord=f\"{name}.ucn\",\n",
    "    saverecord=[(\"CONCENTRATION\", \"ALL\"),],\n",
    "    printrecord=[(\"BUDGET\", \"ALL\")],\n",
    ")\n",
    "# write the model datasets\n",
    "sim_gwt.write_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc7c4f0c-1dc6-4a1b-b933-caa6ba5c360f",
   "metadata": {},
   "source": [
    "### Run the Solute Transport Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a98e64c-cbbb-4f74-9dc6-5f72b23bf13b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_gwt.run_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab370b7-c023-402a-bf3d-792cc82ad4de",
   "metadata": {},
   "source": [
    "### Post-process the results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b64389-c690-4c21-bf74-2d4a5479de23",
   "metadata": {},
   "source": [
    "Use `gwt.output.` method to get the concentrations. Make an animation the concentrations using `flopy.plot` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0cd47c4-f591-4acd-aeed-7ac399e7447c",
   "metadata": {},
   "outputs": [],
   "source": [
    "head = gwf.output.head().get_data()\n",
    "spdis = gwf.output.budget().get_data(text=\"DATA-SPDIS\")[0]\n",
    "qx, qy, qz = flopy.utils.postprocessing.get_specific_discharge(\n",
    "    spdis, gwf, head=gwf.output.head().get_data(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d872228-ac57-4cba-8873-fb6430dac670",
   "metadata": {},
   "outputs": [],
   "source": [
    "times = gwt.output.concentration().get_times()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36492a7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4, 6), constrained_layout=True)\n",
    "ax.set_aspect(1)\n",
    "ax.set_xlabel(r'x')\n",
    "ax.set_ylabel(r'y')\n",
    "title = ax.set_title(f\"Time = {times[0]} days\")\n",
    "\n",
    "# plot persistent items\n",
    "vmin, vmax = 1e-3, 100.\n",
    "norm = colors.LogNorm(vmin=vmin, vmax=vmax)\n",
    "\n",
    "pmv = flopy.plot.PlotMapView(gwt, ax=ax)\n",
    "pmv.plot_grid(lw=0.5, color=\"0.5\")\n",
    "pmv.contour_array(\n",
    "    head, \n",
    "    levels=np.linspace(0, 30, 30), \n",
    "    tri_mask=True,\n",
    "    linestyles=\"-\",\n",
    "    colors=\"blue\",\n",
    "    linewidths=0.5,\n",
    ")\n",
    "ca_dict = {\n",
    "    \"vmin\": vmin,\n",
    "    \"vmax\": vmax,\n",
    "    \"norm\": norm,\n",
    "    \"masked_values\": [0],\n",
    "}\n",
    "conc_alldata = gwt.output.concentration().get_alldata()\n",
    "c = conc_alldata[0]\n",
    "c[c < vmin] = 0.\n",
    "cont = pmv.plot_array(c, **ca_dict)\n",
    "clb = fig.colorbar(\n",
    "    cont, \n",
    "    shrink=0.5, \n",
    ")\n",
    "\n",
    "def animate(i):\n",
    "    c = conc_alldata[i].flatten()\n",
    "    c[c < vmin] = 0.\n",
    "    cont.set_array(c)\n",
    "    title = ax.set_title(f\"Time = {times[i]} days\")\n",
    "    return cont\n",
    "\n",
    "import matplotlib.animation\n",
    "ani = matplotlib.animation.FuncAnimation(fig, animate, frames=conc_alldata.shape[0])\n",
    "plt.close()\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(ani.to_jshtml())\n",
    "\n",
    "# can use this command to write animation to file\n",
    "#ani.save(\"voronoi-conc-animation.avi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7694f069-60e1-4dd1-a150-ecd79872dae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "line = [(0, 8100), (2000, 8300), (5000, 8100)]\n",
    "xs = flopy.plot.PlotCrossSection(model=gwt, line={\"line\": line}) \n",
    "cb = xs.plot_array(conc_alldata[-1],\n",
    "                   head=gwf.output.head().get_data(),\n",
    "                   norm=norm,\n",
    "                   masked_values=[0], vmin=vmin, vmax=vmax)\n",
    "xs.plot_grid(lw=0.5, color=\"0.5\")\n",
    "plt.colorbar(cb, ax=xs.ax)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb2c863a",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}