diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index 7b386414da..d071b9590f 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -120,6 +120,30 @@ ``` +### drying_slope + +```{eval-rst} +.. currentmodule:: polaris.ocean.tasks.drying_slope + +.. autosummary:: + :toctree: generated/ + + add_drying_slope_tasks + + baroclinic.Baroclinic + baroclinic.Baroclinic.configure + + barotropic.Barotropic + + convergence.Convergence + convergence.analysis.Analysis + convergence.analysis.Analysis.exact_solution + convergence.forward.Forward + convergence.forward.Forward.compute_cell_count + convergence.forward.Forward.dynamic_model_config + + decomp.Decomp + ### geostrophic ```{eval-rst} diff --git a/docs/developers_guide/ocean/tasks/drying_slope.md b/docs/developers_guide/ocean/tasks/drying_slope.md new file mode 100644 index 0000000000..c53b7fc5f4 --- /dev/null +++ b/docs/developers_guide/ocean/tasks/drying_slope.md @@ -0,0 +1,100 @@ +(dev-ocean-drying-slope)= + +# drying_slope + +The drying slope tests in `polaris.ocean.tasks.drying_slope` are +variants of the drying slope test case (see {ref}`ocean-drying-slope`). +Here, we describe the test cases and their shared framework. + +(dev-ocean-drying-slope-framework)= + +## framework + +The shared config options for `drying_slope` tests are described in +{ref}`ocean-drying-slope` in the User's Guide. + +Additionally, the tests share a `forward.yaml` file with a few common model +config options related to tidal forcing and wetting and drying, as well as +defining `mesh`, `input`, `restart`, and `output` streams. + +### init + +The class {py:class}`polaris.ocean.tasks.drying_slope.init.Init` +defines a step for setting up the initial state for each test case. + +First, a mesh appropriate for the resolution is generated using +{py:func}`mpas_tools.planar_hex.make_planar_hex_mesh()`. Then, the mesh is +culled to remove periodicity in the y direction. The bathymetry is generated +with a linearly sloping bed and ssh is set to be consistent with the initial +tidal forcing. A vertical grid is then generated, according to the +`vertical_grid` config options for each task. Next, the initial +temperature field is computed with a horizontal anomaly. +The salinity field is uniform in the case of the `barotropic` test and has +vertical gradients in the case of the `baroclinic` test. The initial velocity +field is zeros. Finally, the initial fields are plotted. + +The same `init` step is shared by all tasks at a given resolution. + +### forward + +The class {py:class}`polaris.ocean.tasks.drying_slope.forward.Forward` +defines a step for running MPAS-Ocean from the initial condition produced in +the `init` step. The time integration scheme, drag options, and tidal forcing +are set here. Namelist and streams files are updated in +{py:meth}`polaris.ocean.tasks.drying_slope.forward.Forward.dynamic_model_config()` +with time steps determined algorithmically based on config options. The +number of cells is approximated from config options in +{py:meth}`polaris.ocean.tasks.drying_slope.forward.Forward.compute_cell_count()` +so that this can be used to constrain the number of MPI tasks that Polaris +tasks have as their target and minimum (if the resources are not explicitly +prescribed). For MPAS-Ocean, PIO namelist options are modified and a +graph partition is generated as part of `runtime_setup()`. Next, the ocean +model is run. + +### validate + +The class {py:class}`polaris.ocean.tasks.drying_slope.validate.Validate` +defines a step for validating outputs in two step directories against one +another. This step ensures that `temperature`, `salinity`, `layerThickness` +and `normalVelocity` are identical in `output.nc` files in the two steps. + +### viz + +The class {py:class}`polaris.ocean.tasks.drying_slope.viz.Viz` +defines a step for visualizing horizontal fields and transects from the +forward step as well as the ssh forcing and reference solutions. + +(dev-ocean-drying-slope-baroclinic)= + +## baroclinic + +The {py:class}`polaris.ocean.tasks.drying_slope.baroclinic.Baroclinic` +test runs a drying test at a linear tidal forcing rate with vertical +stratification. + +(dev-ocean-drying-slope-barotropic)= + +## barotropic + +The {py:class}`polaris.ocean.tasks.drying_slope.barotropic.Barotropic` +test runs one sinusoidal cycle of tidal forcing. + +(dev-ocean-drying-slope-decomp)= + +## decomp + +The {py:class}`polaris.ocean.tasks.drying_slope.decomp.Decomp` +performs a 3-time-step run once on 4 cores and once on 8 cores. The +`validate` step ensures that the two runs produce identical results. + +(dev-ocean-drying-slope-convergence)= + +## convergence + +The {py:class}`polaris.ocean.tasks.drying_slope.convergence.Convergence` +tests the convergence of the barotropic test with resolution and time step +against a reference solution. + +The `analysis` step defined by +{py:class}`polaris.ocean.tasks.drying_slope.convergence.analysis.Analysis` +produces a figure of the RMSE as a function of resolution. diff --git a/docs/developers_guide/ocean/tasks/index.md b/docs/developers_guide/ocean/tasks/index.md index fb6448ad51..d8cc681d6a 100644 --- a/docs/developers_guide/ocean/tasks/index.md +++ b/docs/developers_guide/ocean/tasks/index.md @@ -10,6 +10,7 @@ barotropic_gyre correlated_tracers_2d cosine_bell external_gravity_wave +drying_slope geostrophic divergent_2d ice_shelf_2d diff --git a/docs/users_guide/ocean/tasks/drying_slope.md b/docs/users_guide/ocean/tasks/drying_slope.md new file mode 100644 index 0000000000..611a87b33f --- /dev/null +++ b/docs/users_guide/ocean/tasks/drying_slope.md @@ -0,0 +1,241 @@ +(ocean-drying-slope)= + +# drying slope + +The drying_slope tasks test wetting-and-drying algorithms in domains designed +to represent a coastline with sloping bathymetry. + +## config options + +There are a number of config options that are common to all drying_slope tasks: + +```cfg +# config options for all drying slope test cases +[drying_slope] + +# time integration scheme +time_integrator = RK4 + +# time step per resolution (s/km), since dt is proportional to resolution +rk4_dt_per_km = 30 + +split_dt_per_km = 30 + +# barotropic time step per resolution (s/km), since btr_dt is proportional to +# resolution +btr_dt_per_km = 1.5 + +# Coriolis parameter +coriolis_parameter = 0.0 +``` + +(ocean-drying-slope-baroclinic)= + +## baroclinic + +The baroclinic drying_slope task is designed to test wetting-and-drying +algorithms in a multi-layer configuration with vertial gradients in scalars. + +### description + +Description of the test case. + +```{image} images/drying_slope_baroclinic.png +:align: center +:width: 500 px +``` + +### mesh + +Specify whether the mesh is global or planar and the resolution(s) tested. If +planar, specify the mesh size. If global, specify whether the mesh is +icosohedral or quasi-uniform. Specify any relevant options in the config file +pertaining to setting up the mesh. + +### vertical grid + +If there are no restrictions on the vertical grid specifications inherent to +the test case, then the config section may be provided without any further +description. + +Examples of restrictions or special conditions warranting description may +include: + +* Whether the topography is variable +* Whether the test pertains to shallow water dynamics, in which case the +minimum number of vertical levels may be used +* Whether there are several test cases in the test group investigating the +effects of different vertical coordinates (`coord_type`) + +The minimum layer thickness for this case is determined by the minimum column thickness +from the config option `drying_slope_baroclinic:min_column_thickness`. + +```cfg +# Options related to the vertical grid +[vertical_grid] + +# the type of vertical grid +grid_type = uniform + +# The type of vertical coordinate (e.g. z-level, z-star) +coord_type = sigma + +# Number of vertical levels +vert_levels = 10 + +# 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 + +# The minimum number of vertical levels for multi-layer cases +min_vert_levels = 3 + +# Minimum thickness of each layer +min_layer_thickness = ${drying_slope_barotropic:thin_film_thickness} +``` + +### initial conditions + +The initial conditions should be specified for all variables requiring +initial conditions (see {ref}`dev-ocean-models`). + +### forcing + +If applicable, specify the forcing applied at each time step of the simulation +(in MPAS-Ocean, these are the variables contained in the `forcing` stream). +If not applicable, keep this section with the notation N/A. + +### time step and run duration + +The time step for forward integration should be specified here for the test +case's resolution. The run duration should also be specified. + +### config options + +The config options define the geometry of the domain. + +```cfg +# config options for barotropic drying slope test cases +[drying_slope_baroclinic] + +# the width of the domain in km +lx = 12. + +# the length of the domain in km +ly = 55. + +# Length over which wetting and drying actually occur in km +ly_analysis = 50. + +# Bottom depth at the right side of the domain in m +right_bottom_depth = 2.5 + +# Bottom depth at the left side of the domain in m +left_bottom_depth = 0. + +# Initial SSH at the right side of the domain in m +right_tidal_height = 0. + +# salinity at the right side of the domain +right_salinity = 35.0 + +# salinity at the left side of the domain +left_salinity = 1.0 + +# manning coefficient used in mannings bottom drag type +manning_coefficient = 5.0e-2 + +# The column thickness in the thin film region, used to determine thin film +# thickness +min_column_thickness = 0.05 +``` + +Include here any further description of each of the config options. + +### cores + +Specify whether the number of cores is determined by `goal_cells_per_core` and +`max_cells_per_core` in the `ocean` section of the config file or whether the +default and minimum number of cores is given in arguments to the forward step, +and what those defaults are. + +## barotropic + +### description + +Description of the test case. + +```{image} images/drying_slope_barotropic.png +:align: center +:width: 500 px +``` +### vertical grid + +The vertical grid is the same as baroclinic except a single layer may be used and +the minimum layer thickness given in `drying_slope_barotropic:thin_film_thickness` +is used. + +### config options + +The config options define the geometry of the domain. + +```cfg +# config options for barotropic drying slope test cases +[drying_slope_barotropic] + +# the width of the domain in km +lx = 6. + +# Length over which wetting and drying actually occur in km +ly_analysis = 25. + +# the length of the domain in km +ly = 28. + +# Bottom depth at the right side of the domain in m +right_bottom_depth = -10. + +# Bottom depth at the left side of the domain in m +left_bottom_depth = 0. + +# Initial SSH at the right side of the domain in m +right_tidal_height = ${drying_slope_barotropic:right_bottom_depth} + +# Plug width as a fraction of the domain +plug_width_frac = 0.0 + +# Plug temperature +plug_temperature = 20.0 + +# Background temperature +background_temperature = 20.0 + +# Background salinity +background_salinity = 35.0 + +# Thickness of each layer in the thin film region +thin_film_thickness = 1.0e-3 +``` + +## convergence + +### description + +```{image} images/drying_slope_convergence.png +:align: center +:width: 500 px +``` +### config options + +The config options define the resolutions to use in the convergence test: + +```cfg +# config options for drying slope convergence task +[drying_slope_convergence] + +# horizontal resolutions in km +resolutions = 0.25, 0.5, 1, 2 +``` + diff --git a/docs/users_guide/ocean/tasks/images/drying_slope_barotropic.png b/docs/users_guide/ocean/tasks/images/drying_slope_barotropic.png new file mode 100644 index 0000000000..1790c0395a Binary files /dev/null and b/docs/users_guide/ocean/tasks/images/drying_slope_barotropic.png differ diff --git a/docs/users_guide/ocean/tasks/images/drying_slope_convergence.png b/docs/users_guide/ocean/tasks/images/drying_slope_convergence.png new file mode 100644 index 0000000000..568e4eef7c Binary files /dev/null and b/docs/users_guide/ocean/tasks/images/drying_slope_convergence.png differ diff --git a/docs/users_guide/ocean/tasks/index.md b/docs/users_guide/ocean/tasks/index.md index ec85fdcda5..fd59a90c7a 100644 --- a/docs/users_guide/ocean/tasks/index.md +++ b/docs/users_guide/ocean/tasks/index.md @@ -10,6 +10,7 @@ barotropic_gyre correlated_tracers_2d cosine_bell external_gravity_wave +drying_slope geostrophic divergent_2d ice_shelf_2d diff --git a/polaris/ocean/__init__.py b/polaris/ocean/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/polaris/ocean/convergence/analysis.py b/polaris/ocean/convergence/analysis.py index 4760174aa3..595ab58f45 100644 --- a/polaris/ocean/convergence/analysis.py +++ b/polaris/ocean/convergence/analysis.py @@ -373,12 +373,19 @@ def compute_error( zidx=zidx, ) diff = field_exact - field_mpas + # We support nans in the fields, used to indicate regions where + # convergence should not be evaluated + mask = np.logical_and(field_exact == field_exact, + field_mpas == field_mpas) + diff = field_exact - field_mpas # Only the L2 norm is area-weighted if error_type == 'l2': area = area_for_field(ds_mesh, diff) - field_exact = field_exact * area - diff = diff * area + mean_area = area.mean() + field_exact = field_exact[mask] * area[mask] / mean_area + diff = diff[mask] * area[mask] / mean_area + error = np.linalg.norm(diff, ord=norm_type[error_type]) # Normalize the error norm by the vector norm of the exact solution diff --git a/polaris/ocean/convergence/forward.py b/polaris/ocean/convergence/forward.py index 58923a8a7a..284afa3198 100644 --- a/polaris/ocean/convergence/forward.py +++ b/polaris/ocean/convergence/forward.py @@ -39,6 +39,7 @@ def __init__( output_filename='output.nc', validate_vars=None, refinement='both', + forcing=False, ): """ Create a new step @@ -113,6 +114,12 @@ def __init__( filename=output_filename, validate_vars=validate_vars ) + if forcing: + self.add_input_file( + filename='forcing.nc', + work_dir_target=f'{init.path}/forcing.nc', + ) + def compute_cell_count(self): """ Compute the approximate number of cells in the mesh, used to constrain diff --git a/polaris/tasks/ocean/add_tasks.py b/polaris/tasks/ocean/add_tasks.py index 741c2ba671..8f8082c12d 100644 --- a/polaris/tasks/ocean/add_tasks.py +++ b/polaris/tasks/ocean/add_tasks.py @@ -4,6 +4,7 @@ from polaris.tasks.ocean.external_gravity_wave import ( add_external_gravity_wave_tasks as add_external_gravity_wave_tasks, ) +from polaris.tasks.ocean.drying_slope import add_drying_slope_tasks from polaris.tasks.ocean.geostrophic import add_geostrophic_tasks from polaris.tasks.ocean.ice_shelf_2d import add_ice_shelf_2d_tasks from polaris.tasks.ocean.inertial_gravity_wave import ( @@ -31,6 +32,7 @@ def add_ocean_tasks(component): # planar tasks add_baroclinic_channel_tasks(component=component) add_barotropic_gyre_tasks(component=component) + add_drying_slope_tasks(component=component) add_ice_shelf_2d_tasks(component=component) add_inertial_gravity_wave_tasks(component=component) add_internal_wave_tasks(component=component) diff --git a/polaris/tasks/ocean/drying_slope/__init__.py b/polaris/tasks/ocean/drying_slope/__init__.py new file mode 100644 index 0000000000..bd07a6f3fd --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/__init__.py @@ -0,0 +1,152 @@ +from polaris.config import PolarisConfigParser +from polaris.resolution import resolution_to_string +from polaris.tasks.ocean.drying_slope.baroclinic import Baroclinic +from polaris.tasks.ocean.drying_slope.barotropic import Barotropic +from polaris.tasks.ocean.drying_slope.convergence import Convergence +from polaris.tasks.ocean.drying_slope.decomp import Decomp +from polaris.tasks.ocean.drying_slope.init import Init + + +def add_drying_slope_tasks(component): + """ + Add tasks for different drying slope tests to the ocean component + + component : polaris.ocean.Ocean + the ocean component that the tasks will be added to + """ + config_filename = 'drying_slope.cfg' + + for coord_type in ['sigma', 'single_layer']: + group_dir = f'planar/drying_slope/{coord_type}' + + # The config file lives in the coord_type directory because + # config options change between coordinate types + config = PolarisConfigParser(filepath=f'{group_dir}/{config_filename}') + config.add_from_package('polaris.ocean.convergence', 'convergence.cfg') + config.add_from_package( + 'polaris.tasks.ocean.drying_slope', config_filename + ) + if coord_type == 'single_layer': + config.set( + 'vertical_grid', + 'vert_levels', + '1', + comment='Number of vertical levels', + ) + config.set('vertical_grid', 'coord_type', 'z-level') + else: + config.set('vertical_grid', 'coord_type', 'sigma') + + case_dir = f'{group_dir}/barotropic' + for resolution in [0.25, 1.0]: + resdir = resolution_to_string(resolution) + indir = f'{case_dir}/{resdir}' + + # Add one init step per resolution in barotropic + init_dir = f'{indir}/init' + if init_dir in component.steps: + init = component.steps[init_dir] + else: + init = Init( + component=component, + resolution=resolution, + name=f'init_{resdir}', + subdir=init_dir, + baroclinic=False, + ) + init.set_shared_config(config, link=config_filename) + + for method in ['standard', 'ramp']: + if coord_type == 'single_layer': + default = Barotropic( + component=component, + resolution=resolution, + subdir=f'{indir}/{method}', + init=init, + method=method, + coord_type=coord_type, + drag_type='constant', + ) + else: + default = Barotropic( + component=component, + resolution=resolution, + subdir=f'{indir}/{method}', + init=init, + method=method, + coord_type=coord_type, + ) + default.set_shared_config(config, link=config_filename) + component.add_task(default) + + if method == 'ramp': + loglaw = Barotropic( + component=component, + resolution=resolution, + subdir=f'{indir}/{method}', + init=init, + method=method, + drag_type='loglaw', + coord_type=coord_type, + ) + loglaw.set_shared_config(config, link=config_filename) + component.add_task(loglaw) + + if ( + method == 'ramp' + and resolution == 1.0 + and coord_type == 'sigma' + ): + decomp = Decomp( + component=component, + resolution=resolution, + indir=indir, + init=init, + method=method, + coord_type=coord_type, + ) + decomp.set_shared_config(config, link=config_filename) + component.add_task(decomp) + + # Add convergence test only for sigma coordinate + if coord_type == 'sigma': + method = 'ramp' + task_dir = f'{case_dir}/convergence/{method}' + convergence = Convergence( + component=component, + subdir=task_dir, + group_dir=group_dir, + config=config, + method=method, + ) + convergence.set_shared_config(config, link=config_filename) + component.add_task(convergence) + + if coord_type != 'single_layer': + case_dir = f'{group_dir}/baroclinic' + forcing_type = 'linear_drying' + resolution = 1.0 + resdir = resolution_to_string(resolution) + indir = f'{case_dir}/{resdir}' + + # Create a new initial condition for the baroclinic case + init = Init( + component=component, + indir=indir, + resolution=resolution, + baroclinic=True, + ) + init.set_shared_config(config, link=config_filename) + + for method in ['standard', 'ramp']: + baroclinic = Baroclinic( + component=component, + resolution=resolution, + init=init, + subdir=f'{indir}/{method}', + coord_type=coord_type, + method=method, + forcing_type=forcing_type, + ) + baroclinic.set_shared_config(config, link=config_filename) + component.add_task(baroclinic) diff --git a/polaris/tasks/ocean/drying_slope/baroclinic.yaml b/polaris/tasks/ocean/drying_slope/baroclinic.yaml new file mode 100644 index 0000000000..f6f5148592 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/baroclinic.yaml @@ -0,0 +1,16 @@ +omega: + tracer_forcing_activeTracers: + config_use_activeTracers: true + tracer_forcing_debugTracers: + config_use_debugTracers: true + debug: + config_compute_active_tracer_budgets: true + config_disable_tr_all_tend: false + streams: + output: + contents: + - wettingVelocityBaroclinic + - wettingVelocityBarotropicSubcycle + forcing: + contents: + - bottomDrag diff --git a/polaris/tasks/ocean/drying_slope/baroclinic/__init__.py b/polaris/tasks/ocean/drying_slope/baroclinic/__init__.py new file mode 100644 index 0000000000..7c5b154e01 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/baroclinic/__init__.py @@ -0,0 +1,114 @@ +from polaris import Task +from polaris.tasks.ocean.drying_slope.forward import Forward +from polaris.tasks.ocean.drying_slope.viz import Viz + + +class Baroclinic(Task): + """ + The baroclinic version of the drying slope test case creates the mesh + and initial condition, then performs a short forward run on 4 cores. + + Attributes + ---------- + coord_type : str, optional + The vertical coordinate type + + config : polaris.config.PolarisConfigParser + A shared config parser + """ + + def __init__( + self, + component, + resolution, + init, + subdir, + coord_type='sigma', + forcing_type='linear_drying', + method='ramp', + drag_type='constant_and_rayleigh', + ): + """ + Create the test case + + Parameters + ---------- + component : polaris.ocean.Ocean + The ocean component that this task belongs to + + resolution : float + The resolution of the test case in km + + init : polaris.tasks.ocean.drying_slope.init.Init + A shared step for creating the initial state + + subdir : str + The subdirectory to put the task group in + + coord_type : str, optional + The vertical coordinate type + + forcing_type : str, optional + The forcing type to apply at the "tidal" boundary as a namelist + option + + method : str, optional + The type of wetting and drying algorithm to use + + drag_type : str, optional + The bottom drag type to apply as a namelist option + """ + name = f'baroclinic_{method}' + if drag_type == 'loglaw': + name = f'{name}_{drag_type}' + subdir = f'{subdir}_{drag_type}' + super().__init__(component=component, name=name, subdir=subdir) + self.coord_type = coord_type + + self.add_step(init, symlink='init') + self.add_step( + Forward( + component=component, + init=init, + indir=subdir, + subdir=None, + ntasks=None, + min_tasks=None, + openmp_threads=1, + resolution=resolution, + forcing_type=forcing_type, + coord_type=coord_type, + drag_type=drag_type, + damping_coeff=1.0e-4, + baroclinic=True, + method=method, + graph_target=f'{init.path}/culled_graph.info', + ) + ) + + self.add_step( + Viz( + component=component, + indir=subdir, + subdir=None, + damping_coeffs=None, + baroclinic=True, + forcing_type=forcing_type, + ) + ) + + def configure(self): + """ + Change config options as needed + """ + config = self.config + hmin = config.getfloat( + 'drying_slope_baroclinic', 'min_column_thickness' + ) + nz = config.getint('vertical_grid', 'vert_levels') + self.config.set( + 'drying_slope', + 'thin_film_thickness', + f'{hmin / nz}', + comment='Thickness of each layer in the thin film region', + ) diff --git a/polaris/tasks/ocean/drying_slope/barotropic/__init__.py b/polaris/tasks/ocean/drying_slope/barotropic/__init__.py new file mode 100644 index 0000000000..278e2d8be6 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/barotropic/__init__.py @@ -0,0 +1,142 @@ +from polaris import Task +from polaris.resolution import resolution_to_string +from polaris.tasks.ocean.drying_slope.forward import Forward +from polaris.tasks.ocean.drying_slope.viz import Viz + + +class Barotropic(Task): + """ + The default drying_slope test case with a sinusoidal forcing function + + Attributes + ---------- + resolution : float + The resolution of the test case in km + + coord_type : str + The type of vertical coordinate (``sigma``, ``single_layer``, etc.) + + damping_coeff : float + The damping coefficient for the rayleigh drag option + """ + + def __init__( + self, + component, + resolution, + init, + subdir, + coord_type='sigma', + drag_type='constant_and_rayleigh', + forcing_type='tidal_cycle', + method='ramp', + ): + """ + Create the test case + + Parameters + ---------- + component : polaris.ocean.Ocean + The ocean component that this task belongs to + + resolution : float + The resolution of the test case in km + + init : polaris.tasks.ocean.drying_slope.init.Init + A shared step for creating the initial state + + subdir : str + The subdirectory to put the task group in + + coord_type : str, optional + The vertical coordinate type + + forcing_type : str, optional + The forcing type to apply at the "tidal" boundary as a namelist + option + + method : str, optional + The type of wetting and drying algorithm to use + + drag_type : str, optional + The bottom drag type to apply as a namelist option + """ + mesh_name = resolution_to_string(resolution) + name = f'barotropic_{method}_{mesh_name}' + if drag_type == 'loglaw': + name = f'{name}_{drag_type}' + subdir = f'{subdir}_{drag_type}' + super().__init__(component=component, name=name, subdir=subdir) + self.resolution = resolution + self.coord_type = coord_type + self.add_step(init, symlink='init') + if drag_type == 'constant_and_rayleigh': + self.damping_coeffs = [0.0025, 0.01] + for damping_coeff in self.damping_coeffs: + step_name = f'forward_{damping_coeff:03g}' + forward_dir = f'{subdir}/{step_name}' + if forward_dir in component.steps: + forward_step = component.steps[forward_dir] + symlink = step_name + else: + forward_step = Forward( + component=component, + init=init, + subdir=f'{subdir}/{step_name}', + ntasks=None, + min_tasks=None, + openmp_threads=1, + name=f'{step_name}_{mesh_name}', + resolution=resolution, + forcing_type=forcing_type, + coord_type=coord_type, + drag_type=drag_type, + damping_coeff=damping_coeff, + baroclinic=False, + method=method, + graph_target=f'{init.path}/culled_graph.info', + ) + symlink = None + self.add_step(forward_step, symlink=symlink) + else: + self.damping_coeffs = [] + forward_step = Forward( + component=component, + init=init, + indir=subdir, + ntasks=None, + min_tasks=None, + openmp_threads=1, + resolution=resolution, + forcing_type=forcing_type, + coord_type=coord_type, + drag_type=drag_type, + damping_coeff=1.0e-4, + baroclinic=False, + method=method, + ) + self.add_step(forward_step) + self.add_step( + Viz( + component=component, + indir=subdir, + damping_coeffs=self.damping_coeffs, + baroclinic=False, + forcing_type=forcing_type, + ) + ) + + # def validate(self): + # """ + # Validate variables against a baseline + # """ + # damping_coeffs = self.damping_coeffs + # variables = ['layerThickness', 'normalVelocity'] + # if damping_coeffs is not None: + # for damping_coeff in damping_coeffs: + # compare_variables(test_case=self, variables=variables, + # filename1=f'forward_{damping_coeff}/' + # 'output.nc') + # else: + # compare_variables(test_case=self, variables=variables, + # filename1='forward/output.nc') diff --git a/polaris/tasks/ocean/drying_slope/convergence/__init__.py b/polaris/tasks/ocean/drying_slope/convergence/__init__.py new file mode 100644 index 0000000000..16bfd40106 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/convergence/__init__.py @@ -0,0 +1,125 @@ +from typing import Dict + +from polaris import Step, Task +from polaris.ocean.convergence import get_resolution_for_task +from polaris.resolution import resolution_to_string +from polaris.tasks.ocean.drying_slope.convergence.analysis import Analysis +from polaris.tasks.ocean.drying_slope.convergence.forward import Forward +from polaris.tasks.ocean.drying_slope.init import Init + + +class Convergence(Task): + """ + The convergence drying_slope test case + + Attributes + ---------- + resolutions : list of float + The resolution of the test case in km + + damping_coeffs : list of float + The damping coefficients at which to evaluate convergence. Must be of + length 1. + """ + + def __init__( + self, + component, + subdir, + group_dir, + config, + coord_type='sigma', + method='ramp', + ): + """ + Create the test case + + Parameters + ---------- + component : polaris.ocean.Ocean + The ocean component that this task belongs to + + init : polaris.tasks.ocean.drying_slope.init.Init + A shared step for creating the initial state + + subdir : str + The subdirectory to put the task in + + group_dir : str + The subdirectory to put the task group in + + config : polaris.config.PolarisConfigParser + A shared config parser + + method: str, optional + The wetting-and-drying method (``standard``, ``ramp``) + coord_type : str, optional + The type of vertical coordinate (``sigma``, ``single_layer``, etc.) + + """ + name = f'convergence_{method}' + config_filename = 'drying_slope.cfg' + self.config = config + super().__init__(component=component, name=name, subdir=subdir) + self.set_shared_config(config, link=config_filename) + + self.damping_coeffs = [] + analysis_dependencies: Dict[str, Dict[float, Step]] = dict( + mesh=dict(), init=dict(), forward=dict() + ) + refinement_factors = config.getlist( + 'convergence', 'refinement_factors_space', dtype=float + ) + for refinement_factor in refinement_factors: + resolution = get_resolution_for_task( + config, refinement_factor, refinement='both' + ) + mesh_name = resolution_to_string(resolution) + init_dir = f'{group_dir}/barotropic/{mesh_name}/init' + if init_dir in component.steps: + init_step = component.steps[init_dir] + symlink = f'init_{mesh_name}' + else: + init_step = Init( + component=component, + resolution=resolution, + name=f'init_{mesh_name}', + subdir=init_dir, + ) + init_step.set_shared_config(config, link=config_filename) + symlink = None + self.add_step(init_step, symlink=symlink) + + damping_coeff = 0.01 + step_name = f'forward_{damping_coeff:03g}' + forward_dir = ( + f'{group_dir}/barotropic/{mesh_name}/{method}/{step_name}' + ) + symlink = f'{step_name}_{mesh_name}' + if forward_dir in component.steps: + forward_step = component.steps[forward_dir] + else: + forward_step = Forward( + component=component, + name=f'{step_name}_{mesh_name}', + refinement_factor=refinement_factor, + subdir=forward_dir, + init=init_step, + damping_coeff=damping_coeff, + coord_type=coord_type, + method=method, + ) + forward_step.set_shared_config(config, link=config_filename) + self.add_step(forward_step, symlink=symlink) + analysis_dependencies['mesh'][resolution] = init_step + analysis_dependencies['init'][resolution] = init_step + analysis_dependencies['forward'][resolution] = forward_step + + self.add_step( + Analysis( + component=component, + subdir=f'{subdir}/analysis', + damping_coeff=damping_coeff, + dependencies=analysis_dependencies, + ) + ) diff --git a/polaris/tasks/ocean/drying_slope/convergence/analysis.py b/polaris/tasks/ocean/drying_slope/convergence/analysis.py new file mode 100644 index 0000000000..4023e702c9 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/convergence/analysis.py @@ -0,0 +1,118 @@ +import numpy as np +import pandas as pd +import xarray as xr +from scipy.interpolate import interp1d + +from polaris.ocean.convergence.analysis import ConvergenceAnalysis + + +class Analysis(ConvergenceAnalysis): + """ + A step for analyzing the convergence of drying slope results and producing + a convergence plot. + + Attributes + ---------- + damping_coeff : float + The Rayleigh damping coefficient used for the forward runs + """ + + def __init__(self, component, subdir, dependencies, damping_coeff): + """ + Create the step + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + subdir : str + The subdirectory that the step resides in + + dependencies : dict of dict of polaris.Steps + The dependencies of this step + + damping_coeff: float + the value of the rayleigh damping coefficient + """ + self.damping_coeff = damping_coeff + convergence_vars = [{'name': 'ssh', 'title': 'SSH', 'zidx': None}] + super().__init__( + component=component, + subdir=subdir, + dependencies=dependencies, + convergence_vars=convergence_vars, + ) + + # We won't use all of these files but we link them all just in case + # the user changes convergence_eval_time + for time in ['0.05', '0.15', '0.25', '0.30', '0.40', '0.50']: + filename = f'r{damping_coeff}d{time}-analytical.csv' + self.add_input_file( + filename=filename, target=filename, database='drying_slope' + ) + + def exact_solution(self, mesh_name, field_name, time, zidx=None): + """ + Get the exact solution + + Parameters + ---------- + mesh_name : str + The mesh name which is the prefix for the initial condition file + + field_name : str + The name of the variable of which we evaluate convergence + For the default method, we use the same convergence rate for all + fields + + time : float + The time at which to evaluate the exact solution in seconds. + For the default method, we always use the initial state. + + zidx : int, optional + The z-index for the vertical level at which to evaluate the exact + solution + + Returns + ------- + solution : xarray.DataArray + The exact solution as derived from the initial condition + """ + if field_name != 'ssh': + raise ValueError(f'{field_name} is not currently supported') + + # Get the MPAS cell locations + init = xr.open_dataset(f'{mesh_name}_init.nc') + y_min = init.yCell.min() + x_offset = y_min.values / 1000.0 + init = init.drop_vars( + np.setdiff1d([j for j in init.variables], ['yCell', 'ssh']) + ) + + init = init.isel(Time=0) + x_mpas = init.yCell / 1000.0 + + # Load the analytical solution + # we need to convert time from seconds to days for filename + day = time / (3600.0 * 24.0) + datafile = f'./r{self.damping_coeff}d{day:.2f}-analytical.csv' + data = pd.read_csv(datafile, header=None) + x_exact = data[0] + x_offset + ssh_exact = -data[1] + + # Set MPAS locations out of analytical bounds to nans + x_min = np.min(x_exact) + x_max = np.max(x_exact) + x_mpas[x_mpas < x_min] = np.nan + x_mpas[x_mpas > x_max] = np.nan + + # In the original version we interpolated mpas data to exact data + # location + # here we do the opposite because we don't want to have to get the + # exact data locations from within the shared convergence step + # we need to interpolate to the mpas mesh locations + f = interp1d(x_exact, ssh_exact) + ssh_exact_interp = f(x_mpas) + + return ssh_exact_interp diff --git a/polaris/tasks/ocean/drying_slope/convergence/forward.py b/polaris/tasks/ocean/drying_slope/convergence/forward.py new file mode 100644 index 0000000000..cb7d7c3b53 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/convergence/forward.py @@ -0,0 +1,165 @@ +import time + +from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.convergence import get_resolution_for_task +from polaris.ocean.convergence.forward import ConvergenceForward + + +class Forward(ConvergenceForward): + """ + A step for performing forward ocean component runs as part of drying slope + test cases. + + Attributes + ---------- + resolution : float + The resolution of the test case in km + """ + + def __init__( + self, + component, + name, + refinement_factor, + subdir, + init, + damping_coeff, + coord_type, + method, + refinement='both', + drag_type='constant_and_rayleigh', + ): + """ + Create a new test case + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + subdir : str + The subdirectory that the step belongs to + + init : polaris.Step + The step which generates the mesh and initial condition + + damping_coeff : float + the vertical coordinate type + + coord_type : str + the vertical coordinate type + + method : str, optional + The type of wetting and drying algorithm to use + + drag_type : str, optional + The bottom drag type to apply as a namelist option + """ + options = dict() + if coord_type == 'single_layer': + options['config_disable_thick_sflux'] = '.true.' + options['config_disable_vel_hmix'] = '.true.' + + if method == 'ramp': + options['config_zero_drying_velocity_ramp'] = '.true.' + + options['config_implicit_bottom_drag_type'] = drag_type + # for drag types not specified here, defaults are used or given in + # forward.yaml + if drag_type == 'constant': + options['config_implicit_constant_bottom_drag_coeff'] = '3.0e-3' + elif drag_type == 'constant_and_rayleigh': + # update the damping coefficient to the requested value *after* + # loading forward.yaml + options['config_Rayleigh_damping_coeff'] = damping_coeff + + options['config_tidal_forcing_model'] = 'monochromatic' + super().__init__( + component=component, + name=name, + subdir=subdir, + refinement_factor=refinement_factor, + mesh=init, + init=init, + package='polaris.ocean.tasks.drying_slope', + yaml_filename='forward.yaml', + graph_target=f'{init.path}/culled_graph.info', + output_filename='output.nc', + forcing=True, + options=options, + validate_vars=['layerThickness', 'normalVelocity'], + ) + + def setup(self): + """ + TEMP: symlink initial condition to name hard-coded in Omega + """ + super().setup() + config = self.config + self.resolution = get_resolution_for_task( + config, self.refinement_factor, refinement=self.refinement + ) + model = 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') + + 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['drying_slope_barotropic'] + 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 + + def dynamic_model_config(self, at_setup): + """ + Set the model time step from config options at setup and 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=at_setup) + options = dict() + config = self.config + section = config['drying_slope'] + time_integrator = section.get('time_integrator') + options['config_time_integrator'] = time_integrator + + if time_integrator == 'RK4': + dt_per_km = section.getfloat('rk4_dt_per_km') + if time_integrator == 'split_explicit': + dt_per_km = section.getfloat('split_dt_per_km') + # btr_dt is also proportional to resolution + btr_dt_per_km = config.getfloat('drying_slope', 'btr_dt_per_km') + btr_dt = btr_dt_per_km * self.resolution + options['config_btr_dt'] = time.strftime( + '%H:%M:%S', time.gmtime(btr_dt) + ) + self.btr_dt = btr_dt + dt = dt_per_km * self.resolution + # https://stackoverflow.com/a/1384565/7728169 + options['config_dt'] = time.strftime('%H:%M:%S', time.gmtime(dt)) + self.dt = dt + + section = self.config['drying_slope_barotropic'] + thin_film_thickness = section.getfloat('thin_film_thickness') + + options['config_drying_min_cell_height'] = thin_film_thickness + options['config_zero_drying_velocity_ramp_hmin'] = thin_film_thickness + options['config_zero_drying_velocity_ramp_hmax'] = ( + thin_film_thickness * 10.0 + ) + self.add_model_config_options(options=options) diff --git a/polaris/tasks/ocean/drying_slope/decomp/__init__.py b/polaris/tasks/ocean/drying_slope/decomp/__init__.py new file mode 100644 index 0000000000..1eb09ae010 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/decomp/__init__.py @@ -0,0 +1,74 @@ +from polaris import Task +from polaris.tasks.ocean.drying_slope.forward import Forward +from polaris.tasks.ocean.drying_slope.validate import Validate + + +class Decomp(Task): + """ + A drying slope decomposition task, which makes sure the model + produces identical results on 1 and 4 cores. + """ + + def __init__( + self, + component, + resolution, + indir, + init, + coord_type='sigma', + method='ramp', + ): + """ + Create the task + + Parameters + ---------- + component : polaris.ocean.Ocean + The ocean component that this task belongs to + + resolution : float + The resolution of the task in km + + indir : str + The directory the task is in, to which ``name`` will be appended + + init : polaris.tasks.ocean.drying_slope.init.Init + A shared step for creating the initial state + + coord_type : str, optional + The vertical coordinate type + + method : str, optional + The type of wetting and drying algorithm to use + """ + + super().__init__(component=component, name='decomp', indir=indir) + + self.add_step(init, symlink='init') + + subdirs = list() + for procs in [4, 8]: + name = f'{procs}proc' + + self.add_step( + Forward( + component=component, + init=init, + name=name, + indir=self.subdir, + ntasks=procs, + min_tasks=procs, + openmp_threads=1, + resolution=resolution, + run_time_steps=3, + damping_coeff=0.001, + coord_type=coord_type, + method=method, + ) + ) + subdirs.append(name) + self.add_step( + Validate( + component=component, step_subdirs=subdirs, indir=self.subdir + ) + ) diff --git a/polaris/tasks/ocean/drying_slope/drying_slope.cfg b/polaris/tasks/ocean/drying_slope/drying_slope.cfg new file mode 100644 index 0000000000..be6341cb7a --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/drying_slope.cfg @@ -0,0 +1,166 @@ +# Options related to the vertical grid +[vertical_grid] + +# the type of vertical grid +grid_type = uniform + +# The type of vertical coordinate (e.g. z-level, z-star) +coord_type = sigma + +# Number of vertical levels +vert_levels = 10 + +# 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 + +# The minimum number of vertical levels for z-star coordinate +min_vert_levels = 3 + +# Minimum thickness of each layer for z-star coordinate +min_layer_thickness = ${drying_slope_barotropic:thin_film_thickness} + +# config options for all drying slope test cases +[drying_slope] + +# time step per resolution (s/km), since dt is proportional to resolution +rk4_dt_per_km = 30 + +split_dt_per_km = 30 + +# barotropic time step per resolution (s/km), since btr_dt is proportional to +# resolution +btr_dt_per_km = 1.5 + +# Coriolis parameter +coriolis_parameter = 0.0 + +# Background temperature +background_temperature = 20.0 + +# config options for barotropic drying slope test cases +[drying_slope_barotropic] + +# time integration scheme +time_integrator = RK4 + +# the width of the domain in km +lx = 6. + +# Length over which wetting and drying actually occur in km +ly_analysis = 25. + +# the length of the domain in km +ly = 28. + +# Bottom depth at the right side of the domain in m +right_bottom_depth = -11. + +# Bottom depth at the left side of the domain in m +left_bottom_depth = -1. + +# Initial SSH at the right side of the domain in m +right_tidal_height = ${drying_slope_barotropic:right_bottom_depth} + +# Plug width as a fraction of the domain +plug_width_frac = 0.0 + +# Plug temperature +plug_temperature = 20.0 + +# Background salinity +background_salinity = 35.0 + +# Thickness of each layer in the thin film region +thin_film_thickness = 1.0e-3 + +# config options for barotropic drying slope test cases +[drying_slope_baroclinic] + +# time integration scheme +time_integrator = split_explicit + +# the width of the domain in km +lx = 12. + +# the length of the domain in km +ly = 55. + +# Length over which wetting and drying actually occur in km +ly_analysis = 50. + +# Bottom depth at the right side of the domain in m +right_bottom_depth = -3.0 + +# Bottom depth at the left side of the domain in m +left_bottom_depth = -0.5 + +# Initial SSH at the right side of the domain in m +right_tidal_height = -0.5 + +# salinity at the right side of the domain +right_salinity = 35.0 + +# salinity at the left side of the domain +left_salinity = 1.0 + +# manning coefficient used in mannings bottom drag type +manning_coefficient = 5.0e-2 + +# The column thickness in the thin film region, used to determine thin film +# thickness +min_column_thickness = 0.05 + +[convergence] + +# Evaluation time for convergence analysis (in hours) +convergence_eval_time = 12.0 + +# Convergence threshold below which a test fails +convergence_thresh = -0.527 + +# Type of error to compute +error_type = l2 + +# the base mesh resolution (km) to which refinement_factors +# are applied if refinement is 'space' or 'both' on a planar mesh +base_resolution = 0.5 + +# refinement factors for a planar mesh applied to either space or time +refinement_factors_space = 4., 2., 1., 0.5 + +# config options for convergence forward steps +[convergence_forward] + +# time integrator: {'split_explicit', 'RK4'} +time_integrator = ${drying_slope:time_integrator} + +# RK4 time step per resolution (s/km), since dt is proportional to resolution +rk4_dt_per_km = ${drying_slope:rk4_dt_per_km} + +# split time step per resolution (s/km), since dt is proportional to resolution +split_dt_per_km = ${drying_slope:split_dt_per_km} + +# the barotropic time step (s/km) for simulations using split time stepping, +# since btr_dt is proportional to resolution +btr_dt_per_km = ${drying_slope:btr_dt_per_km} + +# Run duration in hours +run_duration = ${convergence:convergence_eval_time} + +# Output interval in hours +output_interval = ${run_duration} + +# config options for visualizing drying slope ouptut +[drying_slope_viz] + +# whether to generate movie +generate_movie = False + +# frames per second for movies +frames_per_second = 30 + +# movie format +movie_format = mp4 diff --git a/polaris/tasks/ocean/drying_slope/forward.py b/polaris/tasks/ocean/drying_slope/forward.py new file mode 100644 index 0000000000..fab2ea59b9 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/forward.py @@ -0,0 +1,290 @@ +import time + +from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.model import OceanModelStep + + +class Forward(OceanModelStep): + """ + A step for performing forward ocean component runs as part of drying + slope tasks. + + Attributes + ---------- + resolution : float + The resolution of the task in km + + baroclinic : logical + Whether this test case is the baroclinic version + + 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 + """ + + def __init__( + self, + component, + resolution, + init, + name='forward', + subdir=None, + indir=None, + ntasks=None, + min_tasks=None, + openmp_threads=1, + damping_coeff=None, + coord_type='sigma', + forcing_type='tidal_cycle', + drag_type='constant_and_rayleigh', + baroclinic=False, + method='ramp', + run_time_steps=None, + graph_target='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 + + name : str + the name of the task + + 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 + + coord_type : str, optional + The vertical coordinate type + + forcing_type : str, optional + The forcing type to apply at the "tidal" boundary as a namelist + option + + method : str, optional + The type of wetting and drying algorithm to use + + drag_type : str, optional + The bottom drag type to apply as a namelist option + + time_integrator : str, optional + The time integration scheme to apply as a namelist option + + run_time_steps : int or None + Number of time steps to run for + """ + if drag_type == 'constant_and_rayleigh': + if coord_type == 'single_layer': + raise ValueError( + f'Drag type {drag_type} is not supported ' + f'with coordinate type {coord_type}' + ) + if damping_coeff is None: + raise ValueError( + 'Damping coefficient must be specified with ' + f'drag type {drag_type}' + ) + + self.damping_coeff = damping_coeff + self.drag_type = drag_type + self.baroclinic = baroclinic + self.coord_type = coord_type + self.forcing_type = forcing_type + self.method = method + self.resolution = resolution + self.run_time_steps = run_time_steps + self.yaml_filename = 'forward.yaml' + 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.add_yaml_file('polaris.ocean.config', 'output.yaml') + if self.coord_type == 'single_layer': + self.add_yaml_file('polaris.ocean.config', 'single_layer.yaml') + if self.baroclinic: + self.add_yaml_file( + 'polaris.tasks.ocean.drying_slope', 'baroclinic.yaml' + ) + self.add_yaml_file( + 'polaris.tasks.ocean.drying_slope', self.yaml_filename + ) + + self.add_input_file( + filename='initial_state.nc', + work_dir_target=f'{init.path}/initial_state.nc', + ) + self.add_input_file( + filename='forcing.nc', work_dir_target=f'{init.path}/forcing.nc' + ) + + self.add_output_file( + filename='output.nc', + validate_vars=[ + 'temperature', + 'salinity', + 'layerThickness', + 'normalVelocity', + ], + ) + + self.dt = None + self.btr_dt = None + + 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 + """ + if self.baroclinic: + section = self.config['drying_slope_baroclinic'] + else: + section = self.config['drying_slope_barotropic'] + 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 + + 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) + + if self.baroclinic: + section = self.config['vertical_grid'] + vert_levels = section.getint('vert_levels') + section = self.config['drying_slope_baroclinic'] + time_integrator = section.get('time_integrator') + thin_film_thickness = ( + section.getfloat('min_column_thickness') / vert_levels + ) + else: + section = self.config['drying_slope_barotropic'] + time_integrator = section.get('time_integrator') + thin_film_thickness = section.getfloat('thin_film_thickness') + + # dt is proportional to resolution: default 30 seconds per km + section = self.config['drying_slope'] + if time_integrator == 'RK4': + dt_per_km = section.getfloat('rk4_dt_per_km') + elif time_integrator == 'split_explicit': + dt_per_km = section.getfloat('split_dt_per_km') + # btr_dt is also proportional to resolution + btr_dt_per_km = section.getfloat('btr_dt_per_km') + btr_dt = btr_dt_per_km * self.resolution + self.btr_dt = btr_dt + else: + print(f'Time integrator {time_integrator} not supported') + btr_dt_str = time.strftime('%H:%M:%S', time.gmtime(self.btr_dt)) + + dt = dt_per_km * self.resolution + # https://stackoverflow.com/a/1384565/7728169 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + self.dt = dt + + if self.run_time_steps is not None: + run_duration_str = time.strftime( + '%H:%M:%S', time.gmtime(dt * self.run_time_steps) + ) + else: + run_duration_str = '0000_12:00:01' + + replacements = dict( + time_integrator=time_integrator, + dt=dt_str, + btr_dt=btr_dt_str, + run_duration=run_duration_str, + hmin=f'{thin_film_thickness}', + ramp_hmin=f'{thin_film_thickness}', + ramp_hmax=f'{thin_film_thickness * 10.0}', + ) + + mpas_options = dict() + if self.method == 'ramp': + mpas_options['config_zero_drying_velocity_ramp'] = True + + mpas_options['config_implicit_bottom_drag_type'] = self.drag_type + # for drag types not specified here, defaults are used or given in + # forward.yaml + if self.drag_type == 'constant': + mpas_options['config_implicit_constant_bottom_drag_coeff'] = 3.0e-3 # type: ignore[assignment] + elif self.drag_type == 'constant_and_rayleigh': + # update the damping coefficient to the requested value *after* + # loading forward.yaml + mpas_options['config_Rayleigh_damping_coeff'] = self.damping_coeff + + forcing_dict = { + 'tidal_cycle': 'monochromatic', + 'linear_drying': 'linear', + } + + mpas_options['config_tidal_forcing_model'] = forcing_dict[ + self.forcing_type + ] # type: ignore[assignment] + if self.forcing_type == 'linear_drying': + if self.baroclinic: + section = self.config['drying_slope_baroclinic'] + else: + section = self.config['drying_slope_barotropic'] + replacements['tidal_min'] = ( + section.getfloat('right_bottom_depth') + 0.5 + ) + replacements['tidal_baseline'] = section.getfloat( + 'right_tidal_height' + ) + + self.add_model_config_options( + options=mpas_options, config_model='mpas-ocean' + ) + self.add_yaml_file( + 'polaris.tasks.ocean.drying_slope', + self.yaml_filename, + template_replacements=replacements, + ) diff --git a/polaris/tasks/ocean/drying_slope/forward.yaml b/polaris/tasks/ocean/drying_slope/forward.yaml new file mode 100644 index 0000000000..5c0d891926 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/forward.yaml @@ -0,0 +1,74 @@ +ocean: + time_management: + config_run_duration: {{ run_duration }} + time_integration: + config_dt: {{ dt }} + config_time_integrator: {{ time_integrator }} + +mpas-ocean: + io: + config_write_output_on_startup: false + split_explicit_ts: + config_btr_dt: {{ btr_dt }} + cvmix: + config_use_cvmix: false + tidal_forcing: + config_use_tidal_forcing: true + config_use_tidal_forcing_tau: 60.0 + config_tidal_forcing_type: direct + config_tidal_forcing_monochromatic_amp: 10.0 + config_tidal_forcing_monochromatic_period: 1.0 + config_tidal_forcing_monochromatic_baseline: 10.0 + config_tidal_forcing_linear_min: {{ tidal_min }} + config_tidal_forcing_linear_rate: -8.0 + config_tidal_forcing_linear_baseline: {{ tidal_baseline }} + Rayleigh_damping: + config_Rayleigh_damping_depth_variable: true + wetting_drying: + config_use_wetting_drying: true + config_prevent_drying: true + config_zero_drying_velocity: true + config_verify_not_dry: true + config_zero_drying_velocity_ramp_hmin: 1e-3 + config_zero_drying_velocity_ramp_hmax: 1e-2 + config_zero_drying_velocity_ramp_factor: 1.0 + ALE_vertical_grid: + config_vert_coord_movement: uniform_stretching + config_ALE_thickness_proportionality: restingThickness_times_weights + advection: + config_thickness_flux_type: upwind + config_vert_advection_method: flux-form + debug: + config_disable_thick_sflux: true + config_disable_vel_hmix: true + streams: + mesh: + filename_template: initial_state.nc + input: + filename_template: initial_state.nc + restart: {} + forcing: + type: input + filename_template: forcing.nc + input_interval: initial_only + contents: + - tidalInputMask + output: + type: output + filename_template: output.nc + output_interval: 0000_00:00:00.1 + clobber_mode: truncate + contents: + - mesh + - tracers + - xtime + - daysSinceStartOfSim + - normalVelocity + - velocityX + - velocityY + - layerThickness + - ssh + - wettingVelocityFactor + - wettingVelocityBaroclinic + - wettingVelocityBarotropicSubcycle +# - upwindFactor diff --git a/polaris/tasks/ocean/drying_slope/init.py b/polaris/tasks/ocean/drying_slope/init.py new file mode 100644 index 0000000000..3457bbfbec --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/init.py @@ -0,0 +1,280 @@ +import cmocean # noqa: F401 +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.ocean.viz.transect import compute_transect, plot_transect +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 +from polaris.viz import plot_horiz_field + + +class Init(Step): + """ + A step for creating a mesh and initial condition for baroclinic channel + tasks + + Attributes + ---------- + resolution : float + The resolution of the task in km + + baroclinic: bool, optional + whether the test case is baroclinic + + drag_type : str, optional + The bottom drag type to apply as a namelist option + """ + + def __init__( + self, + component, + resolution, + indir=None, + subdir=None, + name='init', + baroclinic=False, + drag_type='constant_and_rayleigh', + ): + """ + Create the step + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + resolution : float + The resolution of the task in km + + indir : str + the directory the step is in, to which ``name`` will be appended + + name : str, optional + the name of the step + + baroclinic: bool, optional + whether the test case is baroclinic + + drag_type : str, optional + The bottom drag type to apply as a namelist option + """ + super().__init__( + component=component, name=name, indir=indir, subdir=subdir + ) + self.resolution = resolution + self.baroclinic = baroclinic + self.drag_type = drag_type + + 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 + resolution = self.resolution + + section = config['drying_slope'] + coriolis_parameter = section.getfloat('coriolis_parameter') + background_temperature = section.getfloat('background_temperature') + vert_levels = config.getint('vertical_grid', 'vert_levels') + + if self.baroclinic: + section = config['drying_slope_baroclinic'] + right_salinity = section.getfloat('right_salinity') + left_salinity = section.getfloat('left_salinity') + manning_coefficient = section.getfloat('manning_coefficient') + thin_film_thickness = ( + section.getfloat('min_column_thickness') / vert_levels + 1.0e-8 + ) + else: + section = config['drying_slope_barotropic'] + plug_width_frac = section.getfloat('plug_width_frac') + plug_temperature = section.getfloat('plug_temperature') + background_salinity = section.getfloat('background_salinity') + thin_film_thickness = ( + section.getfloat('thin_film_thickness') + 1.0e-8 + ) + + # config options used in both configurations but which have different + # values in each + drying_length = section.getfloat('ly_analysis') * 1e3 + lx = section.getfloat('lx') + ly = section.getfloat('ly') + right_bottom_depth = section.getfloat('right_bottom_depth') + left_bottom_depth = section.getfloat('left_bottom_depth') + right_tidal_height = section.getfloat('right_tidal_height') + + domain_length = ly * 1e3 + # Check config options + if domain_length < drying_length: + raise ValueError( + 'Domain is not long enough to capture wetting front' + ) + if right_bottom_depth > left_bottom_depth: + raise ValueError( + 'Right boundary must be deeper than left boundary' + ) + + nx, ny = compute_planar_hex_nx_ny(lx, ly, resolution) + 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() + ds_forcing = ds_mesh.copy() + + y_cell = ds.yCell + y_min = y_cell.min() + y_max = y_cell.max() + dc_edge_min = ds.dcEdge.min() + + bottom_depth = -right_bottom_depth - ( + y_max - y_cell + ) / drying_length * (-right_bottom_depth + left_bottom_depth) + ds['bottomDepth'] = bottom_depth + + # SSH is constant except when it would result in a water column less + # than the minimum thickness + ds['ssh'] = np.maximum( + right_tidal_height, + -bottom_depth + thin_film_thickness * vert_levels, + ) + + init_vertical_coord(config, ds) + + if self.baroclinic: + temperature = background_temperature * xr.ones_like(y_cell) + else: + plug_width = domain_length * plug_width_frac + y_plug_boundary = y_min + plug_width + temperature = xr.where( + y_cell < y_plug_boundary, + plug_temperature, + background_temperature, + ) + temperature, _ = xr.broadcast(temperature, ds.refBottomDepth) + ds['temperature'] = temperature.expand_dims(dim='Time', axis=0) + if self.baroclinic: + salinity = right_salinity - (y_max - y_cell) / drying_length * ( + right_salinity - left_salinity + ) + # Use a debug tracer for validation + ds['tracer1'] = xr.ones_like(ds.temperature) + else: + salinity = background_salinity * xr.ones_like(y_cell) + salinity, _ = xr.broadcast(salinity, ds.refBottomDepth) + ds['salinity'] = salinity.expand_dims(dim='Time', axis=0) + + normalVelocity = xr.zeros_like(ds_mesh.xEdge) + normalVelocity, _ = xr.broadcast(normalVelocity, ds.refBottomDepth) + normalVelocity = normalVelocity.transpose('nEdges', 'nVertLevels') + ds['normalVelocity'] = normalVelocity.expand_dims(dim='Time', axis=0) + ds['fCell'] = coriolis_parameter * xr.ones_like(ds.xCell) + ds['fEdge'] = coriolis_parameter * xr.ones_like(ds.xEdge) + ds['fVertex'] = coriolis_parameter * xr.ones_like(ds.xVertex) + + write_netcdf(ds, 'initial_state.nc') + + # Define the tidal boundary condition over 1-cell width + y_tidal_boundary = y_max - dc_edge_min / 2.0 + tidal_forcing_mask = xr.where(y_cell > y_tidal_boundary, 1.0, 0.0) + if tidal_forcing_mask.sum() <= 0: + raise ValueError('Input mask for tidal case is not set!') + ds_forcing['tidalInputMask'] = tidal_forcing_mask + if self.baroclinic and self.drag_type == 'mannings': + ds_forcing['bottomDrag'] = manning_coefficient * xr.ones_like( + tidal_forcing_mask + ) + write_netcdf(ds_forcing, 'forcing.nc') + + x_mid = ds_mesh.xCell.median() + y_min = ds_mesh.yCell.min() + y_max = ds_mesh.yCell.max() + x = xr.DataArray(data=[x_mid, x_mid], dims=('nPoints',)) + y = xr.DataArray(data=[y_min, y_max], dims=('nPoints',)) + ds_transect = compute_transect( + x=x, + y=y, + ds_horiz_mesh=ds_mesh, + layer_thickness=ds.layerThickness.isel(Time=0), + bottom_depth=ds.bottomDepth, + min_level_cell=ds.minLevelCell - 1, + max_level_cell=ds.maxLevelCell - 1, + spherical=False, + ) + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.layerThickness.isel(Time=0), + out_filename='layerThickness_depth_init.png', + title='layer thickness', + outline_color=None, + ssh_color='blue', + seafloor_color='black', + interface_color='grey', + colorbar_label=r'm', + cmap='cmo.thermal', + ) + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.salinity.isel(Time=0), + out_filename='salinity_depth_init.png', + title='salinity', + colorbar_label='PSU', + cmap='cmo.haline', + ) + + cell_mask = ds.maxLevelCell >= 1 + + cellsOnEdge1 = ds_mesh.cellsOnEdge.isel(TWO=0) + cellsOnEdge2 = ds_mesh.cellsOnEdge.isel(TWO=1) + cell1_is_valid = cell_mask[cellsOnEdge1 - 1].values == 1 + cell2_is_valid = cell_mask[cellsOnEdge2 - 1].values == 1 + edge_mask = xr.where( + np.logical_and(cell1_is_valid, cell2_is_valid), 1, 0 + ) + + plot_horiz_field( + ds_mesh, + ds.salinity, + cmap_title='S', + out_file_name='initial_salinity.png', + field_mask=cell_mask, + show_patch_edges=True, + transect_x=x, + transect_y=y, + ) + plot_horiz_field( + ds_mesh, + ds.normalVelocity, + out_file_name='initial_velocity.png', + field_mask=edge_mask, + show_patch_edges=True, + cmap_title='u_{normal}', + transect_x=x, + transect_y=y, + ) diff --git a/polaris/tasks/ocean/drying_slope/validate.py b/polaris/tasks/ocean/drying_slope/validate.py new file mode 100644 index 0000000000..c41462d4f2 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/validate.py @@ -0,0 +1,62 @@ +from polaris import Step +from polaris.validate import compare_variables + + +class Validate(Step): + """ + A step for comparing outputs between steps in a drying slope run + + Attributes + ---------- + step_subdirs : list of str + The number of processors used in each run + """ + + def __init__(self, component, step_subdirs, indir): + """ + Create the step + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + step_subdirs : list of str + The number of processors used in each run + + indir : str + the directory the step is in, to which ``name`` will be appended + """ + super().__init__(component=component, name='validate', indir=indir) + + self.step_subdirs = step_subdirs + + for subdir in step_subdirs: + self.add_input_file( + filename=f'output_{subdir}.nc', target=f'../{subdir}/output.nc' + ) + + def run(self): + """ + Compare ``temperature``, ``salinity``, ``layerThickness`` and + ``normalVelocity`` in the outputs of two previous steps with each other + """ + super().run() + step_subdirs = self.step_subdirs + variables = [ + 'temperature', + 'salinity', + 'layerThickness', + 'normalVelocity', + ] + all_pass = compare_variables( + variables=variables, + filename1=self.inputs[0], + filename2=self.inputs[1], + logger=self.logger, + ) + if not all_pass: + raise ValueError( + f'Validation failed comparing outputs between ' + f'{step_subdirs[0]} and {step_subdirs[1]}.' + ) diff --git a/polaris/tasks/ocean/drying_slope/viz.py b/polaris/tasks/ocean/drying_slope/viz.py new file mode 100644 index 0000000000..50f2de4a98 --- /dev/null +++ b/polaris/tasks/ocean/drying_slope/viz.py @@ -0,0 +1,516 @@ +import os + +import cmocean # noqa: F401 +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from mpas_tools.ocean.viz.transect import compute_transect, plot_transect + +from polaris import Step +from polaris.viz import plot_horiz_field +from polaris.viz.style import use_mplstyle + + +class Viz(Step): + """ + A step for plotting the results of a series of drying_slope runs + + Attributes + ---------- + damping_coeffs : list of float + The damping coefficients that correspond to each forward run to be + plotted + + coord_type : str, optional + The vertical coordinate type + + forcing_type : str, optional + The forcing type to apply at the "tidal" boundary as a namelist + option + + baroclinic : logical + Whether this test case is the baroclinic version + + times : list of float + The solution times at which to plot + + datatypes : list of str + The datatypes to plot + """ + + def __init__( + self, + component, + indir=None, + subdir=None, + name='viz', + damping_coeffs=None, + coord_type='sigma', + baroclinic=False, + forcing_type='tidal_cycle', + ): + """ + 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 + + damping_coeffs : list of float + The damping coefficients that correspond to each forward run to be + plotted + + coord_type : str, optional + The vertical coordinate type + + forcing_type : str, optional + The forcing type to apply at the "tidal" boundary as a namelist + option + + baroclinic : logical + Whether this test case is the baroclinic version + """ + super().__init__( + component=component, name=name, indir=indir, subdir=subdir + ) + + self.damping_coeffs = damping_coeffs + self.coord_type = coord_type + self.forcing_type = forcing_type + self.baroclinic = baroclinic + self.times = ['0.05', '0.15', '0.25', '0.30', '0.40', '0.50'] + self.datatypes = ['analytical', 'ROMS'] + + 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='forcing.nc', target='../init/forcing.nc') + if damping_coeffs is None: + self.add_input_file( + filename='output.nc', target='../forward/output.nc' + ) + else: + for damping_coeff in damping_coeffs: + self.add_input_file( + filename=f'output_{damping_coeff:03g}.nc', + target=f'../forward_{damping_coeff:03g}/output.nc', + ) + for time in self.times: + for datatype in self.datatypes: + for damping_coeff in damping_coeffs: + filename = ( + f'r{damping_coeff}d{time}-{datatype.lower()}.csv' + ) + self.add_input_file( + filename=filename, + target=filename, + database='drying_slope', + ) + + def run(self): + """ + Run this step of the task + """ + use_mplstyle() + ds_forcing = xr.open_dataset('forcing.nc') + self._plot_ssh_time_series( + ds_forcing=ds_forcing, forcing_type=self.forcing_type + ) + + if not self.baroclinic: + self._plot_ssh_validation() + + ds_mesh = xr.load_dataset('init.nc') + cell_mask = ds_mesh.maxLevelCell >= 1 + cellsOnEdge1 = ds_mesh.cellsOnEdge.isel(TWO=0) + cellsOnEdge2 = ds_mesh.cellsOnEdge.isel(TWO=1) + cell1_is_valid = cell_mask[cellsOnEdge1 - 1].values == 1 + cell2_is_valid = cell_mask[cellsOnEdge2 - 1].values == 1 + edge_mask = xr.where( + np.logical_and(cell1_is_valid, cell2_is_valid), 1, 0 + ) + + vert_levels = self.config.getint('vertical_grid', 'vert_levels') + hmin = self.config.getfloat( + 'drying_slope_barotropic', 'thin_film_thickness' + ) + hmax = vert_levels * hmin * 10.0 + out_filenames = [] + if not self.damping_coeffs: + out_filenames.append('output.nc') + else: + for damping_coeff in self.damping_coeffs: + out_filenames.append(f'output_{damping_coeff:04g}.nc') + + for out_filename in out_filenames: + ds = xr.load_dataset(out_filename) + x, y = self._plot_transects(ds_mesh=ds_mesh, ds=ds) + + nt = int(ds.sizes['Time']) + for tidx in np.arange(0, nt, nt / 10, dtype=int): + # for tidx in np.arange(-10,0,1): + # for tidx in [-2]: + # for atime in self.times: + # # Plot MPAS-O data + # mpastime = ds.daysSinceStartOfSim.values + # simtime = pd.to_timedelta(mpastime) + # s_day = 86400.0 + # time = simtime.total_seconds() + # tidx = np.argmin(np.abs(time / s_day - float(atime))) + for zidx in [0, 9]: + plot_horiz_field( + ds_mesh, + ds.wettingVelocityFactor, + cmap_title=r'$\Phi$', + out_file_name=f'wetting_factor_t{tidx:04g}_z{zidx:02g}.png', + transect_x=x, + transect_y=y, + show_patch_edges=True, + t_index=tidx, + z_index=zidx, + field_mask=edge_mask, + vmin=0, + vmax=1, + ) + + if 'upwindFactor' in ds.keys(): + plot_horiz_field( + ds_mesh, + ds.upwindFactor, + cmap_title=r'$\Psi$', + out_file_name=f'upwind_factor_t{tidx:04g}.png', + transect_x=x, + transect_y=y, + show_patch_edges=False, + t_index=tidx, + field_mask=edge_mask, + vmin=0, + vmax=1, + ) + + if self.baroclinic: + if 'wettingVelocityBaroclinic' in ds.keys(): + plot_horiz_field( + ds_mesh, + ds.wettingVelocityBaroclinic, + cmap_title=r'$\Phi_{bcl}$', + out_file_name=f'wetting_baroclinic_t{tidx:04g}.png', + t_index=tidx, + field_mask=edge_mask, + vmin=0, + vmax=1, + ) + if 'wettingVelocityBarotropicSubcycle' in ds.keys(): + plot_horiz_field( + ds_mesh, + ds.wettingVelocityBarotropicSubcycle, + cmap_title=r'$\Phi_{btr}$', + out_file_name=f'wetting_barotropic_t{tidx:04g}.png', + t_index=tidx, + field_mask=edge_mask, + vmin=0, + vmax=1, + ) + self._plot_salinity(tidx=-1, y_distance=45.0) + + plot_horiz_field( + ds_mesh, + ds.ssh + ds_mesh.bottomDepth, + out_file_name=f'wct_horiz_t{tidx:04g}.png', + t_index=tidx, + field_mask=cell_mask, + vmin=hmax, + vmax=2.0, + cmap='viridis', + cmap_title='H', + cmap_set_under='r', + ) + + def _plot_transects(self, ds_mesh, ds): + # Note: capability currently only works for cell-centered quantities + + x_mid = ds_mesh.xCell.median() + y_min = ds_mesh.yCell.min() + y_max = ds_mesh.yCell.max() + x = xr.DataArray(data=[x_mid, x_mid], dims=('nPoints',)) + y = xr.DataArray(data=[y_min, y_max], dims=('nPoints',)) + + ymin = -ds_mesh.bottomDepth.max() + ymax = ds.ssh.max() + vmin_layer_thickness = max(ds.layerThickness.min().values, 0.01) + vmax_layer_thickness = min(ds.layerThickness.max().values, 1.0) + vmax_velocity = min(np.nanmax(np.abs(ds.velocityY.values)), 0.5) + nrows = 2 + if self.baroclinic: + nrows += 1 + section = self.config['drying_slope_baroclinic'] + vmin_salinity = section.getfloat('left_salinity') + vmax_salinity = section.getfloat('right_salinity') + mpastime = ds.daysSinceStartOfSim.values + simtime = pd.to_timedelta(mpastime) + time_hours = simtime.total_seconds() / 3600.0 + nt = int(ds.sizes['Time']) + for tidx in np.arange(0, nt, nt / 10, dtype=int): + # for tidx in range(ds.sizes['Time']): + if ( + np.nanmax(np.abs(ds.layerThickness.isel(Time=tidx).values)) + > 1.0e3 + ): + print(f'layerThickness corrupted at time {tidx}') + ds_transect = compute_transect( + x=x, + y=y, + ds_horiz_mesh=ds_mesh, + layer_thickness=ds.layerThickness.isel(Time=tidx), + bottom_depth=ds.bottomDepth, + min_level_cell=ds.minLevelCell - 1, + max_level_cell=ds.maxLevelCell - 1, + spherical=False, + ) + + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.velocityX.isel(Time=tidx), + title='Across-slope velocity', + out_filename=f'velocityX_depth_t{tidx:04g}.png', + vmin=-vmax_velocity, + vmax=vmax_velocity, + colorbar_label=r'', + cmap='cmo.balance', + ) + + fig, axs = plt.subplots(nrows, sharex=True, figsize=(5, 3 * nrows)) + fig.suptitle(f'Time: {time_hours[tidx]:2.1f} hours') + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.layerThickness.isel(Time=tidx), + ax=axs[0], + outline_color=None, + ssh_color='blue', + seafloor_color='black', + interface_color='grey', + vmin=vmin_layer_thickness, + vmax=vmax_layer_thickness, + colorbar_label=r'm', + cmap='cmo.thermal', + ) + axs[0].set_title('Layer thickness') + axs[0].set_ylim([ymin, ymax]) + axs[0].set_xlabel(None) + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.velocityY.isel(Time=tidx), + ax=axs[1], + outline_color=None, + ssh_color='blue', + seafloor_color='black', + vmin=-vmax_velocity, + vmax=vmax_velocity, + colorbar_label='m/s', + cmap='cmo.balance', + ) + axs[1].set_title('Along-slope velocity') + axs[1].set_ylim([ymin, ymax]) + axs[1].set_xlabel(None) + if self.baroclinic: + plot_transect( + ds_transect=ds_transect, + mpas_field=ds.salinity.isel(Time=tidx), + ax=axs[2], + outline_color=None, + ssh_color='blue', + seafloor_color='black', + vmin=vmin_salinity, + vmax=vmax_salinity, + colorbar_label=r'PSU', + cmap='cmo.haline', + ) + axs[2].set_title('Salinity') + axs[2].set_ylim([ymin, ymax]) + plt.savefig( + f'layerThickness_velocityY_depth_t{tidx:04g}.png', + bbox_inches='tight', + ) + plt.close() + + return x, y + + def _forcing(self, t): + ssh = 10.0 * np.sin(t * np.pi / 12.0) - 10.0 + return ssh + + def _plot_salinity(self, tidx=None, y_distance=0.0, outFolder='.'): + """ + Plot salinity at a point location as a function of vertical levels + y_distance distance in meters along y-axis + """ + ds = xr.open_dataset('output.nc') + y_cell = ds.yCell + cell_idx = np.argmin(y_cell.values - y_distance / 1e3) + salinity = ds['salinity'][tidx, cell_idx, :] + fig = plt.figure() + vert_levels = ds.dims['nVertLevels'] + plt.plot(salinity, np.arange(vert_levels), '.-') + plt.xlabel('Salinity') + plt.ylabel('Vertical level') + fig.savefig('salinity_levels.png', bbox_inches='tight', dpi=200) + plt.close(fig) + + def _plot_ssh_time_series( + self, ds_forcing, outFolder='.', forcing_type='tidal_cycle' + ): + """ + Plot ssh forcing on the right x boundary as a function of time against + the analytical solution. The agreement should be within machine + precision if the namelist options are consistent with the Warner et al. + (2013) test case. + """ + figsize = [6.4, 4.8] + + damping_coeffs = self.damping_coeffs + if not damping_coeffs: + naxes = 1 + ncFilename = ['output.nc'] + else: + naxes = len(damping_coeffs) + ncFilename = [ + f'output_{damping_coeff}.nc' + for damping_coeff in damping_coeffs + ] + fig, _ = plt.subplots(nrows=naxes, ncols=1, figsize=figsize, dpi=100) + + mask = ds_forcing.tidalInputMask + for i in range(naxes): + ax = plt.subplot(naxes, 1, i + 1) + ds = xr.open_dataset(ncFilename[i]) + ympas = ds.ssh.where(mask).mean('nCells').values + xmpas = np.linspace(0, 1.0, len(ds.xtime)) * 12.0 + ax.plot( + xmpas, ympas, marker='o', label='MPAS-O forward', color='k' + ) + if forcing_type == 'tidal_cycle': + xSsh = np.linspace(0, 12.0, 100) + ySsh = self._forcing(t=xSsh) + ax.plot(xSsh, ySsh, lw=3, label='analytical', color='b') + ax.set_ylabel('Tidal amplitude (m)') + ax.set_xlabel('Time (hrs)') + ax.legend(frameon=False) + ax.label_outer() + ds.close() + + fig.suptitle('Tidal amplitude forcing (right side)') + fig.savefig(f'{outFolder}/ssh_t.png', bbox_inches='tight', dpi=200) + + plt.close(fig) + + def _plot_ssh_validation(self, tidx=None, outFolder='.'): + """ + Plot ssh as a function of along-channel distance for all times for + which there is validation data + """ + datatypes = ['analytical', 'ROMS'] + colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'} + + locs = [7.2, 2.2, 0.2, 1.2, 4.2, 9.3] + locs = 0.92 - np.divide(locs, 11.0) + + damping_coeffs = self.damping_coeffs + + if damping_coeffs is None: + raise ValueError( + 'ssh validation is only supported for damping' + 'coefficient comparison' + ) + naxes = len(damping_coeffs) + nhandles = len(datatypes) + 1 + + ds_mesh = xr.open_dataset('init.nc') + # Note: capability currently only works for cell-centered quantities + + x_mid = ds_mesh.xCell.median() + y_min = ds_mesh.yCell.min() + y_max = ds_mesh.yCell.max() + x = xr.DataArray(data=[x_mid, x_mid], dims=('nPoints',)) + y = xr.DataArray(data=[y_min, y_max], dims=('nPoints',)) + ymin = -ds_mesh.bottomDepth.max() + + section = self.config['drying_slope_barotropic'] + drying_length = section.getfloat('ly_analysis') + # we need to add x_offset to observational datasets + x_offset = y_min.values / 1000.0 + s_day = 86400.0 + + for _, atime in enumerate(self.times): + fig, axs = plt.subplots( + nrows=naxes, ncols=1, sharex=True, figsize=(5, 3 * naxes) + ) + fig.suptitle(f'Time: {float(atime) * 24.0:2.1f} hours') + for i, damping_coeff in enumerate(self.damping_coeffs): + ds = xr.load_dataset(f'output_{damping_coeff:03g}.nc') + # Plot MPAS-O data + mpastime = ds.daysSinceStartOfSim.values + simtime = pd.to_timedelta(mpastime) + time = simtime.total_seconds() + tidx = np.argmin(np.abs(time / s_day - float(atime))) + ds_transect = compute_transect( + x=x, + y=y, + ds_horiz_mesh=ds_mesh, + layer_thickness=ds.layerThickness.isel(Time=tidx), + bottom_depth=ds.bottomDepth, + min_level_cell=ds.minLevelCell - 1, + max_level_cell=ds.maxLevelCell - 1, + spherical=False, + ) + plot_transect( + ds_transect=ds_transect, + mpas_field=None, + ax=axs[i], + outline_color=None, + ssh_color='blue', + seafloor_color='black', + ) + ymax = ds.ssh.max() + axs[i].set_xlim([x_offset, drying_length + x_offset]) + axs[i].set_ylim([ymin.values, ymax.values]) + if i == naxes - 1: + axs[i].set_xlabel('Along channel distance (km)') + else: + axs[i].set_xlabel(None) + + axs[i].set_title('r = ' + str(damping_coeffs[i])) + for datatype in datatypes: + datafile = ( + f'./r{damping_coeffs[i]}d{atime}-' + f'{datatype.lower()}.csv' + ) + if os.path.exists(datafile): + data = pd.read_csv(datafile, header=None) + axs[i].scatter( + data[0] + x_offset, + -data[1], + marker='.', + color=colors[datatype], + label=datatype, + ) + ds.close() + + h, l0 = axs[i].get_legend_handles_labels() + axs[i].legend( + h[:nhandles], + l0[:nhandles], + frameon=False, + loc='lower left', + ) + + filename = f'{outFolder}/ssh_depth_section_t{tidx:04d}' + fig.savefig(f'{filename}.png', dpi=200, format='png') + plt.close(fig)