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))