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