numpyを継承する

28219 ワード
Python numpy Python テキストリンク

公式ドキュメントのここらへんに全部載っています。余裕があれば参照してください。

なぜnumpyを継承するのか

numpyを継承する理由は3つ挙げられます。

1. 配列にまつわるメタ情報を残したい

実用上numpyを使っている人であれば分かると思うのですが、csvなどからnumpyを読み込み、解析し、またcsvで保存する。普通同じフォルダ内やその付近に保存したいですが、複数のファイルから読み込んだ場合このパスの管理が意外と面倒です。arr.pathみたいに配列にパス情報を紐づけしたいですよね。

2. 異種の配列を同時に使いたい

同じ「配列」でも、生データとフィッティング結果、実空間とフーリエ空間、画像とラベルとマスクといったように、明らかに性質が異なる配列を同時に扱うことがあります。これらをすべてnumpy.ndarrayとしてまとめてしまうと、変数名以外で区別しづらく不便です。例えば自分の中で「生データはplt.scatterに入れて、フィッティング結果はplt.plotで...」みたく取り決めをする必要があり、対話的な解析の効率が大幅に低下します。

配列が別のクラスに属していれば、isinstanceやPython 3.10から始まるmatch構文で対処してディスパッチ (同じ関数で異なる型のデータに異なる動作をさせる) できます。

3. numpy組み込みの関数を魔改造したい

np.meanなど、実は自分用にカスタマイズできてしまいます。私が個人的に開発している画像解析ライブラリでは、例えば画像のz軸方向の平均 (Z-projection/Z投影) を計算する際、毎回z軸が何番目の次元に対応するか数えるのは面倒なので、画像ファイルの軸情報を読み取って、np.mean(img, axis="z")ができるようにしました。標準のままではaxis="z"は当然エラーです。このように、組み込みの関数を改造しておくとコードが簡潔になり作業も捗ります。

継承のポイント(1): __new__

Pythonは動的に変数を追加できます。素朴にメタ情報をmetaに格納してみましょう。

import numpy as np
arr = np.arange(5)
arr.meta = "metadata"

AttributeError: 'numpy.ndarray' object has no attribute 'meta'

はい。numpyでは変数の追加が禁止されています。

メタ情報を付加するには、numpyを継承して__new__メソッドをオーバーロードする必要があります。ndarrayのviewメソッドでクラスを変更します。

class MetaArray(np.ndarray):
    def __new__(cls, obj, dtype=None, meta=None):
        self = np.asarray(obj, dtype=dtype).view(cls)
        self.meta = meta
        return self

これでもう使えます。

arr = MetaArray([1,2,3], meta="2021/08/14")
arr.meta

'2021/08/14'

ndarrayのサブクラスなので、当然numpy系の計算はすべてできます。

arr1 = MetaArray([1,2,3], meta="metadata-1")
arr2 = MetaArray([6,7,8], meta="metadata-2")
arr1 + arr2   # OK
arr1.max()    # OK
np.mean(arr1) # OK

継承のポイント(2): __array_finalize__

__new__だけでは足りません。なぜなら、関数に入れて新しい配列ができたときにmetaを受け継ぐ動作は定義していないからです。

out = np.sort(arr)
out.meta

AttributeError: 'MetaArray' object has no attribute 'meta'

そこで重要になるのが__array_finalize__メソッドです。配列が作られたときに毎回呼び出されるメソッドなので、ここでメタ情報を引き継げばよいです。

class MetaArray(np.ndarray):
    def __new__(...): ...

    def __array_finalize__(self, obj): # objからselfが作られたときに呼び出される
        if obj is None:
            return None
        self.meta = getattr(obj, "meta", None) # 引き継ぎ

ポイントは、objが配列とは限らないことです。エラーにならないようgetattrを使います。

継承のポイント(3): __array_ufunc__

__new__と__array_finalize__でとりあえず正しい値を返すプログラムは書けます。しかし、思うような動作をしないパターンがあります。

arr1 = MetaArray([1,2,3], meta="metadata-1")
arr2 = MetaArray([6,7,8], meta="metadata-2")
np.mean(arr) # MetaArray(2.) ... スカラーにならない
(arr1 + arr2).meta # 'metadata-1' ... 両方のメタ情報を受け継ぎたい

これらの問題は__array_ufunc__メソッドをオーバーロードすることで解決します。"ufunc"とはuniversal functionのことで、数学の基本的な関数は大抵これに属します。

class MetaArray(np.ndarray):
    def __new__(...): ...
    def __array_finalize__(...): ...

    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
        metalist = [] # メタ情報のリスト
        args_ = [] # 入力引数のリスト
        for arg in args:
            # 可能ならメタ情報をリストに追加
            if isinstance(arg, self.__class__) and hasattr(arg, "meta"):
                metalist.append(arg.meta)
            # MetaArrayはndarrayに直す
            arg = arg.view(np.ndarray) if isinstance(arg, MetaArray) else arg
            args_.append(arg)
        # 関数を呼び出す
        out = getattr(ufunc, method)(*args_, **kwargs)

        # なんか必要らしい
        if out is NotImplemented:
            return NotImplemented

        # MetaArrayに戻す。このとき、スカラー(np.float64など)は変化しない。
        out = out.view(self.__class__)

        # MetaArrayのときのみメタ情報を引き継ぐ。このとき、入力したメタ情報を連結する。
        if isinstance(out, self.__class__):
            out.meta = ", ".join(metalist)
        return out

ポイントは、もとの関数をラップする形をとっている点です。関数を呼び出す前にndarrayに直すと同時にメタ情報をすべてリストに格納しておき、出力に応じてMetaArrayへの変換、メタ情報の引継ぎを行います。

これで目標は達成です。

np.mean(arr1) # 2.0
(arr1+arr2).meta # 'metadata-1, metadata-2'

継承のポイント(4): __array_function__

組み込みの関数の魔改造用メソッドです。セットで

  • numpyの関数をデコレートするクラスメソッドimplements(これはほかの名前でもよい)
  • デコレートした関数ともとの関数を対応づけるNP_DISPATCH(これもほかの名前でもよい)
  • @MetaArray.implement(...)でnumpyの関数をデコレートするコード

も必要になります。

クラスの定義

まずはクラスの定義から見ていきます。

class MetaArray(np.ndarray):
    NP_DISPATCH = {} # ディスパッチする関数のマップ

    def __new__(...): ...
    def __array_finalize__(...): ...
    def __array_ufunc__(...): ...

    def __array_function__(self, func, types, args, kwargs):
        if (func in self.__class__.NP_DISPATCH and 
            all(issubclass(t, MetaArray) for t in types)):
            # ディスパッチするべきなら、NP_DISPATCHに登録された関数を呼び出す
            return self.__class__.NP_DISPATCH[func](*args, **kwargs)

        # そうでなければ__array_ufunc__と同じことをする
        metalist = [] # メタ情報のリスト
        args_ = [] # 入力引数のリスト
        for arg in args:
            if isinstance(arg, self.__class__) and hasattr(arg, "meta"):
                metalist.append(arg.meta)
            arg = arg.view(np.ndarray) if isinstance(arg, MetaArray) else arg
            args_.append(arg)

        out = func(*args_, **kwargs)

        if out is NotImplemented:
            return NotImplemented

        if isinstance(out, np.ndarray): # 今回は様々な型があり得るので
            out = out.view(self.__class__)

        if isinstance(out, self.__class__):
            out.meta = ", ".join(metalist)
        return out

    @classmethod
    def implements(cls, numpy_function):
        def wrapped_func(func):
            cls.NP_DISPATCH[numpy_function] = func  # ディスパッチする関数を登録
            return func
        return wrapped_func

__array_function__はシンプルで、改造済として登録された関数であれば登録した関数を呼び出し、そうでなければ__array_ufunc__とほぼ同じことをするだけです。ですから実際は

class MetaArray(np.ndarray):
    ...
    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
        args, kwargs = _process_input(args, kwargs)
        out = getattr(ufunc, method)(*args_, **kwargs)
        out = _process_output(out, args, kwargs)
        return out

    def __array_function__(...): ... # 同様

みたく_process_inputと_process_outputで関数呼び出し前後の作業を統一することになるかと思います。

implementsはとてもシンプルなデコレータで、関数をNP_DISPATCH登録して返しているだけです。ちなみにwrapsで関数のシグネチャをコピーしなくてもQtConsoleなどではdocstringが表示されます (numpyの関数を拡張しているだけなので当然)。

ディスパッチする関数の登録

implementsメソッドを使って関数の登録をしていきます。このとき、公式ドキュメントでは関数を命名していますが、あとから名前を参照する場合を除き名前自体に意味はないので、ここでは楽して_で定義します。

@MetaArray.implements(np.mean)
def _(arr, *args, **kwargs):
    print("dispatched!!")
    out = np.mean(arr.view(np.ndarray), *args, **kwargs)
    return _process_output(out, args, kwargs)

@MetaArray.implements(np.std)
def _(arr, *args, **kwargs):
    print("dispatched!!")
    out = np.std(arr.view(np.ndarray), *args, **kwargs)
    return _process_output(out, args, kwargs)

この例では単に"dispatched!!"と出てくるだけですが、これで使えるようになりました。

np.mean(arr1)

dispatched!!
Out: 0.816496580927726

終わりに

numpyを継承しましょう。とても便利です。

ndarrayをまるまる継承するのが怖かったら、numpy.lib.mixins.NDArrayOperatorsMixinで機能を部分的に継承することもできます。いろいろ試してみましょう。

JavaはBlockingQueueで簡単なメッセージキュー(生産者と消費者モデル)を実現
Scalaの基本タイプ