NumPy でつくる俺々データ構造

はじめに

Python での数値計算の基盤をなす NumPy 、直感的なスライスやブロードキャスト、関数のベクトル適用など大変便利だ。

import numpy as np
np.__version__
# '1.9.2'

np.array([1, 2, 3])
# array([1, 2, 3])

np.array([1, 2, 3])[:2]
# array([1, 2])

np.array([1, 2, 3]) + 1
# array([2, 3, 4])

が、用途によっては NumPy 標準ではその機能を実現できない場合がある。例えば、

配列とメタデータをひとつのクラスで扱いたい
配列への入力や型を制約/検証したい
自作クラスを NumPy の Universal Functions (ufunc) に対応させたい
新しい型 ( dtype ) を作りたい

こういったとき、NumPy の動作を継承した自作のクラスが作れるとうれしい。方法として、大きく二つがある。

NumPy ndarray を継承する
NumPy Array Interface を実装する

どちらを使うべきかは目的による。サブクラス化しないほうが柔軟な処理ができるが、互換をとるための実装は多くなる。pandas の主要データ構造では柔軟性を重視して ndarray の継承をやめ、現在は 2 の方法を取っている。

ndarray のサブクラスを作成する方法は以下のドキュメントに記載されている。きちんとやりたい場合はこちらを読むのがよい。

Subclassing ndarray — NumPy v1.10 Manual

このエントリでは 2. Array Interface を使って最低限の動作をさせる方法を書く。

クラスの定義

サンプルとして、名前をもつ配列クラス NamedArray を考える。__init__での初期化時に、名前をあらわす name と配列の値 values を渡すことにする。

class NamedArray(object):
    
    def __init__(self, name, values):
        self.name = name
        self.values = np.array(values)
        
    def __repr__(self):
        return 'NamedArray: {0}: {1}'.format(self.name, self.values)
    
    def __len__(self):
        return len(self.values)
        
n = NamedArray('x', [1, 2, 3])
n
# NamedArray: x: [1 2 3]

len(n)
# 3

このクラスに ndarray と互換の挙動をさせたい。まず、NamedArray インスタンスから ndarray が直接作成できるとうれしい。が、NamedArray を np.array に渡すと shape のない ndarray (0-dimmentional array) が作成されてしまう。

np.array(n)
# array(NamedArray: x: [1 2 3], dtype=object)

np.array(n).shape
# ()

理由は np.array が NamedArray をスカラーとして扱うため。期待の動作をさせるためには、np.array に NamedArray クラスが配列であることを教える必要がある。

具体的には NamedArray に Array Interface ( __array__ ) を定義すればよい。以下のように __array__ メソッドを追加したクラス NPNamedArray を作成すると、期待通り values の値から ndarray が作成されていることがわかる。

class NPNamedArray(NamedArray):
    
    def __array__(self):
        return self.values
    
n = NPNamedArray('x', [1, 2, 3])
np.array(n)
# array([1, 2, 3])

np.array(n).shape
# (3,)

ufunc 処理のオーバーライド

また、Array Interface をもつインスタンスには ufunc が適用できるようになる。引数は ufunc 内で ndarray に変換されるため、返り値は ndarray になる。

np.sin(n)
# array([ 0.84147098,  0.90929743,  0.14112001])

このときの内部の動きをみてみる。ufunc 適用後の処理は __array_wrap__ メソッドによって制御することができる。

__array_wrap__ メソッドには out_arr と context 2 つの引数を定義する必要がある。これらに何が渡されているのか、print して確認する。

class NPNamedArray(NamedArray):
    
    def __array__(self):
        return self.values
    
    def __array_wrap__(self, out_arr, context=None):
        print('out_arr: ', out_arr)
        print('context: ', context)
        return out_arr
    
n = NPNamedArray('x', [1, 2, 3])
np.sin(n)
# ('out_arr: ', array([ 0.84147098,  0.90929743,  0.14112001]))
# ('context: ', (<ufunc 'sin'>, (NamedArray: x: [1 2 3],), 0))
# array([ 0.84147098,  0.90929743,  0.14112001])

out_arr には ufunc が適用された結果の ndarray が、 context には ufunc を呼び出した際のシグネチャ ( ufunc 自体, ufunc の引数, ufunc のドメイン ) が tuple で渡ってきている。__array_wrap__ の返り値を NPNamedArray に書き換えてやれば、期待する動作ができそうだ。

補足 np.isreal や np.ones_like など、Universal Functions に記載されていても __array_wrap__ を通らないものもある。

補足 ufunc のドメインは NumPy の場合 (おそらく) 全部 0。

class NPNamedArray(NamedArray):
    
    def __array__(self):
        return self.values
    
    def __array_wrap__(self, out_arr, context=None):
        return NPNamedArray(self.name, out_arr)
    
n = NPNamedArray('x', [1, 2, 3])
np.sin(n) 
# NamedArray: x: [ 0.84147098  0.90929743  0.14112001]

ここまでの時点で、NPNamedArray には __add__ メソッドを定義していないため、通常の加算はエラーになる。

n + 1
# TypeError: unsupported operand type(s) for +: 'NPNamedArray' and 'int'

が、ufunc である np.add や ndarray を含む演算は利用できるようになる。これは先日のエントリに記載した通り ndarray 側に Array Interface をもつクラスとの演算の定義があるため。

np.add(n, 1)
# NamedArray: x: [2 3 4]

n + np.array([5, 5, 5])
# NamedArray: x: [6 7 8]

また、__array_wrap__ の context 引数の値によって、特定の ufunc の処理をオーバーライドすることができる。例えば np.sin を適用したときに任意の値を返すこともできる。

class NPNamedArray(NamedArray):
    
    def __array__(self):
        return self.values
    
    def __array_wrap__(self, out_arr, context=None):
        if context[0] is np.sin:
            return u'返したい値'
        return NPNamedArray(self.name, out_arr)
    
n = NPNamedArray('x', [1, 2, 3])
np.sin(n)
# u'ひまわり'

とはいえ、上のように個々の関数について条件分岐を作成するのは手間だ。こんなとき、 ufunc 自体のプロパティを参照することでより簡単に書ける場合がある。

例えば、特定の型への処理をサポートする ufunc のみを利用する場合は .types プロパティを参照すればよい。このプロパティは ufunc がどのような引数/返り値の型を取るかを示す Array-protocol Type String を返す。

参考

np.sin.types
# ['e->e', 'f->f', 'd->d', 'g->g', 'F->F', 'D->D', 'G->G', 'O->O']

np.cos.types
# ['e->e', 'f->f', 'd->d', 'g->g', 'F->F', 'D->D', 'G->G', 'O->O']

np.add.types
# ['??->?', 'bb->b', 'BB->B', 'hh->h', 'HH->H', 'ii->i', 'II->I', 'll->l', 'LL->L',
#  'qq->q', 'QQ->Q', 'ee->e', 'ff->f', 'dd->d', 'gg->g', 'FF->F', 'DD->D', 'GG->G',
#  'Mm->M', 'mm->m', 'mM->M', 'OO->O']