{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 05: Unstructred Grids\n",
    "\n",
    "This notebook demonstrates FloPy functionality for creating different types of model grids.  Six different model grids are constructed for the same domain.  The domain is defined by a polygon intended to represent a hydrologic basin.\n",
    "\n",
    "The grids generated in this notebook include the following:\n",
    "\n",
    "* Regular MODFLOW grid\n",
    "* Irregular MODFLOW grid with variable row and column spacings\n",
    "* Nested grid, consisting of parent and child grids\n",
    "* Quadtree grid created using the FloPy wrapper for the [USGS Gridgen program](https://www.usgs.gov/software/gridgen-program-generating-unstructured-finite-volume-grids)\n",
    "* Triangular grid created using the FloPy wrapper for the [Triangle program](https://www.cs.cmu.edu/~quake/triangle.html).\n",
    "* Voronoi grid created using the [Triangle program](https://www.cs.cmu.edu/~quake/triangle.html) and the [Scipy voronoi class](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.Voronoi.html#scipy.spatial.Voronoi)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib as pl\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from shapely.geometry import Polygon\n",
    "import flopy\n",
    "from flopy.discretization import StructuredGrid, VertexGrid\n",
    "from flopy.utils.triangle import Triangle\n",
    "from flopy.utils.voronoi import VoronoiGrid\n",
    "from flopy.utils.gridgen import Gridgen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import basin data and utilities from basin.py\n",
    "import basin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_path = pl.Path(\"./temp\")\n",
    "temp_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basin Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boundary_polygon = basin.string2geom(basin.boundary)\n",
    "print(\"len boundary\", len(boundary_polygon))\n",
    "bp = np.array(boundary_polygon)\n",
    "\n",
    "sgs = [\n",
    "    basin.string2geom(sg) for sg in (\n",
    "        basin.streamseg1, basin.streamseg2, basin.streamseg3, basin.streamseg4\n",
    "        )\n",
    "]\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "ax.set_aspect(\"equal\")\n",
    "\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"ks-\")\n",
    "colors = (\"blue\", \"cyan\", \"green\", \"limegreen\")\n",
    "for idx, sg in enumerate(sgs):\n",
    "    print(\"Len segment: \", len(sg))\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], marker=\"o\", color=colors[idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FloPy `StructuredGrid`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a regular MODFLOW grid\n",
    "Lx = 180000\n",
    "Ly = 100000\n",
    "dx = dy = 2000\n",
    "nrow = int(Ly / dy)\n",
    "ncol = int(Lx / dx)\n",
    "print(Lx, Ly, nrow, ncol)\n",
    "delr = np.array(ncol * [dx])\n",
    "delc = np.array(nrow * [dy])\n",
    "regular_grid = StructuredGrid(delr=delr, delc=delc, xoff=0.0, yoff=0.0)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "ax.set_aspect(\"equal\")\n",
    "regular_grid.plot(ax=ax)\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "for sg in sgs:\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_idomain(grid, boundary):\n",
    "    from flopy.utils.gridintersect import GridIntersect\n",
    "    from shapely.geometry import Polygon\n",
    "\n",
    "    ix = GridIntersect(grid, method=\"vertex\", rtree=True)\n",
    "    result = ix.intersect(Polygon(boundary))\n",
    "    idx = [coords for coords in result.cellids]\n",
    "    idx = np.array(idx, dtype=int)\n",
    "    nr = idx.shape[0]\n",
    "    if idx.ndim == 1:\n",
    "        idx = idx.reshape((nr, 1))\n",
    "    print(idx.shape, idx.ndim)\n",
    "    idx = tuple([idx[:, i] for i in range(idx.shape[1])])\n",
    "    # idx = (idx[:, 0], idx[:, 1])\n",
    "    idomain = np.zeros(grid.shape[1:], dtype=int)\n",
    "    idomain[idx] = 1\n",
    "    idomain = idomain.reshape(grid.shape)\n",
    "    grid.idomain = idomain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Regular MODFLOW Grid (DIS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a regular MODFLOW grid\n",
    "Lx = 180000\n",
    "Ly = 100000\n",
    "dx = dy = 2000\n",
    "nrow = int(Ly / dy)\n",
    "ncol = int(Lx / dx)\n",
    "print(Lx, Ly, nrow, ncol)\n",
    "delr = np.array(ncol * [dx])\n",
    "delc = np.array(nrow * [dy])\n",
    "regular_grid = StructuredGrid(nlay=1, delr=delr, delc=delc, xoff=0.0, yoff=0.0)\n",
    "\n",
    "set_idomain(regular_grid, boundary_polygon)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "pmv = flopy.plot.PlotMapView(modelgrid=regular_grid)\n",
    "ax.set_aspect(\"equal\")\n",
    "pmv.plot_grid()\n",
    "pmv.plot_inactive()\n",
    "# regular_grid.plot(ax=ax, )\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "for sg in sgs:\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Irregular Grid (DIS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an irregular MODFLOW grid\n",
    "Lx = 180000\n",
    "Ly = 100000\n",
    "dx = dy = 5000\n",
    "\n",
    "smooth = [5000 / 1.2**i for i in range(9)]\n",
    "smoothr = smooth.copy()\n",
    "smoothr.reverse()\n",
    "dx = 12 * [5000] + smooth + 12 * [1000] + smoothr + 12 * [5000]\n",
    "dy = 4 * [5000] + smooth + 12 * [1000] + smoothr + 4 * [5000]\n",
    "\n",
    "ncol = len(dx)\n",
    "nrow = len(dy)\n",
    "\n",
    "delr = np.array(dx)\n",
    "delc = np.array(dy)\n",
    "irregular_grid = StructuredGrid(\n",
    "    nlay=1, delr=delr, delc=delc, xoff=0.0, yoff=0.0\n",
    ")\n",
    "set_idomain(irregular_grid, boundary_polygon)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "pmv = flopy.plot.PlotMapView(modelgrid=irregular_grid)\n",
    "ax.set_aspect(\"equal\")\n",
    "pmv.plot_grid()\n",
    "pmv.plot_inactive()\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "for sg in sgs:\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nested grid - two regular grids (DIS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nested grid\n",
    "from flopy.utils.lgrutil import Lgr\n",
    "\n",
    "# define parent grid information\n",
    "nlayp = 1\n",
    "dx = 5000\n",
    "nrowp = int(Ly / dx)\n",
    "ncolp = int(Lx / dx)\n",
    "delrp = dx\n",
    "delcp = dx\n",
    "topp = 1.0\n",
    "botmp = [0.0]\n",
    "idomainp = np.ones((nlayp, nrowp, ncolp), dtype=int)\n",
    "idomainp[0, 8:12, 13:18] = 0\n",
    "\n",
    "# define child grid resolution parameters\n",
    "ncpp = 3\n",
    "ncppl = [1]\n",
    "\n",
    "lgr = Lgr(\n",
    "    nlayp,\n",
    "    nrowp,\n",
    "    ncolp,\n",
    "    delrp,\n",
    "    delcp,\n",
    "    topp,\n",
    "    botmp,\n",
    "    idomainp,\n",
    "    ncpp=ncpp,\n",
    "    ncppl=ncppl,\n",
    "    xllp=0.0,\n",
    "    yllp=0.0,\n",
    ")\n",
    "\n",
    "delr = np.array(ncolp * [dx])\n",
    "delc = np.array(nrowp * [dx])\n",
    "regular_gridp = StructuredGrid(nlay=1, delr=delr, delc=delc, idomain=idomainp)\n",
    "set_idomain(regular_gridp, boundary_polygon)\n",
    "\n",
    "delr, delc = lgr.get_delr_delc()\n",
    "xoff, yoff = lgr.get_lower_left()\n",
    "regular_gridc = StructuredGrid(\n",
    "    delr=delr, delc=delc, xoff=xoff, yoff=yoff, idomain=idomainp\n",
    ")\n",
    "\n",
    "nested_grid = [regular_gridp, regular_gridc]\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "pmv = flopy.plot.PlotMapView(modelgrid=regular_gridp)\n",
    "pmv.plot_inactive()\n",
    "ax.set_aspect(\"equal\")\n",
    "regular_gridc.plot(ax=ax)\n",
    "pmv.plot_grid()\n",
    "\n",
    "# regular_gridp.plot(ax=ax)\n",
    "\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "for sg in sgs:\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FloPy `VertexGrid`\n",
    "\n",
    "### Quadtree Grid (DISV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quadtree grid\n",
    "sim = flopy.mf6.MFSimulation()\n",
    "gwf = flopy.mf6.ModflowGwf(sim)\n",
    "dx = dy = 5000.0\n",
    "nr = int(Ly / dy)\n",
    "nc = int(Lx / dx)\n",
    "dis6 = flopy.mf6.ModflowGwfdis(\n",
    "    gwf,\n",
    "    nrow=nr,\n",
    "    ncol=nc,\n",
    "    delr=dy,\n",
    "    delc=dx,\n",
    ")\n",
    "\n",
    "# Create gridgen object, add refinement features, and build grid\n",
    "g = Gridgen(gwf.modelgrid, model_ws=temp_path)\n",
    "refine_line = sgs\n",
    "g.add_refinement_features(refine_line, \"line\", 2, range(1))\n",
    "g.build(verbose=False)\n",
    "\n",
    "gridprops_vg = g.get_gridprops_vertexgrid()\n",
    "quadtree_grid = flopy.discretization.VertexGrid(**gridprops_vg)\n",
    "set_idomain(quadtree_grid, boundary_polygon)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "pmv = flopy.plot.PlotMapView(modelgrid=quadtree_grid)\n",
    "pmv.plot_grid()\n",
    "pmv.plot_inactive()\n",
    "ax.set_aspect(\"equal\")\n",
    "\n",
    "ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "for sg in sgs:\n",
    "    sa = np.array(sg)\n",
    "    ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Triangular grid (DISV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set maximum cell area\n",
    "maximum_area = 5000 * 5000\n",
    "\n",
    "nodes = []\n",
    "for sg in sgs:\n",
    "    sg_densify = basin.densify_geometry(sg, 2000)\n",
    "    nodes += sg_densify\n",
    "nodes = np.array(nodes)\n",
    "\n",
    "# Use the flopy Triangle class to build a triangular mesh\n",
    "tri = Triangle(angle=30, maximum_area=maximum_area, nodes=nodes, model_ws=temp_path)\n",
    "poly = bp\n",
    "tri.add_polygon(poly)\n",
    "tri.build(verbose=False)\n",
    "\n",
    "# Create a flopy VertexGrid\n",
    "cell2d = tri.get_cell2d()\n",
    "vertices = tri.get_vertices()\n",
    "idomain = np.ones((1, tri.ncpl), dtype=int)\n",
    "triangular_grid = VertexGrid(vertices=vertices, cell2d=cell2d, idomain=idomain)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "ax.set_aspect(\"equal\")\n",
    "triangular_grid.plot(ax=ax)\n",
    "\n",
    "if False:\n",
    "    ax.plot(bp[:, 0], bp[:, 1], \"k-\")\n",
    "    for sg in sgs:\n",
    "        sa = np.array(sg)\n",
    "        ax.plot(sa[:, 0], sa[:, 1], \"b-\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Voronoi Grid (DISV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maximum_area = 5000 * 5000\n",
    "\n",
    "nodes = []\n",
    "for sg in sgs:\n",
    "    sg_densify = basin.densify_geometry(sg, 2000)\n",
    "    nodes += sg_densify\n",
    "nodes = np.array(nodes)\n",
    "\n",
    "# Use the flopy Triangle class to build a triangular mesh\n",
    "tri = Triangle(angle=30, maximum_area=maximum_area / 1, nodes=nodes, model_ws=temp_path)\n",
    "poly = bp\n",
    "tri.add_polygon(poly)\n",
    "tri.build(verbose=False)\n",
    "\n",
    "# Create the flopy Voronoi grid object and the flopy VertexGrid\n",
    "vor = VoronoiGrid(tri)\n",
    "gridprops = vor.get_gridprops_vertexgrid()\n",
    "idomain = np.ones((1, vor.ncpl), dtype=int)\n",
    "voronoi_grid = VertexGrid(**gridprops, nlay=1, idomain=idomain)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "ax.set_aspect(\"equal\")\n",
    "voronoi_grid.plot(ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}