Skip to content

Commit

Permalink
feat(mp7particledata): add localz option for PRT PRP conversions (#2166)
Browse files Browse the repository at this point in the history 
* determines whether to return local z coordinates
* useful to set release elevation w.r.t. water table
* minor docstring improvements
  • Loading branch information
@wpbonelli
wpbonelli authored Apr 24, 2024
1 parent bb5461b commit e50ab9a
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 53 deletions.
64 changes: 41 additions & 23 deletions autotest/test_particledata.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from functools import reduce
from itertools import chain
from itertools import chain, repeat

import matplotlib.pyplot as plt
import numpy as np
Expand Down Expand Up @@ -276,42 +276,47 @@ def test_particledata_to_prp_dis_9():
assert np.allclose(rpts_prt, rpts_exp, atol=1e-3)


@pytest.mark.parametrize("localx", [None, 0.5, 0.25])
@pytest.mark.parametrize("localy", [None, 0.5, 0.25])
def test_particledata_to_prp_disv_1(localx, localy):
@pytest.mark.parametrize("lx", [None, 0.5, 0.25]) # local x coord
@pytest.mark.parametrize("ly", [None, 0.5, 0.25]) # local y coord
@pytest.mark.parametrize(
"localz", [False, True]
) # whether to return local z coords
def test_particledata_to_prp_disv_1(lx, ly, localz):
"""
1 particle in bottom left cell, testing with default
location (middle), explicitly specifying middle, and
offset in x and y directions
offset in x and y directions. Also test the `localz`
switch, which determines whether to return local z
coordinates (on interval [0, 1]) rather than global.
"""

# model grid
grid = GridCases().vertex_small()

# particle data
locs = [4]
localx = [localx] if localx else None
localy = [localy] if localy else None
lx = [lx] if lx else None
ly = [ly] if ly else None
part_data = ParticleData(
partlocs=locs,
structured=False,
particleids=range(len(locs)),
localx=localx,
localy=localy,
localx=lx,
localy=ly,
)

# convert to global coordinates
rpts_prt = flatten(list(part_data.to_prp(grid)))
rpts_prt = flatten(list(part_data.to_prp(grid, localz=localz)))

# check conversion succeeded
assert len(rpts_prt) == len(locs)
assert all(
len(c) == 6 for c in rpts_prt
) # each coord should be a tuple (irpt, k, j, x, y, z)
for ci, c in enumerate(rpts_prt):
assert np.isclose(c[3], localx[0] if localx else 0.5) # check x
assert np.isclose(c[4], localy[0] if localy else 0.5) # check y
assert np.isclose(c[5], 7.5) # check z
for rpt in rpts_prt:
assert np.isclose(rpt[3], lx[0] if lx else 0.5) # check x
assert np.isclose(rpt[4], ly[0] if ly else 0.5) # check y
assert np.isclose(rpt[5], 0.5 if localz else 7.5) # check z

# debugging: plot grid, cell centers, and particle location
# fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
Expand All @@ -333,7 +338,10 @@ def test_particledata_to_prp_disv_1(localx, localy):
# plt.show()


def test_particledata_to_prp_disv_9():
@pytest.mark.parametrize(
"localz", [False, True]
) # whether to return local z coords
def test_particledata_to_prp_disv_9(localz):
# minimal vertex grid
grid = GridCases().vertex_small()

Expand Down Expand Up @@ -367,24 +375,27 @@ def test_particledata_to_prp_disv_9():
0,
float(f"0.{i + 1}"),
float(f"2.{i + 1}"),
(grid.xyzextent[5] - grid.xyzextent[4]) / 2,
0.5 if localz else ((grid.xyzextent[5] - grid.xyzextent[4]) / 2),
)
for i in range(9)
]

# convert to prt format
rpts_prt = flatten(list(part_data.to_prp(grid)))
rpts_prt = flatten(list(part_data.to_prp(grid, localz=localz)))
assert np.allclose(rpts_prt, rpts_exp, atol=1e-3)


def test_lrcparticledata_to_prp_divisions_defaults():
@pytest.mark.parametrize(
"localz", [False, True]
) # whether to return local z coords
def test_lrcparticledata_to_prp_divisions_defaults(localz):
sd_data = CellDataType()
regions = [[0, 0, 1, 0, 1, 1]]
part_data = LRCParticleData(
subdivisiondata=[sd_data], lrcregions=[regions]
)
grid = GridCases().structured_small()
rpts_prt = flatten(list(part_data.to_prp(grid)))
rpts_prt = flatten(list(part_data.to_prp(grid, localz=localz)))
rpts_exp = [
[0, 0, 0, 1, 1.166666, 1.166666, 5.833333],
[1, 0, 0, 1, 1.166666, 1.166666, 7.5],
Expand Down Expand Up @@ -441,6 +452,9 @@ def test_lrcparticledata_to_prp_divisions_defaults():
[52, 0, 1, 1, 1.833333, 0.833333, 7.5],
[53, 0, 1, 1, 1.833333, 0.833333, 9.166666],
]
if localz:
locz = [0.166666, 0.5, 0.833333] * 18
rpts_exp = [(*rpt[0:-1], z) for rpt, z in zip(rpts_exp, locz)]

num_cells = reduce(
sum,
Expand Down Expand Up @@ -605,8 +619,11 @@ def test_lrcparticledata_to_prp_1_per_face():
assert np.allclose(rpts_prt, rpts_exp)


@pytest.mark.parametrize(
"localz", [False, True]
) # whether to return local z coords
def test_nodeparticledata_to_prp_disv_defaults(
function_tmpdir, example_data_path
function_tmpdir, example_data_path, localz
):
"""
This test loads a GWF simulation, runs it, and feeds it to an MP7 simulation
Expand Down Expand Up @@ -669,14 +686,15 @@ def test_nodeparticledata_to_prp_disv_defaults(
mp7_pls = pd.concat([pd.DataFrame(ra) for ra in pldata])
mp7_pls = mp7_pls.sort_values(by=["time", "particleid"]).head(27)
mp7_rpts = [
[0, r.k, r.x, r.y, r.z] for r in mp7_pls.itertuples()
[0, r.k, r.x, r.y, r.zloc if localz else r.z]
for r in mp7_pls.itertuples()
] # omit rpt index
mp7_rpts.sort()

# convert particle data to prt format, flatten (remove cell ID tuples),
# remove irpt as it is not guaranteed to match, and sort
prt_rpts = flatten(list(pdat.to_prp(grid)))
prt_rpts = [r[1:] for r in prt_rpts] #
prt_rpts = flatten(list(pdat.to_prp(grid, localz=localz)))
prt_rpts = [r[1:] for r in prt_rpts]
prt_rpts.sort()
assert np.allclose(prt_rpts, mp7_rpts)

Expand Down
92 changes: 62 additions & 30 deletions flopy/modpath/mp7particledata.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
"""
mp7particledata module. Contains the ParticleData, CellDataType,
FaceDataType, and NodeParticleData classes.
Support for MODPATH 7 particle release configurations. Contains the
ParticleData, CellDataType, FaceDataType, and NodeParticleData classes.
"""

from itertools import product
from typing import Generator, Iterator, Tuple
from typing import Iterator, Tuple

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -363,13 +363,16 @@ def write(self, f=None):
for v in d:
f.write(fmt.format(*v))

def to_coords(self, grid) -> Iterator[tuple]:
def to_coords(self, grid, localz=False) -> Iterator[tuple]:
"""
Compute global particle coordinates on the given grid.
Compute particle coordinates on the given grid.
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand Down Expand Up @@ -401,7 +404,9 @@ def convert(row) -> Tuple[float, float, float]:
return [
cvt_xy(row.localx, xs),
cvt_xy(row.localy, ys),
cvt_z(row.localz, row.k, row.i, row.j),
row.localz
if localz
else cvt_z(row.localz, row.k, row.i, row.j),
]

else:
Expand All @@ -425,13 +430,13 @@ def convert(row) -> Tuple[float, float, float]:
return [
cvt_xy(row.localx, xs),
cvt_xy(row.localy, ys),
cvt_z(row.localz, row.node),
row.localz if localz else cvt_z(row.localz, row.node),
]

for t in self.particledata.itertuples():
yield convert(t)

def to_prp(self, grid) -> Iterator[tuple]:
def to_prp(self, grid, localz=False) -> Iterator[tuple]:
"""
Convert particle data to PRT particle release point (PRP)
package data entries for the given grid. A model grid is
Expand All @@ -441,7 +446,10 @@ def to_prp(self, grid) -> Iterator[tuple]:
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand All @@ -454,7 +462,7 @@ def to_prp(self, grid) -> Iterator[tuple]:
for i, (t, c) in enumerate(
zip(
self.particledata.itertuples(index=False),
self.to_coords(grid),
self.to_coords(grid, localz),
)
):
row = [i] # release point index (irpt)
Expand Down Expand Up @@ -773,7 +781,9 @@ def write(self, f=None):
f.write(line)


def get_release_points(subdivisiondata, grid, k=None, i=None, j=None, nn=None):
def get_release_points(
subdivisiondata, grid, k=None, i=None, j=None, nn=None, localz=False
):
if nn is None and (k is None or i is None or j is None):
raise ValueError(
"A cell (node) must be specified by indices (for structured grids) or node number (for vertex/unstructured)"
Expand All @@ -785,7 +795,7 @@ def get_release_points(subdivisiondata, grid, k=None, i=None, j=None, nn=None):
# 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)
minz, maxz = grid.botm[k, i, j], grid.top[i, j]
minz, maxz = (0, 1) if localz else (grid.botm[k, i, j], grid.top[i, j])
elif nn is not None:
verts = grid.get_cell_vertices(nn)
if grid.grid_type == "structured":
Expand All @@ -797,8 +807,12 @@ def get_release_points(subdivisiondata, grid, k=None, i=None, j=None, nn=None):
else:
k, j = grid.get_lni([nn])[0]
minz, maxz = (
grid.botm[k, j],
grid.top[j] if k == 0 else grid.botm[k - 1, j],
(0, 1)
if localz
else (
grid.botm[k, j],
grid.top[j] if k == 0 else grid.botm[k - 1, j],
)
)
else:
raise ValueError(
Expand Down Expand Up @@ -1086,13 +1100,16 @@ def write(self, f=None):
line += "\n"
f.write(line)

def to_coords(self, grid) -> Iterator[tuple]:
def to_coords(self, grid, localz=False) -> Iterator[tuple]:
"""
Compute global particle coordinates on the given grid.
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand All @@ -1107,11 +1124,11 @@ def to_coords(self, grid) -> Iterator[tuple]:
for j in range(minj, maxj + 1):
for sd in self.subdivisiondata:
for rpt in get_release_points(
sd, grid, k, i, j
sd, grid, k, i, j, localz=localz
):
yield (rpt[3], rpt[4], rpt[5])
yield (*rpt[3:6],)

def to_prp(self, grid) -> Iterator[tuple]:
def to_prp(self, grid, localz=False) -> Iterator[tuple]:
"""
Convert particle data to PRT particle release point (PRP)
package data entries for the given grid. A model grid is
Expand All @@ -1121,7 +1138,10 @@ def to_prp(self, grid) -> Iterator[tuple]:
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand All @@ -1143,7 +1163,9 @@ def to_prp(self, grid) -> Iterator[tuple]:
for j in range(minj, maxj + 1):
for sd in self.subdivisiondata:
for irpt, rpt in enumerate(
get_release_points(sd, grid, k, i, j)
get_release_points(
sd, grid, k, i, j, localz=localz
)
):
assert rpt[0] == k
assert rpt[1] == i
Expand Down Expand Up @@ -1314,13 +1336,16 @@ def write(self, f=None):
line += "\n"
f.write(line)

def to_coords(self, grid) -> Iterator[tuple]:
def to_coords(self, grid, localz=False) -> Iterator[tuple]:
"""
Compute global particle coordinates on the given grid.
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand All @@ -1329,11 +1354,13 @@ def to_coords(self, grid) -> Iterator[tuple]:

for sd in self.subdivisiondata:
for nd in self.nodedata:
rpts = get_release_points(sd, grid, nn=int(nd[0]))
for i, rpt in enumerate(rpts):
yield (rpt[1], rpt[2], rpt[3])
rpts = get_release_points(
sd, grid, nn=int(nd[0]), localz=localz
)
for rpt in rpts:
yield (*rpt[1:4],)

def to_prp(self, grid) -> Iterator[tuple]:
def to_prp(self, grid, localz=False) -> Iterator[tuple]:
"""
Convert particle data to PRT particle release point (PRP)
package data entries for the given grid. A model grid is
Expand All @@ -1343,7 +1370,10 @@ def to_prp(self, grid) -> Iterator[tuple]:
Parameters
----------
grid : The grid on which to locate particle release points.
grid : flopy.discretization.grid.Grid
The grid on which to locate particle release points.
localz : bool, optional
Whether to return local z coordinates.
Returns
-------
Expand All @@ -1353,7 +1383,9 @@ def to_prp(self, grid) -> Iterator[tuple]:

for sd in self.subdivisiondata:
for nd in self.nodedata:
rpts = get_release_points(sd, grid, nn=int(nd[0]))
rpts = get_release_points(
sd, grid, nn=int(nd[0]), localz=localz
)
for irpt, rpt in enumerate(rpts):
row = [irpt]
if grid.grid_type == "structured":
Expand Down

0 comments on commit e50ab9a

Please sign in to comment.