Mask on LazyMultibandRasters directly (#89)

geotrellis · Apr 10, 2019 · 899f6e0 · 899f6e0
1 parent f95bb3b
commit 899f6e0
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 7 deletions.
diff --git a/jvm/src/main/scala/eval/directive/OpDirectives.scala b/jvm/src/main/scala/eval/directive/OpDirectives.scala
@@ -360,7 +360,7 @@ object OpDirectives {
     }).andThen({ case (lzRaster, geom) =>
       geom.as[MultiPolygon] match {
         case Some(mp) =>
-          Valid(ImageResult(LazyMultibandRaster(List(MaskingNode(lzRaster.bands.values.toList, mp)))))
+          Valid(ImageResult(lzRaster.mask(mp)))
         case None =>
           Invalid(NEL.of(NonEvaluableNode(mask, Some("Masking operation requires its vector argument to be a multipolygon"))))
       }

diff --git a/jvm/src/main/scala/eval/tile/LazyMultibandRaster.scala b/jvm/src/main/scala/eval/tile/LazyMultibandRaster.scala
@@ -87,6 +87,13 @@ case class LazyMultibandRaster(val bands: Map[String, LazyRaster]) {
     val lztiles = bands.mapValues({ lt => LazyRaster.Hillshade(List(lt), gridbounds, zFactor, cs, azimuth, altitude) })
     LazyMultibandRaster(lztiles)
   }
+
+  def mask(
+    maskPoly: MultiPolygon
+  ): LazyMultibandRaster = {
+    val lztiles = bands.mapValues({ lt => MaskingNode(List(lt), maskPoly) })
+    LazyMultibandRaster(lztiles)
+  }
 }
 
 object LazyMultibandRaster {

diff --git a/jvm/src/test/scala/eval/ResultSpec.scala b/jvm/src/test/scala/eval/ResultSpec.scala
@@ -28,7 +28,7 @@ class ResultSpec extends FunSpec with Matchers {
 
   it("Evaluate to desired output (tile)") {
     val anImage = ImageResult(LazyMultibandRaster(List(
-      LazyRaster(IntArrayTile(1 to 4 toArray, 2, 2), Extent(0,0,0,0), WebMercator))
+      LazyRaster(IntArrayTile(1 to 4 toArray, 2, 2), Extent(0,0,1,1), WebMercator))
     ))
     val anInt = IntResult(1)
 
@@ -37,7 +37,7 @@ class ResultSpec extends FunSpec with Matchers {
     anInt.as[MultibandTile] should matchPattern { case Invalid(_) => }
 
     val complexImage = ImageResult(LazyMultibandRaster(List(
-      LazyRaster.MapInt(List(LazyRaster(IntArrayTile(1 to 4 toArray, 2, 2), Extent(0,0,0,0), WebMercator)), { i: Int => i + 4 })
+      LazyRaster.MapInt(List(LazyRaster(IntArrayTile(1 to 4 toArray, 2, 2), Extent(0,0,1,1), WebMercator)), { i: Int => i + 4 })
     )))
     complexImage.as[MultibandTile] should matchPattern { case Valid(_) => }
   }
@@ -62,15 +62,23 @@ class ResultSpec extends FunSpec with Matchers {
     val mask = Extent(0, 0, 1, 1).as[Polygon].map(MultiPolygon(_)).get
     val maskResult = ImageResult(LazyMultibandRaster(List(MaskingNode(List(rasterOnes), mask))))
 
+
+    val maskResultSB = ImageResult(LazyMultibandRaster(List(MaskingNode(List(rasterOnes), mask))))
+    val maskResultRGB = ImageResult(LazyMultibandRaster(
+                                      List(rasterOnes, rasterOnes, rasterOnes)).mask(mask))
+
     for {
       x <- (0 to 3 toArray)
       y <- (0 to 3 toArray)
     } yield {
-      val fetched = maskResult.res.bands.head._2.get(x, y)
+      val fetchedSB = maskResultSB.res.bands.head._2.get(x, y)
+      val fetchedRGB = maskResultRGB.res.bands.toList map { _._2.get(x, y) }
       if ((x, y) == (0, 3)) {
-        isData(fetched) should be (true)
+        isData(fetchedSB) should be (true)
+        fetchedRGB map { isData(_) } should be (List(true, true, true))
       } else {
-        isData(fetched) should be (false)
+        isData(fetchedSB) should be (false)
+        fetchedRGB map { isData(_) } should be (List(false, false, false))
       }
     }
   }

diff --git a/jvm/src/test/scala/eval/ScopedEvaluationSpec.scala b/jvm/src/test/scala/eval/ScopedEvaluationSpec.scala
@@ -21,7 +21,7 @@ import scala.reflect._
 class ScopedEvaluationSpec extends FunSpec with Matchers {
 
   def tileToLit(tile: Tile): RasterLit[ProjectedRaster[MultibandTile]] =
-    RasterLit(ProjectedRaster(MultibandTile(tile), Extent(0, 0, 0, 0), WebMercator))
+    RasterLit(ProjectedRaster(MultibandTile(tile), Extent(0, 0, 1, 1), WebMercator))
 
   implicit class TypeRefinement(self: Interpreted[Result]) {
     def as[T](implicit ct: ClassTag[T]): Interpreted[T] = self match {