Skip to content

Commit

Permalink
Remove IC staging from setup
Browse files Browse the repository at this point in the history 
Remove the functions and calls to set up
symlinks to ICs in ROTDIR.

Refs NOAA-EMC#2475
  • Loading branch information
@KateFriedman-NOAA
KateFriedman-NOAA committed May 14, 2024
1 parent 8d2203d commit ac9fe80
Showing 1 changed file with 1 addition and 227 deletions.
228 changes: 1 addition & 227 deletions workflow/setup_expt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,231 +29,6 @@ def makedirs_if_missing(dirname):
os.makedirs(dirname)


def fill_ROTDIR(host, inputs):
"""
Method to populate the ROTDIR for supported modes.
INPUTS:
host: host object from class Host
inputs: user inputs to setup_expt.py
"""

fill_modes = {
'cycled': fill_ROTDIR_cycled,
'forecast-only': fill_ROTDIR_forecasts
}

try:
fill_modes[inputs.mode](host, inputs)
except KeyError:
raise NotImplementedError(f'{inputs.mode} is not a supported mode.\n' +
'Currently supported modes are:\n' +
f'{" | ".join(fill_modes.keys())}')

return


def fill_ROTDIR_cycled(host, inputs):
"""
Implementation of 'fill_ROTDIR' for cycled mode
"""

rotdir = os.path.join(inputs.comroot, inputs.pslot)

do_ocean = do_ice = do_med = False

if 'S2S' in inputs.app:
do_ocean = do_ice = do_med = True

if inputs.icsdir is None:
warnings.warn("User did not provide '--icsdir' to stage initial conditions")
return

rdatestr = datetime_to_YMDH(inputs.idate - to_timedelta('T06H'))
idatestr = datetime_to_YMDH(inputs.idate)

# Test if we are using the new COM structure or the old flat one for ICs
if inputs.start in ['warm']:
pathstr = os.path.join(inputs.icsdir, f'{inputs.cdump}.{rdatestr[:8]}',
rdatestr[8:], 'model_data', 'atmos')
else:
pathstr = os.path.join(inputs.icsdir, f'{inputs.cdump}.{idatestr[:8]}',
idatestr[8:], 'model_data', 'atmos')

if os.path.isdir(pathstr):
flat_structure = False
else:
flat_structure = True

# Destination always uses the new COM structure
# These should match the templates defined in config.com
if inputs.start in ['warm']:
dst_atm_dir = os.path.join('model_data', 'atmos', 'restart')
dst_med_dir = os.path.join('model_data', 'med', 'restart')
else:
dst_atm_dir = os.path.join('model_data', 'atmos', 'input')
dst_med_dir = '' # no mediator files for a "cold start"
do_med = False
dst_ocn_rst_dir = os.path.join('model_data', 'ocean', 'restart')
dst_ocn_anl_dir = os.path.join('analysis', 'ocean')
dst_ice_rst_dir = os.path.join('model_data', 'ice', 'restart')
dst_atm_anl_dir = os.path.join('analysis', 'atmos')

if flat_structure:
# ICs are in the old flat COM structure
if inputs.start in ['warm']: # This is warm start experiment
src_atm_dir = os.path.join('atmos', 'RESTART')
src_med_dir = os.path.join('med', 'RESTART')
elif inputs.start in ['cold']: # This is a cold start experiment
src_atm_dir = os.path.join('atmos', 'INPUT')
src_med_dir = '' # no mediator files for a "cold start"
do_med = False
# ocean and ice have the same filenames for warm and cold
src_ocn_rst_dir = os.path.join('ocean', 'RESTART')
src_ocn_anl_dir = 'ocean'
src_ice_rst_dir = os.path.join('ice', 'RESTART')
src_atm_anl_dir = 'atmos'
else:
src_atm_dir = dst_atm_dir
src_med_dir = dst_med_dir
src_ocn_rst_dir = dst_ocn_rst_dir
src_ocn_anl_dir = dst_ocn_anl_dir
src_ice_rst_dir = dst_ice_rst_dir
src_atm_anl_dir = dst_atm_anl_dir

def link_files_from_src_to_dst(src_dir, dst_dir):
files = os.listdir(src_dir)
for fname in files:
os.symlink(os.path.join(src_dir, fname),
os.path.join(dst_dir, fname))
return

# Link ensemble member initial conditions
if inputs.nens > 0:
previous_cycle_dir = f'enkf{inputs.cdump}.{rdatestr[:8]}/{rdatestr[8:]}'
current_cycle_dir = f'enkf{inputs.cdump}.{idatestr[:8]}/{idatestr[8:]}'

for ii in range(1, inputs.nens + 1):
memdir = f'mem{ii:03d}'
# Link atmospheric files
if inputs.start in ['warm']:
dst_dir = os.path.join(rotdir, previous_cycle_dir, memdir, dst_atm_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, memdir, src_atm_dir)
elif inputs.start in ['cold']:
dst_dir = os.path.join(rotdir, current_cycle_dir, memdir, dst_atm_dir)
src_dir = os.path.join(inputs.icsdir, current_cycle_dir, memdir, src_atm_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link ocean files
if do_ocean:
dst_dir = os.path.join(rotdir, previous_cycle_dir, memdir, dst_ocn_rst_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, memdir, src_ocn_rst_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# First 1/2 cycle needs a MOM6 increment
incfile = f'enkf{inputs.cdump}.t{idatestr[8:]}z.ocninc.nc'
src_file = os.path.join(inputs.icsdir, current_cycle_dir, memdir, src_ocn_anl_dir, incfile)
dst_file = os.path.join(rotdir, current_cycle_dir, memdir, dst_ocn_anl_dir, incfile)
makedirs_if_missing(os.path.join(rotdir, current_cycle_dir, memdir, dst_ocn_anl_dir))
os.symlink(src_file, dst_file)

# Link ice files
if do_ice:
dst_dir = os.path.join(rotdir, previous_cycle_dir, memdir, dst_ice_rst_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, memdir, src_ice_rst_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link mediator files
if do_med:
dst_dir = os.path.join(rotdir, previous_cycle_dir, memdir, dst_med_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, memdir, src_med_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link deterministic initial conditions
previous_cycle_dir = f'{inputs.cdump}.{rdatestr[:8]}/{rdatestr[8:]}'
current_cycle_dir = f'{inputs.cdump}.{idatestr[:8]}/{idatestr[8:]}'

# Link atmospheric files
if inputs.start in ['warm']:
dst_dir = os.path.join(rotdir, previous_cycle_dir, dst_atm_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, src_atm_dir)
elif inputs.start in ['cold']:
dst_dir = os.path.join(rotdir, current_cycle_dir, dst_atm_dir)
src_dir = os.path.join(inputs.icsdir, current_cycle_dir, src_atm_dir)

makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link ocean files
if do_ocean:
dst_dir = os.path.join(rotdir, previous_cycle_dir, dst_ocn_rst_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, src_ocn_rst_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# First 1/2 cycle needs a MOM6 increment
incfile = f'{inputs.cdump}.t{idatestr[8:]}z.ocninc.nc'
src_file = os.path.join(inputs.icsdir, current_cycle_dir, src_ocn_anl_dir, incfile)
dst_file = os.path.join(rotdir, current_cycle_dir, dst_ocn_anl_dir, incfile)
makedirs_if_missing(os.path.join(rotdir, current_cycle_dir, dst_ocn_anl_dir))
os.symlink(src_file, dst_file)

# Link ice files
if do_ice:
dst_dir = os.path.join(rotdir, previous_cycle_dir, dst_ice_rst_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, src_ice_rst_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link mediator files
if do_med:
dst_dir = os.path.join(rotdir, previous_cycle_dir, dst_med_dir)
src_dir = os.path.join(inputs.icsdir, previous_cycle_dir, src_med_dir)
makedirs_if_missing(dst_dir)
link_files_from_src_to_dst(src_dir, dst_dir)

# Link bias correction and radiance diagnostics files
src_dir = os.path.join(inputs.icsdir, current_cycle_dir, src_atm_anl_dir)
dst_dir = os.path.join(rotdir, current_cycle_dir, dst_atm_anl_dir)
makedirs_if_missing(dst_dir)
for ftype in ['abias', 'abias_pc', 'abias_air', 'radstat']:
fname = f'{inputs.cdump}.t{idatestr[8:]}z.{ftype}'
src_file = os.path.join(src_dir, fname)
if os.path.exists(src_file):
os.symlink(src_file, os.path.join(dst_dir, fname))
# First 1/2 cycle also needs a atmos increment if doing warm start
if inputs.start in ['warm']:
for ftype in ['atmi003.nc', 'atminc.nc', 'atmi009.nc']:
fname = f'{inputs.cdump}.t{idatestr[8:]}z.{ftype}'
src_file = os.path.join(src_dir, fname)
if os.path.exists(src_file):
os.symlink(src_file, os.path.join(dst_dir, fname))
if inputs.nens > 0:
current_cycle_dir = f'enkf{inputs.cdump}.{idatestr[:8]}/{idatestr[8:]}'
for ii in range(1, inputs.nens + 1):
memdir = f'mem{ii:03d}'
src_dir = os.path.join(inputs.icsdir, current_cycle_dir, memdir, src_atm_anl_dir)
dst_dir = os.path.join(rotdir, current_cycle_dir, memdir, dst_atm_anl_dir)
makedirs_if_missing(dst_dir)
for ftype in ['ratmi003.nc', 'ratminc.nc', 'ratmi009.nc']:
fname = f'enkf{inputs.cdump}.t{idatestr[8:]}z.{ftype}'
src_file = os.path.join(src_dir, fname)
if os.path.exists(src_file):
os.symlink(src_file, os.path.join(dst_dir, fname))

return


def fill_ROTDIR_forecasts(host, inputs):
"""
Implementation of 'fill_ROTDIR' for forecast-only mode
"""
print('forecast-only mode treats ICs differently and cannot be staged here')


def fill_EXPDIR(inputs):
"""
Method to copy config files from workflow to experiment directory
Expand Down Expand Up @@ -464,7 +239,7 @@ def _gefs_args(parser):
description = """
Setup files and directories to start a GFS parallel.\n
Create EXPDIR, copy config files.\n
Create ROTDIR experiment directory structure,
Create ROTDIR experiment directory,
"""

parser = ArgumentParser(description=description,
Expand Down Expand Up @@ -576,7 +351,6 @@ def main(*argv):

if create_rotdir:
makedirs_if_missing(rotdir)
fill_ROTDIR(host, user_inputs)

if create_expdir:
makedirs_if_missing(expdir)
Expand Down

0 comments on commit ac9fe80

Please sign in to comment.