diff --git a/polaris/ocean/convergence/analysis.py b/polaris/ocean/convergence/analysis.py index 3a0add0975..e231a686ac 100644 --- a/polaris/ocean/convergence/analysis.py +++ b/polaris/ocean/convergence/analysis.py @@ -426,7 +426,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None): The exact solution as derived from the initial condition """ - ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc') + ds_init = self.open_model_dataset( + f'init_r{refinement_factor:02g}.nc', self.config + ) ds_init = ds_init.isel(Time=0) if zidx is not None: ds_init = ds_init.isel(nVertLevels=zidx) @@ -458,7 +460,7 @@ def get_output_field(self, refinement_factor, field_name, time, zidx=None): """ config = self.config ds_out = self.open_model_dataset( - f'output_r{refinement_factor:02g}.nc', decode_times=False + f'output_r{refinement_factor:02g}.nc', config, decode_times=False ) model = config.get('ocean', 'model') diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 360b8b5f9c..e72c9334bf 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/polaris/ocean/model/ocean_io_step.py @@ -32,7 +32,7 @@ def map_to_native_model_vars(self, ds): """ return self.component.map_to_native_model_vars(ds) - def write_model_dataset(self, ds, filename): + def write_model_dataset(self, ds, filename, config=None): """ Write out the given dataset, mapping dimension and variable names from MPAS-Ocean to Omega names if appropriate @@ -45,7 +45,7 @@ def write_model_dataset(self, ds, filename): filename : str The path for the NetCDF file to write """ - self.component.write_model_dataset(ds, filename) + self.component.write_model_dataset(ds, filename, config=config) def map_from_native_model_vars(self, ds): """ @@ -65,7 +65,7 @@ def map_from_native_model_vars(self, ds): """ return self.component.map_from_native_model_vars(ds) - def open_model_dataset(self, filename, **kwargs): + def open_model_dataset(self, filename, config=None, **kwargs): """ Open the given dataset, mapping variable and dimension names from Omega to MPAS-Ocean names if appropriate @@ -83,4 +83,6 @@ def open_model_dataset(self, filename, **kwargs): ds : xarray.Dataset The dataset with variables named as expected in MPAS-Ocean """ - return self.component.open_model_dataset(filename, **kwargs) + return self.component.open_model_dataset( + filename, config=config, **kwargs + ) diff --git a/polaris/ocean/ocean.cfg b/polaris/ocean/ocean.cfg index 6a3f26a9dd..7f9942e3bc 100644 --- a/polaris/ocean/ocean.cfg +++ b/polaris/ocean/ocean.cfg @@ -55,6 +55,9 @@ iterations = 10 # Options related to the vertical grid [vertical_grid] +# The initial surface pressure when not spatially-varying in Pascals +surface_pressure = 0. + # The minimum number of vertical levels for z-star coordinate min_vert_levels = 1 diff --git a/polaris/ocean/vertical/diagnostics.py b/polaris/ocean/vertical/diagnostics.py index d93dfafe13..a8de8dc225 100644 --- a/polaris/ocean/vertical/diagnostics.py +++ b/polaris/ocean/vertical/diagnostics.py @@ -1,6 +1,52 @@ import numpy as np import xarray as xr +from polaris.constants import get_constant +from polaris.ocean.vertical.ztilde import ( + pressure_and_spec_vol_from_state_at_geom_height, + pseudothickness_from_pressure, +) + +RhoSw = get_constant('seawater_density_reference') + + +def geom_thickness_from_ds(ds, config): + if 'layerThickness' in ds.keys(): + return ds['layerThickness'] + elif 'SpecVol' in ds.keys() and 'PseudoThickness' in ds.keys(): + return RhoSw * ds['SpecVol'] * ds['layerThickness'] + else: + raise ValueError( + 'Geometric layerThickness is not present in the ' + 'initial condition and SpecVol is not present ' + 'to compute it' + ) + + +def pseudothickness_from_ds(ds, config): + if 'temperature' not in ds.keys() or 'salinity' not in ds.keys(): + print( + 'PseudoThickness is not present in the ' + 'initial condition and T,S are not present ' + 'to compute it' + ) + return + + surface_pressure = config.getfloat('vertical_grid', 'surface_pressure') + + p_interface, _, _ = pressure_and_spec_vol_from_state_at_geom_height( + config, + ds.layerThickness, + ds.temperature, + ds.salinity, + surface_pressure * xr.ones_like(ds.ssh), + iter_count=1, + ) + + pseudothickness = pseudothickness_from_pressure(p_interface) + + return pseudothickness + def depth_from_thickness(ds): """ diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index d94ccbaba8..373c8fa3c0 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -2,9 +2,9 @@ Conversions between Omega pseudo-height and pressure. Omega's vertical coordinate is pseudo-height - z_tilde = -p / (rho0 * g) + z_tilde = -p / (RhoSw * g) with z_tilde positive upward. Here, ``p`` is sea pressure in Pascals (Pa), -``rho0`` is a reference density (kg m^-3), and ``g`` is gravitational +``RhoSw`` is a reference density (kg m^-3), and ``g`` is gravitational acceleration as defined by ``polaris.constants`` via the Physical Constants Dictionary. """ @@ -20,6 +20,7 @@ __all__ = [ 'z_tilde_from_pressure', + 'pseudothickness_from_pressure', 'pressure_from_z_tilde', 'pressure_and_spec_vol_from_state_at_geom_height', 'pressure_from_geom_thickness', @@ -29,11 +30,48 @@ RhoSw = get_constant('seawater_density_reference') +def pseudothickness_from_pressure( + p: xr.DataArray, +) -> xr.DataArray: + """ + Convert sea pressure to pseudo-thickness. + + z_tilde = -p / (RhoSw * g) + + Parameters + ---------- + p : xarray.DataArray + Sea pressure in Pascals (Pa) at layer interfaces. + + Returns + ------- + xarray.DataArray + Pseudo-thickness at layer mid-points with dimensions NCells by + NVertLayers (one less layer than ``p``) (units: m). + """ + + p_top = p.isel(nVertLevelsP1=slice(0, -1)) + p_bot = p.isel(nVertLevelsP1=slice(1, None)) + dims = list(p.dims) + dims = [item.replace('nVertLevelsP1', 'nVertLevels') for item in dims] + h = xr.DataArray( + (p_bot - p_top) / (RhoSw * Gravity), + dims=tuple(dims), + ) + return h.assign_attrs( + { + 'long_name': 'pseudo-thickness', + 'units': 'm', + 'note': 'h_tilde = -dp / (RhoSw * g)', + } + ) + + def z_tilde_from_pressure(p: xr.DataArray) -> xr.DataArray: """ Convert sea pressure to pseudo-height. - z_tilde = -p / (rho0 * g) + z_tilde = -p / (RhoSw * g) Parameters ---------- @@ -51,7 +89,7 @@ def z_tilde_from_pressure(p: xr.DataArray) -> xr.DataArray: { 'long_name': 'pseudo-height', 'units': 'm', - 'note': 'z_tilde = -p / (rho0 * g)', + 'note': 'z_tilde = -p / (RhoSw * g)', } ) @@ -60,7 +98,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray) -> xr.DataArray: """ Convert pseudo-height to sea pressure. - p = -z_tilde * (rho0 * g) + p = -z_tilde * (RhoSw * g) Parameters ---------- @@ -78,7 +116,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray) -> xr.DataArray: { 'long_name': 'sea pressure', 'units': 'Pa', - 'note': 'p = -rho0 * g * z_tilde', + 'note': 'p = -RhoSw * g * z_tilde', } ) diff --git a/polaris/tasks/ocean/__init__.py b/polaris/tasks/ocean/__init__.py index 243aa09f18..248dee1a00 100644 --- a/polaris/tasks/ocean/__init__.py +++ b/polaris/tasks/ocean/__init__.py @@ -7,6 +7,10 @@ from ruamel.yaml import YAML from polaris import Component +from polaris.ocean.vertical.diagnostics import ( + geom_thickness_from_ds, + pseudothickness_from_ds, +) class Ocean(Component): @@ -154,7 +158,7 @@ def map_var_list_to_native_model(self, var_list): ] return renamed_vars - def write_model_dataset(self, ds, filename): + def write_model_dataset(self, ds, filename, config=None): """ Write out the given dataset, mapping dimension and variable names from MPAS-Ocean to Omega names if appropriate @@ -167,6 +171,14 @@ def write_model_dataset(self, ds, filename): filename : str The path for the NetCDF file to write """ + if ( + self.model == 'omega' + and 'layerThickness' in ds.keys() + and 'PseudoThickness' not in ds.keys() + and config is not None + ): + ds['PseudoThickness'] = pseudothickness_from_ds(ds, config=config) + ds['layerThickness'] = ds['PseudoThickness'].copy() ds = self.map_to_native_model_vars(ds) write_netcdf(ds=ds, fileName=filename) @@ -233,7 +245,7 @@ def map_var_list_from_native_model(self, var_list): ] return renamed_vars - def open_model_dataset(self, filename, **kwargs): + def open_model_dataset(self, filename, config=None, **kwargs): """ Open the given dataset, mapping variable and dimension names from Omega to MPAS-Ocean names if appropriate @@ -252,6 +264,13 @@ def open_model_dataset(self, filename, **kwargs): The dataset with variables named as expected in MPAS-Ocean """ ds = xr.open_dataset(filename, **kwargs) + if ( + self.model == 'omega' + and 'LayerThickness' in ds.keys() + and config is not None + ): + ds['PseudoThickness'] = ds.LayerThickness + ds['LayerThickness'] = geom_thickness_from_ds(ds, config=config) ds = self.map_from_native_model_vars(ds) return ds diff --git a/polaris/tasks/ocean/barotropic_channel/init.py b/polaris/tasks/ocean/barotropic_channel/init.py index 71063a825c..665b1383e7 100644 --- a/polaris/tasks/ocean/barotropic_channel/init.py +++ b/polaris/tasks/ocean/barotropic_channel/init.py @@ -4,12 +4,12 @@ from mpas_tools.mesh.conversion import convert, cull from mpas_tools.planar_hex import make_planar_hex_mesh -from polaris import Step from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.model import OceanIOStep from polaris.ocean.vertical import init_vertical_coord -class Init(Step): +class Init(OceanIOStep): """ A step for creating a mesh and initial condition for barotropic channel tasks @@ -102,7 +102,7 @@ def run(self): ds['fCell'] = xr.zeros_like(ds.xCell) ds['fEdge'] = xr.zeros_like(ds.xEdge) ds['fVertex'] = xr.zeros_like(ds.xVertex) - write_netcdf(ds, 'initial_state.nc') + self.write_model_dataset(ds, 'initial_state.nc', config) # set the wind stress forcing ds_forcing = xr.Dataset() @@ -114,4 +114,4 @@ def run(self): ds_forcing['windStressMeridional'] = ( wind_stress_meridional.expand_dims(dim='Time', axis=0) ) - write_netcdf(ds_forcing, 'forcing.nc') + self.write_model_dataset(ds_forcing, 'forcing.nc') diff --git a/polaris/tasks/ocean/barotropic_channel/viz.py b/polaris/tasks/ocean/barotropic_channel/viz.py index f7c1db77f8..e413af5d0f 100644 --- a/polaris/tasks/ocean/barotropic_channel/viz.py +++ b/polaris/tasks/ocean/barotropic_channel/viz.py @@ -40,8 +40,8 @@ def run(self): Run this step of the task """ ds_mesh = self.open_model_dataset('mesh.nc') - ds_init = self.open_model_dataset('init.nc') - ds_out = self.open_model_dataset('output.nc') + ds_init = self.open_model_dataset('init.nc', self.config) + ds_out = self.open_model_dataset('output.nc', self.config) cell_mask = ds_init.maxLevelCell >= 1 vertex_mask = ds_init.boundaryVertex == 0 diff --git a/polaris/tasks/ocean/barotropic_gyre/init.py b/polaris/tasks/ocean/barotropic_gyre/init.py index c3ad1d8f7b..f7aa67ff4a 100644 --- a/polaris/tasks/ocean/barotropic_gyre/init.py +++ b/polaris/tasks/ocean/barotropic_gyre/init.py @@ -109,6 +109,8 @@ def run(self): # use polaris framework functions to initialize the vertical coordinate init_vertical_coord(config, ds) + ds['temperature'] = 20.0 * xr.ones_like(ds.zMid) + ds['salinity'] = 35.0 * xr.ones_like(ds.zMid) # set the coriolis values for loc in ['Cell', 'Edge', 'Vertex']: @@ -140,7 +142,7 @@ def run(self): ) # write the initial condition file - self.write_model_dataset(ds, 'init.nc') + self.write_model_dataset(ds, 'init.nc', config) cell_mask = ds.maxLevelCell >= 1 diff --git a/polaris/tasks/ocean/cosine_bell/analysis.py b/polaris/tasks/ocean/cosine_bell/analysis.py index 2af7da0233..38e24f8a81 100644 --- a/polaris/tasks/ocean/cosine_bell/analysis.py +++ b/polaris/tasks/ocean/cosine_bell/analysis.py @@ -83,7 +83,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None): ds_mesh = xr.open_dataset(f'mesh_r{refinement_factor:02g}.nc') sphere_radius = ds_mesh.sphere_radius - ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc') + ds_init = self.open_model_dataset( + f'init_r{refinement_factor:02g}.nc', config + ) latCell = ds_init.latCell.values lonCell = ds_init.lonCell.values diff --git a/polaris/tasks/ocean/cosine_bell/init.py b/polaris/tasks/ocean/cosine_bell/init.py index 0d09abd137..a3a2d0dcea 100644 --- a/polaris/tasks/ocean/cosine_bell/init.py +++ b/polaris/tasks/ocean/cosine_bell/init.py @@ -112,7 +112,7 @@ def run(self): ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge) ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex) - self.write_model_dataset(ds, 'initial_state.nc') + self.write_model_dataset(ds, 'initial_state.nc', config) def cosine_bell(max_value, ri, r): diff --git a/polaris/tasks/ocean/cosine_bell/validate.py b/polaris/tasks/ocean/cosine_bell/validate.py index 2bf88fe025..0d9d0157ec 100644 --- a/polaris/tasks/ocean/cosine_bell/validate.py +++ b/polaris/tasks/ocean/cosine_bell/validate.py @@ -58,8 +58,8 @@ def run(self): all_pass = False if all_pass: - ds1 = self.open_model_dataset(filename1) - ds2 = self.open_model_dataset(filename2) + ds1 = self.open_model_dataset(filename1, self.config) + ds2 = self.open_model_dataset(filename2, self.config) all_pass = compare_variables( component=self.component, diff --git a/polaris/tasks/ocean/cosine_bell/viz.py b/polaris/tasks/ocean/cosine_bell/viz.py index 147654a733..e041adac24 100644 --- a/polaris/tasks/ocean/cosine_bell/viz.py +++ b/polaris/tasks/ocean/cosine_bell/viz.py @@ -67,7 +67,7 @@ def run(self): mesh_name = self.mesh_name run_duration = config.getfloat('convergence_forward', 'run_duration') - ds_init = self.open_model_dataset('initial_state.nc') + ds_init = self.open_model_dataset('initial_state.nc', config) da = ds_init['tracer1'].isel(Time=0, nVertLevels=0) plot_global_mpas_field( diff --git a/polaris/tasks/ocean/manufactured_solution/analysis.py b/polaris/tasks/ocean/manufactured_solution/analysis.py index e1d914945b..a11f86c74c 100644 --- a/polaris/tasks/ocean/manufactured_solution/analysis.py +++ b/polaris/tasks/ocean/manufactured_solution/analysis.py @@ -67,7 +67,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None): solution : xarray.DataArray The exact solution as derived from the initial condition """ - init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc') + init = self.open_model_dataset( + f'init_r{refinement_factor:02g}.nc', self.config + ) exact = ExactSolution(self.config, init) if field_name != 'ssh': raise ValueError(f'{field_name} is not currently supported') diff --git a/polaris/tasks/ocean/manufactured_solution/init.py b/polaris/tasks/ocean/manufactured_solution/init.py index 3553405dff..1f8acfb56c 100644 --- a/polaris/tasks/ocean/manufactured_solution/init.py +++ b/polaris/tasks/ocean/manufactured_solution/init.py @@ -109,5 +109,7 @@ def run(self): 'Time', 'nCells', 'nVertLevels' ) ds['layerThickness'] = layer_thickness + ds['temperature'] = xr.zeros_like(layer_thickness) + ds['salinity'] = xr.ones_like(layer_thickness) - self.write_model_dataset(ds, 'initial_state.nc') + self.write_model_dataset(ds, 'initial_state.nc', config) diff --git a/polaris/tasks/ocean/manufactured_solution/viz.py b/polaris/tasks/ocean/manufactured_solution/viz.py index 1bdacc893c..d82fe91f83 100644 --- a/polaris/tasks/ocean/manufactured_solution/viz.py +++ b/polaris/tasks/ocean/manufactured_solution/viz.py @@ -149,13 +149,15 @@ def run(self): for i, refinement_factor in enumerate(refinement_factors): ds_mesh = self.open_model_dataset( - f'mesh_r{refinement_factor:02g}.nc' + f'mesh_r{refinement_factor:02g}.nc', config ) ds_init = self.open_model_dataset( - f'init_r{refinement_factor:02g}.nc' + f'init_r{refinement_factor:02g}.nc', config ) ds = self.open_model_dataset( - f'output_r{refinement_factor:02g}.nc', decode_times=False + f'output_r{refinement_factor:02g}.nc', + config, + decode_times=False, ) exact = ExactSolution(config, ds_init) diff --git a/polaris/tasks/ocean/sphere_transport/init.py b/polaris/tasks/ocean/sphere_transport/init.py index 576daf7dd3..e62cd2c75a 100644 --- a/polaris/tasks/ocean/sphere_transport/init.py +++ b/polaris/tasks/ocean/sphere_transport/init.py @@ -197,4 +197,4 @@ def run(self): ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge) ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex) - self.write_model_dataset(ds, 'initial_state.nc') + self.write_model_dataset(ds, 'initial_state.nc', config) diff --git a/polaris/tasks/ocean/sphere_transport/viz.py b/polaris/tasks/ocean/sphere_transport/viz.py index ef85d45b63..5c003c6eaa 100644 --- a/polaris/tasks/ocean/sphere_transport/viz.py +++ b/polaris/tasks/ocean/sphere_transport/viz.py @@ -82,10 +82,12 @@ def run(self): run_duration = config.getfloat('convergence_forward', 'run_duration') variables_to_plot = self.variables_to_plot - ds_init = self.open_model_dataset('initial_state.nc') + ds_init = self.open_model_dataset('initial_state.nc', config) ds_init = ds_init[variables_to_plot.keys()].isel(Time=0, nVertLevels=0) - ds_out = self.open_model_dataset('output.nc', decode_times=False) + ds_out = self.open_model_dataset( + 'output.nc', config, decode_times=False + ) s_per_hour = 3600.0 # Visualization at halfway around the globe (provided run duration is