def derive_SST_avg_index(self, run, index, dsdt='ds_dt', time=None):
""" generates all area avg indices from detrended SST data """
assert run in ['ctrl', 'rcp', 'lpd', 'lpi', 'had']
assert time is None or len(time)==2
if run=='had': domain, dims, ts = 'ocn_had' , ('latitude', 'longitude'), ''
elif run in ['ctrl', 'rcp']: domain, dims, ts = 'ocn_rect', ('t_lat', 't_lon'), f'_{time[0]}_{time[1]}'
elif run in ['lpd', 'lpi']: domain, dims, ts = 'ocn_low' , ('nlat', 'nlon') , f'_{time[0]}_{time[1]}'
fn_monthly = f'{path_prace}/SST/SST_monthly_{dsdt}_{run}{ts}.nc'
SST_monthly = xr.open_dataarray(fn_monthly, decode_times=False)
if index in ['AMO', 'SOM', 'SMV']:
blats, blons, mask_nr = self.bounding_lats_lons(index)
MASK = mask_box_in_region(domain=domain, mask_nr=mask_nr, bounding_lats=blats, bounding_lons=blons)
AREA = xr_AREA(domain=domain).where(MASK)
SST_index = self.SST_area_average(xa_SST=SST_monthly, AREA=AREA, AREA_index=AREA.sum(), MASK=MASK, dims=dims)
if index=='TPI': # monthly data
for i, TPI_i in enumerate(['TPI1', 'TPI2', 'TPI3']):
blats, blons, mask_nr = self.bounding_lats_lons(TPI_i)
MASK = mask_box_in_region(domain=domain, mask_nr=mask_nr, bounding_lats=blats, bounding_lons=blons)
AREA = xr_AREA(domain=domain).where(MASK)
TPI_ = self.SST_area_average(xa_SST=SST_monthly, AREA=AREA, AREA_index=AREA.sum(), MASK=MASK, dims=dims)
TPI_.to_netcdf(f'{path_prace}/SST/{TPI_i}_ds_dt_raw_{run}{ts}.nc')
if i==0: SST_index = -0.5*TPI_
elif i==1: SST_index = SST_index + TPI_
elif i==2: SST_index = SST_index - 0.5*TPI_
SST_index.to_netcdf(f'{path_prace}/SST/{index}_{dsdt}_raw_{run}{ts}.nc')
return SST_index
def spatial_correlation(self,
field_A,
field_B,
method=None,
selection=None):
""" correlate two 2D fields """
if np.shape(field_A) != np.shape(field_B): # have to regrid
A, B = self.regrid_to_lower_resolution(field_A, field_B)
else:
A, B = field_A, field_B
assert np.shape(A) == np.shape(B)
domain = self.determine_domain(A)
AREA = xr_AREA(domain)
MASK = boolean_mask(domain=domain, mask_nr=0)
if type(selection) == int:
MASK = boolean_mask(domain=domain, mask_nr=selection)
elif type(selection) == dict:
MASK, AREA = MASK.sel(selection), AREA.sel(selection)
A, B = A.sel(selection), B.sel(selection)
D = np.any(np.array(
[np.isnan(A).values,
np.isnan(B).values, (MASK == 0).values]),
axis=0)
A = xr.where(D, np.nan,
A).stack(z=('latitude', 'longitude')).dropna(dim='z')
B = xr.where(D, np.nan,
B).stack(z=('latitude', 'longitude')).dropna(dim='z')
C = xr.where(D, np.nan,
AREA).stack(z=('latitude', 'longitude')).dropna(dim='z')
d = DescrStatsW(np.array([A.values, B.values]).T, weights=C)
spatial_corr_coef = d.corrcoef[0, 1]
return spatial_corr_coef
def make_SALT_vol_integrals(run):
""" calculate SALT volume integrals for specific regions from depth integrated SALT*DZT maps
[g->kg] * (SALT*DZT).sum('z_t') * rho_w * \int \int dx dy
1kg/1000g * g_S/kg_W*m * 1000kg_W/1m^3 * m^2 = kg
output:
netcdf with timeseries and trends for all depth levels
"""
# Atlantic + Labrador + GIN, no Med
dd, ddd = {}, []
if run in ['ctrl', 'rcp']: AREA = xr_AREA('ocn')
elif run in ['lpd', 'lr1']: AREA = xr_AREA('ocn_low')
dm = [
xr.open_dataarray(f'{path_prace}/SALT/SALT_dz_0-100m_{run}.nc'),
xr.open_dataarray(f'{path_prace}/SALT/SALT_dz_0-1000m_{run}.nc'),
xr.open_dataarray(f'{path_prace}/SALT/SALT_dz_below_1000m_{run}.nc')
]
dt = [
xr.open_dataarray(f'{path_prace}/SALT/SALT_dz_0-100m_trend_{run}.nc'),
xr.open_dataarray(f'{path_prace}/SALT/SALT_dz_0-1000m_trend_{run}.nc'),
xr.open_dataarray(
f'{path_prace}/SALT/SALT_dz_below_1000m_trend_{run}.nc')
]
for j, (latS, latN) in enumerate(lat_bands):
MASK = make_Mask(run, latS, latN)
for d, depth in notebook.tqdm(
enumerate(['0-100m', '0-1000m', 'below_1000m'])):
dm_, dt_ = dm[d], dt[d]
tseries = (dm_ * AREA).where(MASK).sum(dim=['nlat', 'nlon'
]) # g/kg*m*m^2=kg
trend = (dt_ * AREA).where(MASK).sum(dim=['nlat', 'nlon'])
tseries.name = f'SALT_{depth}_timeseries_{latS}N_{latN}N'
trend.name = f'SALT_{depth}_trend_{latS}N_{latN}N'
dd[f'SALT_{depth}_timeseries_{latS}N_{latN}N'] = tseries
dd[f'SALT_{depth}_trend_{latS}N_{latN}N'] = trend
ddd.append(tseries)
ddd.append(trend)
# print(f'{run:4}', f'{latS:4}', f'{latN:4}', f'{salt.values:4.1e}')
xr.merge(ddd).to_netcdf(f'{path_results}/SALT/SALT_integrals_{run}.nc')
return
def generate_yrly_global_mean_SST(self, run):
""" calcaultes the global mean sea surface temperature
ca. 37 sec for ctrl """
assert run in ['ctrl', 'lpd']
da = xr.open_dataarray(f'{path_prace}/SST/SST_yrly_{run}.nc',
decode_times=False)
if run == 'ctrl':
AREA = xr_AREA(domain='ocn')
REGION_MASK = xr.open_dataset(file_ex_ocn_ctrl,
decode_times=False).REGION_MASK
elif run == 'lpd':
AREA = xr_AREA(domain='ocn_low')
REGION_MASK = xr.open_dataset(file_ex_ocn_lpd,
decode_times=False).REGION_MASK
AREA_total = AREA.where(REGION_MASK > 0).sum(dim=['nlat', 'nlon'],
skipna=True)
print(AREA_total)
da_new = (da * AREA).where(REGION_MASK > 0).sum(
dim=['nlat', 'nlon'], skipna=True) / AREA_total
da_new.to_netcdf(f'{path_prace}/SST/GMSST_yrly_{run}.nc')
return
def CESM_xlfca(run, basin, dsdt, test=False):
""" performing the LFCA via xlfca of our model output """
print(run, basin, dsdt)
# performance had: North Pacific 20 N, time steps 10: 5.47 s, 100: 27.8 s, 1788 (all): 477s -> 142 MB
if run == 'had': dt, domain = '', 'ocn_had'
elif run == 'ctrl': dt, domain = '_51_301', 'ocn_rect'
elif run == 'lpd': dt, domain = '_154_404', 'ocn_low'
if basin == 'North_Pacific':
mask_nr, bounding_lats, bounding_lons = 2, (20, 68), (110, 255)
if basin == 'full_Pacific':
mask_nr, bounding_lats, bounding_lons = 2, (-38, 68), (110, 290)
if basin == 'Southern_Ocean':
mask_nr, bounding_lats, bounding_lons = 1, None, None
if basin == 'North_Atlantic':
mask_nr, bounding_lats, bounding_lons = 6, (0, 60), (-80, 0)
fn = f'{path_prace}/SST/SST_monthly_{dsdt}_{run}{dt}.nc'
MASK = mask_box_in_region(domain=domain,
mask_nr=mask_nr,
bounding_lats=bounding_lats,
bounding_lons=bounding_lons)
AREA = xr_AREA(domain=domain).where(MASK)
SST = xr.open_dataarray(fn, decode_times=False).where(MASK)
if basin in ['North_Pacific', 'full_Pacific'] and run == 'had': # shifting
AREA = DS().shift_had(AREA)
SST = DS().shift_had(SST)
if basin == 'North_Atlantic' and run == 'ctrl':
AREA = DS().shift_ocn_rect(AREA)
SST = DS().shift_ocn_rect(SST)
if basin == 'North_Atlantic' and run == 'lpd':
AREA = DS().shift_ocn_low(AREA)
SST = DS().shift_ocn_low(SST)
AREA = AREA.where(np.isnan(SST[0, :, :]) == False, drop=True)
SST = SST.where(np.isnan(SST[0, :, :]) == False, drop=True)
scale = AREA / AREA.sum()
scale = xr.apply_ufunc(np.sqrt, scale)
for n_EOFs in [3, 30]:
fn_lfca = f'{path_prace}/LFCA/LFCA_{run}_{basin}_{dsdt}_n{n_EOFs}.nc'
if os.path.exists(fn_lfca): continue
if test:
lfca = xlfca(x=SST.isel(time=slice(0, 40)),
cutoff=120,
truncation=n_EOFs,
scale=scale)
else:
lfca = xlfca(x=SST, cutoff=120, truncation=n_EOFs, scale=scale)
lfca.to_netcdf(fn_lfca)
return lfca
def SST_index_from_monthly(run, index_loc, MASK):
""" loads monthly SST data, calculated SST_index, returns raw timeseries"""
assert run in ['ctrl', 'rcp', 'lpd', 'lpi']
if run in ['ctrl', 'rcp']:
domain = 'ocn'
elif run in ['lpd', 'lpi']:
domain = 'ocn_low'
AREA = xr_AREA(domain)
AREA_index = AREA.where(MASK).sum()
for i, (y, m, s) in enumerate(
IterateOutputCESM(domain=domain, run=run, tavg='monthly')):
if m == 1: print(y)
xa_SST = xr.open_dataset(s, decode_times=False).TEMP[0, 0, :, :]
SSTi = SST_index(xa_SST, AREA, index_loc, AREA_index, MASK)
if i == 0:
new_SSTi = SSTi.copy()
else:
new_SSTi = xr.concat([new_SSTi, SSTi], dim='time')
return new_SSTi
def surface_heat_flux(self, run):
""" total surface heat flux into ocean basins """
# 32:20 min ctrl
# 1min 4s lpd
if run == 'ctrl': domain = 'ocn'
elif run in ['lc1', 'lpd']: domain = 'ocn_low'
da = xr.open_mfdataset(f'{path_prace}/{run}/ocn_yrly_SHF_0*.nc',
combine='nested',
concat_dim='time').SHF
print(len(da.time))
AREA = xr_AREA(domain=domain)
SHF = spy * (da * AREA).sum(dim=['nlat', 'nlon'])
SHF.name = 'Global_Ocean'
for nr in tqdm(np.arange(1, 12)):
MASK = boolean_mask(domain=domain, mask_nr=nr)
temp = spy * (da * AREA).where(MASK).sum(dim=['nlat', 'nlon'])
temp.name = regions_dict[nr]
SHF = xr.merge([SHF, temp])
SHF.attrs[
'quantity'] = 'yrly averaged total surface heat flux, positive down'
SHF.attrs['units'] = '[J/yr]'
SHF.to_netcdf(f'{path_prace}/OHC/SHF_{run}.nc')
return
def all_transports(self, run, quantity):
""" computes heat or salt fluxes """
assert run in ['ctrl', 'lpd']
assert quantity in ['SALT', 'OHC']
if quantity == 'OHC':
VN, UE = 'VNT', 'UET'
conversion = rho_sw * cp_sw
qstr = 'heat'
unit_out = 'W'
elif quantity == 'SALT':
VN, UE = 'VNS', 'UES'
conversion = rho_sw * 1e-3
qstr = 'salt'
unit_out = 'kg/s'
if run == 'ctrl':
domain = 'ocn'
all_transports_list = []
elif run == 'lpd':
domain = 'ocn_low'
mf_fn = f'{path_prace}/{run}/ocn_yrly_{VN}_{UE}_*.nc'
kwargs = {
'concat_dim': 'time',
'decode_times': False,
'drop_variables': ['TLONG', 'TLAT', 'ULONG', 'ULAT'],
'parallel': True
}
ds = xr.open_mfdataset(mf_fn, **kwargs)
DZ = xr_DZ(domain=domain)
adv = self.all_advection_cells(domain=domain)
AREA = xr_AREA(domain=domain)
dims = [dim for dim in dll_dims_names(domain=domain)]
for i, pair in enumerate(tqdm(neighbours)):
name = f'{qstr}_flux_{regions_dict[pair[0]]}_to_{regions_dict[pair[1]]}'
# if i>2: continue
adv_E = adv[
f'adv_E_{regions_dict[pair[0]]}_to_{regions_dict[pair[1]]}']
adv_N = adv[
f'adv_N_{regions_dict[pair[0]]}_to_{regions_dict[pair[1]]}']
MASK = ((abs(adv_E) + abs(adv_N)) /
(abs(adv_E) + abs(adv_N))).copy()
adv_E = adv_E.where(MASK == 1, drop=True)
adv_N = adv_N.where(MASK == 1, drop=True)
DZ_ = DZ.where(MASK == 1, drop=True)
AREA_ = AREA.where(MASK == 1, drop=True)
if run == 'ctrl':
for j, (y,m,f) in tqdm(enumerate(IterateOutputCESM(domain='ocn', run='ctrl',\
tavg='yrly', name=f'{VN}_{UE}'))):
# if j>1: continue
ds = xr.open_dataset(f,
decode_times=False).where(MASK == 1,
drop=True)
transport = ((adv_E * ds[UE] + adv_N * ds[VN]) * AREA_ *
DZ_).sum(dim=dims) * conversion
transport.name = name
transport.attrs['units'] = unit_out
if j == 0: transport_t = transport
else:
transport_t = xr.concat([transport_t, transport],
dim='time')
all_transports_list.append(transport_t)
elif run == 'lpd':
ds_ = ds.where(MASK == 1, drop=True)
transport = ((adv_E * ds_[UE] + adv_N * ds_[VN]) * AREA_ *
DZ_).sum(dim=dims) * conversion
transport.name = name
transport.attrs['units'] = unit_out
if i == 0: all_transports = transport
else: all_transports = xr.merge([all_transports, transport])
if run == 'ctrl': all_transports = xr.merge(all_transports_list)
all_transports.to_netcdf(
f'{path_prace}/{quantity}/{quantity}_fluxes_{run}.nc')
return all_transports
def generate_OHC_files(self, run, year=None, pwqd=False):
""" non-detrended OHC files for full length of simulations
One file contains integrals (all global and by basin):
x,y,z .. scalars
x,y .. vertical profiles
x .. "zonal" integrals
A separate file each for 4 different depth levels
z .. 2D maps, global only, but for different vertical levels
# (ocn: takes about 45 seconds per year: 70 yrs approx 55 mins)
(ocn: takes about 14 min per year)
(ocn_rect: takes about 3 seconds per year: 70 yrs approx 3 mins)
"""
def t2da(da, t):
"""adds time dimension to xr DataArray, then sets time value to t"""
da = da.expand_dims('time')
da = da.assign_coords(time=[t])
return da
def t2ds(da, name, t):
"""
adds time dimension to xr DataArray, then sets time value to t,
and then returns as array in xr dataset
"""
da = t2da(da, t)
ds = da.to_dataset(name=name)
return ds
start = datetime.datetime.now()
def tss(): # time since start
return datetime.datetime.now()-start
print(f'{start} start OHC calculation: run={run}')
assert run in ['ctrl', 'rcp', 'lpd', 'lpi']
if run=='rcp':
domain = 'ocn'
elif run=='ctrl':
domain = 'ocn_rect'
elif run in ['lpd', 'lpi']:
domain = 'ocn_low'
(z, lat, lon) = dll_dims_names(domain)
# geometry
DZT = xr_DZ(domain)#.chunk(chunks={z:1})
AREA = xr_AREA(domain)
HTN = xr_HTN(domain)
LATS = xr_LATS(domain)
def round_tlatlon(das):
""" rounds TLAT and TLONG to 2 decimals
some files' coordinates differ in their last digit
rounding them avoids problems in concatonating
"""
das['TLAT'] = das['TLAT'].round(decimals=2)
das['TLONG'] = das['TLONG'].round(decimals=2)
return das
if domain=='ocn':
round_tlatlon(HTN)
round_tlatlon(LATS)
MASK = boolean_mask(domain, mask_nr=0)
DZT = DZT.where(MASK)#.chunk(chunks={z:1})
# with chunking for ctrl_rect: 21 sec per iteration, 15 sec without
# for k in range(42):
# DZT[k,:,:] = DZT[k,:,:].where(MASK)
AREA = AREA.where(MASK)
HTN = HTN.where(MASK)
LATS = LATS.where(MASK)
# print(f'{datetime.datetime.now()} done with geometry')
if pwqd: name = 'TEMP_pwqd'
else: name = 'TEMP_PD'
# print(run, domain, name)
for y,m,file in IterateOutputCESM(domain=domain, run=run, tavg='yrly', name=name):
# print(tss(), y)
# break
if year!=None: # select specific year
if year==y:
pass
else:
continue
if pwqd: file_out = f'{path_samoc}/OHC/OHC_integrals_{run}_{y:04d}_pwqd.nc'
else: file_out = f'{path_samoc}/OHC/OHC_integrals_{run}_{y:04d}.nc'
if os.path.exists(file_out) and year is None:
# # should check here if all the fields exist
# print(f'{datetime.datetime.now()} {y} skipped as files exists already')
# if y not in [250,251,252,253,254,255,273,274,275]:
continue
print(f'{tss()} {y}, {file}')
t = y*365 # time in days since year 0, for consistency with CESM date output
# ds = xr.open_dataset(file, decode_times=False, chunks={z:1}).TEMP
ds = xr.open_dataset(file, decode_times=False).TEMP
print(f'{tss()} opened dataset')
if domain=='ocn':
ds = ds.drop(['ULONG', 'ULAT'])
ds = round_tlatlon(ds)
# if ds.PD[0,150,200].round(decimals=0)==0:
# ds['PD'] = ds['PD']*1000 + rho_sw
# elif ds.PD[0,150,200].round(decimals=0)==1:
# ds['PD'] = ds['PD']*1000
# else:
# print('density [g/cm^3] is neither close to 0 or 1')
# OHC = ds.TEMP*ds.PD*cp_sw
OHC = ds*rho_sw*cp_sw
ds.close()
OHC = OHC.where(MASK)
OHC_DZT = OHC*DZT
print(f'{tss()} {y} calculated OHC & OHC_DZT')
# global, global levels, zonal, zonal levels integrals for different regions
for mask_nr in tqdm([0,1,2,3,6,7,8,9,10]):
# for mask_nr in [0,1,2,3,6,7,8,9,10]:
name = regions_dict[mask_nr]
da = OHC.where(boolean_mask(domain, mask_nr=mask_nr))
da_g = (da*AREA*DZT).sum(dim=[z, lat, lon])
da_g.attrs['units'] = '[J]'
ds_g = t2ds(da_g , f'OHC_{name}', t)
da_gl = (da*AREA).sum(dim=[lat, lon])
da_gl.attrs['units'] = '[J m^-1]'
ds_gl = t2ds(da_gl, f'OHC_levels_{name}', t)
if domain=='ocn': da_z = xr_int_zonal(da=da, HTN=HTN, LATS=LATS, AREA=AREA, DZ=DZT)
else: da_z = (da*HTN*DZT).sum(dim=[z, lon])
da_z.attrs['units'] = '[J m^-1]'
ds_z = t2ds(da_z , f'OHC_zonal_{name}', t)
if domain=='ocn': da_zl = xr_int_zonal_level(da=da, HTN=HTN, LATS=LATS, AREA=AREA, DZ=DZT)
else: da_zl = (da*HTN).sum(dim=[lon])
da_zl.attrs['units'] = '[J m^-2]'
ds_zl = t2ds(da_zl, f'OHC_zonal_levels_{name}', t)
if mask_nr==0: ds_new = xr.merge([ds_g, ds_gl, ds_z, ds_zl])
else: ds_new = xr.merge([ds_new, ds_g, ds_gl, ds_z, ds_zl])
print(f'{tss()} done with horizontal calculations')
# vertical integrals
# full depth
da_v = OHC_DZT.sum(dim=z) # 0-6000 m
da_v.attrs = {'depths':f'{OHC_DZT[z][0]-OHC_DZT[z][-1]}',
'units':'[J m^-2]'}
if domain in ['ocn', 'ocn_rect']: zsel = [[0,9], [0,20], [20,26]]
elif domain=='ocn_low': zsel = [[0,9], [0,36], [36,45]]
# 0- 100 m
da_va = OHC_DZT.isel({z:slice(zsel[0][0], zsel[0][1])}).sum(dim=z)
da_va.attrs = {'depths':f'{OHC_DZT[z][zsel[0][0]].values:.0f}-{OHC_DZT[z][zsel[0][1]].values:.0f}',
'units':'[J m^-2]'}
# 0- 700 m
da_vb = OHC_DZT.isel({z:slice(zsel[1][0],zsel[1][1])}).sum(dim=z)
da_vb.attrs = {'depths':f'{OHC_DZT[z][zsel[1][0]].values:.0f}-{OHC_DZT[z][zsel[1][1]].values:.0f}',
'units':'[J m^-2]'}
# 700-2000 m
da_vc = OHC_DZT.isel({z:slice(zsel[2][0],zsel[2][1])}).sum(dim=z)
da_vc.attrs = {'depths':f'{OHC_DZT[z][zsel[2][0]].values:.0f}-{OHC_DZT[z][zsel[2][1]].values:.0f}',
'units':'[J m^-2]'}
ds_v = t2ds(da_v , 'OHC_vertical_0_6000m' , t)
ds_va = t2ds(da_va, 'OHC_vertical_0_100m' , t)
ds_vb = t2ds(da_vb, 'OHC_vertical_0_700m' , t)
ds_vc = t2ds(da_vc, 'OHC_vertical_700_2000m', t)
ds_new = xr.merge([ds_new, ds_v, ds_va, ds_vb, ds_vc])
print(f'{tss()} done making datasets')
# print(f'output: {file_out}\n')
ds_new.to_netcdf(path=file_out, mode='w')
ds_new.close()
# if y in [2002, 102, 156, 1602]: break # for testing only
# combining yearly files
print(f'{datetime.datetime.now()} done\n')
# if run=='ctrl': print('year 205 is wrong and should be averaged by executing `fix_ctrl_year_205()`')
return
def GMST_timeseries(run):
""" builds a timesries of the GMST and saves it to a netCDF
input:
run .. (str) ctrl or cp
output:
ds_new .. xr Dataset containing GMST and T_zonal
"""
domain = 'atm'
tavg = 'yrly'
name = 'T_T850_U_V'
if run in ['ctrl', 'rcp', 'hq']: AREA = xr_AREA('atm')
elif run=='lpi': AREA = xr_AREA('atm_f19')
elif run in ['lpd', 'lc1', 'lr1', 'lq']: AREA = xr_AREA('atm_f09')
AREA_lat = AREA.sum(dim='lon')
AREA_total = AREA.sum(dim=('lat','lon'))
if run in ['lpd']: name = None
ny = len(IterateOutputCESM(domain=domain, run=run, tavg=tavg, name=name))
first_yr = IterateOutputCESM(domain=domain, run=run, tavg=tavg, name=name).year
iterator = IterateOutputCESM(domain=domain, run=run, tavg=tavg, name=name)
years = (np.arange(ny) + first_yr)*365 # this is consistent with CESM output
for i, (y, m, file) in enumerate(iterator):
print(y)
assert os.path.exists(file)
if run in ['ctrl', 'rcp', 'lpi', 'hq']:
da = xr.open_dataset(file, decode_times=False)['T'][-1,:,:]
elif run in ['lpd', 'lr1', 'lq', 'ld']:
da = xr.open_dataset(file, decode_times=False)['T'][0,-1,:,:]
if i==0: # create new xr Dataset
lats = da.lat.values
ds_new = xr.Dataset()
ds_new['GMST'] = xr.DataArray(data=np.zeros((ny)),
coords={'time': years},
dims=('time'))
ds_new['T_zonal'] = xr.DataArray(data=np.zeros((ny, len(lats))),
coords={'time': years, 'lat': lats},
dims=('time', 'lat'))
ds_new['GMST'][i] = (da*AREA).sum(dim=('lat','lon'))/AREA_total
ds_new['T_zonal'][i,:] = (da*AREA).sum(dim='lon')/AREA_lat
# [K] to [degC]
for field in ['GMST', 'T_zonal']:
ds_new[field] = ds_new[field] + abs_zero
# rolling linear trends [degC/yr]
ds_new = rolling_lin_trends(ds=ds_new, ny=ny, years=years)
# fits
lfit = np.polyfit(np.arange(ny), ds_new.GMST, 1)
qfit = np.polyfit(np.arange(ny), ds_new.GMST, 2)
ds_new[f'lin_fit'] = xr.DataArray(data=np.empty((len(ds_new['GMST']))),
coords={'time': years},
dims=('time'),
attrs={'lin_fit_params':lfit})
ds_new[f'quad_fit'] = xr.DataArray(data=np.empty((len(ds_new['GMST']))),
coords={'time': years},
dims=('time'),
attrs={'quad_fit_params':qfit})
for t in range(ny):
ds_new[f'lin_fit'][t] = lfit[0]*t + lfit[1]
ds_new[f'quad_fit'][t] = qfit[0]*t**2 + qfit[1]*t + qfit[2]
ds_new.to_netcdf(path=f'{path_prace}/GMST/GMST_{run}.nc', mode='w')
return ds_new
def SST_index(index, run, detrend_signal='GMST', time_slice='full'):
""" calcalates SST time series from yearly detrended SST dataset """
assert index in ['AMO', 'SOM', 'TPI1', 'TPI2', 'TPI3']
assert run in ['ctrl', 'rcp', 'lpd', 'lpi', 'had']
assert detrend_signal in ['GMST', 'AMO', 'SOM', 'TPI1', 'TPI2', 'TPI3']
print(index, run)
if run in ['ctrl', 'rcp']:
domain = 'ocn' #check this
dims = ('nlat', 'nlon')
elif run in ['lpd', 'lpi']:
domain = 'ocn_low'
dims = ('nlat', 'nlon')
elif run == 'had':
domain = 'ocn_had'
dims = ('latitude', 'longitude')
blats, blons, mask_nr = bounding_lats_lons(index)
MASK = mask_box_in_region(domain=domain,
mask_nr=mask_nr,
bounding_lats=blats,
bounding_lons=blons)
AREA = xr_AREA(domain=domain).where(MASK)
index_area = AREA.sum()
if run == 'had' and detrend_signal in ['AMO', 'SOM'
] or detrend_signal == 'GMST':
print(
f'underlying SST field: detrended with {detrend_signal}, no filtering'
)
if detrend_signal == 'GMST':
print('GMST(t) signal scaled at each grid point\n')
else:
print(
f'{detrend_signal}(t) removed from all SST gridpoints without scaling\n'
)
if time_slice == 'full':
fn = f'{path_samoc}/SST/SST_{detrend_signal}_dt_yrly_{run}.nc'
trange = ''
else:
first_year, last_year = determine_years_from_slice(
run=run, tres='yrly', time_slice=time_slice)
trange = f'_{first_year}_{last_year}'
fn = f'{path_samoc}/SST/SST_{detrend_signal}_dt_yrly{trange}_{run}.nc'
assert os.path.exists(fn)
SST_yrly = xr.open_dataarray(fn).where(MASK)
detr = f'_{detrend_signal}_dt'
else: # run=='had' and detrend_signal!='GMST'
print('underlying SST field: no detrending, no filtering')
if detrend_signal in ['AMO', 'SOM']:
print(
f'{detrend_signal} must subsequently be detrended with polynomial\n'
)
else:
print(
f'{detrend_signal} must not be detrended since forcing signal compensated in TPI\n'
)
SST_yrly = xr.open_dataarray(
f'{path_samoc}/SST/SST_yrly_{run}.nc').where(MASK)
detr = ''
SSTindex = SST_area_average(xa_SST=SST_yrly,
AREA=AREA,
AREA_index=index_area,
MASK=MASK,
dims=dims)
SSTindex.to_netcdf(f'{path_samoc}/SST/{index}{detr}_raw{trange}_{run}.nc')
return SSTindex
def PMV_EOF_indices(run, extent):
""" perform EOF of monthly, deseasonalized SST (detrended with global mean SST)
NB: we will use 60S to 60N SST data as polar data is limited in observations
1. for detrending: compute monthly global mean SST time series, deseasonalize them
>>> see `SST_data_generation.py` file for scripts
2. North Pacific monthly output fields
2.1. create monthly SST field
(if appropriate: determine extend of grid, limit all coordinates)
save as single file
a) North of 38 deg S
b) North of Equator
b) North of 20 deg N
2.2. deseasonalize
2.3. detrend global mean, deseasonalized SST,
2.4. (remove mean at each point)
3. EOF analysis
--> index is first principal component
4. regress time series on global maps
goal: perform EOF of monthly, deseasonalized SST (detrended with global mean SST)
NB: we will use 60S to 60N SST data as polar data is limited in observations
4. regress time series on global maps
>>> in correspoinding .ipynb files
"""
assert run in ['ctrl', 'rcp', 'lpd', 'lpi', 'had']
assert extent in ['38S', 'Eq', '20N']
if run in ['ctrl', 'rcp']:
domain = 'ocn_rect'
run_name = f'rect_{run}'
elif run in ['lpd', 'lpi']:
domain = 'ocn_low'
run_name = run
elif run == 'had':
domain = 'ocn_had'
run_name = run
# load fields (generated in `SST_data_generatoin.py`)
# <SST>(t): 60S-60N (tropical/extratropical) monthly timeseries
SST_xm = xr.open_dataarray(
f'{path_samoc}/SST/SST_60S_60N_mean_monthly_{run_name}.nc',
decode_times=False)
# SST(t,x,y): monthly SST field limited to Pacific North of `extent` (*)
SST = xr.open_dataarray(f'{path_samoc}/SST/SST_monthly_{run_name}.nc',
decode_times=False)
print('opened datasets')
# deseasonalize
SST_xm_ds = deseasonalize(SST_xm)
SST_ds = deseasonalize(SST)
print('deseasonalized datasets')
# some of the time series are not the same length
if run == 'ctrl': SST_xm_ds = SST_xm_ds[:-7]
elif run == 'rcp': SST_xm_ds = SST_xm_ds[:-1]
# detrend
SST_ds_dt = SST_ds - SST_xm_ds
print('detrended SSTs')
# remove mean at each point
SST_ds_dt_dm = SST_ds_dt - SST_ds_dt.mean('time')
# EOF analysis
# N.B. cut off two year on either end as the arbitrary start month biases the filtered time series
if extent == '38S':
latS, lonE = -38, 300
elif extent == 'Eq':
latS, lonE = 0, 285
elif extent == '20N':
latS, lonE = 20, 255
AREA = xr_AREA(domain=domain)
Pac_MASK = mask_box_in_region(domain=domain,
mask_nr=2,
bounding_lats=(latS, 68),
bounding_lons=(110, lonE))
Pac_area = AREA.where(Pac_MASK)
fn = f'{path_samoc}/SST/SST_PDO_EOF_{extent}_{run}.nc'
print('prepared EOF')
eof, pc = EOF_SST_analysis(xa=SST_ds_dt_dm[24:-24, :, :].where(Pac_MASK),
weights=Pac_area,
fn=fn)
print('performed EOF')
return eof, pc
def OHC_parallel(run, mask_nr=0):
""" ocean heat content calculation """
print('***************************************************')
print('* should be run with a dask scheduler *')
print('`from dask.distributed import Client, LocalCluster`')
print('`cluster = LocalCluster(n_workers=2)`')
print('`client = Client(cluster)`')
print('***************************************************')
print(
f'\n{datetime.datetime.now()} start OHC calculation: run={run} mask_nr={mask_nr}'
)
assert run in ['ctrl', 'rcp', 'lpd', 'lpi']
assert type(mask_nr) == int
assert mask_nr >= 0 and mask_nr < 13
# file_out = f'{path_samoc}/OHC/OHC_test.nc'
file_out = f'{path_samoc}/OHC/OHC_integrals_{regions_dict[mask_nr]}_{run}.nc'
if run in ['ctrl', 'rcp']:
domain = 'ocn'
elif run in ['lpd', 'lpi']:
domain = 'ocn_low'
MASK = boolean_mask(domain, mask_nr)
# geometry
DZT = xr_DZ(domain)
AREA = xr_AREA(domain)
HTN = xr_HTN(domain)
LATS = xr_LATS(domain)
print(f'{datetime.datetime.now()} done with geometry')
# multi-file
file_list = ncfile_list(domain='ocn', run=run, tavg='yrly', name='TEMP_PD')
OHC = xr.open_mfdataset(paths=file_list,
concat_dim='time',
decode_times=False,
compat='minimal',
parallel=True).drop(['ULAT', 'ULONG'
]).TEMP * cp_sw * rho_sw
if mask_nr != 0:
OHC = OHC.where(MASK)
print(f'{datetime.datetime.now()} done loading data')
for ds in [OHC, HTN, LATS]:
round_tlatlon(ds)
OHC_DZT = OHC * DZT
print(f'{datetime.datetime.now()} done OHC_DZT')
# xr DataArrays
da_g = xr_int_global(da=OHC, AREA=AREA, DZ=DZT)
da_gl = xr_int_global_level(da=OHC, AREA=AREA, DZ=DZT)
da_v = OHC_DZT.sum(dim='z_t') #xr_int_vertical(da=OHC, DZ=DZT)
da_va = OHC_DZT.isel(z_t=slice(0, 9)).sum(dim='z_t') # above 100 m
da_vb = OHC_DZT.isel(z_t=slice(9, 42)).sum(dim='z_t') # below 100 m
da_z = xr_int_zonal(da=OHC, HTN=HTN, LATS=LATS, AREA=AREA, DZ=DZT)
da_zl = xr_int_zonal_level(da=OHC, HTN=HTN, LATS=LATS, AREA=AREA, DZ=DZT)
print(f'{datetime.datetime.now()} done calculations')
# xr Datasets
ds_g = da_g.to_dataset(name='OHC_global')
ds_gl = da_gl.to_dataset(name='OHC_global_levels')
ds_v = da_v.to_dataset(name='OHC_vertical')
ds_va = da_va.to_dataset(name='OHC_vertical_above_100m')
ds_vb = da_vb.to_dataset(name='OHC_vertical_below_100m')
ds_z = da_z.to_dataset(name='OHC_zonal')
ds_zl = da_zl.to_dataset(name='OHC_zonal_levels')
print(f'{datetime.datetime.now()} done dataset')
print(f'output: {file_out}')
ds_new = xr.merge([ds_g, ds_gl, ds_z, ds_zl, ds_v, ds_va, ds_vb])
ds_new.to_netcdf(path=file_out, mode='w')
# ds_new.close()
print(f'{datetime.datetime.now()} done\n')
return ds_new
def make_SFWF_surface_integrals():
""" calculate SFWF surface integrals for specific regions
E/P/R/T .. evap, precip, runoff, total
m/t .. mean/trend
! output as a pickled dictionary
"""
ds_ctrl = xr.open_dataset(
f'{path_prace}/ctrl/EVAP_F_PREC_F_ROFF_F_ctrl_mean_200-229.nc')
ds_lpd = xr.open_dataset(
f'{path_prace}/lpd/EVAP_F_PREC_F_ROFF_F_lpd_mean_500-529.nc')
Tm_ctrl = xr.open_dataarray(f'{path_prace}/ctrl/SFWF_ctrl_mean_200-229.nc')
Tm_lpd = xr.open_dataarray(f'{path_prace}/lpd/SFWF_lpd_mean_500-529.nc')
Et_rcp = xr.open_dataarray(
f'{path_prace}/rcp/EVAP_F_yrly_trend_rcp.nc') # 2D data
Pt_rcp = xr.open_dataarray(f'{path_prace}/rcp/PREC_F_yrly_trend_rcp.nc')
Rt_rcp = xr.open_dataarray(f'{path_prace}/rcp/ROFF_F_yrly_trend_rcp.nc')
Tt_rcp = xr.open_dataarray(f'{path_prace}/rcp/SFWF_yrly_trend_rcp.nc')
Et_lr1 = xr.open_dataarray(f'{path_prace}/lr1/EVAP_F_yrly_trend_lr1.nc')
Pt_lr1 = xr.open_dataarray(f'{path_prace}/lr1/PREC_F_yrly_trend_lr1.nc')
Rt_lr1 = xr.open_dataarray(f'{path_prace}/lr1/ROFF_F_yrly_trend_lr1.nc')
Tt_lr1 = xr.open_dataarray(
f'{path_prace}/lr1/SFWF_yrly_trend_lr1.nc') # for now only sum
for i, sim in enumerate(['HIGH', 'LOW']):
print('-----', sim, '-----')
d = {}
(Em, Pm, Rm,
Tm) = [(ds_ctrl.EVAP_F, ds_ctrl.PREC_F, ds_ctrl.ROFF_F, Tm_ctrl),
(ds_lpd.EVAP_F, ds_lpd.PREC_F, ds_lpd.ROFF_F, Tm_lpd)][i]
(Et, Pt, Rt, Tt) = [(Et_rcp, Pt_rcp, Rt_rcp, Tt_rcp),
(Et_lr1, Pt_lr1, Rt_lr1, Tt_lr1)][i]
AREA = xr_AREA(domain=['ocn', 'ocn_low'][i])
for (latS, latN) in lat_bands:
MASK = make_Mask(run=['ctrl', 'lpd'][i], latS=latS, latN=latN)
AREA_total = AREA.where(MASK).sum()
# integrals of mean
Pmi = (Pm.where(MASK) * AREA).sum().values
Emi = (Em.where(MASK) * AREA).sum().values
Rmi = (Rm.where(MASK) * AREA).sum().values
Tmi = (Tm.where(MASK) * AREA).sum().values
d['Pmi_mmd'] = Pmi / AREA_total.values * 24 * 3600 # [kg/s] -> [mm/d]
d['Emi_mmd'] = Emi / AREA_total.values * 24 * 3600
d['Rmi_mmd'] = Rmi / AREA_total.values * 24 * 3600
d['Tmi_mmd'] = Tmi / AREA_total.values * 24 * 3600
d['Pmi_Sv'] = Pmi / 1e9 # [kg/m^2/s] -> [Sv]
d['Emi_Sv'] = Emi / 1e9
d['Rmi_Sv'] = Rmi / 1e9
d['Tmi_Sv'] = Tmi / 1e9
# integrals of trends
Pti = (Pt.where(MASK) * AREA).sum().values
Eti = (Et.where(MASK) * AREA).sum().values
Rti = (Rt.where(MASK) * AREA).sum().values
Tti = (Tt.where(MASK) * AREA).sum().values
d['Pti_mmd'] = Pti / AREA_total.values * 24 * 3600 * 365 * 100 # [mm/d/100yr]
d['Eti_mmd'] = Eti / AREA_total.values * 24 * 3600 * 365 * 100
d['Rti_mmd'] = Rti / AREA_total.values * 24 * 3600 * 365 * 100
d['Tti_mmd'] = Tti / AREA_total.values * 24 * 3600 * 365 * 100
d['Pti_Sv'] = Pti / 1e9 * 365 * 100 # [Sv/100yr]
d['Eti_Sv'] = Eti / 1e9 * 365 * 100
d['Rti_Sv'] = Rti / 1e9 * 365 * 100
d['Tti_Sv'] = Tti / 1e9 * 365 * 100
print(f'\n{latS}N to {latN}N, {AREA_total.values:4.2E} m^2\n')
print(' PREC EVAP ROFF TOTAL DIFF')
print(
f'[mm/d] {d["Pmi_mmd"]:5.2f} {d["Emi_mmd"]:5.2f} {d["Rmi_mmd"]:5.2f} {d["Tmi_mmd"]:7.4f} {d["Tmi_mmd"]-(d["Pmi_mmd"]+d["Emi_mmd"]+d["Rmi_mmd"]):7.4f}'
)
print(
f'[mm/d/100y] {d["Pti_mmd"]:5.2f} {d["Eti_mmd"]:5.2f} {d["Rti_mmd"]:5.2f} {d["Tti_mmd"]:7.4f} {d["Tti_mmd"]-(d["Pti_mmd"]+d["Eti_mmd"]+d["Rti_mmd"]):7.4f}'
)
print(
f'[Sv] {d["Pmi_Sv"]:5.2f} {d["Emi_Sv"]:5.2f} {d["Rmi_Sv"]:5.2f} {d["Tmi_Sv"]:7.4f} {d["Tmi_Sv"]-(d["Pmi_Sv"]+d["Emi_Sv"]+d["Rmi_Sv"]):7.4f}'
)
print(
f'[Sv/100y] {d["Pti_Sv"]:5.2f} {d["Eti_Sv"]:5.2f} {d["Rti_Sv"]:5.2f} {d["Tti_Sv"]:7.4f} {d["Tti_Sv"]-(d["Pti_Sv"]+d["Eti_Sv"]+d["Rti_Sv"]):7.4f}'
)
print(
f'[%/100y] {d["Pti_Sv"]/d["Pmi_Sv"]*100:5.1f} {d["Eti_Sv"]/d["Emi_Sv"]*100:5.1f} {d["Rti_Sv"]/d["Rmi_Sv"]*100:5.1f} {d["Tti_Sv"]/d["Tmi_Sv"]*100:5.1f}\n'
)
print(
f'total surface flux: {d["Tmi_Sv"]:5.2f} Sv {d["Tti_Sv"]:5.2f} Sv/100yr {d["Tti_Sv"]/d["Tmi_Sv"]*100:5.1f} %/100yr'
)
print('\n\n\n')
fn = f'{path_results}/SFWF/Atlantic_SFWF_integrals_{sim}_{latS}N_{latN}N'
save_obj(d, fn)
return
import sys
import datetime
sys.path.append("..")
import xarray as xr
from paths import path_prace
from regions import boolean_mask, Atlantic_mask
from timeseries import IterateOutputCESM
from xr_DataArrays import xr_AREA
for j, run in enumerate(['ctrl','lc1']):
# if j==1: break
domain = ['ocn', 'ocn_low'][j]
TAREA = xr_AREA(domain=domain)
mask_A = Atlantic_mask(domain=domain)
mask_P = boolean_mask(domain=domain, mask_nr=2)
mask_S = boolean_mask(domain=domain, mask_nr=1)
shf, shf_A, shf_P, shf_S = [], [], [], []
for i, (y,m,f) in enumerate(IterateOutputCESM(domain=domain, run=run, tavg='monthly')):
da = xr.open_dataset(f, decode_times=False).SHF*TAREA
shf.append(da.sum())
shf_A.append(da.where(mask_A).sum())
shf_P.append(da.where(mask_P).sum())
shf_S.append(da.where(mask_S).sum())
# if i==24: break
shf = xr.concat(shf, dim='time')
shf_A = xr.concat(shf_A, dim='time')
shf_P = xr.concat(shf_P, dim='time')
shf_S = xr.concat(shf_S, dim='time')
shf.name = 'Global_Ocean'
shf_A.name = 'Atlantic_Ocean'
