feat(ParticleData): support to_coords for quadtree/patch vertex grids

.docs/Notebooks/mf6_prt_tutorial01.py

@@ -34,11 +34,12 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
-import flopy
 import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+
+import flopy
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
 
@@ -86,7 +87,8 @@
 #
 # In this tutorial we do the latter. First we define a GWF model and simulation, then a separate PRT model/simulation.
 #
-# Define a function to build a MODFLOW 6 GWF model. 
+# Define a function to build a MODFLOW 6 GWF model.
+
 
 def build_gwf_sim(ws, mf6):
     # create simulation
@@ -157,6 +159,7 @@ def build_gwf_sim(ws, mf6):
 
 # Define the PRT simulation.
 
+
 def build_prt_sim(ws, mf6):
     # create simulation
     sim = flopy.mf6.MFSimulation(
@@ -300,5 +303,9 @@ def build_prt_sim(ws, mf6):
 pmv.plot_grid()
 pmv.plot_array(hds[0], alpha=0.1)
 pmv.plot_vector(qx, qy, normalize=True, color="white")
-colors = cm.plasma(np.linspace(0, 1, pls.groupby(["imdl", "iprp", "irpt", "trelease"]).ngroups))
+colors = cm.plasma(
+    np.linspace(
+        0, 1, pls.groupby(["imdl", "iprp", "irpt", "trelease"]).ngroups
+    )
+)
 pmv.plot_pathline(pls, layer="all", colors=colors, linewidth=2)

.docs/Notebooks/modpath7_to_prt_migration_example.py

@@ -33,11 +33,12 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
-import flopy
 import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+
+import flopy
 from flopy.utils import PathlineFile
 from flopy.utils.binaryfile import HeadFile
 
@@ -290,6 +291,7 @@ def build_mp7_sim(ws, mp7, gwf):
 
     return mp
 
+
 # ## Running the models
 
 # Construct and run the models.

autotest/test_export.py

@@ -51,7 +51,6 @@
 from flopy.utils.crs import get_authority_crs
 from flopy.utils.geometry import Polygon
 
-
 HAS_PYPROJ = has_pkg("pyproj", strict=True)
 if HAS_PYPROJ:
     import pyproj

autotest/test_particledata.py

@@ -1,4 +1,6 @@
+import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import pytest
 from autotest.test_grid_cases import GridCases
 
@@ -11,6 +13,8 @@
     ParticleData,
 )
 
+# basic structured grid info
+
 structured_cells = [(0, 1, 1), (0, 1, 2)]
 structured_dtype = np.dtype(
     [
@@ -39,7 +43,10 @@ def test_particledata_structured_ctor_with_partlocs_as_list_of_tuples():
 
     assert data.particlecount == 2
     assert data.dtype == structured_dtype
-    assert np.array_equal(data.particledata, structured_array)
+    assert isinstance(data.particledata, pd.DataFrame)
+    assert np.array_equal(
+        data.particledata.to_records(index=False), structured_array
+    )
 
 
 def test_particledata_structured_ctor_with_partlocs_as_ndarray():
@@ -48,7 +55,10 @@ def test_particledata_structured_ctor_with_partlocs_as_ndarray():
 
     assert data.particlecount == 2
     assert data.dtype == structured_dtype
-    assert np.array_equal(data.particledata, structured_array)
+    assert isinstance(data.particledata, pd.DataFrame)
+    assert np.array_equal(
+        data.particledata.to_records(index=False), structured_array
+    )
 
 
 def test_particledata_structured_ctor_with_partlocs_as_list_of_lists():
@@ -57,31 +67,23 @@ def test_particledata_structured_ctor_with_partlocs_as_list_of_lists():
 
     assert data.particlecount == 2
     assert data.dtype == structured_dtype
-    assert np.array_equal(data.particledata, structured_array)
-
-
-def test_particledata_structured_to_coords_with_ids():
-    data = ParticleData(
-        partlocs=structured_cells,
-        structured=True,
-        particleids=range(len(structured_cells)),
+    assert isinstance(data.particledata, pd.DataFrame)
+    assert np.array_equal(
+        data.particledata.to_records(index=False), structured_array
     )
-    grid = GridCases().structured_small()
-    coords = data.to_coords(grid)
 
-    assert len(coords) == len(structured_cells)
-    assert all(
-        len(c) == 4 for c in coords
-    )  # each coord should be a tuple (id, x, y, z)
-    for ci, c in enumerate(coords):
-        assert ci == c[0]
 
+def test_particledata_structured_to_coords():
+    # model grid
+    grid = GridCases().structured_small()
 
-def test_particledata_structured_to_coords_with_no_ids():
+    # particle data
     data = ParticleData(partlocs=structured_cells, structured=True)
-    grid = GridCases().structured_small()
-    coords = data.to_coords(grid)
 
+    # convert to global coordinates
+    coords = list(data.to_coords(grid))
+
+    # check conversion
     assert len(coords) == len(structured_cells)
     assert all(
         len(c) == 3 for c in coords
@@ -110,6 +112,63 @@ def test_particledata_structured_to_coords_with_no_ids():
         )
 
 
+@pytest.mark.parametrize("localx", [None, 0.5, 0.25])
+@pytest.mark.parametrize("localy", [None, 0.5, 0.25])
+def test_particledata_vertex_to_coords(localx, localy):
+    """
+    1 particle in bottom left cell, testing with default
+    location (middle), explicitly specifying middle, and
+    offset in x and y directions
+    """
+
+    # model grid
+    grid = GridCases().vertex_small()
+
+    # particle data
+    locs = [4]
+    localx = [localx] if localx else None
+    localy = [localy] if localy else None
+    data = ParticleData(
+        partlocs=locs,
+        structured=False,
+        particleids=range(len(locs)),
+        localx=localx,
+        localy=localy,
+    )
+
+    # convert to global coordinates
+    coords = list(data.to_coords(grid))
+
+    # check conversion succeeded
+    assert len(coords) == len(locs)
+    assert all(
+        len(c) == 3 for c in coords
+    )  # each coord should be a tuple (x, y, z)
+    for ci, c in enumerate(coords):
+        assert np.isclose(c[0], localx[0] if localx else 0.5)  # check x
+        assert np.isclose(c[1], localy[0] if localy else 0.5)  # check y
+        assert np.isclose(c[2], 7.5)  # check z
+
+    # debugging: plot grid, cell centers, and particle location
+    # fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
+    # ax.set_aspect("equal")
+    # grid.plot() # plot grid
+    # xc, yc = ( # plot cell centers
+    #     grid.get_xcellcenters_for_layer(0),
+    #     grid.get_ycellcenters_for_layer(0)
+    # )
+    # xc = xc.flatten()
+    # yc = yc.flatten()
+    # for i in range(grid.ncpl):
+    #     x, y = xc[i], yc[i]
+    #     nn = grid.intersect(x, y, 0)[1]
+    #     ax.plot(x, y, "ro")
+    #     ax.annotate(str(nn + 1), (x, y), color="r") # 1-based node numbering
+    # for c in coords:  # plot particle location(s)
+    #     ax.plot(c[0], c[1], "bo")
+    # plt.show()
+
+
 def test_lrcparticledata_to_coords_with_particle_on_top_and_bottom_faces():
     rd = 1
     cd = 1

flopy/modpath/mp7particledata.py

@@ -5,8 +5,10 @@
 
 """
 from itertools import product
+from typing import Tuple
 
 import numpy as np
+import pandas as pd
 from numpy.lib.recfunctions import unstructured_to_structured
 
 from ..utils.recarray_utils import create_empty_recarray
@@ -315,7 +317,8 @@ def __init__(
         self.particlecount = particledata.shape[0]
         self.particleidoption = particleidoption
         self.locationstyle = locationstyle
-        self.particledata = particledata
+        self.dtype = particledata.dtype
+        self.particledata = pd.DataFrame.from_records(particledata)
 
     def write(self, f=None):
         """
@@ -334,7 +337,7 @@ def write(self, f=None):
             )
 
         # particle data item 4 and 5
-        d = np.recarray.copy(self.particledata)
+        d = np.recarray.copy(self.particledata.to_records(index=False))
         lnames = [name.lower() for name in d.dtype.names]
         # Add one to the kij and node indices
         for idx in (
@@ -356,67 +359,65 @@ def write(self, f=None):
 
     def to_coords(self, grid):
         """
-        Convert the particle representation to a list of global coordinates.
+        Convert particle data to global coordinates.
 
         Parameters
         ----------
         grid : The grid to use for locating particle coordinates
 
         Returns
         -------
-            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.
+            A list of coordinate tuples (x, y, z)
         """
 
+        def cvt_xy(p, vs):
+            mn, mx = min(vs), max(vs)
+            span = mx - mn
+            return mn + span * p
+
         if grid.grid_type == "structured":
             if not hasattr(self.particledata, "k"):
                 raise ValueError(
                     f"Particle representation is not structured but grid is"
                 )
 
-            def to_x(p):
-                i, j = p[1], p[2]
-                local = p[3]
-                verts = grid.get_cell_vertices(i, j)
-                xs, _ = list(zip(*verts))
-                minx, maxx = min(xs), max(xs)
-                span = maxx - minx
-                return minx + span * local
-
-            def to_y(p):
-                i, j = p[1], p[2]
-                local = p[4]
-                verts = grid.get_cell_vertices(i, j)
-                _, ys = list(zip(*verts))
-                miny, maxy = min(ys), max(ys)
-                span = maxy - miny
-                return miny + span * local
-
-            def to_z(p):
-                k, i, j = p[0], p[1], p[2]
-                local = p[5]
-                minz, maxz = grid.botm[k, i, j], grid.top[i, j]
-                span = maxz - minz
-                return minz + span * local
-
-            pdl = self.particledata.tolist()
-            if hasattr(self.particledata, "id"):
-                ids = [p[0] for p in pdl]
-                particles = [r[1:7] for r in pdl]
-                return [
-                    (id, to_x(p), to_y(p), to_z(p))
-                    for id, p in zip(ids, particles)
-                ]
-            else:
-                particles = [r[0:6] for r in pdl]
-                return [(to_x(p), to_y(p), to_z(p)) for p in particles]
+            def cvt_z(p, k, i, j):
+                mn, mx = grid.botm[k, i, j], grid.top[i, j]
+                span = mx - mn
+                return mn + span * p
+
+            def convert(row) -> Tuple[float, float, float]:
+                verts = grid.get_cell_vertices(row.i, row.j)
+                xs, ys = list(zip(*verts))
+                return (
+                    cvt_xy(row.localx, xs),
+                    cvt_xy(row.localy, ys),
+                    cvt_z(row.localz, row.k, row.i, row.j),
+                )
+
         else:
             if hasattr(self.particledata, "k"):
                 raise ValueError(
                     f"Particle representation is structured but grid is not"
                 )
 
-            raise NotImplementedError()
+            def cvt_z(p, nn):
+                k, j = grid.get_lni([nn])[0]
+                mn, mx = grid.botm[k, j], grid.top[j]
+                span = mx - mn
+                return mn + span * p
+
+            def convert(row) -> Tuple[float, float, float]:
+                verts = grid.get_cell_vertices(row.node)
+                xs, ys = list(zip(*verts))
+                return (
+                    cvt_xy(row.localx, xs),
+                    cvt_xy(row.localy, ys),
+                    cvt_z(row.localz, row.node),
+                )
+
+        for row in self.particledata.itertuples():
+            yield convert(row)
 
     def _get_dtype(self, structured, particleid):
         """
@@ -479,7 +480,7 @@ def _fmt_string(self):
 
         """
         fmts = []
-        for field in self.particledata.dtype.descr:
+        for field in self.dtype.descr:
             vtype = field[1][1].lower()
             if vtype == "i" or vtype == "b":
                 fmts.append("{:9d}")