「それNArrayでできるよ」をもっと便利にした - 札幌市中央区Ruby会議01

もう先週の話になってしまいましたが、Tricknotesさん主催の札幌市中央区Ruby会議01に参加してきました。コンパクトな規模ながらも「非常に良いRubyistの活動が観測」（byしまださん）された、有意義かつ楽しい会議でした。

どの発表も聴き応えのあるすばらしいものでしたが、中でもtmaedaさんの「それNArrayでできるよ」は数値計算ライブラリのNArrayをbetter Excelとして使ってしまうコペルニクス的発想で衝撃的でした。

この発表のなかで、NArrayの軸をHuman-Readableな形式で扱えるラッパークラスを作ると便利だよ、というお話がありましたが、

これをもっと便利にするNArrayCubeクラスを試験的に実装してみたので紹介します。

使い方

# 店、世代、性別の3軸からなる売上表を作成
cube = NArrayCube.new(:int,
                      { title: :shop,       labels: ['A', 'B', 'C'] },
                      { title: :generation, labels: ['teens', 'twenties', 'thirties', 'fourties'] },
                      { title: :gender,     labels: ['men', 'women'] }
                     )

# 店A、20代、男性の売上を入力
cube.set({ shop: 'A', generation: 'teens', gender: 'men' }, 100)

# 店A、10代の売上を入力
cube.set({ shop: 'A', generation: 'teens' }, [90, 110])

# 店Bの売上を一気に入力（4世代 × 2性別のArray）
cube.set({ shop: 'B'}, [ [50, 80, 60, 75], [70, 40, 35, 50] ])

# 店Cの売上を一気に入力：第2引数を一次元のArrayで与えると勝手に4x2に変換してくれる
cube.set({ shop: 'C'}, [70, 80, 55, 60, 30, 45, 70, 60])

# 10代の売上を取得
cube.get({ generation: 'teens' })
#=>
# #<NArrayCube: ...
# [ [ 90, 50, 70 ], 
#   [ 110, 70, 30 ] ]>

# 10代の売上の合計
cube.get({ generation: 'teens' }).sum
#=> 420

# 10代の売上を店ごとに集計 => 男性＋女性の値を計算すればよい
cube.get({ generation: 'teens' }).sum(:gender)
#=>
# NArray.int(3): 
# [ 200, 120, 100 ]

ツボ

NArrayCube#setで複数のセルに一気に値を流し込める
NArray#sumをNArrayCube#sumでラップしたことにより、軸の指定をHuman-Readableに行える

TODO

spec書く
document書く
NArray#sumだけじゃなくprod/max/minとかその他のメソッドにも対応したい => method_missingとか使えるかな？
「10代の売上を店ごとに集計」はcube.get({ generation: 'teens' }).sum_by(:shop)みたいに書けたほうが直感的
getで軸ラベルを複数指定できたほうがいいcube.get({ shop: ['A', 'B'] })
tmaedaさんがいいよって言ってくれたらGem化して公開

コード

require 'narray'

class NArrayCube
  NARRAY_TYPES = %i(byte sint int sfloat float scomplex complex object)

  class NArrayAxis
    def initialize(id, title, labels)
      @id = id
      @title = title.to_s
      @labels = labels.map(&:to_s)
    end

    def value_index(value)
      @labels.index(value.to_s)
    end

    attr_reader :id, :title, :labels
  end

  def initialize(type, *axes)
    if NARRAY_TYPES.include?(type)
      @type = type
    else
      raise ArgumentError
    end

    @axes = axes.map.with_index do |a, i|
      NArrayAxis.new(i, a[:title], a[:labels])
    end

    @narray = NArray.method(type).call(*@axes.map { |a| a.labels.length })
  end

  def get(conditions)
    conditions = build_condition(conditions)
    values = @narray[*conditions]

    axes = conditions.map.with_index do |c, id|
      if c.class == Array
        a = @axes[id]
        { title: a.title, labels: a.labels.select.with_index { |v, i| c.include? i } }
      else
        nil
      end
    end

    NArrayCube.new(@type, *axes.compact).set({}, values)
  end

  def set(conditions, values)
    conditions = build_condition(conditions)

    unless [Array, NArray].include?(values.class)
      @narray[*conditions] = values
    else
      partial = @narray[*conditions]
      values = NArray.to_na(values)

      if partial.shape == values.shape
        @narray[*conditions] = values
      elsif partial.total == values.total && values.shape.length == 1
        @narray[*conditions] = values.reshape(*partial.shape)
      else
        raise ArgumentError
      end
    end

    self
  end

  def sum(*titles)
    if titles.length > 0
      @narray.sum(*titles.map { |t| axis(t).id })
    else
      @narray.sum
    end
  end

  private
  def build_condition(conditions)
    cond = @axes.map { |a| (0..(a.labels.length - 1)).to_a }
    conditions.each do |title, label|
      a = axis(title)
      cond[a.id] = a.value_index(label)
    end
    cond
  end

  def axis(title)
    @axes.select { |a| a.title == title.to_s }.first
  end
end