From 02749f09b1edface66dc822b91bf1fcd905c6c71 Mon Sep 17 00:00:00 2001 From: Mike Taves Date: Thu, 20 Oct 2022 11:25:47 +1300 Subject: [PATCH] refactor(get_lrc, get_node): use numpy methods, add examples (#1591) --- autotest/test_grid.py | 29 ++++++++---- autotest/test_gridintersect.py | 11 ++--- flopy/discretization/structuredgrid.py | 65 +++++++++++++++++--------- 3 files changed, 66 insertions(+), 39 deletions(-) diff --git a/autotest/test_grid.py b/autotest/test_grid.py index a26bee4d22..4a3e3eb82f 100644 --- a/autotest/test_grid.py +++ b/autotest/test_grid.py @@ -100,27 +100,36 @@ def test_get_vertices(): def test_get_lrc_get_node(): nlay, nrow, ncol = 3, 4, 5 + nnodes = nlay * nrow * ncol ml = Modflow() dis = ModflowDis( ml, nlay=nlay, nrow=nrow, ncol=ncol, top=50, botm=[0, -1, -2] ) - nodes = list(range(nlay * nrow * ncol)) + nodes = list(range(nnodes)) indices = np.indices((nlay, nrow, ncol)) layers = indices[0].flatten() rows = indices[1].flatten() cols = indices[2].flatten() for node, (l, r, c) in enumerate(zip(layers, rows, cols)): # ensure get_lrc returns zero-based layer row col - lrc = dis.get_lrc(node)[0] - assert lrc == ( - l, - r, - c, - ), f"get_lrc() returned {lrc}, expecting {l, r, c}" + assert dis.get_lrc(node)[0] == (l, r, c) # ensure get_node returns zero-based node number - n = dis.get_node((l, r, c))[0] - assert node == n, f"get_node() returned {n}, expecting {node}" - return + assert dis.get_node((l, r, c))[0] == node + + # check full list + lrc_list = list(zip(layers, rows, cols)) + assert dis.get_lrc(nodes) == lrc_list + assert dis.get_node(lrc_list) == nodes + + # check array-like input + assert dis.get_lrc(np.arange(nnodes)) == lrc_list + # dis.get_node does not accept array-like inputs, just tuple or list + + # check out of bounds errors + with pytest.raises(ValueError, match="index 60 is out of bounds for"): + dis.get_lrc(nnodes) + with pytest.raises(ValueError, match="invalid entry in coordinates array"): + dis.get_node((4, 4, 4)) def test_get_rc_from_node_coordinates(): diff --git a/autotest/test_gridintersect.py b/autotest/test_gridintersect.py index 0c4bd9ab03..c5a726f37a 100644 --- a/autotest/test_gridintersect.py +++ b/autotest/test_gridintersect.py @@ -1195,10 +1195,7 @@ def test_rasters(example_data_path): ws = str(example_data_path / "options") raster_name = "dem.img" - try: - rio = Raster.load(os.path.join(ws, "dem", raster_name)) - except: - return + rio = Raster.load(os.path.join(ws, "dem", raster_name)) ml = Modflow.load( "sagehen.nam", version="mfnwt", model_ws=os.path.join(ws, "sagehen") @@ -1254,14 +1251,12 @@ def test_rasters(example_data_path): # %% test raster sampling methods @pytest.mark.slow +@requires_pkg("rasterstats") def test_raster_sampling_methods(example_data_path): ws = str(example_data_path / "options") raster_name = "dem.img" - try: - rio = Raster.load(os.path.join(ws, "dem", raster_name)) - except: - return + rio = Raster.load(os.path.join(ws, "dem", raster_name)) ml = Modflow.load( "sagehen.nam", version="mfnwt", model_ws=os.path.join(ws, "sagehen") diff --git a/flopy/discretization/structuredgrid.py b/flopy/discretization/structuredgrid.py index 7e671ace84..9969f2f6fa 100644 --- a/flopy/discretization/structuredgrid.py +++ b/flopy/discretization/structuredgrid.py @@ -944,45 +944,68 @@ def get_cell_vertices(self, *args, **kwargs): def get_lrc(self, nodes): """ - Get layer, row, column from a list of zero based + Get layer, row, column from a list of zero-based MODFLOW node numbers. + Parameters + ---------- + nodes : int, list or array_like + Zero-based node number + Returns ------- - v : list of tuples containing the layer (k), row (i), + list + list of tuples containing the layer (k), row (i), and column (j) for each node in the input list + + Examples + -------- + >>> import flopy + >>> sg = flopy.discretization.StructuredGrid(nlay=20, nrow=30, ncol=40) + >>> sg.get_lrc(100) + [(0, 2, 20)] + >>> sg.get_lrc([100, 1000, 10_000]) + [(0, 2, 20), (0, 25, 0), (8, 10, 0)] """ - if not isinstance(nodes, list): + if isinstance(nodes, int): nodes = [nodes] - ncpl = self.ncpl - v = [] - for node in nodes: - k = int(np.floor(node / ncpl)) - ij = int((node) - (ncpl * k)) - i = int(np.floor(ij / self.__ncol)) - j = int(ij - (i * self.__ncol)) - - v.append((k, i, j)) - return v + shape = self.shape + if shape[0] is None: + shape = tuple(dim or 1 for dim in shape) + return list(zip(*np.unravel_index(nodes, shape))) def get_node(self, lrc_list): """ - Get node number from a list of zero based MODFLOW + Get node number from a list of zero-based MODFLOW layer, row, column tuples. + Parameters + ---------- + lrc_list : tuple of int or list of tuple of int + Zero-based layer, row, column tuples + Returns ------- - v : list of MODFLOW nodes for each layer (k), row (i), + list + list of MODFLOW nodes for each layer (k), row (i), and column (j) tuple in the input list + + Examples + -------- + >>> import flopy + >>> sg = flopy.discretization.StructuredGrid(nlay=20, nrow=30, ncol=40) + >>> sg.get_node((0, 2, 20)) + [100] + >>> sg.get_node([(0, 2, 20), (0, 25, 0), (8, 10, 0)]) + [100, 1000, 10000] """ if not isinstance(lrc_list, list): lrc_list = [lrc_list] - nrc = self.__nrow * self.__ncol - v = [] - for [k, i, j] in lrc_list: - node = int(((k) * nrc) + ((i) * self.__ncol) + j) - v.append(node) - return v + multi_index = tuple(np.array(lrc_list).T) + shape = self.shape + if shape[0] is None: + shape = tuple(dim or 1 for dim in shape) + return np.ravel_multi_index(multi_index, shape).tolist() def plot(self, **kwargs): """