diff --git a/autotest/t007_test.py b/autotest/t007_test.py index 7d05947b50..10e6caaec0 100644 --- a/autotest/t007_test.py +++ b/autotest/t007_test.py @@ -406,11 +406,11 @@ def test_mt_modelgrid(): mt = flopy.mt3d.Mt3dms(modelname='test_mt', modflowmodel=ml, model_ws=ml.model_ws, verbose=True) btn = flopy.mt3d.Mt3dBtn(mt, icbund=ml.bas6.ibound.array) - + # reload swt swt = flopy.seawat.Seawat(modelname='test_swt', modflowmodel=ml, mt3dmodel=mt, model_ws=ml.model_ws, verbose=True) - + assert \ ml.modelgrid.xoffset == mt.modelgrid.xoffset == swt.modelgrid.xoffset assert \ @@ -903,6 +903,42 @@ def check_vertices(): plt.close() +def test_tricontour_NaN(): + from flopy.plot import PlotMapView + import numpy as np + from flopy.discretization import StructuredGrid + + arr = np.random.rand(10, 10) * 100 + arr[-1, :] = np.nan + delc = np.array([10] * 10, dtype=float) + delr = np.array([8] * 10, dtype=float) + top = np.ones((10, 10), dtype=float) + botm = np.ones((3, 10, 10), dtype=float) + botm[0] = 0.75 + botm[1] = 0.5 + botm[2] = 0.25 + idomain = np.ones((3, 10, 10)) + idomain[0, 0, :] = 0 + vmin = np.nanmin(arr) + vmax = np.nanmax(arr) + levels = np.linspace(vmin, vmax, 7) + + grid = StructuredGrid(delc=delc, + delr=delr, + top=top, + botm=botm, + idomain=idomain, + lenuni=1, + nlay=3, nrow=10, ncol=10) + + pmv = PlotMapView(modelgrid=grid, layer=0) + contours = pmv.contour_array(a=arr) + + for ix, lev in enumerate(contours.levels): + if not np.allclose(lev, levels[ix]): + raise AssertionError("TriContour NaN catch Failed") + + def test_get_vertices(): from flopy.utils.reference import SpatialReference from flopy.discretization import StructuredGrid @@ -1191,7 +1227,7 @@ def test_export_array_contours(): # build_sfr_netcdf() # test_mg() # test_mbase_modelgrid() - test_mt_modelgrid() + # test_mt_modelgrid() # test_rotation() # test_model_dot_plot() # test_vertex_model_dot_plot() @@ -1212,5 +1248,6 @@ def test_export_array_contours(): #test_wkt_parse() #test_get_rc_from_node_coordinates() # test_export_array() - # test_export_array_contours() + test_export_array_contours() + test_tricontour_NaN() pass diff --git a/flopy/plot/map.py b/flopy/plot/map.py index e325445e64..43e580af25 100644 --- a/flopy/plot/map.py +++ b/flopy/plot/map.py @@ -125,7 +125,13 @@ def plot_array(self, a, masked_values=None, **kwargs): raise Exception('Array must be of dimension 1, 2, or 3') elif self.mg.grid_type == "vertex": - if a.ndim == 2: + if a.ndim == 3: + if a.shape[0] == 1: + a = np.squeeze(a, axis=1) + plotarray = a[self.layer, :] + else: + raise Exception("Array must be of dimension 1 or 2") + elif a.ndim == 2: plotarray = a[self.layer, :] elif a.ndim == 1: plotarray = a @@ -215,6 +221,7 @@ def contour_array(self, a, masked_values=None, **kwargs): err_msg = "Matplotlib must be updated to use contour_array" raise ImportError(err_msg) + a = np.copy(a) if not isinstance(a, np.ndarray): a = np.array(a) @@ -232,7 +239,13 @@ def contour_array(self, a, masked_values=None, **kwargs): raise Exception('Array must be of dimension 1, 2 or 3') elif self.mg.grid_type == "vertex": - if a.ndim == 2: + if a.ndim == 3: + if a.shape[0] == 1: + a = np.squeeze(a, axis=1) + plotarray = a[self.layer, :] + else: + raise Exception("Array must be of dimension 1 or 2") + elif a.ndim == 2: plotarray = a[self.layer, :] elif a.ndim == 1: plotarray = a @@ -242,9 +255,39 @@ def contour_array(self, a, masked_values=None, **kwargs): else: plotarray = a + # work around for tri-contour ignore vmin & vmax + # necessary block for tri-contour NaN issue + if "levels" not in kwargs: + if "vmin" not in kwargs: + vmin = np.nanmin(plotarray) + else: + vmin = kwargs.pop("vmin") + if "vmax" not in kwargs: + vmax = np.nanmax(plotarray) + else: + vmax = kwargs.pop('vmax') + + 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: for mval in masked_values: - plotarray = np.ma.masked_equal(plotarray, mval) + if ismasked is None: + ismasked = np.equal(plotarray, mval) + else: + t = np.equal(plotarray, mval) + ismasked += t if 'ax' in kwargs: ax = kwargs.pop('ax') @@ -276,16 +319,11 @@ def contour_array(self, a, masked_values=None, **kwargs): ycentergrid = ycentergrid.flatten() triang = tri.Triangulation(xcentergrid, ycentergrid) - mask = None - try: - amask = plotarray.mask - mask = [False for i in range(triang.triangles.shape[0])] - for ipos, (n0, n1, n2) in enumerate(triang.triangles): - if amask[n0] or amask[n1] or amask[n2]: - mask[ipos] = True + if ismasked is not None: + ismasked = ismasked.flatten() + mask = np.any(np.where(ismasked[triang.triangles], + True, False), axis=1) triang.set_mask(mask) - except (AttributeError, IndexError): - pass contour_set = ax.tricontour(triang, plotarray, **kwargs) diff --git a/flopy/plot/plotutil.py b/flopy/plot/plotutil.py index 379cecbb73..efbaf4b7a8 100644 --- a/flopy/plot/plotutil.py +++ b/flopy/plot/plotutil.py @@ -1833,6 +1833,56 @@ def line_intersect_grid(ptsin, xgrid, ygrid): return vdict + @staticmethod + def irregular_shape_patch(xverts, yverts): + """ + Patch for vertex cross section plotting when + we have an irregular shape type throughout the + model grid or multiple shape types. + + Parameters + ---------- + xverts : list + xvertices + yverts : list + yvertices + + Returns + ------- + xverts, yverts as np.ndarray + + """ + max_verts = 0 + + for xv in xverts: + if len(xv) > max_verts: + max_verts = len(xv) + + for yv in yverts: + if len(yv) > max_verts: + max_verts = len(yv) + + adj_xverts = [] + for xv in xverts: + if len(xv) < max_verts: + n = max_verts - len(xv) + adj_xverts.append(xv + [xv[-1]] * n) + else: + adj_xverts.append(xv) + + adj_yverts = [] + for yv in yverts: + if len(yv) < max_verts: + n = max_verts - len(yv) + adj_yverts.append(yv + [yv[-1]] * n) + else: + adj_yverts.append(yv) + + xverts = np.array(adj_xverts) + yverts = np.array(adj_yverts) + + return xverts, yverts + @staticmethod def arctan2(verts): """ diff --git a/flopy/plot/vcrosssection.py b/flopy/plot/vcrosssection.py index 494dad6b68..9efc6b9f9f 100644 --- a/flopy/plot/vcrosssection.py +++ b/flopy/plot/vcrosssection.py @@ -90,11 +90,22 @@ def __init__(self, ax=None, model=None, modelgrid=None, self.mg.xoffset, self.mg.yoffset, self.mg.angrot_radians, inverse=True) - self.xvertices, self.yvertices = \ - geometry.transform(self.mg.xvertices, - self.mg.yvertices, - self.mg.xoffset, self.mg.yoffset, - self.mg.angrot_radians, inverse=True) + try: + self.xvertices, self.yvertices = \ + geometry.transform(self.mg.xvertices, + self.mg.yvertices, + self.mg.xoffset, self.mg.yoffset, + self.mg.angrot_radians, inverse=True) + except ValueError: + # irregular shapes in vertex grid ie. squares and triangles + xverts, yverts = plotutil.UnstructuredPlotUtilities.\ + irregular_shape_patch(self.mg.xvertices, self.mg.yvertices) + + self.xvertices, self.yvertices = \ + geometry.transform(xverts, yverts, + self.mg.xoffset, + self.mg.yoffset, + self.mg.angrot_radians, inverse=True) pts = [(xt, yt) for xt, yt in zip(xp, yp)] self.pts = np.array(pts) @@ -430,9 +441,38 @@ def contour_array(self, a, masked_values=None, head=None, **kwargs): plotarray = np.array([a[cell] for cell in sorted(self.projpts)]) + # work around for tri-contour ignore vmin & vmax + # necessary for the tri-contour NaN issue fix + if "levels" not in kwargs: + if "vmin" not in kwargs: + vmin = np.nanmin(plotarray) + else: + vmin = kwargs.pop("vmin") + if "vmax" not in kwargs: + vmax = np.nanmax(plotarray) + else: + vmax = kwargs.pop('vmax') + + levels = np.linspace(vmin, vmax, 7) + kwargs['levels'] = levels + + # workaround for tri-contour nan issue + 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: for mval in masked_values: - plotarray = np.ma.masked_equal(plotarray, mval) + if ismasked is None: + ismasked = np.equal(plotarray, mval) + else: + t = np.equal(plotarray, mval) + ismasked += t if isinstance(head, np.ndarray): zcenters = self.set_zcentergrid(np.ravel(head)) @@ -457,15 +497,11 @@ def contour_array(self, a, masked_values=None, head=None, **kwargs): triang = tri.Triangulation(xcenters, zcenters) - try: - amask = plotarray.mask - mask = [False for _ in range(triang.triangles.shape[0])] - for ipos, (n0, n1, n2) in enumerate(triang.triangles): - if amask[n0] or amask[n1] or amask[n2]: - mask[ipos] = True + if ismasked is not None: + ismasked = ismasked.flatten() + mask = np.any(np.where(ismasked[triang.triangles], + True, False), axis=1) triang.set_mask(mask) - except (AttributeError, IndexError): - pass contour_set = ax.tricontour(triang, plotarray, **kwargs) diff --git a/flopy/utils/geometry.py b/flopy/utils/geometry.py index 87c112a6cd..366e6c99ec 100644 --- a/flopy/utils/geometry.py +++ b/flopy/utils/geometry.py @@ -360,9 +360,9 @@ def transform(x, y, xoff, yoff, angrot_radians, """ if isinstance(x, list): - x = np.array(x) + x = np.array(x, dtype=float) if isinstance(y, list): - y = np.array(y) + y = np.array(y, dtype=float) if not np.isscalar(x): x, y = x.copy(), y.copy()