flopy_henry.py

import os
import sys
from tempfile import TemporaryDirectory

import matplotlib.pyplot as plt
import numpy as np

import flopy


def run(workspace, quiet):
    fext = "png"
    narg = len(sys.argv)
    iarg = 0
    if narg > 1:
        while iarg < narg - 1:
            iarg += 1
            basearg = sys.argv[iarg].lower()
            if basearg == "--pdf":
                fext = "pdf"

    # Input variables for the Henry Problem
    Lx = 2.0
    Lz = 1.0
    nlay = 50
    nrow = 1
    ncol = 100
    delr = Lx / ncol
    delc = 1.0
    delv = Lz / nlay
    henry_top = 1.0
    henry_botm = np.linspace(henry_top - delv, 0.0, nlay)
    qinflow = 5.702  # m3/day
    dmcoef = 0.57024  # m2/day  Could also try 1.62925 as another case of the Henry problem
    hk = 864.0  # m/day

    # Create the basic MODFLOW model data
    modelname = "flopy_henry"
    m = flopy.seawat.Seawat(modelname, exe_name="swtv4", model_ws=workspace)

    # Add DIS package to the MODFLOW model
    dis = flopy.modflow.ModflowDis(
        m,
        nlay,
        nrow,
        ncol,
        nper=1,
        delr=delr,
        delc=delc,
        laycbd=0,
        top=henry_top,
        botm=henry_botm,
        perlen=1.5,
        nstp=15,
    )

    # Variables for the BAS package
    ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
    ibound[:, :, -1] = -1
    bas = flopy.modflow.ModflowBas(m, ibound, 0)

    # Add LPF package to the MODFLOW model
    lpf = flopy.modflow.ModflowLpf(m, hk=hk, vka=hk, ipakcb=53)

    # Add PCG Package to the MODFLOW model
    pcg = flopy.modflow.ModflowPcg(m, hclose=1.0e-8)

    # Add OC package to the MODFLOW model
    oc = flopy.modflow.ModflowOc(
        m,
        stress_period_data={(0, 0): ["save head", "save budget"]},
        compact=True,
    )

    # Create WEL and SSM data
    itype = flopy.mt3d.Mt3dSsm.itype_dict()
    wel_data = {}
    ssm_data = {}
    wel_sp1 = []
    ssm_sp1 = []
    for k in range(nlay):
        wel_sp1.append([k, 0, 0, qinflow / nlay])
        ssm_sp1.append([k, 0, 0, 0.0, itype["WEL"]])
        ssm_sp1.append([k, 0, ncol - 1, 35.0, itype["BAS6"]])
    wel_data[0] = wel_sp1
    ssm_data[0] = ssm_sp1
    wel = flopy.modflow.ModflowWel(m, stress_period_data=wel_data)

    # Create the basic MT3DMS model data
    btn = flopy.mt3d.Mt3dBtn(
        m,
        nprs=-5,
        prsity=0.35,
        sconc=35.0,
        ifmtcn=0,
        chkmas=False,
        nprobs=10,
        nprmas=10,
        dt0=0.001,
    )
    adv = flopy.mt3d.Mt3dAdv(m, mixelm=0)
    dsp = flopy.mt3d.Mt3dDsp(m, al=0.0, trpt=1.0, trpv=1.0, dmcoef=dmcoef)
    gcg = flopy.mt3d.Mt3dGcg(m, iter1=500, mxiter=1, isolve=1, cclose=1e-7)
    ssm = flopy.mt3d.Mt3dSsm(m, stress_period_data=ssm_data)

    # Create the SEAWAT model data
    vdf = flopy.seawat.SeawatVdf(
        m,
        iwtable=0,
        densemin=0,
        densemax=0,
        denseref=1000.0,
        denseslp=0.7143,
        firstdt=1e-3,
    )

    # Write the input files
    m.write_input()

    # Try to delete the output files, to prevent accidental use of older files
    try:
        os.remove(os.path.join(workspace, "MT3D001.UCN"))
        os.remove(os.path.join(workspace, f"{modelname}.hds"))
        os.remove(os.path.join(workspace, f"{modelname}.cbc"))
    except:
        pass

    # run the model
    m.run_model(silent=quiet)

    # Post-process the results

    # Load data
    ucnobj = flopy.utils.binaryfile.UcnFile(
        os.path.join(workspace, "MT3D001.UCN"), model=m
    )
    times = ucnobj.get_times()
    concentration = ucnobj.get_data(totim=times[-1])
    cbbobj = flopy.utils.binaryfile.CellBudgetFile(
        os.path.join(workspace, f"{modelname}.cbc")
    )
    times = cbbobj.get_times()
    qx = cbbobj.get_data(text="flow right face", totim=times[-1])[0]
    qz = cbbobj.get_data(text="flow lower face", totim=times[-1])[0]

    # Average flows to cell centers
    qx_avg = np.empty(qx.shape, dtype=qx.dtype)
    qx_avg[:, :, 1:] = 0.5 * (qx[:, :, 0 : ncol - 1] + qx[:, :, 1:ncol])
    qx_avg[:, :, 0] = 0.5 * qx[:, :, 0]
    qz_avg = np.empty(qz.shape, dtype=qz.dtype)
    qz_avg[1:, :, :] = 0.5 * (qz[0 : nlay - 1, :, :] + qz[1:nlay, :, :])
    qz_avg[0, :, :] = 0.5 * qz[0, :, :]

    # Make the plot
    # import matplotlib.pyplot as plt
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(1, 1, 1, aspect="equal")
    ax.imshow(
        concentration[:, 0, :], interpolation="nearest", extent=(0, Lx, 0, Lz)
    )
    y, x, z = dis.get_node_coordinates()
    X, Z = np.meshgrid(x, z[:, 0, 0])
    iskip = 3
    ax.quiver(
        X[::iskip, ::iskip],
        Z[::iskip, ::iskip],
        qx_avg[::iskip, 0, ::iskip],
        -qz_avg[::iskip, 0, ::iskip],
        color="w",
        scale=5,
        headwidth=3,
        headlength=2,
        headaxislength=2,
        width=0.0025,
    )

    outfig = os.path.join(workspace, f"henry_flows.{fext}")
    fig.savefig(outfig, dpi=300)
    print("created...", outfig)

    # Extract the heads
    fname = os.path.join(workspace, f"{modelname}.hds")
    headobj = flopy.utils.binaryfile.HeadFile(fname)
    times = headobj.get_times()
    head = headobj.get_data(totim=times[-1])

    # Make a simple head plot
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(1, 1, 1, aspect="equal")
    im = ax.imshow(
        head[:, 0, :], interpolation="nearest", extent=(0, Lx, 0, Lz)
    )
    ax.set_title("Simulated Heads")

    outfig = os.path.join(workspace, f"henry_heads.{fext}")
    fig.savefig(outfig, dpi=300)
    print("created...", outfig)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--keep", help="output directory")
    parser.add_argument(
        "--quiet", action="store_false", help="don't show model output"
    )
    args = vars(parser.parse_args())

    workspace = args.get("keep", None)
    quiet = args.get("quiet", False)

    if workspace is not None:
        run(workspace, quiet)
    else:
        try:
            with TemporaryDirectory() as workspace:
                run(workspace, quiet)
        except PermissionError:
            # can occur on windows: https://docs.python.org/3/library/tempfile.html#tempfile.TemporaryDirectory
            pass