refactor(get_lrc, get_node): use numpy methods, add examples (#1591)

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():

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")

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):
         """