examples/Notebooks/flopy3_mt3d-usgs_example_with_sft_lkt_uzt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstrate use of SFR/LAK/UZF with SFT/LKT/UZT\n",
    "A more comprehensive demonstration of setting up an MT3D-USGS model that uses all of the new packages included in the first release of MT3D-USGS.  Also includes RCT.\n",
    "\n",
    "#### Problem Description:\n",
    "* 300 row x 300 col x 3 layer x 2 stress period model\n",
    "* Flow model uses SFR, LAK, and UZF with connections between all three\n",
    "* Transport model simulates streamflow transport (SFT), with connction to a single lake (LKT)\n",
    "* Transport model simulates overland runoff and spring discharge (UZT) to surface water network\n",
    "\n",
    "Start by importing some libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.8.11 (default, Aug  6 2021, 08:56:27) \n",
      "[Clang 10.0.0 ]\n",
      "numpy version: 1.19.2\n",
      "matplotlib version: 3.4.2\n",
      "flopy version: 3.3.5\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import platform\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "# run installed version of flopy or add local path\n",
    "try:\n",
    "    import flopy\n",
    "except:\n",
    "    fpth = os.path.abspath(os.path.join(\"..\", \"..\"))\n",
    "    sys.path.append(fpth)\n",
    "    import flopy\n",
    "\n",
    "print(sys.version)\n",
    "print(\"numpy version: {}\".format(np.__version__))\n",
    "print(\"matplotlib version: {}\".format(mpl.__version__))\n",
    "print(\"flopy version: {}\".format(flopy.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a MODFLOW model and store it, in this case in the variable 'mf'. \n",
    "The modelname will be the name given to all MODFLOW files.\n",
    "The exe_name should be the name of the MODFLOW executable. \n",
    "In this case, we want to use version: 'mfnwt' for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelpth = os.path.join(\".\", \"temp\", \"no3\")\n",
    "modelname = \"no3\"\n",
    "mfexe = \"mfnwt\"\n",
    "mtexe = \"mt3dusgs\"\n",
    "\n",
    "if platform.system() == \"Windows\":\n",
    "    mfexe += \".exe\"\n",
    "    mtexe += \".exe\"\n",
    "\n",
    "# Make sure modelpth directory exists\n",
    "if not os.path.isdir(modelpth):\n",
    "    os.makedirs(modelpth, exist_ok=True)\n",
    "\n",
    "# Instantiate MODFLOW object in flopy\n",
    "mf = flopy.modflow.Modflow(\n",
    "    modelname=modelname, exe_name=mfexe, model_ws=modelpth, version=\"mfnwt\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up model discretization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "Lx = 90000.0\n",
    "Ly = 90000.0\n",
    "nrow = 300\n",
    "ncol = 300\n",
    "nlay = 3\n",
    "\n",
    "delr = Lx / ncol\n",
    "delc = Ly / nrow\n",
    "\n",
    "xmax = ncol * delr\n",
    "ymax = nrow * delc\n",
    "\n",
    "X, Y = np.meshgrid(\n",
    "    np.linspace(delr / 2, xmax - delr / 2, ncol),\n",
    "    np.linspace(ymax - delc / 2, 0 + delc / 2, nrow),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate output control (oc) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "oc = flopy.modflow.ModflowOc(mf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate solver package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Newton-Raphson Solver: Create a flopy nwt package object\n",
    "\n",
    "headtol = 1.0e-4\n",
    "fluxtol = 5\n",
    "maxiterout = 5000\n",
    "thickfact = 1e-06\n",
    "linmeth = 2\n",
    "iprnwt = 1\n",
    "ibotav = 1\n",
    "\n",
    "nwt = flopy.modflow.ModflowNwt(\n",
    "    mf,\n",
    "    headtol=headtol,\n",
    "    fluxtol=fluxtol,\n",
    "    maxiterout=maxiterout,\n",
    "    thickfact=thickfact,\n",
    "    linmeth=linmeth,\n",
    "    iprnwt=iprnwt,\n",
    "    ibotav=ibotav,\n",
    "    options=\"SIMPLE\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate discretization (DIS) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "elv_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"dis_arrays\", \"grnd_elv.txt\"\n",
    ")\n",
    "\n",
    "# Top of Layer 1 elevation determined using GW Vistas and stored locally\n",
    "grndElv = np.loadtxt(elv_pth)\n",
    "\n",
    "# Bottom of layer 1 elevation also determined from use of GUI and stored locally\n",
    "bt1_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"dis_arrays\", \"bot1.txt\"\n",
    ")\n",
    "bot1Elv = np.loadtxt(bt1_pth)\n",
    "\n",
    "bot2Elv = np.ones(bot1Elv.shape) * 100\n",
    "bot3Elv = np.zeros(bot2Elv.shape)\n",
    "\n",
    "botm = [bot1Elv, bot2Elv, bot3Elv]\n",
    "botm = np.array(botm)\n",
    "Steady = [False, False]\n",
    "nstp = [1, 1]\n",
    "tsmult = [1.0, 1.0]\n",
    "\n",
    "# Stress periods\n",
    "perlen = [9131.25, 9131.25]\n",
    "\n",
    "# Create the discretization object\n",
    "# itmuni = 4 (days); lenuni = 1 (feet)\n",
    "dis = flopy.modflow.ModflowDis(\n",
    "    mf,\n",
    "    nlay,\n",
    "    nrow,\n",
    "    ncol,\n",
    "    nper=2,\n",
    "    delr=delr,\n",
    "    delc=delc,\n",
    "    top=grndElv,\n",
    "    botm=botm,\n",
    "    laycbd=0,\n",
    "    itmuni=4,\n",
    "    lenuni=1,\n",
    "    steady=Steady,\n",
    "    nstp=nstp,\n",
    "    tsmult=tsmult,\n",
    "    perlen=perlen,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate upstream weighting (UPW) flow package for MODFLOW-NWT\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UPW must be instantiated after DIS.  Otherwise, during the mf.write_input() procedures,\n",
    "# flopy will crash.\n",
    "\n",
    "# First line of UPW input is: IUPWCB HDRY NPUPW IPHDRY\n",
    "hdry = -1.00e30\n",
    "iphdry = 0\n",
    "\n",
    "# Next variables are: LAYTYP, LAYAVG, CHANI, LAYVKA, LAYWET\n",
    "laytyp = [1, 3, 3]  # >0: convertible\n",
    "layavg = 0  #  0: harmonic mean\n",
    "chani = 1.0  # >0: CHANI is the horizontal anisotropy for the entire layer\n",
    "layvka = 0  # =0: indicates VKA is vertical hydraulic conductivity\n",
    "laywet = 0  # Always set equal to zero in UPW package\n",
    "hk = 20\n",
    "# hani = 1          # Not needed because CHANI > 1\n",
    "vka = 0.5  # Is equal to vert. K b/c LAYVKA = 0\n",
    "ss = 0.00001\n",
    "sy = 0.20\n",
    "\n",
    "upw = flopy.modflow.ModflowUpw(\n",
    "    mf,\n",
    "    laytyp=laytyp,\n",
    "    layavg=layavg,\n",
    "    chani=chani,\n",
    "    layvka=layvka,\n",
    "    laywet=laywet,\n",
    "    ipakcb=53,\n",
    "    hdry=hdry,\n",
    "    iphdry=iphdry,\n",
    "    hk=hk,\n",
    "    vka=vka,\n",
    "    ss=ss,\n",
    "    sy=sy,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate basic (BAS or BA6) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ibnd1_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"bas_arrays\", \"ibnd_lay1.txt\"\n",
    ")\n",
    "ibnd1 = np.loadtxt(ibnd1_pth)\n",
    "ibnd2 = np.ones(ibnd1.shape)\n",
    "ibnd3 = np.ones(ibnd2.shape)\n",
    "\n",
    "ibnd = [ibnd1, ibnd2, ibnd3]\n",
    "ibnd = np.array(ibnd)\n",
    "\n",
    "StHd1_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"bas_arrays\", \"strthd1.txt\"\n",
    ")\n",
    "StHd1 = np.loadtxt(StHd1_pth)\n",
    "\n",
    "StHd2_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"bas_arrays\", \"strthd2.txt\"\n",
    ")\n",
    "StHd2 = np.loadtxt(StHd2_pth)\n",
    "\n",
    "StHd3_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"bas_arrays\", \"strthd3.txt\"\n",
    ")\n",
    "StHd3 = np.loadtxt(StHd3_pth)\n",
    "\n",
    "strtElev = [StHd1, StHd2, StHd3]\n",
    "strtElev = np.array(strtElev)\n",
    "\n",
    "hdry = 999.0\n",
    "\n",
    "bas = flopy.modflow.ModflowBas(mf, ibound=ibnd, hnoflo=hdry, strt=strtElev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate general head boundary (GHB) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GHB boundaries are located along the top (north) and bottom (south)\n",
    "# edges of the domain, all 3 layers.\n",
    "\n",
    "elev_stpt_row1 = 308.82281\n",
    "elev_stpt_row300 = 239.13811\n",
    "elev_slp = (308.82281 - 298.83649) / (ncol - 1)\n",
    "\n",
    "sp = []\n",
    "for k in [0, 1, 2]:  # These indices need to be adjusted for 0-based moronicism\n",
    "    for i in [\n",
    "        0,\n",
    "        299,\n",
    "    ]:  # These indices need to be adjusted for 0-based silliness\n",
    "        for j in np.arange(\n",
    "            0, 300, 1\n",
    "        ):  # These indices need to be adjusted for 0-based foolishness\n",
    "            # Skipping cells not satisfying the conditions below\n",
    "            if (i == 1 and (j < 27 or j > 31)) or (\n",
    "                i == 299 and (j < 26 or j > 31)\n",
    "            ):\n",
    "                if i % 2 == 0:\n",
    "                    sp.append(\n",
    "                        [\n",
    "                            k,\n",
    "                            i,\n",
    "                            j,\n",
    "                            elev_stpt_row1 - (elev_slp * (j - 1)),\n",
    "                            11.3636,\n",
    "                        ]\n",
    "                    )\n",
    "                else:\n",
    "                    sp.append(\n",
    "                        [\n",
    "                            k,\n",
    "                            i,\n",
    "                            j,\n",
    "                            elev_stpt_row300 - (elev_slp * (j - 1)),\n",
    "                            11.3636,\n",
    "                        ]\n",
    "                    )\n",
    "\n",
    "\n",
    "for k in [0, 1, 2]:\n",
    "    for j in np.arange(26, 32, 1):\n",
    "        sp.append([k, 299, j, 238.20, 3409.0801])\n",
    "\n",
    "ghb = flopy.modflow.ModflowGhb(mf, stress_period_data=sp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate streamflow routing (SFR2) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read pre-prepared reach data into numpy recarrays using numpy.genfromtxt()\n",
    "# Remember that the cell indices stored in the pre-prepared NO3_ReachInput.csv file are based on 0-based indexing.\n",
    "# Flopy will convert to 1-based when it writes the files\n",
    "\n",
    "rpth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"sfr_data\", \"no3_reachinput.csv\"\n",
    ")\n",
    "reach_data = np.genfromtxt(rpth, delimiter=\",\", names=True)\n",
    "reach_data\n",
    "\n",
    "# Read pre-prepared segment data into numpy recarrays using numpy.genfromtxt()\n",
    "\n",
    "spth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"sfr_data\", \"no3_segmentdata.csv\"\n",
    ")\n",
    "ss_segment_data = np.genfromtxt(spth, delimiter=\",\", names=True)\n",
    "segment_data = {0: ss_segment_data, 1: ss_segment_data}\n",
    "segment_data[0][0:1][\"width1\"]\n",
    "\n",
    "nstrm = len(reach_data)\n",
    "nss = len(segment_data[0])\n",
    "nsfrpar = 0\n",
    "const = 128390.4  # constant for manning's equation, units of cfs\n",
    "dleak = 0.0001\n",
    "ipakcb = 53  # flag for writing SFR output to cell-by-cell budget (on unit 53)\n",
    "istcb2 = 37  # flag for writing SFR output to text file\n",
    "isfropt = 1\n",
    "dataset_5 = {\n",
    "    0: [nss, 0, 0],\n",
    "    1: [-1, 0, 0],\n",
    "}  # dataset 5 (see online guide) (ITMP, IRDFLG, IPTFLG)\n",
    "\n",
    "# Input arguments generally follow the variable names defined in the Online Guide to MODFLOW\n",
    "sfr = flopy.modflow.ModflowSfr2(\n",
    "    mf,\n",
    "    nstrm=nstrm,\n",
    "    nss=nss,\n",
    "    const=const,\n",
    "    dleak=dleak,\n",
    "    ipakcb=ipakcb,\n",
    "    istcb2=istcb2,\n",
    "    isfropt=isfropt,\n",
    "    reachinput=True,\n",
    "    reach_data=reach_data,\n",
    "    segment_data=segment_data,\n",
    "    dataset_5=dataset_5,\n",
    "    unit_number=15,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Lake (LAK) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read pre-prepared lake arrays\n",
    "LakArr_pth = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"lak_arrays\", \"lakarr1.txt\"\n",
    ")\n",
    "LakArr_lyr1 = np.loadtxt(LakArr_pth)\n",
    "LakArr_lyr2 = np.zeros(LakArr_lyr1.shape)\n",
    "LakArr_lyr3 = np.zeros(LakArr_lyr2.shape)\n",
    "\n",
    "LakArr = [LakArr_lyr1, LakArr_lyr2, LakArr_lyr3]\n",
    "LakArr = np.array(LakArr)\n",
    "\n",
    "nlakes = int(np.max(LakArr))\n",
    "ipakcb = ipakcb  # From above\n",
    "theta = -1.0  # Implicit\n",
    "nssitr = 10  # Maximum number of iterations for Newton’s method\n",
    "sscncr = 1.000e-03  # Convergence criterion for equilibrium lake stage solution\n",
    "surfdep = 2.000e00  # Height of small topological variations in lake-bottom\n",
    "stages = 268.00  # Initial stage of each lake at the beginning of the run\n",
    "\n",
    "# ITMP  > 0, read lake definition data\n",
    "# ITMP1 ≥ 0, read new recharge, evaporation, runoff, and withdrawal data for each lake\n",
    "# LWRT  > 0, suppresses printout from the lake package\n",
    "\n",
    "bdlknc_lyr1 = LakArr_lyr1.copy()\n",
    "bdlknc_lyr2 = LakArr_lyr1.copy()\n",
    "bdlknc_lyr3 = np.zeros((LakArr_lyr1.shape))\n",
    "\n",
    "# Need to expand bdlknc_lyr1 non-zero values by 1 in either direction\n",
    "# (left/right and up/down)\n",
    "for i in np.arange(0, LakArr_lyr1.shape[0]):\n",
    "    for j in np.arange(0, LakArr_lyr1.shape[1]):\n",
    "        im1 = i - 1\n",
    "        ip1 = i + 1\n",
    "        jm1 = j - 1\n",
    "        jp1 = j + 1\n",
    "\n",
    "        if im1 >= 0:\n",
    "            if LakArr_lyr1[i, j] == 1 and LakArr_lyr1[im1, j] == 0:\n",
    "                bdlknc_lyr1[im1, j] = 1\n",
    "\n",
    "        if ip1 < LakArr_lyr1.shape[0]:\n",
    "            if LakArr_lyr1[i, j] == 1 and LakArr_lyr1[ip1, j] == 0:\n",
    "                bdlknc_lyr1[ip1, j] = 1\n",
    "\n",
    "        if jm1 >= 0:\n",
    "            if LakArr_lyr1[i, j] == 1 and LakArr_lyr1[i, jm1] == 0:\n",
    "                bdlknc_lyr1[i, jm1] = 1\n",
    "\n",
    "        if jp1 < LakArr_lyr1.shape[1]:\n",
    "            if LakArr_lyr1[i, j] == 1 and LakArr_lyr1[i, jp1] == 0:\n",
    "                bdlknc_lyr1[i, jp1] = 1\n",
    "\n",
    "\n",
    "bdlknc = [bdlknc_lyr1, bdlknc_lyr2, bdlknc_lyr3]\n",
    "bdlknc = np.array(bdlknc)\n",
    "\n",
    "flux_data = {0: [[0.0073, 0.0073, 0.0, 0.0]], 1: [[0.0073, 0.0073, 0.0, 0.0]]}\n",
    "\n",
    "lak = flopy.modflow.ModflowLak(\n",
    "    mf,\n",
    "    nlakes=nlakes,\n",
    "    ipakcb=ipakcb,\n",
    "    theta=theta,\n",
    "    nssitr=nssitr,\n",
    "    sscncr=sscncr,\n",
    "    surfdep=surfdep,\n",
    "    stages=stages,\n",
    "    lakarr=LakArr,\n",
    "    bdlknc=bdlknc,\n",
    "    flux_data=flux_data,\n",
    "    unit_number=16,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate gage package for use with MODFLOW-NWT package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "gages = [\n",
    "    [1, 225, 90, 3],\n",
    "    [2, 68, 91, 3],\n",
    "    [3, 33, 92, 3],\n",
    "    [4, 165, 93, 3],\n",
    "    [5, 123, 94, 3],\n",
    "    [6, 77, 95, 3],\n",
    "    [7, 173, 96, 3],\n",
    "    [8, 328, 97, 3],\n",
    "    [9, 115, 98, 3],\n",
    "    [-1, -101, 1],\n",
    "]\n",
    "\n",
    "# gages = [[1,38,61,1],[2,67,62,1], [3,176,63,1], [4,152,64,1], [5,186,65,1], [6,31,66,1]]\n",
    "files = [\n",
    "    \"no3.gag\",\n",
    "    \"seg1_gag.out\",\n",
    "    \"seg2_gag.out\",\n",
    "    \"seg3_gag.out\",\n",
    "    \"seg4_gag.out\",\n",
    "    \"seg5_gag.out\",\n",
    "    \"seg6_gag.out\",\n",
    "    \"seg7_gag.out\",\n",
    "    \"seg8_gag.out\",\n",
    "    \"seg9_gag.out\",\n",
    "    \"lak1_gag.out\",\n",
    "]\n",
    "\n",
    "numgage = len(gages)\n",
    "gage = flopy.modflow.ModflowGage(\n",
    "    mf, numgage=numgage, gage_data=gages, filenames=files\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Unsaturated-Zone Flow (UZF) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "nuztop = 2\n",
    "iuzfopt = 2\n",
    "irunflg = 1\n",
    "ietflg = 0\n",
    "iuzfcb = 52\n",
    "iuzfcb2 = 0\n",
    "ntrail2 = 20\n",
    "nsets2 = 20\n",
    "nuzgag = 2\n",
    "surfdep = 2.0\n",
    "\n",
    "eps = 3.0\n",
    "thts = 0.30\n",
    "thti = 0.13079\n",
    "\n",
    "fname_uzbnd = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"uzf_arrays\", \"iuzbnd.txt\"\n",
    ")\n",
    "fname_runbnd = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"uzf_arrays\", \"irunbnd.txt\"\n",
    ")\n",
    "\n",
    "iuzfbnd = np.loadtxt(fname_uzbnd)\n",
    "irunbnd = np.loadtxt(fname_runbnd)\n",
    "\n",
    "uzgag = [[106, 160, 121, 3], [1, 1, -122, 1]]\n",
    "\n",
    "finf = {0: 1.8250e-03, 1: 1.8250e-03}\n",
    "\n",
    "uzf = flopy.modflow.ModflowUzf1(\n",
    "    mf,\n",
    "    nuztop=nuztop,\n",
    "    iuzfopt=iuzfopt,\n",
    "    irunflg=irunflg,\n",
    "    ietflg=ietflg,\n",
    "    ipakcb=iuzfcb,\n",
    "    iuzfcb2=iuzfcb2,\n",
    "    ntrail2=ntrail2,\n",
    "    nsets=nsets2,\n",
    "    surfdep=surfdep,\n",
    "    uzgag=uzgag,\n",
    "    iuzfbnd=1,\n",
    "    irunbnd=0,\n",
    "    vks=1.0e-6,\n",
    "    eps=3.5,\n",
    "    thts=0.35,\n",
    "    thtr=0.15,\n",
    "    thti=0.20,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Drain (DRN) package for MODFLOW-NWT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_drnElv = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"drn_arrays\", \"elv.txt\"\n",
    ")\n",
    "fname_drnCond = os.path.join(\n",
    "    \"..\", \"data\", \"mt3d_example_sft_lkt_uzt\", \"drn_arrays\", \"cond.txt\"\n",
    ")\n",
    "\n",
    "drnElv = np.loadtxt(fname_drnElv)\n",
    "drnCond = np.loadtxt(fname_drnCond)\n",
    "\n",
    "drnElv_lst = pd.DataFrame(\n",
    "    {\n",
    "        \"lay\": 1,\n",
    "        \"row\": np.nonzero(drnElv)[0] + 1,\n",
    "        \"col\": np.nonzero(drnElv)[1] + 1,\n",
    "        \"elv\": drnElv[np.nonzero(drnElv)],\n",
    "        \"cond\": drnCond[np.nonzero(drnCond)],\n",
    "    },\n",
    "    columns=[\"lay\", \"row\", \"col\", \"elv\", \"cond\"],\n",
    ")\n",
    "\n",
    "# Convert the DataFrame into a list of lists for the drn constructor\n",
    "stress_period_data = drnElv_lst.values.tolist()\n",
    "\n",
    "# Create a dictionary, 1 entry for each of the two stress periods.\n",
    "stress_period_data = {0: stress_period_data, 1: stress_period_data}\n",
    "\n",
    "drn = flopy.modflow.ModflowDrn(\n",
    "    mf, ipakcb=ipakcb, stress_period_data=stress_period_data\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate linkage with mass transport routing (LMT) package for MODFLOW-NWT (generates linker file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmt = flopy.modflow.ModflowLmt(\n",
    "    mf,\n",
    "    output_file_name=\"NO3.ftl\",\n",
    "    output_file_header=\"extended\",\n",
    "    output_file_format=\"formatted\",\n",
    "    package_flows=[\"all\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now work on MT3D-USGS file creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start by setting up MT3D-USGS class and pass in MODFLOW-NWT object for setting up a number of the BTN arrays\n",
    "\n",
    "mt = flopy.mt3d.Mt3dms(\n",
    "    modflowmodel=mf,\n",
    "    modelname=modelname,\n",
    "    model_ws=modelpth,\n",
    "    version=\"mt3d-usgs\",\n",
    "    namefile_ext=\"mtnam\",\n",
    "    exe_name=mtexe,\n",
    "    ftlfilename=\"no3.ftl\",\n",
    "    ftlfree=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate basic transport (BTN) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncomp = 1\n",
    "lunit = \"FT\"\n",
    "sconc = 0.0\n",
    "prsity = 0.3\n",
    "cinact = -1.0\n",
    "thkmin = 0.000001\n",
    "nprs = -2\n",
    "nprobs = 10\n",
    "nprmas = 10\n",
    "dt0 = 0.1\n",
    "nstp = 1\n",
    "mxstrn = 500\n",
    "ttsmult = 1.2\n",
    "ttsmax = 100\n",
    "\n",
    "# These observations need to be entered with 0-based indexing\n",
    "obs = [[0, 104, 158], [1, 104, 158], [2, 104, 158]]\n",
    "\n",
    "btn = flopy.mt3d.Mt3dBtn(\n",
    "    mt,\n",
    "    MFStyleArr=True,\n",
    "    DRYCell=True,\n",
    "    lunit=lunit,\n",
    "    sconc=sconc,\n",
    "    ncomp=ncomp,\n",
    "    prsity=prsity,\n",
    "    cinact=cinact,\n",
    "    obs=obs,\n",
    "    thkmin=thkmin,\n",
    "    nprs=nprs,\n",
    "    nprobs=nprobs,\n",
    "    chkmas=True,\n",
    "    nprmas=nprmas,\n",
    "    dt0=dt0,\n",
    "    nstp=nstp,\n",
    "    mxstrn=mxstrn,\n",
    "    ttsmult=ttsmult,\n",
    "    ttsmax=ttsmax,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate advection (ADV) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixelm = 0\n",
    "percel = 1.0000\n",
    "mxpart = 5000\n",
    "nadvfd = 1  # (1 = Upstream weighting)\n",
    "\n",
    "adv = flopy.mt3d.Mt3dAdv(\n",
    "    mt, mixelm=mixelm, percel=percel, mxpart=mxpart, nadvfd=nadvfd\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate generalized conjugate gradient solver (GCG) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "mxiter = 1\n",
    "iter1 = 50\n",
    "isolve = 3\n",
    "ncrs = 0\n",
    "accl = 1.000000\n",
    "cclose = 1.00e-06\n",
    "iprgcg = 5\n",
    "\n",
    "gcg = flopy.mt3d.Mt3dGcg(\n",
    "    mt,\n",
    "    mxiter=mxiter,\n",
    "    iter1=iter1,\n",
    "    isolve=isolve,\n",
    "    ncrs=ncrs,\n",
    "    accl=accl,\n",
    "    cclose=cclose,\n",
    "    iprgcg=iprgcg,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate dispersion (DSP) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "al = 0.1  # longitudinal dispersivity\n",
    "trpt = 0.1  # ratio of the horizontal transverse dispersivity to 'AL'\n",
    "trpv = 0.1  # ratio of the vertical transverse dispersitvity to 'AL'\n",
    "dmcoef = 1.0000e-10\n",
    "\n",
    "dsp = flopy.mt3d.Mt3dDsp(\n",
    "    mt, al=al, trpt=trpt, trpv=trpv, dmcoef=dmcoef, multiDiff=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate source-sink mixing (SSM) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# no user-specified concentrations associated with boundary conditions\n",
    "\n",
    "mxss = 11199\n",
    "\n",
    "ssm = flopy.mt3d.Mt3dSsm(mt, mxss=mxss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate reaction (RCT) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "isothm = 0\n",
    "ireact = 1\n",
    "irctop = 2\n",
    "igetsc = 0\n",
    "ireaction = 0\n",
    "\n",
    "rc1 = 6.3258e-04  # first-order reaction rate for the dissolved phase\n",
    "rc2 = 0.0  # Decay on Soil Layer\n",
    "\n",
    "rct = flopy.mt3d.Mt3dRct(\n",
    "    mt, isothm=isothm, ireact=ireact, igetsc=igetsc, rc1=rc1, rc2=rc2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate streamflow transport (SFT) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsfinit = len(reach_data)\n",
    "mxsfbc = len(reach_data)\n",
    "icbcsf = 0\n",
    "ioutobs = 92\n",
    "isfsolv = 1\n",
    "wimp = 0.5\n",
    "wups = 1.0\n",
    "cclosesf = 1.0e-6\n",
    "mxitersf = 10\n",
    "crntsf = 1.0\n",
    "iprtxmd = 0\n",
    "coldsf = 0\n",
    "dispsf = 0\n",
    "obs_sf = [225, 293, 326, 491, 614, 691, 864, 1192, 1307]\n",
    "sf_stress_period_data = {0: [0, 0, 0], 1: [0, 0, 0], 2: [0, 0, 0]}\n",
    "\n",
    "gage_output = [None, None, \"no3.sftobs\"]\n",
    "\n",
    "sft = flopy.mt3d.Mt3dSft(\n",
    "    mt,\n",
    "    nsfinit=nsfinit,\n",
    "    mxsfbc=mxsfbc,\n",
    "    icbcsf=icbcsf,\n",
    "    ioutobs=ioutobs,\n",
    "    isfsolv=isfsolv,\n",
    "    wimp=wimp,\n",
    "    wups=wups,\n",
    "    cclosesf=cclosesf,\n",
    "    mxitersf=mxitersf,\n",
    "    crntsf=crntsf,\n",
    "    iprtxmd=iprtxmd,\n",
    "    coldsf=coldsf,\n",
    "    dispsf=dispsf,\n",
    "    nobssf=len(obs_sf),\n",
    "    obs_sf=obs_sf,\n",
    "    sf_stress_period_data=sf_stress_period_data,\n",
    "    filenames=gage_output,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate unsaturated-zone transport (UZT) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "mxuzcon = np.count_nonzero(irunbnd)\n",
    "icbcuz = 45\n",
    "iet = 0\n",
    "wc = np.ones((nlay, nrow, ncol)) * 0.29\n",
    "sdh = np.ones((nlay, nrow, ncol))\n",
    "\n",
    "uzt = flopy.mt3d.Mt3dUzt(\n",
    "    mt,\n",
    "    mxuzcon=mxuzcon,\n",
    "    icbcuz=icbcuz,\n",
    "    iet=iet,\n",
    "    iuzfbnd=iuzfbnd,\n",
    "    sdh=sdh,\n",
    "    cuzinf=1.4158e-03,\n",
    "    filenames=\"no3\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate lake transport (LKT) package for MT3D-USGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlkinit = 1\n",
    "mxlkbc = 720\n",
    "icbclk = 81\n",
    "ietlak = 1\n",
    "coldlak = 1\n",
    "\n",
    "lkt_flux_data = {0: [[0, 1, 0.01667]], 1: [[0, 1, 0.02667]]}\n",
    "\n",
    "lkt = flopy.mt3d.Mt3dLkt(\n",
    "    mt,\n",
    "    nlkinit=nlkinit,\n",
    "    mxlkbc=mxlkbc,\n",
    "    icbclk=icbclk,\n",
    "    ietlak=ietlak,\n",
    "    coldlak=coldlak,\n",
    "    lk_stress_period_data=lkt_flux_data,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Write the MT3D-USGS input files for inspecting and running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/jdhughes/Documents/Development/flopy_git/flopy_fork/examples/Notebooks\n"
     ]
    }
   ],
   "source": [
    "pth = os.getcwd()\n",
    "print(pth)\n",
    "\n",
    "mf.write_input()\n",
    "mt.write_input()\n",
    "\n",
    "# mf.run_model()\n",
    "# mt.run_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}