Update mask functions with inverse arg, add rf_mask_by_values function

Signed-off-by: Jason T. Brown <[email protected]>

Author: vpipkt

core/src/main/scala/org/locationtech/rasterframes/RasterFunctions.scala

@@ -292,12 +292,26 @@ trait RasterFunctions {
   }
 
   /** Where the rf_mask tile contains NODATA, replace values in the source tile with NODATA */
-  def rf_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] =
-    Mask.MaskByDefined(sourceTile, maskTile)
+  def rf_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] = rf_mask(sourceTile, maskTile, false)
+
+  /** Where the rf_mask tile contains NODATA, replace values in the source tile with NODATA */
+  def rf_mask(sourceTile: Column, maskTile: Column, inverse: Boolean=false): TypedColumn[Any, Tile] =
+    if(!inverse) Mask.MaskByDefined(sourceTile, maskTile)
+    else Mask.InverseMaskByDefined(sourceTile, maskTile)
 
   /** Where the `maskTile` equals `maskValue`, replace values in the source tile with `NoData` */
-  def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Column): TypedColumn[Any, Tile] =
-    Mask.MaskByValue(sourceTile, maskTile, maskValue)
+  def rf_mask_by_value(sourceTile: Column, maskTile: Column, maskValue: Column, inverse: Boolean=false): TypedColumn[Any, Tile] =
+    if (!inverse) Mask.MaskByValue(sourceTile, maskTile, maskValue)
+    else Mask.InverseMaskByValue(sourceTile, maskTile, maskValue)
+
+  /** Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
+       list, replace the value with NODATA.
+       If `inverse` is True, the cells in `mask_tile` that are not in `mask_values` list become NODATA */
+  def rf_mask_by_values(sourceTile: Column, maskTile: Column, maskValues: Column, inverse: Boolean=false): TypedColumn[Any, Tile] =
+    if (!inverse)
+      Mask.MaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(1))
+    else
+      Mask.InverseMaskByValue(sourceTile, rf_local_is_in(maskTile, maskValues), lit(0))
 
   /** Where the `maskTile` does **not** contain `NoData`, replace values in the source tile with `NoData` */
   def rf_inverse_mask(sourceTile: Column, maskTile: Column): TypedColumn[Any, Tile] =

docs/src/main/paradox/reference.md

@@ -215,27 +215,45 @@ Masking is a raster operation that sets specific cells to NoData based on the va
 
 ### rf_mask
 
-    Tile rf_mask(Tile tile, Tile mask)
+    Tile rf_mask(Tile tile, Tile mask, bool inverse)
 
 Where the `mask` contains NoData, replace values in the `tile` with NoData.
 
 Returned `tile` cell type will be coerced to one supporting NoData if it does not already.
 
+`inverse` is a literal not a Column. If `inverse` is true, return the `tile` with NoData in locations where the `mask` _does not_ contain NoData. Equivalent to @ref:[`rf_inverse_mask`](reference.md#rf-inverse-mask).
+
 See also @ref:[`rf_rasterize`](reference.md#rf-rasterize).
 
+### rf_mask_by_value
+
+    Tile rf_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value, bool inverse)
+
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is equal to `mask_value`.
+
+`inverse` is a literal not a Column. If `inverse` is true, return the `data_tile` with NoData in locations where the `mask_tile` value is _not equal_ to `mask_value`. Equivalent to @ref:[`rf_inverse_mask_by_value`](reference.md#rf-inverse-mask-by-value).
+
+### rf_mask_by_values
+
+    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, Array mask_values, bool inverse)
+    Tile rf_mask_by_values(Tile data_tile, Tile mask_tile, seq mask_values, bool inverse)
+
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is in the `mask_values` Array or list. `mask_values` can be a [`pyspark.sql.ArrayType`][Array] or a `list`.
+
+`inverse` is a literal not a Column. If it is True, the `data_tile` cells are set to NoData where the `mask_tile` cells are __not__ in `mask_values`.
 
 ### rf_inverse_mask
 
     Tile rf_inverse_mask(Tile tile, Tile mask)
 
 Where the `mask` _does not_ contain NoData, replace values in `tile` with NoData.
 
-### rf_mask_by_value
 
-    Tile rf_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value)
+### rf_inverse_mask_by_value
 
-Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is equal to `mask_value`.
+    Tile rf_inverse_mask_by_value(Tile data_tile, Tile mask_tile, Int mask_value)
 
+Generate a `tile` with the values from `data_tile`, with NoData in cells where the `mask_tile` is not equal to `mask_value`. In other words, only keep `data_tile` cells in locations where the `mask_tile` is equal to `mask_value`.
 
 ### rf_is_no_data_tile
 

docs/src/main/paradox/release-notes.md

@@ -6,7 +6,7 @@
 
 * _Breaking_ (potentially): removed `GeoTiffCollectionRelation` due to usage limitation and overlap with `RasterSourceDataSource` functionality.
 * Upgraded to Spark 2.4.4
- * Add `rf_local_is_in` raster function
+ * Add `rf_mask_by_values` and `rf_local_is_in` raster functions; added optional `inverse` argument to `rf_mask` functions
 
 ### 0.8.3
 

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -25,6 +25,7 @@
 """
 from __future__ import absolute_import
 from pyspark.sql.column import Column, _to_java_column
+from pyspark.sql.functions import lit
 from .rf_context import RFContext
 from .rf_types import CellType
 
@@ -137,20 +138,6 @@ def rf_explode_tiles_sample(sample_frac, seed, *tile_cols):
     return Column(jfcn(sample_frac, seed, RFContext.active().list_to_seq(jcols)))
 
 
-def rf_mask_by_value(data_tile, mask_tile, mask_value):
-    """Generate a tile with the values from the data tile, but where cells in the masking tile contain the masking
-    value, replace the data value with NODATA. """
-    jfcn = RFContext.active().lookup('rf_mask_by_value')
-    return Column(jfcn(_to_java_column(data_tile), _to_java_column(mask_tile), _to_java_column(mask_value)))
-
-
-def rf_inverse_mask_by_value(data_tile, mask_tile, mask_value):
-    """Generate a tile with the values from the data tile, but where cells in the masking tile do not contain the
-    masking value, replace the data value with NODATA. """
-    jfcn = RFContext.active().lookup('rf_inverse_mask_by_value')
-    return Column(jfcn(_to_java_column(data_tile), _to_java_column(mask_tile), _to_java_column(mask_value)))
-
-
 def _apply_scalar_to_tile(name, tile_col, scalar):
     jfcn = RFContext.active().lookup(name)
     return Column(jfcn(_to_java_column(tile_col), scalar))
@@ -270,14 +257,16 @@ def rf_local_data(tile_col):
     """Return a tile with zeros where the input is NoData, otherwise one."""
     return _apply_column_function('rf_local_data', tile_col)
 
+
 def rf_local_is_in(tile_col, array):
     """Return a tile with cell values of 1 where the `tile_col` cell is in the provided array."""
-    from pyspark.sql.functions import array as sql_array, lit
+    from pyspark.sql.functions import array as sql_array
     if isinstance(array, list):
         array = sql_array([lit(v) for v in array])
 
     return _apply_column_function('rf_local_is_in', tile_col, array)
 
+
 def _apply_column_function(name, *args):
     jfcn = RFContext.active().lookup(name)
     jcols = [_to_java_column(arg) for arg in args]
@@ -459,16 +448,54 @@ def rf_agg_local_stats(tile_col):
     return _apply_column_function('rf_agg_local_stats', tile_col)
 
 
-def rf_mask(src_tile_col, mask_tile_col):
-    """Where the rf_mask (second) tile contains NODATA, replace values in the source (first) tile with NODATA."""
-    return _apply_column_function('rf_mask', src_tile_col, mask_tile_col)
+def rf_mask(src_tile_col, mask_tile_col, inverse=False):
+    """Where the rf_mask (second) tile contains NODATA, replace values in the source (first) tile with NODATA.
+       If `inverse` is true, replaces values in the source tile with NODATA where the mask tile contains valid data.
+    """
+    if not inverse:
+        return _apply_column_function('rf_mask', src_tile_col, mask_tile_col)
+    else:
+        rf_inverse_mask(src_tile_col, mask_tile_col)
 
 
 def rf_inverse_mask(src_tile_col, mask_tile_col):
-    """Where the rf_mask (second) tile DOES NOT contain NODATA, replace values in the source (first) tile with NODATA."""
+    """Where the rf_mask (second) tile DOES NOT contain NODATA, replace values in the source
+       (first) tile with NODATA."""
     return _apply_column_function('rf_inverse_mask', src_tile_col, mask_tile_col)
 
 
+def rf_mask_by_value(data_tile, mask_tile, mask_value, inverse=False):
+    """Generate a tile with the values from the data tile, but where cells in the masking tile contain the masking
+    value, replace the data value with NODATA. """
+    if isinstance(mask_value, (int, float)):
+        mask_value = lit(mask_value)
+    jfcn = RFContext.active().lookup('rf_mask_by_value')
+
+    return Column(jfcn(_to_java_column(data_tile), _to_java_column(mask_tile), _to_java_column(mask_value), inverse))
+
+
+def rf_mask_by_values(data_tile, mask_tile, mask_values, inverse=False):
+    """Generate a tile with the values from `data_tile`, but where cells in the `mask_tile` are in the `mask_values`
+       list, replace the value with NODATA.
+       If `inverse` is True, the cells in `mask_tile` that are not in `mask_values` list become NODATA
+    """
+    from pyspark.sql.functions import array as sql_array
+    if isinstance(mask_values, list):
+        mask_values = sql_array([lit(v) for v in mask_values])
+
+    jfcn = RFContext.active().lookup('rf_mask_by_values')
+    col_args = [_to_java_column(c) for c in [data_tile, mask_tile, mask_values]]
+    return Column(jfcn(*col_args, inverse))
+
+
+def rf_inverse_mask_by_value(data_tile, mask_tile, mask_value):
+    """Generate a tile with the values from the data tile, but where cells in the masking tile do not contain the
+    masking value, replace the data value with NODATA. """
+    if isinstance(mask_value, (int, float)):
+        mask_value = lit(mask_value)
+    return _apply_column_function('rf_inverse_mask_by_value', data_tile, mask_tile, mask_value)
+
+
 def rf_local_less(left_tile_col, right_tile_col):
     """Cellwise less than comparison between two tiles"""
     return _apply_column_function('rf_local_less', left_tile_col, right_tile_col)

pyrasterframes/src/main/python/tests/RasterFunctionsTests.py

@@ -24,6 +24,8 @@
 from pyspark import Row
 from pyspark.sql.functions import *
 
+import numpy as np
+from numpy.testing import assert_equal
 
 from . import TestEnvironment
 
@@ -103,7 +105,6 @@ def test_agg_mean(self):
     def test_agg_local_mean(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile
-        import numpy as np
 
         # this is really testing the nodata propagation in the agg  local summation
         ct = CellType.int8().with_no_data_value(4)
@@ -221,20 +222,51 @@ def test_mask_by_value(self):
                                      rf_local_greater_int(self.rf.tile, 25000),
                                      "uint8"),
                                  lit(mask_value)).alias('mask'))
-        rf2 = rf1.select(rf1.tile, rf_mask_by_value(rf1.tile, rf1.mask, lit(mask_value)).alias('masked'))
+        rf2 = rf1.select(rf1.tile, rf_mask_by_value(rf1.tile, rf1.mask, lit(mask_value), False).alias('masked'))
         result = rf2.agg(rf_agg_no_data_cells(rf2.tile) < rf_agg_no_data_cells(rf2.masked)) \
             .collect()[0][0]
         self.assertTrue(result)
 
-        rf3 = rf1.select(rf1.tile, rf_inverse_mask_by_value(rf1.tile, rf1.mask, lit(mask_value)).alias('masked'))
-        result = rf3.agg(rf_agg_no_data_cells(rf3.tile) < rf_agg_no_data_cells(rf3.masked)) \
-            .collect()[0][0]
-        self.assertTrue(result)
+        # note supplying a `int` here, not a column to mask value
+        rf3 = rf1.select(
+            rf1.tile,
+            rf_inverse_mask_by_value(rf1.tile, rf1.mask, mask_value).alias('masked'),
+            rf_mask_by_value(rf1.tile, rf1.mask, mask_value, True).alias('masked2'),
+        )
+        result = rf3.agg(
+            rf_agg_no_data_cells(rf3.tile) < rf_agg_no_data_cells(rf3.masked),
+            rf_agg_no_data_cells(rf3.tile) < rf_agg_no_data_cells(rf3.masked2),
+        ) \
+            .first()
+        self.assertTrue(result[0])
+        self.assertTrue(result[1])  # inverse mask arg gives equivalent result
+
+        result_equiv_tiles = rf3.select(rf_for_all(rf_local_equal(rf3.masked, rf3.masked2))).first()[0]
+        self.assertTrue(result_equiv_tiles)  # inverse fn and inverse arg produce same Tile
+
+    def test_mask_by_values(self):
+
+        tile = Tile(np.random.randint(1, 100, (5, 5)), CellType.uint8())
+        mask_tile = Tile(np.array(range(1, 26), 'uint8').reshape(5, 5))
+        expected_diag_nd = Tile(np.ma.masked_array(tile.cells, mask=np.eye(5)))
+        expected_off_diag_nd = Tile(np.ma.masked_array(tile.cells, mask=1 - np.eye(5)))
+
+        df = self.spark.createDataFrame([Row(t=tile, m=mask_tile)]) \
+            .select(rf_mask_by_values('t', 'm', [0, 6, 12, 18, 24]))  # values on the diagonal
+        result0 = df.first()
+        # assert_equal(result0[0].cells, expected_diag_nd)
+        self.assertTrue(result0[0] == expected_diag_nd)
+
+        # mask values off the diagonal! (inverse=True)
+        result1 = self.spark.createDataFrame([Row(t=tile, m=mask_tile)]) \
+            .select(rf_mask_by_values('t', 'm', [0, 6, 12, 18, 24], True)) \
+            .first()
+        # assert_equal(result1[0].cells, expected_off_diag_nd)
+        self.assertTrue(result1[0] == expected_off_diag_nd)
 
     def test_mask(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile, CellType
-        import numpy as np
 
         np.random.seed(999)
         ma = np.ma.array(np.random.randint(0, 10, (5, 5), dtype='int8'), mask=np.random.rand(5, 5) > 0.7)
@@ -326,7 +358,6 @@ def test_render_composite(self):
     def test_rf_interpret_cell_type_as(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile
-        import numpy as np
 
         df = self.spark.createDataFrame([
             Row(t=Tile(np.array([[1, 3, 4], [5, 0, 3]]), CellType.uint8().with_no_data_value(5)))
@@ -341,8 +372,6 @@ def test_rf_interpret_cell_type_as(self):
     def test_rf_local_data_and_no_data(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile
-        import numpy as np
-        from numpy.testing import assert_equal
 
         nd = 5
         t = Tile(
@@ -363,8 +392,6 @@ def test_rf_local_data_and_no_data(self):
     def test_rf_local_is_in(self):
         from pyspark.sql.functions import lit, array, col
         from pyspark.sql import Row
-        import numpy as np
-        from numpy.testing import assert_equal
 
         nd = 5
         t = Tile(