Merge pull request #43 from phobson/cleanup-masking

Cleanup polygon masking

pygridtools/__init__.py

@@ -2,4 +2,5 @@
 from .core import *
 from . import iotools
 from . import viz
+from . import qa
 from . import testing

pygridtools/core.py

@@ -104,12 +104,12 @@ def yn(self):
 
     @property
     def xc(self):
-        """shortcut to x-coords of nodes"""
+        """shortcut to x-coords of cells/centroids"""
         return self.cells_x
 
     @property
     def yc(self):
-        """shortcut to y-coords of nodes"""
+        """shortcut to y-coords of cells/centroids"""
         return self.cells_y
 
     @property
@@ -185,12 +185,68 @@ def flipud(self):
         return self.transform(np.flipud)
 
     def merge(self, other, how='vert', where='+', shift=0):
-        """ Merge with another grid.
+        """ Merge with another grid using pygridtools.misc.padded_stack.
 
         Parameters
         ----------
         other : ModelGrid
             The other ModelGrid object.
+        how : optional string (default = 'vert')
+            The method through wich the arrays should be stacked.
+            `'Vert'` is analogous to `np.vstack`. `'Horiz'` maps to
+            `np.hstack`.
+        where : optional string (default = '+')
+            The placement of the arrays relative to each other. Keeping
+            in mind that the origin of an array's index is in the
+            upper-left corner, `'+'` indicates that the second array
+            will be placed at higher index relative to the first array.
+            Essentially:
+              - if how == 'vert'
+                - `'+'` -> `a` is above (higher index) `b`
+                - `'-'` -> `a` is below (lower index) `b`
+              - if how == 'horiz'
+                - `'+'` -> `a` is to the left of `b`
+                - `'-'` -> `a` is to the right of `b`
+            See the examples and pygridtools.misc.padded_stack for more
+            info.
+        shift : int (default = 0)
+            The number of indices the second array should be shifted in
+            axis other than the one being merged. In other words,
+            vertically stacked arrays can be shifted horizontally,
+            and horizontally stacked arrays can be shifted vertically.
+
+        Returns
+        -------
+        self (operates in-place)
+
+        Notes
+        -----
+        The ``cell_mask`` attribute is not automatically updated
+        following merge operates. See the Examples section on handling
+        this manually.
+
+        Examples
+        --------
+        >>> domain1 = pandas.DataFrame({
+            'x': [2, 5, 5, 2],
+            'y': [6, 6, 4, 4],
+            'beta': [1, 1, 1, 1]
+        })
+        >>> domain2 = pandas.DataFrame({
+            'x': [6, 11, 11, 5],
+            'y': [5, 5, 3, 3],
+            'beta': [1, 1, 1, 1]
+        })
+        >>> grid1 = pgt.makeGrid(domain=domain1, nx=6, ny=5, rawgrid=False)
+        >>> grid2 = pgt.makeGrid(domain=domain2, nx=8, ny=7, rawgrid=False)
+        >>> grid1.merge(grid2, how='horiz')
+        >>> # update the cell mask to include new NA points:
+        >>> grid1.cell_mask = np.ma.masked_invalid(grid1.xc).mask
+
+        See Also
+        --------
+        pygridtools.padded_stack
+
         """
 
         self.nodes_x = self.nodes_x.merge(other.nodes_x, how=how,
@@ -199,43 +255,72 @@ def merge(self, other, how='vert', where='+', shift=0):
                                           where=where, shift=shift)
         return self
 
-    def mask_cells_with_polygon(self, polyverts, use_cells=True,
-                                inside=True, triangles=False,
-                                min_nodes=3, inplace=True):
-        polyverts = np.asarray(polyverts)
-        if polyverts.ndim != 2:
-            raise ValueError('polyverts must be a 2D array, or a '
-                             'similar sequence')
-
-        if polyverts.shape[1] != 2:
-            raise ValueError('polyverts must be two columns of points')
-
-        if polyverts.shape[0] < 3:
-            raise ValueError('polyverts must contain at least 3 points')
-
-        if use_cells:
-            cells = self.as_coord_pairs(which='cells')
-            cell_mask = misc.points_inside_poly(
-                cells, polyverts
-            ).reshape(self.cell_shape)
+    def mask_cells_with_polygon(self, polyverts, use_centroids=True,
+                                min_nodes=3, inside=True,
+                                use_existing=True, triangles=False,
+                                inplace=True):
+
+        """ Create mask for the cells of the ModelGrid with a polygon.
+
+        Parameters
+        ----------
+        polyverts : sequence of a polygon's vertices
+            A sequence of x-y pairs for each vertex of the polygon.
+        use_centroids : bool (default = True)
+            When True, the cell centroid will be used to determine
+            whether the cell is "inside" the polygon. If False, the
+            nodes are used instead.
+        min_nodes : int (default = 3)
+            Only used when ``use_centroids`` is False. This is the
+            minimum number of nodes inside the polygon required to mark
+            the cell as "inside". Must be greater than 0, but no more
+            than 4.
+        inside : bool (default = True)
+            Toggles masking of cells either *inside* (True) or *outside*
+            (False) the polygon.
+        triangles : bool
+            Not yet implemented.
+        use_existing : bool (default = True)
+            When True, the newly computed mask is combined (via a
+            bit-wise `or` opteration) with the existing ``cell_mask``
+            attribute of the MdoelGrid.
+        inplace : bool (default = True):
+            If True, the ``cell_mask`` attribute of the ModelGrid is set
+            to the returned masked and None is returned. Otherwise, the
+            a new mask is returned the ``cell_mask`` attribute of the
+            ModelGrid is unchanged.
+
+        Returns
+        -------
+        cell_mask : np.array of bools
+            The final mask to be applied to the cells of the ModelGrid.
+
+        """
+
+        if triangles:
+            raise NotImplementedError("triangles are not yet implemented.")
+
+        if use_centroids:
+            cell_mask = misc.mask_with_polygon(self.xc, self.yc, polyverts, inside=inside)
         else:
-            nodes = self.as_coord_pairs(which='nodes')
-            _node_mask = misc.points_inside_poly(
-                nodes, polyverts
-            ).reshape(self.shape)
+            if min_nodes <= 0 or min_nodes > 4:
+                raise ValueError("`min_nodes` must be greater than 0 and no more than 4.")
 
+            _node_mask = misc.mask_with_polygon(self.xn, self.yn, polyverts, inside=inside).astype(int)
             cell_mask = (
                 _node_mask[1:, 1:] + _node_mask[:-1, :-1] +
                 _node_mask[:-1, 1:] + _node_mask[1:, :-1]
-            ) < min_nodes,
-        if not inside:
-            cell_mask = ~cell_mask
+            ) >= min_nodes
+            cell_mask = cell_mask.astype(bool)
+
+
+        if use_existing:
+            cell_mask = np.bitwise_or(self.cell_mask, cell_mask)
 
         if inplace:
-            self.cell_mask = np.bitwise_or(self.cell_mask, cell_mask)
+            self.cell_mask = cell_mask
 
-        else:
-            return cell_mask
+        return cell_mask
 
     def writeGEFDCControlFile(self, outputdir=None, filename='gefdc.inp',
                               bathyrows=0, title='test'):

pygridtools/misc.py

@@ -4,12 +4,9 @@
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.path as mpath
-import matplotlib.mlab as mlab
 import pandas
 
-
-def points_inside_poly(points, polyverts):
-    return mpath.Path(polyverts).contains_points(points)
+from pygridtools import qa
 
 
 def makePolyCoords(xarr, yarr, zpnt=None, triangles=False):
@@ -274,6 +271,39 @@ def padded_stack(a, b, how='vert', where='+', shift=0, padval=np.nan):
     return stacked
 
 
+def mask_with_polygon(x, y, polyverts, inside=True):
+    """ Mask x-y arrays inside or outside a polygon
+
+    Parameters
+    ----------
+    x, y : array-like
+        NxM arrays of x- and y-coordinates.
+    polyverts : sequence of a polygon's vertices
+        A sequence of x-y pairs for each vertex of the polygon.
+    inside : bool (default = True)
+        Toggles returning a mask *inside* or *outside* the polygon.
+
+    Returns
+    -------
+    mask : bool array
+        The NxM mask that can be applied to ``x`` and ``y``.
+
+    """
+
+    # validate input
+    polyverts = qa._validate_polygon(polyverts)
+    points = qa._validate_xy_array(x, y, as_pairs=True)
+
+    # compute the mask
+    mask = mpath.Path(polyverts).contains_points(points).reshape(x.shape)
+
+    # invert if we're masking things outside the polygon
+    if not inside:
+        mask = ~mask
+
+    return mask
+
+
 def make_gefdc_cells(node_mask, cell_mask=None, triangles=False):
     """ Take an array defining the nodes as wet (1) or dry (0) create
     the array of cell values needed for GEFDC.

pygridtools/qa.py

@@ -0,0 +1,44 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def _check_ax(ax):
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.figure
+
+    ax.set_aspect('equal')
+    return fig, ax
+
+
+def _validate_polygon(polyverts, min_points=3):
+    polyverts_array = np.asarray(polyverts)
+    if polyverts_array.ndim != 2:
+        raise ValueError('polyverts must be a 2D array, or '
+                         'similar sequence')
+
+    if polyverts_array.shape[1] != 2:
+        raise ValueError('polyverts must be two columns of points')
+
+    if polyverts_array.shape[0] < min_points:
+        raise ValueError('polyverts must contain at least {} points'.format(min_points))
+
+    return polyverts_array
+
+
+def _validate_xy_array(x, y, as_pairs=True):
+    x, y = np.asanyarray(x), np.asanyarray(y)
+    if x.shape != y.shape:
+        raise ValueError("x and y must have the same shape.")
+
+    if hasattr(x, 'mask') != hasattr(y, 'mask'):
+        raise ValueError("only 1 of x and y have masks. Must be both or neither.")
+
+    if hasattr(x, 'mask') and not np.all(x.mask == y.mask):
+        raise ValueError("x and y has different masks.")
+
+    if as_pairs:
+        return np.array(list(zip(x.flatten(), y.flatten())))
+    else:
+        return x, y

pygridtools/tests/test_core.py

@@ -251,6 +251,12 @@ def setup(self):
         self.mg = core.ModelGrid(self.xn, self.yn)
         self.g1 = core.ModelGrid(self.xn[:, :3], self.yn[:, :3])
         self.g2 = core.ModelGrid(self.xn[2:5, 3:], self.yn[2:5, 3:])
+        self.polyverts = [
+            (2.4, 0.9),
+            (3.6, 0.9),
+            (3.6, 2.4),
+            (2.4, 2.4),
+        ]
 
         self.template = 'pygridtools/tests/test_data/schema_template.shp'
         self.g1.template = self.template
@@ -486,6 +492,82 @@ def test_merge(self):
         nptest.assert_array_equal(g3.xn, g4.xn)
         nptest.assert_array_equal(g3.xc, g4.xc)
 
+    def test_mask_cells_with_polygon_inside_not_inplace(self):
+        orig_mask = self.mg.cell_mask.copy()
+        cell_mask = self.mg.mask_cells_with_polygon(self.polyverts, inplace=False, use_existing=False)
+        known_inside_mask = np.array([
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False,  True,  True, False],
+            [False, False, False,  True,  True, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False]
+        ], dtype=bool)
+        nptest.assert_array_equal(cell_mask, known_inside_mask)
+        nptest.assert_array_equal(self.mg.cell_mask, orig_mask)
+
+    def test_mask_cells_with_polygon_inside_use_existing(self):
+        self.mg.cell_mask = np.ma.masked_invalid(self.mg.xc).mask
+        cell_mask = self.mg.mask_cells_with_polygon(self.polyverts, use_existing=True)
+        known_inside_mask = np.array([
+            [False, False,  True,  True,  True,  True],
+            [False, False,  True,  True,  True,  True],
+            [False, False, False,  True,  True, False],
+            [False, False, False,  True,  True, False],
+            [False, False,  True,  True,  True,  True],
+            [False, False,  True,  True,  True,  True],
+            [False, False,  True,  True,  True,  True],
+            [False, False,  True,  True,  True,  True]
+        ], dtype=bool)
+        nptest.assert_array_equal(cell_mask, known_inside_mask)
+        nptest.assert_array_equal(cell_mask, self.mg.cell_mask)
+
+    def test_mask_cells_with_polygon_outside(self):
+        cell_mask = self.mg.mask_cells_with_polygon(self.polyverts, inside=False, use_existing=False)
+        known_outside_mask = np.array([
+            [ True,  True,  True,  True,  True, True],
+            [ True,  True,  True,  True,  True, True],
+            [ True,  True,  True, False, False, True],
+            [ True,  True,  True, False, False, True],
+            [ True,  True,  True,  True,  True, True],
+            [ True,  True,  True,  True,  True, True],
+            [ True,  True,  True,  True,  True, True],
+            [ True,  True,  True,  True,  True, True]
+        ], dtype=bool)
+        nptest.assert_array_equal(cell_mask, known_outside_mask)
+
+    def test_mask_cells_with_polygon_use_nodes(self):
+        cell_mask = self.mg.mask_cells_with_polygon(self.polyverts, use_centroids=False, use_existing=False)
+        known_node_mask = np.array([
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False,  True,  True, False],
+            [False, False, False,  True,  True, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False],
+            [False, False, False, False, False, False]
+        ], dtype=bool)
+        nptest.assert_array_equal(cell_mask, known_node_mask)
+
+    @nt.raises(ValueError)
+    def test_mask_cells_with_polygon_nodes_too_few_nodes(self):
+        cell_mask = self.mg.mask_cells_with_polygon(
+            self.polyverts, use_centroids=False, min_nodes=0
+        )
+
+    @nt.raises(ValueError)
+    def test_mask_cells_with_polygon_nodes_too_many_nodes(self):
+        cell_mask = self.mg.mask_cells_with_polygon(
+            self.polyverts, use_centroids=False, min_nodes=5
+        )
+
+    @nt.raises(NotImplementedError)
+    def test_mask_cells_with_polygon_triangles(self):
+        cell_mask = self.mg.mask_cells_with_polygon(self.polyverts, triangles=True)
+
     @nt.raises(ValueError)
     def test_to_shapefile_bad_geom(self):
         self.g1.to_shapefile('junk', geom='Line')
@@ -620,7 +702,6 @@ def test_writeGEFDCGridextFiles(self):
             known_filename
         )
 
-
 @image_comparison(
     baseline_images=[
         'test_ModelGrid_plots_basic',