import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D
import xarray as xr
import datetime
from pandas.plotting import register_matplotlib_converters #これを描かないと警告文がでる（オマジナイ）
register_matplotlib_converters()
# データの読み込み
# 以下を参考にしました
# https://arg.usask.ca/docs/LOTUS_regression/_modules/LOTUS_regression/predictors/download.html
file = "data/qbo_1953_2021.dat"
#--> "data/dd/qbo_1953_2021.dat"に変更のこと
def date_parser(s):
s = int(s)
return datetime.datetime(2000 + s // 100 if (s // 100) < 50 else 1900 + s // 100, s % 100, 1)
dataORG = pd.read_fwf(file,
skiprows=92, header=None,
colspecs=[(0, 5), (6, 10), (12, 16), (19, 23), (26, 30), (33, 37), (40, 44), (47, 51), (54, 58)],
delim_whitespace=True, index_col=1, parse_dates=True, date_parser=date_parser,
names=['station', 'month', '70', '50', '40', '30', '20', '15', '10'])
#dataORG.index = dataORG.index.to_period(freq='M')
#最初の方（1950年代は歯抜けが多いので、skiprows=92として読み飛ばしています）
data_sub = dataORG.loc['1960-01':'2020-12',:].copy()
data_sub.drop('station', axis=1, inplace=True)
data_sub /= 10.0 # units: m/s に変換
print(data_sub.head()) #列方向は気圧座標が入っている（高度に換算すると 70 hPa~20 km, 10 hPa~ 30 km, に相当）
年々変動に注目するので、季節変化成分を予め取り除いておく。
色々やり方はあるが、各月でコンポジットを取って月平均の東西風（の高度分布）を算出し、オリジナルデータから差し引くのが簡単（他にも調和関数で季節変化を定義して取り除く方法などがある）。
コンポジットについては、例えば演習DBの内容を思い出してください(演習DB [2.時系列データの解析] )
軸データやデータ行列は以下のように取り出せる（前回の復習）:
data_sub.index: 行方向のインデックス（上のデータの場合、時刻）
data_sub.columns: 列方向のインデックス（上のデータの場合、気圧座標）
data_sub.values: データ行列（軸は落とす）
【Fig.3】 行列の積を計算する際に、縦ベクトルと横ベクトルを区別したいときがある。uを大きさnの配列とすると：
uh = np.reshape(u, (1,n)):横ベクトルにするuv = np.reshape(u, (n,1)):縦ベクトルにするFig.3: 横軸：時刻(x)、縦軸(y):気圧座標として、東西風速(z)をコンター表示
ax.contourf(x,y,z) (z = z(y, x))Fig.4: PC1 (横軸）➖ PC2（縦軸）の散布図
ax.scatter(x,y,....)Fig.5: 横軸: 緯度(x)、縦軸(y):気圧座標として、回帰係数をコンター表示
月平均東西風データはdata/ncep.reanalysis.derived/pressure/uwnd.mon.mean.ncを使用。
データの読み込み・切り出しは、演習DBの内容を参考にしてください（演習DB[3.全球大気再解析データの解析] )
以下、データの読み込み例（の一部）を示す:
file = "Fortran/data/uwnd.mon.mean.nc"
#月平均東西風データを指定
#"data/ncep.reanalysis.derived/pressure/"以下に変更
data = xr.open_dataset(file)
print(data)
uwnd = data['uwnd'] #東西風(uwnd)読み込み。この段階ではまだ軸が付いているので、以下解析の為にそれらを分離する
plevel = uwnd['level'].values #軸データ：気圧座標
lat = uwnd['lat'].values # 軸データ：緯度
print('Pressure levels = ', plevel)
print('Latitudes = ',lat)
uwnd_v = uwnd.values #データ行列(東西風の値のみ)