diff --git a/dartwrf/update_IC.py b/dartwrf/update_IC.py index 8158cc902ecc47d50736daa69bd21c6bf2c994ba..38760eb66443c039206499d068625dd1ab6c29db 100755 --- a/dartwrf/update_IC.py +++ b/dartwrf/update_IC.py @@ -9,24 +9,22 @@ Updates initial condition (wrfinput/wrfrst files) in the run_WRF directories wit # assumes T = THM (dry potential temperature as prognostic variable) """ -use_wrfrst = True # recommended to be True -if use_wrfrst: - initials_fmt = '/wrfrst_d01_%Y-%m-%d_%H:%M:%S' -else: - initials_fmt = '/wrfinput_d01' - - def update_initials_in_WRF_rundir(time): - """updates wrfrst in run_WRF directory with posterior state from ./filter + """Updates wrfrst-files in `/run_WRF/` directory + with posterior state from ./filter output, e.g. filter_restart_d01.0001 Args: time (dt.datetime): time of assimilation (directory preceeding ./assim_stage0/...) """ + use_wrfrst = True # if wrfrst is used to restart (recommended) + if use_wrfrst: + initials_fmt = '/wrfrst_d01_%Y-%m-%d_%H:%M:%S' + else: + initials_fmt = '/wrfinput_d01' # which WRF variables will be updated? update_vars = ['Times',] update_vars.extend(exp.update_vars) - updates = ','.join(update_vars) for iens in range(1, exp.n_ens+1): ic_file = cluster.wrf_rundir(iens) + time.strftime(initials_fmt) @@ -35,22 +33,26 @@ def update_initials_in_WRF_rundir(time): else: # overwrite DA updated variables filter_out = cluster.archivedir+time.strftime('/%Y-%m-%d_%H:%M/assim_stage0/filter_restart_d01.'+str(iens).zfill(4)) - print('update assimilated variables => overwrite', updates, 'from', filter_out) - os.system(cluster.ncks+' -A -v '+updates+' '+filter_out+' '+ic_file) - - # assumes T = THM (dry potential temperature as prognostic variable) - print('writing T into THM') - with nc.Dataset(filter_out, 'r') as ds_filter: - if use_wrfrst: - with nc.Dataset(ic_file, 'r+') as ds_wrfrst: - ds_wrfrst.variables['THM_1'][:] = ds_filter.variables['T'][:] - ds_wrfrst.variables['THM_2'][:] = ds_filter.variables['T'][:] - else: - with nc.Dataset(ic_file, 'r+') as ds_wrfout: - ds_wrfout.variables['THM'][:] = ds_filter.variables['T'][:] - - print(ic_file, 'updated.') + with nc.Dataset(filter_out, 'r') as ds_filter: + with nc.Dataset(ic_file, 'r+') as ds_new: + + # assumes T = THM (dry potential temperature as prognostic variable) + if use_wrfrst: + ds_new.variables['THM_2'][:] = ds_filter.variables['T'][:] + else: + ds_new.variables['THM'][:] = ds_filter.variables['T'][:] + + # update all other variables + for var in update_vars: + if var in ds_new.variables: + var_new = var + else: + var_new = var+'_2' # e.g. U_2, W_2, THM_2 + + ds_new.variables[var_new][:] = ds_filter.variables[var][:] + + print(ic_file, 'created, updated from', filter_out) if __name__ == '__main__':