How it works(22) Geotrellis是如何在Spark上计算的(C) 计算模型初探

1.引入

在上一章结尾我们最终生成了MultibandTileLayerRDD[SpatialKey]对象,一切都是为了最重要的步骤——计算——做准备.

2. 计算NDVI

首先我们回顾一下计算代码:

val ndviTiledRDD: TileLayerRDD[SpatialKey] =
  tiledRDD.withContext { rdd =>
    rdd.mapValues { tile =>
      tile.convert(DoubleConstantNoDataCellType).combineDouble(0, 1) { (r, nir) =>
        (nir - r) / (nir + r)
      }
    }.reduceByKey(_.localMax(_)) // spark计算优化,可省略
  }

我们知道NDVI的算法是(nir-r)/(nir+r),但如何将算法应用到Geotrellis中去呢?换句话说,我们需要明确,算法操作的对象是谁.

2.1 分离出Tile

通过阅读官方文档,我们可以明确在Geotrellis中进行spark并行计算的两个重要信息:

Tile是我们精细操作的核心.
需要在计算中保留元数据,使之不丢失空间信息.

所以,我们的首要任务是从上文的结果中找到Tile.

我们在上一步生成了一个MultibandTileLayerRDD[SpatialKey]对象,即RDD[(SpatialKey, MultibandTile)] with Metadata[TileLayerMetadata[SpatialKey]],想要得到其中的MultiBandTile,需要两个步骤:

将RDD[K,V] with Metadata还原为RDD[K,V]
从RDD[K,V]中获取我们所需要的V,即MultibandTile.在本例中,实际是其子类:ArrayMultiBandTile.

同时,因为要兼顾文档中提及的保留元数据的要求.最终得到了文中的代码:

val ndviTiledRDD: TileLayerRDD[SpatialKey] =
  tiledRDD.withContext { rdd: RDD[(SpatialKey, MultibandTile)] =>
    // 得到了RDD就可以直接用Spark的mapValues方法来提取Tile了
    rdd.mapValues { tile: MultibandTile =>
      // ...具体操作Tile
    }
}

这部分的主体是withContext方法,代码如下:

[代码位于spark/Implicit.scala]

implicit class WithContextWrapper[K, V, M](val rdd: RDD[(K, V)] with Metadata[M]) {
    def withContext[K2, V2](f: RDD[(K, V)] => RDD[(K2, V2)]) =
      new ContextRDD(f(rdd), rdd.metadata)
}

可以看出,仅仅是较为简单的隐式转换,实现了将TileLayerRDD类型数据中的RDD抽出来单独计算,最终和原有的元数据重新打包为一个ContextRDD对象的功能.

最终,我们又得到了一个熟悉的ContextRDD对象,它既可以被称为ContextRDD[SpatialKey, Tile, TileLayerMetadata[SpatialKey]],也可以用RDD[(SpatialKey, Tile)] with Metadata[TileLayerMetadata[SpatialKey]]或者更方便的TileLayerRDD[SpatialKey]来称呼.

2.2 Tile与计算核心

回到最核心的NDVI计算代码:

tile.convert(DoubleConstantNoDataCellType)
    .combineDouble(0, 1) { (r, nir) => (nir - r) / (nir + r) }

可以看出,代码主要有两个步骤:

类型转换:convert
按指定的方式合并波段:combineDouble

按顺序,先从类型转换开始.

2.2.1 类型转换

在这一步,我们将数据原有的类型转换为DoubleConstantNoDataCellType类型.为何要做这一步的转换?

首先,我们无法确定原始数据的类型.在本例中计算NDVI时需要保证数据类型为浮点型.我们实际选择了Double类型,也可以根据实际需求选择float类型.
其次,我们想使用Geotrellis类型默认值作为识别Nodata值时的策略,因此选用了DoubleConstantNoDataCellType作为目标类型.

这部分的主体是Tile.convert方法,整理后代码如下:

[代码位于ArrayMultibandTile.scala]

def convert(newCellType: CellType): MultibandTile = {
    // 创建容纳新数据的数组
    val newBands = Array.ofDim[Tile](bandCount)
    cfor(0)(_ < bandCount, _ + 1) { i =>
      // 逐个处理每个波段
      newBands(i) = band(i).convert(newCellType)
    }
    ArrayMultibandTile(newBands)
}

def band(bandIndex: Int): Tile = {
    _bands(bandIndex)
}

// 定义ArrayMultibandTile时传入
class ArrayMultibandTile(_bands: Array[Tile]) extends MultibandTile with MacroMultibandCombiners {
  val bandCount = _bands.size
  ...
}

从代码中可以看出,类型转换的核心是对band的convert操作.我们回顾定义ArrayMultibandTile的代码,可以看出,所谓一个band,即是一个[celltype]ArrayTile:

[代码位于GeoTiffMultibandTile.scala]

case class Chip(
  window: GridBounds[Int],
  bands: Array[MutableArrayTile],
  intersectingSegments: Int,
  var segmentsBurned: Int = 0
)
...
val bands = Array.fill(bandSubsetLength)(ArrayTile.empty(cellType, window.width, window.height))
val chip = Chip(window, bands, segments.length)

上方代码的核心是ArrayTile.empty方法:

[代码位于ArrayTile.scala]

def empty(t: CellType, cols: Int, rows: Int): MutableArrayTile =
    t match {
      // 根据实际情况调用不同的类进行初始化.
      case _: BitCells => BitArrayTile.empty(cols, rows)
      case ct: ByteCells => ByteArrayTile.empty(cols, rows, ct)
      ...
    }

上文所说,类型转换的核心是对band的convert操作,既然知道了[celltype]ArrayTile(即band)是如何产生的,我们就可以看一下它的convert是如何实现的:

[代码位于ArrayTile.scala]

def convert(targetCellType: CellType): ArrayTile = {
    // 创建一个指定类型的空ArrayTile
    val tile = ArrayTile.alloc(targetCellType, cols, rows)
    // 区分指定的是否浮点类型.并进行赋值
    if(!cellType.isFloatingPoint) {
      cfor(0)(_ < rows, _ + 1) { row =>
        cfor(0)(_ < cols, _ + 1) { col =>
          tile.set(col, row, get(col, row))
        }
      }
    } else {
      cfor(0)(_ < rows, _ + 1) { row =>
        cfor(0)(_ < cols, _ + 1) { col =>
          tile.setDouble(col, row, getDouble(col, row))
        }
      }
    }

tile
}

// 预定义的方法,applyDouble的实现在具体子类中
def getDouble(col: Int, row: Int) = applyDouble(row * cols + col)

// 创建一个指定类型的空ArrayTile
def alloc(t: CellType, cols: Int, rows: Int): MutableArrayTile =
    t match {
      case _: BitCells => BitArrayTile.ofDim(cols, rows)
      ...
      case ct: DoubleCells => DoubleArrayTile.ofDim(cols, rows, ct)
    }

代码的主体是ofDim方法,我们假设源数据为DoubleArrayTile类型,看看是如何实现的:

[代码位于DoubleArrayTile.scala]

// 根据具体类型判断Nodata策略
def ofDim(cols: Int, rows: Int, cellType: DoubleCells with NoDataHandling): DoubleArrayTile =
cellType match {
  // 默认策略
  case DoubleCellType =>
    new DoubleRawArrayTile(Array.ofDim[Double](cols * rows), cols, rows)
  // 使用默认Nodata值
  case DoubleConstantNoDataCellType =>
    new DoubleConstantNoDataArrayTile(Array.ofDim[Double](cols * rows), cols, rows)
  // 使用用户自定义Nodata值
  case udct: DoubleUserDefinedNoDataCellType =>
    new DoubleUserDefinedNoDataArrayTile(Array.ofDim[Double](cols * rows), cols, rows, udct)
}
// 对应了不同策略下的applyDouble方法

final case class DoubleConstantNoDataArrayTile(arr: Array[Double], val cols: Int, val rows: Int)
    extends DoubleArrayTile(arr, cols, rows) {
  ...
  // 在遇到是Nodata值时不做特殊处理
  def applyDouble(i: Int): Double = arr(i)
  ...
}

final case class DoubleUserDefinedNoDataArrayTile(arr: Array[Double], val cols: Int, val rows: Int, val cellType: DoubleUserDefinedNoDataCellType)
    extends DoubleArrayTile(arr, cols, rows)
       with UserDefinedDoubleNoDataConversions {
  ...
  // 在遇到是用户预定值时当做Nodata值处理
  def applyDouble(i: Int): Double = udd2d(arr(i))
  ...
}

从上面的代码可知,类型转换的逻辑很简单:将多波段影像中的每一个波段按照按照用户指定的数据类型重新取值,放到新的容器中去.我们也可以一窥多波段影像存储数据的结构:
在这里插入图片描述

接下来,就可以分析最关键的任务:按指定的方式(在本例中是NDVI算法)合并波段.

2.2.2 合并波段

这部分的主体是combine方法:

[代码位于ArrayMultibandTile.scala]

def combineDouble(b0: Int, b1: Int)(f: (Double, Double) => Double): Tile = {
    val band1 = band(b0)
    val band2 = band(b1)
    // 结果容器
    val result = ArrayTile.empty(cellType, cols, rows)
    cfor(0)(_ < rows, _ + 1) { row =>
      cfor(0)(_ < cols, _ + 1) { col =>
        // 计算后赋值
        result.setDouble(col, row, f(band1.getDouble(col, row), band2.getDouble(col, row)))
      }
    }
    result
}

可以看出,这部分的代码逻辑很简单.但其中有一个更深层的问题:它能满足我们的实际使用需求吗?

很显然,在实际操作中,我们可能用到多个波段融合,仅支持绑定两个波段是不够的.仅观察代码,可以发现Geotrellis提供了更加一般化的合并波段方法:

[代码位于ArrayMultibandTile.scala]

  def combineDouble(subset: Seq[Int])(f: Seq[Double] => Double): Tile = {
    subset.foreach({ b => require(0 <= b && b < bandCount, "All elements of subset must be present") })
    val subsetSize = subset.size
    val subsetArray = subset.toArray

    val result = ArrayTile.empty(cellType, cols, rows)
    val values: Array[Double] = Array.ofDim(subsetSize)
    // 遍历指定的波段
    cfor(0)(_ < rows, _ + 1) { row =>
      cfor(0)(_ < cols, _ + 1) { col =>
        cfor(0)(_ < subsetSize, _ + 1) { i =>
          values(i) = _bands(subsetArray(i)).getDouble(col, row)
        }
        result.setDouble(col, row, f(values))
      }
    }
    result
  }

但在我们实际使用Geotrellis时可以发现,我们可以使用诸如如下的方式进行波段合并操作:

tile.combineDouble(0, 1, 2, 3) { (b0, b1, b2, b3) =>  /* 具体操作 */ }

实际支持的波段还可以更多,显然这种方式的调用更加优雅.但我们并不能在实际代码中找到这些方法的定义,他们来自于ArrayMutibandTile的特质.

3. ArrayMultibandTile类及MultibandTile类的特质

ArrayMultibandTile的继承结构如下:
在这里插入图片描述

其中:

绿色的为类继承
红色的为特征实现

但我们可以发现,其中的三个特质是找不到对应的代码文件的.因为这些代码都是编译前通过文本模板生成的,定义于project/Boilerplate.scala文件中.我们无需了解模板是如何运作的,只需要知道模板最终生成的内容是什么,再按照常规的方式分析即可.

3.1 MacroCombinableMultibandTile和MacroCombineFunctions特质

首先来看一下生成的MacroCombinableMultibandTile特质:

[代码将生成于geotrellis/macros/MacroCombinableMultibandTile.scala]

package geotrellis.macros
trait MacroCombinableMultibandTile[T] {
  def combineIntTileCombiner(combiner: IntTileCombiner3): T
  def combineDoubleTileCombiner(combiner: DoubleTileCombiner3): T
  // ...以此类推
  def combineIntTileCombiner(combiner: IntTileCombiner10): T
  def combineDoubleTileCombiner(combiner: DoubleTileCombiner10): T
}

可以看出,MacroCombinableMultibandTile特质为Int和Double类型各预定义了10个接收TileCombiner类型的方法.也由此可以看出,因为是大量重复性的工作,使用模板生成代码是比较方便的实现方法.

代码中的[type]TileCombiner[arity]类型同样来自于模板(以Double类型为例,Int类型类似):

[代码将生成于geotrellis/macros/DoubleTileCombiners.scala]

package geotrellis.macros
  trait DoubleTileCombiner3 {
      val b0: Int; val b1: Int; val b2: Int
      def apply(b0: Double, b1: Double, b2: Double): Double
  }
  // ...以此类推
  trait DoubleTileCombiner10 {
      val b0: Int; val b1: Int; val b2: Int; val b3: Int; val b4: Int; val b5: Int; val b6: Int; val b7: Int; val b8: Int; val b9: Int
      def apply(b0: Double, b1: Double, b2: Double, b3: Double, b4: Double, b5: Double, b6: Double, b7: Double, b8: Double, b9: Double): Double
  }

再来看看生成的MacroCombineFunctions特质:

[代码将生成于geotrellis/macros/MacroCombinableMultibandTile.scala]

package geotrellis.macros
import scala.language.experimental.macros
trait MacroCombineFunctions[T, MBT <: MacroCombinableMultibandTile[T]] {
  def combine(b0: Int, b1: Int, b2: Int)(f: (Int, Int, Int) => Int): T =
    macro MultibandTileMacros.intCombine3_impl[T, MBT]
  def combineDouble(b0: Int, b1: Int, b2: Int)(f: (Double, Double, Double) => Double): T =
    macro MultibandTileMacros.doubleCombine3_impl[T, MBT]
  // ...以此类推
  def combine(b0: Int, b1: Int, b2: Int, b3: Int, b4: Int, b5: Int, b6: Int, b7: Int, b8: Int, b9: Int)(f: (Int, Int, Int, Int, Int, Int, Int, Int, Int, Int) => Int): T =
    macro MultibandTileMacros.intCombine10_impl[T, MBT]
  def combineDouble(b0: Int, b1: Int, b2: Int, b3: Int, b4: Int, b5: Int, b6: Int, b7: Int, b8: Int, b9: Int)(f: (Double, Double, Double, Double, Double, Double, Double, Double, Double, Double) => Double): T =
    macro MultibandTileMacros.doubleCombine10_impl[T, MBT]
}

代码中的[type]Combine[arity]_impl类型依然来自于模板:

[代码将生成于geotrellis/macros/MultibandTileMacros.scala]

package geotrellis.macros
import spire.macros.InlineUtil
import scala.reflect.macros.whitebox.Context
import scala.language.experimental.macros
object MultibandTileMacros {
  def intCombine3_impl[T, MBT <: MacroCombinableMultibandTile[T]](c: Context)(b0: c.Expr[Int], b1: c.Expr[Int], b2: c.Expr[Int])(f: c.Expr[(Int, Int, Int) => Int]): c.Expr[T] = {
    import c.universe._
    val self = c.Expr[MacroCombinableMultibandTile[T]](c.prefix.tree)
    val tree =
    q"""$self.combineIntTileCombiner(new geotrellis.macros.IntTileCombiner3 {
       val b0 = $b0; val b1 = $b1; val b2 = $b2
       def apply(b0: Int, b1: Int, b2: Int): Int = $f(b0, b1, b2)
    })"""
    new InlineUtil[c.type](c).inlineAndReset[T](tree)
  }
  // ...以此类推
  def doubleCombine10_impl[T, MBT <: MacroCombinableMultibandTile[T]](c: Context)(b0: c.Expr[Int], b1: c.Expr[Int], b2: c.Expr[Int], b3: c.Expr[Int], b4: c.Expr[Int], b5: c.Expr[Int], b6: c.Expr[Int], b7: c.Expr[Int], b8: c.Expr[Int], b9: c.Expr[Int])(f: c.Expr[(Double, Double, Double, Double, Double, Double, Double, Double, Double, Double) => Double]): c.Expr[T] = {
    import c.universe._
    val self = c.Expr[MacroCombinableMultibandTile[T]](c.prefix.tree)
    val tree =
    q"""$self.combineDoubleTileCombiner(new geotrellis.macros.DoubleTileCombiner10 {
       val b0 = $b0; val b1 = $b1; val b2 = $b2; val b3 = $b3; val b4 = $b4; val b5 = $b5; val b6 = $b6; val b7 = $b7; val b8 = $b8; val b9 = $b9
       def apply(b0: Double, b1: Double, b2: Double, b3: Double, b4: Double, b5: Double, b6: Double, b7: Double, b8: Double, b9: Double): Double = $f(b0, b1, b2, b3, b4, b5, b6, b7, b8, b9)
    })"""
    new InlineUtil[c.type](c).inlineAndReset[T](tree)
  }
}

至此,我们也能推断出这两个特质究竟做了些什么:

MacroCombineFunctions特质实现了这样一种功能:它通过宏生成了若干combine/combineDouble方法的重载,以支持从3个波段参数到10个波段参数的融合方法.
这个宏是这样实现的:
1. 将上下文限定在继承MacroCombinableMultibandTile的类中
2. 调用该上下文中的combineIntTileCombiner/combineDoubleTileCombiner方法,并构建其需要的参数
3. 通过宏将方法挂载到对应的combine/combineDouble方法的重载上
在宏定义中,需要接收指定参数的combineIntTileCombiner/combineDoubleTileCombiner方法
- 该方法在MacroCombinableMultibandTile中定义
- 需要的不同参数在IntTileCombiners/DoubleTileCombiners.scala中定义

即MacroCombineFunctions定义行为,MacroCombinableMultibandTile定义类型,因此这两个特质必须同时被类继承.

3.2 MacroMultibandCombiners特质

从MacroCombinableMultibandTile和MacroCombineFunctions特质的代码中可以看出,继承了这两个特质的类只需要实现具体的combineIntTileCombiner/combineDoubleTileCombiner方法,就能实现combine/combineDouble方法重载.

MacroMultibandCombiners特质就具体实现了combineIntTileCombiner/combineDoubleTileCombiner方法,也因此让ArryMultiBand具有支持同时融合至多10个波段的combine方法:

[代码将生成于geotrellis/raster/MacroMultibandCombiners.scala]

package geotrellis.raster
import geotrellis.macros._
import spire.syntax.cfor._
trait MacroMultibandCombiners { self: MultibandTile =>
  def combineIntTileCombiner(combiner: IntTileCombiner3): Tile = {
    val band0 = band(combiner.b0); val band1 = band(combiner.b1); val band2 = band(combiner.b2)
    val result = ArrayTile.empty(targetCellType, cols, rows)
    val arr = Array.ofDim[Int](bandCount)
    cfor(0)(_ < rows, _ + 1) { row =>
      cfor(0)(_ < cols, _ + 1) { col =>
        result.set(col, row, combiner(band0.get(col, row), band1.get(col, row), band2.get(col, row)))
      }
    }
    result
  }
  // ...以此类推
  def combineDoubleTileCombiner(combiner: DoubleTileCombiner10): Tile = {
    val band0 = band(combiner.b0); val band1 = band(combiner.b1); val band2 = band(combiner.b2); val band3 = band(combiner.b3); val band4 = band(combiner.b4); val band5 = band(combiner.b5); val band6 = band(combiner.b6); val band7 = band(combiner.b7); val band8 = band(combiner.b8); val band9 = band(combiner.b9)
    val result = ArrayTile.empty(targetCellType, cols, rows)
    val arr = Array.ofDim[Int](bandCount)
    cfor(0)(_ < rows, _ + 1) { row =>
      cfor(0)(_ < cols, _ + 1) { col =>
        result.setDouble(col, row, combiner(band0.getDouble(col, row), band1.getDouble(col, row), band2.getDouble(col, row), band3.getDouble(col, row), band4.getDouble(col, row), band5.getDouble(col, row), band6.getDouble(col, row), band7.getDouble(col, row), band8.getDouble(col, row), band9.getDouble(col, row)))
      }
    }
    result
  }
}

不难看出,其原理与上文提到的双波段融合一致,只是用[type]combiner[arity]代替了f(b0,b1)

为什么需要多此一举,专门定义一个新的类型来包装一下f(b0,b1...)方法呢?其实很简单,因为f(b0,b1...)是在运行时定义的,而宏是在编译期定义的,不专门定义一个包装无法通过编译,因此为了实现这个功能,需要专门定义一个包装对象.

4. 总结

本篇我们从计算NDVI入手,探讨了Geotrellis是如何实现基本的计算功能的.

ArrayMultibandTile提供了map/foreach/combine这3种主要的元算子,已经可以完成许多常用的简单计算工作了.在接下来的文章中,我们将继续分析Geotrellis是如何实现诸如波段卷积等更加复杂的计算操作的.