geotrellis.spark.resample.ZoomResample.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of geotrellis-spark_2.11 Show documentation
GeoTrellis is an open source geographic data processing engine for high performance applications.
The newest version!
package geotrellis.spark.resample

import geotrellis.raster._
import geotrellis.raster.resample._
import geotrellis.spark._
import geotrellis.spark.tiling._
import geotrellis.util._

import org.apache.spark.rdd.RDD

object ZoomResample {
  private def gridBoundsAtZoom(sourceZoom: Int, spatialKey: SpatialKey, targetZoom: Int): GridBounds = {
    val SpatialKey(col, row) = spatialKey
    val zoomDiff = targetZoom - sourceZoom
    val factor = math.pow(2, zoomDiff).toInt
    val (minCol, minRow) = (col * factor, row * factor)
    val (maxCol, maxRow) = (((col + 1) * factor) - 1, ((row + 1) * factor) - 1)
    GridBounds(minCol, minRow, maxCol, maxRow)
  }

  private def boundsAtZoom[K: SpatialComponent](sourceZoom: Int, bounds: Bounds[K], targetZoom: Int): Bounds[K] =
    bounds match {
      case KeyBounds(minKey, maxKey) =>
        val min = {
          val gb = gridBoundsAtZoom(sourceZoom, minKey.getComponent[SpatialKey], targetZoom)
          minKey.setComponent(SpatialKey(gb.colMin, gb.rowMin))
        }

        val max = {
          val gb = gridBoundsAtZoom(sourceZoom, maxKey.getComponent[SpatialKey], targetZoom)
          maxKey.setComponent(SpatialKey(gb.colMax, gb.rowMax))
        }
        KeyBounds(min, max)
      case EmptyBounds =>
        EmptyBounds
    }

  /** Resamples a tile layer from a lower zoom level to a higher zoom level.
    * The levels are based on the ZoomedLayoutScheme.
    *
    * @param       rdd              The RDD to be resampled.
    * @param       sourceZoom       The zoom level of the rdd.
    * @param       targetZoom       The zoom level we want to resample to.
    * @param       targetGridBounds Optionally, a grid bounds in the target zoom level we want to filter by.
    * @param       method           The resample method to use for resampling.
    */
  def apply[K: SpatialComponent](
    rdd: TileLayerRDD[K],
    sourceZoom: Int,
    targetZoom: Int,
    targetGridBounds: Option[GridBounds] = None,
    method: ResampleMethod = NearestNeighbor
  ): TileLayerRDD[K] = {
    require(sourceZoom < targetZoom, "This resample call requires that the target zoom level be greater than the source zoom level")
    val tileSize = rdd.metadata.layout.tileLayout.tileCols
    val targetLayoutDefinition =
      ZoomedLayoutScheme.layoutForZoom(targetZoom, rdd.metadata.layout.extent, tileSize)
    val targetMapTransform = targetLayoutDefinition.mapTransform
    val sourceMapTransform = rdd.metadata.mapTransform
    val (resampledRdd: RDD[(K, Tile)], md) =
      targetGridBounds match {
        case Some(tgb) =>
          val resampleKeyBounds = boundsAtZoom(sourceZoom, rdd.metadata.bounds, targetZoom).get
          resampleKeyBounds.toGridBounds.intersection(tgb) match {
            case Some(resampleGridBounds) =>
              val resampled =
                rdd
                  .mapPartitions { partition =>
                    partition
                      .map { case (key, tile) =>
                        gridBoundsAtZoom(sourceZoom, key.getComponent[SpatialKey], targetZoom).intersection(resampleGridBounds).map { gb =>
                          gb
                            .coords
                            .map { case (col, row) =>
                              val sourceExtent = sourceMapTransform(key)
                              val targetExtent = targetMapTransform(col, row)
                              val resampled =
                                tile.resample(sourceExtent, RasterExtent(targetExtent, tileSize, tileSize), method)
                              (key.setComponent(SpatialKey(col, row)), resampled)
                            }
                          }
                      }
                      .flatten
                      .flatten
                    }
              val extent = targetMapTransform(resampleGridBounds).intersection(rdd.metadata.extent).get
              val md = rdd.metadata.copy(layout = targetLayoutDefinition, bounds = resampleKeyBounds.setSpatialBounds(resampleGridBounds), extent = extent)
              (resampled, md)
            case None =>
              val md = rdd.metadata.copy(layout = targetLayoutDefinition, bounds = (EmptyBounds: Bounds[K]))
              (rdd.sparkContext.parallelize(Seq()), md)
          }
        case None =>
          val resampled =
            rdd
              .flatMap { case (key, tile) =>
                gridBoundsAtZoom(sourceZoom, key.getComponent[SpatialKey], targetZoom)
                  .coords
                  .map { case (col, row) =>
                    val sourceExtent = sourceMapTransform(key)
                    val targetExtent = targetMapTransform(col, row)
                    val resampled =
                      tile.resample(sourceExtent, RasterExtent(targetExtent, tileSize, tileSize), method)
                    (key.setComponent(SpatialKey(col, row)), resampled)
                  }
               }
            val md = rdd.metadata.copy(layout = targetLayoutDefinition, bounds = boundsAtZoom(sourceZoom, rdd.metadata.bounds, targetZoom))
          (resampled, md)
        }
    ContextRDD(resampledRdd, md)
  }
}