Crystalでk-d木を実装してRubyと実行速度を比べてみた

はじめに

こんにちは、Misocaの洋食（@yoshoku）です。この記事は、Misoca+弥生 Advent Calendar 2019の4日目の記事です。

Misocaは、WebフレームワークにRuby on Railsを採用しており、Rubyと関わりの深い会社です。そのRubyとよく似た文法で、静的型付けな言語にCrystalがあります。どの様な言語か試してみました。

インストールと動作確認

Crystalのインストール方法は、Macであればhomebrewでインストールするのが簡単です（公式ページに各OSでのインストール方法が掲載されています）。

$ brew install crystal

挨拶するだけのコードを書いてみました。公式ページには「Crystal is statically type checked...Crystal has built-in type inference」とあり、型推論により型を明示的に書くことは少なそうですね。

class Greeting
  def hello
    puts "Hello, World."
  end
end

g = Greeting.new
g.hello

動かすだけの場合：

$ crystal hello.cr
Hello, World.

実行ファイルを作って動かす場合：

$ crystal build hoge.cr
$ ./hoge
Hello, World.

buildでは、--releaseオプションをつけると、最適化した実行ファイルを得られます。ただし、コンパイルは遅くなります。

エディタの設定

公式Wikiに各エディタのプラグインなどが紹介されています。私はneovim+vim-plugな環境でプログラムを書いているので、以下のようなinit.vimを書いてPlugInstallしました。

call plug#begin()

" ...省略
Plug 'rhysd/vim-crystal'

call plug#end()

k-d木を実装する

k-d木は、二分木による空間分割データ構造で、近傍探索に利用されます。k次元の空間に配置されるデータを、各軸の中央値で分割することを繰り返し、木構造で表現します。この記事の実装では、k次元のデータは、k個の要素からなるfloatの配列で表すとします。

# 4個の3次元なデータの例
data = [[1.2, 0.5, 3.0], 
        [0.2, 0.3, 2.0],
        [2.3, 1.8, 1.7],
        [1.6, 1.4, 2.9]]

Crystalでの実装

まず、木のノードを実装します。どの軸を分割の基準に用いたか、中央値の値、木の深さ、二分木なので左右のノード、そして、葉であった場合のデータ点を保持します。

class Node
  getter :axis, :threshold, :depth, :left, :right, :element

  def initialize(@axis : Int32, @threshold : Float32, @depth : Int32,
                 @left : Node | Nil, @right : Node | Nil, 
                 @element : Array(Float32) | Nil)
  end
end

インスタンス変数を直接的にinitializeに書けるのがおもしろいですね。getterは、Rubyのattr_accessorに相当するものです。ただ値を保持するだけなので、Structでも良いのですが、後でRubyと比較したいのでClassにしました。

次にk-d木の構築と、最近傍探索を実装していきましょう。各ノードでの分割以外は、以下の様な感じになります。検索は、葉に到達するまで木をたどるだけです。

class KdTree
  @tree : Node | Nil = nil

  # 木の構築
  def build(points)
    @tree = grow(points, 0)
  end

  # 最近傍データの探索
  def search(point)
    apply(point, @tree)
  end

  private def grow(points, depth)
    # ...あとで実装する。
  end

  private def apply(point, node)
    # 葉に到達したら、葉がもつデータを返す.
    return nil if node.nil?
    return node.element unless node.element.nil?

    # 基準となる軸と閾値で、次にたどる左右のノードを選択する.
    if point[node.axis] >= node.threshold
      apply(point, node.left)
    else
      apply(point, node.right)
    end
  end
end

k次元データの軸を分割することで、木を成長させていく部分を実装します。分割の基準となる軸の選択は、木の深さに応じて順番に選択されます。そして、選択された軸の中央値を閾値とします。閾値でデータを分割して、左右のノードに割り当てます。これをデータが1個になるまで繰り返します（葉に割り当てるデータの数は、1個と決まっているわけではありませんが、この記事では簡単のため1個としました）。

private def grow(points, depth)
  # データ数が0ならばNodeを作らない.
  n_points = points.size
  return nil if n_points.zero?

  # データ数が1ならば葉としてデータを割り当てたNodeを返す.
  if n_points == 1
    return Node.new(-1, 0.0, depth, nil, nil, points[0])
  end

  # 分割の基準となる軸を選択する.
  n_axes = points[0].size
  target_axis = depth % n_axes

  # その軸上の値の中央値を閾値とする.
  threshold = points.map { |e| e[target_axis] }.sort[(n_points / 2).to_i]

  # 閾値でデータを左右に分割する.
  Node.new(
    target_axis, threshold, depth + 1,
    grow(points.select { |e| e[target_axis] >= threshold }, depth + 1),
    grow(points.select { |e| e[target_axis] < threshold }, depth + 1),
    nil
  )
end

Crystalは型推論を行うので、型を明示する必要がない場合は、上記のように、ほぼRubyと同じようなものになります。

Rubyでの実装

同じものをRubyで実装してみました。よく似てますね。

class Node
  attr_accessor :axis, :threshold, :depth, :left, :right, :element

  def initialize(axis, threshold, depth, left, right, element)
    @axis = axis
    @threshold = threshold
    @depth = depth
    @left = left
    @right = right
    @element = element
  end
end

class KdTree
  def initialize
    @tree = nil
  end

  def build(points)
    @tree = grow(points, 0)
  end

  def search(point)
    apply(point, @tree)
  end

  private def grow(points, depth)
    n_points = points.size
    return nil if n_points.zero?

    if n_points == 1
      return Node.new(-1, 0.0, depth, nil, nil, points[0])
    end

    n_axes = points[0].size

    target_axis = depth % n_axes

    threshold = points.map { |e| e[target_axis] }.sort[(n_points / 2).to_i]
    Node.new(
      target_axis, threshold, depth + 1,
      grow(points.select { |e| e[target_axis] >= threshold }, depth + 1),
      grow(points.select { |e| e[target_axis] < threshold }, depth + 1),
      nil
    )
  end

  private def apply(point, node)
    return nil if node.nil?
    return node.element unless node.element.nil?

    if point[node.axis] >= node.threshold
      apply(point, node.left)
    else
      apply(point, node.right)
    end
  end
end

実験

準備

pendigitsという16次元のデータで検索を行ってみて、その実効速度を、CrystalとRubyで比較してみましょう。csvファイルのpendigitsデータセットを取得します。

$ wget https://datahub.io/machine-learning/pendigits/r/pendigits.csv

このpendigitsデータセットを読み込んで配列にするコードは、以下のようになります。

require "csv"

points = Array(Array(Float32)).new
skip_header = true

CSV.each_row(File.open("pendigits.csv")) do |row|
  if skip_header
    skip_header = false
    next
  end

  pt = Array(Float32).new
  row[0...-1].each { |e| pt.push(e.to_f32) }
  points.push(pt)
end

さらにBenchmarkで実行速度を計測するコードは、以下のようになります。

require "benchmark"

Benchmark.bm do |x|
  x.report("kd-tree") do
    tree = KdTree.new
    tree.build(points)
    points.each { |pt| tree.search(pt) }
  end
end

掲載しませんが、Rubyでも同様のコードを実装しました。

実験結果

実験環境は、MacBook Early 2016（Intel Core m3 1.1GHz, 8GB DDR3）
です。まずは、Crystalで実装したものから実行してみます。

$ crystal -v
Crystal 0.31.1 (2019-10-02)

LLVM: 8.0.1
Default target: x86_64-apple-macosx
$ crystal build --release kdtree.cr
$ ./kdtree
              user     system      total        real
kd-tree   0.031911   0.002470   0.034381 (  0.028745)

次に、Rubyで実装したものを実行します。

$ ruby -v
ruby 2.6.5p114 (2019-10-01 revision 67812) [x86_64-darwin19]
$ ruby kdtree.rb
       user     system      total        real
kd-tree  0.164685   0.003777   0.168462 (  0.170641)

さすがにコンパイルした実行ファイルなだけあって、Crystalのほうが速いですね（コンパイルの時間を必要としますが）。

ただし、ちょっと例が良くなかったのか、例えば以下のようなk-d木の構築・検索では、Rubyでも遅くは感じませんでした。

tree = KdTree.new
tree.build(points)
p tree.search(points[0])
p tree.search(points[1])
p tree.search(points[2])

$ time ruby kdtree.rb
[47.0, 100.0, 27.0, 81.0, 57.0, 37.0, 26.0, 0.0, 0.0, 23.0, 56.0, 53.0, 100.0, 90.0, 40.0, 98.0]
[0.0, 89.0, 27.0, 100.0, 42.0, 75.0, 29.0, 45.0, 15.0, 15.0, 37.0, 0.0, 69.0, 2.0, 100.0, 6.0]
[0.0, 57.0, 31.0, 68.0, 72.0, 90.0, 100.0, 100.0, 76.0, 75.0, 50.0, 51.0, 28.0, 25.0, 16.0, 0.0]
ruby kdtree.rb  0.44s user 0.10s system 95% cpu 0.562 total

おわりに

Crystalの書き味は、Rubyとよく似たものでした。Rubyで開発したもので、外部ライブラリに依存がなく、実行速度が欲しいものは、Crystalに書き換えるのを試しても良いかもしれません。