feat(vtk): improve vtk particle track export (#2132)

* support PRT pathline format
* support dataframe or recarray
* support merged or separate paths
* attach layer array to exported points

autotest/test_export.py

@@ -1455,13 +1455,14 @@ def test_vtk_unstructured(function_tmpdir, unstructured_grid):
 
 
 @requires_pkg("vtk", "pyvista")
-def test_vtk_to_pyvista(function_tmpdir, example_data_path):
+def test_vtk_to_pyvista(function_tmpdir):
     from autotest.test_mp7_cases import Mp7Cases
+    from pprint import pformat
 
     case_mf6 = Mp7Cases.mp7_mf6(function_tmpdir)
     case_mf6.write_input()
     success, buff = case_mf6.run_model()
-    assert success, f"MP7 model ({case_mf6.name}) failed"
+    assert success, f"MP7 model ({case_mf6.name}) failed: {pformat(buff)}"
 
     gwf = case_mf6.flowmodel
     plf = PathlineFile(Path(case_mf6.model_ws) / f"{case_mf6.name}.mppth")
@@ -1479,6 +1480,9 @@ def test_vtk_to_pyvista(function_tmpdir, example_data_path):
     n_pts = sum([pl.shape[0] for pl in pls])
     assert pathlines.n_points == n_pts
     assert pathlines.n_cells == n_pts + len(pls)
+    assert "particleid" in pathlines.point_data
+    assert "time" in pathlines.point_data
+    assert "k" in pathlines.point_data
 
     # uncomment to debug
     # grid.plot()

flopy/export/vtk.py

@@ -9,6 +9,7 @@
 from typing import Union
 
 import numpy as np
+import pandas as pd
 
 from ..datbase import DataInterface, DataType
 from ..utils import Util3d, import_optional_dependency
@@ -1079,23 +1080,75 @@ def add_model(self, model, selpaklist=None, masked_values=None):
 
     def add_pathline_points(self, pathlines, timeseries=False):
         """
-        Method to add Modpath output from a pathline or timeseries file
-        to the grid. Colors will be representative of totim.
+        Method to add particle pathlines to the grid, with points
+        colored by travel-time. Supports MODPATH or MODFLOW6 PRT
+        pathline format, or MODPATH timeseries format.
 
         Parameters
         ----------
-        pathlines : np.recarray or list
-            pathlines accepts a numpy recarray of a particle pathline or
-            a list of numpy reccarrays associated with pathlines
-        timeseries : bool
-            method to plot data as a series of vtk timeseries files for
-            animation or as a single static vtk file. Default is false
+        pathlines : pd.dataframe, np.recarray or list
+            Particle pathlines, either as a single dataframe or recarray
+            or a list of such, separated by particle ID. If pathlines are
+            not provided separately, the dataframe or recarray must have
+            columns: 'time', 'k' & 'particleid' for MODPATH 3, 5, 6 or 7,
+            and 'irpt', 'iprp', 'imdl', and 'trelease' for MODFLOW 6 PRT,
+            so particle pathlines can be distinguished.
+        timeseries : bool, optional
+            Whether to plot data as a series of vtk timeseries files for
+            animation or as a single static vtk file. Default is false.
         """
 
-        if isinstance(pathlines, (np.recarray, np.ndarray)):
-            pathlines = [pathlines]
+        mpx_keys = ["particleid", "time", "k"]
+        prt_keys = ["imdl", "iprp", "irpt", "trelease", "ilay"]
+
+        if isinstance(pathlines, list):
+            if len(pathlines) == 0:
+                return
+            pathlines = [
+                (
+                    pl.to_records(index=False)
+                    if isinstance(pl, pd.DataFrame)
+                    else pl
+                )
+                for pl in pathlines
+            ]
+            if all(k in pathlines[0].dtype.names for k in mpx_keys):
+                keys = mpx_keys
+            elif all(k in pathlines[0].dtype.names for k in prt_keys):
+                keys = prt_keys
+            else:
+                raise ValueError("Unrecognized pathline dtype")
+        elif isinstance(pathlines, (np.recarray, np.ndarray, pd.DataFrame)):
+            if isinstance(pathlines, pd.DataFrame):
+                pathlines = pathlines.to_records(index=False)
+            if all(k in pathlines.dtype.names for k in mpx_keys):
+                keys = mpx_keys
+                pids = np.unique(pathlines.particleid)
+                pathlines = [
+                    pathlines[pathlines.particleid == pid] for pid in pids
+                ]
+            elif all(k in pathlines.dtype.names for k in prt_keys):
+                keys = prt_keys
+                pls = []
+                for imdl in np.unique(pathlines.imdl):
+                    for iprp in np.unique(pathlines.iprp):
+                        for irpt in np.unique(pathlines.irpt):
+                            pl = pathlines[
+                                (pathlines.imdl == imdl)
+                                & (pathlines.iprp == iprp)
+                                & (pathlines.irpt == irpt)
+                            ]
+                            pls.extend(
+                                [pl[pl.trelease == t] for t in np.unique(pl.t)]
+                            )
+                pathlines = pls
+            else:
+                raise ValueError("Unrecognized pathline dtype")
+        else:
+            raise ValueError(
+                "Unsupported pathline format, expected array, recarray, dataframe, or list"
+            )
 
-        keys = ["particleid", "time"]
         if not timeseries:
             arrays = {key: [] for key in keys}
             points = []