feat(prp): support local, offset and transformed coords

autotest/test_prt_coord_transform.py

@@ -0,0 +1,235 @@
+"""
+Tests ability to transform release point coordinates from
+global, local, or local offset coordinate representations
+to model coordinates, as well as to transform output data
+from model to global cooordinates.
+
+The grid is a 10x10 square with a single layer,
+the same flow system shown on the FloPy readme.
+
+Test cases are defined for each coordinate system option.
+These are:
+
+ - local_xy
+ - local_xy_offset
+ - dev_global_xy
+
+"""
+
+from pathlib import Path
+
+import flopy
+import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import pytest
+from flopy.utils import PathlineFile
+from flopy.utils.binaryfile import HeadFile
+from flopy.utils.gridintersect import GridIntersect
+from shapely.geometry import Point
+from framework import TestFramework
+from prt_test_utils import (
+    FlopyReadmeCase,
+    all_equal,
+    check_budget_data,
+    check_track_data,
+    get_model_name,
+    get_partdata,
+    has_default_boundnames,
+    DEFAULT_EXIT_SOLVE_TOL,
+)
+
+simname = "prtcrd"
+cases = [f"{simname}mdl", f"{simname}lcl", f"{simname}ofs"]  # todo global
+nodes = list(range(100))
+
+
+def get_start_points():
+    return np.add(
+        np.transpose(
+            np.array(np.meshgrid(range(10), range(10))).reshape(2, -1)
+        ),
+        0.5,
+    )
+
+
+def get_prp(prt):
+    gi = GridIntersect(prt.modelgrid, method="vertex", rtree=True)
+    if "mdl" in prt.name:
+        startpts = get_start_points()
+        releasepts = [
+            (i, (0, *gi.intersect(Point(rpt))[0].cellids), *rpt, 0.5)
+            for i, rpt in enumerate(startpts)
+        ]
+    elif "lcl" in prt.name:
+        releasepts = [
+            (nn, tuple([*prt.modelgrid.get_lrc([nn])[0]]), 0.5, 0.5, 0.5)
+            for nn in nodes
+        ]
+    elif "ofs" in prt.name:
+        releasepts = [
+            (nn, tuple([*prt.modelgrid.get_lrc([nn])[0]]), 0.0, 0.0, 0.0)
+            for nn in nodes
+        ]
+    elif "gbl" in prt.name:
+        # todo global
+        pass
+
+    prp_track_file = f"{prt.name}_prt.prp.trk"
+    prp_track_csv_file = f"{prt.name}_prt.prp.trk.csv"
+    return flopy.mf6.ModflowPrtprp(
+        prt,
+        pname="prp1",
+        filename=f"{prt.name}_1.prp",
+        nreleasepts=len(releasepts),
+        packagedata=releasepts,
+        perioddata={0: ["FIRST"]},
+        track_filerecord=[prp_track_file],
+        trackcsv_filerecord=[prp_track_csv_file],
+        local_xy="lcl" in prt.name,
+        local_xy_offset="ofs" in prt.name,
+        exit_solve_tolerance=DEFAULT_EXIT_SOLVE_TOL,
+        extend_tracking=True,
+    )
+
+
+def build_prt_sim(name, gwf_ws, prt_ws, mf6):
+    # create simulation
+    sim = flopy.mf6.MFSimulation(
+        sim_name=name,
+        exe_name=mf6,
+        version="mf6",
+        sim_ws=prt_ws,
+    )
+
+    # create tdis package
+    flopy.mf6.modflow.mftdis.ModflowTdis(
+        sim,
+        pname="tdis",
+        time_units="DAYS",
+        nper=FlopyReadmeCase.nper,
+        perioddata=[
+            (
+                FlopyReadmeCase.perlen,
+                FlopyReadmeCase.nstp,
+                FlopyReadmeCase.tsmult,
+            )
+        ],
+    )
+
+    # create prt model
+    prt_name = get_model_name(name, "prt")
+    prt = flopy.mf6.ModflowPrt(sim, modelname=prt_name, save_flows=True)
+
+    # create prt discretization
+    flopy.mf6.modflow.mfgwfdis.ModflowGwfdis(
+        prt,
+        pname="dis",
+        nlay=FlopyReadmeCase.nlay,
+        nrow=FlopyReadmeCase.nrow,
+        ncol=FlopyReadmeCase.ncol,
+    )
+
+    # create mip package
+    flopy.mf6.ModflowPrtmip(
+        prt, pname="mip", porosity=FlopyReadmeCase.porosity
+    )
+
+    # create prp package
+    prp = get_prp(prt)
+
+    # create output control package
+    prt_budget_file = f"{prt_name}.bud"
+    prt_track_file = f"{prt_name}.trk"
+    prt_track_csv_file = f"{prt_name}.trk.csv"
+    flopy.mf6.ModflowPrtoc(
+        prt,
+        pname="oc",
+        budget_filerecord=[prt_budget_file],
+        track_filerecord=[prt_track_file],
+        trackcsv_filerecord=[prt_track_csv_file],
+        saverecord=[("BUDGET", "ALL")],
+    )
+
+    # create the flow model interface
+    gwf_name = get_model_name(name, "gwf")
+    gwf_budget_file = gwf_ws / f"{gwf_name}.bud"
+    gwf_head_file = gwf_ws / f"{gwf_name}.hds"
+    flopy.mf6.ModflowPrtfmi(
+        prt,
+        packagedata=[
+            ("GWFHEAD", gwf_head_file),
+            ("GWFBUDGET", gwf_budget_file),
+        ],
+    )
+
+    # add explicit model solution
+    ems = flopy.mf6.ModflowEms(
+        sim,
+        pname="ems",
+        filename=f"{prt_name}.ems",
+    )
+    sim.register_solution_package(ems, [prt.name])
+
+    return sim
+
+
+def build_models(idx, test):
+    gwf_ws = test.workspace / "gwf"
+    prt_ws = test.workspace / "prt"
+    gwf_sim = FlopyReadmeCase.get_gwf_sim(
+        test.name, gwf_ws, test.targets["mf6"]
+    )
+    prt_sim = build_prt_sim(test.name, gwf_ws, prt_ws, test.targets["mf6"])
+    return gwf_sim, prt_sim
+
+
+def check_output(idx, test):
+    name = test.name
+    gwf_ws = test.workspace
+    prt_ws = test.workspace / "prt"
+    gwf_sim, prt_sim = test.sims[:2]
+    gwf = gwf_sim.get_model()
+    prt = prt_sim.get_model()
+    hds = gwf.output.head().get_data()
+    mg = gwf.modelgrid
+    pathlines = pd.read_csv(prt_ws / f"{prt.name}.trk.csv")
+    startpts = pathlines[pathlines.ireason == 0]
+    endpts = pathlines[pathlines.ireason == 3]
+
+    # check start points
+    assert np.allclose(
+        np.sort(startpts[["x", "y"]].to_numpy(), axis=0),
+        np.sort(np.array(get_start_points()), axis=0),
+    )
+
+    # setup plot
+    plot_data = False
+    if plot_data:
+        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 10))
+        for a in ax:
+            a.set_aspect("equal")
+
+        # plot pathline in map view
+        pmv = flopy.plot.PlotMapView(modelgrid=mg, ax=ax[0])
+        pmv.plot_grid()
+        pmv.plot_array(hds[0], alpha=0.1, cmap="jet")
+        pmv.plot_pathline(pathlines)
+
+        # view/save plot
+        plt.show()
+        plt.savefig(gwf_ws / f"test_{simname}.png")
+
+
+@pytest.mark.parametrize("idx, name", enumerate(cases))
+def test_mf6model(idx, name, function_tmpdir, targets):
+    test = TestFramework(
+        name=name,
+        workspace=function_tmpdir,
+        build=lambda t: build_models(idx, t),
+        check=lambda t: check_output(idx, t),
+        targets=targets,
+        compare=None,
+    )
+    test.run()

autotest/test_prt_fmi.py

@@ -1,7 +1,6 @@
 """
 Tests ability to run a GWF model then a PRT model
 in separate simulations via flow model interface,
-as well as
 
 The grid is a 10x10 square with a single layer,
 the same flow system shown on the FloPy readme.

doc/ReleaseNotes/develop.tex

@@ -46,7 +46,7 @@
 		\item The Stress Package Concentration (SPC) utility available with the SSM Package is now the referred to as the Stress Package Component utility in the MF6IO guide.  Additionally, some relatively minor refactoring of the code facilitates use of the SPC utility with the GWE model type so that TEMPERATURE arrays may be read by the utility.  The SPC acronym was maintained to preserve backward compatibility.
 		\item The GWF-GWF Exchange has been fixed to support the configuration of the Mover package (MVR) also in cases where the number of exchanges equals zero (NEXG = 0). In earlier versions this has caused MODFLOW to terminate with an error.
 		\item A PRT bug was fixed in which array-based input read from the RCH (recharge) or EVT (evapotranspiration) packages could fail to be processed correctly by the PRT FMI (flow model interface) package, causing a crash.
-
+		\item Previously the PRT model's PRP package required release points to be specified in the model coordinate system (except for the z coordinate via the LOCAL\_Z option). PRT also previously reported pathlines in the model coordinate system. The PRT model now supports coordinate transformations for release point input and pathline output via keyword options for the PRP package, while preserving the existing default model coordinate system. Release point x and y coordinates may now be specified for structured grids with reference to the local rectilinear cell (with coordinates scaled to the unit interval) via the LOCAL\_XY keyword option. The LOCAL\_XY\_OFFSET option, supported for structured or unstructured grids, can be used to locate release point x and y coordinates at an offset from the cell center (with no distance rescaling).
 	\end{itemize}
 
 	%\underline{INTERNAL FLOW PACKAGES}

doc/mf6io/mf6ivar/dfn/prt-prp.dfn

@@ -40,7 +40,31 @@ type keyword
 reader urword
 optional true
 longname whether to use local z coordinates
-description indicates that ``zrpt'' defines the local z coordinate of the release point within the cell, with value of 0 at the bottom and 1 at the top of the cell.  If the cell is partially saturated at release time, the top of the cell is considered to be the water table elevation (the head in the cell) rather than the top defined by the user.
+description indicates that ``zrpt'' defines the local z coordinate of the release point within the cell scaled to the unit interval, with value of 0 at the bottom and 1 at the top of the cell.  If the cell is partially saturated at release time, the top of the cell is considered to be the water table elevation (the head in the cell) rather than the top defined by the user.
+
+block options
+name local_xy
+type keyword
+reader urword
+optional true
+longname whether to use local x and y coordinates
+description indicates that ``xrpt'' and ``yrpt'' define the local x and y coordinates of the release point within the cell scaled to the unit interval, with value of 0 at the west and south faces and 1 at the east and north faces of the cell. This option may only be used with structured models.
+
+block options
+name local_xy_offset
+type keyword
+reader urword
+optional true
+longname whether to use local x and y offsets
+description indicates that ``xrpt'' and ``rpt'' define the local x and y offsets of the release point within the cell, where the local coordinate system's origin is the cell center and distances are not rescaled. This option may be used with structured or unstructured models.
+
+block options
+name dev_global_xy
+type keyword
+reader urword
+optional true
+longname whether to transform x and y coordinates
+description indicates that ``xrpt'' and ``yrpt'' define global coordinates and should be transformed to the model coordinate system if georeferencing information (origin and rotation) is available on the model grid. By default, release points are assumed to be in model coordinates. If this option is enabled, particle pathlines will also be reported in global coordinates in output files.
 
 block options
 name extend_tracking

src/Model/ModelUtilities/TrackControl.f90

@@ -110,10 +110,11 @@ end subroutine expand
   !! any PRP-level files with PRP index matching the particle's
   !! PRP index.
   !<
-  subroutine save(this, particle, kper, kstp, reason, level)
+  subroutine save(this, particle, dis, kper, kstp, reason, level)
     ! dummy
     class(TrackControlType), intent(inout) :: this
     type(ParticleType), pointer, intent(in) :: particle
+    class(DisBaseType), pointer, intent(in) :: dis
     integer(I4B), intent(in) :: kper
     integer(I4B), intent(in) :: kstp
     integer(I4B), intent(in) :: reason
@@ -146,7 +147,7 @@ subroutine save(this, particle, kper, kstp, reason, level)
       if (file%iun > 0 .and. &
           (file%iprp == -1 .or. &
            file%iprp == particle%iprp)) &
-        call save_record(file%iun, particle, &
+        call save_record(file%iun, particle, dis, &
                          kper, kstp, reason, csv=file%csv)
     end do
   end subroutine save

src/Model/ModelUtilities/TrackFile.f90

@@ -2,6 +2,7 @@ module TrackFileModule
   use KindModule, only: DP, I4B, LGP
   use ConstantsModule, only: DZERO, DPIO180
   use ParticleModule, only: ParticleType
+  use BaseDisModule, only: DisBaseType
   use GeomUtilModule, only: transform
 
   implicit none
@@ -85,21 +86,39 @@ module TrackFileModule
 contains
 
   !> @brief Save a particle track record to a binary or CSV file.
-  subroutine save_record(iun, particle, kper, kstp, reason, csv)
+  subroutine save_record(iun, particle, dis, kper, kstp, reason, csv)
     ! dummy
     integer(I4B), intent(in) :: iun
     type(ParticleType), pointer, intent(in) :: particle
+    class(DisBaseType), pointer, intent(in) :: dis
     integer(I4B), intent(in) :: kper
     integer(I4B), intent(in) :: kstp
     integer(I4B), intent(in) :: reason
     logical(LGP), intent(in) :: csv
     ! local
-    real(DP) :: x, y, z
+    real(DP) :: x, y, z, xout, yout, zout
     integer(I4B) :: status
 
     ! Convert from cell-local to model coordinates if needed
     call particle%get_model_coords(x, y, z)
 
+    ! Apply model grid offset/rotation if needed
+    if (particle%icoords == 1 .and. &
+        (dis%angrot /= DZERO .or. &
+         dis%xorigin /= DZERO .or. &
+         dis%yorigin /= DZERO)) then
+      call transform(x, y, z, &
+                     xout, yout, zout, &
+                     dis%xorigin, &
+                     dis%yorigin, &
+                     DZERO, &
+                     sin(dis%angrot * DPIO180), &
+                     cos(dis%angrot * DPIO180), &
+                     invert=.true.)
+      x = xout
+      y = yout
+    end if
+
     ! Set status
     if (particle%istatus .lt. 0) then
       status = 1