From 98febc7329ff013ca657f5628647eea438e3dd4c Mon Sep 17 00:00:00 2001 From: Wes Bonelli Date: Tue, 29 Aug 2023 12:32:58 -0400 Subject: [PATCH] feat(ParticleData): support to_coords for quadtree/patch vertex grids --- autotest/test_particledata.py | 101 ++++++++++++++++++++++++------- flopy/modpath/mp7particledata.py | 84 ++++++++++++------------- 2 files changed, 122 insertions(+), 63 deletions(-) diff --git a/autotest/test_particledata.py b/autotest/test_particledata.py index e1a640d64..e5fcd994b 100644 --- a/autotest/test_particledata.py +++ b/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 diff --git a/flopy/modpath/mp7particledata.py b/flopy/modpath/mp7particledata.py index a780a4638..0f33b8ec8 100644 --- a/flopy/modpath/mp7particledata.py +++ b/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,7 @@ def __init__( self.particlecount = particledata.shape[0] self.particleidoption = particleidoption self.locationstyle = locationstyle - self.particledata = particledata + self.particledata = pd.DataFrame.from_records(particledata) def write(self, f=None): """ @@ -334,7 +336,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,7 +358,7 @@ 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 ---------- @@ -364,59 +366,57 @@ def to_coords(self, grid): 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): """