#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Feb 25 11:57:10 AM EST 2026
# https://ajdawson.github.io/eofs/examples/elnino_xarray.html
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
from eofs.xarray import Eof
from eofs.examples import example_data_path
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
# Read SST anomalies using the xarray module. The file contains November-March
# averages of SST anomaly in the central and northern Pacific.
filename = example_data_path('sst_ndjfm_anom.nc')
sst = xr.open_dataset(filename)['sst']

# Create an EOF solver to do the EOF analysis. Square-root of cosine of
# latitude weights are applied before the computation of EOFs.
coslat = np.cos(np.deg2rad(sst.coords['latitude'].values))
wgts = np.sqrt(coslat)[..., np.newaxis]
solver = Eof(sst, weights=wgts)

# Retrieve the leading EOF, expressed as the correlation between the leading
# PC time series and the input SST anomalies at each grid point, and the
# leading PC time series itself.
eof1 = solver.eofsAsCorrelation(neofs=1)
pc1 = solver.pcs(npcs=1, pcscaling=1)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #import cartopy.crs as ccrs
    #import cartopy.feature as cfeature
    from geoplots import mapplot
    # Plot the leading EOF expressed as correlation in the Pacific domain.
    clevs = np.linspace(-1, 1, 11)
    #ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=190))
    #fill = eof1[0].plot.contourf(ax=ax, levels=clevs, cmap=plt.cm.RdBu_r,
    #                             add_colorbar=False, transform=ccrs.PlateCarree())
    eof1[0].plot.contourf(levels=21)
    #ax.add_feature(cfeature.LAND, facecolor='w', edgecolor='k')
    #cb = plt.colorbar(fill, orientation='horizontal')
    #cb.set_label('correlation coefficient', fontsize=12)
    mapplot()
    ax = plt.gca()
    ax.set_title('EOF1 expressed as correlation', fontsize=16)
    
    # Plot the leading PC time series.
    plt.figure()
    #pc1[:, 0].plot(color='b', linewidth=2)
    pc1[:, 0].plot()
    ax = plt.gca()
    ax.axhline(0, color='gray', ls='--')
    ax.set_ylim(-3, 3)
    ax.set_xlabel('Year')
    ax.set_ylabel('Normalized Units')
    ax.set_title('PC1 Time Series', fontsize=16)
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()