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