From 00757a4dc7ea03ba3582242ba4d2590f407c0d39 Mon Sep 17 00:00:00 2001 From: w-bonelli Date: Thu, 10 Nov 2022 15:59:18 -0500 Subject: [PATCH] feat(contours): use standard matplotlib contours for StructuredGrid map view plots (#1615) --- autotest/test_export.py | 11 ++-- autotest/test_plot.py | 10 +-- flopy/plot/map.py | 141 +++++++++++++++++++++------------------- 3 files changed, 87 insertions(+), 75 deletions(-) diff --git a/autotest/test_export.py b/autotest/test_export.py index 6709e7d43..e004f87e3 100644 --- a/autotest/test_export.py +++ b/autotest/test_export.py @@ -631,10 +631,9 @@ def test_export_contourf(tmpdir, example_data_path): with Reader(filename) as r: shapes = r.shapes() - if len(shapes) != 65: - raise AssertionError( - "multipolygons were skipped in contourf routine" - ) + # expect 65 with standard mpl contours (structured grids), 86 with tricontours + assert len(shapes) >= 65, "multipolygons were skipped in contourf routine" + @pytest.mark.mf6 @@ -662,7 +661,9 @@ def test_export_contours(tmpdir, example_data_path): with Reader(filename) as r: shapes = r.shapes() - assert len(shapes) == 65 + # expect 65 with standard mpl contours (structured grids), 86 with tricontours + assert len(shapes) >= 65 + @pytest.mark.mf6 diff --git a/autotest/test_plot.py b/autotest/test_plot.py index 40e89ffc4..f742ab80a 100644 --- a/autotest/test_plot.py +++ b/autotest/test_plot.py @@ -242,7 +242,7 @@ def test_cross_section_bc_UZF_3lay(example_data_path): ), f"Unexpected collection type: {type(col)}" -def test_map_view_tricontour_nan(): +def test_map_view_contour(tmpdir): arr = np.random.rand(10, 10) * 100 arr[-1, :] = np.nan delc = np.array([10] * 10, dtype=float) @@ -272,10 +272,12 @@ def test_map_view_tricontour_nan(): pmv = PlotMapView(modelgrid=grid, layer=0) contours = pmv.contour_array(a=arr) + plt.savefig(tmpdir / "map_view_contour.png") - for ix, lev in enumerate(contours.levels): - if not np.allclose(lev, levels[ix]): - raise AssertionError("TriContour NaN catch Failed") + # if we ever revert from standard contours to tricontours, restore this nan check + # for ix, lev in enumerate(contours.levels): + # if not np.allclose(lev, levels[ix]): + # raise AssertionError("TriContour NaN catch Failed") @pytest.mark.mf6 diff --git a/flopy/plot/map.py b/flopy/plot/map.py index 9d16d51b9..e6b48f81e 100644 --- a/flopy/plot/map.py +++ b/flopy/plot/map.py @@ -175,56 +175,27 @@ def contour_array(self, a, masked_values=None, **kwargs): """ import matplotlib.tri as tri + # coerce array to ndarray of floats a = np.copy(a) if not isinstance(a, np.ndarray): a = np.array(a) - a = a.astype(float) + # Use the model grid to pass back an array of the correct shape plotarray = self.mg.get_plottable_layer_array(a, self.layer) - # work around for tri-contour ignore vmin & vmax - # necessary block for tri-contour NaN issue - if "levels" not in kwargs: - vmin = kwargs.pop("vmin", np.nanmin(plotarray)) - vmax = kwargs.pop("vmax", np.nanmax(plotarray)) - levels = np.linspace(vmin, vmax, 7) - kwargs["levels"] = levels - - # workaround for tri-contour nan issue - # use -2**31 to allow for 32 bit int arrays - plotarray[np.isnan(plotarray)] = -(2**31) - if masked_values is None: - masked_values = [-(2**31)] - else: - masked_values = list(masked_values) - if -(2**31) not in masked_values: - masked_values.append(-(2**31)) - - ismasked = None - if masked_values is not None: - self._masked_values.extend(list(masked_values)) - - for mval in self._masked_values: - if ismasked is None: - ismasked = np.isclose(plotarray, mval) - else: - t = np.isclose(plotarray, mval) - ismasked += t + # Get vertices for the selected layer + xcentergrid = self.mg.get_xcellcenters_for_layer(self.layer) + ycentergrid = self.mg.get_ycellcenters_for_layer(self.layer) ax = kwargs.pop("ax", self.ax) - - if "colors" in kwargs.keys(): - if "cmap" in kwargs.keys(): - kwargs.pop("cmap") - filled = kwargs.pop("filled", False) plot_triplot = kwargs.pop("plot_triplot", False) tri_mask = kwargs.pop("tri_mask", False) - # Get vertices for the selected layer - xcentergrid = self.mg.get_xcellcenters_for_layer(self.layer) - ycentergrid = self.mg.get_ycellcenters_for_layer(self.layer) + if "colors" in kwargs.keys(): + if "cmap" in kwargs.keys(): + kwargs.pop("cmap") if "extent" in kwargs: extent = kwargs.pop("extent") @@ -239,39 +210,77 @@ def contour_array(self, a, masked_values=None, **kwargs): xcentergrid = xcentergrid[idx] ycentergrid = ycentergrid[idx] - plotarray = plotarray.flatten() - xcentergrid = xcentergrid.flatten() - ycentergrid = ycentergrid.flatten() - triang = tri.Triangulation(xcentergrid, ycentergrid) - analyze = tri.TriAnalyzer(triang) - mask = analyze.get_flat_tri_mask(rescale=False) - - # mask out holes, optional??? - if tri_mask: - triangles = triang.triangles - for i in range(2): - for ix, nodes in enumerate(triangles): - neighbors = self.mg.neighbors(nodes[i], as_nodes=True) - isin = np.isin(nodes[i + 1 :], neighbors) - if not np.alltrue(isin): - mask[ix] = True - - if ismasked is not None: - ismasked = ismasked.flatten() - mask2 = np.any( - np.where(ismasked[triang.triangles], True, False), axis=1 + # use standard contours for structured grid, otherwise tricontours + if self.mg.grid_type == "structured": + contour_set = ( + ax.contourf(xcentergrid, ycentergrid, plotarray, **kwargs) + if filled + else ax.contour(xcentergrid, ycentergrid, plotarray, **kwargs) ) - mask[mask2] = True + else: + # work around for tri-contour ignore vmin & vmax + # necessary block for tri-contour NaN issue + if "levels" not in kwargs: + vmin = kwargs.pop("vmin", np.nanmin(plotarray)) + vmax = kwargs.pop("vmax", np.nanmax(plotarray)) + levels = np.linspace(vmin, vmax, 7) + kwargs["levels"] = levels + + # workaround for tri-contour nan issue + # use -2**31 to allow for 32 bit int arrays + plotarray[np.isnan(plotarray)] = -(2**31) + if masked_values is None: + masked_values = [-(2**31)] + else: + masked_values = list(masked_values) + if -(2**31) not in masked_values: + masked_values.append(-(2**31)) - triang.set_mask(mask) + ismasked = None + if masked_values is not None: + self._masked_values.extend(list(masked_values)) - if filled: - contour_set = ax.tricontourf(triang, plotarray, **kwargs) - else: - contour_set = ax.tricontour(triang, plotarray, **kwargs) + for mval in self._masked_values: + if ismasked is None: + ismasked = np.isclose(plotarray, mval) + else: + t = np.isclose(plotarray, mval) + ismasked += t + + plotarray = plotarray.flatten() + xcentergrid = xcentergrid.flatten() + ycentergrid = ycentergrid.flatten() + triang = tri.Triangulation(xcentergrid, ycentergrid) + analyze = tri.TriAnalyzer(triang) + mask = analyze.get_flat_tri_mask(rescale=False) + + # mask out holes, optional??? + if tri_mask: + triangles = triang.triangles + for i in range(2): + for ix, nodes in enumerate(triangles): + neighbors = self.mg.neighbors(nodes[i], as_nodes=True) + isin = np.isin(nodes[i + 1 :], neighbors) + if not np.alltrue(isin): + mask[ix] = True + + if ismasked is not None: + ismasked = ismasked.flatten() + mask2 = np.any( + np.where(ismasked[triang.triangles], True, False), axis=1 + ) + mask[mask2] = True + + triang.set_mask(mask) + + contour_set = ( + ax.tricontourf(triang, plotarray.flatten(), **kwargs) + if filled + else ax.tricontour(triang, plotarray.flatten(), **kwargs) + ) - if plot_triplot: - ax.triplot(triang, color="black", marker="o", lw=0.75) + if plot_triplot: + ax.triplot(triang, color="black", marker="o", lw=0.75) ax.set_xlim(self.extent[0], self.extent[1]) ax.set_ylim(self.extent[2], self.extent[3])