From 942a4cd0eadec1201b430867f84412a03751747e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dav=C3=ADd=20Brakenhoff?= Date: Mon, 23 Dec 2024 15:18:46 +0100 Subject: [PATCH] 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. --- autotest/test_grid_cases.py | 4 ++- autotest/test_particledata.py | 27 ++++++++++++++ flopy/modpath/mp7particledata.py | 61 +++++++++++++++++++++++--------- 3 files changed, 75 insertions(+), 17 deletions(-) diff --git a/autotest/test_grid_cases.py b/autotest/test_grid_cases.py index 7e6298153..9415a237b 100644 --- a/autotest/test_grid_cases.py +++ b/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 diff --git a/autotest/test_particledata.py b/autotest/test_particledata.py index e183bc962..bc3036339 100644 --- a/autotest/test_particledata.py +++ b/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() diff --git a/flopy/modpath/mp7particledata.py b/flopy/modpath/mp7particledata.py index a191315c9..b57bb2701 100644 --- a/flopy/modpath/mp7particledata.py +++ b/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,9 +424,12 @@ 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), @@ -431,9 +437,9 @@ def convert(row) -> tuple[float, float, float]: ] 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,16 +1392,18 @@ 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 ---------- @@ -1386,6 +1411,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, 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":