diff --git a/polaris/tasks/ocean/add_tasks.py b/polaris/tasks/ocean/add_tasks.py index 975e3d277..4d1dc2cc0 100644 --- a/polaris/tasks/ocean/add_tasks.py +++ b/polaris/tasks/ocean/add_tasks.py @@ -1,4 +1,5 @@ from polaris.tasks.ocean.baroclinic_channel import add_baroclinic_channel_tasks +from polaris.tasks.ocean.barotropic_channel import add_barotropic_channel_tasks from polaris.tasks.ocean.barotropic_gyre import add_barotropic_gyre_tasks from polaris.tasks.ocean.cosine_bell import add_cosine_bell_tasks from polaris.tasks.ocean.external_gravity_wave import ( @@ -31,6 +32,7 @@ def add_ocean_tasks(component): """ # planar tasks add_baroclinic_channel_tasks(component=component) + add_barotropic_channel_tasks(component=component) add_barotropic_gyre_tasks(component=component) add_ice_shelf_2d_tasks(component=component) add_inertial_gravity_wave_tasks(component=component) diff --git a/polaris/tasks/ocean/barotropic_channel/__init__.py b/polaris/tasks/ocean/barotropic_channel/__init__.py new file mode 100644 index 000000000..ae3c62da0 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/__init__.py @@ -0,0 +1,21 @@ +from polaris.config import PolarisConfigParser as PolarisConfigParser +from polaris.tasks.ocean.barotropic_channel.default import Default as Default + + +def add_barotropic_channel_tasks(component): + """ + Add tasks for barotropic channel tests to the ocean component + + component : polaris.tasks.ocean.Ocean + the ocean component that the tasks will be added to + """ + + group_name = 'barotropic_channel' + config_filename = f'{group_name}.cfg' + config = PolarisConfigParser(filepath=f'polaris.tasks.{config_filename}') + config.add_from_package( + f'polaris.tasks.ocean.{group_name}', config_filename + ) + default = Default(component=component) + default.set_shared_config(config, link=config_filename) + component.add_task(default) diff --git a/polaris/tasks/ocean/barotropic_channel/barotropic_channel.cfg b/polaris/tasks/ocean/barotropic_channel/barotropic_channel.cfg new file mode 100644 index 000000000..d5ef11603 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/barotropic_channel.cfg @@ -0,0 +1,39 @@ +# Options related to the vertical grid +[vertical_grid] + +# the type of vertical grid +grid_type = uniform + +# Number of vertical levels +vert_levels = 3 + +# Depth of the bottom of the ocean +bottom_depth = 100.0 + +# The type of vertical coordinate (e.g. z-level, z-star) +coord_type = z-level + +# Whether to use "partial" or "full", or "None" to not alter the topography +partial_cell_type = None + +# The minimum fraction of a layer for partial cells +min_pc_fraction = 0.1 + +# config options for baroclinic channel testcases +[barotropic_channel] + +# the size of the domain in km in the x and y directions +ly = 30.0 +lx = 100.0 + +zonal_velocity = 0.0 +meridional_velocity = 0.0 + +zonal_wind_stress = 0.1 +meridional_wind_stress = 0.0 + +resolution = 10. + +horizontal_viscosity = 1.e-2 + +bottom_drag = 1.e2 diff --git a/polaris/tasks/ocean/barotropic_channel/default/__init__.py b/polaris/tasks/ocean/barotropic_channel/default/__init__.py new file mode 100644 index 000000000..481d515d1 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/default/__init__.py @@ -0,0 +1,43 @@ +from polaris import Task as Task +from polaris.tasks.ocean.barotropic_channel.forward import Forward as Forward +from polaris.tasks.ocean.barotropic_channel.init import Init as Init +from polaris.tasks.ocean.barotropic_channel.viz import Viz as Viz + + +class Default(Task): + """ + The default barotropic channel test case simply creates the mesh and + initial condition, then performs a short forward run on 4 cores. + """ + + def __init__(self, component, indir=None): + """ + Create the test case + + Parameters + ---------- + component : polaris.tasks.ocean.Ocean + The ocean component that this task belongs to + + indir : str + The directory the task is in, to which ``name`` will be appended + """ + group_name = 'barotropic_channel' + base_dir = f'planar/{group_name}' + if indir is None: + indir = base_dir + test_name = 'default' + super().__init__(component=component, name=test_name, indir=indir) + + init_step = Init(component=component, indir=f'{indir}/{test_name}') + self.add_step(init_step) + + self.add_step( + Forward( + component=component, + indir=self.subdir, + graph_target=f'{init_step.path}/culled_graph.info', + ) + ) + + self.add_step(Viz(component=component, indir=self.subdir)) diff --git a/polaris/tasks/ocean/barotropic_channel/forward.py b/polaris/tasks/ocean/barotropic_channel/forward.py new file mode 100644 index 000000000..ea0485a36 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/forward.py @@ -0,0 +1,132 @@ +from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.model import OceanModelStep + + +class Forward(OceanModelStep): + """ + A step for performing forward MPAS-Ocean runs as part of overflow + test cases. + + Attributes + ---------- + yaml_filename : str + The name of the yaml file for this forward step + """ + + def __init__( + self, + component, + graph_target='culled_graph.info', + yaml_filename='forward.yaml', + name='forward', + subdir=None, + indir=None, + ntasks=None, + min_tasks=None, + openmp_threads=1, + ): + """ + Create a new test case + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + name : str + the name of the task + + task_name : str + The name of the task that this step belongs to + + yaml_filename : str + The name of the yaml file for this forward step + + subdir : str, optional + the subdirectory for the step. The default is ``name`` + + 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 the test group) + """ + if min_tasks is None: + min_tasks = ntasks + super().__init__( + component=component, + name=name, + subdir=subdir, + indir=indir, + ntasks=ntasks, + min_tasks=min_tasks, + openmp_threads=openmp_threads, + graph_target=graph_target, + ) + self.yaml_filename = yaml_filename + + # make sure output is double precision + self.add_yaml_file('polaris.ocean.config', 'output.yaml') + + self.add_input_file( + filename='initial_state.nc', + target='../init/initial_state.nc', + ) + + self.add_input_file( + filename='forcing.nc', + target='../init/forcing.nc', + ) + + self.add_output_file( + filename='output.nc', + validate_vars=[ + 'layerThickness', + 'normalVelocity', + 'temperature', + ], + ) + + def dynamic_model_config(self, at_setup): + super().dynamic_model_config(at_setup=at_setup) + + config = self.config + nu = config.getfloat('barotropic_channel', 'horizontal_viscosity') + drag = config.getfloat('barotropic_channel', 'bottom_drag') + replacements = dict(nu=nu, drag=drag) + self.add_yaml_file( + 'polaris.tasks.ocean.barotropic_channel', + self.yaml_filename, + template_replacements=replacements, + ) + model = config.get('ocean', 'model') + vert_levels = config.getfloat('vertical_grid', 'vert_levels') + if model == 'mpas-ocean' and vert_levels == 1: + self.add_yaml_file('polaris.ocean.config', 'single_layer.yaml') + + 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['barotropic_channel'] + lx = section.getfloat('lx') + ly = section.getfloat('ly') + resolution = section.getfloat('resolution') + nx, ny = compute_planar_hex_nx_ny(lx, ly, resolution) + cell_count = nx * ny + return cell_count diff --git a/polaris/tasks/ocean/barotropic_channel/forward.yaml b/polaris/tasks/ocean/barotropic_channel/forward.yaml new file mode 100644 index 000000000..9bed20f86 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/forward.yaml @@ -0,0 +1,64 @@ +mpas-ocean: + time_management: + config_stop_time: none + config_run_duration: 0010_00:00:00 + time_integration: + config_dt: 00:01:00 + config_time_integrator: split_explicit_ab2 + #config_time_integrator: RK4 + io: + config_write_output_on_startup: false + lateral_walls: + config_wall_slip_factor: 0.0 + hmix_del2: + config_use_mom_del2: true + config_mom_del2: {{ nu }} + pressure_gradient: + config_pressure_gradient_type: pressure_and_zmid + bottom_drag: + config_implicit_bottom_drag_type: constant + config_implicit_constant_bottom_drag_coeff: {{ drag }} + ALE_vertical_grid: + config_vert_coord_movement: uniform_stretching + config_ALE_thickness_proportionality: restingThickness_times_weights + cvmix: + config_cvmix_background_scheme: none + config_use_cvmix_convection: true + split_explicit_ts: + config_btr_dt: 00:00:05 + forcing: + config_use_bulk_wind_stress: true + debug: + config_disable_vel_coriolis: true + streams: + mesh: + filename_template: initial_state.nc + input: + filename_template: initial_state.nc + restart: + output_interval: 0000_01:00:00 + forcing: + filename_template: forcing.nc + input_interval: initial_only + type: input + contents: + - windStressZonal + - windStressMeridional + output: + type: output + filename_template: output.nc + output_interval: 0000_01:00:00 + clobber_mode: truncate + contents: + - tracers + - xtime + - normalVelocity + - velocityZonal + - velocityMeridional + - layerThickness + - ssh + - vertVelocityTop + - density + - daysSinceStartOfSim + - circulation + - relativeVorticity diff --git a/polaris/tasks/ocean/barotropic_channel/init.py b/polaris/tasks/ocean/barotropic_channel/init.py new file mode 100644 index 000000000..71063a825 --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/init.py @@ -0,0 +1,117 @@ +import numpy as np +import xarray as xr +from mpas_tools.io import write_netcdf +from mpas_tools.mesh.conversion import convert, cull +from mpas_tools.planar_hex import make_planar_hex_mesh + +from polaris import Step +from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.vertical import init_vertical_coord + + +class Init(Step): + """ + A step for creating a mesh and initial condition for barotropic channel + tasks + + Attributes + ---------- + """ + + def __init__(self, component, indir): + """ + Create the step + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + indir : str + the directory the step is in, to which ``name`` will be appended + """ + super().__init__(component=component, name='init', indir=indir) + + for file in [ + 'base_mesh.nc', + 'culled_mesh.nc', + 'culled_graph.info', + 'forcing.nc', + ]: + self.add_output_file(file) + self.add_output_file( + 'initial_state.nc', + validate_vars=['temperature', 'salinity', 'layerThickness'], + ) + + def run(self): + """ + Run this step of the task + """ + config = self.config + logger = self.logger + + section = config['barotropic_channel'] + resolution = section.getfloat('resolution') + lx = section.getfloat('lx') + ly = section.getfloat('ly') + u = section.getfloat('zonal_velocity') + v = section.getfloat('meridional_velocity') + u_wind = section.getfloat('zonal_wind_stress') + v_wind = section.getfloat('meridional_wind_stress') + + # these could be hard-coded as functions of specific supported + # resolutions but it is preferable to make them algorithmic like here + # for greater flexibility + nx, ny = compute_planar_hex_nx_ny(lx, ly, resolution) + ny = 4 + dc = 1e3 * resolution + + ds_mesh = make_planar_hex_mesh( + nx=nx, ny=ny, dc=dc, nonperiodic_x=False, nonperiodic_y=True + ) + write_netcdf(ds_mesh, 'base_mesh.nc') + + ds_mesh = cull(ds_mesh, logger=logger) + ds_mesh = convert( + ds_mesh, graphInfoFileName='culled_graph.info', logger=logger + ) + write_netcdf(ds_mesh, 'culled_mesh.nc') + + ds = ds_mesh.copy() + + bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') + y_min = ds.yCell.min().values + y_max = ds.yCell.max().values + y_cell = ds.yCell + frac = xr.where((y_cell <= y_min) | (y_cell >= y_max), 2.0 / 3.0, 1.0) + ds['bottomDepth'] = bottom_depth * frac + ds['ssh'] = xr.zeros_like(ds.xCell) + init_vertical_coord(config, ds) + + cell_field = xr.ones_like(ds.xCell) + cell_field, _ = xr.broadcast(cell_field, ds.refBottomDepth) + ds['temperature'] = cell_field.expand_dims(dim='Time', axis=0) + ds['salinity'] = 35.0 * cell_field.expand_dims(dim='Time', axis=0) + normal_velocity = u * np.cos(ds_mesh.angleEdge) + v * np.sin( + ds_mesh.angleEdge + ) + normal_velocity, _ = xr.broadcast(normal_velocity, ds.refBottomDepth) + normal_velocity = normal_velocity.transpose('nEdges', 'nVertLevels') + ds['normalVelocity'] = normal_velocity.expand_dims(dim='Time', axis=0) + ds['fCell'] = xr.zeros_like(ds.xCell) + ds['fEdge'] = xr.zeros_like(ds.xEdge) + ds['fVertex'] = xr.zeros_like(ds.xVertex) + write_netcdf(ds, 'initial_state.nc') + + # set the wind stress forcing + ds_forcing = xr.Dataset() + wind_stress_zonal = u_wind * xr.ones_like(ds.xCell) + wind_stress_meridional = v_wind * xr.ones_like(ds.xCell) + ds_forcing['windStressZonal'] = wind_stress_zonal.expand_dims( + dim='Time', axis=0 + ) + ds_forcing['windStressMeridional'] = ( + wind_stress_meridional.expand_dims(dim='Time', axis=0) + ) + write_netcdf(ds_forcing, 'forcing.nc') diff --git a/polaris/tasks/ocean/barotropic_channel/viz.py b/polaris/tasks/ocean/barotropic_channel/viz.py new file mode 100644 index 000000000..fe78b217e --- /dev/null +++ b/polaris/tasks/ocean/barotropic_channel/viz.py @@ -0,0 +1,98 @@ +import cmocean # noqa: F401 +import numpy as np +import xarray as xr + +from polaris import Step + +# from polaris.mpas import cell_mask_to_edge_mask +from polaris.viz import plot_horiz_field + + +class Viz(Step): + """ + A step for plotting the results of a series of baroclinic channel RPE runs + """ + + def __init__(self, component, indir): + """ + Create the step + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + indir : str + the directory the step is in, to which ``name`` will be appended + """ + super().__init__(component=component, name='viz', indir=indir) + self.add_input_file( + filename='mesh.nc', target='../init/culled_mesh.nc' + ) + self.add_input_file( + filename='init.nc', target='../init/initial_state.nc' + ) + self.add_input_file( + filename='output.nc', target='../forward/output.nc' + ) + + def run(self): + """ + Run this step of the task + """ + ds_mesh = xr.load_dataset('mesh.nc') + ds_init = xr.load_dataset('init.nc') + ds_out = xr.load_dataset('output.nc') + ds_out = ds_out.isel(nVertLevels=-1) + + cell_mask = ds_init.maxLevelCell >= 1 + vertex_mask = ds_init.boundaryVertex == 0 + + # Uncomment these lines to get 10 evenly spaced time slices + # nt = ds_out.sizes['Time'] + # for t_index in np.arange(0, nt, int(np.floor(nt / 10))): + + # These indices correspond to the first and last time step + for t_index in [0, -1]: + ds = ds_out.isel(Time=t_index) + vmax = np.max(np.abs(ds.velocityZonal.values)) + plot_horiz_field( + ds_mesh, + ds['velocityZonal'], + f'velocity_zonal_t{t_index}_zbot.png', + vmin=-vmax, + vmax=vmax, + cmap='cmo.balance', + field_mask=cell_mask, + ) + plot_horiz_field( + ds_mesh, + ds['velocityMeridional'], + f'velocity_meridional_t{t_index}_zbot.png', + vmin=-vmax, + vmax=vmax, + cmap='cmo.balance', + field_mask=cell_mask, + ) + + vmax = np.max(np.abs(ds.relativeVorticity.values)) + plot_horiz_field( + ds_mesh, + ds['relativeVorticity'], + f'relative_vorticity_t{t_index}_zbot.png', + vmin=-vmax, + vmax=vmax, + cmap='cmo.balance', + field_mask=vertex_mask, + ) + + vmax = np.max(np.abs(ds.circulation.values)) + plot_horiz_field( + ds_mesh, + ds['circulation'], + f'circulation_t{t_index}_zbot.png', + vmin=-vmax, + vmax=vmax, + cmap='cmo.balance', + field_mask=vertex_mask, + )