support PRT output in plotting routines (WIP)

autotest/test_plot.py

@@ -358,7 +358,7 @@ def test_model_dot_plot_export(function_tmpdir, example_data_path):
 
 
 @pytest.fixture
-def modpath_model(function_tmpdir, example_data_path):
+def mp6_model(function_tmpdir, example_data_path):
     # test with multi-layer example
     load_ws = example_data_path / "mp6"
 
@@ -388,8 +388,8 @@ def modpath_model(function_tmpdir, example_data_path):
 
 
 @requires_exe("mf2005", "mp6")
-def test_plot_map_view_mp6_plot_pathline(modpath_model):
-    ml, mp, sim = modpath_model
+def test_plot_map_view_mp6_pathline(mp6_model):
+    ml, mp, sim = mp6_model
     mp.write_input()
     mp.run_model(silent=False)
 
@@ -421,8 +421,8 @@ def test_plot(pl):
 
 
 @requires_exe("mf2005", "mp6")
-def test_plot_cross_section_mp6_plot_pathline(modpath_model):
-    ml, mp, sim = modpath_model
+def test_plot_cross_section_mp6_pathline(mp6_model):
+    ml, mp, sim = mp6_model
     mp.write_input()
     mp.run_model(silent=False)
 
@@ -452,8 +452,8 @@ def test_plot(pl):
 
 
 @requires_exe("mf2005", "mp6")
-def test_plot_map_view_mp6_endpoint(modpath_model):
-    ml, mp, sim = modpath_model
+def test_plot_map_view_mp6_endpoint(mp6_model):
+    ml, mp, sim = mp6_model
     mp.write_input()
     mp.run_model(silent=False)
 
@@ -518,6 +518,46 @@ def test_plot_map_view_mp6_endpoint(modpath_model):
     assert isinstance(ep, PathCollection)
 
 
+@pytest.fixture
+def mp7_model(function_tmpdir, example_data_path):
+    pass
+
+
+@requires_exe("mf6", "mp7")
+def test_plot_map_view_mp7_pathline(mp7_model):
+    pass
+
+
+@requires_exe("mf6", "mp7")
+def test_plot_cross_section_mp7_pathline(mp7_model):
+    pass
+
+
+@pytest.fixture
+def prt_model(function_tmpdir, example_data_path):
+    pass
+
+
+@requires_exe("mf6")
+def test_plot_map_view_prt_pathline(prt_model):
+    pass
+
+
+@requires_exe("mf6")
+def test_plot_cross_section_prt_pathline(prt_model):
+    pass
+
+
+@requires_exe("mf6")
+def test_to_mp7_pathlines(prt_model):
+    pass
+
+
+@requires_exe("mf6")
+def test_to_mp7_endpoints(prt_model):
+    pass
+
+
 @pytest.fixture
 def quasi3d_model(function_tmpdir):
     mf = Modflow("model_mf", model_ws=function_tmpdir, exe_name="mf2005")

flopy/modpath/mp7particledata.py

@@ -364,7 +364,8 @@ def to_coords(self, grid):
 
         Returns
         -------
-            A list of particle tuples (particle ID, x coord, y coord, z coord)
+            A list of particle tuples ([optional particle ID], x coord, y coord, z coord)
+            where particle ID is only included if provided to the particle data instance.
         """
 
         if grid.grid_type == "structured":
@@ -375,7 +376,7 @@ def to_coords(self, grid):
 
             def to_x(p):
                 i, j = p[1], p[2]
-                local = p[4]
+                local = p[3]
                 verts = grid.get_cell_vertices(i, j)
                 xs, _ = list(zip(*verts))
                 minx, maxx = min(xs), max(xs)
@@ -384,7 +385,7 @@ def to_x(p):
 
             def to_y(p):
                 i, j = p[1], p[2]
-                local = p[5]
+                local = p[4]
                 verts = grid.get_cell_vertices(i, j)
                 _, ys = list(zip(*verts))
                 miny, maxy = min(ys), max(ys)
@@ -393,7 +394,7 @@ def to_y(p):
 
             def to_z(p):
                 k, i, j = p[0], p[1], p[2]
-                local = p[3]
+                local = p[5]
                 minz, maxz = grid.botm[k, i, j], grid.top[i, j]
                 span = maxz - minz
                 return minz + span * local

flopy/plot/map.py

@@ -696,19 +696,27 @@ def plot_vector(
 
     def plot_pathline(self, pl, travel_time=None, **kwargs):
         """
-        Plot MODPATH pathlines.
+        Plot particle pathlines. Compatible with MODFLOW 6 PRT particle track
+        data format, or MODPATH 6 or 7 pathline data format.
 
         Parameters
         ----------
         pl : list of recarrays or dataframes, or a single recarray or dataframe
             Particle pathline data. If a list of recarrays or dataframes,
             each must contain the path of only a single particle. If just
             one recarray or dataframe, it should contain the paths of all
-            particles. Pathline data returned from PathlineFile.get_data()
-            or get_alldata() can be passed directly as this argument. Data
-            columns should be 'x', 'y', 'z', 'time', 'k', and 'particleid'
-            at minimum. Additional columns are ignored. The 'particleid'
-            column must be unique to each particle path.
+            particles. The flopy.utils.modpathfile.PathlineFile.get_data()
+            or get_alldata() return value may be passed directly as this
+            argument.
+            
+            For MODPATH 6 or 7 pathlines, columns must include 'x', 'y', 'z',
+            'time', 'k', and 'particleid'. Additional columns are ignored.
+
+            For MODFLOW 6 PRT pathlines, columns must include 'x', 'y', 'z',
+            't', 'trelease', 'imdl', 'iprp', 'irpt', and 'ilay'. Additional
+            columns are ignored. Note that MODFLOW 6 PRT does not assign to
+            particles a unique ID, but infers particle identity from 'imdl',
+            'iprp', 'irpt', and 'trelease' combos (i.e. via composite key).
         travel_time : float or str
             Travel time selection. If a float, then pathlines with total
             time less than or equal to the given value are plotted. If a
@@ -729,6 +737,9 @@ def plot_pathline(self, pl, travel_time=None, **kwargs):
 
         from matplotlib.collections import LineCollection
 
+        # todo check provided data format
+        # todo convert prt to mp7 format
+
         # make sure pathlines is a list
         if not isinstance(pl, list):
             pids = np.unique(pl["particleid"])
@@ -817,7 +828,7 @@ def plot_pathline(self, pl, travel_time=None, **kwargs):
 
     def plot_timeseries(self, ts, travel_time=None, **kwargs):
         """
-        Plot MODPATH timeseries.
+        Plot MODPATH 6 or 7 timeseries. Incompatible with MODFLOW 6 PRT.
 
         Parameters
         ----------
@@ -861,13 +872,20 @@ def plot_endpoint(
         **kwargs,
     ):
         """
-        Plot MODPATH endpoints.
+        Plot particle endpoints. Compatible with MODFLOW 6 PRT particle
+        track data format, or MODPATH 6 or 7 endpoint data format.
 
         Parameters
         ----------
         ep : recarray or dataframe
             A numpy recarray with the endpoint particle data from the
-            MODPATH endpoint file
+            MODPATH endpoint file.
+
+            For MODFLOW 6 PRT pathlines, columns must include 'x', 'y', 'z',
+            't', 'trelease', 'imdl', 'iprp', 'irpt', and 'ilay'. Additional
+            columns are ignored. Note that MODFLOW 6 PRT does not assign to
+            particles a unique ID, but infers particle identity from 'imdl',
+            'iprp', 'irpt', and 'trelease' combos (i.e. via composite key).
         direction : str
             String defining if starting or ending particle locations should be
             considered. (default is 'ending')
@@ -898,6 +916,9 @@ def plot_endpoint(
 
         """
 
+        # todo check provided data format
+        # todo convert prt to mp7 format
+
         ax = kwargs.pop("ax", self.ax)
         tep, _, xp, yp = plotutil.parse_modpath_selection_options(
             ep, direction, selection, selection_direction

flopy/plot/plotutil.py

@@ -10,6 +10,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 
 from ..datbase import DataInterface, DataType
 from ..utils import Util3d, import_optional_dependency
@@ -2593,7 +2594,7 @@ def parse_modpath_selection_options(
     return tep, istart, xp, yp
 
 
-def to_mp7_pathlines(data):
+def to_mp7_pathlines(data) -> pd.DataFrame:
     """
     Convert MODFLOW 6 PRT pathline data to MODPATH 7 pathline format.
 
@@ -2607,17 +2608,16 @@ def to_mp7_pathlines(data):
     pd.DataFrame
     """
 
-    pd = import_optional_dependency("pandas")
-
     # convert to dataframe if needed
     if isinstance(data, np.recarray):
         data = pd.DataFrame(data)
 
     # assign a unique particle index column incrementing an integer
     # for each unique combination of irpt, iprp, imdl, and trelease
     data = data.sort_values(["imdl", "iprp", "irpt", "trelease"])
-    seqn = data.groupby(["imdl", "iprp", "irpt", "trelease"])
-    data["sequencenumber"] = seqn.ngroup()
+    particles = data.groupby(["imdl", "iprp", "irpt", "trelease"])
+    seqn_key = "sequencenumber"
+    data[seqn_key] = particles.ngroup()
 
     # convert to recarray
     data = data.to_records(index=False)
@@ -2628,7 +2628,7 @@ def to_mp7_pathlines(data):
             ("particleid", np.int32),  # same as sequencenumber
             ("particlegroup", np.int32),
             (
-                "sequencenumber",
+                seqn_key,
                 np.int32,
             ),  # mp7 sequencenumber (globally unique auto-generated ID)
             (
@@ -2653,9 +2653,9 @@ def to_mp7_pathlines(data):
     return pd.DataFrame(
         np.core.records.fromarrays(
             [
-                data["sequencenumber"],
+                data[seqn_key],
                 data["iprp"],
-                data["sequencenumber"],
+                data[seqn_key],
                 data["irpt"],
                 data["t"],
                 data["x"],
@@ -2675,7 +2675,7 @@ def to_mp7_pathlines(data):
     )
 
 
-def to_mp7_endpoints(data):
+def to_mp7_endpoints(data) -> pd.DataFrame:
     """
     Convert MODFLOW 6 PRT pathline data to MODPATH 7 endpoint format.
 
@@ -2689,33 +2689,36 @@ def to_mp7_endpoints(data):
     pd.DataFrame
     """
 
-    pd = import_optional_dependency("pandas")
-
-    #  convert to dataframe if needed
+    # convert to dataframe if needed
     if isinstance(data, np.recarray):
         data = pd.DataFrame(data)
 
     # assign a unique particle index column incrementing an integer
     # for each unique combination of irpt, iprp, imdl, and trelease
     data = data.sort_values(["imdl", "iprp", "irpt", "trelease"])
-    seqn = data.groupby(["imdl", "iprp", "irpt", "trelease"])
-    data["sequencenumber"] = seqn.ngroup()
+    particles = data.groupby(["imdl", "iprp", "irpt", "trelease"])
+    seqn_key = "sequencenumber"
+    data[seqn_key] = particles.ngroup()
 
-    # select (and sort) startpoints and endpoints
+    # select startpoints and endpoints, sort by sequencenumber
     startpts = (
         data.sort_values("t")
-        .groupby(["irpt", "iprp"])
+        .groupby(seqn_key)
         .head(1)
-        .sort_values(["irpt", "iprp"])
+        .sort_values(seqn_key)
     )
     endpts = (
         data.sort_values("t")
-        .groupby(["irpt", "iprp"])
+        .groupby(seqn_key)
         .tail(1)
-        .sort_values(["irpt", "iprp"])
+        .sort_values(seqn_key)
     )
 
-    # add new columns for initial coordinates, initial zone, and initial cell (node) number
+    # add new columns for:
+    #   - initial coordinates
+    #   - initial zone
+    #   - initial node number
+    #   - initial layer
     starttimes = startpts["t"].to_numpy()
     startxs = startpts["x"].to_numpy()
     startys = startpts["y"].to_numpy()
@@ -2724,7 +2727,7 @@ def to_mp7_endpoints(data):
     startnodes = startpts["icell"].to_numpy()
     startlayers = startpts["ilay"].to_numpy()
     conditions = [
-        startpts["irpt"].eq(row["irpt"]) & startpts["iprp"].eq(row["iprp"])
+        startpts[seqn_key].eq(row[seqn_key])
         for i, row in startpts.iterrows()
     ]
     endpts["x0"] = np.select(conditions, startxs)
@@ -2734,7 +2737,7 @@ def to_mp7_endpoints(data):
     endpts["node0"] = np.select(conditions, startnodes)
     endpts["k0"] = np.select(conditions, startlayers)
 
-    # convert endpoints to recarray
+    # convert to recarray
     endpts = endpts.to_records(index=False)
 
     # define mp7 dtype