import os
from polaris.mesh.planar import compute_planar_hex_nx_ny
from polaris.ocean.model import OceanModelStep, get_time_interval_string
[docs]
class Forward(OceanModelStep):
"""
A step for performing forward ocean component runs as part of baroclinic
channel tasks.
Attributes
----------
resolution : float
The resolution of the task in km
dt : float
The model time step in seconds
btr_dt : float
The model barotropic time step in seconds
run_time_steps : int or None
Number of time steps to run for
"""
[docs]
def __init__(
self,
component,
resolution,
init,
name='forward',
subdir=None,
indir=None,
ntasks=None,
min_tasks=None,
openmp_threads=1,
nu=None,
run_time_steps=None,
start_time_steps=None,
restart_time_steps=None,
graph_target='culled_graph.info',
):
"""
Create a new task
Parameters
----------
component : polaris.Component
The component the step belongs to
resolution : km
The resolution of the task in km
init : polaris.step
The init step for the task
name : str
the name of the step
subdir : str, optional
the subdirectory for the step. If neither this nor ``indir``
are provided, the directory is the ``name``
indir : str, optional
the directory the step is in, to which ``name`` will be appended
ntasks : int, optional
the number of tasks the step would ideally use. If fewer tasks
are available on the system, the step will run on all available
tasks as long as this is not below ``min_tasks``
min_tasks : int, optional
the number of tasks the step requires. If the system has fewer
than this number of tasks, the step will fail
openmp_threads : int, optional
the number of OpenMP threads the step will use
nu : float, optional
the viscosity (if different from the default for baroclinic channel
tests)
run_time_steps : int, optional
Number of time steps to run for
start_time_steps : int, optional
Number of time steps to start from
graph_target : str, optional
The graph file name (relative to the base work directory).
If none, it will be created.
"""
self.nu = nu
if self.nu is not None:
self.is_rpe_step = True
else:
self.is_rpe_step = False
self.resolution = resolution
self.run_time_steps = run_time_steps
self.start_time_steps = start_time_steps
if start_time_steps is not None:
self.do_restart = True
else:
self.do_restart = False
self.restart_time_steps = restart_time_steps
super().__init__(
component=component,
name=name,
subdir=subdir,
indir=indir,
ntasks=ntasks,
min_tasks=min_tasks,
openmp_threads=openmp_threads,
graph_target=graph_target,
)
# make sure output is double precision
self.add_yaml_file('polaris.ocean.config', 'output.yaml')
self.add_horiz_mesh_input_file(
work_dir_target=f'{init.path}/culled_mesh.nc'
)
self.add_vert_coord_input_file(
work_dir_target=f'{init.path}/vert_coord.nc'
)
self.add_init_input_file(work_dir_target=f'{init.path}/init.nc')
self.add_output_file(
filename='output.nc',
validate_class='state',
)
self.dt = None
self.btr_dt = None
def setup(self):
"""
Set up the step, adding the OmegaMesh.nc symlink for Omega
"""
super().setup()
model = self.config.get('ocean', 'model')
# TODO: remove as soon as Omega no longer hard-codes this file
if model == 'omega':
self.add_input_file(filename='OmegaMesh.nc', target='init.nc')
[docs]
def compute_cell_count(self):
"""
Compute the approximate number of cells in the mesh, used to constrain
resources
Returns
-------
cell_count : int or None
The approximate number of cells in the mesh
"""
section = self.config['baroclinic_channel']
lx = section.getfloat('lx')
ly = section.getfloat('ly')
nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
cell_count = nx * ny
return cell_count
[docs]
def dynamic_model_config(self, at_setup):
"""
Add model config options, namelist, streams and yaml files using config
options or template replacements that need to be set both during step
setup and at runtime
Parameters
----------
at_setup : bool
Whether this method is being run during setup of the step, as
opposed to at runtime
"""
super().dynamic_model_config(at_setup)
config = self.config
model = config.get('ocean', 'model')
time_integrator = config.get('baroclinic_channel', 'time_integrator')
time_integrator_map = dict(
[('RK4', 'RungeKutta4'), ('split_explicit', 'SplitExplicitRK2')]
)
if model == 'omega':
if time_integrator in time_integrator_map.keys():
time_integrator = time_integrator_map[time_integrator]
else:
print(
'Warning: mapping from time integrator '
f'{time_integrator} to omega not found, '
'retaining name given in config'
)
# dt is proportional to resolution
dt_per_km = config.getfloat('baroclinic_channel', 'dt_per_km')
dt = dt_per_km * self.resolution
dt_str = get_time_interval_string(seconds=dt)
# btr_dt is shared by Omega and MPAS-Ocean
btr_dt_per_km = config.getfloat('baroclinic_channel', 'btr_dt_per_km')
btr_dt = btr_dt_per_km * self.resolution
btr_dt_str = get_time_interval_string(seconds=btr_dt)
mpaso_options = {
'config_do_restart': False,
}
# Set dt and default run duration, which may be changed below
ocean_options = {
'config_dt': dt_str,
'config_btr_dt': btr_dt_str,
}
start_str = '0001-01-01_00:00:00'
if self.start_time_steps is not None:
start_seconds = self.start_time_steps * dt
start_str = get_time_interval_string(seconds=start_seconds)
# Assume the default start time for the full_run
start_str = f'0001-01-01_{start_str.split("_")[1]}'
ocean_options['config_start_time'] = start_str
if self.name == 'long_forward':
run_duration = config.getfloat(
'baroclinic_channel_long', 'run_duration'
)
run_duration_str = get_time_interval_string(days=run_duration)
output_freq = config.getfloat(
'baroclinic_channel_long', 'output_interval'
)
output_freq_units = config.get(
'baroclinic_channel_long', 'output_interval_units'
)
output_freq = int(output_freq)
elif self.is_rpe_step:
run_duration = config.getfloat(
'baroclinic_channel_rpe', 'run_duration'
)
run_duration_str = get_time_interval_string(days=run_duration)
output_freq = 1
output_freq_units = 'days'
elif self.run_time_steps is not None:
run_seconds = self.run_time_steps * dt
run_duration_str = get_time_interval_string(seconds=run_seconds)
output_freq = int(run_seconds)
output_freq_units = 'seconds'
else:
raise ValueError(
'Could not determine run duration and output frequency for run'
)
ocean_options['config_run_duration'] = run_duration_str
# Get output interval for mpas-ocean
if model == 'mpas-ocean':
if output_freq_units == 'seconds':
seconds = output_freq
elif output_freq_units == 'minutes':
seconds = output_freq * 60
elif output_freq_units == 'hours':
seconds = output_freq * 3600
elif output_freq_units == 'days':
seconds = output_freq * 86400
else:
raise ValueError(
'Warning: output frequency units '
f'{output_freq_units} not supported'
)
seconds = 1.0
output_interval = get_time_interval_string(seconds=seconds)
else:
output_interval = ''
if self.restart_time_steps is not None:
# restart units are set to seconds below
restart_interval = get_time_interval_string(
seconds=dt * self.restart_time_steps
)
# We store restarts one level up from the step to be easily
# accessed by multiple steps
restart_dir = os.path.abspath(
os.path.join(self.work_dir, '..', 'restarts')
)
os.makedirs(restart_dir, exist_ok=True)
restart_filename = f'{restart_dir}/ocn.rst.$Y-$M-$D_$h.$m.$s'
restart_freq = int(dt)
else:
restart_interval = get_time_interval_string(days=1)
restart_filename = 'ocn.rst.$Y-$M-$D_$h.$m.$s'
restart_freq = int(24 * 3600)
if self.do_restart:
init_freq = 'never'
start_type = 'Continue'
mpaso_options['config_do_restart'] = True
else:
init_freq = 'OnStartup'
start_type = 'StartUp'
if self.nu is not None:
# update the viscosity to the requested value *after* loading
# forward.yaml
ocean_options['config_mom_del2'] = self.nu
else:
ocean_options['config_mom_del2'] = config.getfloat(
'baroclinic_channel', 'default_nu'
)
self.add_model_config_options(
options=ocean_options, config_model='ocean'
)
self.add_model_config_options(
options=mpaso_options, config_model='mpas-ocean'
)
replacements = dict(
init_freq=init_freq,
restart_filename=restart_filename,
start_type=start_type,
output_freq=f'{output_freq}',
output_freq_units=output_freq_units,
time_integrator=time_integrator,
output_interval=output_interval,
restart_interval=restart_interval,
restart_freq=restart_freq,
restart_freq_units='seconds',
)
self.add_yaml_file(
'polaris.tasks.ocean.baroclinic_channel',
'forward.yaml',
template_replacements=replacements,
)