Add PRT support

docs/examples/00_model_from_scratch.ipynb

@@ -125,7 +125,9 @@
     "ims = nlmod.sim.ims(sim, complexity=\"SIMPLE\")\n",
     "gwf = nlmod.gwf.gwf(ds, sim)\n",
     "dis = nlmod.gwf.dis(ds, gwf)\n",
-    "npf = nlmod.gwf.npf(ds, gwf)\n",
+    "npf = nlmod.gwf.npf(\n",
+    "    ds, gwf, save_flows=True, save_specific_discharge=True, save_saturation=True\n",
+    ")\n",
     "ic = nlmod.gwf.ic(ds, gwf, starting_head=1.0)\n",
     "oc = nlmod.gwf.oc(ds, gwf, save_head=True)"
    ]

docs/examples/18_particle_tracking_prt.ipynb

@@ -0,0 +1,168 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Particle Tracking with PRT"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import flopy as fp\n",
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "\n",
+    "import nlmod"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nlmod.util.get_color_logger(\"INFO\")\n",
+    "nlmod.show_versions()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_ws = Path(\"./scratch_model\")\n",
+    "model_name = \"from_scratch\"\n",
+    "\n",
+    "# create new subdir for PRT model\n",
+    "prt_ws = Path(model_ws) / \"prt\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = xr.open_dataset(model_ws / f\"{model_name}.nc\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a new Simulation object\n",
+    "simprt = nlmod.sim.sim(ds, sim_ws=prt_ws)\n",
+    "tdis = nlmod.sim.tdis(ds, simprt)\n",
+    "\n",
+    "# Add PRT model\n",
+    "prt = nlmod.prt.prt(ds, simprt)\n",
+    "\n",
+    "# DIS: discretization package\n",
+    "dis = nlmod.prt.dis(ds, prt)\n",
+    "\n",
+    "# MIP: matrix input package(?)\n",
+    "mip = nlmod.prt.mip(ds, prt, porosity=0.3)\n",
+    "\n",
+    "# PRP: particle release package\n",
+    "pd = fp.modpath.ParticleData(\n",
+    "    [\n",
+    "        (k, i, j)\n",
+    "        for i in np.arange(0, ds.sizes[\"y\"], 11)\n",
+    "        for j in np.arange(0, ds.sizes[\"x\"], 11)\n",
+    "        for k in [0, 1, 2]\n",
+    "    ],\n",
+    "    structured=True,\n",
+    ")\n",
+    "releasepts = list(pd.to_prp(prt.modelgrid, global_xy=False))\n",
+    "prp = nlmod.prt.prp(ds, prt, packagedata=releasepts, perioddata={0: [\"FIRST\"]})\n",
+    "\n",
+    "# FMI: flow model interface\n",
+    "gwf_budget_file = Path(ds.model_ws) / f\"{model_name}.cbc\"\n",
+    "gwf_head_file = Path(ds.model_ws) / f\"{model_name}.hds\"\n",
+    "fmi = fp.mf6.ModflowPrtfmi(\n",
+    "    prt,\n",
+    "    packagedata=[\n",
+    "        (\"GWFHEAD\", gwf_head_file.absolute().resolve()),\n",
+    "        (\"GWFBUDGET\", gwf_budget_file.absolute().resolve()),\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "# OC: output control\n",
+    "oc = nlmod.prt.oc(ds, prt)\n",
+    "\n",
+    "# EMS: explicit model solution\n",
+    "ems = nlmod.sim.ems(simprt, model=prt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "simprt.write_simulation()\n",
+    "simprt.run_simulation()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load pathline data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from flopy.utils.prtfile import PathlineFile\n",
+    "\n",
+    "trkfile = PathlineFile(\n",
+    "    Path(prt.model_ws) / f\"{prt.name}.trk\",\n",
+    "    header_filename=Path(prt.model_ws) / f\"{prt.name}.trk.hdr\",\n",
+    ")\n",
+    "trk = trkfile.get_alldata()\n",
+    "trk_df = trkfile.get_dataframe()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pmv = fp.plot.PlotMapView(prt)\n",
+    "# loop to avoid different tracks get connected in plot\n",
+    "for i in range(len(trk)):\n",
+    "    pmv.plot_pathline(trk[i:i+1], layer=\"all\", colors=\"k\", lw=1.0)\n",
+    "pmv.plot_endpoint(trk_df, color=\"r\", marker=\".\", direction=\"starting\", zorder=10);"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

nlmod/__init__.py

@@ -3,7 +3,7 @@
 
 NLMOD_DATADIR = os.path.join(os.path.dirname(__file__), "data")
 
-from . import dims, gis, gwf, gwt, modpath, plot, read, sim, util
+from . import dims, gis, gwf, gwt, modpath, plot, prt, read, sim, util
 from .dims import base, get_ds, grid, layers, resample, time, to_model_ds
 from .util import download_mfbinaries
 from .version import __version__, show_versions

nlmod/gwf/gwf.py

@@ -114,46 +114,34 @@ def _dis(ds, model, length_units="METERS", pname="dis", **kwargs):
 
     idomain = get_idomain(ds).data
     if model.model_type == "gwf6":
-        dis = flopy.mf6.ModflowGwfdis(
-            model,
-            pname=pname,
-            length_units=length_units,
-            xorigin=xorigin,
-            yorigin=yorigin,
-            angrot=angrot,
-            nlay=ds.sizes["layer"],
-            nrow=ds.sizes["y"],
-            ncol=ds.sizes["x"],
-            delr=get_delr(ds),
-            delc=get_delc(ds),
-            top=ds["top"].data,
-            botm=ds["botm"].data,
-            idomain=idomain,
-            filename=f"{ds.model_name}.dis",
-            **kwargs,
-        )
+        filename = f"{ds.model_name}.dis"
+        klass = flopy.mf6.ModflowGwfdis
     elif model.model_type == "gwt6":
-        dis = flopy.mf6.ModflowGwtdis(
-            model,
-            pname=pname,
-            length_units=length_units,
-            xorigin=xorigin,
-            yorigin=yorigin,
-            angrot=angrot,
-            nlay=ds.sizes["layer"],
-            nrow=ds.sizes["y"],
-            ncol=ds.sizes["x"],
-            delr=get_delr(ds),
-            delc=get_delc(ds),
-            top=ds["top"].data,
-            botm=ds["botm"].data,
-            idomain=idomain,
-            filename=f"{ds.model_name}_gwt.dis",
-            **kwargs,
-        )
+        filename = f"{ds.model_name}_gwt.dis"
+        klass = flopy.mf6.ModflowGwtdis
+    elif model.model_type == "prt6":
+        filename = f"{ds.model_name}_prt.dis"
+        klass = flopy.mf6.ModflowPrtdis
     else:
         raise ValueError("Unknown model type.")
-
+    dis = klass(
+        model,
+        pname=pname,
+        length_units=length_units,
+        xorigin=xorigin,
+        yorigin=yorigin,
+        angrot=angrot,
+        nlay=ds.sizes["layer"],
+        nrow=ds.sizes["y"],
+        ncol=ds.sizes["x"],
+        delr=get_delr(ds),
+        delc=get_delc(ds),
+        top=ds["top"].data,
+        botm=ds["botm"].data,
+        idomain=idomain,
+        filename=filename,
+        **kwargs,
+    )
     return dis
 
 
@@ -213,45 +201,34 @@ def _disv(ds, model, length_units="METERS", pname="disv", **kwargs):
     cell2d = grid.get_cell2d_from_ds(ds)
     idomain = get_idomain(ds).data
     if model.model_type == "gwf6":
-        disv = flopy.mf6.ModflowGwfdisv(
-            model,
-            idomain=idomain,
-            xorigin=xorigin,
-            yorigin=yorigin,
-            length_units=length_units,
-            angrot=angrot,
-            nlay=len(ds.layer),
-            ncpl=len(ds.icell2d),
-            nvert=len(ds.iv),
-            top=ds["top"].data,
-            botm=ds["botm"].data,
-            vertices=vertices,
-            cell2d=cell2d,
-            pname=pname,
-            **kwargs,
-        )
+        klass = flopy.mf6.ModflowGwfdisv
+        filename = f"{ds.model_name}.disv"
     elif model.model_type == "gwt6":
-        disv = flopy.mf6.ModflowGwtdisv(
-            model,
-            idomain=idomain,
-            xorigin=xorigin,
-            yorigin=yorigin,
-            length_units=length_units,
-            angrot=angrot,
-            nlay=len(ds.layer),
-            ncpl=len(ds.icell2d),
-            nvert=len(ds.iv),
-            top=ds["top"].data,
-            botm=ds["botm"].data,
-            vertices=vertices,
-            cell2d=cell2d,
-            pname=pname,
-            filename=f"{ds.model_name}_gwt.disv",
-            **kwargs,
-        )
+        klass = flopy.mf6.ModflowGwtdisv
+        filename = f"{ds.model_name}.disv"
+    elif model.model_type == "prt6":
+        klass = flopy.mf6.ModflowPrtdisv
+        filename = (f"{ds.model_name}_prt.disv",)
     else:
         raise ValueError("Unknown model type.")
-
+    disv = klass(
+        model,
+        idomain=idomain,
+        xorigin=xorigin,
+        yorigin=yorigin,
+        length_units=length_units,
+        angrot=angrot,
+        nlay=len(ds.layer),
+        ncpl=len(ds.icell2d),
+        nvert=len(ds.iv),
+        top=ds["top"].data,
+        botm=ds["botm"].data,
+        vertices=vertices,
+        cell2d=cell2d,
+        pname=pname,
+        filename=filename,
+        **kwargs,
+    )
     if "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
         model.modelgrid.set_coord_info(xoff=xorigin, yoff=yorigin, angrot=angrot)
 
@@ -875,7 +852,7 @@ def _set_record(out, budget, output="head"):
     return record
 
 
-def buy(ds, gwf, pname="buy", **kwargs):
+def buy(ds, gwf, pname="buy", gwt_modelname=None, **kwargs):
     """Create buoyancy package from model dataset.
 
     Parameters
@@ -886,6 +863,9 @@ def buy(ds, gwf, pname="buy", **kwargs):
         groundwaterflow object.
     pname : str, optional
         package name, by default "buy"
+    gwt_modelname : str, optional
+        name of the groundwater transport model, by default None,
+        which uses gwf modelname with '_gwt' appended.
 
     Returns
     -------
@@ -903,6 +883,7 @@ def buy(ds, gwf, pname="buy", **kwargs):
             "BUY package requires a groundwater transport model. "
             "Set 'transport' to True in model dataset."
         )
+    logger.info("creating mf6 BUY")
 
     drhodc = _get_value_from_ds_attr(
         ds, "drhodc", attr="drhodc", value=kwargs.pop("drhodc", None)
@@ -914,7 +895,9 @@ def buy(ds, gwf, pname="buy", **kwargs):
         ds, "denseref", attr="denseref", value=kwargs.pop("denseref", None)
     )
 
-    pdata = [(0, drhodc, crhoref, f"{ds.model_name}_gwt", "CONCENTRATION")]
+    if gwt_modelname is None:
+        gwt_modelname = f"{ds.model_name}_gwt"
+    pdata = [(0, drhodc, crhoref, gwt_modelname, "CONCENTRATION")]
 
     buy = flopy.mf6.ModflowGwfbuy(
         gwf,

nlmod/mfoutput/mfoutput.py

@@ -252,8 +252,7 @@ def _get_budget_da(
 def _get_flopy_data_object(
     var, ds=None, gwml=None, fname=None, grb_file=None, **kwargs
 ):
-    """Get modflow HeadFile or CellBudgetFile object, containg heads, budgets or
-    concentrations.
+    """Get flopy data object, containing heads, budgets or concentrations.
 
     Provide one of ds, gwf or fname.