import numpy as np
from supersvd import supersvd
import matplotlib.pyplot as plt

ts = np.fromfile('ts.std', dtype=np.float32).reshape(1147, 28, 31)
ps = np.fromfile('ps.std', dtype=np.float32).reshape(1147, 28, 31)

w = np.zeros((28, 31), dtype=np.float32)
lat = np.deg2rad(35 + 1.5 * np.arange(31))
w[:, :] = np.cos(lat).reshape(1, -1)

svd = supersvd(ts, ps, 4, WX=w, WY=w)

first_mode_ts = svd.x_vect[0]

plt.matshow(first_mode_ts)
plt.savefig('first_mode_ts.png')

first_ps_timeseries = svd.x_coeff[0]

plt.figure()
plt.plot(first_ps_timeseries)
plt.savefig('first_ps_timeseries.png')

print('Correlation coefficient between first_ts_timeseries and first_ps_timeseries:', svd.corrcoeff[0])