import numpy as np

def eoft(trmat,nm=None):

'''

evals_perc,evecs_norm,amp = eoft(trmat)

Computes eof in time (or depth) for general cases

normally trmat is a matrix containing each time series

(or vertical profiles) as a row.

That is, for N stations having m data points,

trmat is a N by m matrix.

'''

#if is masked the data

try:

Trmat = trmat.data

Trmat[trmat.mask] = np.nan

trmat = Trmat.copy()

except:

None

#% demeans trmat prior to doing eof

trmat = (trmat.T-np.nanmean(trmat,axis=1)).T

#% computes zero-lag cross-covariance matrix

mcov = np.dot(trmat,trmat.T)/(trmat.shape[1]-1)

#% computes eigenvalues and eigenvectors

evals,evecs = np.linalg.eig(mcov)

#find the descending order of eigenvalues

ind = np.argsort(evals)[::-1]

#% normalize eigenvectors

evecs_norm = evecs[:,ind]/np.linalg.norm(evecs[:,ind],axis=0)

#% sort eigenvalues and computes percent variance explained by each mode

evals_perc = evals[ind]/np.sum(evals)

#% computes amplitude functions

amp = np.dot(evecs_norm.T,trmat)

#if was chosen a number of eigenvectors

if nm!=None:

evecs_norm = evecs_norm[:,:nm]

evals_perc = evals_perc[:nm]

amp = amp[:nm,:]

return evals_perc,evecs_norm,amp

def my_eof_interp(M,nmodes,errmin=1e-15,repmax=None):

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% vi = my_eof_interp(M,nmodes)

#% This function try to fill gappy data using its EOFs

#%

#% M is the data Matrix (generaly the diferent stations

#% are on the columns);

#%

#% nmodes is the number of statiscally significant EOF modes (to be used for the series reconstruction)

#%

#%% The method -- based on Beckers & Rixen (2003, JAOTech)

#%

#% The algorithm proposed by the authors is as follows. (see Eqs. (13) and (14) from Beckers & Rixen)

#%

#% (I) Put zero in the GAPPY positions ("unbiased INITIAL GUESS, by assuming that the removed mean was unbiased).

#%

#% (II) Then you can have a first estimative of the EOFs.

#%

#% (III) You can reconstruct the series using the statiscally significant modes (e.g.: determined using Monte Carlo

#% simulations as in Preisendorfer [1988] -- OBS: the authors propose another method using cross validation

#% during the iterative process)

#%

#% (IV) Substitute the reconstructed series ONLY in the GAPPY positions in your original matrix.

#%

#% (V) Run steps (II), (III) and (IV) iteractively till its convergence (i.e. compare the reconstructed series at time

#% t+1 with that at time t).

#%

#% OBS: It doesn't work for filling gaps simultaneously in all depths.

Mmean = np.nanmean(M,axis=1)

M = M.T - np.nanmean(M,axis=1)

M = M.T

f = np.isnan(M)

M[f] = 0 #%% initial guess in gappy positions

rep = 0

while True:

print(rep+1)

Ud,Dd,Vd = np.linalg.svd(M,'econ')

Dd = np.diagflat(Dd)

#% compute the explained variance (total variance is the sum of the Dd elements squared -total energy-)

evd = np.diag(Dd)**2 #%% remember that the singular values (Dd) are the sqrt of

evd = Dd/np.sum(Dd) #%% the eigenvalue of the cov. matrix.

#%% truncated EOF reconstruction

Ud = Ud[:,0:nmodes]

Vd = Vd[:,0:nmodes]

Dd = Dd[0:nmodes,0:nmodes]

#Dd = Dd[0:nmodes]

#% reconstructing the series -- truncated reconstruction (SVD: Mrec = Ud * Dd *Vd')

Xa = np.dot(np.dot(Ud,Dd),Vd.T)

Mold = M.copy() #% saving for compute the error

M[f] = Xa[f]

#% testing the error

if rep>=1:

err = np.nanmean(np.abs(M - Mold))/(np.nanmax(M)-np.nanmin(M))

print(err)

if err < errmin:

break

if repmax!=None:

if rep>=repmax:

break

rep += 1

#%% Saving the last the series with the EOF reconstructed series in the gappy positions;

vi = (M.T+Mmean).T

return vi