flopy_swi2_ex2.py

import os
import sys
from pathlib import Path
from tempfile import TemporaryDirectory

import matplotlib.pyplot as plt
import numpy as np

import flopy

# --modify default matplotlib settings
updates = {
    "font.family": ["Univers 57 Condensed", "Arial"],
    "mathtext.default": "regular",
    "pdf.compression": 0,
    "pdf.fonttype": 42,
    "legend.fontsize": 7,
    "axes.labelsize": 8,
    "xtick.labelsize": 7,
    "ytick.labelsize": 7,
}
plt.rcParams.update(updates)


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

    dirs = [os.path.join(workspace, "SWI2"), os.path.join(workspace, "SEAWAT")]

    if cleanFiles:
        print("cleaning all files")
        print("excluding *.py files")
        file_dict = {}
        file_dict[0] = os.listdir(dirs[0])
        file_dict[1] = os.listdir(dirs[1])
        file_dict[-1] = os.listdir(workspace)
        for key, files in list(file_dict.items()):
            pth = "."
            if key >= 0:
                pth = dirs[key]
            for f in files:
                fpth = os.path.join(pth, f)
                if os.path.isdir(fpth):
                    continue
                if ".py" != os.path.splitext(f)[1].lower():
                    print(f"  removing...{os.path.basename(f)}")
                    try:
                        os.remove(fpth)
                    except:
                        pass
        for d in dirs:
            if os.path.exists(d):
                os.rmdir(d)
        return 1

    # make working directories
    for d in dirs:
        if not os.path.exists(d):
            os.mkdir(d)

    modelname = "swiex2"
    mf_name = "mf2005"

    # problem data
    nper = 1
    perlen = 2000
    nstp = 1000
    nlay, nrow, ncol = 1, 1, 60
    delr = 5.0
    nsurf = 2
    x = np.arange(0.5 * delr, ncol * delr, delr)
    xedge = np.linspace(0, float(ncol) * delr, len(x) + 1)
    ibound = np.ones((nrow, ncol), int)
    ibound[0, 0] = -1
    # swi2 data
    z0 = np.zeros((nlay, nrow, ncol), float)
    z1 = np.zeros((nlay, nrow, ncol), float)
    z0[0, 0, 30:38] = np.arange(-2.5, -40, -5)
    z0[0, 0, 38:] = -40
    z1[0, 0, 22:30] = np.arange(-2.5, -40, -5)
    z1[0, 0, 30:] = -40
    z = []
    z.append(z0)
    z.append(z1)
    ssz = 0.2
    isource = np.ones((nrow, ncol), "int")
    isource[0, 0] = 2

    # stratified model
    modelname = "swiex2_strat"
    print("creating...", modelname)
    ml = flopy.modflow.Modflow(
        modelname, version="mf2005", exe_name=mf_name, model_ws=dirs[0]
    )
    discret = flopy.modflow.ModflowDis(
        ml,
        nlay=1,
        ncol=ncol,
        nrow=nrow,
        delr=delr,
        delc=1,
        top=0,
        botm=[-40.0],
        nper=nper,
        perlen=perlen,
        nstp=nstp,
    )
    bas = flopy.modflow.ModflowBas(ml, ibound=ibound, strt=0.05)
    bcf = flopy.modflow.ModflowBcf(ml, laycon=0, tran=2 * 40)
    swi = flopy.modflow.ModflowSwi2(
        ml,
        iswizt=55,
        nsrf=nsurf,
        istrat=1,
        toeslope=0.2,
        tipslope=0.2,
        nu=[0, 0.0125, 0.025],
        zeta=z,
        ssz=ssz,
        isource=isource,
        nsolver=1,
    )
    oc = flopy.modflow.ModflowOc(
        ml, stress_period_data={(0, 999): ["save head"]}
    )
    pcg = flopy.modflow.ModflowPcg(ml)
    ml.write_input()
    # run stratified model
    if not skipRuns:
        m = ml.run_model(silent=quiet)

    # read stratified results
    zetafile = os.path.join(dirs[0], f"{modelname}.zta")
    zobj = flopy.utils.CellBudgetFile(zetafile)
    zkstpkper = zobj.get_kstpkper()
    zeta = zobj.get_data(kstpkper=zkstpkper[-1], text="ZETASRF  1")[0]
    zeta2 = zobj.get_data(kstpkper=zkstpkper[-1], text="ZETASRF  2")[0]
    #
    # vd model
    modelname = "swiex2_vd"
    print("creating...", modelname)
    ml = flopy.modflow.Modflow(
        modelname, version="mf2005", exe_name=mf_name, model_ws=dirs[0]
    )
    discret = flopy.modflow.ModflowDis(
        ml,
        nlay=1,
        ncol=ncol,
        nrow=nrow,
        delr=delr,
        delc=1,
        top=0,
        botm=[-40.0],
        nper=nper,
        perlen=perlen,
        nstp=nstp,
    )
    bas = flopy.modflow.ModflowBas(ml, ibound=ibound, strt=0.05)
    bcf = flopy.modflow.ModflowBcf(ml, laycon=0, tran=2 * 40)
    swi = flopy.modflow.ModflowSwi2(
        ml,
        iswizt=55,
        nsrf=nsurf,
        istrat=0,
        toeslope=0.2,
        tipslope=0.2,
        nu=[0, 0, 0.025, 0.025],
        zeta=z,
        ssz=ssz,
        isource=isource,
        nsolver=1,
    )
    oc = flopy.modflow.ModflowOc(
        ml, stress_period_data={(0, 999): ["save head"]}
    )
    pcg = flopy.modflow.ModflowPcg(ml)
    ml.write_input()
    # run vd model
    if not skipRuns:
        m = ml.run_model(silent=quiet)

    # read vd model data
    zetafile = os.path.join(dirs[0], f"{modelname}.zta")
    zobj = flopy.utils.CellBudgetFile(zetafile)
    zkstpkper = zobj.get_kstpkper()
    zetavd = zobj.get_data(kstpkper=zkstpkper[-1], text="ZETASRF  1")[0]
    zetavd2 = zobj.get_data(kstpkper=zkstpkper[-1], text="ZETASRF  2")[0]
    #
    # seawat model
    swtexe_name = "swtv4"
    modelname = "swiex2_swt"
    print("creating...", modelname)
    swt_xmax = 300.0
    swt_zmax = 40.0
    swt_delr = 1.0
    swt_delc = 1.0
    swt_delz = 0.5
    swt_ncol = int(swt_xmax / swt_delr)  # 300
    swt_nrow = 1
    swt_nlay = int(swt_zmax / swt_delz)  # 80
    print(swt_nlay, swt_nrow, swt_ncol)
    swt_ibound = np.ones((swt_nlay, swt_nrow, swt_ncol), int)
    # swt_ibound[0, swt_ncol-1, 0] = -1
    swt_ibound[0, 0, 0] = -1
    swt_x = np.arange(0.5 * swt_delr, swt_ncol * swt_delr, swt_delr)
    swt_xedge = np.linspace(0, float(ncol) * delr, len(swt_x) + 1)
    swt_top = 0.0
    z0 = swt_top
    swt_botm = np.zeros((swt_nlay), float)
    swt_z = np.zeros((swt_nlay), float)
    zcell = -swt_delz / 2.0
    for ilay in range(0, swt_nlay):
        z0 -= swt_delz
        swt_botm[ilay] = z0
        swt_z[ilay] = zcell
        zcell -= swt_delz
    # swt_X, swt_Z = np.meshgrid(swt_x, swt_botm)
    swt_X, swt_Z = np.meshgrid(swt_x, swt_z)
    # mt3d
    # mt3d boundary array set to all active
    icbund = np.ones((swt_nlay, swt_nrow, swt_ncol), int)
    # create initial concentrations for MT3D
    sconc = np.ones((swt_nlay, swt_nrow, swt_ncol), float)
    sconcp = np.zeros((swt_nlay, swt_ncol), float)
    xsb = 110
    xbf = 150
    for ilay in range(0, swt_nlay):
        for icol in range(0, swt_ncol):
            if swt_x[icol] > xsb:
                sconc[ilay, 0, icol] = 0.5
            if swt_x[icol] > xbf:
                sconc[ilay, 0, icol] = 0.0
        for icol in range(0, swt_ncol):
            sconcp[ilay, icol] = sconc[ilay, 0, icol]
        xsb += swt_delz
        xbf += swt_delz

    # ssm data
    itype = flopy.mt3d.Mt3dSsm.itype_dict()
    ssm_data = {0: [0, 0, 0, 35.0, itype["BAS6"]]}

    # print sconcp
    # mt3d print times
    timprs = (np.arange(5) + 1) * 2000.0
    nprs = len(timprs)
    # create the MODFLOW files
    m = flopy.seawat.Seawat(modelname, exe_name=swtexe_name, model_ws=dirs[1])
    discret = flopy.modflow.ModflowDis(
        m,
        nrow=swt_nrow,
        ncol=swt_ncol,
        nlay=swt_nlay,
        delr=swt_delr,
        delc=swt_delc,
        laycbd=0,
        top=swt_top,
        botm=swt_botm,
        nper=nper,
        perlen=perlen,
        nstp=1,
        steady=False,
    )
    bas = flopy.modflow.ModflowBas(m, ibound=swt_ibound, strt=0.05)
    lpf = flopy.modflow.ModflowLpf(
        m, hk=2.0, vka=2.0, ss=0.0, sy=0.0, laytyp=0, layavg=0
    )
    oc = flopy.modflow.ModflowOc(m, save_every=1, save_types=["save head"])
    pcg = flopy.modflow.ModflowPcg(m)
    # Create the MT3DMS model files
    adv = flopy.mt3d.Mt3dAdv(
        m,
        mixelm=-1,  # -1 is TVD
        percel=0.05,
        nadvfd=0,
        # 0 or 1 is upstream; 2 is central in space
        # particle based methods
        nplane=4,
        mxpart=1e7,
        itrack=2,
        dceps=1e-4,
        npl=16,
        nph=16,
        npmin=8,
        npmax=256,
    )
    btn = flopy.mt3d.Mt3dBtn(
        m,
        icbund=1,
        prsity=ssz,
        sconc=sconc,
        ifmtcn=-1,
        chkmas=False,
        nprobs=10,
        nprmas=10,
        dt0=0.0,
        ttsmult=1.2,
        ttsmax=100.0,
        ncomp=1,
        nprs=nprs,
        timprs=timprs,
        mxstrn=1e8,
    )
    dsp = flopy.mt3d.Mt3dDsp(m, al=0.0, trpt=1.0, trpv=1.0, dmcoef=0.0)
    gcg = flopy.mt3d.Mt3dGcg(
        m, mxiter=1, iter1=50, isolve=3, cclose=1e-6, iprgcg=5
    )
    ssm = flopy.mt3d.Mt3dSsm(m, stress_period_data=ssm_data)
    # Create the SEAWAT model files
    vdf = flopy.seawat.SeawatVdf(
        m,
        nswtcpl=1,
        iwtable=0,
        densemin=0,
        densemax=0,
        denseref=1000.0,
        denseslp=25.0,
        firstdt=1.0e-03,
    )
    m.write_input()
    # run seawat model
    if not skipRuns:
        m = m.run_model(silent=quiet)

    # read seawat model data
    ucnfile = os.path.join(dirs[1], "MT3D001.UCN")
    uobj = flopy.utils.UcnFile(ucnfile)
    times = uobj.get_times()
    print(times)
    ukstpkper = uobj.get_kstpkper()
    print(ukstpkper)
    c = uobj.get_data(totim=times[-1])
    conc = np.zeros((swt_nlay, swt_ncol), float)
    for icol in range(0, swt_ncol):
        for ilay in range(0, swt_nlay):
            conc[ilay, icol] = c[ilay, 0, icol]
    #
    # figure
    fwid = 7.0  # 6.50
    fhgt = 4.5  # 6.75
    flft = 0.125
    frgt = 0.95
    fbot = 0.125
    ftop = 0.925

    print("creating  cross-section figure...")
    xsf, axes = plt.subplots(3, 1, figsize=(fwid, fhgt), facecolor="w")
    xsf.subplots_adjust(
        wspace=0.25, hspace=0.25, left=flft, right=frgt, bottom=fbot, top=ftop
    )
    # plot initial conditions
    ax = axes[0]
    ax.text(
        -0.075,
        1.05,
        "A",
        transform=ax.transAxes,
        va="center",
        ha="center",
        size="8",
    )
    # text(.975, .1, '(a)', transform = ax.transAxes, va = 'center', ha = 'center')
    ax.plot([110, 150], [0, -40], "k")
    ax.plot([150, 190], [0, -40], "k")
    ax.set_xlim(0, 300)
    ax.set_ylim(-40, 0)
    ax.set_yticks(np.arange(-40, 1, 10))
    ax.text(50, -20, "salt", va="center", ha="center")
    ax.text(150, -20, "brackish", va="center", ha="center")
    ax.text(250, -20, "fresh", va="center", ha="center")
    ax.set_ylabel("Elevation, in meters")
    # plot stratified swi2 and seawat results
    ax = axes[1]
    ax.text(
        -0.075,
        1.05,
        "B",
        transform=ax.transAxes,
        va="center",
        ha="center",
        size="8",
    )
    #
    zp = zeta[0, 0, :]
    p = (zp < 0.0) & (zp > -40.0)
    ax.plot(x[p], zp[p], "b", linewidth=1.5, drawstyle="steps-mid")
    zp = zeta2[0, 0, :]
    p = (zp < 0.0) & (zp > -40.0)
    ax.plot(x[p], zp[p], "b", linewidth=1.5, drawstyle="steps-mid")
    # seawat data
    cc = ax.contour(
        swt_X,
        swt_Z,
        conc,
        levels=[0.25, 0.75],
        colors="k",
        linestyles="solid",
        linewidths=0.75,
        zorder=101,
    )
    # fake figures
    ax.plot([-100.0, -100], [-100.0, -100], "b", linewidth=1.5, label="SWI2")
    ax.plot(
        [-100.0, -100], [-100.0, -100], "k", linewidth=0.75, label="SEAWAT"
    )
    # legend
    leg = ax.legend(loc="lower left", numpoints=1)
    leg._drawFrame = False
    # axes
    ax.set_xlim(0, 300)
    ax.set_ylim(-40, 0)
    ax.set_yticks(np.arange(-40, 1, 10))
    ax.set_ylabel("Elevation, in meters")
    # plot vd model
    ax = axes[2]
    ax.text(
        -0.075,
        1.05,
        "C",
        transform=ax.transAxes,
        va="center",
        ha="center",
        size="8",
    )
    dr = zeta[0, 0, :]
    ax.plot(x, dr, "b", linewidth=1.5, drawstyle="steps-mid")
    dr = zeta2[0, 0, :]
    ax.plot(x, dr, "b", linewidth=1.5, drawstyle="steps-mid")
    dr = zetavd[0, 0, :]
    ax.plot(x, dr, "r", linewidth=0.75, drawstyle="steps-mid")
    dr = zetavd2[0, 0, :]
    ax.plot(x, dr, "r", linewidth=0.75, drawstyle="steps-mid")
    # fake figures
    ax.plot(
        [-100.0, -100],
        [-100.0, -100],
        "b",
        linewidth=1.5,
        label="SWI2 stratified option",
    )
    ax.plot(
        [-100.0, -100],
        [-100.0, -100],
        "r",
        linewidth=0.75,
        label="SWI2 continuous option",
    )
    # legend
    leg = ax.legend(loc="lower left", numpoints=1)
    leg._drawFrame = False
    # axes
    ax.set_xlim(0, 300)
    ax.set_ylim(-40, 0)
    ax.set_yticks(np.arange(-40, 1, 10))
    ax.set_xlabel("Horizontal distance, in meters")
    ax.set_ylabel("Elevation, in meters")

    outfig = os.path.join(workspace, f"Figure07_swi2ex2.{fext}")
    xsf.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:
        temp_dir = TemporaryDirectory()
        workspace = temp_dir.name

        run(workspace, quiet)

        try:
            temp_dir.cleanup()
        except:
            # prevent permission errors on windows
            pass