From ce8560f4e19578f647df0e1a8eed570e1f47dc11 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 26 Sep 2023 11:56:53 +0200 Subject: [PATCH] Update RPE analysis step and framework Fix the place where nu is added to model config options. --- polaris/ocean/rpe.py | 32 ++-- .../baroclinic_channel/baroclinic_channel.cfg | 17 +- .../ocean/tasks/baroclinic_channel/forward.py | 9 +- .../tasks/baroclinic_channel/rpe/__init__.py | 3 +- .../tasks/baroclinic_channel/rpe/analysis.py | 152 +++++++----------- polaris/ocean/tasks/cosine_bell/__init__.py | 2 + polaris/ocean/tasks/cosine_bell/analysis.py | 2 +- 7 files changed, 97 insertions(+), 120 deletions(-) diff --git a/polaris/ocean/rpe.py b/polaris/ocean/rpe.py index 7f98fef52..8c0901513 100755 --- a/polaris/ocean/rpe.py +++ b/polaris/ocean/rpe.py @@ -5,16 +5,19 @@ from mpas_tools.cime.constants import constants -def compute_rpe(initial_state_file_name, output_files): +def compute_rpe(mesh_filename, initial_state_filename, output_filenames): """ Computes the reference (resting) potential energy for the whole domain Parameters ---------- - initial_state_file_name : str + mesh_filename : str + Name of the netCDF file containing the MPAS horizontal mesh variables + + initial_state_filename : str Name of the netCDF file containing the initial state - output_files : list + output_filenames : list List of netCDF files containing output of forward steps Returns @@ -22,21 +25,22 @@ def compute_rpe(initial_state_file_name, output_files): rpe : numpy.ndarray the reference potential energy of size ``Time`` x ``len(output_files)`` """ - num_files = len(output_files) + num_files = len(output_filenames) if num_files == 0: raise ValueError('Must provide at least one output filename') gravity = constants['SHR_CONST_G'] - dsInit = xr.open_dataset(initial_state_file_name) - nVertLevels = dsInit.sizes['nVertLevels'] + ds_mesh = xr.open_dataset(mesh_filename) + ds_init = xr.open_dataset(initial_state_filename) + nVertLevels = ds_init.sizes['nVertLevels'] - xEdge = dsInit.xEdge - yEdge = dsInit.yEdge - areaCell = dsInit.areaCell - minLevelCell = dsInit.minLevelCell - 1 - maxLevelCell = dsInit.maxLevelCell - 1 - bottomDepth = dsInit.bottomDepth + xEdge = ds_mesh.xEdge + yEdge = ds_mesh.yEdge + areaCell = ds_mesh.areaCell + minLevelCell = ds_init.minLevelCell - 1 + maxLevelCell = ds_init.maxLevelCell - 1 + bottomDepth = ds_init.bottomDepth areaCellMatrix = np.tile(areaCell, (nVertLevels, 1)).transpose() bottomMax = np.max(bottomDepth.values) @@ -54,13 +58,13 @@ def compute_rpe(initial_state_file_name, output_files): vert_index <= maxLevelCell) cell_mask = np.swapaxes(cell_mask, 0, 1) - with xr.open_dataset(output_files[0]) as ds: + with xr.open_dataset(output_filenames[0]) as ds: nt = ds.sizes['Time'] xtime = ds.xtime.values rpe = np.ones((num_files, nt)) - for file_index, out_filename in enumerate(output_files): + for file_index, out_filename in enumerate(output_filenames): ds = xr.open_dataset(out_filename) diff --git a/polaris/ocean/tasks/baroclinic_channel/baroclinic_channel.cfg b/polaris/ocean/tasks/baroclinic_channel/baroclinic_channel.cfg index 99488cbca..a5286ea5e 100644 --- a/polaris/ocean/tasks/baroclinic_channel/baroclinic_channel.cfg +++ b/polaris/ocean/tasks/baroclinic_channel/baroclinic_channel.cfg @@ -37,9 +37,6 @@ btr_dt_per_km = 1.5 # or fractions. False means fractions. use_distances = False -# Viscosity values to test for rpe test case -viscosities = 1, 5, 10, 20, 200 - # Temperature of the surface in the northern half of the domain. surface_temperature = 13.1 @@ -63,3 +60,17 @@ salinity = 35.0 # Coriolis parameter for entire domain. coriolis_parameter = -1.2e-4 + + +# config options for the baroclinic channel RPE tasks +[baroclinic_channel_rpe] + +# Viscosity values to test for rpe test case +viscosities = 1.0, 5.0, 10.0, 20.0, 200.0 + +# plot time (days) +plot_time = 20.0 + +# min and max temperature range +min_temp = 11.8 +max_temp = 13.0 diff --git a/polaris/ocean/tasks/baroclinic_channel/forward.py b/polaris/ocean/tasks/baroclinic_channel/forward.py index fc2fa7752..2d5eed775 100644 --- a/polaris/ocean/tasks/baroclinic_channel/forward.py +++ b/polaris/ocean/tasks/baroclinic_channel/forward.py @@ -72,10 +72,6 @@ def __init__(self, component, resolution, name='forward', subdir=None, indir=indir, ntasks=ntasks, min_tasks=min_tasks, openmp_threads=openmp_threads) - if nu is not None: - # update the viscosity to the requested value - self.add_model_config_options(options=dict(config_mom_del2=nu)) - # make sure output is double precision self.add_yaml_file('polaris.ocean.config', 'output.yaml') @@ -87,6 +83,11 @@ def __init__(self, component, resolution, name='forward', subdir=None, self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel', 'forward.yaml') + if nu is not None: + # update the viscosity to the requested value *after* loading + # forward.yaml + self.add_model_config_options(options=dict(config_mom_del2=nu)) + self.add_output_file( filename='output.nc', validate_vars=['temperature', 'salinity', 'layerThickness', diff --git a/polaris/ocean/tasks/baroclinic_channel/rpe/__init__.py b/polaris/ocean/tasks/baroclinic_channel/rpe/__init__.py index 4997f25c5..e16ceaa17 100644 --- a/polaris/ocean/tasks/baroclinic_channel/rpe/__init__.py +++ b/polaris/ocean/tasks/baroclinic_channel/rpe/__init__.py @@ -68,7 +68,8 @@ def _add_rpe_and_analysis_steps(self, config=None): component = self.component resolution = self.resolution - nus = config.getlist('baroclinic_channel', 'viscosities', dtype=float) + nus = config.getlist('baroclinic_channel_rpe', 'viscosities', + dtype=float) for nu in nus: name = f'nu_{nu:g}' step = Forward( diff --git a/polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py b/polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py index 8ce538106..6eaab45b9 100644 --- a/polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py +++ b/polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py @@ -5,6 +5,7 @@ from polaris import Step from polaris.ocean.rpe import compute_rpe +from polaris.viz import plot_horiz_field class Analysis(Step): @@ -38,7 +39,11 @@ def __init__(self, component, indir, resolution, nus): self.nus = nus self.add_input_file( - filename='initial_state.nc', + filename='mesh.nc', + target='../../init/culled_mesh.nc') + + self.add_input_file( + filename='init.nc', target='../../init/initial_state.nc') for nu in nus: @@ -55,100 +60,53 @@ def run(self): """ Run this step of the test case """ - section = self.config['baroclinic_channel'] - lx = section.getfloat('lx') - ly = section.getfloat('ly') - init_filename = self.inputs[0] - rpe = compute_rpe(initial_state_file_name=init_filename, - output_files=self.inputs[1:]) - with xr.open_dataset(init_filename) as ds_init: - nx = ds_init.attrs['nx'] - ny = ds_init.attrs['ny'] - _plot(nx, ny, lx, ly, self.outputs[0], self.nus, rpe) - - -def _plot(nx, ny, lx, ly, filename, nus, rpe): - """ - Plot section of the baroclinic channel at different viscosities - - Parameters - ---------- - nx : int - The number of cells in the x direction - - ny : int - The number of cells in the y direction (before culling) - - lx : float - The size of the domain in km in the x direction - - ly : int - The size of the domain in km in the y direction - - filename : str - The output file name - - nus : list - The viscosity values - - rpe : numpy.ndarray - The reference potential energy with size len(nu) x len(time) - """ - - plt.switch_backend('Agg') - num_files = len(nus) - time = 20 - - ds = xr.open_dataset(f'output_nu_{nus[0]:g}.nc', decode_times=False) - times = ds.daysSinceStartOfSim.values - - fig = plt.figure() - for i in range(num_files): - rpe_norm = np.divide((rpe[i, :] - rpe[i, 0]), rpe[i, 0]) - plt.plot(times, rpe_norm, - label=f"$\\nu_h=${nus[i]}") - plt.xlabel('Time, days') - plt.ylabel('RPE-RPE(0)/RPE(0)') - plt.legend() - plt.savefig('rpe_t.png') - plt.close(fig) - - fig, axs = plt.subplots(1, num_files, figsize=( - 2.1 * num_files, 5.0), constrained_layout=True) - - # ***NOTE***: This is a quick-and-dirty plotting technique for regular - # planar hex meshes that we do not recommend adopting in other tasks - for iCol, nu in enumerate(nus): - ds = xr.open_dataset(f'output_nu_{nu:g}.nc', decode_times=False) + mesh_filename = 'mesh.nc' + init_filename = 'init.nc' + output_filename = self.outputs[0] + nus = self.nus + section = self.config['baroclinic_channel_rpe'] + + rpe = compute_rpe(mesh_filename=mesh_filename, + initial_state_filename=init_filename, + output_filenames=self.inputs[2:]) + + plt.switch_backend('Agg') + sim_count = len(nus) + time = section.getfloat('plot_time') + min_temp = section.getfloat('min_temp') + max_temp = section.getfloat('max_temp') + + ds = xr.open_dataset(f'output_nu_{nus[0]:g}.nc', decode_times=False) times = ds.daysSinceStartOfSim.values - time_index = np.argmin(np.abs(times - time)) - var = ds.temperature.values - var1 = np.reshape(var[time_index, :, 0], [ny, nx]) - # flip in y-dir - var = np.flipud(var1) - - # Every other row in y needs to average two neighbors in x on - # planar hex mesh - var_avg = var - for j in range(0, ny, 2): - for i in range(1, nx - 2): - var_avg[j, i] = (var[j, i + 1] + var[j, i]) / 2.0 - - ax = axs[iCol] - dis = ax.imshow( - var_avg, - extent=[0, lx, 0, ly], - cmap='cmo.thermal', - vmin=11.8, - vmax=13.0) - ax.set_title(f'day {times[time_index]}, ' - f'$\\nu_h=${nus[iCol]}') - ax.set_xticks(np.linspace(0, lx, 5)) - ax.set_yticks(np.arange(0, ly, 11)) - - ax.set_xlabel('x, km') - if iCol == 0: - ax.set_ylabel('y, km') - if iCol == num_files - 1: - fig.colorbar(dis, ax=axs[num_files - 1], aspect=40) - plt.savefig(filename) + + fig = plt.figure() + for i in range(sim_count): + rpe_norm = np.divide((rpe[i, :] - rpe[i, 0]), rpe[i, 0]) + plt.plot(times, rpe_norm, label=f'$\\nu_h=${nus[i]}') + plt.xlabel('Time, days') + plt.ylabel('RPE-RPE(0)/RPE(0)') + plt.legend() + plt.savefig('rpe_t.png') + plt.close(fig) + + ds_mesh = xr.open_dataset(mesh_filename) + ds_init = xr.open_dataset(init_filename) + + fig, axes = plt.subplots(1, sim_count, figsize=( + 3 * sim_count, 5.0), constrained_layout=True) + + for row_index, nu in enumerate(nus): + ax = axes[row_index] + ds = xr.open_dataset(f'output_nu_{nu:g}.nc', decode_times=False) + ds = ds.isel(nVertLevels=0) + ds['maxLevelCell'] = ds_init.maxLevelCell + times = ds.daysSinceStartOfSim.values + time_index = np.argmin(np.abs(times - time)) + + plot_horiz_field(ds, ds_mesh, 'temperature', ax=ax, + cmap='cmo.thermal', t_index=time_index, + vmin=min_temp, vmax=max_temp, + cmap_title='SST (C)') + ax.set_title(f'day {times[time_index]:g}, $\\nu_h=${nu:g}') + + plt.savefig(output_filename) diff --git a/polaris/ocean/tasks/cosine_bell/__init__.py b/polaris/ocean/tasks/cosine_bell/__init__.py index ae012cbf9..db90f1183 100644 --- a/polaris/ocean/tasks/cosine_bell/__init__.py +++ b/polaris/ocean/tasks/cosine_bell/__init__.py @@ -78,6 +78,8 @@ def configure(self): super().configure() config = self.config config.add_from_package('polaris.mesh', 'mesh.cfg') + config.add_from_package('polaris.ocean.tasks.cosine_bell', + 'cosine_bell.cfg') # set up the steps again in case a user has provided new resolutions self._setup_steps(config) diff --git a/polaris/ocean/tasks/cosine_bell/analysis.py b/polaris/ocean/tasks/cosine_bell/analysis.py index 7fc436d08..6e88a812b 100644 --- a/polaris/ocean/tasks/cosine_bell/analysis.py +++ b/polaris/ocean/tasks/cosine_bell/analysis.py @@ -9,7 +9,7 @@ class Analysis(Step): """ - A step for visualizing the output from the cosine bell test case + A step for analyzing the output from the cosine bell test case Attributes ----------