From 4bec2ce8d0a0788774b1d5d3bd8f935060a9d47c Mon Sep 17 00:00:00 2001 From: Sam Hartke Date: Mon, 27 Feb 2023 16:27:13 -0700 Subject: [PATCH] Test --- helpers/writeERA5file.py | 119 +++++++++++++++++++++++++++++++-------- 1 file changed, 95 insertions(+), 24 deletions(-) diff --git a/helpers/writeERA5file.py b/helpers/writeERA5file.py index 3fb7b6d..a615422 100644 --- a/helpers/writeERA5file.py +++ b/helpers/writeERA5file.py @@ -1,29 +1,43 @@ +from dask_jobqueue import PBSCluster +from dask.distributed import Client +import time + +cluster = PBSCluster( + queue="casper", + walltime="03:00:00", + project="P48500028", + memory="30GB", + cores=1, + processes=1, +) + +cluster.scale(12) + +client=Client(cluster) +time.sleep(30) # wait 30 seconds to give all dask workers time to populate +print('Cluster created and assigned to dask client') + import os import numpy as np import pandas as pd import xarray as xr -#import qfgrib from datetime import timedelta,datetime from dateutil.relativedelta import relativedelta -era5pl_template = "/glade/collections/rda/data/ds633.0/e5.oper.an.pl/{yr}{mth}/e5.oper.an.pl.{var_id}.{dt}00_{dt}23.grb" -era5sfc_template = "/glade/collections/rda/data/ds633.0/e5.oper.an.sfc/{yr}{mth}/e5.oper.an.sfc.{var_id}.{dt}00_{dt2}23.grb" -w_format = "128_135_w.ll025sc" -v_format = "128_132_v.ll025uv" -u_format = "128_131_u.ll025uv" -q_format = "128_133_q.ll025sc" -t_format = "128_130_t.ll025sc" -tcrw_format = "228_089_tcrw.ll025sc" +#-------------------------------------------------------------------------- +# Function for writing ERA5 data to yearly input files for GARD +era5pl_template = "/glade/collections/rda/data/ds633.0/e5.oper.an.pl/{yr}{mth}/e5.oper.an.pl.{var_id}.{dt}00_{dt}23.grb" +era5sfc_template = "/glade/collections/rda/data/ds633.0/e5.oper.an.sfc/{yr}{mth}/e5.oper.an.sfc.{var_id}.{dt}00_{dt2}23.grb" -var_formats = {'v':v_format, - 'u':u_format, - 'w':w_format, - 'q':q_format, - 't':t_format, - 'tcrw':tcrw_format, +var_formats = {'v':"128_132_v.ll025uv", + 'u':"128_131_u.ll025uv", + 'w':"128_135_w.ll025sc", + 'q':"128_133_q.ll025sc", + 't':"128_130_t.ll025sc", + 'tcrw':"228_089_tcrw.ll025sc", } # new variable names to match those from CESM LENS 2 @@ -35,7 +49,7 @@ 'tcrw':'PRECT', } -def createInputDataset(yr_st,yr_end,varlist): +def createERA5Dataset(yr_st,yr_end,varlist): m=0 for var in varlist: @@ -45,7 +59,8 @@ def createInputDataset(yr_st,yr_end,varlist): yr = dt.year mth = "%.02d"%dt.month if var in ('u','v','w','q','t'): - var_files.append(era5pl_template.format(yr=yr,mth=mth,var_id=var_formats[var],dt=dt.strftime('%Y%m%d'))) + var_files.append(era5pl_template.format(yr=yr,mth=mth,var_id=var_formats[var], + dt=dt.strftime('%Y%m%d'))) dt = dt + timedelta(days=1) elif var == 'tcrw': @@ -56,7 +71,7 @@ def createInputDataset(yr_st,yr_end,varlist): dt = dt + relativedelta(months=+1) vardata = xr.open_mfdataset(var_files,concat_dim='time',combine='nested', - backend_kwargs={"indexpath":""}).sel(latitude=slice(50,20), + backend_kwargs={"indexpath":""},parallel=True).sel(latitude=slice(50,20), longitude=slice(360-120,360-60))[var] if var in ('u','v','w','q','t'): @@ -69,13 +84,15 @@ def createInputDataset(yr_st,yr_end,varlist): if m==0: era5_ds = vardata.to_dataset() + era5_ds = era5_ds.rename({var:newvarname[var]}) else: era5_ds = era5_ds.assign(var=vardata) - - era5_ds = era5_ds.rename({'var':newvarname[var]}) + era5_ds = era5_ds.rename({'var':newvarname[var]}) m+=1 + era5_ds = era5_ds.drop_vars(('number','step','surface')) + if yr_st!=yr_end: outfile = '/glade/scratch/shartke/gard/era5/era5_daily_%d_%d.nc'%(yr_st,yr_end) else: @@ -83,9 +100,63 @@ def createInputDataset(yr_st,yr_end,varlist): era5_ds.to_netcdf(outfile) + +#-------------------------------------------------------------------------- +# Function for writing CESM LENS2 data to decadal input files for GARD + + +cesm_template = "/glade/campaign/cgd/cesm/CESM2-LE/atm/proc/tseries/day_1/{var}/b.e21.B{scen}{forcing}.f09_g17.LE2-{styr}.0{ens}.cam.h{i}.{var}.{yr1}0101-{yr2}1231.nc" +scen="HIST" +f="cmip6" + +def createCESM2Dataset(yr,styr,varlist,enslist): + + yr1 = yr-yr%10 + + for e in enslist: + m=0 + for var in varlist: + if var in ('U','V','T','Q'): + ds = xr.open_dataset(cesm_template.format(var=var,scen=scen,forcing=f,styr=styr, + ens="%.02d"%e,yr1=yr1,yr2=yr1+9,i=6))[var] + # select data over CONUS at ~450 mb level + vardata = ds.sel(lev=ds.lev[19],lat=slice(20,50),lon=slice(360.-120.,360.-60.)) # ,time=slice(str(yr),str(yr)) + vardata = vardata.drop_vars('lev') + elif var in ('PSL','PRECT'): + ds = xr.open_dataset(cesm_template.format(var=var,scen=scen,forcing=f,styr=styr, + ens="%.02d"%e,yr1=yr1,yr2=yr1+9,i=1))[var] + vardata = ds.sel(lat=slice(20,50),lon=slice(360.-120.,360.-60.)) # ,time=slice(str(yr),str(yr)) + # convert m/s to mm/d + vardata = vardata*3600*24*1000 + + if m==0: + cesm_ds = vardata.to_dataset() + else: + cesm_ds = cesm_ds.assign(var=vardata) + cesm_ds = cesm_ds.rename({'var':var}) + + m+=1 + + outfile = '/glade/scratch/shartke/gard/cesmlens2/cesm_daily_%d_%d_%d_%.02d.nc'%(yr,yr+9,styr,e) + cesm_ds.to_netcdf(outfile) + + +#-------------------------------------------------------------------------- +# Note: Generating the ERA5 datasets will take the bulk of the time for this program + print(datetime.now()) -createInputDataset(1982,1982,['u','v','w','q','t','tcrw']) -print('1982 dataset complete at: ',datetime.now()) -createInputDataset(1983,1983,['u','v','w','q','t','tcrw']) -print('1983 dataset complete at: ',datetime.now()) \ No newline at end of file +styr = 1301 # 1231, 1251, 1281, or 1301 +createCESM2Dataset(1960,styr,['U','V','W','Q','T','PRECT'],np.arange(1,3)) +createCESM2Dataset(1970,styr,['U','V','W','Q','T','PRECT'],np.arange(1,3)) +print('CESM LENS2 datasets complete at: ',datetime.now()) + +for yr in (1980,1981,1982,1983): + createERA5Dataset(yr,yr,['u','v','w','q','t','tcrw']) + print('ERA5 %s dataset complete at: '%yr,datetime.now()) + + + + +# now you should be able to train GARD using 1980-1999 ERA5 data +# and predict downscaled 1960-1979 precip or temp using CESM LENS2 data