Update(PlotCrossSection): updates for unstructured grid plotting (#1088)

flopy/discretization/grid.py

@@ -381,6 +381,116 @@ def xyzvertices(self):
     def cross_section_vertices(self):
         return self.xyzvertices[0], self.xyzvertices[1]
 
+    def cross_section_lay_ncpl_ncb(self, ncb):
+        """
+        Get PlotCrossSection compatible layers, ncpl, and ncb
+        variables
+
+        Parameters
+        ----------
+        ncb : int
+            number of confining beds
+
+        Returns
+        -------
+            tuple : (int, int, int) layers, ncpl, ncb
+        """
+        return self.nlay, self.ncpl, ncb
+
+    def cross_section_nodeskip(self, nlay, xypts):
+        """
+        Get a nodeskip list for PlotCrossSection. This is a correction
+        for UnstructuredGridPlotting
+
+        Parameters
+        ----------
+        nlay : int
+            nlay is nlay + ncb
+        xypts : dict
+            dictionary of node number and xyvertices of a cross-section
+
+        Returns
+        -------
+            list : n-dimensional list of nodes to not plot for each layer
+        """
+        return [[] for _ in range(nlay)]
+
+    def cross_section_adjust_indicies(self, k, cbcnt):
+        """
+        Method to get adjusted indicies by layer and confining bed
+        for PlotCrossSection plotting
+
+        Parameters
+        ----------
+        k : int
+            zero based layer number
+        cbcnt : int
+            confining bed counter
+
+        Returns
+        -------
+            tuple: (int, int, int) (adjusted layer, nodeskip layer, node
+            adjustment value based on number of confining beds and the layer)
+        """
+        adjnn = k * self.ncpl
+        ncbnn = adjnn - (cbcnt * self.ncpl)
+        return k + 1, k + 1, ncbnn
+
+    def cross_section_set_contour_arrays(
+        self, plotarray, xcenters, head, elev, projpts
+    ):
+        """
+        Method to set countour array centers for rare instances where
+        matplotlib contouring is prefered over trimesh plotting
+
+        Parameters
+        ----------
+        plotarray : np.ndarray
+            array of data for contouring
+        xcenters : np.ndarray
+            xcenters array
+        zcenters : np.ndarray
+            zcenters array
+        head : np.ndarray
+            head array to adjust cell centers location
+        elev : np.ndarray
+            cell elevation array
+        projpts : dict
+            dictionary of projected cross sectional vertices
+
+        Returns
+        -------
+            tuple: (np.ndarray, np.ndarray, np.ndarray, bool)
+            plotarray, xcenter array, ycenter array, and a boolean flag
+            for contouring
+        """
+        if self.nlay != 1:
+            return plotarray, xcenters, None, False
+        else:
+            zcenters = []
+            if isinstance(head, np.ndarray):
+                head = head.reshape(1, self.ncpl)
+                head = np.vstack((head, head))
+            else:
+                head = elev.reshape(2, self.ncpl)
+
+            elev = elev.reshape(2, self.ncpl)
+            for k, ev in enumerate(elev):
+                if k == 0:
+                    zc = [
+                        ev[i] if head[k][i] > ev[i] else head[k][i]
+                        for i in sorted(projpts)
+                    ]
+                else:
+                    zc = [ev[i] for i in sorted(projpts)]
+                zcenters.append(zc)
+
+            plotarray = np.vstack((plotarray, plotarray))
+            xcenters = np.vstack((xcenters, xcenters))
+            zcenters = np.array(zcenters)
+
+            return plotarray, xcenters, zcenters, True
+
     @property
     def map_polygons(self):
         """

flopy/discretization/unstructuredgrid.py

@@ -299,6 +299,102 @@ def xyzvertices(self):
         else:
             return self._cache_dict[cache_index].data_nocopy
 
+    def cross_section_lay_ncpl_ncb(self, ncb):
+        """
+        Get PlotCrossSection compatible layers, ncpl, and ncb
+        variables
+
+        Parameters
+        ----------
+        ncb : int
+            number of confining beds
+
+        Returns
+        -------
+            tuple : (int, int, int) layers, ncpl, ncb
+        """
+        return 1, self.nnodes, 0
+
+    def cross_section_nodeskip(self, nlay, xypts):
+        """
+        Get a nodeskip list for PlotCrossSection. This is a correction
+        for UnstructuredGridPlotting
+
+        Parameters
+        ----------
+        nlay : int
+            nlay is nlay + ncb
+        xypts : dict
+            dictionary of node number and xyvertices of a cross-section
+
+        Returns
+        -------
+            list : n-dimensional list of nodes to not plot for each layer
+        """
+        strt = 0
+        end = 0
+        nodeskip = []
+        for ncpl in self.ncpl:
+            end += ncpl
+            layskip = []
+            for nn, verts in xypts.items():
+                if strt <= nn < end:
+                    continue
+                else:
+                    layskip.append(nn)
+
+            strt += ncpl
+            nodeskip.append(layskip)
+
+        return nodeskip
+
+    def cross_section_adjust_indicies(self, k, cbcnt):
+        """
+        Method to get adjusted indicies by layer and confining bed
+        for PlotCrossSection plotting
+
+        Parameters
+        ----------
+        k : int
+            zero based model layer
+        cbcnt : int
+            confining bed counter
+
+        Returns
+        -------
+            tuple: (int, int, int) (adjusted layer, nodeskip layer, node
+            adjustment value based on number of confining beds and the layer)
+        """
+        return 1, k + 1, 0
+
+    def cross_section_set_contour_arrays(
+        self, plotarray, xcenters, head, elev, projpts
+    ):
+        """
+        Method to set countour array centers for rare instances where
+        matplotlib contouring is prefered over trimesh plotting
+
+        Parameters
+        ----------
+        plotarray : np.ndarray
+            array of data for contouring
+        xcenters : np.ndarray
+            xcenters array
+        head : np.ndarray
+            head array to adjust cell centers location
+        elev : np.ndarray
+            cell elevation array
+        projpts : dict
+            dictionary of projected cross sectional vertices
+
+        Returns
+        -------
+            tuple: (np.ndarray, np.ndarray, np.ndarray, bool)
+            plotarray, xcenter array, ycenter array, and a boolean flag
+            for contouring
+        """
+        return plotarray, xcenters, None, False
+
     @property
     def map_polygons(self):
         """

flopy/plot/crosssection.py

@@ -217,8 +217,12 @@ def __init__(
         else:
             self.active = np.ones(self.mg.nlay, dtype=int)
 
-        top = self.mg.top.reshape(1, self.mg.ncpl)
-        botm = self.mg.botm.reshape(self.mg.nlay + self.ncb, self.mg.ncpl)
+        self._nlay, self._ncpl, self.ncb = self.mg.cross_section_lay_ncpl_ncb(
+            self.ncb
+        )
+
+        top = self.mg.top.reshape(1, self._ncpl)
+        botm = self.mg.botm.reshape(self._nlay + self.ncb, self._ncpl)
 
         self.elev = np.concatenate((top, botm), axis=0)
 
@@ -390,7 +394,7 @@ def plot_surface(self, a, masked_values=None, **kwargs):
         if a.ndim > 1:
             a = np.ravel(a)
 
-        if a.size % self.mg.ncpl != 0:
+        if a.size % self._ncpl != 0:
             raise AssertionError("Array size must be a multiple of ncpl")
 
         if masked_values is not None:
@@ -520,30 +524,16 @@ def contour_array(self, a, masked_values=None, head=None, **kwargs):
         xcenters = self.xcenters
         plotarray = np.array([a[cell] for cell in sorted(self.projpts)])
 
-        if self.mg.nlay == 1 and not isinstance(self.mg.ncpl, np.ndarray):
-            zcenters = []
-            if isinstance(head, np.ndarray):
-                head = head.reshape(1, self.mg.ncpl)
-                head = np.vstack((head, head))
-            else:
-                head = self.elev.reshape(2, self.mg.ncpl)
-
-            elev = self.elev.reshape(2, self.mg.ncpl)
-            for k, ev in enumerate(elev):
-                if k == 0:
-                    zc = [
-                        ev[i] if head[k][i] > ev[i] else head[k][i]
-                        for i in sorted(self.projpts)
-                    ]
-                else:
-                    zc = [ev[i] for i in sorted(self.projpts)]
-                zcenters.append(zc)
-
-            plotarray = np.vstack((plotarray, plotarray))
-            xcenters = np.vstack((xcenters, xcenters))
-            zcenters = np.array(zcenters)
+        (
+            plotarray,
+            xcenters,
+            zcenters,
+            mplcontour,
+        ) = self.mg.cross_section_set_contour_arrays(
+            plotarray, xcenters, head, self.elev, self.projpts
+        )
 
-        else:
+        if not mplcontour:
             if isinstance(head, np.ndarray):
                 zcenters = self.set_zcentergrid(np.ravel(head))
             else:
@@ -595,7 +585,7 @@ def contour_array(self, a, masked_values=None, head=None, **kwargs):
             xcenters = xcenters[idx].flatten()
             zcenters = zcenters[idx].flatten()
 
-        if self.mg.nlay == 1 and not isinstance(self.mg.ncpl, np.ndarray):
+        if mplcontour:
             plotarray = np.ma.masked_array(plotarray, ismasked)
             contour_set = ax.contour(xcenters, zcenters, plotarray, **kwargs)
         else:
@@ -833,7 +823,7 @@ def plot_bc(
                     idx = mflist["node"]
 
         if len(self.mg.shape) != 3:
-            plotarray = np.zeros((self.mg.nlay, self.mg.ncpl), dtype=np.int)
+            plotarray = np.zeros((self._nlay, self._ncpl), dtype=np.int)
             plotarray[tuple(idx)] = 1
         else:
             plotarray = np.zeros(
@@ -953,7 +943,7 @@ def plot_vector(
         projpts = {
             key: value
             for key, value in self.projpts.items()
-            if (key // self.mg.ncpl) % kstep == 0
+            if (key // self._ncpl) % kstep == 0
         }
 
         # set x and z centers
@@ -1056,8 +1046,8 @@ def plot_specific_discharge(
             )
             spdis = spdis[-1]
 
-        ncpl = self.mg.ncpl
-        nlay = self.mg.nlay
+        ncpl = self._ncpl
+        nlay = self._nlay
 
         qx = np.zeros((nlay * ncpl))
         qy = np.zeros((nlay * ncpl))
@@ -1500,24 +1490,28 @@ def set_zpts(self, vs):
 
         projpts = {}
 
-        nlay = 1
-        if len(self.xvertices) != self.mg.nnodes:
-            nlay = self.mg.nlay + self.ncb
+        nlay = self.mg.nlay + self.ncb
+
+        nodeskip = self.mg.cross_section_nodeskip(nlay, self.xypts)
 
         cbcnt = 0
         for k in range(1, nlay + 1):
+
             if not self.active[k - 1]:
                 cbcnt += 1
                 continue
+
+            k, ns, ncbnn = self.mg.cross_section_adjust_indicies(k - 1, cbcnt)
             top = self.elev[k - 1, :]
             botm = self.elev[k, :]
-            adjnn = (k - 1) * self.mg.ncpl
-            ncbnn = adjnn - (cbcnt * self.mg.ncpl)
             d0 = 0
 
             # trap to split multipolygons
             xypts = []
             for nn, verts in self.xypts.items():
+                if nn in nodeskip[ns - 1]:
+                    continue
+
                 if len(verts) > 2:
                     i0 = 2
                     for ix in range(len(verts)):
@@ -1566,10 +1560,11 @@ def set_zpts(self, vs):
                     d0 += c
 
                 projpt = projt + projb
-                if nn + ncbnn not in projpts:
-                    projpts[nn + ncbnn] = projpt
+                node = nn + ncbnn
+                if node not in projpts:
+                    projpts[node] = projpt
                 else:
-                    projpts[nn + ncbnn] += projpt
+                    projpts[node] += projpt
 
         return projpts
 
@@ -1594,7 +1589,7 @@ def set_zcentergrid(self, vs, kstep=1):
         zcenters = [
             np.mean(np.array(v).T[1])
             for i, v in sorted(verts.items())
-            if (i // self.mg.ncpl) % kstep == 0
+            if (i // self._ncpl) % kstep == 0
         ]
         return zcenters
 

flopy/plot/plotutil.py

@@ -1889,6 +1889,7 @@ def irregular_shape_patch(xverts, yverts):
         adj_xverts = []
         for xv in xverts:
             if len(xv) < max_verts:
+                xv = list(xv)
                 n = max_verts - len(xv)
                 adj_xverts.append(xv + [xv[-1]] * n)
             else:
@@ -1897,6 +1898,7 @@ def irregular_shape_patch(xverts, yverts):
         adj_yverts = []
         for yv in yverts:
             if len(yv) < max_verts:
+                yv = list(yv)
                 n = max_verts - len(yv)
                 adj_yverts.append(yv + [yv[-1]] * n)
             else: