def check_ENSO_sign(eofs, pcs, lon, lat, verbose=True):
"""
checks sign of EOF for ENSO and flips sign by check the sum over
conventional ENSO boxes (see below for more information)
checks to see if the sum within the region and flips if sum < 0
inputs:
1) eofs [space (latxlon), PC] EOF spatial pattern from eof_simple
2) pcs [time, PC] PC timeseries calculated from eof_simple
3) lat [lat] Latitude values
4) lon [lon] Longitude values
outputs:
1) eofs [space, PC]
2) pcs [time, PC]
"""
# Set EOF boxes to check over (west,east,south,north)
eofbox1 = [190, 240, -5,
5] #EOF1 sign check Lat/Lon Box (Currently using nino3.4)
eofbox2 = [
190, 240, -5, 5
] #EOF2 (Frankignoul et al. 2011; extend of positive contours offshore S. Am > 0.1)
eofbox3 = [
200, 260, 5, 5
] #EOF3 (Frankignoul et al. 2011; narrow band of positive value around equator)
chkboxes = [eofbox1, eofbox2, eofbox3]
# Find dimensions and separate out space
nlon = lon.shape[0]
nlat = lat.shape[0]
npcs = eofs.shape[1]
eofs = eofs.reshape(nlat, nlon, npcs)
eofs = eofs.transpose(1, 0, 2) # [lon x lat x npcs]
for n in range(npcs):
chk = proc.sel_region(eofs, lon, lat, chkboxes[n], reg_sum=1)[n]
if chk < 0:
if verbose:
print("Flipping EOF %i because sum is %.2f" % (n, chk))
eofs[:, :, n] *= -1
pcs[:, n] *= -1
eofs = eofs.transpose(1, 0,
2).reshape(nlat * nlon,
npcs) # Switch back to lat x lon x pcs
return eofs, pcs
# 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] fullratio = entr / fullr slabratio = nentr / slabr proc.sel_region(fullvar1, lon180, lat, bbox) bboxplot = [-80, 0, 0, 65] cmap = cmocean.cm.balance
cesmauto = np.load(projpath + "01_Data/Autocorrelation_Region.npy",
allow_pickle=True).item()
print("Data loaded in %.2fs" % (time.time() - start))
#%% Get Regional Data
nregion = len(regions)
sstregion = {}
for r in range(nregion):
bbox = bboxes[r]
sstr = {}
for model in range(4):
tsmodel = sst[model]
sstr[model], _, _ = proc.sel_region(tsmodel, lonr, latr, bbox)
sstregion[r] = sstr
#%% Calculate autocorrelation and SST averaged time series for a given region
kmonths = {}
autocorr_region = {}
sstavg_region = {}
for r in range(nregion):
bbox = bboxes[r]
autocorr = {}
sstavg = {}
for model in range(4):
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()
plt.plot(acreg[kmonth,:,klonr,klatr])
# Append to new plot
cesmacproc.append(acreg)
#%% Calculate RMSE with selected simulation
dampmat = 'ensavg_nhflxdamping_monwin3_sig020_dof082_mode4.mat'
loaddamp = loadmat(input_path+dampmat)
LON = np.squeeze(loaddamp['LON1'])
LAT = np.squeeze(loaddamp['LAT'])
damping = loaddamp['ensavg']
# Load Mixed layer variables (preprocessed in prep_mld.py)
mld = np.load(input_path+"HMXL_hclim.npy") # Climatological MLD
kprevall = np.load(input_path+"HMXL_kprev.npy") # Entraining Month
# ------------------
# % Restrict to region --------------------------------------------------
# ------------------
# Note: what is the second dimension for?
dampingr,lonr,latr = proc.sel_region(damping,LON,LAT,bboxsim)
hclim,_,_ = proc.sel_region(mld,LON,LAT,bboxsim)
kprev,_,_ = proc.sel_region(kprevall,LON,LAT,bboxsim)
if mldavg == 1:
hclim = np.ones(hclim.shape) * np.nanmean(hclim,(0,1,2))
if lbdavg == 1:
dampingr = np.ones(hclim.shape) * np.nanmean(dampingr,(0,1,2))
# Set artificial MLD
# hclim = np.ones(hclim.shape[0:-1])
# # Set artificial MLD (remove later)
# hclima=np.array([1,1,1,
# 1,1,1,
# 1,100,100,
# 100,1,1])
lons = dssum.lon.values
lats = dssum.lat.values
sstac = dssum.SST.values # [ens x mon x lats x lons]
nens, nmon, nlatr, nlonr = sstac.shape
#%% Prepare to Cut damping variables to the region
bboxreg = [lons[0], lons[-1], lats[0], lats[-1]]
# Reshape damping (also flip longitude, and cut region)
nens, nlag, nmon, nlat, nlon = damping.shape
damping_rs = damping.reshape(np.array(damping.shape[:-2]).prod(), nlat,
nlon) # Combine dims
damping_rs = damping_rs.transpose(2, 1, 0) # Flip to [lon x lat x otherdim]
lon1, damping_rs180 = proc.lon360to180(lon, damping_rs) # Flip longitude
dampingr, lonr, latr, = proc.sel_region(damping_rs180, lon1, lat,
bboxreg) # Cut Region
dampingr = dampingr.transpose(2, 1, 0) # Flip to [otherdim x lat x lon]
dampingr = dampingr.reshape(damping.shape[:-2] +
(len(latr), len(lonr))) # Uncombine dims
# Now compute the pattern correlation
dampingr_in = dampingr[:,
0, :, :, :] # Select first lag # [ens x mon x lat x lon]
patcorr_t2 = np.zeros((nens, nmon)) * np.nan
for e in tqdm(range(nens)):
for m in range(12):
patcorr_t2[e, m] = proc.patterncorr(dampingr_in[e, m, :, :],
sstac[e, m, :, :])
#%% Plot Intermember Variability (Stdev)
# 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
bbox = bboxes[r]
datr, _, _ = proc.sel_region(fstd, clon1, clat, bbox)
# Make Data 1D and remove NaN points
datr = datr.flatten()
datr, _, _ = proc.find_nan(datr, 0)
# Append data
stdboxes.append(datr)
stdval.append(np.mean(datr))
# Create Plot
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
plt.style.use("seaborn")
bp = ax.boxplot(stdboxes, 0, '',
labels=regions) # Note Outlier POints are not shown
ax.set_xlabel("Region")
# 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)
#%% Save Variable as input into stochastic Model
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]
# Compute autocorrelation and save data for region
acs[:,r] = proc.calc_lagcovar(tsmodel2,tsmodel2,lags,kmonth+1,0)
np.save("%sCESM-SLAB_PIC_autocorrelation_Regions.npy"%(outpath),acs)
def remapcluster(inclust,
lat5,
lon5,
regiondict,
printmsg=True,
returnremap=False):
# Remap an input cluster [inclust] according
# to a regiondict.
# Searches within each region and assigns
# value to most frequent class in a given region
nlat, nlon = inclust.shape
clusternew = inclust.copy()
clusternewflat = clusternew.flatten()
clusteroldflat = inclust.flatten()
assigned = []
remapdict = {}
for r in regiondict.keys():
#print(r)
# Get Region
bbox = regiondict[r].copy()
for i in range(2): # Just check Longitudes
if bbox[i] < 0:
bbox[i] += 360
varr, lonr, latr, = proc.sel_region(inclust.T,
lon5,
lat5,
bbox,
warn=printmsg)
# Get rid of NaNs
varrok = varr.flatten().copy()
varrok = varrok[~np.isnan(varrok)]
# Get unique elements and counts, sort by count
eles, freqs = np.unique(varrok, return_counts=True)
sortid = np.argsort(freqs)[::-1]
eles = eles[sortid]
done = False
for ele in eles:
if done: # Skip if already assigned
continue
if ele in assigned: # Skip if class has already be reassigned
continue
# Assign new cluster
clusternewflat[clusteroldflat == ele] = r
if printmsg:
print("Reassigned Class %i to %i" % (ele, r))
assigned.append(int(ele))
remapdict[int(ele)] = r
done = True
if done is False: # When no cluster is assigned...
# Get unassigned regions, and assign first one
unassigned = np.setdiff1d(list(regiondict.keys()), assigned)
ele = unassigned[0]
clusternewflat[clusteroldflat == ele] = r
assigned.append(int(ele))
remapdict[int(ele)] = r
if printmsg:
print(
"Reassigned (Leftover) Class %i to %i because nothing was found"
% (ele, r))
clusternew = clusternewflat.reshape(nlat, nlon)
if returnremap:
return clusternew, remapdict
return clusternew
for i in tqdm(range(len(clusts))):
k = np.arange(i, i + winsize)
grabsla = sla_lp[k, :, :]
sla_std[i, :, :] = grabsla.std(0)
#%% Try Shfting, then calculating pattern correlation
test = clusts[0]
bbtest = [55, 95, -25, -5]
bbtest2 = [55, 95, -30, -10]
test1 = test.copy()
klon, klat = proc.sel_region(test, lon5, lat5, bbtest, returnidx=True)
selpoints = test[klat[:, None], klon[None, :]]
klon2, klat2 = proc.sel_region(test, lon5, lat5, bbtest2, returnidx=True)
test1[klat2[:, None], klon2[None, :]] = selpoints
# Calculate Pattern Correlation
R = slutil.patterncorr(test, test1)
print(R)
# Plot some results (Clusters Themselves)
proj = ccrs.PlateCarree(central_longitude=180)
fig, axs = plt.subplots(2, 1, subplot_kw={'projection': proj}, figsize=(4, 5))
tests = [test, test1]
names = ['Before', 'After']
edgecolor="k",
color=fcolors[4],
alpha=falpha)
axs[2].set_title("EAP+NAO DJFM max=%03d" % h3[0].max())
axs[2].set_ylim([0, 5000])
axs[2].set_yticks(np.arange(0, 7500, 2500))
axs[2].set_xlim((-5, 5))
axs[2].set_xlabel("Forcing (W/$m^{2})$")
axs[1].set_ylabel("Frequency")
plt.tight_layout()
plt.savefig(outpath + "Forcing_Histograms.png", dpi=200)
#%% Plot Seasonal Damping and MLD parameters, averaged over the NAtl
mons3 = np.arange(1, 13, 1)
damppt = proc.sel_region(dampingr, lonr, latr, bbox_NA, reg_avg=1)
hpt = proc.sel_region(hclim, lonr, latr, bbox_NA, reg_avg=1)
fig, ax1 = plt.subplots(1, 1, figsize=(4, 3))
plt.style.use("seaborn-bright")
color = 'tab:red'
ax1.set_xlabel('Month')
ax1.set_ylabel('$\lambda_{net} (W m^{-2} K^{-1}$)')
ln1 = ax1.plot(mons3, damppt, color='r', label=r'$\lambda_{net}$')
ax1.tick_params(axis='y', labelcolor=color)
ax1.grid(None)
ax2 = ax1.twinx()
color = 'tab:blue'
ax2.set_ylabel('Mixed Layer Depth (m)', color=color)
def check_NAO_sign(eofs, pcs, lon, lat, tol=5, verbose=True):
"""
checks sign of EOF for NAO/EAP and flips sign by check the sum over
several boxes (see below for more information)
checks to see if the sum within the region and flips if sum < 0
inputs:
1) eofs [space (latxlon), PC] EOF spatial pattern from eof_simple
2) pcs [time, PC] PC timeseries calculated from eof_simple
3) lat [lat] Latitude values
4) lon [lon] Longitude values
5) tol [INT] Tolerance (in degrees) for searching
outputs:
1) eofs [space, PC]
2) pcs [time, PC]
"""
# Set EOF boxes to check over (west,east,south,north)
lonLisb = -9.1398 + 360
latLisb = 38.7223
lonReyk = -21.9426 + 360
latReyk = 64.146621
naobox1 = [lonReyk - tol, lonReyk + tol, latReyk - tol, latReyk + tol
] #EOF1 sign check Lat/Lon Box (Currently using nino3.4)
naobox2 = [lonLisb - tol, lonLisb + tol, latLisb - tol, latLisb + tol]
eapbox = [320, 15, 35, 60] #EAP Box
# Find dimensions and separate out space
nlon = lon.shape[0]
nlat = lat.shape[0]
npcs = eofs.shape[1]
eofs = eofs.reshape(nlat, nlon, npcs)
eofs = eofs.transpose(1, 0, 2) # [lon x lat x npcs]
for n in range(npcs):
if n == 0: # EOF 1
chk = proc.sel_region(eofs, lon, lat, naobox1, reg_sum=1)[n]
if chk > 0:
if verbose:
print(
"Flipping EOF %i (Reykjavik Check) because sum is %.2f"
% (n, chk))
eofs[:, :, n] *= -1
pcs[:, n] *= -1
chk = proc.sel_region(eofs, lon, lat, naobox2, reg_sum=1)[n]
if chk < 0:
if verbose:
print(
"Flipping EOF %i (Lisbon Check) because sum is %.2f" %
(n, chk))
eofs[:, :, n] *= -1
pcs[:, n] *= -1
else:
chk = proc.sel_region(eofs, lon, lat, eapbox, reg_sum=1)[n]
if chk > 0:
if verbose:
print("Flipping EOF %i (EAP Check) because sum is %.2f" %
(n, chk))
eofs[:, :, n] *= -1
pcs[:, n] *= -1
eofs = eofs.transpose(1, 0,
2).reshape(nlat * nlon,
npcs) # Switch back to lat x lon x pcs
return eofs, pcs
hpt = hclim[klon, klat, :]
kprev = kprevall[klon, klat, :]
if funiform > 2:
naopt = naoforce[klon, klat, :]
elif funiform == 2:
naopt = naoforce[klon, klat]
# Set string labels
loctitle = "LON: %02d LAT: %02d" % (lonf, latf)
locfname = "LON%03d_LAT%03d" % (lonf, latf)
elif regionmode == 1:
# Limit parameters to region and return average
hpt = proc.sel_region(hclim, lon, lat, rbbox, reg_avg=1)
kprev = proc.sel_region(kprevall, lon, lat, rbbox, reg_avg=1)
damppt = proc.sel_region(damping, lon, lat, rbbox, reg_avg=1)
if funiform >= 2:
naopt = proc.sel_region(naoforce, lon, lat, rbbox, reg_avg=1)
# Set string labels
loctitle = "%s (Avg.)" % (regionname)
locfname = regionname + "AVG"
# Determine month for autocorrelation calculation
kmon = hpt.argmax()
# Introduce Artificial seasonal cycle in damping
if seasonal_damping == 1:
# Get dimensions
nlon = len(lon)
nlat = len(lat)
ntime = len(times)
nens = len(mnum)
nyr = int(ntime / 12)
# Reshape to [lon,lat,time x ens]
invar = np.reshape(pslglo.transpose(3, 2, 0, 1), (nlon, nlat, ntime * nens))
# Correct if expressed in degrees west
if bboxNAO[0] < 0: bboxNAO[0] += 360
if bboxNAO[1] < 0: bboxNAO[1] += 360
# Select NAO region and reshape to original dimensions
pslnao, rlon, rlat = proc.sel_region(invar, lon, lat, bboxNAO, warn=0)
nlatr = len(rlat)
nlonr = len(rlon)
pslnao = pslnao.reshape(nlonr, nlatr, 1032,
42).transpose(3, 2, 1,
0) # reshape to [ens x time x lat x lon]
# Apply Area Weights
wgt = np.sqrt(np.cos(np.radians(rlat)))
pslnao = pslnao * wgt[None, None, :, None]
# Separate out the month and year dimensions and combine lat/lon, then transpose to [ens x space x yr x mon]
pslnao = pslnao.reshape(nens, nyr, 12,
nlatr * nlonr).transpose(0, 3, 1,
2) #[ens x space x yr x mon]
pslglo = pslglo.reshape(nens, nyr, 12, nlat * nlon).transpose(0, 3, 1, 2)
len(s_in),idx))
clusterout,knan,okpts = proc.find_nan((clusterin).flatten(),0)
# Make Silhouette Map
silmap = np.zeros(nlat5*nlon5)*np.nan
silmap[okpts] = s_in
silmap = silmap.reshape(nlat5,nlon5)
silmap_all[idx,:,:] = silmap
# Reassign clusters
#%% Find some extreme values
# Calculate averaged silhouette value over a particular region
bboxsil = regiondict[1]
sreg,slon,slat = proc.sel_region(silmap_all.transpose(2,1,0),lon5,lat5,bboxsil)
s_allin = np.nanmean(sreg,(0,1))
# Flatte the region, and count the number of unique values within that region
creg,_,_ = proc.sel_region(clusts.transpose(2,1,0),lon5,lat5,bboxsil)
nlatr,nlonr,nrngs = creg.shape
cregflat = creg.reshape(nlatr*nlonr,nrngs)
okdata,knan,okpts = proc.find_nan(cregflat,1) # Pick only non-nan points
npts = okdata.shape[0]
rcount = np.zeros(nrngs)
for t in range(nrngs):
indata = okdata[:,t]
rcount[t] = len(np.unique(indata))
def findk(arr,k,max=True):
mld = np.load(input_path + "HMXL_hclim.npy") # Climatological MLD
kprevall = np.load(input_path + "HMXL_kprev.npy") # Entraining Month
dampmat = 'ensavg_nhflxdamping_monwin3_sig020_dof082_mode4.mat'
loaddamp = loadmat(input_path + dampmat)
LON = np.squeeze(loaddamp['LON1'])
LAT = np.squeeze(loaddamp['LAT'])
if mconfig == "SLAB_PIC":
damping = np.load(input_path + mconfig +
"_NHFLX_Damping_monwin3_sig005_dof894_mode4.npy")
elif mconfig == "FULL_HTR":
damping = np.load(input_path + mconfig +
"_NHFLX_Damping_monwin3_sig020_dof082_mode4.npy")
# Restrict to Region
dampingr, lonr, latr = proc.sel_region(damping, LON, LAT, bboxsim)
hclim, _, _ = proc.sel_region(mld, LON, LAT, bboxsim)
kprev, _, _ = proc.sel_region(kprevall, LON, LAT, bboxsim)
# Get lat and long sizes
lonsize = lonr.shape[0]
latsize = latr.shape[0]
# Calculate values
o, a = proc.find_latlon(-30, 50, lonr, latr)
if usetau:
lbd, lbd_entr, FAC, beta = scm.set_stochparams(hclim[o, a, :],
tauall.mean(1),
dt,
ND=0,
hfix=hfix)
def stochmod_region(pointmode,funiform,fscale,runid,genrand,nyr,fstd,bboxsim,stormtrack,
points=[-30,50],mconfig='FULL_HTR',applyfac=1,parallel=False):
# --------------
# %% Set Parameters--------------------------------------------------------
# --------------
# Unpack Points if in pointmode
lonf,latf = points
# Other intengration Options (not set by user)
t_end = 12*nyr # Calculates Integration Period
dt = 60*60*24*30 # Timestep size (Will be used to multiply lambda)
T0 = 0 # Initial temperature [degC]
hfix = 50 # Fixed MLD value (meters)
# Set Constants
cp0 = 3996 # Specific Heat [J/(kg*C)]
rho = 1026 # Density of Seawater [kg/m3]
# Set Integration Region
lonW,lonE,latS,latN = bboxsim
# Save Option
saveforcing = 0 # Save Forcing for each point (after scaling, etc)
#Set Paths (stormtrack and local)
if stormtrack == 0:
projpath = "/Users/gliu/Downloads/02_Research/01_Projects/01_AMV/02_stochmod/"
datpath = projpath + '01_Data/'
sys.path.append("/Users/gliu/Downloads/02_Research/01_Projects/01_AMV/02_stochmod/03_Scripts/stochmod/model/")
sys.path.append("/Users/gliu/Downloads/02_Research/01_Projects/01_AMV/00_Commons/03_Scripts/")
elif stormtrack == 1:
datpath = "/stormtrack/data3/glliu/01_Data/02_AMV_Project/02_stochmod/Model_Data/"
sys.path.append("/home/glliu/00_Scripts/01_Projects/00_Commons/")
sys.path.append("/home/glliu/00_Scripts/01_Projects/01_AMV/02_stochmod/stochmod/model/")
import scm
from amv import proc
input_path = datpath + 'model_input/'
output_path = datpath + 'model_output/'
## ------------ Script Start -------------------------------------------------
print("Now Running stochmod_region with the following settings: \n")
print("mconfig = " + mconfig)
print("funiform = " + str(funiform))
print("genrand = " + str(genrand))
print("fstd = " + str(fstd))
print("runid = " + runid)
print("pointmode = " + str(pointmode))
print("fscale = " + str(fscale))
print("nyr = " + str(nyr))
print("bbox = " + str(bboxsim))
print("Data will be saved to %s" % datpath)
allstart = time.time()
# Set experiment ID
#expid = "%iyr_funiform%i_run%s_fscale%03d" %(nyr,funiform,runid,fscale)
expid = "%s_%iyr_funiform%i_run%s_fscale%03d_applyfac%i" %(mconfig,nyr,funiform,runid,fscale,applyfac)
# --------------
# %% Load Variables ------------------------------------------------------
# --------------
# Load Latitude and Longitude
dampmat = 'ensavg_nhflxdamping_monwin3_sig020_dof082_mode4.mat'
loaddamp = loadmat(input_path+dampmat)
LON = np.squeeze(loaddamp['LON1'])
LAT = np.squeeze(loaddamp['LAT'])
# Load Atmospheric Heat Flux Feedback/Damping
if mconfig == "FULL_HTR":
damping = np.load(input_path+mconfig+"_NHFLX_Damping_monwin3_sig020_dof082_mode4.npy")
elif mconfig == "SLAB_PIC":
damping = np.load(input_path+mconfig+"_NHFLX_Damping_monwin3_sig005_dof894_mode4.npy")
# Load Mixed layer variables (preprocessed in prep_mld.py)
mld = np.load(input_path+"FULL_PIC_HMXL_hclim.npy") # Climatological MLD
kprevall = np.load(input_path+"FULL_PIC_HMXL_kprev.npy") # Entraining Month
# Load MLD from SLAB run
hblt = np.load(input_path+"SLAB_PIC_hblt.npy")
# ------------------
# %% Restrict to region --------------------------------------------------
# ------------------
# Note: what is the second dimension for?
dampingr,lonr,latr = proc.sel_region(damping,LON,LAT,bboxsim)
hclim,_,_ = proc.sel_region(mld,LON,LAT,bboxsim)
kprev,_,_ = proc.sel_region(kprevall,LON,LAT,bboxsim)
hbltr,_,_ = proc.sel_region(hblt,LON,LAT,bboxsim)
hbltr = hbltr.mean(2) # Take mean along month dimensions
# Get lat and long sizes
lonsize = lonr.shape[0]
latsize = latr.shape[0]
np.save(datpath+"lat.npy",latr)
np.save(datpath+"lon.npy",lonr)
# ------------------
# %% Prep NAO Forcing ----------------------------------------------------
# ------------------
# Consider moving this section to another script?
if funiform > 1: # For NAO-like forcings (and EAP forcings, load in data and setup)
if funiform == 1.5: # Scale by standard deviation of NHFLX
# [lon x lat x mon]
NAO1 = np.load(input_path+mconfig+"_NHFLXSTD_Forcing_Mon.npy")
# # Load Longitude for processing
# lon360 = np.load(datpath+"CESM_lon360.npy")
if funiform == 2: # Load (NAO-NHFLX)_DJFM Forcing
# [lon x lat x pc]
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM.npy") #[PC x Ens x Lat x Lon]
# Select PC1 # [lon x lat x 1]
NAO1 = naoforcing[:,:,[0]]
elif funiform == 3: # NAO (DJFM) regressed to monthly NHFLX
# [lon x lat x pc x mon]
#naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON.npy") #[PC x Ens x Lat x Lon]
# Calculated from calc-NAO_PIC_monhtly.py.... Fixed NAO Pattern
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON_Fix.npy") # New file, magnitude does not vary at each point
# Select PC 1 and 2 # [lon x lat x mon]
NAO1 = naoforcing[:,:,0,:]
elif funiform == 4: # Monthly NAO and NHFLX
# NOTE: THESE HAVE NOT BEEN RECALCULATED. NEED TO REDO FOR PIC SLAB ----
# # Load Forcing and take ensemble average
# naoforcing = np.load(datpath+"NAO_Monthly_Regression_PC.npz")['eofall'] #[Ens x Mon x Lat x Lon]
# NAO1 = np.nanmean(naoforcing,0)
# Load Forcing and take ensemble average
naoforcing = np.load(datpath+"NAO_Monthly_Regression_PC123.npz")['flxpattern'] #[Ens x Mon x Lat x Lon]
# Select PC1 Take ensemble average
NAO1 = naoforcing[:,:,:,:,0].mean(0)
elif funiform == 5: # EAP (DJFM) ONLY
# [lon x lat x pc]
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM.npy") #[PC x Ens x Lat x Lon]
# Select PC 2 # [lon x lat x 1]
NAO1 = naoforcing[:,:,[1]]
elif funiform == 5.5: # EAP (DJFM-MON)
# [lon x lat x pc x mon]
#naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON.npy") #[PC x Ens x Lat x Lon]
# Calculated from calc-NAO_PIC_monhtly.py.... Fixed NAO Pattern
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON_Fix.npy") # New file, magnitude does not vary at each point
# Select PC 2 # [lon x lat x 2 x mon]
NAO1 = naoforcing[:,:,1,:]
elif funiform == 6: # DJFM NAO and EAP
# [lon x lat x pc]
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM.npy") #[PC x Ens x Lat x Lon]
# Select PC 1 and 2 # [lon x lat x 2]
NAO1 = naoforcing[:,:,[0,1]]
elif funiform == 7: # DJFM Index, Monthly NHFLX
# [lon x lat x pc x mon]
#naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON.npy") #[PC x Ens x Lat x Lon]
# Calculated from calc-NAO_PIC_monhtly.py.... Fixed NAO Pattern
naoforcing = np.load(input_path+mconfig+"_NAO_EAP_NHFLX_Forcing_DJFM-MON_Fix.npy") # New file, magnitude does not vary at each point
# Select PC 1 and 2 # [lon x lat x 2 x mon]
NAO1 = naoforcing[:,:,[0,1],:]
# Restrict to region
NAO1,_,_ = proc.sel_region(NAO1,LON,LAT,bboxsim,autoreshape=True)
else: # For funiform= uniform or random forcing, just make array of ones
NAO1 = np.ones(hclim.shape)
# Convert NAO from W/m2 to degC/sec. Returns dict with keys 0-2
NAOF = {}
NAOF1 = {}
if applyfac == 0: # Don't Apply MLD Cycle
if funiform > 1:
NAO1 = NAO1 * dt / rho / cp0 # Do conversions (minus MLD)
for i in range(3):
if funiform > 5.5: # Separate NAO and EAP Forcing
NAOF[i] = NAO1[:,:,0,...].copy() # NAO Forcing
NAOF1[i] = NAO1[:,:,1,...].copy() # EAP Forcing
else:
NAOF[i] = NAO1.copy()
else: # Apply seasonal MLD cycle and convert
if funiform >= 6: # Separately convert NAO and EAP forcing
NAOF = scm.convert_NAO(hclim,NAO1[:,:,0],dt,rho=rho,cp0=cp0,hfix=hfix,hmean=hbltr[:,:,None]) # NAO Forcing
NAOF1 = scm.convert_NAO(hclim,NAO1[:,:,1],dt,rho=rho,cp0=cp0,hfix=hfix,hmean=hbltr[:,:,None]) # EAP Forcing
else:
NAOF = scm.convert_NAO(hclim,NAO1,dt,rho=rho,cp0=cp0,hfix=hfix,hmean=hbltr[:,:,None])
# Out: Dict. (keys 0-2) with [lon x lat x mon]
"""
# Outformat: Dict. (keys 0-2, representing MLD type) with [lon x lat x mon]
# We have prepared NAO forcing patterns for the 3 different MLD treatments (if
# applyfac is set. All it requires now is scaling by both the chosen factor and
# white nosie timeseries)
"""
# ----------------------------
# %% Set-up damping parameters
# ----------------------------
lbd,lbd_entr,FAC,beta = scm.set_stochparams(hclim,dampingr,dt,ND=1,rho=rho,cp0=cp0,hfix=hfix,hmean=hbltr[:,:,None])
"""
Out Format:
lbd -> Dict (keys 0-3) representing each mode, damping parameter
lbd_entr -> array of entrainment damping
FAC -> Dict (keys 0-3) representing each model, integration factor
beta ->array [Lon x Lat x Mon]
"""
# ----------------------------
# %% Set Up Forcing ------------------------------------------------
# ----------------------------
startf = time.time()
# Prepare or load random time series
if genrand == 1: # Generate new time series
print("Generating New Time Series")
# Create and save entire forcing array [lon x lat x time] and apply scaling factor
if funiform == 0:
F = np.random.normal(0,fstd,size=(lonsize,latsize,t_end)) * fscale # Removed Divide by 4 to scale between -1 and 1
np.save(output_path+"stoch_output_%s_Forcing.npy"%(expid),F)
randts = np.random.normal(0,fstd,size=t_end) # Just generate dummy randts
# Just generate the time series
else:
randts = np.random.normal(0,fstd,size=t_end) # Removed Divide by 4 to scale between -1 and 1
np.save(output_path+"stoch_output_%iyr_run%s_randts.npy"%(nyr,runid),randts)
else: # Load old data
print("Loading Old Data")
if funiform == 0:# Directly load full forcing
F = np.load(output_path+"stoch_output_%s_Forcing.npy"%(expid))
randts = np.random.normal(0,fstd,size=t_end) # Just generate dummy randts
else: # Load random time series
randts = np.load(output_path+"stoch_output_%iyr_run%s_randts.npy"%(nyr,runid))
# Generate extra time series for EAP forcing
if funiform in [5,6,7]:
numforce = 1 # In the future, incoporate forcing for other EOFs
# Generate newtimeseries if it is missing
if (genrand == 1) | (len(glob.glob(output_path+"stoch_output_%iyr_run%s_randts_%03d.npy"%(nyr,runid,numforce)))==0):
print("Generating Additional New Time Series for EAP")
randts1 = np.random.normal(0,fstd,size=t_end) # Removed Divide by 4 to scale between -1 and 1
np.save(output_path+"stoch_output_%iyr_run%s_randts_%03d.npy"%(nyr,runid,numforce),randts)
else:
print("Loading Additional New Time Series for EAP")
randts1 = np.load(output_path+"stoch_output_%iyr_run%s_randts_%03d.npy"%(nyr,runid,numforce))
if funiform in [5,5.5]: # Assign EAP Forcing white noise time series
randts = randts1
# Use random time series to scale the forcing pattern
if funiform != 0:
if funiform in [6,7]: # NAO + EAP Forcing
F,Fseas = scm.make_naoforcing(NAOF,randts,fscale,nyr) # Scale NAO Focing
F1,Fseas1 = scm.make_naoforcing(NAOF1,randts1,fscale,nyr) # Scale EAP forcing
# Add the two forcings together
for hi in range(3):
F[hi] += F1[hi]
Fseas[hi] += Fseas1[hi]
else: # NAO Like Forcing of funiform with mld/lbd factors, apply scaling and randts
F,Fseas = scm.make_naoforcing(NAOF,randts,fscale,nyr)
# Save Forcing if option is set
if saveforcing == 1:
np.save(output_path+"stoch_output_%s_Forcing.npy"%(runid),F)
else: # Duplicate for uniform forcing
F0 = F.copy()
F={}
for hi in range(3):
F[hi] = F0
print("Forcing Setup in %.2fs" % (time.time() - startf))
"""
Output:
F - dict (keys = 0-2, representing each MLD treatment) [ lon x lat x time (simulation length)]
Fseas - dict (keys = 0-2, representing each MLD treatment) [ lon x lat x month]
"""
# ----------------------------
# %% Additional setup based on pointmode ------------------------------------------------
# ----------------------------
if pointmode == 1: # Find indices for pointmode
# Get point indices
ko,ka = proc.find_latlon(lonf,latf,lonr,latr)
locstring = "lon%02d_lat%02d" % (lonf,latf)
# Select variable at point
hclima = hclim[ko,ka,:]
dampinga = dampingr[ko,ka,:]
kpreva = kprev[ko,ka,:]
lbd_entr = lbd_entr[ko,ka,:]
beta = beta[ko,ka,:]
hblta = hbltr[ko,ka]
#naoa = NAO1[ko,ka,...]
# Select forcing at point
Fa = {} # Forcing
Fseasa = {} # Seasonal Forcing pattern
for hi in range(3):
Fa[hi] = F[hi][ko,ka,:]
Fseasa = Fseas[hi][ko,ka,:]
F = Fa.copy()
Fseas=Fseasa.copy()
# Do the same but for each model type (hfdamping and FAC)
lbda = {}
FACa = {}
for model in range(4):
FACa[model] = FAC[model][ko,ka,:]
lbda[model] = lbd[model][ko,ka,:]
lbd = lbda.copy()
FAC = FACa.copy()
if pointmode == 2: # Take regionally averaged parameters (need to recalculate some things)
# Make string for plotting
locstring = "lon%02d_%02d_lat%02d_%02d" % (lonW,lonE,latS,latN)
# Current setup: Average raw variables, assuming
# that bboxsim is the region you want to average over
hclima = np.nanmean(hclim,(0,1)) # Take lon,lat mean, ignoring nans
kpreva = scm.find_kprev(hclima)[0] # Recalculate entrainment month
dampinga = np.nanmean(dampingr,(0,1)) # Repeat for damping
hblta = np.nanmean(hbltr,(0,1)) # Repeat for slab mld
#naoa = np.nanmean(NAO1,(0,1)) # Repeat for nao forcing
# Get regionally averaged forcing based on mld config
rNAOF = {}
rF = {}
for hi in range(3):
rNAOF[hi] = proc.sel_region(NAOF[hi],lonr,latr,bboxsim,reg_avg=1)
rF[hi] = randts * np.tile(rNAOF[hi],nyr)
# Add in EAP Forcing [consider making separate file to save?]
if funiform in [6,7]: # NAO + EAP Forcing
for hi in range(3):
rNAOF1 = proc.sel_region(NAOF1[hi],lonr,latr,bboxsim,reg_avg=1)
rF1 = randts1 * np.tile(rNAOF1,nyr)
# Add to forcing
rNAOF[hi] += rNAOF1
rF[hi] += rF1
# Copy over forcing
F = rF.copy()
Fseas = rNAOF.copy()
# Convert units
lbd,lbd_entr,FAC,beta = scm.set_stochparams(hclima,dampinga,dt,ND=0,rho=rho,cp0=cp0,hfix=hfix,hmean=hblta)
"""
Output:
Dict with keys 0-2 for MLD configuation
- F (Forcing, full timeseries)
- Fseas (Forcing, seasonal pattern)
Dict with keys 0-3 for Model Type
- lbd (damping parameter)
- FAC (integration factor)
Just Arrays...
- beta (entrainment velocity)
- dampinga (atmospheric damping)
- hclima (mixed layer depth)
- kpreva (entraining month)
- naoa (NAO forcing pattern)
"""
# ----------
# %%RUN MODELS -----------------------------------------------------------------
# ----------
# Set mulFAC condition based on applyfac
if applyfac == 2:
multFAC = 1 # Don't apply integrationreduction factor if applyfac is set to 0 or 1
else:
multFAC = 0
# Run Model Without Entrainment
sst = {}
#Loop for each Mixed Layer Depth Treatment
for hi in range(3):
start = time.time()
# Select damping and FAC based on MLD
FACh = FAC[hi]
lbdh = lbd[hi]
# Select Forcing
Fh = F[hi]
# # Match Forcing and FAC shape
# if (len(Fh.shape)>2) & (Fh.shape[2] != FACh.shape[2]):
# FACh = np.tile(FACh,int(t_end/12))
if pointmode == 0: #simulate all points
# Match Forcing and FAC shape
if (len(Fh.shape)>2) & (Fh.shape[2] != FACh.shape[2]):
FACh = np.tile(FACh,int(t_end/12))
sst[hi] = scm.noentrain_2d(randts,lbdh,T0,Fh,FACh,multFAC=multFAC)
print("\nSimulation for No Entrain Model, hvarmode %s completed in %s" % (hi,time.time() - start))
else: # simulate for 1 point (or regionally averaged case)
start = time.time()
# Run Point Model
sst[hi],_,_=scm.noentrain(t_end,lbdh,T0,Fh,FACh,multFAC=multFAC)
elapsed = time.time() - start
tprint = "\nNo Entrain Model, hvarmode %i, ran in %.2fs" % (hi,elapsed)
print(tprint)
#%%
# Run Model With Entrainment
start = time.time()
icount = 0
Fh = F[2] # Forcing with varying MLD
FACh = FAC[3] # Integration Factor with entrainment
if pointmode == 0: # All Points
if parallel:
st = time.time()
lonsize,latsize,_ = F[2].shape
Fin = F[2].reshape(lonsize*latsize,t_end)
FACin = FAC[3].reshape(lonsize*latsize,12)
dampingpt = dampingr.reshape(FACin.shape)
lbdin = lbd[3].reshape(FACin.shape)
betain = beta.reshape(FACin.shape)
hin = hclim.reshape(FACin.shape)
kprevin = kprev.reshape(FACin.shape)
#T_entr1 = np.zeros(Fin.shape) * np.nan
results = []
T_entr1 = np.array([])
for i in trange(lonsize*latsize):
if np.isnan(np.mean(dampingpt[i,:])):
results.append(np.zeros(t_end)*np.nan)
continue
inputs = (t_end,lbdin[i],T0,Fin[i],betain[i],hin[i],kprevin[i],FACin[i],multFAC)
#T_entr1[i,:] = dask.delayed(scm.entrain_parallel)(inputs)
result = dask.delayed(scm.entrain_parallel)(inputs)
results.append(result)
#T_entr1.append(dask.delayed(scm.entrain_parallel)(inputs))
dask.compute(*results)
x = T_entr1.compute()
end = time.time()
print("Finished in %.2fs"%(end-st))
else: # Regular Loop without parallelization
T_entr1 = np.ones((lonsize,latsize,t_end))*np.nan
for o in trange(lonsize,desc="Longitude"):
for a in range(latsize):
# Skip if the point is land
if np.isnan(np.mean(dampingr[o,a,:])):
#msg = "Land Point @ lon %f lat %f" % (lonf,latf)
icount += 1
continue
#print(msg)
else:
# T_entr1[o,a,:] = scm.entrain(t_end,lbd[3][o,a,:],
# T0,Fh[o,a,:],beta[o,a,:],
# hclim[o,a,:],kprev[o,a,:],
# FACh[o,a,:],multFAC=multFAC,
# debug=False,debugprint=False)
T_entr1[o,a,:] = scm.entrain(t_end,lbd[3][o,a,:],T0,Fh[o,a,:],beta[o,a,:],hclim[o,a,:],kprev[o,a,:],FACh[o,a,:],multFAC=multFAC)
icount += 1
#msg = '\rCompleted Entrain Run for %i of %i points' % (icount,lonsize*latsize)
#print(msg,end="\r",flush=True)
#End Latitude Loop
#End Longitude Loop
else: # Single point/average region
T_entr1= scm.entrain(t_end,lbd[3],T0,Fh,beta,hclima,kpreva,FACh,multFAC=multFAC)
# Copy over to sst dictionary
sst[3] = T_entr1.copy()
elapsed = time.time() - start
tprint = "\nEntrain Model ran in %.2fs" % (elapsed)
print(tprint)
#%% save output
if pointmode > 0:
np.savez(output_path+"stoch_output_point%s_%s.npz"%(locstring,expid),
sst=sst,
hclim=hclima,
kprev=kpreva,
dampping=dampinga,
F=F,
lbd=lbd,
lbd_entr=lbd_entr,
beta=beta,
FAC=FAC,
NAO1=NAO1,
NAOF=NAOF
)
else:
# SAVE ALL in 1
np.save(output_path+"stoch_output_%s.npy"%(expid),sst)
print("stochmod_region.py ran in %.2fs"% (time.time()-allstart))
print("Output saved as" + output_path + "stoch_output_%s.npy"%(expid))