diff --git a/polaris/ocean/__init__.py b/polaris/ocean/__init__.py index 1ca8e7361e..1cd7683f14 100644 --- a/polaris/ocean/__init__.py +++ b/polaris/ocean/__init__.py @@ -1,5 +1,6 @@ from polaris import Component from polaris.ocean.tests.baroclinic_channel import BaroclinicChannel +from polaris.ocean.tests.galewsky_jet import GalewskyJet from polaris.ocean.tests.global_convergence import GlobalConvergence from polaris.ocean.tests.single_column import SingleColumn @@ -17,6 +18,7 @@ def __init__(self): # please keep these in alphabetical order self.add_test_group(BaroclinicChannel(component=self)) + self.add_test_group(GalewskyJet(component=self)) self.add_test_group(GlobalConvergence(component=self)) self.add_test_group(SingleColumn(component=self)) diff --git a/polaris/ocean/tests/galewsky_jet/__init__.py b/polaris/ocean/tests/galewsky_jet/__init__.py new file mode 100644 index 0000000000..e1dbf98ac8 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/__init__.py @@ -0,0 +1,24 @@ +from polaris import TestGroup +from polaris.ocean.tests.galewsky_jet.test_balance import TestBalance +from polaris.ocean.tests.galewsky_jet.short_adjustment import ShortAdjustment +from polaris.ocean.tests.galewsky_jet.barotropic_instability import BarotropicInstability + + +class GalewskyJet(TestGroup): + """ + A test group for "galewsky jet" test cases + """ + def __init__(self, component): + """ + component : polaris.ocean.Ocean + the ocean component that this test group belongs to + """ + super().__init__(component=component, name='galewsky_jet') + + for resolution in [120.]: + self.add_test_case( + TestBalance(test_group=self, resolution=resolution)) + self.add_test_case( + ShortAdjustment(test_group=self, resolution=resolution)) + self.add_test_case( + BarotropicInstability(test_group=self, resolution=resolution)) diff --git a/polaris/ocean/tests/galewsky_jet/barotropic_instability/__init__.py b/polaris/ocean/tests/galewsky_jet/barotropic_instability/__init__.py new file mode 100644 index 0000000000..ac2374b0c1 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/barotropic_instability/__init__.py @@ -0,0 +1,65 @@ +from polaris.mesh.spherical import IcosahedralMeshStep +from polaris.ocean.tests.galewsky_jet.barotropic_instability.forward import ( + Forward, +) +from polaris.ocean.tests.galewsky_jet.barotropic_instability.initial_state import ( + InitialState, +) +from polaris.ocean.tests.galewsky_jet.barotropic_instability.viz import Viz +from polaris.testcase import TestCase + + +class BarotropicInstability(TestCase): + """ + A class to define the Galewsky jet barotropic instability + test case (6 days for instbaility development) + + Attributes + ---------- + resolution : str + The resolution of the test case + """ + + def __init__(self, test_group, resolution): + """ + Create the test case + + Parameters + ---------- + test_group : + polaris.ocean.tests.galewsky_jet.GalewskyJet + The test group that this test case belongs to + + resolution : float + The resolution of the test case (in km) + """ + name = 'barotropic_instability' + self.resolution = resolution + res = int(resolution) + subdir = f'{res}km/{name}' + super().__init__(test_group=test_group, name=name, + subdir=subdir) + + self.add_step(IcosahedralMeshStep( + test_case=self, cell_width=resolution)) + self.add_step( + InitialState(test_case=self, resolution=resolution)) + self.add_step( + Forward(test_case=self, resolution=resolution)) + + mesh_name = f'Icos{res}' + name = f'{mesh_name}_viz' + subdir = 'viz' + self.add_step(Viz(test_case=self, name=name, subdir=subdir, + mesh_name=mesh_name)) + + def configure(self): + """ + Set config options for the test case + """ + super().configure() + config = self.config + config.add_from_package('polaris.mesh', 'mesh.cfg') + + config.set('spherical_mesh', 'mpas_mesh_filename', + 'mesh.nc') diff --git a/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.py b/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.py new file mode 100644 index 0000000000..e8fed19db3 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.py @@ -0,0 +1,99 @@ +import time + +import xarray as xr + +from polaris.ocean.model import OceanModelStep + + +class Forward(OceanModelStep): + """ + A step for performing forward ocean component runs as part of + the galewsky jet - barotropic instability test case + + Attributes + ---------- + resolution : int + The resolution of the (uniform) mesh in km + """ + + def __init__(self, test_case, resolution): + """ + Create a new step + + Parameters + ---------- + test_case : polaris.ocean.tests.galewsky_jet.barotropic_instability.BarotropicInstability # noqa: E501 + The test case this step belongs to + + resolution : int + The resolution of the (uniform) mesh in km + """ + super().__init__(test_case=test_case, + name='forward', + subdir='forward', + openmp_threads=1) + + self.resolution = resolution + + # make sure output is double precision + self.add_yaml_file('polaris.ocean.config', 'output.yaml') + + self.add_yaml_file( + 'polaris.ocean.tests.galewsky_jet.barotropic_instability', + 'forward.yaml') + + self.add_input_file(filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file(filename='graph.info', + target='../base_mesh/graph.info') + + self.add_output_file(filename='output.nc') + + def compute_cell_count(self, at_setup): + """ + Compute the approximate number of cells in the mesh, used to constrain + resources + + Parameters + ---------- + at_setup : bool + Whether this method is being run during setup of the step, as + opposed to at runtime + + Returns + ------- + cell_count : int or None + The approximate number of cells in the mesh + """ + if at_setup: + # use a heuristic based on QU30 (65275 cells) and QU240 (10383 + # cells) + cell_count = 6e8 / self.resolution**2 + else: + # get nCells from the input file + with xr.open_dataset('initial_state.nc') as ds: + cell_count = ds.sizes['nCells'] + 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) + + config = self.config + # dt is proportional to resolution: default 30 seconds per km + dt_per_km = config.getfloat('galewsky_jet', 'dt_per_km') + + dt = dt_per_km * self.resolution + # https://stackoverflow.com/a/1384565/7728169 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + + options = dict(config_dt=dt_str) + self.add_model_config_options(options) diff --git a/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.yaml b/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.yaml new file mode 100644 index 0000000000..4418e7a1af --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.yaml @@ -0,0 +1,75 @@ +omega: + run_modes: + config_ocean_run_mode: forward + time_management: + config_run_duration: 0006_00:00:00 + ALE_vertical_grid: + config_vert_coord_movement: impermeable_interfaces + config_ALE_thickness_proportionality: weights_only + decomposition: + config_block_decomp_file_prefix: graph.info.part. + time_integration: + config_time_integrator: RK4 + debug: + config_disable_thick_all_tend: false + config_disable_thick_hadv: false + config_disable_thick_vadv: true + config_disable_thick_sflux: true + config_disable_vel_all_tend: false + config_disable_vel_coriolis: false + config_disable_vel_pgrad: false + config_disable_vel_hmix: false + config_disable_vel_surface_stress: true + config_disable_vel_explicit_bottom_drag: true + config_disable_vel_vmix: true + config_disable_vel_vadv: true + config_disable_tr_all_tend: true + config_disable_tr_hmix: true + config_disable_tr_vmix: true + config_disable_tr_sflux: true + config_disable_tr_nonlocalflux: true + config_check_ssh_consistency: false + cvmix: + config_use_cvmix: false + eos: + config_eos_type: linear + forcing: + config_use_bulk_wind_stress: false + config_use_bulk_thickness_flux: false + pressure_gradient: + config_pressure_gradient_type: ssh_gradient + tracer_forcing_activeTracers: + config_use_activeTracers_surface_restoring: false + config_use_activeTracers_interior_restoring: false + config_use_activeTracers: true + config_use_activeTracers_surface_bulk_forcing: false + tracer_forcing_debugTracers: + config_use_debugTracers: false + AM_mixedLayerDepths: + config_AM_mixedLayerDepths_enable: false + streams: + mesh: + filename_template: initial_state.nc + input: + filename_template: initial_state.nc + restart: + output_interval: 0000_01:00:00 + output: + type: output + filename_template: output.nc + output_interval: 0001_00:00:00 + clobber_mode: truncate + reference_time: 0001-01-01_00:00:00 + contents: + - mesh + - xtime + - normalVelocity + - layerThickness + - refZMid + - refLayerThickness + - ssh + - kineticEnergyCell + - relativeVorticityCell + - relativeVorticity + - normalizedPlanetaryVorticityEdge + - normalizedRelativeVorticityEdge diff --git a/polaris/ocean/tests/galewsky_jet/barotropic_instability/initial_state.py b/polaris/ocean/tests/galewsky_jet/barotropic_instability/initial_state.py new file mode 100644 index 0000000000..8764b1ded5 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/barotropic_instability/initial_state.py @@ -0,0 +1,99 @@ +import cmocean # noqa: F401 +import xarray as xr +from mpas_tools.io import write_netcdf +from mpas_tools.mesh.jet import init as jet_init + +from polaris import Step +from polaris.ocean.vertical import init_vertical_coord + + +class InitialState(Step): + """ + A step for creating a mesh and initial condition + for galewsky jet barotropic instability + test case + + Attributes + ---------- + resolution : float + The resolution of the test case in km + """ + def __init__(self, test_case, resolution): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + resolution : float + The resolution of the test case in km + """ + super().__init__(test_case=test_case, name='initial_state') + self.resolution = resolution + + self.add_input_file( + filename='mesh.nc', + target='../base_mesh/mesh.nc') + self.add_input_file( + filename='graph.info', + target='../base_mesh/graph.info') + # self.add_input_file( + # filename='init.nc') + self.add_output_file('initial_state.nc') + + def run(self): + """ + Run this step of the test case + """ + config = self.config +# update rsph from cime.constants + # jet_init(name='mesh.nc', save='velocity_ic.nc', + # rsph=6371220.0, pert=False) + # ds2 = xr.open_dataset('velocity_ic.nc') + + dsMesh = xr.open_dataset('mesh.nc') + + ds = dsMesh.copy() + x_cell = ds.xCell + + bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') + + ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell) + ds['ssh'] = xr.zeros_like(x_cell) + # ds['ssh'] = ds2.h - ds['bottomDepth'] + + init_vertical_coord(config, ds) + + # resolution = self.resolution + section = config['galewsky_jet'] + temperature = section.getfloat('temperature') + salinity = section.getfloat('salinity') + # coriolis_parameter = section.getfloat('coriolis_parameter') + + temperature_array = temperature * xr.ones_like(x_cell) + temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid) + ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0) + ds['salinity'] = salinity * xr.ones_like(ds.temperature) + + # ds2 = xr.open_dataset('init.nc') + jet_init(name='mesh.nc', save='velocity_ic.nc', + rsph=6371220.0, pert=True) + ds2 = xr.open_dataset('velocity_ic.nc') + + unrm_array, _ = xr.broadcast(ds2.u, ds.refZMid) + ds['normalVelocity'] = unrm_array + # ds['normalVelocity'] = ds2.u + h_array, _ = xr.broadcast(ds2.h, ds.refZMid) + ds['layerThickness'] = h_array + ds['fCell'] = ds2.fCell + ds['fEdge'] = ds2.fEdge + ds['fVertex'] = ds2.fVertex + ds['ssh'] = ds2.h - ds['bottomDepth'] + + # if (config.getfloat('vertical_grid', 'grid_type') == 'uniform'): + # nlev = config.getfloat('vertical_grid', 'vert_levels') + # ds['layerThickness'], _ = xr.broadcast(ds2.h / nlev, ds.refZMid) + + write_netcdf(ds, 'initial_state.nc') diff --git a/polaris/ocean/tests/galewsky_jet/barotropic_instability/viz.py b/polaris/ocean/tests/galewsky_jet/barotropic_instability/viz.py new file mode 100644 index 0000000000..ec054ca6fd --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/barotropic_instability/viz.py @@ -0,0 +1,74 @@ +import cmocean # noqa: F401 +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from scipy.interpolate import griddata + +from polaris import Step + + +class Viz(Step): + """ + A step for plotting fields from the galewsky jet output + + Attributes + ---------- + mesh_name : str + The name of the mesh + """ + def __init__(self, test_case, name, subdir, mesh_name): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + name : str + The name of the step + + subdir : str + The subdirectory in the test case's work directory for the step + + mesh_name : str + The name of the mesh + """ + super().__init__(test_case=test_case, name=name, subdir=subdir) + self.add_input_file( + filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file( + filename='output.nc', + target='../forward/output.nc') + self.mesh_name = mesh_name + self.add_output_file('vorticity.png') + + def run(self): + """ + Run this step of the test case + """ + plt.figure(1, figsize=(12.0, 6.0)) + + ncfileIC = Dataset('initial_state.nc', 'r') + ncfile = Dataset('output.nc', 'r') + var1 = ncfile.variables['normalizedRelativeVorticityEdge'][6, :, 0] + var2 = ncfile.variables['normalizedPlanetaryVorticityEdge'][6, :, 0] + var = var1 + var2 + latEdge = ncfileIC.variables['latEdge'][:] + lonEdge = ncfileIC.variables['lonEdge'][:] + latEdge = latEdge * 180 / np.pi + lonEdge = lonEdge * 180 / np.pi - 180 + xi = np.linspace(-180, 180, 720) + yi = np.linspace(0, 80, 180) + X, Y = np.meshgrid(xi, yi) + Z = griddata((lonEdge, latEdge), var, (X, Y), method='linear') + Z = np.clip(Z, 0, 2.6e-8) + plt.figure(figsize=(9.6, 4.8)) + plt.contourf(X, Y, Z, 200) + plt.xlabel('longitude') + plt.ylabel('latitude') + plt.axis('scaled') + plt.colorbar + + plt.savefig('vorticity.png') diff --git a/polaris/ocean/tests/galewsky_jet/galewsky_jet.cfg b/polaris/ocean/tests/galewsky_jet/galewsky_jet.cfg new file mode 100644 index 0000000000..aa7b4504a0 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/galewsky_jet.cfg @@ -0,0 +1,81 @@ +# Options related to the vertical grid +[vertical_grid] + +# the type of vertical grid +grid_type = uniform + +# Number of vertical levels +vert_levels = 1 + +# Depth of the bottom of the ocean +bottom_depth = 10000.0 + +# The type of vertical coordinate (e.g. z-level, z-star) +coord_type = z-star + +# 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 +[galewsky_jet] + +# time step per resolution (s/km), since dt is proportional to resolution +dt_per_km = 0.25 + +# reference temperature +temperature = 15.0 + +# reference salinity +salinity = 35.0 + +# options for visualization for the layer thickness of galewsky jet +[galewsky_jet_vizmap] + +# visualization latitude and longitude resolution +dlon = 0.5 +dlat = 0.5 + +# remapping method ('bilinear', 'neareststod', 'conserve') +remap_method = conserve + +[galewsky_jet_viz_thickness] +# colormap options +# colormap +colormap_name = viridis + +# the type of norm used in the colormap +norm_type = linear + +# A dictionary with keywords for the norm +norm_args = {'vmin': 9000., 'vmax': 10200.} + +# We could provide colorbar tick marks but we'll leave the defaults +# colorbar_ticks = np.linspace(0., 1., 9) + +[galewsky_jet_viz_kineticenergy] +# colormap options +# colormap +colormap_name = viridis + +# the type of norm used in the colormap +norm_type = linear + +# A dictionary with keywords for the norm +norm_args = {'vmin': 0., 'vmax': None} + +# We could provide colorbar tick marks but we'll leave the defaults +# colorbar_ticks = np.linspace(0., 1., 9) + +[galewsky_jet_viz_kineticenergy_diff] +# colormap options +# colormap +colormap_name = viridis + +# the type of norm used in the colormap +norm_type = linear + +# A dictionary with keywords for the norm +norm_args = {'vmin': -300., 'vmax': 300.} diff --git a/polaris/ocean/tests/galewsky_jet/short_adjustment/__init__.py b/polaris/ocean/tests/galewsky_jet/short_adjustment/__init__.py new file mode 100644 index 0000000000..01baa95f78 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/short_adjustment/__init__.py @@ -0,0 +1,68 @@ +from polaris.mesh.spherical import IcosahedralMeshStep +from polaris.ocean.tests.galewsky_jet.short_adjustment.forward import Forward +from polaris.ocean.tests.galewsky_jet.short_adjustment.initial_state import ( + InitialState, +) +from polaris.ocean.tests.galewsky_jet.short_adjustment.viz import Viz +from polaris.testcase import TestCase + + +class ShortAdjustment(TestCase): + """ + A class to define the Galewsky jet short adjustment test case + (6 hrs for gravity wave adjustment) + + Attributes + ---------- + resolution : str + The resolution of the test case + """ + + def __init__(self, test_group, resolution): + """ + Create the test case + + Parameters + ---------- + test_group : + polaris.ocean.tests.galewsky_jet.GalewskyJet + The test group that this test case belongs to + + resolution : float + The resolution of the test case (in km) + """ + name = 'short_adjustment' + self.resolution = resolution + res = int(resolution) + subdir = f'{res}km/{name}' + super().__init__(test_group=test_group, name=name, + subdir=subdir) + + self.add_step(IcosahedralMeshStep( + test_case=self, cell_width=resolution)) + self.add_step( + InitialState(test_case=self, resolution=resolution)) + self.add_step( + Forward(test_case=self, resolution=resolution)) + + mesh_name = f'Icos{res}' + # name = f'{mesh_name}_map' + # subdir = 'map' + # viz_map = VizMap(test_case=self, name=name, subdir=subdir, + # mesh_name=mesh_name) + # self.add_step(viz_map) + + name = f'{mesh_name}_viz' + subdir = 'viz' + self.add_step(Viz(test_case=self, name=name, subdir=subdir, + mesh_name=mesh_name)) + + def configure(self): + """ + Set config options for the test case + """ + super().configure() + config = self.config + config.add_from_package('polaris.mesh', 'mesh.cfg') + + config.set('spherical_mesh', 'mpas_mesh_filename', 'mesh.nc') diff --git a/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.py b/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.py new file mode 100644 index 0000000000..15d7d2502a --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.py @@ -0,0 +1,99 @@ +import time + +import xarray as xr + +from polaris.ocean.model import OceanModelStep + + +class Forward(OceanModelStep): + """ + A step for performing forward ocean component runs as part of + the galewsky jet - short adjustment test case + + Attributes + ---------- + resolution : int + The resolution of the (uniform) mesh in km + """ + + def __init__(self, test_case, resolution): + """ + Create a new step + + Parameters + ---------- + test_case : polaris.ocean.tests.galewsky_jet.short_adjustment.ShortAdjustment # noqa: E501 + The test case this step belongs to + + resolution : int + The resolution of the (uniform) mesh in km + """ + super().__init__(test_case=test_case, + name='forward', + subdir='forward', + openmp_threads=1) + + self.resolution = resolution + + # make sure output is double precision + self.add_yaml_file('polaris.ocean.config', 'output.yaml') + + self.add_yaml_file( + 'polaris.ocean.tests.galewsky_jet.short_adjustment', + 'forward.yaml') + + self.add_input_file(filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file(filename='graph.info', + target='../base_mesh/graph.info') + + self.add_output_file(filename='output.nc') + + def compute_cell_count(self, at_setup): + """ + Compute the approximate number of cells in the mesh, used to constrain + resources + + Parameters + ---------- + at_setup : bool + Whether this method is being run during setup of the step, as + opposed to at runtime + + Returns + ------- + cell_count : int or None + The approximate number of cells in the mesh + """ + if at_setup: + # use a heuristic based on QU30 (65275 cells) and QU240 (10383 + # cells) + cell_count = 6e8 / self.resolution**2 + else: + # get nCells from the input file + with xr.open_dataset('initial_state.nc') as ds: + cell_count = ds.sizes['nCells'] + 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) + + config = self.config + # dt is proportional to resolution: default 30 seconds per km + dt_per_km = config.getfloat('galewsky_jet', 'dt_per_km') + + dt = dt_per_km * self.resolution + # https://stackoverflow.com/a/1384565/7728169 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + + options = dict(config_dt=dt_str) + self.add_model_config_options(options) diff --git a/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.yaml b/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.yaml new file mode 100644 index 0000000000..153ef2ab79 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/short_adjustment/forward.yaml @@ -0,0 +1,75 @@ +omega: + run_modes: + config_ocean_run_mode: forward + time_management: + config_run_duration: 0000_06:00:00 + ALE_vertical_grid: + config_vert_coord_movement: impermeable_interfaces + config_ALE_thickness_proportionality: weights_only + decomposition: + config_block_decomp_file_prefix: graph.info.part. + time_integration: + config_time_integrator: RK4 + debug: + config_disable_thick_all_tend: false + config_disable_thick_hadv: false + config_disable_thick_vadv: true + config_disable_thick_sflux: true + config_disable_vel_all_tend: false + config_disable_vel_coriolis: false + config_disable_vel_pgrad: false + config_disable_vel_hmix: false + config_disable_vel_surface_stress: true + config_disable_vel_explicit_bottom_drag: true + config_disable_vel_vmix: true + config_disable_vel_vadv: true + config_disable_tr_all_tend: true + config_disable_tr_hmix: true + config_disable_tr_vmix: true + config_disable_tr_sflux: true + config_disable_tr_nonlocalflux: true + config_check_ssh_consistency: false + cvmix: + config_use_cvmix: false + eos: + config_eos_type: linear + forcing: + config_use_bulk_wind_stress: false + config_use_bulk_thickness_flux: false + pressure_gradient: + config_pressure_gradient_type: ssh_gradient + tracer_forcing_activeTracers: + config_use_activeTracers_surface_restoring: false + config_use_activeTracers_interior_restoring: false + config_use_activeTracers: true + config_use_activeTracers_surface_bulk_forcing: false + tracer_forcing_debugTracers: + config_use_debugTracers: false + AM_mixedLayerDepths: + config_AM_mixedLayerDepths_enable: false + streams: + mesh: + filename_template: initial_state.nc + input: + filename_template: initial_state.nc + restart: + output_interval: 0000_01:00:00 + output: + type: output + filename_template: output.nc + output_interval: 0000_00:10:00 + clobber_mode: truncate + reference_time: 0001-01-01_00:00:00 + contents: + - mesh + - xtime + - normalVelocity + - layerThickness + - refZMid + - refLayerThickness + - ssh + - kineticEnergyCell + - relativeVorticityCell + - relativeVorticity + - normalizedPlanetaryVorticityEdge + - normalizedRelativeVorticityEdge diff --git a/polaris/ocean/tests/galewsky_jet/short_adjustment/initial_state.py b/polaris/ocean/tests/galewsky_jet/short_adjustment/initial_state.py new file mode 100644 index 0000000000..09d7caf3ff --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/short_adjustment/initial_state.py @@ -0,0 +1,99 @@ +import cmocean # noqa: F401 +import xarray as xr +from mpas_tools.io import write_netcdf +from mpas_tools.mesh.jet import init as jet_init + +from polaris import Step +from polaris.ocean.vertical import init_vertical_coord + + +class InitialState(Step): + """ + A step for creating a mesh and initial condition + for galewsky jet - short adjustment + test case + + Attributes + ---------- + resolution : float + The resolution of the test case in km + """ + def __init__(self, test_case, resolution): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + resolution : float + The resolution of the test case in km + """ + super().__init__(test_case=test_case, name='initial_state') + self.resolution = resolution + + self.add_input_file( + filename='mesh.nc', + target='../base_mesh/mesh.nc') + self.add_input_file( + filename='graph.info', + target='../base_mesh/graph.info') + # self.add_input_file( + # filename='init.nc') + self.add_output_file('initial_state.nc') + + def run(self): + """ + Run this step of the test case + """ + config = self.config +# update rsph from cime.constants + # jet_init(name='mesh.nc', save='velocity_ic.nc', + # rsph=6371220.0, pert=False) + # ds2 = xr.open_dataset('velocity_ic.nc') + + dsMesh = xr.open_dataset('mesh.nc') + + ds = dsMesh.copy() + x_cell = ds.xCell + + bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') + + ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell) + ds['ssh'] = xr.zeros_like(x_cell) + # ds['ssh'] = ds2.h - ds['bottomDepth'] + + init_vertical_coord(config, ds) + + # resolution = self.resolution + section = config['galewsky_jet'] + temperature = section.getfloat('temperature') + salinity = section.getfloat('salinity') + # coriolis_parameter = section.getfloat('coriolis_parameter') + + temperature_array = temperature * xr.ones_like(x_cell) + temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid) + ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0) + ds['salinity'] = salinity * xr.ones_like(ds.temperature) + + # ds2 = xr.open_dataset('init.nc') + jet_init(name='mesh.nc', save='velocity_ic.nc', + rsph=6371220.0, pert=True) + ds2 = xr.open_dataset('velocity_ic.nc') + + unrm_array, _ = xr.broadcast(ds2.u, ds.refZMid) + ds['normalVelocity'] = unrm_array + # ds['normalVelocity'] = ds2.u + h_array, _ = xr.broadcast(ds2.h, ds.refZMid) + ds['layerThickness'] = h_array + ds['fCell'] = ds2.fCell + ds['fEdge'] = ds2.fEdge + ds['fVertex'] = ds2.fVertex + ds['ssh'] = ds2.h - ds['bottomDepth'] + + # if (config.getfloat('vertical_grid', 'grid_type') == 'uniform'): + # nlev = config.getfloat('vertical_grid', 'vert_levels') + # ds['layerThickness'], _ = xr.broadcast(ds2.h / nlev, ds.refZMid) + + write_netcdf(ds, 'initial_state.nc') diff --git a/polaris/ocean/tests/galewsky_jet/short_adjustment/viz.py b/polaris/ocean/tests/galewsky_jet/short_adjustment/viz.py new file mode 100644 index 0000000000..6ef13b4ebc --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/short_adjustment/viz.py @@ -0,0 +1,101 @@ +import cmocean # noqa: F401 +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from mpas_tools.io import write_netcdf +from mpas_tools.mesh.jet import init as jet_init +from netCDF4 import Dataset +from scipy.interpolate import griddata + +from polaris import Step +from polaris.ocean.vertical import init_vertical_coord + + +class Viz(Step): + """ + A step for plotting fields from the galewsky jet output + + Attributes + ---------- + mesh_name : str + The name of the mesh + """ + def __init__(self, test_case, name, subdir, mesh_name): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + name : str + The name of the step + + subdir : str + The subdirectory in the test case's work directory for the step + + viz_map : polaris.ocean.tests.galewsky_jet.test_balance.viz.VizMap + The step for creating a mapping files, also used to remap data + from the MPAS mesh to a lon-lat grid + + mesh_name : str + The name of the mesh + """ + super().__init__(test_case=test_case, name=name, subdir=subdir) + self.add_input_file( + filename='mesh.nc', + target='../base_mesh/mesh.nc') + self.add_input_file( + filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file( + filename='output.nc', + target='../forward/output.nc') + self.mesh_name = mesh_name + self.add_output_file('init_unperturbed.nc') + self.add_output_file('height.png') + + def run(self): + """ + Run this step of the test case + """ + config = self.config + ds = xr.open_dataset('mesh.nc') + x_cell = ds.xCell + bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') + ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell) + ds['ssh'] = xr.zeros_like(x_cell) + init_vertical_coord(config, ds) + jet_init(name='mesh.nc', save='velocity_ic.nc', + rsph=6371220.0, pert=False) + ds2 = xr.open_dataset('velocity_ic.nc') + unrm_array, _ = xr.broadcast(ds2.u, ds.refZMid) + ds['normalVelocity'] = unrm_array + h_array, _ = xr.broadcast(ds2.h, ds.refZMid) + ds['layerThickness'] = h_array + write_netcdf(ds, 'init_unperturbed.nc') + + ncfileIC = Dataset('initial_state.nc', 'r') + ncfileIC2 = Dataset('init_unperturbed.nc', 'r') + ncfile = Dataset('output.nc', 'r') + var1 = ncfileIC2.variables['layerThickness'][0, :, 0] + var2 = ncfile.variables['layerThickness'][36, :, 0] + var = var2 - var1 + latCell = ncfileIC.variables['latCell'][:] + lonCell = ncfileIC.variables['lonCell'][:] + latCell = latCell * 180 / np.pi + lonCell = lonCell * 180 / np.pi - 180 + xi = np.linspace(-180, 180, 720) + yi = np.linspace(-80, 80, 180) + X, Y = np.meshgrid(xi, yi) + Z = griddata((lonCell, latCell), var, (X, Y), method='linear') + # Z = np.clip(Z, 0, 2.6e-8) + plt.figure(figsize=(9.6, 4.8)) + plt.contour(X, Y, Z, 20) + plt.xlabel('longitude') + plt.ylabel('latitude') + plt.axis('scaled') + plt.colorbar + + plt.savefig('height.png') diff --git a/polaris/ocean/tests/galewsky_jet/test_balance/__init__.py b/polaris/ocean/tests/galewsky_jet/test_balance/__init__.py new file mode 100644 index 0000000000..af80627b5f --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/test_balance/__init__.py @@ -0,0 +1,67 @@ +from polaris.mesh.spherical import IcosahedralMeshStep +from polaris.ocean.tests.galewsky_jet.test_balance.forward import Forward +from polaris.ocean.tests.galewsky_jet.test_balance.initial_state import ( + InitialState, +) +from polaris.ocean.tests.galewsky_jet.test_balance.viz import Viz, VizMap +from polaris.testcase import TestCase + + +class TestBalance(TestCase): + """ + A class to define the Galewsky jet test balance test cases + + Attributes + ---------- + resolution : str + The resolution of the test case + """ + + def __init__(self, test_group, resolution): + """ + Create the test case + + Parameters + ---------- + test_group : + polaris.ocean.tests.galewsky_jet.GalewskyJet + The test group that this test case belongs to + + resolution : float + The resolution of the test case (in km) + """ + name = 'steady_state' + self.resolution = resolution + res = int(resolution) + subdir = f'{res}km/{name}' + super().__init__(test_group=test_group, name=name, + subdir=subdir) + + self.add_step(IcosahedralMeshStep( + test_case=self, cell_width=resolution)) + self.add_step( + InitialState(test_case=self, resolution=resolution)) + self.add_step( + Forward(test_case=self, resolution=resolution)) + + mesh_name = f'Icos{res}' + name = f'{mesh_name}_map' + subdir = 'map' + viz_map = VizMap(test_case=self, name=name, subdir=subdir, + mesh_name=mesh_name) + self.add_step(viz_map) + + name = f'{mesh_name}_viz' + subdir = 'viz' + self.add_step(Viz(test_case=self, name=name, subdir=subdir, + viz_map=viz_map, mesh_name=mesh_name)) + + def configure(self): + """ + Set config options for the test case + """ + super().configure() + config = self.config + config.add_from_package('polaris.mesh', 'mesh.cfg') + + config.set('spherical_mesh', 'mpas_mesh_filename', 'mesh.nc') diff --git a/polaris/ocean/tests/galewsky_jet/test_balance/forward.py b/polaris/ocean/tests/galewsky_jet/test_balance/forward.py new file mode 100644 index 0000000000..715edc70f3 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/test_balance/forward.py @@ -0,0 +1,99 @@ +import time + +import xarray as xr + +from polaris.ocean.model import OceanModelStep + + +class Forward(OceanModelStep): + """ + A step for performing forward ocean component runs as part of + the galewsky jet - balanced test case + + Attributes + ---------- + resolution : int + The resolution of the (uniform) mesh in km + """ + + def __init__(self, test_case, resolution): + """ + Create a new step + + Parameters + ---------- + test_case : polaris.ocean.tests.galewsky_jet.test_balance.TestBalance # noqa: E501 + The test case this step belongs to + + resolution : int + The resolution of the (uniform) mesh in km + """ + super().__init__(test_case=test_case, + name='forward', + subdir='forward', + openmp_threads=1) + + self.resolution = resolution + + # make sure output is double precision + self.add_yaml_file('polaris.ocean.config', 'output.yaml') + + self.add_yaml_file( + 'polaris.ocean.tests.galewsky_jet.test_balance', + 'forward.yaml') + + self.add_input_file(filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file(filename='graph.info', + target='../base_mesh/graph.info') + + self.add_output_file(filename='output.nc') + + def compute_cell_count(self, at_setup): + """ + Compute the approximate number of cells in the mesh, used to constrain + resources + + Parameters + ---------- + at_setup : bool + Whether this method is being run during setup of the step, as + opposed to at runtime + + Returns + ------- + cell_count : int or None + The approximate number of cells in the mesh + """ + if at_setup: + # use a heuristic based on QU30 (65275 cells) and QU240 (10383 + # cells) + cell_count = 6e8 / self.resolution**2 + else: + # get nCells from the input file + with xr.open_dataset('initial_state.nc') as ds: + cell_count = ds.sizes['nCells'] + 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) + + config = self.config + # dt is proportional to resolution: default 30 seconds per km + dt_per_km = config.getfloat('galewsky_jet', 'dt_per_km') + + dt = dt_per_km * self.resolution + # https://stackoverflow.com/a/1384565/7728169 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + + options = dict(config_dt=dt_str) + self.add_model_config_options(options) diff --git a/polaris/ocean/tests/galewsky_jet/test_balance/forward.yaml b/polaris/ocean/tests/galewsky_jet/test_balance/forward.yaml new file mode 100644 index 0000000000..aebe682b23 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/test_balance/forward.yaml @@ -0,0 +1,75 @@ +omega: + run_modes: + config_ocean_run_mode: forward + time_management: + config_run_duration: 0006_00:00:00 + ALE_vertical_grid: + config_vert_coord_movement: impermeable_interfaces + config_ALE_thickness_proportionality: weights_only + decomposition: + config_block_decomp_file_prefix: graph.info.part. + time_integration: + config_time_integrator: RK4 + debug: + config_disable_thick_all_tend: false + config_disable_thick_hadv: false + config_disable_thick_vadv: true + config_disable_thick_sflux: true + config_disable_vel_all_tend: false + config_disable_vel_coriolis: false + config_disable_vel_pgrad: false + config_disable_vel_hmix: false + config_disable_vel_surface_stress: true + config_disable_vel_explicit_bottom_drag: true + config_disable_vel_vmix: true + config_disable_vel_vadv: true + config_disable_tr_all_tend: true + config_disable_tr_hmix: true + config_disable_tr_vmix: true + config_disable_tr_sflux: true + config_disable_tr_nonlocalflux: true + config_check_ssh_consistency: false + cvmix: + config_use_cvmix: false + eos: + config_eos_type: linear + forcing: + config_use_bulk_wind_stress: false + config_use_bulk_thickness_flux: false + pressure_gradient: + config_pressure_gradient_type: ssh_gradient + tracer_forcing_activeTracers: + config_use_activeTracers_surface_restoring: false + config_use_activeTracers_interior_restoring: false + config_use_activeTracers: true + config_use_activeTracers_surface_bulk_forcing: false + tracer_forcing_debugTracers: + config_use_debugTracers: false + AM_mixedLayerDepths: + config_AM_mixedLayerDepths_enable: false + streams: + mesh: + filename_template: initial_state.nc + input: + filename_template: initial_state.nc + restart: + output_interval: 0000_01:00:00 + output: + type: output + filename_template: output.nc + output_interval: 0000_06:00:00 + clobber_mode: truncate + reference_time: 0001-01-01_00:00:00 + contents: + - mesh + - xtime + - normalVelocity + - layerThickness + - refZMid + - refLayerThickness + - ssh + - kineticEnergyCell + - relativeVorticityCell + - relativeVorticity + - normalizedPlanetaryVorticityEdge + - normalizedRelativeVorticityEdge diff --git a/polaris/ocean/tests/galewsky_jet/test_balance/initial_state.py b/polaris/ocean/tests/galewsky_jet/test_balance/initial_state.py new file mode 100644 index 0000000000..f1edf7a5e0 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/test_balance/initial_state.py @@ -0,0 +1,98 @@ +import cmocean # noqa: F401 +import xarray as xr +from mpas_tools.io import write_netcdf +from mpas_tools.mesh.jet import init as jet_init + +from polaris import Step +from polaris.ocean.vertical import init_vertical_coord + + +class InitialState(Step): + """ + A step for creating a mesh and initial condition for baroclinic channel + test cases + + Attributes + ---------- + resolution : float + The resolution of the test case in km + """ + def __init__(self, test_case, resolution): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + resolution : float + The resolution of the test case in km + """ + super().__init__(test_case=test_case, name='initial_state') + self.resolution = resolution + + self.add_input_file( + filename='mesh.nc', + target='../base_mesh/mesh.nc') + self.add_input_file( + filename='graph.info', + target='../base_mesh/graph.info') + # self.add_input_file( + # filename='init.nc') + self.add_output_file('initial_state.nc') + + def run(self): + """ + Run this step of the test case + """ + config = self.config +# update rsph from cime.constants + # jet_init(name='mesh.nc', save='velocity_ic.nc', + # rsph=6371220.0, pert=False) + # ds2 = xr.open_dataset('velocity_ic.nc') + + dsMesh = xr.open_dataset('mesh.nc') + + ds = dsMesh.copy() + x_cell = ds.xCell + + bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') + + ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell) + ds['ssh'] = xr.zeros_like(x_cell) + # ds['ssh'] = ds2.h - ds['bottomDepth'] + + init_vertical_coord(config, ds) + + # resolution = self.resolution + section = config['galewsky_jet'] + temperature = section.getfloat('temperature') + salinity = section.getfloat('salinity') + # coriolis_parameter = section.getfloat('coriolis_parameter') + + temperature_array = temperature * xr.ones_like(x_cell) + temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid) + ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0) + ds['salinity'] = salinity * xr.ones_like(ds.temperature) + + # ds2 = xr.open_dataset('init.nc') + jet_init(name='mesh.nc', save='velocity_ic.nc', + rsph=6371220.0, pert=False) + ds2 = xr.open_dataset('velocity_ic.nc') + + unrm_array, _ = xr.broadcast(ds2.u, ds.refZMid) + ds['normalVelocity'] = unrm_array + # ds['normalVelocity'] = ds2.u + h_array, _ = xr.broadcast(ds2.h, ds.refZMid) + ds['layerThickness'] = h_array + ds['fCell'] = ds2.fCell + ds['fEdge'] = ds2.fEdge + ds['fVertex'] = ds2.fVertex + ds['ssh'] = ds2.h - ds['bottomDepth'] + + # if (config.getfloat('vertical_grid', 'grid_type') == 'uniform'): + # nlev = config.getfloat('vertical_grid', 'vert_levels') + # ds['layerThickness'], _ = xr.broadcast(ds2.h / nlev, ds.refZMid) + + write_netcdf(ds, 'initial_state.nc') diff --git a/polaris/ocean/tests/galewsky_jet/test_balance/viz.py b/polaris/ocean/tests/galewsky_jet/test_balance/viz.py new file mode 100644 index 0000000000..48377fdc48 --- /dev/null +++ b/polaris/ocean/tests/galewsky_jet/test_balance/viz.py @@ -0,0 +1,156 @@ +import cmocean # noqa: F401 +import xarray as xr + +from polaris import Step +from polaris.remap import MappingFileStep +from polaris.viz.globe import plot_global + + +class VizMap(MappingFileStep): + """ + A step for making a mapping file for galewsky jet viz + + Attributes + ---------- + mesh_name : str + The name of the mesh + """ + def __init__(self, test_case, name, subdir, mesh_name): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + name : str + The name of the step + + subdir : str + The subdirectory in the test case's work directory + for the step + + mesh_name : str + The name of the mesh + """ + super().__init__(test_case=test_case, name=name, subdir=subdir, + ntasks=128, min_tasks=1) + self.mesh_name = mesh_name + self.add_input_file(filename='mesh.nc', + target='../base_mesh/mesh.nc') + + def run(self): + """ + Set up the source and destination grids for this step + """ + config = self.config + section = config['galewsky_jet_vizmap'] + dlon = section.getfloat('dlon') + dlat = section.getfloat('dlat') + method = section.get('remap_method') + self.src_from_mpas(filename='mesh.nc', mesh_name=self.mesh_name) + self.dst_global_lon_lat(dlon=dlon, dlat=dlat) + self.method = method + + super().run() + + +class Viz(Step): + """ + A step for plotting fields from the galewsky jet output + + Attributes + ---------- + mesh_name : str + The name of the mesh + """ + def __init__(self, test_case, name, subdir, viz_map, mesh_name): + """ + Create the step + + Parameters + ---------- + test_case : polaris.TestCase + The test case this step belongs to + + name : str + The name of the step + + subdir : str + The subdirectory in the test case's work directory for the step + + viz_map : polaris.ocean.tests.galewsky_jet.test_balance.viz.VizMap + The step for creating a mapping files, also used to remap data + from the MPAS mesh to a lon-lat grid + + mesh_name : str + The name of the mesh + """ + super().__init__(test_case=test_case, name=name, subdir=subdir) + self.add_input_file( + filename='initial_state.nc', + target='../initial_state/initial_state.nc') + self.add_input_file( + filename='output.nc', + target='../forward/output.nc') + self.add_dependency(viz_map, name='viz_map') + self.mesh_name = mesh_name + self.add_output_file('init.png') + self.add_output_file('final.png') + + def run(self): + """ + Run this step of the test case + """ + config = self.config + mesh_name = self.mesh_name + + viz_map = self.dependencies['viz_map'] + + remapper = viz_map.get_remapper() + + ds_init = xr.open_dataset('initial_state.nc') + ds_init = ds_init[['layerThickness', ]].isel(Time=0, nVertLevels=0) + ds_init = remapper.remap(ds_init) + + plot_global(ds_init.lon.values, ds_init.lat.values, + ds_init.layerThickness.values, + out_filename='init.png', config=config, + colormap_section='galewsky_jet_viz_thickness', + title=f'{mesh_name} layerThickness at init', + plot_land=False) + + ds_out = xr.open_dataset('output.nc') + ds_out = ds_out[['layerThickness', ]].isel(Time=-1, nVertLevels=0) + ds_out = remapper.remap(ds_out) + + plot_global(ds_out.lon.values, ds_out.lat.values, + ds_out.layerThickness.values, + out_filename='final.png', config=config, + colormap_section='galewsky_jet_viz_thickness', + title=f'{mesh_name} layerThickness at end of run', + plot_land=False) + + ds_out = xr.open_dataset('output.nc') + ds_out = ds_out[['kineticEnergyCell', ]].isel(Time=-1, nVertLevels=0) + ds_out = remapper.remap(ds_out) + + plot_global(ds_out.lon.values, ds_out.lat.values, + ds_out.kineticEnergyCell.values, + out_filename='final_ke.png', config=config, + colormap_section='galewsky_jet_viz_kineticenergy', + title=f'{mesh_name} kinetic energy at end of run', + plot_land=False) + + ds_out2 = xr.open_dataset('output.nc') + ds_out2 = ds_out2[['kineticEnergyCell', ]].isel(Time=1, nVertLevels=0) + ds_out2 = remapper.remap(ds_out2) + + plot_global(ds_out.lon.values, ds_out.lat.values, + ds_out.kineticEnergyCell.values - + ds_out2.kineticEnergyCell.values, + out_filename='diff_ke.png', config=config, + colormap_section='galewsky_jet_viz_kineticenergy_diff', + title=f'{mesh_name} kinetic energy diff', + plot_land=False) diff --git a/polaris/ocean/tests/yet_another_channel/__init__.py b/polaris/ocean/tests/yet_another_channel/__init__.py new file mode 100644 index 0000000000..95279daf59 --- /dev/null +++ b/polaris/ocean/tests/yet_another_channel/__init__.py @@ -0,0 +1,19 @@ +from polaris.ocean.tests.yet_another_channel.default import Default +from polaris import TestGroup + + +class YetAnotherChannel(TestGroup): + """ + A test group for "yet another channel" test cases + """ + def __init__(self, component): + """ + component : polaris.ocean.Ocean + the ocean component that this test group belongs to + """ + super().__init__(component=component, + name='yet_another_channel') + + for resolution in [1., 4., 10.]: + self.add_test_case( + Default(test_group=self, resolution=resolution))