numpyを継承する

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

https://numpy.org/doc/stable/user/basics.subclassing.html

https://numpy.org/doc/stable/reference/arrays.classes.html

https://numpy.org/doc/stable/user/basics.dispatch.html#basics-dispatch

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

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

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

```python
import numpy as np
arr = np.arange(5)
arr.meta = "metadata"
```
```
AttributeError: 'numpy.ndarray' object has no attribute 'meta'
```

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

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

```python
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
```

これでもう使えます。

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

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

```python
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`を受け継ぐ動作は定義していないからです。

```python
out = np.sort(arr)
out.meta
```
```
AttributeError: 'MetaArray' object has no attribute 'meta'
```

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

```python
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__`でとりあえず正しい値を返すプログラムは書けます。しかし、思うような動作をしないパターンがあります。

```python
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のことで、数学の基本的な関数は大抵これに属します。

```python
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`への変換、メタ情報の引継ぎを行います。

これで目標は達成です。

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

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

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

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

も必要になります。

### クラスの定義

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

```python
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__`とほぼ同じことをするだけです。ですから実際は

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

```python
@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!!"と出てくるだけですが、これで使えるようになりました。

```python
np.mean(arr1)
```
```
dispatched!!
Out: 0.816496580927726
```

# 終わりに

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

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