ld = np.load(datpath+"../FULL_PIC_ENSOREM_TS_lag1_pcs2_monwin3.npz" ,allow_pickle=True)
sstfull = ld['TS']
ld2 = np.load(datpath+"../SLAB_PIC_ENSOREM_TS_lag1_pcs2_monwin3.npz" ,allow_pickle=True)
sstslab = ld2['TS']
query = [lonf,latf]
locstring = "lon%i_lat%i" % (query[0],query[1])
locstringtitle = "Lon: %.1f Lat: %.1f" % (query[0],query[1])
mons3=('Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec')
labels=["MLD Fixed","MLD Mean","MLD Seasonal","MLD Entrain"]
expcolors = ('blue','orange','magenta','red')
#%% Test for a single point
# Get data for point
klon,klat = proc.find_latlon(lonf,latf,lon,lat) # Globa, 180 lon
klon360,_ = proc.find_latlon(lonf+360,latf,lon360,lat) # Globa, 360 lon
klonr,klatr = proc.find_latlon(lonf,latf,lonr,latr) # NAtl, 180lon
# Get MLD Information
mldpt=mld[klon,klat,:]
kmonth = np.argmax(mldpt)
print("Kmonth is %i"%kmonth)
# Get comparison data from CESM SLAB and Full Simulations
slabauto= cesmslabac[kmonth,lags,klat,klon360]
fullauto= cesmfullac[kmonth,lags,klon360,klat]
# Get point value for simulation output
ftype = config['ftype']
locstring = "lon%i_lat%i" % (query[0], query[1])
locstringtitle = "Lon: %.1f Lat: %.1f" % (query[0], query[1])
# Run Model
#config['Fpt'] = np.roll(Fpt,1)
ac, sst, dmp, frc, ent, Td, kmonth, params = scm.synth_stochmod(
config, projpath=projpath)
[o, a], damppt, mldpt, kprev, Fpt = params
# Read in CESM autocorrelation for all points'
kmonth = np.argmax(mldpt)
print("Kmonth is %i" % kmonth)
_, _, lon, lat, lon360, cesmslabac, damping, _, _ = scm.load_data(
mconfig, ftype)
ko, ka = proc.find_latlon(query[0] + 360, query[1], lon360, lat)
cesmauto2 = cesmslabac[kmonth, :, ka, ko]
cesmauto = cesmauto2[lags]
# Plot some differences
xtk2 = np.arange(0, 37, 2)
fig, ax = plt.subplots(1, 1)
title = "SST Autocorrelation at %s (Lag 0 = %s)" % (locstringtitle,
mons3[mldpt.argmax()])
ax, ax2, ax3 = viz.init_acplot(kmonth,
xtk2,
lags,
ax=ax,
loopvar=params[2],
title=title)
ax.plot(lags, cesmauto2[lags], label="CESM SLAB", color='k')
#
#%% Remove GMSL
#
if rem_gmsl:
print("Removing GMSL")
gmslrem = np.nanmean(sla_5deg, (1, 2))
sla_5deg_ori = sla_5deg.copy()
sla_5deg = sla_5deg - gmslrem[:, None, None]
if debug:
# Test plot point
lonf = 330
latf = 40
klon, klat = proc.find_latlon(lonf, latf, lon5, lat5)
fig, ax = plt.subplots(1, 1)
ax.set_xticks(np.arange(0, 240, 12))
ax.set_xticklabels(timesyr[::12], rotation=45)
ax.grid(True, ls='dotted')
ax.plot(sla_5deg_ori[:, klat, klon], label="Original", color='k')
ax.plot(sla_5deg[:, klat, klon], label="Post-Removal")
ax.plot(gmslrem, label="GMSL")
ax
ax.legend()
ax.set_title("GMSL Removal at Lon %.2f Lat %.2f % (%s to %s)" %
(lon5[klon], lat5[klat], start, end))
ax.set_ylabel()
plt.savefig(outfigpath + "GMSL_Removal_lon%i_lat%i.png" % (lonf, latf),
"""
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 indices for selected point and make string for plotting
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,:]
naoa = NAO1[ko,ka,...]
# Do the same for dictionaries indexed by MLD config
Fa = {} # Forcing
Fseasa = {} # Seasonal Forcing pattern
for hi in range(3):
Fa[hi] = F[hi][ko,ka,:]
mconfig = config['mconfig']
lags = config['lags']
ftype = config['ftype']
locstring = "lon%i_lat%i" % (query[0],query[1])
locstringtitle = "Lon: %.1f Lat: %.1f" % (query[0],query[1])
# Run Model
#config['Fpt'] = np.roll(Fpt,1)
ac,sst,dmp,frc,ent,Td,kmonth,params=scm.synth_stochmod(config,projpath=projpath)
[o,a],damppt,mldpt,kprev,Fpt =params
# Read in CESM autocorrelation for all points'
kmonth = np.argmax(mldpt)
print("Kmonth is %i"%kmonth)
_,_,lon,lat,lon360,cesmslabac,damping,_,_ = scm.load_data(mconfig,ftype)
ko,ka = proc.find_latlon(query[0]+360,query[1],lon360,lat)
cesmauto2 = cesmslabac[kmonth,:,ka,ko]
cesmauto = cesmauto2[lags]
cesmautofull = fullauto[kmonth,lags,ko,ka]
# Plot some differences
xtk2 = np.arange(0,37,2)
fig,ax = plt.subplots(1,1)
title = "SST Autocorrelation at %s (Lag 0 = %s)" % (locstringtitle,mons3[mldpt.argmax()])
ax,ax2,ax3 = viz.init_acplot(kmonth,xtk2,lags,ax=ax,loopvar=params[2],title=title)
ax.plot(lags,cesmauto2[lags],label="CESM SLAB",color='k')
ax.plot(lags,cesmautofull,color='k',label='CESM Full',ls='dashdot')
for i in range(1,4):
ax.plot(lags,ac[i],label=labels[i],color=expcolors[i])
flx = dsflx.NHFLX.values
lon = dsflx.lon.values
lat = dsflx.lat.values
print("Loaded NHFLX in %.2fs"% (time.time()-st))
# Mixed Layer Depths
st = time.time()
dsmld = xr.open_dataset(datpath+"HMXL_PIC.nc")
mld = dsmld.HMXL.values/100 # Convert to meters
lon = dsmld.lon.values
lat = dsmld.lat.values
print("Loaded MLD in %.2fs"% (time.time()-st))
#%% Preprocessing
klon,klat = proc.find_latlon(lonf,latf,lon,lat)
loctitle = "Lon %.2f Lat %.2f" % (lon[klon],lat[klat])
locfn = "lon%i_lat%i" % (lonf,latf)
# Calculate 1000 year mean and seasonal MLDs
nlon,nlat,ntimef = mld.shape
hclim = mld.reshape(nlon,nlat,int(ntimef/12),12)
mldcycle = hclim.mean(2)
mld_1kyr = mld[:,:,:1000]
mld_1kmean = mld_1kyr.mean(2)
print("1000 yr mean mld is : %.3f m" % (mld_1kmean[klon,klat]))
print("Mean of scycle is : %.3f m" % (mldcycle.mean(2)[klon,klat]))
#hclim = np.zeros((nlon,nlat,12))*np.nan
if debug:
# Also save the linear model
ymodall = np.zeros(hsstnew.shape) * np.nan
ymodall[:, okpts] = (beta[:, None] * tper + b[:, None]).T
# Reshape again
dt_hsst = np.reshape(hsstall.T, (360, 180, 1176))
hsstnew = np.reshape(hsstnew.T, (360, 180, 1176))
ymodall = np.reshape(ymodall.T, (360, 180, 1176))
print("Detrended in %.2fs" % (time.time() - start))
#%% Visualize detrending for point [lonf,latf] (Visualize Detrending effect)
# Find Point
lonf = -30
latf = 64
klon, klat = proc.find_latlon(lonf, latf, hlon, hlatnew)
# Get values at point
tempts = hsstnew[klon, klat, :]
dtts = dt_hsst[klon, klat, :]
ymodts = ymodall[klon, klat, :]
# Test other ways of detrending
olddt = detrendlin(tempts)
from scipy import signal
scidt = signal.detrend(tempts)
# Make time eriod
tper = np.arange(0, len(tempts), 1)
#% Plot Detrended and undetrended lines
sla_5deg_ori = sla_5deg.copy()
sla_5deg = sla_5deg - gmslrem[:,None,None]
if np.all(gmslrem>1e-10):
print("Saving GMSL")
np.save(outpath+"AVISO_GMSL_%s_%s.npy"%(start,end),gmslrem)
else:
print("GMSL is already removed")
plt.plot(gmslrem)
if debug:
# Test plot point
lonf = 330
latf = 40
klon,klat = proc.find_latlon(lonf,latf,lon5,lat5)
fig,ax = plt.subplots(1,1)
ax.set_xticks(np.arange(0,240,12))
ax.set_xticklabels(timesyr[::12],rotation = 45)
ax.grid(True,ls='dotted')
ax.plot(sla_5deg_ori[:,klat,klon],label="Original",color='k')
ax.plot(sla_5deg[:,klat,klon],label="Post-Removal")
ax.plot(gmslrem,label="GMSL")
ax
ax.legend()
ax.set_title("GMSL Removal at Lon %.2f Lat %.2f (%s to %s)" % (lon5[klon],lat5[klat],start,end))
ax.set_ylabel("SSH (m)")
plt.savefig(outfigpath+"GMSL_Removal_lon%i_lat%i.png"% (lonf,latf),dpi=200)
plt.legend()
#plt.title(titlestr)
plt.tight_layout()
#plt.grid(True)
plt.savefig(outpath + "Damping_entr_RhoCPH_MLD_NATAVG.png", dpi=200)
#%% Visualize a point
damping_hist = damping.copy()
damping_slab = np.load(
input_path + "SLAB_PIC_NHFLX_Damping_monwin3_sig005_dof894_mode4.npy")
damping_slab = damping_slab[klon[:, None], klat[None, :], :]
nlon, nlat, nmon = damping_slab.shape
klon, klat = proc.find_latlon(-30, 50, lonr, latr)
mons3 = ('Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct',
'Nov', 'Dec')
# Plot Heat Flux Feedback
fig, ax = plt.subplots(1, 1)
ax.plot(mons3,
damping_slab.reshape(nlon * nlat, 12).T,
color='k',
alpha=0.01,
label="")
ax.plot(mons3, damping_slab[klon, klat, :], label="Lon -30, Lat 50", color='r')
ax.plot(mons3, np.nanmean(damping_slab, (0, 1)), color='b', label="N Atl. Avg")
ax.legend()
ax.set_ylabel("W/m2/K")
# Save the masks as well
outname1 = "%sSLAB_PIC_NHFLX_Damping_monwin3_sig005_dof%i_mode%i_mask_%s.npy" % (picout,dofs[0],mode,lagstr)
outname0 = "%sFULL_PIC_NHFLX_Damping_monwin3_sig005_dof%i_mode%i_mask_%s.npy" % (picout,dofs[1],mode,lagstr)
outnames = [outname1,outname0]
for i in range(2):
np.save(outnames[i],malls[i])
print("Saved to "+outnames[i])
#%% Select point for comparison
# -----------------------------
lonf = -30
latf = 50
klon,klat = proc.find_latlon(lonf,latf,lon180,lat)
klon360,_ = proc.find_latlon(lonf+360,latf,lon,lat)
#%% Make some plots (Annual Heat Flux without the mask)
# -----------------------------------------------------
bboxplot = [-85,0,0,65]
proj = ccrs.PlateCarree()
fig,axs = plt.subplots(1,2,subplot_kw={'projection':proj})
clvls = np.arange(-40,41,1)
for i in range(2):
ax = axs[i]
ax = viz.add_coast_grid(ax,bbox=bboxplot)
def remove_GMSL(ssh,
lat,
lon,
times,
tol=1e-10,
viz=False,
testpoint=[330, 50]):
"""
Parameters
----------
ssh : ARRAY [time x lat x lon]
Sea Surface Height Anomalies to process
lat : ARRAY [lat]
Latitude values
lon : ARRAY [lon]
Longitude values
times : ARRAY [time]
Time values to plot (years)
tol : FLOAT, optional
Tolerance to check if GMSL is zero. The default is 1e-10.
viz : BOOL, optional
Visualize GMSL removal. The default is False.
testpoint : LIST [lon,lat]
Longitude and latitude points to visualize removal at. The default is [330,50].
Returns
-------
ssh : ARRAY [time x lat x lon]
Sea surface height anomalies with GMSL removed
gmslrem: ARRAY [time]
Time series that was removed
Additional outputs for viz == True:
fig,ax that was visualized
Note: Add latitude weights for future update...
"""
# Calculate GMSL (Not Area Weighted)
gmslrem = np.nanmean(ssh, (1, 2))
if np.any(gmslrem > tol):
# Remove GMSL
ssh_ori = ssh.copy()
ssh = ssh - gmslrem[:, None, None]
# Plot Results
if viz:
lonf, latf = testpoint
klon, klat = proc.find_latlon(lonf, latf, lon, lat)
fig, ax = plt.subplots(1, 1)
#ax.set_xticks(np.arange(0,len(times)+1,12))
ax.set_xticks(np.arange(0, len(times), 12))
ax.set_xticklabels(times[::12], rotation=45)
ax.grid(True, ls='dotted')
ax.plot(ssh_ori[:, klat, klon], label="Original", color='k')
ax.plot(ssh[:, klat, klon], label="Post-Removal")
ax.plot(gmslrem, label="GMSL")
ax.legend()
ax.set_title("GMSL Removal at Lon %.2f Lat %.2f (%s to %s)" %
(lon[klon], lat[klat], times[0], times[-1]))
ax.set_ylabel("SSH (m)")
return ssh, gmslrem, fig, ax
else:
print("GMSL has already been removed, largest value is %e" %
(gmslrem.max()))
return ssh, gmslrem
r"p = %.2f | $\rho$ > %.2f " % (p, corrthres),
fontsize=12)
plt.colorbar(pcm, ax=ax, orientation='horizontal', fraction=0.040, pad=0.05)
plt.savefig(outpath + "%s_Damping__mode%i_monwin%i_lags%i_sig%03d.png" %
(flux, mode, monwin, lagmax, p * 100),
dpi=200)
#%% Plot Damping Seasonal Cycle at a Point
lonf = -30
latf = 65
mons = np.arange(1, 13, 1)
locstring = "Lon%03d_Lat%03d" % (lonf, latf)
# Find point and get data
klon, klat = proc.find_latlon(lonf + 360, latf, lon, lat)
damppt = dampchoose[klon, klat, :, :] # [ ens x mon ]
if len(damppt.shape) > 2:
damppt = np.nanmean(damppt, 2)
# Plot seasonal cycle
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
for e in range(42):
ax.plot(mons, damppt[e, :], color=[.75, .75, .75])
ln1 = ax.plot(mons, damppt[-1, :], color=[.75, .75, .75], label="Indv. Member")
ln2 = ax.plot(mons, np.nanmean(damppt, 0), color='k', label="Ens. Avg.")
#ax.set_ylim([-30,90])
lns = ln1 + ln2
labs = [l.get_label() for l in lns]
nsmooths = [500,250] # Set Smothing
# Other Params
pct = 0.10
opt = 1
dt = 3600*24*30
tunit = "Months"
clvl = [0.95]
axopt = 3
clopt = 1
# Retrieve point
lonf,latf = config['query']
if lonf < 0:
lonf += 360
klon360,klat = proc.find_latlon(lonf,latf,lon360,lat)
fullpt = sstfull[:,klat,klon360]
slabpt = sstslab[:,klat,klon360]
# Calculate spectra
freq1s,P1s,CLs = [],[],[]
for i,sstin in enumerate([fullpt,slabpt]):
# Calculate and Plot
sps = ybx.yo_spec(sstin,opt,nsmooths[i],pct,debug=False)
P,freq,dof,r1=sps
# Plot if option is set
if plotcesm:
pps = ybx.yo_specplot(freq,P,dof,r1,tunit,dt=dt,clvl=clvl,axopt=axopt,clopt=clopt)
fig,ax,h,hcl,htax,hleg = pps
"_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)
yolbd = 1 / tauall.mean(1)
yofac = (1 - np.exp(-yolbd)) / yolbd
else:
lbd, lbd_entr, FAC, beta = scm.set_stochparams(hclim,
dampingr,
dt,
ND=1,
hfix=hfix)
sshss = np.zeros(sshc.shape) * np.nan
sshss[:, okpts] = ssh_ss
sshss = sshss.reshape(ntime, nlat5, nlon5)
# Try another removal method
ssha2 = ssha.copy()
clim, ssha2 = proc.calc_clim(ssha2, 0, returnts=1)
ssha2 = ssha2 - clim[None, :, :, :]
ssha2 = ssha2.reshape(ssha2.shape[0] * 12, nlat5, nlon5)
#ssha2 = ssha2.reshape(int(ntime/12)
# Plot sample removal
plotmons = 60
#klonss,klatss = proc.find_latlon(325,5,lon5,lat5)
klonss, klatss = proc.find_latlon(330, 50, lon5, lat5)
fig, axs = plt.subplots(2, 1)
ax = axs[0]
ax.plot(ssha[:plotmons, klatss, klonss], color='gray', label="Original")
ax.plot(sshnew[:plotmons, klatss, klonss],
color='k',
label="Deseasoned (Sinusoid)")
ax.plot(sshss[:plotmons, klatss, klonss],
color='red',
ls='dotted',
label="Sinusoid Fit")
ax.legend(ncol=3, fontsize=8)
ax.grid(True, ls='dotted')
ax.set_ylim([-.2, .2])
ax.set_ylabel("SSH (cm)")
orientation='horizontal',
fraction=0.030,
pad=0.02)
cb.set_label("Atmospheric Heat Flux Feedback ($Wm^{-2}K^{-1}$)")
# if notitle is False:
# plt.suptitle("Seasonal Mean Mixed Layer Depth Differences in meters \n (CESM1 - WOA 1994)",fontsize=14,y=.94)
plt.savefig("%sHFLX_values_FULL-SLAB_Savg_regional.png" % (figpath),
dpi=150,
bbox_inches='tight')
# ------------------------------
#%% Check Conditions at 1 point
# ------------------------------
lonf = -30 + 360
latf = 50
klon, klat = proc.find_latlon(lonf, latf, lon, lat)
flip = -1
locstring = "Lon %.2f ; Lat %.2f" % (lon[klon] - 360, lat[klat])
locfstring = "Lon_%i_Lat%i" % (lon[klon] - 360, lat[klat])
plotmean = True
# Plot heat flux feedback
fig, axs = plt.subplots(2, 1, figsize=(8, 8), sharey=True)
dampingpts = []
for imcf in range(2):
ax = axs.flatten()[imcf]
dampingpts.append(dampings[imcf][:, :, klat, klon])
for ilag in range(3):
ax.plot(mons3,
dampingpts[imcf][:, ilag] * -1,
label="Lag %i" % (ilag + 1),
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))
