多次元データ解析ライブラリ xarrayで複素数を利用する際の注意点　（つまずいた点)

はじめに

どうも，こんにちは．かわつあんです．

これまで私はpandasモジュールを使ってデータの管理を行ってきましたが，最近になってxarrayモジュールを使い始めました．個人的にはpandasよりもデータの管理がしやすく気に入っています．

xarray.pydata.org

特にpandasと比べて優れている点は，

標準で2つ以上の軸に対応していること
軸自体に値を代入できること

だと個人的に思っています．

xarrayに関する導入についてはxarray へのcontributorの方が日本語でレビューしてくださっていますので，以下のリンクより確認ください．

qiita.com

私はデータ解析では主にフーリエ変換などを利用するために複素数のデータが多いのですが, xarrayで複素数を扱おうとした際に少しつまずいたことがありました．

この記事でこのつまずきについて共有しようと思います．

ちなみにこの記事ではバージョン0.15.1を利用しています．

DataArrayを複素数配列でインスタンス化

xarrayのDataArrayを使って複素数配列でインスタンス化してみます．今回は簡単のため3x3の複素数データを扱います．

import numpy as np
import xarray as xr
complex_array1 = np.array([
    [1 + 5.j, 2 + 4.j,3 + 3.j],
    [2 + 4.j, 3 + 3.j,4 + 2.j],
    [3 + 3.j, 4 + 2.j,5 + 1.j],
])

da1 = xr.DataArray(
    complex_array1,
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
)
print(da1)

出力結果

<xarray.DataArray (t: 3, x: 3)>
array([[1.+5.j, 2.+4.j, 3.+3.j],
       [2.+4.j, 3.+3.j, 4.+2.j],
       [3.+3.j, 4.+2.j, 5.+1.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

このように，インスタンス化については直感的に利用できます． DataArrayは出力結果が綺麗でいいですね．特に軸のデータも一緒に見れるのが好きです．

DataArrayに複素数を代入する時につまずいた点)

インスタンス化するのはいいのですが，代入したりする時は注意が必要です．というのもインスタンス化する時に，DataArrayの中に入れるデータの型を明示しないといけないようです．

ダメな例

da2 = xr.DataArray(
    np.zeros((3,3)),
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
)
print(da2)

da2[0,0] = 1 + 1j
print("---after---")
print(da2)

出力結果（ダメな例）

<xarray.DataArray (t: 3, x: 3)>
array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---after---
<xarray.DataArray (t: 3, x: 3)>
array([[1., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

要素が0の配列でインスタンス化したのちに[0,0]成分に1 + 1jを代入していますが，代入されたDataArrayを出力しても実数として反映されちゃっています．

これに気づかずに解析したら地獄をみますね...

良い例: dataの型を明示する

da3 = xr.DataArray(
    np.zeros((3,3)),
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
).astype("complex")       #<======  データの型を明示
print(da3)

da3[0,0] = 1 + 1j
print("---after---")
print(da3)

出力結果(良い例)

<xarray.DataArray (t: 3, x: 3)>
array([[0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---after---
<xarray.DataArray (t: 3, x: 3)>
array([[1.+1.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

このように代入した値もきちんと反映されています．インスタンス化してすぐのデータもきちんと複素数型になっています．

ちなみに，以下のように書いても問題なく複素数を扱えます．

da4 = xr.DataArray(
    np.zeros((3,3), dtype="complex"),  #<======  データの型を明示
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
)       
print(da4)

da4[0,0] = 1 + 1j
print("---after---")
print(da4)

出力結果

<xarray.DataArray (t: 3, x: 3)>
array([[0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---after---
<xarray.DataArray (t: 3, x: 3)>
array([[1.+1.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

np.zerosで要素が0の配列を作る時点でデータ型を複素数にしてしまえば，当然代入でも複素数は使えます．

要はインスタンス化の時点でDataArrayに入れるデータ型が決まってしまうということですね．

以上のつまずいた点はpython以外ではメモリの確保等のために型を明示するため常識だと思いますが，pythonで楽しようと思ってぽけーっと解析していたのでこのようにつまずきました．．．危なかった．．．

おまけ1（複素数の操作）

複素数といえば，複素共役や実部虚部を取り出したいときありますよね．これらの操作はnumpyのndarrayとほぼ同様にDataArrayで利用できます．

complex_array2 = np.array([
    [1 + 5.j, 2 + 0.j,3 + 3.j],
    [2 + 0.j, 3 + 3.j,4 + 0.j],
    [3 + 3.j, 4 + 0.j,5 + 1.j],
])                              

da5 = xr.DataArray(
    complex_array2,
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
)

print(da5)
print("---complex conjgate ( .conj ) ---")
print(da5.conj)
print("---complex conjgate ( .conj() ) ---")
print(da5.conj())
print("---imaginary part---")
print(da5.real)
print("---real part---")
print(da5.imag)

出力結果

<xarray.DataArray (t: 3, x: 3)>
array([[1.+5.j, 2.+0.j, 3.+3.j],
       [2.+0.j, 3.+3.j, 4.+0.j],
       [3.+3.j, 4.+0.j, 5.+1.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---complex conjgate ( .conj ) ---
<bound method _method_wrapper.<locals>.func of <xarray.DataArray (t: 3, x: 3)>
array([[1.+5.j, 2.+0.j, 3.+3.j],
       [2.+0.j, 3.+3.j, 4.+0.j],
       [3.+3.j, 4.+0.j, 5.+1.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1>
---complex conjgate ( .conj() ) ---
<xarray.DataArray (t: 3, x: 3)>
array([[1.-5.j, 2.-0.j, 3.-3.j],
       [2.-0.j, 3.-3.j, 4.-0.j],
       [3.-3.j, 4.-0.j, 5.-1.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---imaginary part---
<xarray.DataArray (t: 3, x: 3)>
array([[1., 2., 3.],
       [2., 3., 4.],
       [3., 4., 5.]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1
---real part---
<xarray.DataArray (t: 3, x: 3)>
array([[5., 0., 3.],
       [0., 3., 0.],
       [3., 0., 1.]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

このようにndarrayと同じようにDataArrayでメソッドや属性(attribute)で呼べば問題なく利用できます．

ただし，複素共役を扱う時は少し注意が必要です．ndarrayの場合ndarray.conjのように属性として複素共役を呼ぼうとしたらエラーがでますが，DataArrayはDataArray.conjを出力しようとするとエラーはでず，しかも複素共役ではなく元のデータが呼ばれてしまうようです．これはバグかもしれないので，githubで報告してみます． ==> すみません勘違いでした!

numpy.ndarary.conjもxarray.DataArray.conjもそれぞれ属性は存在せずメソッドのみ存在しているようです．出力だけではエラーはでず，numpy.ndarary.conjの場合はのようにメソッドですよという内容が出力される一方で，xarrayはメソッドですよって知らせてくれる内容に加えてDataArrayの中身も同時に出力してくれるみたいです．これですこしこんがらがってしまいました．ちなみにnumpy.ndarary.conjやxarray.DataArray.conjに何かしらの演算をしようとするとちゃんとエラーが出ますのでご安心ください．

おまけ2（複素数を含むデータの保存）

DataArrayでデータを保存する時はnetCDFという形式のファイルで保存するそうです．拡張子は.ncです．実数の場合は，以下の通り書きます

da6 = xr.DataArray(
    np.zeros((3,3)),
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
)
da6.to_netcdf("da6.nc")

複素数の場合，以下の通り少し工夫が必要なようです．

da7 = xr.DataArray(
    np.zeros((3,3)),
    dims = ["t", "x"],
    coords = {
        "t":[0, 1, 2],
        "x":[-1, 0, 1],}
).astype("complex")

# da7.to_netcdf("da7.nc")
# >>ValueError: NetCDF 3 does not support type complex128

da7.to_netcdf("da7_1.nc", engine="h5netcdf",invalid_netcdf=True)

netCDFにもバージョンがあって，対応するバージョンを指定しないといけないようですね．これにはpythonのモジュール"h5netcdf"が必要なのでpipなどでインストールしてください．詳しくは公式ドキュメントを参照してください．

xarray.pydata.org

.ncデータを読み込む時はopen_dataarrayを使います．保存した2つのデータを読み込んでみます．

da6_read = xr.open_dataarray("da6.nc")
da7_read = xr.open_dataarray("da7.nc", engine="h5netcdf")
print(da6_read)
print("------")
print(da7_read)

出力結果

<xarray.DataArray (t: 3, x: 3)>
array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]])
Coordinates:
  * t        (t) int32 0 1 2
  * x        (x) int32 -1 0 1
------
<xarray.DataArray (t: 3, x: 3)>
array([[0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
Coordinates:
  * t        (t) int64 0 1 2
  * x        (x) int64 -1 0 1

読み込みの際にも，オプションでengine="h5netcdf"を指定しないといけないのが注意点ですね．