Skip to content

Commit

Permalink
feat(ParticleData): support to_coords for quadtree/patch vertex grids
Browse files Browse the repository at this point in the history
  • Loading branch information
wpbonelli committed Aug 29, 2023
1 parent 264bb4f commit 98febc7
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 63 deletions.
101 changes: 80 additions & 21 deletions autotest/test_particledata.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -11,6 +13,8 @@
ParticleData,
)

# basic structured grid info

structured_cells = [(0, 1, 1), (0, 1, 2)]
structured_dtype = np.dtype(
[
Expand Down Expand Up @@ -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():
Expand All @@ -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():
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
84 changes: 42 additions & 42 deletions flopy/modpath/mp7particledata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand All @@ -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 (
Expand All @@ -356,67 +358,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):
"""
Expand Down

0 comments on commit 98febc7

Please sign in to comment.