{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Circle\n\nThis example illustrates how to setup a very simple unstructured groundwater\nmodel using the ``imod`` package and associated packages.\n\nIn overview, we'll set the following steps:\n\n    * Create a triangular mesh for a disk geometry.\n    * Create the xugrid UgridDataArrays containg the MODFLOW6 parameters.\n    * Feed these arrays into the imod mf6 classes.\n    * Write to modflow6 files.\n    * Run the model.\n    * Open the results back into UgridDataArrays.\n    * Visualize the results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll start with the following imports:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nimport xarray as xr\nimport xugrid as xu\n\nimport imod"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create a mesh\n\nThe first steps consists of generating a mesh. In this example, we'll use data\nincluded with iMOD Python for a circular mesh. Note that this is a [Ugrid2D\nobject.](https://deltares.github.io/xugrid/api/xugrid.Ugrid2d.html)\nFor more information on working with unstructured grids see the\n[Xugrid documentation](https://deltares.github.io/xugrid/index.html)\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "grid = imod.data.circle()\n\ngrid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can plot this object as follows:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots()\nxu.plot.line(grid, ax=ax)\nax.set_aspect(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create UgridDataArray\n\nNow that we have defined the grid, we can start defining the model parameter\ndata.\n\nOur goal here is to define a steady-state model with:\n\n* Uniform conductivities of 1.0 m/d;\n* Two layers of 5.0 m thick;\n* Uniform recharge of 0.001 m/d on the top layer;\n* Constant heads of 1.0 m along the exterior edges of the mesh.\n\nFrom these boundary conditions, we would expect circular mounding of the\ngroundwater; with small flows in the center and larger flows as the recharge\naccumulates while the groundwater flows towards the exterior boundary.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nface = grid.n_face\nnlayer = 2\n\nidomain = xu.UgridDataArray(\n    xr.DataArray(\n        np.ones((nlayer, nface), dtype=np.int32),\n        coords={\"layer\": [1, 2]},\n        dims=[\"layer\", grid.face_dimension],\n    ),\n    grid=grid,\n)\nicelltype = xu.full_like(idomain, 0)\nk = xu.full_like(idomain, 1.0, dtype=float)\nk33 = k.copy()\nrch_rate = xu.full_like(idomain.sel(layer=1), 0.001, dtype=float)\nbottom = idomain * xr.DataArray([5.0, 0.0], dims=[\"layer\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "All the data above have been constants over the grid. For the constant head\nboundary, we'd like to only set values on the external border. We can\n`py:method:xugrid.UgridDataset.binary_dilation` to easily find these cells:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "chd_location = xu.zeros_like(idomain.sel(layer=2), dtype=bool).ugrid.binary_dilation(\n    border_value=True\n)\nconstant_head = xu.full_like(idomain.sel(layer=2), 1.0, dtype=float).where(chd_location)\n\nfig, ax = plt.subplots()\nconstant_head.ugrid.plot(ax=ax)\nxu.plot.line(grid, ax=ax, color=\"black\")\nax.set_aspect(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Write the model\n\nThe first step is to define an empty model, the parameters and boundary\nconditions are added in the form of the familiar MODFLOW packages.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model = imod.mf6.GroundwaterFlowModel()\ngwf_model[\"disv\"] = imod.mf6.VerticesDiscretization(\n    top=10.0, bottom=bottom, idomain=idomain\n)\ngwf_model[\"chd\"] = imod.mf6.ConstantHead(\n    constant_head, print_input=True, print_flows=True, save_flows=True\n)\ngwf_model[\"ic\"] = imod.mf6.InitialConditions(head=0.0)\ngwf_model[\"npf\"] = imod.mf6.NodePropertyFlow(\n    icelltype=icelltype,\n    k=k,\n    k33=k33,\n    save_flows=True,\n)\ngwf_model[\"sto\"] = imod.mf6.SpecificStorage(\n    specific_storage=1.0e-5,\n    specific_yield=0.15,\n    transient=False,\n    convertible=0,\n)\ngwf_model[\"oc\"] = imod.mf6.OutputControl(save_head=\"all\", save_budget=\"all\")\ngwf_model[\"rch\"] = imod.mf6.Recharge(rch_rate)\n\nsimulation = imod.mf6.Modflow6Simulation(\"circle\")\nsimulation[\"GWF_1\"] = gwf_model\nsimulation[\"solver\"] = imod.mf6.Solution(\n    modelnames=[\"GWF_1\"],\n    print_option=\"summary\",\n    csv_output=False,\n    no_ptc=True,\n    outer_dvclose=1.0e-4,\n    outer_maximum=500,\n    under_relaxation=None,\n    inner_dvclose=1.0e-4,\n    inner_rclose=0.001,\n    inner_maximum=100,\n    linear_acceleration=\"cg\",\n    scaling_method=None,\n    reordering_method=None,\n    relaxation_factor=0.97,\n)\nsimulation.create_time_discretization(additional_times=[\"2000-01-01\", \"2000-01-02\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll create a new directory in which we will write and run the model.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "modeldir = imod.util.temporary_directory()\nsimulation.write(modeldir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Run the model\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The following lines assume the ``mf6`` executable is available on your PATH.\n  `The Modflow 6 examples introduction <mf6-introduction>` shortly\n  describes how to add it to yours.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "simulation.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Open the results\n\nFirst, we'll open the heads (.hds) file.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "head = imod.mf6.out.open_hds(\n    modeldir / \"GWF_1/GWF_1.hds\",\n    modeldir / \"GWF_1/disv.disv.grb\",\n)\nhead"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For a DISV MODFLOW6 model, the heads are returned as a UgridDataArray.  While\nall layers are timesteps are available, they are only loaded into memory as\nneeded.\n\nWe may also open the cell-by-cell flows (.cbc) file.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "cbc = imod.mf6.open_cbc(\n    modeldir / \"GWF_1/GWF_1.cbc\",\n    modeldir / \"GWF_1/disv.disv.grb\",\n)\nprint(cbc.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The flows are returned as a dictionary of UgridDataArrays. This dictionary\ncontains all entries that are stored in the CBC file, but like for the heads\nfile the data are only loaded into memory when needed.\n\nThe horizontal flows are stored on the edges of the UgridDataArray topology.\nThe other flows are generally stored on the faces; this includes the\nflow-lower-face.\n\nWe'll create a dataset for the horizontal flows for further analysis.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds = xu.UgridDataset(grids=grid)\nds[\"u\"] = cbc[\"flow-horizontal-face-x\"]\nds[\"v\"] = cbc[\"flow-horizontal-face-y\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize the results\n\nWe can quickly and easily visualize the output with the plotting functions\nprovided by xarray and xugrid. We'll add some some edge coordinates to the\ndataset so that they can be used to place the arrows in the quiver plot.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds = ds.ugrid.assign_edge_coords()\nfig, ax = plt.subplots()\nhead.isel(time=0, layer=0).compute().ugrid.plot(ax=ax)\nds.isel(time=0, layer=0).plot.quiver(\n    x=\"mesh2d_edge_x\", y=\"mesh2d_edge_y\", u=\"u\", v=\"v\", color=\"white\"\n)\nax.set_aspect(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As would be expected from our model input, we observe circular groundwater\nmounding and increasing flows as we move from the center to the exterior.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}