import xarray as xr import numpy as np import os import cartopy.crs as ccrs from scipy import signal from NorESM_utils import NorESM_utils as nutils def crf_calc_icevar(dpath,data,cmip,case,model,ens,icevar,ntime): ''' this is to calculate the Barents Sea ice concentration ''' BSice = data[model+'_'+case+'_'+ens+'_'+icevar][icevar] ice_dims = BSice.dims ice_coords = BSice.coords # figure out what the x,y dims are called # and what are the respective coordinates for dim in ice_dims: if dim in ['lon','nlon','ni','i','x','longitude']: xname = dim elif dim in ['lat','nlat','nj','j','y','latitude']: yname = dim # for coord in ice_coords: if 'lon' in coord: lonname = coord elif 'lat' in coord: latname = coord # if 'area' in list(data[model+'_'+case+'_'+ens+'_'+icevar].variables): areacello = data[model+'_'+case+'_'+ens+'_'+icevar].area elif os.path.exists(dpath+'/'+cmip+'/fixed/'+case+'/areacello/'+model+'/') and (model not in ['HadGEM3-GC31-HM', 'HadGEM3-GC31-LL', 'HadGEM3-GC31-MM']): areacello=xr.open_mfdataset(dpath+'/'+cmip+'/fixed/'+case+'/areacello/'+model+'/*.nc',combine='nested',concat_dim='dum').areacello.isel(dum=0) else: print(model) areacello=BSice.isel(time=0).drop('time').notnull() areacello=areacello.where(areacello==1)*np.cos(np.radians(areacello[latname])) # if 'time' in areacello.dims: areacello = areacello.isel(time=0).drop('time') # if (areacello.dims[0] not in BSice.dims) or (areacello.dims[1] not in BSice.dims): area_dims = areacello.dims areacello = areacello.rename({area_dims[0]:yname,area_dims[1]:xname}) # areacello_BS = areacello.where(BSice.isel(time=0).drop('time').notnull()) # # if len(BSice[lonname].shape)>2: mask_ice = np.logical_and(np.logical_and(BSice[lonname].isel(time=0).values>15,BSice[lonname].isel(time=0).values<60),np.logical_and(BSice[latname].isel(time=0).values>70,BSice[latname].isel(time=0).values<80)) else: if len(BSice[lonname].shape)<2: lon2,lat2 = np.meshgrid(BSice[lonname].values,BSice[latname].values) mask_ice = np.logical_and(np.logical_and(lon2>15,lon2<60),np.logical_and(lat2>70,lat2<80)) else: mask_ice = np.logical_and(np.logical_and(BSice[lonname].values>15,BSice[lonname].values<60),np.logical_and(BSice[latname].values>70,BSice[latname].values<80)) # BSice = BSice.where(mask_ice).dropna(dim=yname,how='all').dropna(dim=xname,how='all') BSarea = areacello_BS.where(mask_ice).dropna(dim=yname,how='all').dropna(dim=xname,how='all') # BSice_mean = ((BSice*BSarea).isel(time=slice(0,ntime,1)).sum(dim=xname).sum(dim=yname)/BSarea.sum()).persist() # return BSice_mean def crf_calc_predictor(dpath,data,cmip,case,model,ens,predictor,latnew,lon0,lon1,ntime): ''' this is to calculate the Barents Sea ice concentration ''' predictor_data = np.zeros((ntime,len(latnew))) if predictor in ['hfx','hfds','sos','tos']: ocn_dims = data[model+'_'+case+'_'+ens+'_'+predictor][predictor].dims ocn_coords = data[model+'_'+case+'_'+ens+'_'+predictor][predictor].coords # figure out what the x,y dims are called # and what are the respective coordinates for dim in ocn_dims: if dim in ['lon','nlon','ni','i','x','longitude']: xname = dim elif dim in ['lat','nlat','nj','j','y','latitude']: yname = dim # for coord in ocn_coords: if ('lon' in coord) and (coord!='nlon'): lonname = coord elif ('lat' in coord) and (coord!='nlat'): latname = coord # if predictor in ['hfx']: # #ytransport = all_out[model+'_hfy'].isel(time=slice(0,120)).fillna(0).load() #xtransport = all_out[model+'_hfx'].isel(time=slice(0,120)).fillna(0).load() # lat2 = data[model+'_'+case+'_'+ens+'_hfds'][latname] lon2 = data[model+'_'+case+'_'+ens+'_hfds'][lonname] if 'time' in lat2.dims: lat2=lat2.isel(time=0) lon2=lon2.isel(time=0) # # move to numpy lon2 = lon2.values lat2 = lat2.values lon2[np.where(lon2>180)] = lon2[np.where(lon2>180)] - 360 # line_lats = [] line_lons = [] # HERE WE ROTATE THE POLE SO THAT WE GET LINES THAT ARE MORE ALIGNED # WITH THE COASTAL CURRENT - NOTE THAT THIS DOES NOT AFFECT THE REST # OF THE CALCULATIONS WHERE WE SIMPLY CALCULATE THE TRANSPORT ACROSS THE LINE rotpole = ccrs.RotatedPole(pole_longitude=80.0, pole_latitude=80.0) out = rotpole.transform_points(ccrs.PlateCarree(),lon2,lat2) rotlon = out[:,:,0] rotlat = out[:,:,1] rotdum = out[:,:,2] # calc_mean = True for l,lat0 in enumerate(latnew): # print(l) iind,jind=nutils.latitude_line(lat0, rotlat,bipolar=False) # line_lat = lat2[jind,iind] line_lon = lon2[jind,iind] line_lon[np.where(line_lon>180)] = line_lon[np.where(line_lon>180)]-360 # ii = np.where(np.logical_and(line_lon>lon0[l],line_lon1: i1=1+np.where(np.diff(ii).max()==np.diff(ii))[0][0] ii=np.concatenate([ii[i1:], ii[:i1]],axis=0) #ii = np.where(np.logical_and(line_lonrot>30,line_lonrot<100))[0] line_lats.append(line_lat[ii]) line_lons.append(line_lon[ii]) # calculate the sum iind2=iind[slice(min(ii),max(ii))] jind2=jind[slice(min(ii),max(ii))] # xtransport = data[model+'_'+case+'_'+ens+'_hfx'].hfx.rename({xname:'x',yname:'y'}) ytransport = data[model+'_'+case+'_'+ens+'_hfy'].hfy.rename({xname:'x',yname:'y'}) # take a halo of 1 in each direction xtransport = xtransport.isel(y=slice(min(jind2)-1,max(jind2)+2),x=slice(min(iind2)-1,max(iind2)+2)).fillna(0).persist() ytransport = ytransport.isel(y=slice(min(jind2)-1,max(jind2)+2),x=slice(min(iind2)-1,max(iind2)+2)).fillna(0).persist() #xtransport = data[model+'_'+case+'_'+ens+'_hfx'].hfx.isel(y=slice(min(jind2),max(jind2)),x=slice(min(iind2),max(iind2))).fillna(0).persist() #.dropna(dim=xname,how='all').persist() #ytransport = data[model+'_'+case+'_'+ens+'_hfy'].hfy.isel(y=slice(min(jind2),max(jind2)),x=slice(min(iind2),max(iind2))).fillna(0).persist() #.dropna(dim=xname,how='all').persist() # # iind2=np.array(iind2)-min(iind2)+1 # jind2=np.array(jind2)-min(jind2)+1 # #if model in ['IPSL-CM6A-LR','CNRM-CM5-2','MPI-ESM-LR','GFDL-ESM2M']: #htransport[:,l]=nutils.heat_transport_NEMO(np.array(iind2),np.array(jind2),xtransport,ytransport) if model not in ['NorESM2-LM']:#['IPSL-CM6A-LR','HadGEM3-GC31-LL','UKESM1-0-LL','CNRM-CM5-2','CNRM-CM6-1-HR','MPI-ESM-LR','GFDL-ESM2M']: sumtot, inds_out = nutils.heat_transport_NEMO(iind2,jind2,xtransport,ytransport) elif model in ['NorESM2-LM']: sumtot, inds_out = nutils.heat_transport_NorESM(iind2,jind2,xtransport,ytransport) # #sumtot = np.reshape(signal.detrend(sumtot.values),(-1,12)) sumtot = signal.detrend(np.reshape(sumtot.values,(-1,12)),axis=0) # predictor_data[:,l] = (sumtot/np.nanstd(sumtot,axis=0)).flatten() #else: # htransport[:,l]=nutils.heat_transport_NorESM(np.array(iind),np.array(jind),xtransport.values,ytransport.values) # elif predictor in ['hfds','sos','tos']: # if 'area' in list(data[model+'_'+case+'_'+ens+'_'+predictor].variables): areacello = data[model+'_'+case+'_'+ens+'_'+predictor].area elif os.path.exists(dpath+'/'+cmip+'/fixed/'+case+'/areacello/'+model+'/'): areacello=xr.open_mfdataset(dpath+'/'+cmip+'/fixed/'+case+'/areacello/'+model+'/*.nc',combine='nested',concat_dim='dum').areacello.isel(dum=0) if 'time' in areacello.dims: areacello = areacello.isel(time=0).drop('time') else: print(model) areacello=data[model+'_'+case+'_'+ens+'_'+predictor][predictor].isel(time=0).drop('time').notnull() areacello=areacello.where(areacello==1).astype(float)*np.cos(np.radians(areacello[latname])) # lat2 = data[model+'_'+case+'_'+ens+'_'+predictor][latname] lon2 = data[model+'_'+case+'_'+ens+'_'+predictor][lonname] # if 'time' in lat2.dims: lat2 = lat2.isel(time=0) lon2 = lon2.isel(time=0) # if len(lon2.shape)<2: lon2,lat2 = np.meshgrid(lon2,lat2) else: lat2 = lat2.values lon2 = lon2.values # lon2[np.where(lon2>180)] = lon2[np.where(lon2>180)] - 360 # # HERE WE ROTATE THE POLE SO THAT WE GET LINES THAT ARE MORE ALIGNED # WITH THE COASTAL CURRENT - NOTE THAT THIS DOES NOT AFFECT THE REST # OF THE CALCULATIONS WHERE WE SIMPLY CALCULATE THE TRANSPORT ACROSS THE LINE rotpole = ccrs.RotatedPole(pole_longitude=80.0, pole_latitude=80.0) out = rotpole.transform_points(ccrs.PlateCarree(),lon2,lat2) rotlon = out[:,:,0] rotlat = out[:,:,1] rotdum = out[:,:,2] # calc_mean = True for l,lat0 in enumerate(latnew): # print(l) mask1 = np.logical_and(rotlat<=lat0+1,rotlat>=lat0-1) mask2 = np.logical_and(lon2>lon0[l],lon2