【python3】pandas で平均と共分散行列をjson に保存

前書き

いつもエイや！でコードを書いていて平均と共分散を出力するのも何度も行っているのですが、メモしておきます。本当に中身は単純です。ひねりはありません。すみません。

状況：pandas データがあり、特定のデータX,Y,Z についての平均と共分散を求めたい

方針

1. x,y,z のみを含むDataFrame を作成する
2. 作成したDataFrame の平均と共分散を求める。
3. json ファイルに出力する。

実装

ポイントは、

• DataFrame でデータ名を指定して抽出するには、.locを使います。
• DataFrame には mean() も cov() も corr() もあり、出力はDataFrame だったりするので、values でndarray を参照します。
• 辞書にlist として登録するには、ndarray　のtolist() を利用できます。

from pandas import DataFrame
from numpy import random
import json

df = DataFrame(random.randint(0,100,size=(252, 4)), columns=list('XYZW'))
output_data = dict()

# 1. extract XYZ
df_xyz = df.loc[:,list("XYZ")]

# 2-1 mean vector
u = df_xyz.mean()
output_data["mean"] = u.values.tolist()

# 2-2 covariance
s = df_xyz.cov()
output_data["covariance"] = s.values.tolist()

# 3
with open("out.json", 'w') as f:
    json.dump(output_data, f, indent=2)

出力したjson ファイルは

{
  "mean": [
    48.34126984126984,
    50.52777777777778,
    51.492063492063494
  ],
  "covariance": [
    [
      877.6360589388478,
      -44.88202744577245,
      -71.94548788971099
    ],
    [
      -44.88202744577245,
      876.4733289065962,
      -32.312527667109336
    ],
    [
      -71.94548788971099,
      -32.312527667109336,
      784.7768291911716
    ]
  ]
}

です。

おわりに

この実装にたどり着くまで、意外と調べました。（汗）
DataFrame のCovariance はAPI document (ここ)に書いてあります。

(2020/05/11)

その後

• NaN (not a number)の処理を書きたいとき、if文をこういう感じにして確認できる。

   from numpy import isnan
   if isnan(x).any():
       x = zeros(3)
   if isnan(S).any():
       S = zeros( (3,3) )