fix(mp7particledata): add global_xy option for to_coords/to_prp (#2405)

Following PRT's expectation of model coordinates, convert to model coords by default if the grid has georef info. Add `global_xy` option to disable this and emit global coords. Default is False.

Note: method does get slow when it has to do the vertices conversion for lots of points. This can undoubtedly be optimized by first computing vertices in model coordinates or something along those lines. But I figured that was outside the scope for this PR.

Author: dbrakenhoff

autotest/test_grid_cases.py

@@ -10,7 +10,7 @@
 
 class GridCases:
     @staticmethod
-    def structured_small():
+    def structured_small(xoff=0.0, yoff=0.0):
         nlay, nrow, ncol = 3, 2, 3
         delc = 1.0 * np.ones(nrow, dtype=float)
         delr = 1.0 * np.ones(ncol, dtype=float)
@@ -29,6 +29,8 @@ def structured_small():
             top=top,
             botm=botm,
             idomain=idomain,
+            xoff=xoff,
+            yoff=yoff,
         )
 
     @staticmethod

autotest/test_particledata.py

@@ -245,6 +245,33 @@ def test_particledata_to_prp_dis_1():
         assert np.isclose(rpt[6], minz + (exp[ci][5] * (maxz - minz)))
 
 
+def test_particledata_to_prp_dis_1_global_xy():
+    # model grid
+    xoff = 100.0
+    yoff = 300.0
+    grid = GridCases().structured_small(xoff=xoff, yoff=yoff)
+
+    # particle data
+    cells = [(0, 1, 1), (0, 1, 2)]
+    part_data = ParticleData(partlocs=cells, structured=True)
+
+    # convert to global coordinates
+    rpts_prt_model_coords = flatten(list(part_data.to_prp(grid)))
+    rpts_prt_global_coords = flatten(list(part_data.to_prp(grid, global_xy=True)))
+
+    # check global and model coords
+    # x
+    assert np.all(
+        np.array(rpts_prt_global_coords)[:, -3] - xoff
+        == np.array(rpts_prt_model_coords)[:, -3]
+    )
+    # y
+    assert np.all(
+        np.array(rpts_prt_global_coords)[:, -2] - yoff
+        == np.array(rpts_prt_model_coords)[:, -2]
+    )
+
+
 def test_particledata_to_prp_dis_9():
     # minimal structured grid
     grid = GridCases().structured_small()

flopy/modpath/mp7particledata.py

@@ -362,7 +362,7 @@ def write(self, f=None):
         for v in d:
             f.write(fmt.format(*v))
 
-    def to_coords(self, grid, localz=False) -> Iterator[tuple]:
+    def to_coords(self, grid, localz=False, global_xy=False) -> Iterator[tuple]:
         """
         Compute particle coordinates on the given grid.
 
@@ -397,9 +397,12 @@ def cvt_z(p, k, i, j):
                 span = mx - mn
                 return mn + span * p
 
-            def convert(row) -> tuple[float, float, float]:
+            def convert(row, global_xy=False) -> tuple[float, float, float]:
                 verts = grid.get_cell_vertices(row.i, row.j)
-                xs, ys = list(zip(*verts))
+                if global_xy:
+                    xs, ys = list(zip(*verts))
+                else:
+                    xs, ys = grid.get_local_coords(*np.array(verts).T)
                 return [
                     cvt_xy(row.localx, xs),
                     cvt_xy(row.localy, ys),
@@ -421,19 +424,22 @@ def cvt_z(p, nn):
                 span = mx - mn
                 return mn + span * p
 
-            def convert(row) -> tuple[float, float, float]:
+            def convert(row, global_xy=False) -> tuple[float, float, float]:
                 verts = grid.get_cell_vertices(row.node)
-                xs, ys = list(zip(*verts))
+                if global_xy:
+                    xs, ys = list(zip(*verts))
+                else:
+                    xs, ys = grid.get_local_coords(*np.array(verts).T)
                 return [
                     cvt_xy(row.localx, xs),
                     cvt_xy(row.localy, ys),
                     row.localz if localz else cvt_z(row.localz, row.node),
                 ]
 
         for t in self.particledata.itertuples():
-            yield convert(t)
+            yield convert(t, global_xy=global_xy)
 
-    def to_prp(self, grid, localz=False) -> Iterator[tuple]:
+    def to_prp(self, grid, localz=False, global_xy=False) -> Iterator[tuple]:
         """
         Convert particle data to PRT particle release point (PRP)
         package data entries for the given grid. A model grid is
@@ -447,6 +453,8 @@ def to_prp(self, grid, localz=False) -> Iterator[tuple]:
             The grid on which to locate particle release points.
         localz : bool, optional
             Whether to return local z coordinates.
+        global_xy : bool, optional
+            Whether to return global x and y coordinates, default is False.
 
         Returns
         -------
@@ -459,7 +467,7 @@ def to_prp(self, grid, localz=False) -> Iterator[tuple]:
         for i, (t, c) in enumerate(
             zip(
                 self.particledata.itertuples(index=False),
-                self.to_coords(grid, localz),
+                self.to_coords(grid, localz, global_xy=global_xy),
             )
         ):
             row = [i]  # release point index (irpt)
@@ -778,10 +786,14 @@ def write(self, f=None):
 )
 
 
-def get_extent(grid, k=None, i=None, j=None, nn=None, localz=False) -> Extent:
+def get_extent(
+    grid, k=None, i=None, j=None, nn=None, localz=False, global_xy=False
+) -> Extent:
     # get cell coords and span in each dimension
     if not (k is None or i is None or j is None):
         verts = grid.get_cell_vertices(i, j)
+        if not global_xy and grid._has_ref_coordinates:
+            verts = list(zip(*grid.get_local_coords(*np.array(verts).T)))
         minz, maxz = (
             (0, 1)
             if localz
@@ -792,6 +804,8 @@ def get_extent(grid, k=None, i=None, j=None, nn=None, localz=False) -> Extent:
         )
     elif nn is not None:
         verts = grid.get_cell_vertices(nn)
+        if not global_xy and grid._has_ref_coordinates:
+            verts = list(zip(*grid.get_local_coords(*np.array(verts).T)))
         if grid.grid_type == "structured":
             k, i, j = grid.get_lrc([nn])[0]
             minz, maxz = (
@@ -967,7 +981,14 @@ def get_cell_release_points(subdivisiondata, cellid, extent) -> Iterator[tuple]:
 
 
 def get_release_points(
-    subdivisiondata, grid, k=None, i=None, j=None, nn=None, localz=False
+    subdivisiondata,
+    grid,
+    k=None,
+    i=None,
+    j=None,
+    nn=None,
+    localz=False,
+    global_xy=False,
 ) -> Iterator[tuple]:
     """
     Get MODPATH 7 release point tuples for the given cell.
@@ -980,7 +1001,7 @@ def get_release_points(
         )
 
     cellid = [k, i, j] if nn is None else [nn]
-    extent = get_extent(grid, k, i, j, nn, localz)
+    extent = get_extent(grid, k, i, j, nn, localz, global_xy=global_xy)
 
     if isinstance(subdivisiondata, FaceDataType):
         return get_face_release_points(subdivisiondata, cellid, extent)
@@ -1351,7 +1372,7 @@ def write(self, f=None):
                     line += "\n"
             f.write(line)
 
-    def to_coords(self, grid, localz=False) -> Iterator[tuple]:
+    def to_coords(self, grid, localz=False, global_xy=False) -> Iterator[tuple]:
         """
         Compute global particle coordinates on the given grid.
 
@@ -1361,6 +1382,8 @@ def to_coords(self, grid, localz=False) -> Iterator[tuple]:
             The grid on which to locate particle release points.
         localz : bool, optional
             Whether to return local z coordinates.
+        global_xy : bool, optional
+            Whether to return global x, y coordinates. Default is False.
 
         Returns
         -------
@@ -1369,23 +1392,27 @@ def to_coords(self, grid, localz=False) -> Iterator[tuple]:
 
         for sd in self.subdivisiondata:
             for nd in self.nodedata:
-                for rpt in get_release_points(sd, grid, nn=int(nd[0]), localz=localz):
+                for rpt in get_release_points(
+                    sd, grid, nn=int(nd[0]), localz=localz, global_xy=global_xy
+                ):
                     yield (*rpt[1:4],)
 
-    def to_prp(self, grid, localz=False) -> Iterator[tuple]:
+    def to_prp(self, grid, localz=False, global_xy=False) -> Iterator[tuple]:
         """
         Convert particle data to PRT particle release point (PRP)
         package data entries for the given grid. A model grid is
         required because MODPATH supports several ways to specify
         particle release locations by cell ID and subdivision info
-        or local coordinates, but PRT expects global coordinates.
+        or local coordinates, but PRT expects model coordinates, by default.
 
         Parameters
         ----------
         grid : flopy.discretization.grid.Grid
             The grid on which to locate particle release points.
         localz : bool, optional
             Whether to return local z coordinates.
+        global_xy : bool, optional
+            Whether to return global x, y coordinates. Default is False.
 
         Returns
         -------
@@ -1396,7 +1423,9 @@ def to_prp(self, grid, localz=False) -> Iterator[tuple]:
         for sd in self.subdivisiondata:
             for nd in self.nodedata:
                 for irpt, rpt in enumerate(
-                    get_release_points(sd, grid, nn=int(nd[0]), localz=localz)
+                    get_release_points(
+                        sd, grid, nn=int(nd[0]), localz=localz, global_xy=global_xy
+                    )
                 ):
                     row = [irpt]
                     if grid.grid_type == "structured":