def GMST_regression(run):
""" calculates the 2D surface temperature trends
input:
run .. (str) ctrl or rcp
output:
da_trend .. 2D xr DataArray with linear trends
(takes about 4:30 minutes)
"""
field = 'T'
lev = -1
domain = 'atm'
tavg = 'yrly'
name = 'T_T850_U_V'
first_year = IterateOutputCESM(domain=domain, run=run, tavg=tavg, name=name).year
first_file = CESM_filename(domain=domain, run=run, y=first_year, m=0, name=name)
da = xr.open_dataset(first_file, decode_times=False)[field][lev,:,:]
da = da.expand_dims('time')
for y, m, file in IterateOutputCESM(domain=domain, run=run, tavg=tavg, name=name):
if y>first_year:
da_new = xr.open_dataset(file, decode_times=False)[field][lev,:,:]
da = xr.concat([da, da_new], dim='time')
print(y)
da_trend = xr_linear_trends_2D(da, ('lat', 'lon'))
da_trend = da_trend*365*100 # change from [K/day] to [K/century]
return da_trend
def make_MOC_file(self):
# # ca 20 sec per file
# # 50 min for ctrl
# # 26 min for rcp
if selfrun in ['ctrl', 'rcp']:
DXU = xr_DXU(self.domain) # [m]
DZU = xr_DZ(self.domain, grid='U') # [m]
# MASK = Atlantic_mask('ocn') # Atlantic
# MASK = boolean_mask('ocn', 2) # Pacific
MASK = boolean_mask(self.domain, 0) # Global OCean
for i, (y,m,s) in enumerate(IterateOutputCESM(domain=self.domain,
run=self.run,
tavg='yrly',
name='UVEL_VVEL')):
# ca. 20 sec per year
print(i, y, s)
ds = xr.open_dataset(s, decode_times=False)
MOC = calculate_MOC(ds=ds, DXU=DXU, DZU=DZU, MASK=MASK)
if i==0:
MOC_out = MOC.copy()
else:
MOC_out = xr.concat([MOC_out, MOC], dim='time')
# if y==202: break
MOC_out.to_netcdf(f'{path_results}/MOC/GMOC_{self.run}.nc')
elif self.run in ['lpd', 'lpi']:
ds = xr.open_mfdataset()
return
def generate_yrly_SST_files(self, run):
""" generate the SST data files from TEMP_PD yearly averaged files """
# ca. 4:30 min for ctrl/rcp, 1:25 for lpi
# stacking files into one xr DataArray object
if run in ['ctrl', 'rcp', 'lpd', 'lc1', 'lpi', 'lr1']:
for i, (y, m, s) in enumerate(
IterateOutputCESM('ocn', run, 'yrly', name='TEMP_PD')):
print(y)
da = xr.open_dataset(s, decode_times=False).TEMP[0, :, :]
da = da.drop(['z_t', 'ULONG', 'ULAT'])
da_time = int(da.time.item())
if run == 'ctrl':
# years 5-50 have different TLAT/TLON grids
# somehow the non-computed boxes changed (in the continents)
if i == 0:
TLAT = da['TLAT'].round(decimals=2)
TLONG = da['TLONG'].round(decimals=2)
da['TLAT'] = TLAT
da['TLONG'] = TLONG
else:
da['TLAT'] = da['TLAT'].round(decimals=2)
da['TLONG'] = da['TLONG'].round(decimals=2)
del da.encoding["contiguous"]
ds = t2ds(da=da, name='SST', t=da_time)
ds.to_netcdf(path=f'{path_prace}/SST/SST_yrly_{run}_{y:04}.nc',
mode='w')
combined = xr.open_mfdataset(
f'{path_prace}/SST/SST_yrly_{run}_*.nc',
concat_dim='time',
autoclose=True,
coords='minimal')
combined.to_netcdf(f'{path_prace}/SST/SST_yrly_{run}.nc')
self.remove_superfluous_files(
f'{path_prace}/SST/SST_yrly_{run}_*.nc')
if run == 'ctrl': # create also ocn_rect file
fn = f'{path_prace}/SST/SST_monthly_ctrl.nc'
fn_out = f'{path_prace}/SST/SST_yrly_rect_ctrl.nc'
monthly_ctrl = xr.open_dataarray(fn, decode_times=False)
t_bins = np.arange(0, len(monthly_ctrl) + 1, 12) / 12
t_coords = np.array(monthly_ctrl.time[0::12].values, dtype=int)
yrly_ctrl = monthly_ctrl.groupby_bins(
'time', t_bins, right=False).mean(dim='time')
yrly_ctrl = yrly_ctrl.assign_coords(time_bins=t_coords).rename(
{'time_bins': 'time'})
yrly_ctrl.to_netcdf(fn_out)
elif run == 'had':
ds2 = xr.open_dataset(file_HadISST)
ds2 = ds2.where(ds2['sst'] != -1000.)
ds2 = ds2.sst.where(np.isnan(ds2.sst) == False, -1.8)
ds2 = ds2.groupby('time.year').mean('time')
ds2 = ds2.rename({'year': 'time'}).isel(time=slice(0, -1))
ds2.coords['time'] = (ds2.coords['time'] - 1870) * 365
ds2.to_netcdf(f'{path_prace}/SST/SST_yrly_had.nc')
return
def generate_monthly_SST_files(self, run, time=None):
""" concatonate monthly files, ocn_rect for high res runs"""
# 8 mins for 200 years of ctrl
if run in ['ctrl', 'rcp', 'hq']: domain = 'ocn_rect'
elif run in ['lpd', 'lpi', 'lc1', 'lr1', 'ld']: domain = 'ocn_low'
for y, m, s in IterateOutputCESM(domain=domain,
tavg='monthly',
run=run):
if time == None:
pass
elif y < time[0] or y >= time[1]:
continue
if m == 1: print(y)
if run in ['ctrl', 'rcp']:
xa = xr.open_dataset(s, decode_times=False).TEMP[0, :, :]
if run in ['lpd', 'lpi', 'lc1', 'lr1', 'ld']:
xa = xr.open_dataset(s, decode_times=False).TEMP[0, 0, :, :]
if m == 1: xa_out = xa.copy()
else: xa_out = xr.concat([xa_out, xa], dim='time')
drops = ['ULAT', 'ULONG', 'TLAT', 'TLONG']
if m == 12:
if type(xa_out) == xr.core.dataarray.Dataset:
if 'z_t' in xa_out.variables: drops.append('z_t')
elif 'z_t' in xa_out.coords: drops.append('z_t')
xa_out.drop(drops).to_netcdf(
f'{path_prace}/SST/SST_monthly_{run}_y{y:04}.nc')
# this also means only full years are written out
combined = xr.open_mfdataset(
f'{path_prace}/SST/SST_monthly_{run}_y*.nc',
combine='nested',
concat_dim='time',
decode_times=False)
if run == 'ctrl':
time_ctrl = np.arange(1 + 1 / 24, 301, 1 / 12)
combined = combined.assign_coords(time=time_ctrl)
# something is wrong in Febraury 99, so I average
n = np.where(np.isclose(time_ctrl, 99.125))[0][0]
print(f'month 99/2 is averaged; n={n}')
combined[n] = (combined[n - 1] + combined[n + 1]) / 2
if time == None:
fn_out = f'{path_prace}/SST/SST_monthly_{run}.nc'
else:
fn_out = f'{path_prace}/SST/SST_monthly_{run}_{time[0]}_{time[1]}.nc'
combined.to_netcdf(fn_out)
combined.close()
self.remove_superfluous_files(
f'{path_prace}/SST/SST_monthly_{run}_y*.nc')
return
def make_SFWF_trends(run):
""" calculate linear trends of SFWF """
for qs in [['SFWF'], ['EVAP_F', 'PREC_F', 'ROFF_F']]:
name = '_'.join(qs)
print(f' making trend of {name}')
# concatenating yearly fields
fn = f'{path_prace}/{run}/{qs[0]}_yrly_trend_{run}.nc'
if os.path.exists(fn):
pass
else:
for i, (y, m, f) in enumerate(
IterateOutputCESM(domain='ocn',
run=run,
tavg='yrly',
name=name)):
ds_ = xr.open_dataset(f, decode_times=False)
if i == 0: ds = []
ds.append(ds_)
ds = xr.concat(ds, dim='time')
print('mfdataset created')
for q in qs:
ds[q].to_netcdf(f'{path_prace}/{run}/{q}_yrly_{run}.nc')
# regression
for i, q in enumerate(qs):
if i < 1: continue
fn = f'{path_prace}/{run}/{q}_yrly_trend_{run}.nc'
if os.path.exists(fn):
pass
else:
print('before regression')
ds = xr.open_dataarray(f'{path_prace}/{run}/{q}_yrly_{run}.nc',
decode_times=False)
trend = ocn_field_regression(xa=ds, run=run)
print('after regression')
trend[0].to_netcdf(fn)
fn_stats = f'{path_prace}/{run}/{q}_yrly_trend_stats_{run}.nc'
if os.path.exists(fn_stats):
pass
else:
ds = xr.open_dataarray(f'{path_prace}/{run}/{q}_yrly_{run}.nc',
decode_times=False)
print('creating {q} scipy stats linregress')
xr_regression_with_stats(ds, fn=fn_stats)
return
def make_SST_yrly_data_file(self):
""" generates annual SST .nc file (t,lat,lon)
ca. 4:30 min for ctrl/rcp, 1:25 for lpi
"""
for i, (y,m,s) in enumerate(IterateOutputCESM('ocn', self.run, 'yrly', name='TEMP_PD')):
print(y)
da = xr.open_dataset(s, decode_times=False).TEMP[0,:,:]
da = da.drop(['z_t', 'ULONG', 'ULAT'])
da['TLAT' ] = da['TLAT' ].round(decimals=2)
da['TLONG'] = da['TLONG'].round(decimals=2)
del da.encoding["contiguous"]
ds = t2ds(da=da, name='SST', t=int(round(da.time.item())))
ds.to_netcdf(path=f'{path_prace}/SST/SST_yrly_{self.run}_{y}.nc', mode='w')
combined = xr.open_mfdataset(f'{path_prace}/SST/SST_yrly_{self.run}_*.nc',
concat_dim='time',
autoclose=True,
coords='minimal')
combined.to_netcdf(f'{path_prace}/SST/SST_yrly_{self.run}.nc')
# remove extra netCDF files
return
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 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 yrly_avg_nc(self, domain, fields, test=False, years=None):
""" creates yearly average file from monthly
input: fields .. list of field names
output: [writes netCDF file]
(takes approx. 2 min for high res ocean data for two 3D fields)
(takes approx. 4 sec for lower res atm data for one 2D and one 3D field)
"""
assert domain in ['ocn', 'ocn_rect', 'atm', 'ice']
assert self.run in ['ctrl', 'rcp', 'hq', 'lpd', 'lc1', 'lpi', 'lr1', 'lr2', 'ld', 'lq']
print(f'yearly averaging of {self.run} {domain}')
for field in fields: print(f' {field}')
name = ''
n_fields = len(fields)
for i, field in enumerate(fields):
name += field
if i<n_fields-1:
name += '_'
ffield = fields[0]
cond = False
if years is None:
cond = True
elif np.all(years<501):
cond = True
if self.run=='lpd' and domain=='atm' and cond:
# there are no monthly files for the lpd run
# but for consistency, I create the same files as for the runs with monthly output only
for y, m, s in IterateOutputCESM(domain=domain, run=self.run, tavg='yrly'):
new_filename = CESM_filename(domain=domain, run=self.run, y=y, m=0, name=name)
ds = xr.open_dataset(s, decode_times=False)
dim = len(np.shape(ds[ffield]))
if dim==3: ds_out = (ds[ffield][0,:,:]).to_dataset()
elif dim==4: ds_out = (ds[ffield][0,:,:,:]).to_dataset()
if len(fields)>1:
for field in fields[1:]:
dim = len(np.shape(ds[field]))
if dim==3: ds_out[field] = ds[field][0,:,:]
elif dim==4: ds_out[field] = ds[field][0,:,:,:]
print(y, new_filename)
ds_out.to_netcdf(path=new_filename, mode='w')
return
else:
first_year = IterateOutputCESM(domain=domain, run=self.run, tavg='monthly').year
for y, m, s in IterateOutputCESM(domain=domain, run=self.run, tavg='monthly'):
if years is not None and y not in years: continue
if m==1:
new_filename = CESM_filename(domain=domain, run=self.run, y=y, m=0, name=name)
if os.path.exists(new_filename):
continue
ds = xr.open_dataset(s, decode_times=False)
if m==1: # create new xr Dataset
dim = len(np.shape(ds[ffield]))
if domain in ['atm', 'ocn', 'ice']:
if dim==3: # 2D field
ds_out = (ds[ffield][0,:,:]/12).to_dataset()
elif dim==4: # 3D
ds_out = (ds[ffield][0,:,:,:]/12).to_dataset()
elif domain=='ocn_rect':
if dim==2: # 2D
ds_out = (ds[ffield][:,:]/12).to_dataset()
elif dim==3: # 3D
ds_out = (ds[ffield][:,:,:]/12).to_dataset()
for field in fields[1:]: # add rest of fields
dim = len(np.shape(ds[field]))
if domain in ['atm', 'ocn', 'ice']:
if dim==3:
ds_out[field] = ds[field][0,:,:]/12
elif dim==4:
ds_out[field] = ds[field][0,:,:,:]/12
elif domain=='ocn_rect':
if dim==2:
ds_out[field] = ds[field][:,:]/12
elif dim==3:
ds_out[field] = ds[field][:,:,:]/12
else: # add subsequent monthly values
for field in fields:
dim = len(np.shape(ds[field]))
if domain in ['atm', 'ocn', 'ice']:
if dim==3: ds_out[field][:,:] += ds[field][0,:,:]/12
elif dim==4: ds_out[field][:,:,:] += ds[field][0,:,:,:]/12
elif domain=='ocn_rect':
if dim==2: ds_out[field][:,:] += ds[field][:,:]/12
elif dim==3: ds_out[field][:,:,:] += ds[field][:,:,:]/12
if m==12: # write to new file
if 'TLONG' in ds_out: ds_out.drop('TLONG')
if 'TLAT' in ds_out: ds_out.drop('TLAT')
if 'ULONG' in ds_out: ds_out.drop('ULONG')
if 'ULAT' in ds_out: ds_out.drop('ULAT')
print(y, new_filename)
ds_out.to_netcdf(path=new_filename, mode='w')
if test==True and y==first_year+2: break
print('done!')
return
import datetime
from grid import generate_lats_lons
from CICE import CICE_XMXL_plots
from paths import CESM_filename, file_ex_ocn_rcp
from timeseries import IterateOutputCESM
if __name__ == "__main__":
print(f'running CICE_XMXL plots run={run}')
print(f'{datetime.datetime.now()}\n\n')
lats, lons = generate_lats_lons('ocn')
MASK = xr.open_dataset(file_ex_ocn_rcp, decode_times=False).REGION_MASK
for i, (y, m, ocn_file) in enumerate(
IterateOutputCESM(domain='ocn', run=run, tavg='monthly')):
if m == 1: print(y)
ice_file = CESM_filename(domain='ice', run=run, y=y, m=m)
aice = xr.open_dataset(ice_file, decode_times=False).aice[0, :, :]
XMXL = xr.open_dataset(ocn_file, decode_times=False).XMXL[0, :, :]
CICE_XMXL_plots(aice=aice,
XMXL=XMXL,
lons=lons,
lats=lats,
MASK=MASK,
run=run,
i=i)
aice.close()
XMXL.close()
print(f'\n\nfinished at\n{datetime.datetime.now()}')
def main(run, ys=None, ye=None):
"""
6:30 for lpd transports only
21 min for 30 years of monthly lpd
input: {run}
run .. file name
output:
ds .. dataset containing FW/SALT transport terms
saved to f'{path_prace}/Mov/FW_SALT_fluxes_{run}_{y}.nc'
"""
print(datetime.now())
if run in ['ctrl', 'rcp']:
if run == 'ctrl':
s1, s2 = path_ocn_ctrl, spinup
# if ys is None and ye is None: ys, ye = 200, 230 #230
elif run == 'rcp':
s1, s2 = path_ocn_rcp, rcpstr
# if ys is None and ye is None: ys, ye = 2070, 2100 #2000, 2101
domain = 'ocn'
elif run in ['lpd', 'lr1']:
if run == 'lpd':
s1, s2 = path_ocn_lpd, lpdstr
if ys is None and ye is None: ys, ye = 500, 530
elif run == 'lr1':
s1, s2 = path_ocn_lr1, lr1str
if ys is None and ye is None: ys, ye = 2000, 2101
domain = 'ocn_low'
else:
raise ValueError(f'`run`={run} not implemented')
ds_geo = geometry_file(domain)
mfkw = dict(decode_times=False, combine='nested', concat_dim='time')
q = ['SALT', 'VNS', 'UES', 'VVEL', 'UVEL']
if domain == 'ocn_low': # can concatenate all files
fns = f'{path_prace}/Mov/FW_Med_Bering_{run}_{ys:04d}-{ye:04d}_monthly.nc'
fnt = f'{path_prace}/Mov/FW_SALT_fluxes_{run}_{ys:04d}-{ye:04d}_monthly_alt.nc'
if os.path.exists(fns) and os.path.exists(fnt):
print(f'files exists:\n{fns}\n{fnt}')
pass
else:
ds_mf = xr.open_mfdataset(
mf_list(f'{s1}/{s2}.pop.h.', y1=ys, y2=ye), **mfkw)[q]
ds = xr.merge([ds_mf, ds_geo])
# section transports
if os.path.exists(fns):
print(f'file exists:\n {fns}')
pass
else:
print(f'create file:\n {fns}')
calc_sections(ds, fns)
# transport terms
if os.path.exists(fnt):
print(f'file exists:\n {fnt}')
pass
else:
print(f'create file:\n {fnt}')
calc_transports(ds, domain, fnt)
elif domain == 'ocn':
for (y, m, f) in IterateOutputCESM(run=run,
domain=domain,
tavg='monthly'):
if m == 1: print(y)
if y < ys or y > ye: continue
fns = f'{path_prace}/Mov/FW_Med_Bering_monthly_{run}_{y:04d}-{m:02d}.nc'
fnt = f'{path_prace}/Mov/FW_SALT_fluxes_monthly_{run}_{y:04d}-{m:02d}_alt.nc'
ds_ = xr.open_dataset(f, decode_times=False)[q]
ds = xr.merge([ds_, ds_geo])
del (ds_)
if os.path.exists(fns) and os.path.exists(fnt):
print(f'files exists:\n{fns}\n{fnt}')
pass
else:
if os.path.exists(fns):
print(f'file exists:\n {fns}')
pass
else:
print(f'create file:\n {fns}')
calc_sections(ds, fns)
# transport terms
if os.path.exists(fnt):
print(f'file exists:\n {fnt}')
pass
else:
print(f'create file:\n {fnt}')
calc_transports(ds, domain, fnt)
gc.collect()
# loop only through pentades because of memory limits
# dy = 1
# y_list = np.arange(ys, ye, dy)
# for y in tqdm.tqdm(y_list):
# fns = f'{path_prace}/Mov/FW_Med_Bering_{run}_{y:04d}-{y+dy:04d}_monthly.nc'
# fnt = f'{path_prace}/Mov/FW_SALT_fluxes_{run}_{y:04d}-{y+dy:04d}_monthly.nc'
# if os.path.exists(fns) and os.path.exists(fnt):
# print(f'files exists:\n{fns}\n{fnt}')
# pass
# else:
# ds_mf = xr.open_mfdataset(mf_list(f'{s1}/{s2}.pop.h.', y1=ys, y2=ye), **mfkw)[q]
# ds = xr.merge([ds_mf, ds_geo])
# # section transports
# if os.path.exists(fns): print(f'file exists:\n {fns}'); pass
# else: print(f'create file:\n {fns}'); calc_sections(ds, fns)
# # transport terms
# if os.path.exists(fnt): print(f'file exists:\n {fnt}'); pass
# else: print(f'create file:\n {fnt}'); calc_transports(ds, domain, fnt)
print(datetime.now())
return
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'
shf_P.name = 'Pacific_Ocean'
shf_S.name = 'Southern_Ocean'
xr.merge([shf, shf_A, shf_P, shf_S]).to_netcdf(f'{path_prace}/SHF/SHF_monthly_{run}.nc')
