feat(prp): support coordinate transformations

autotest/test_prt_coord_transform.py

@@ -0,0 +1,219 @@
+"""
+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
+
+"""
+
+import flopy
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import pytest
+from flopy.utils.gridintersect import GridIntersect
+from framework import TestFramework
+from prt_test_utils import (
+    DEFAULT_EXIT_SOLVE_TOL,
+    FlopyReadmeCase,
+    get_model_name,
+)
+from shapely.geometry import Point
+
+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,8 @@
 		\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). The PRT model now supports coordinate transformations for release point input via keyword options for the PRP package. 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. (A future version will likely add support for LOCAL\_XY on unstructured grids, where "local" is reinterpreted in the context of the generalized tracking method). Release point x and y coordinates may also now be specified at an offset from the cell center (with no distance rescaling) via the LOCAL\_XY\_OFFSET option, supported for structured or unstructured grids.
+	%	\item xxx
 	\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