def plot_hfdamping(plotlag,
plotmon,
lon,
lat,
damping,
msk,
ax=None,
cints=None,
select_id=True):
# damping : [mon x lag x lat x lon]
# msk : [mon x lag x lat x lon]
if cints is None:
cints = np.arange(-50, 55, 5)
if ax is None:
ax = plt.gca()
# Select what to plot
if select_id:
plotvar = damping[plotmon, plotlag, :, :]
plotmsk = msk[plotmon, plotlag, :, :]
else:
plotvar = damping
plotmsk = msk
# Flip Longitude
lon1, plotvar1 = proc.lon360to180(lon, plotvar.T)
_, plotmsk1 = proc.lon360to180(lon, plotmsk.T)
# Plot the contours
pcm = ax.contourf(lon1,
lat,
plotvar1.T,
levels=cints,
cmap='cmo.balance',
extend='both')
# Plot significant points
viz.plot_mask(lon1,
lat,
plotmsk1,
reverse=False,
ax=ax,
markersize=.75,
color='gray')
return ax, pcm
def preproc_CESMPIC(sst):
# Apply Land/Ice Mask
mask = np.load(datpath + "landicemask_enssum.npy")
sst = sst * mask[None, :, :]
# Adjust dimensions [time x lat x lon] --> [lon x lat x time]
sst = sst.transpose(2, 1, 0)
# Flip longitude
st = time.time()
lon180, sst = proc.lon360to180(lon360, sst)
print("Flipped Longitude in %.2fs" % (time.time() - st))
# Remove monthly anomalies
st = time.time()
nlon, nlat, ntime = sst.shape
sst = sst.reshape(nlon, nlat, int(ntime / 12), 12)
ssta = sst - sst.mean(2)[:, :, None, :]
print("Deseasoned in %.2fs" % (time.time() - st))
print("Mean was %e" % (np.nanmax(ssta.mean(2))))
ssta = ssta.reshape(nlon, nlat, int(ntime / 12) * 12)
return ssta, lon180
def postprocess(dampingmasked,limask,sellags,lon):
# Inputs
## Dampingmasked [month x lag x lat x lon]
## limask [lat x lon]
# Select lags, apply landice mask
mchoose = dampingmasked[:,sellags,:,:] * limask[None,:,:]
# Flip longiude coordinates ([mon lat lon] --> [lon x lat x mon])
lon1,dampingw = proc.lon360to180(lon,mchoose.transpose(2,1,0))
# Multiple by 1 to make positive upwards
dampingw *= -1
return dampingw
6,
figsize=(16, 16),
facecolor='white',
subplot_kw={'projection': ccrs.PlateCarree(central_longitude=0)})
cints = np.arange(-50, 55, 5)
for e in tqdm(range(42)):
ax = axs.flatten()[e]
ax.coastlines()
ax.set_extent(bboxplot)
plotvar = damping[e, plotlag, plotmon, :, :]
plotmsk = mall[e, plotlag, plotmon, :, :]
# Flip Longitude
lon1, plotvar1 = proc.lon360to180(lon, plotvar.T)
_, plotmsk1 = proc.lon360to180(lon, (plotmsk * limask).T)
#_,plotmsk1 = proc.lon360to180(lon,limask.T)
pcm = ax.contourf(lon1,
lat,
plotvar1.T,
levels=cints,
cmap='cmo.balance',
extend='both')
# Plot significant points
viz.plot_mask(lon1,
lat,
plotmsk1,
reverse=False,
regions = ("SPG", "STG", "TRO", "NAT")
bboxes = (bbox_SP, bbox_ST, bbox_TR, bbox_NA)
regioncolor = ('b', 'r', [0, 1, 0], 'k')
nregion = len(regions)
lags = np.arange(0, 61, 1)
#%% Check STDEV of Forcing
cesmforce = xr.open_dataset(
datpath + 'NHFLX_ForcingStats_Converted_NotMonthly_hvarmode2_ensavg.nc')
fstd = cesmforce.fstd.values
clon = cesmforce.lon.values
clat = cesmforce.lat.values
clon1, fstd = proc.lon360to180(clon, fstd.transpose(1, 0))
# Get values for different regions
bbox_SP = [-60, -15, 40, 65]
bbox_ST = [-80, -10, 20, 40]
bbox_TR = [-75, -15, 0, 20]
bbox_NA = [-80, 0, 0, 65]
regions = ("SPG", "STG", "TRO", "NAT")
bboxes = (bbox_SP, bbox_ST, bbox_TR, bbox_NA)
stdboxes = []
stdval = []
for r in range(4):
# Select data from region
# Load in the data
naoforcing = np.load(input_path+"NAO_EAP_NHFLX_ForcingDJFM.npy") #[PC x Ens x Lat x Lon]
naof_hist = naoforcing.mean(1).squeeze()
naof_hist = naof_hist.transpose(2,1,0)# --> [lon x lat x pc]
npzdata = np.load(datpath+"Manual_EOF_Calc_NAO_corr.npz")
eofall = npzdata['eofall'] # [ens x lat x lon x pc]
nao_hist = eofall.mean(0)
nao_hist = nao_hist.transpose(1,0,2) # --> [lon x lat x pc]
varexpall = npzdata['varexpall'] # [ens x pc]
varexp_hist = varexpall.mean(0)
# Apply Mask
naof_hist *= msk.T[:,:,None]
# Flip longitude
lon180,naof_hist = proc.lon360to180(lon360,naof_hist)
_,nao_hist = proc.lon360to180(lon360,nao_hist)
# Visualize Results
if debug:
cint = np.arange(-50,55,5)
fig,axs = plt.subplots(2,1,subplot_kw={"projection":ccrs.PlateCarree()})
ax = axs[0]
ax = viz.init_map(bbox,ax=ax)
cf1 = ax.contourf(lon180,lat,naof_hist[:,:,0].T,levels=cint,cmap=cmbal)
ax = viz.plot_contoursign(naof_hist[:,:,0].T,lon180,lat,cint,ax=ax,clab=True,lw=0.5)
fig.colorbar(cf1,ax=ax)
ax.set_title("$NAO_{DJFM}$ (EOF 1, Variance Explained = %.2f)"% (varexp_hist[0]*100)+r"%")
#%% #%% Load in TAU files and process tauname = "tauall_1940-2015_40ens_lag3.nc" dstau = xr.open_dataarray(datpath+tauname) tlon = dstau['lon'].values tlat = dstau['lat'].values tau = dstau.values #[mon x lat x lon x ens] # Take ensemble mean and flip longitude coordinates tauavg = np.nanmean(tau,3) tauavg = np.transpose(tauavg,(2,1,0)) #lon x lat x mon tlon1,tauavg = proc.lon360to180(tlon,tauavg) tauavg = np.transpose(tauavg,(2,1,0)) # back to mon x lat x lon tauavg[~np.isfinite(tauavg)] = np.nan #%% Load in MLD and Damping data... # Load Damping Data dampmat = 'ensavg_nhflxdamping_monwin3_sig020_dof082_mode4.mat' loaddamp = loadmat(damppath+dampmat) lon = np.squeeze(loaddamp['LON1']) lat = np.squeeze(loaddamp['LAT']) damping = loaddamp['ensavg'] #[lon x lat x mon] #damping = np.transpose(damping,(2,1,0)) # [mon x lat x lon] # Load MLD Data (preprocessed in prep_mld.py) mld = np.load(datpath+"HMXL_hclim.npy") # Climatological MLD
subplot_kw={'projection': ccrs.PlateCarree()})
cints = np.arange(-1, 1.1, .1)
plotlag = 1
plotmon = 0
plotmcf = 1
imsk = 1 # 0 = SST, 1 = Crosscorr
ax.coastlines()
#ax.set_extent(bboxplot)
plotvar = rhovalues[imsk][plotmcf][plotlag, plotmon, :, :]
plotmsk = rmasks[plotmon, plotlag, :, :, plotmcf, imsk]
# Flip Longitude
lon1, plotvar1 = proc.lon360to180(lon, plotvar.T)
_, plotmsk1 = proc.lon360to180(lon, plotmsk.T)
#_,plotmsk1 = proc.lon360to180(lon,limask.T)
plotvar1 *= limask180
plotmsk1 *= limask180
pcm = ax.contourf(lon1,
lat,
plotvar1.T,
levels=cints,
cmap='cmo.balance',
extend='both')
# Plot significant points
viz.plot_mask(lon1, lat, plotmsk1, reverse=False, ax=ax, markersize=.75)
proj = ccrs.PlateCarree()
fig, axs = plt.subplots(4,
3,
subplot_kw={'projection': proj},
figsize=(12, 12),
constrained_layout=True)
plotmon = np.roll(np.arange(0, 12), 1)
for im in range(12):
monid = plotmon[im]
ax = axs.flatten()[im]
plotvar = damping[monid, il, :, :]
lon1, plotvar1 = proc.lon360to180(lon, (plotvar.T)[..., None])
blabel = [0, 0, 0, 0]
if im % 3 == 0:
blabel[0] = 1
if im > 8:
blabel[-1] = 1
ax = viz.add_coast_grid(ax,
bbox=[-80, 0, -10, 62],
fill_color='gray',
blabels=blabel,
ignore_error=True)
pcm = ax.contourf(lon1,
lat,
plotvar1.squeeze().T * -1,
frames = 12 #Indicate number of frames
figsize = (8, 6)
vm = [-5, 5]
interval = 0.1
bbox = [-80, 0, 0, 80]
fps = 10
savetype = "mp4"
dpi = 100
yrstrs = []
for i in tqdm(range(ntime - 240)):
rng = np.arange(i, i + winsize + 1)
yrstr = "%s to %s" % (timesmon[rng[0]][:4], timesmon[rng[-1]][:4])
yrstrs.append(yrstr)
lon180, remap180 = proc.lon360to180(lon5, remapclusts.transpose(2, 1, 0))
invar = remap180.transpose(1, 0, 2)
#invar = remapclusts.transpose(1,2,0) # [lat x lon x time]
def make_figure():
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
ax = viz.add_coast_grid(ax)
return fig, ax
start = time.time()
fig, ax = make_figure() # Make the basemap
pcm = ax.pcolormesh(lon180, lat5, invar[..., i], cmap=cmapn)
outpath = projpath + '02_Figures/20200820/'
# Path to SST data from obsv
datpath2 = "/Users/gliu/Downloads/02_Research/01_Projects/01_AMV/00_Commons/01_Data/"
#%% Load/Process Data
# Load NAO Forcing DJFM (Monthyl)
naoforcing = np.load(datpath +
"Monthly_NAO_Regression.npy") #[Ens x Mon x Lat x Lon]
NAO1 = np.nanmean(naoforcing, axis=0) # Take PC1, Ens Mean and Transpose
NAO1 = np.transpose(NAO1, (2, 1, 0))
# Convert Longitude from degrees East
lon360 = np.load(datpath + "CESM_lon360.npy")
lon180, NAO1 = proc.lon360to180(lon360, NAO1)
# Load NAO Forcing Seasonally Varying
naomon = np.load(datpath + "NAO_Monthly_Regression_PC.npz")
pcs = naomon['pcall']
naomon = naomon['eofall']
naomon = np.nanmean(naomon, 0) # Take ens mean
naomon = np.transpose(naomon, (2, 1, 0))
_, naomon = proc.lon360to180(lon360, naomon)
# Create quick nan mask and apply to NAO forcing
nanmask = np.sum(naomon, 2)
nanmask = np.where(~np.isnan(nanmask), np.ones(nanmask.shape), np.nan)
NAO1 = NAO1 * nanmask[:, :, None]
#%% Load data from CESM-FULL Historical
ds = xr.open_dataset("/Users/gliu/Downloads/02_Research/01_Projects/01_AMV/01_hfdamping/01_Data/allens_nhflxdamping_monwin3_sig020_dof082_mode4_lag1.nc")
lbdpt = ds.sel(lon=330,lat=50,method='nearest')
dampfull = lbdpt.NHFLX_Damping.values
# Load all data in for basinwide comparison
dampfullall = ds.NHFLX_Damping.values # Lat x Lon360 x Mon x Ens
# Flip the longitude around
nlat,nlon,nmon,nens = dampfullall.shape
dampfullall = dampfullall.reshape(nlat,nlon,nmon*nens).transpose(1,0,2)
_,dampfullall = proc.lon360to180(lon,dampfullall)
dampfullall = dampfullall.reshape(nlon,nlat,nmon,nens)
#%% Plot the comparison at the selected point
fig,ax = plt.subplots(1,1)
for i in range(42):
ax.plot(dampfull[:,i],label="",color='gray',alpha=0.25)
ax.plot(dampfull[:,-1],label="HTR-FULL (member)",color='gray',alpha=0.25)
ax.plot(dampfull.mean(1),label="HTR-FULL (ens-avg)",color='k',alpha=1)
for i in range(2):
div = 1
# if i == 1:
# div = 2
ax.plot(dampingfin[i][klon,klat,:]/div,label=mconfigs[i])
'ensemble': np.arange(1, 43, 1)
},
name="NHFLX_Damping")
da.to_netcdf(datpath +
"allens_nhflxdamping_monwin3_sig020_dof082_mode4.nc")
# Take ensemble and lag average
if len(lags) > 1:
dampseason = np.mean(dampchoose, (2, 4))
dampseason1 = np.mean(dampchoose, (2, 4))
else:
dampseason = np.mean(dampchoose, (2))
# Save data if option is set
if savevar == 1:
LON1, ensavg = proc.lon360to180(lon, dampseason)
# Save variables, matching matlab script output
savedict = {'ensavg': ensavg, 'LON1': LON1, 'LAT': lat}
outname = "%sensavg_%sdamping_monwin%i_sig%03d_dof%03d_mode%i.mat" % (
datpath, flux.lower(), monwin, p * 100, dof, mode)
savemat(outname, savedict)
print("Saved data to %s" % outname)
#%% Make some frequency plots
mchoose = mtot[:, :, :, :, lagidx]
if len(lags) > 1:
mfreq = np.nansum(mchoose, (2, 3, 4))
else:
mfreq = np.nansum(mchoose, (2, 3))
# DJFM NHFLX Regressed to DJFM NAOIndex
if funiform == 2:
naoforce = np.load(datpath + "NAO_Forcing_EnsAvg.npy") #lon x lat
# Monthly NHFLX regressed to DJFM NAOIndex
elif funiform == 3:
naoforcing = np.load(
datpath + "Monthly_NAO_Regression.npy") #[Ens x Mon x Lat x Lon]
NAO1 = np.nanmean(naoforcing, axis=0) # Take PC1, Ens Mean and Transpose
NAO1 = np.transpose(NAO1, (2, 1, 0))
# Convert Longitude from degrees East
lon360 = np.load(datpath + "CESM_lon360.npy")
lon180, NAO1 = proc.lon360to180(lon360, NAO1)
naoforce = NAO1 * -1 #lon x lat x mon
# Monthly NHFLX regressed to Monthly NAOIndex
elif funiform == 4:
lon360 = np.load(datpath + "CESM_lon360.npy")
naomon = np.load(datpath + "NAO_Monthly_Regression_PC.npz")
naomon = naomon['eofall']
naomon = np.nanmean(naomon, 0) # Take ens mean
naomon = np.transpose(naomon, (2, 1, 0))
_, naoforce = proc.lon360to180(lon360, naomon) #lon x lat x mon
else:
naopt = 1
# # Load Forcing and take ensemble average
# naoforcing = np.load(datpath+"NAO_Monthly_Regression_PC123.npz")['flxpattern'] #[Ens x Mon x Lat x Lon]
# # Take ensemble average, then sum EOF 1 and EOF2
# NAO1 = naoforcing[:,:,:,:,:2].mean(0) # [ PC x Mon x Lat x Lon]
# # Temporarily reshape to combine PC and mon
# if funiform > 6:
# NAO1 = NAO1.reshape(24,192,288) # NOTE: need to uncombine later
# Transpose to [Lon x Lat x Time]
NAO1 = np.transpose(NAO1,(2,1,0))
# Convert Longitude to Degrees East
lon180,NAO1 = proc.lon360to180(lon360,NAO1)
# Restrict to region
NAO1,_,_ = proc.sel_region(NAO1,LON,LAT,bboxsim)
else: # For funiform= uniform or random forcing, just make array of ones
NAO1 = np.ones(hclim.shape)
# for ____ NAO1.shape = ...
# funiform1 --> 37x27x12, array of ones
# funiform2 --> 37x27x1, fixed pattern
# funiform3 --> 37x27x12, Monthly Pattern
# funiform5 --> 37x27x1, fixed pattern (EAP-DJFM)
# funiform6 --> 36x27x2, fixed pattern (EAP and NAO like forcing)
# ----------------------------------------------
# Preprocess ACs
cesmslabac1 = cesmslabac[:,lags,:,:].transpose(0,1,3,2) # [mon x lag x lon x lat]
cesmacs = [cesmslabac1,cesmfullac]
cesmacproc = []
for ac in cesmacs:
# Make into [lon x lat x otherdims]
_,nlag,nlon,nlat = ac.shape
ac = ac.reshape(12*nlag,nlon,nlat)
ac = ac.transpose(1,2,0)
# Flip longitude
lon1,acflip=proc.lon360to180(lon360,ac)
if debug:
plt.pcolormesh(acflip[:,:,22].T)
plt.show()
# Restrict to region, apply mask
acreg,_,_ = proc.sel_region(acflip,lon1,lat,bboxsim)
acreg *= msk[:,:,None]
# Make into [month x lag x lon x lat]
acreg = acreg.transpose(2,0,1)
acreg = acreg.reshape(12,nlag,nlonr,nlatr)
if debug:
plt.pcolormesh(acreg[1,11,:,:].T),plt.colorbar()
plt.show()
dampingmasked = dampingt.copy()
#%% Further Processing
# Average over selected lags
ilags = [l - 1 for l in lags] # get indices
dampingt = dampingt[:, ilags, :, :].mean(1).squeeze()
lagstr = ""
for l in lags:
lagstr += str(l)
# Apply Land/Ice Mask
dampingm = dampingt * limask[None, :, :]
# Flip longiude coordinates ([mon lat lon] --> [lon x lat x mon])
lon1, dampingw = proc.lon360to180(lon, dampingm.transpose(2, 1, 0))
# Multiple by 1 to make positive upwards
dampingw *= -1
# Save Result
if savevar == 1:
# Save variables, matching matlab script output
savedict = {'damping': dampingw, 'LON1': lon1, 'LAT': lat}
outname = "%s%s_PIC_%sdamping_monwin%i_sig%03d_dof%03d_mode%i.mat" % (
datpath, mconfig, flux.lower(), monwin, p * 100, dof, mode)
savemat(outname, savedict)
print("Saved data to %s" % outname)
if plotfigs:
#%% Vizualize the results ----
sstvars = []
for i in range(4):
sstin = sst[i]
sstvar = sstin.var(2)
sstvars.append(sstvar)
# Flip the longitude
fullvar = fullpt.var(0)
slabvar = slabpt.var(0)
# Apply the mask
fullvar *= msk
slabvar *= msk
# Flip the longitude
lon180, fullvar1 = proc.lon360to180(lon, fullvar.T[:, :, None])
_, slabvar1 = proc.lon360to180(lon, slabvar.T[:, :, None])
fullvar1 = np.squeeze(fullvar1)
slabvar1 = np.squeeze(slabvar1)
# Select region
bbox = [lonr[0], lonr[-1], latr[0], latr[-1]]
# Limit to region
fullr, _, _ = proc.sel_region(fullvar1, lon180, lat, bbox)
slabr, _, _ = proc.sel_region(slabvar1, lon180, lat, bbox)
#%% Take the ratios
entr = sstvars[3]
nentr = sstvars[1]
bbox_ST = [-80,-10,20,40]
bbox_TR = [-75,-15,0,20]
bbox_NA = [-80,0 ,0,65]
regions = ("SPG","STG","TRO","NAT") # Region Names
bboxes = (bbox_SP,bbox_ST,bbox_TR,bbox_NA) # Bounding Boxes
TS = TS.transpose(2,1,0) # Lon Lat Time
# Apply landice mask
limask = np.load(outpath+"landicemask_enssum.npy")
TS *= limask.T[:,:,None]
plt.pcolormesh(lon,lat,TS[:,:,0].T),plt.show()
# Flip Longitude
lon180,sst = proc.lon360to180(lon,TS)
plt.pcolormesh(lon180,lat,sst[:,:,0].T),plt.show()
nregions = len(regions)
acs = np.zeros((len(lags),nregions)) # Lag x Region
for r in range(nregions):
sstr,_,_ = proc.sel_region(sst,lon180,lat,bboxes[r])
tsmodel = np.nanmean(sstr,(0,1))
tsmodel = proc.year2mon(tsmodel) # mon x year
# Deseason
tsmodel2 = tsmodel - np.mean(tsmodel,1)[:,None]
#%% Load in mean SST from slab
fn1 = datpath + "../TS_SLAB_withENSO.npy"
fn2 = datpath + "../ENSOREM_TS_lag1_pcs2_monwin3.npz"
tsenso1 = np.load(fn1)
ld2 = np.load(fn2)
tsenso0 = ld2['TS']
lon360 = ld2['lon']
lat = ld2['lat']
# Reshape to [yr x mon x lat x lon]
nmon,nlat,nlon = tsenso0.shape
tsenso0 = tsenso0.transpose(2,1,0) # lon x lat x mon
lon1,tsenso180 = proc.lon360to180(lon360,tsenso0) # Flip longitudes
#tsenso0 = tsenso0.reshape(int(nmon/12),12,nlat,nlon)
#%% Integrate the model
flxin = flx
qdpin= qdpa
if quadrature:
flxin = flxnew
qdpin = np.roll(qdpa,1,)
mldmean = hblt.mean(2)#mld_1kmean
vstats = quickloader(vname,mconfig,datpath)
# For each variable, flip the longitude
vstatreg = []
vstatglob = []
vmaxmon = []
maxmonid = []
maxmonidglob = []
for v in vstats:
# Find month of maximum
kmon = np.argmax(v,axis=0).flatten()
maxmonidglob.append(kmon)
# Preprocess (flip longitude, cut to region)
v = v.transpose(2,1,0) * (msk.T)[:,:,None]
lon1,v1 = proc.lon360to180(lon,v)
vrr,lonr,latr = proc.sel_region(v1,lon1,lat,bbox=bbox)
# Do the same for the mask
_,msk180 = proc.lon360to180(lon,msk.T)
mskr,_,_ = proc.sel_region(msk180,lon1,lat,bbox=bbox)
# Find the month of maximum
vrrmax = np.argmax(vrr,axis=2).T+1 * mskr.T
kmonr = np.argmax(vrr,axis=2).flatten()
vstatglob.append(v1)
vstatreg.append(vrr)
vmaxmon.append(vrrmax)
maxmonid.append(kmonr)