diff --git a/compass/ocean/iceshelf.py b/compass/ocean/iceshelf.py index 19bf6d90ea..8e0662ba60 100644 --- a/compass/ocean/iceshelf.py +++ b/compass/ocean/iceshelf.py @@ -42,7 +42,7 @@ def compute_land_ice_pressure_and_draft(ssh, modify_mask, ref_density): def adjust_ssh(variable, iteration_count, step, update_pio=True, - convert_to_cdf5=False): + convert_to_cdf5=False, delta_ssh_threshold=None): """ Adjust the sea surface height or land-ice pressure to be dynamically consistent with one another. A series of short model runs are performed, @@ -168,6 +168,7 @@ def adjust_ssh(variable, iteration_count, step, update_pio=True, iCell = numpy.abs(deltaSSH.where(mask)).argmax().values ds_cell = ds.isel(nCells=iCell) + deltaSSHMax = deltaSSH.isel(nCells=iCell).values if on_a_sphere: coords = (f'lon/lat: ' @@ -177,7 +178,7 @@ def adjust_ssh(variable, iteration_count, step, update_pio=True, coords = (f'x/y: {1e-3 * ds_cell.xCell.values:f} ' f'{1e-3 * ds_cell.yCell.values:f}') string = (f'deltaSSHMax: ' - f'{deltaSSH.isel(nCells=iCell).values:g}, {coords}') + f'{deltaSSHMax:g}, {coords}') logger.info(f' {string}') log_file.write(f'{string}\n') string = (f'ssh: {finalSSH.isel(nCells=iCell).values:g}, ' @@ -186,6 +187,10 @@ def adjust_ssh(variable, iteration_count, step, update_pio=True, logger.info(f' {string}') log_file.write(f'{string}\n') + if delta_ssh_threshold is not None: + if abs(deltaSSHMax) < delta_ssh_threshold: + break + logger.info(" - Complete\n") if out_filename is not None: diff --git a/compass/ocean/tests/ice_shelf_2d/__init__.py b/compass/ocean/tests/ice_shelf_2d/__init__.py index afe1cdf1cd..f79b3fce5c 100644 --- a/compass/ocean/tests/ice_shelf_2d/__init__.py +++ b/compass/ocean/tests/ice_shelf_2d/__init__.py @@ -14,11 +14,14 @@ def __init__(self, mpas_core): """ super().__init__(mpas_core=mpas_core, name='ice_shelf_2d') - for resolution in ['5km']: - for coord_type in ['z-star', 'z-level']: + for resolution in [2e3, 5e3]: + for coord_type in ['z-star', 'z-level', 'single_layer']: self.add_test_case( Default(test_group=self, resolution=resolution, coord_type=coord_type)) + self.add_test_case( + Default(test_group=self, resolution=resolution, + coord_type=coord_type, tidal_forcing=True)) self.add_test_case( RestartTest(test_group=self, resolution=resolution, coord_type=coord_type)) @@ -30,8 +33,8 @@ def configure(resolution, coord_type, config): Parameters ---------- - resolution : str - The resolution of the test case + resolution : float + The resolution of the test case in meters coord_type : str The type of vertical coordinate (``z-star``, ``z-level``, etc.) @@ -39,16 +42,22 @@ def configure(resolution, coord_type, config): config : compass.config.CompassConfigParser Configuration options for this test case """ - res_params = {'5km': {'nx': 10, 'ny': 44, 'dc': 5e3}} + dx = 50e3 # width of domain in m + dy = 220e3 # length of domain in m + dc = resolution + nx = int(dx / resolution) + # ny needs to be even because it is nonperiodic + ny = 2 * int(dy / (2. * resolution)) - if resolution not in res_params: - raise ValueError('Unsupported resolution {}. Supported values are: ' - '{}'.format(resolution, list(res_params))) - res_params = res_params[resolution] - for param in res_params: - config.set('ice_shelf_2d', param, '{}'.format(res_params[param])) + config.set('ice_shelf_2d', 'nx', f'{nx}') + config.set('ice_shelf_2d', 'ny', f'{ny}') + config.set('ice_shelf_2d', 'dc', f'{dc}') config.set('vertical_grid', 'coord_type', coord_type) if coord_type == 'z-level': # we need more vertical resolution config.set('vertical_grid', 'vert_levels', '100') + elif coord_type == 'single_layer': + config.set('vertical_grid', 'vert_levels', '1') + config.set('vertical_grid', 'coord_type', 'z-level') + config.set('vertical_grid', 'partial_cell_type', 'None') diff --git a/compass/ocean/tests/ice_shelf_2d/default/__init__.py b/compass/ocean/tests/ice_shelf_2d/default/__init__.py index 991d358a81..c556b51931 100644 --- a/compass/ocean/tests/ice_shelf_2d/default/__init__.py +++ b/compass/ocean/tests/ice_shelf_2d/default/__init__.py @@ -15,14 +15,15 @@ class Default(TestCase): Attributes ---------- - resolution : str - The horizontal resolution of the test case + resolution : float + The horizontal resolution of the test case in m coord_type : str The type of vertical coordinate (``z-star``, ``z-level``, etc.) """ - def __init__(self, test_group, resolution, coord_type): + def __init__(self, test_group, resolution, coord_type, + tidal_forcing=False): """ Create the test case @@ -38,19 +39,31 @@ def __init__(self, test_group, resolution, coord_type): The type of vertical coordinate (``z-star``, ``z-level``, etc.) """ name = 'default' + if tidal_forcing: + name = 'tidal_forcing' + with_frazil = False + else: + with_frazil = True self.resolution = resolution self.coord_type = coord_type - subdir = '{}/{}/{}'.format(resolution, coord_type, name) + if resolution >= 1e3: + res_name = f'{resolution / 1e3:g}km' + else: + res_name = f'{resolution:g}m' + subdir = f'{res_name}/{coord_type}/{name}' super().__init__(test_group=test_group, name=name, subdir=subdir) self.add_step( InitialState(test_case=self, resolution=resolution)) self.add_step( - SshAdjustment(test_case=self, ntasks=4, openmp_threads=1)) + SshAdjustment(test_case=self, coord_type=coord_type, + resolution=resolution, ntasks=4, openmp_threads=1, + tidal_forcing=tidal_forcing)) self.add_step( Forward(test_case=self, ntasks=4, openmp_threads=1, - resolution=resolution, with_frazil=True)) + coord_type=coord_type, resolution=resolution, + with_frazil=with_frazil, tidal_forcing=tidal_forcing)) self.add_step(Viz(test_case=self), run_by_default=False) def configure(self): diff --git a/compass/ocean/tests/ice_shelf_2d/forward.py b/compass/ocean/tests/ice_shelf_2d/forward.py index f28c647ed6..33381b4719 100644 --- a/compass/ocean/tests/ice_shelf_2d/forward.py +++ b/compass/ocean/tests/ice_shelf_2d/forward.py @@ -1,3 +1,5 @@ +import time + from compass.model import run_model from compass.step import Step @@ -9,11 +11,16 @@ class Forward(Step): Attributes ---------- - resolution : str - The resolution of the test case + resolution : float + The resolution of the test case in m + + coord_type: str + The coordinate type (e.g., 'z-star', 'single_layer', etc.) + """ - def __init__(self, test_case, resolution, name='forward', subdir=None, - ntasks=1, min_tasks=None, openmp_threads=1, with_frazil=True): + def __init__(self, test_case, resolution, coord_type, name='forward', + subdir=None, ntasks=1, min_tasks=None, openmp_threads=1, + with_frazil=True, tidal_forcing=False): """ Create a new test case @@ -22,8 +29,11 @@ def __init__(self, test_case, resolution, name='forward', subdir=None, test_case : compass.TestCase The test case this step belongs to - resolution : str - The resolution of the test case + coord_type: str + The coordinate type (e.g., 'z-star', 'single_layer', etc.) + + resolution : float + The resolution of the test case in m name : str the name of the test case @@ -58,6 +68,13 @@ def __init__(self, test_case, resolution, name='forward', subdir=None, self.add_namelist_file('compass.ocean.tests.ice_shelf_2d', 'namelist.forward') + if coord_type == 'single_layer': + self.add_namelist_file( + 'compass.ocean.tests.ice_shelf_2d', + 'namelist.single_layer.forward_and_ssh_adjust') + if tidal_forcing: + self.add_namelist_file('compass.ocean.tests.ice_shelf_2d', + 'namelist.tidal_forcing.forward') if with_frazil: options = {'config_use_frazil_ice_formation': '.true.', 'config_frazil_maximum_depth': '2000.0'} @@ -71,6 +88,9 @@ def __init__(self, test_case, resolution, name='forward', subdir=None, self.add_streams_file('compass.ocean.tests.ice_shelf_2d', 'streams.forward') + self.add_input_file(filename='forcing_data.nc', + target=('../initial_state/' + 'init_mode_forcing_data.nc')) self.add_input_file(filename='init.nc', target='../ssh_adjustment/adjusted_init.nc') self.add_input_file(filename='graph.info', @@ -87,4 +107,9 @@ def run(self): """ Run this step of the test case """ + config = self.config + dt_per_km = config.getfloat('ice_shelf_2d', 'dt_per_km') + dt = dt_per_km * self.resolution / 1.e3 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + self.update_namelist_at_runtime({'config_dt': dt_str}) run_model(self) diff --git a/compass/ocean/tests/ice_shelf_2d/ice_shelf_2d.cfg b/compass/ocean/tests/ice_shelf_2d/ice_shelf_2d.cfg index 1e8039589c..844d5b7d9d 100644 --- a/compass/ocean/tests/ice_shelf_2d/ice_shelf_2d.cfg +++ b/compass/ocean/tests/ice_shelf_2d/ice_shelf_2d.cfg @@ -8,7 +8,7 @@ grid_type = uniform vert_levels = 20 # Depth of the bottom of the ocean -bottom_depth = 2000.0 +bottom_depth = 1050.0 # The type of vertical coordinate (e.g. z-level, z-star) coord_type = z-star @@ -23,20 +23,23 @@ min_pc_fraction = 0.1 # config options for 2D ice-shelf testcases [ice_shelf_2d] -# Vertical thickness of ocean sub-ice cavity -cavity_thickness = 50.0 +# Time step in seconds as a function of resolution +dt_per_km = 5 -# Vertical thickness of fixed slope -slope_height = 500.0 +# Vertical thickness of ocean sub-ice cavity at GL +y1_water_column_thickness = 50.0 -# Horizontal width of angled part of the ice -edge_width = 15.0e3 +# Vertical thickness of water column thickness at y2 +y2_water_column_thickness = 1050.0 -# cavity edge in y +# GL location in y y1 = 30.0e3 -# shelf edge in y -y2 = 60.0e3 +# ice shelf inflection point in y +y2 = 90.0e3 + +# ice shelf front location in y +y3 = 90.0e3 # Temperature of the surface in the northern half of the domain temperature = 1.0 diff --git a/compass/ocean/tests/ice_shelf_2d/initial_state.py b/compass/ocean/tests/ice_shelf_2d/initial_state.py index ac795c40ed..57a9763896 100644 --- a/compass/ocean/tests/ice_shelf_2d/initial_state.py +++ b/compass/ocean/tests/ice_shelf_2d/initial_state.py @@ -1,3 +1,4 @@ +import numpy as np import xarray from mpas_tools.cime.constants import constants from mpas_tools.io import write_netcdf @@ -71,9 +72,9 @@ def run(self): # d variables are total water-column thickness below ice shelf y1 = section.getfloat('y1') y2 = section.getfloat('y2') - y3 = y2 + section.getfloat('edge_width') - d1 = section.getfloat('cavity_thickness') - d2 = d1 + section.getfloat('slope_height') + y3 = section.getfloat('y3') + d1 = section.getfloat('y1_water_column_thickness') + d2 = section.getfloat('y2_water_column_thickness') d3 = bottom_depth ds = dsMesh.copy() @@ -130,3 +131,10 @@ def run(self): ds['landIceDraft'] = landIceDraft write_netcdf(ds, 'initial_state.nc') + + # Generate the tidal forcing dataset whether it is used or not + ds_forcing = xarray.Dataset() + y_max = np.max(ds.yCell.values) + ds_forcing['tidalInputMask'] = xarray.where( + ds.yCell > (y_max - 0.6 * 5.0e3), 1.0, 0.0) + write_netcdf(ds_forcing, 'init_mode_forcing_data.nc') diff --git a/compass/ocean/tests/ice_shelf_2d/namelist.forward b/compass/ocean/tests/ice_shelf_2d/namelist.forward index b3355cb7ba..45c1b33261 100644 --- a/compass/ocean/tests/ice_shelf_2d/namelist.forward +++ b/compass/ocean/tests/ice_shelf_2d/namelist.forward @@ -1,4 +1,3 @@ -config_dt = '00:05:00' config_btr_dt = '00:00:15' config_run_duration = '0000_00:10:00' config_use_mom_del2 = .true. diff --git a/compass/ocean/tests/ice_shelf_2d/namelist.single_layer.forward_and_ssh_adjust b/compass/ocean/tests/ice_shelf_2d/namelist.single_layer.forward_and_ssh_adjust new file mode 100644 index 0000000000..aecbfefd83 --- /dev/null +++ b/compass/ocean/tests/ice_shelf_2d/namelist.single_layer.forward_and_ssh_adjust @@ -0,0 +1,14 @@ +config_time_integrator='RK4' +config_bottom_drag_mode = 'explicit' +config_explicit_bottom_drag_coeff = 3.0e-3 +config_AM_globalStats_enable = .false. +config_AM_globalStats_compute_on_startup = .false. +config_AM_globalStats_write_on_startup = .false. +config_compute_active_tracer_budgets = .false. +config_disable_thick_vadv = .true. +config_disable_thick_sflux = .true. +config_disable_vel_hmix = .true. +config_disable_vel_vmix = .true. +config_disable_vel_vadv = .true. +config_disable_tr_all_tend = .true. +config_disable_redi_k33 = .true. diff --git a/compass/ocean/tests/ice_shelf_2d/namelist.tidal_forcing.forward b/compass/ocean/tests/ice_shelf_2d/namelist.tidal_forcing.forward new file mode 100644 index 0000000000..d658d3a75b --- /dev/null +++ b/compass/ocean/tests/ice_shelf_2d/namelist.tidal_forcing.forward @@ -0,0 +1,8 @@ +config_use_tidal_forcing = .true. +config_use_tidal_forcing_tau = 10000 +config_tidal_forcing_model = 'monochromatic' +config_tidal_forcing_type = 'thickness_source' +config_tidal_forcing_monochromatic_amp = 1.0 +config_tidal_forcing_monochromatic_period = 10.0 +config_eos_type = 'linear' +config_land_ice_flux_mode = 'pressure_only' diff --git a/compass/ocean/tests/ice_shelf_2d/restart_test/__init__.py b/compass/ocean/tests/ice_shelf_2d/restart_test/__init__.py index ce437b7033..0729a564cc 100644 --- a/compass/ocean/tests/ice_shelf_2d/restart_test/__init__.py +++ b/compass/ocean/tests/ice_shelf_2d/restart_test/__init__.py @@ -40,20 +40,26 @@ def __init__(self, test_group, resolution, coord_type): name = 'restart_test' self.resolution = resolution self.coord_type = coord_type - subdir = '{}/{}/{}'.format(resolution, coord_type, name) + if resolution >= 1e3: + res_name = f'{resolution / 1e3:g}km' + else: + res_name = f'{resolution:g}m' + subdir = f'{res_name}/{coord_type}/{name}' super().__init__(test_group=test_group, name=name, subdir=subdir) self.add_step( InitialState(test_case=self, resolution=resolution)) self.add_step( - SshAdjustment(test_case=self, ntasks=4, openmp_threads=1)) + SshAdjustment(test_case=self, resolution=resolution, + coord_type=coord_type, ntasks=4, + openmp_threads=1)) for part in ['full', 'restart']: name = '{}_run'.format(part) - step = Forward(test_case=self, name=name, subdir=name, ntasks=4, - openmp_threads=1, resolution=resolution, - with_frazil=True) + step = Forward(test_case=self, name=name, subdir=name, + coord_type=coord_type, ntasks=4, openmp_threads=1, + resolution=resolution, with_frazil=True) step.add_namelist_file( 'compass.ocean.tests.ice_shelf_2d.restart_test', diff --git a/compass/ocean/tests/ice_shelf_2d/ssh_adjustment.py b/compass/ocean/tests/ice_shelf_2d/ssh_adjustment.py index cdba398f86..13e15d99a1 100644 --- a/compass/ocean/tests/ice_shelf_2d/ssh_adjustment.py +++ b/compass/ocean/tests/ice_shelf_2d/ssh_adjustment.py @@ -1,3 +1,5 @@ +import time + from compass.ocean.iceshelf import adjust_ssh from compass.step import Step @@ -7,7 +9,8 @@ class SshAdjustment(Step): A step for iteratively adjusting the pressure from the weight of the ice shelf to match the sea-surface height as part of ice-shelf 2D test cases """ - def __init__(self, test_case, ntasks=1, min_tasks=None, openmp_threads=1): + def __init__(self, test_case, resolution, coord_type, ntasks=1, + min_tasks=None, openmp_threads=1, tidal_forcing=False): """ Create the step @@ -16,6 +19,12 @@ def __init__(self, test_case, ntasks=1, min_tasks=None, openmp_threads=1): test_case : compass.TestCase The test case this step belongs to + resolution : float + The resolution of the test case in m + + coord_type: str + The coordinate type (e.g., 'z-star', 'single_layer', etc.) + 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 @@ -29,6 +38,7 @@ def __init__(self, test_case, ntasks=1, min_tasks=None, openmp_threads=1): the number of OpenMP threads the step will use """ + self.resolution = resolution if min_tasks is None: min_tasks = ntasks super().__init__(test_case=test_case, name='ssh_adjustment', @@ -39,6 +49,10 @@ def __init__(self, test_case, ntasks=1, min_tasks=None, openmp_threads=1): # start with the same namelist settings as the forward run self.add_namelist_file('compass.ocean.tests.ice_shelf_2d', 'namelist.forward') + if coord_type == 'single_layer': + self.add_namelist_file( + 'compass.ocean.tests.ice_shelf_2d', + 'namelist.single_layer.forward_and_ssh_adjust') # we don't want the global stats AM for this run self.add_namelist_options({'config_AM_globalStats_enable': '.false.'}) @@ -67,6 +81,10 @@ def run(self): Run this step of the test case """ config = self.config + dt_per_km = config.getfloat('ice_shelf_2d', 'dt_per_km') + dt = dt_per_km * self.resolution / 1.e3 + dt_str = time.strftime('%H:%M:%S', time.gmtime(dt)) + self.update_namelist_at_runtime({'config_dt': dt_str}) iteration_count = config.getint('ssh_adjustment', 'iterations') adjust_ssh(variable='landIcePressure', iteration_count=iteration_count, - step=self) + step=self, delta_ssh_threshold=1.e-10) diff --git a/compass/ocean/tests/ice_shelf_2d/streams.forward b/compass/ocean/tests/ice_shelf_2d/streams.forward index 15f3cab2ca..67dc4e3915 100644 --- a/compass/ocean/tests/ice_shelf_2d/streams.forward +++ b/compass/ocean/tests/ice_shelf_2d/streams.forward @@ -8,6 +8,14 @@ + + + + + + + + + @@ -33,7 +45,7 @@ + output_interval="0000_00:10:00" > diff --git a/compass/ocean/tests/ice_shelf_2d/viz.py b/compass/ocean/tests/ice_shelf_2d/viz.py index 972f303705..0afb7ad197 100644 --- a/compass/ocean/tests/ice_shelf_2d/viz.py +++ b/compass/ocean/tests/ice_shelf_2d/viz.py @@ -1,7 +1,9 @@ import matplotlib.pyplot as plt -import numpy -import xarray +import numpy as np +import pandas as pd +import xarray as xr +from compass.ocean.tests.isomip_plus.viz.plot import MoviePlotter from compass.step import Step @@ -22,14 +24,27 @@ def __init__(self, test_case): 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_input_file(filename='output_ssh.nc', + target='../ssh_adjustment/output_ssh.nc') self.add_output_file('vert_grid.png') + self.add_output_file('velocity_max_t.png') def run(self): """ Run this step of the test case """ - ds = xarray.open_dataset('initial_state.nc') + ds = xr.open_dataset('initial_state.nc') + dsOut = xr.open_dataset('output.nc') + dsAdj = xr.open_dataset('output_ssh.nc') + plotter = MoviePlotter(inFolder='../forward', + streamfunctionFolder='', + outFolder='./plots', + expt='', sectionY=20. * 5.0e3, + dsMesh=ds, ds=dsOut, + showProgress=False) if 'Time' in ds.dims: ds = ds.isel(Time=0) @@ -39,24 +54,24 @@ def run(self): nCells = ds.sizes['yCell'] nVertLevels = ds.sizes['nVertLevels'] - zIndex = xarray.DataArray(data=numpy.arange(nVertLevels), - dims='nVertLevels') + zIndex = xr.DataArray(data=np.arange(nVertLevels), + dims='nVertLevels') minLevelCell = ds.minLevelCell - 1 maxLevelCell = ds.maxLevelCell - 1 - cellMask = numpy.logical_and(zIndex >= minLevelCell, - zIndex <= maxLevelCell) + cellMask = np.logical_and(zIndex >= minLevelCell, + zIndex <= maxLevelCell) - zIndex = xarray.DataArray(data=numpy.arange(nVertLevels + 1), - dims='nVertLevelsP1') + zIndex = xr.DataArray(data=np.arange(nVertLevels + 1), + dims='nVertLevelsP1') - interfaceMask = numpy.logical_and(zIndex >= minLevelCell, - zIndex <= maxLevelCell + 1) + interfaceMask = np.logical_and(zIndex >= minLevelCell, + zIndex <= maxLevelCell + 1) zMid = ds.zMid.where(cellMask) - zInterface = numpy.zeros((nCells, nVertLevels + 1)) + zInterface = np.zeros((nCells, nVertLevels + 1)) zInterface[:, 0] = ds.ssh.values for zIndex in range(nVertLevels): thickness = ds.layerThickness.isel(nVertLevels=zIndex) @@ -64,14 +79,89 @@ def run(self): zInterface[:, zIndex + 1] = \ zInterface[:, zIndex] - thickness.values - zInterface = xarray.DataArray(data=zInterface, - dims=('yCell', 'nVertLevelsP1')) + zInterface = xr.DataArray(data=zInterface, + dims=('yCell', 'nVertLevelsP1')) zInterface = zInterface.where(interfaceMask) - plt.figure(figsize=[24, 12], dpi=100) - plt.plot(ds.yCell.values, zMid.values, 'b') - plt.plot(ds.yCell.values, zInterface.values, 'k') - plt.plot(ds.yCell.values, ds.ssh.values, 'g') - plt.plot(ds.yCell.values, -ds.bottomDepth.values, 'g') + y = ds.yCell.values / 1.e3 + plt.figure(figsize=[12, 6], dpi=100) + plt.plot(y, zMid.values, 'b', label='zMid') + plt.plot(y, zInterface.values, 'k', label='zTop') + plt.plot(y, ds.ssh.values, '--g', label='SSH') + plt.plot(y, -ds.bottomDepth.values, '--r', label='zBed') + plt.xlabel('Distance (km)') + plt.ylabel('Depth (m)') plt.savefig('vert_grid.png', dpi=200) plt.close() + + ds = xr.open_dataset('initial_state.nc') + + # Plot the time series of max velocity + plt.figure(figsize=[12, 6], dpi=100) + umax = np.amax(dsOut.velocityX[:, :, 0].values, axis=1) + vmax = np.amax(dsOut.velocityY[:, :, 0].values, axis=1) + t = dsOut.daysSinceStartOfSim.values + time = pd.to_timedelta(t) / 1.e9 + plt.plot(time, umax, 'k', label='u') + plt.plot(time, vmax, 'b', label='v') + plt.xlabel('Time (s)') + plt.ylabel('Velocity (m/s)') + plt.legend() + plt.savefig('velocity_max_t.png', dpi=200) + plt.close() + min_column_thickness = self.config.getfloat( + 'ice_shelf_2d', 'y1_water_column_thickness') + + figsize = (4, 8) + delssh = dsOut.ssh - dsOut.ssh[0, :] + s_vmin = np.nanmin(delssh.values) + s_vmax = np.nanmax(delssh.values) + plotter.plot_horiz_series(da=delssh, nameInTitle='delssh', + prefix='delssh', oceanDomain=True, + cmap='cmo.curl', + vmin=-1. * max(abs(s_vmin), abs(s_vmax)), + vmax=max(abs(s_vmin), abs(s_vmax)), + figsize=figsize) + + u_vmin = np.nanmin(dsOut.velocityX[:, :, 0].values) + u_vmax = np.nanmax(dsOut.velocityX[:, :, 0].values) + plotter.plot_horiz_series(da=dsOut.velocityX[:, :, 0], nameInTitle='u', + prefix='u', oceanDomain=True, + cmap='cmo.balance', + vmin=-1. * max(abs(u_vmin), abs(u_vmax)), + vmax=max(abs(u_vmin), abs(u_vmax)), + figsize=figsize) + + v_vmin = np.nanmin(dsOut.velocityY[:, :, 0].values) + v_vmax = np.nanmax(dsOut.velocityY[:, :, 0].values) + plotter.plot_horiz_series(da=dsOut.velocityY[:, :, 0], nameInTitle='v', + prefix='v', oceanDomain=True, + cmap='cmo.balance', + vmin=-1. * max(abs(v_vmin), abs(v_vmax)), + vmax=max(abs(v_vmin), abs(v_vmax)), + figsize=figsize) + + plotter.plot_horiz_series(da=dsOut.landIcePressure, + nameInTitle='landIcePressure', + prefix='landIcePressure', oceanDomain=True, + vmin=1e3, vmax=1e7, cmap_set_under='r', + cmap_scale='log', figsize=figsize) + + plotter.plot_horiz_series(da=dsOut.ssh + ds.bottomDepth, + nameInTitle='H', prefix='H', + oceanDomain=True, + vmin=min_column_thickness + 1e-10, + vmax=2000, cmap_set_under='r', + cmap_scale='log', figsize=figsize) + + delssh = dsAdj.ssh - ds.ssh + s_vmin = np.nanmin(delssh.values) + s_vmax = np.nanmax(delssh.values) + plotter.plot_horiz_series(da=delssh, + nameInTitle='delssh_adjust', + prefix='delssh_adjust', + oceanDomain=True, + cmap='cmo.curl', time_indices=[0], + vmin=-1. * max(abs(s_vmin), abs(s_vmax)), + vmax=max(abs(s_vmin), abs(s_vmax)), + figsize=figsize) diff --git a/compass/ocean/tests/isomip_plus/viz/plot.py b/compass/ocean/tests/isomip_plus/viz/plot.py index dc1e298573..049e2f675a 100644 --- a/compass/ocean/tests/isomip_plus/viz/plot.py +++ b/compass/ocean/tests/isomip_plus/viz/plot.py @@ -103,11 +103,13 @@ def plot_melt_time_series(self, sshMax=None): self.plot_time_series(1e-6 * totalMeltFlux, 'total melt flux', 'totalMeltFlux', 'kT/yr') - blTempName = 'timeMonthly_avg_landIceBoundaryLayerTracers_' \ - 'landIceBoundaryLayerTemperature' - interTempName = 'timeMonthly_avg_landIceInterfaceTracers_' \ - 'landIceInterfaceTemperature' - da = (self.ds[blTempName] - self.ds[interTempName]) + prefix = 'timeMonthly_avg_landIceBoundaryLayerTracers_' + boundary_layer_temperature = \ + self.ds[f'{prefix}landIceBoundaryLayerTemperature'] + prefix = 'timeMonthly_avg_landIceInterfaceTracers_' + interface_temperature = \ + self.ds[f'{prefix}landIceInterfaceTemperature'] + da = boundary_layer_temperature - interface_temperature da = (da * areaCell * iceMask).sum(dim='nCells') / totalArea self.plot_time_series(da, 'mean thermal driving', @@ -377,23 +379,23 @@ def plot_ice_shelf_boundary_variables(self): oceanDomain=False, units='W/s', vmin=-1e1, vmax=1e1, cmap='cmo.curl') - blTempName = 'timeMonthly_avg_landIceBoundaryLayerTracers_' \ - 'landIceBoundaryLayerTemperature' - interTempName = 'timeMonthly_avg_landIceInterfaceTracers_' \ - 'landIceInterfaceTemperature' - - da = (self.ds[blTempName] - self.ds[interTempName]) + prefix = 'timeMonthly_avg_landIceBoundaryLayerTracers_' + boundary_layer_temperature = \ + self.ds[f'{prefix}landIceBoundaryLayerTemperature'] + boundary_layer_salinity = \ + self.ds[f'{prefix}landIceBoundaryLayerSalinity'] + prefix = 'timeMonthly_avg_landIceInterfaceTracers_' + interface_temperature = \ + self.ds[f'{prefix}landIceInterfaceTemperature'] + interface_salinity = \ + self.ds[f'{prefix}landIceInterfaceSalinity'] + da = boundary_layer_temperature - interface_temperature self.plot_horiz_series(da, 'thermal driving', prefix='thermalDriving', oceanDomain=False, units='deg C', vmin=-2, vmax=2, cmap='cmo.thermal') - blSalinName = 'timeMonthly_avg_landIceBoundaryLayerTracers_' \ - 'landIceBoundaryLayerSalinity' - interSalinName = 'timeMonthly_avg_landIceInterfaceTracers_' \ - 'landIceInterfaceSalinity' - - da = (self.ds[blSalinName] - self.ds[interSalinName]) + da = boundary_layer_salinity - interface_salinity self.plot_horiz_series(da, 'haline driving', prefix='halineDriving', oceanDomain=False, units='PSU', @@ -467,7 +469,8 @@ def plot_haney_number(self, haneyFolder=None): def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain, units=None, vmin=None, vmax=None, cmap=None, cmap_set_under=None, cmap_set_over=None, - cmap_scale='linear'): + cmap_scale='linear', time_indices=None, + figsize=(9, 3)): """ Plot a series of image of a given variable @@ -500,6 +503,9 @@ def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain, cmap_scale : {'log', 'linear'}, optional Whether the colormap is logarithmic or linear + + time_indices : list of int, optional + The time indices at which to plot. If not provided, set to all. """ nTime = self.ds.sizes['Time'] @@ -512,7 +518,9 @@ def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain, else: bar = None - for tIndex in range(nTime): + if time_indices is None: + time_indices = range(nTime) + for tIndex in time_indices: self.update_date(tIndex) field = da.isel(Time=tIndex).values outFileName = '{}/{}/{}_{:04d}.png'.format( @@ -526,7 +534,7 @@ def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain, vmax=vmax, cmap=cmap, cmap_set_under=cmap_set_under, cmap_set_over=cmap_set_over, - cmap_scale=cmap_scale) + cmap_scale=cmap_scale, figsize=figsize) if self.showProgress: bar.update(tIndex + 1) if self.showProgress: @@ -859,8 +867,12 @@ def _plot_vert_field(self, inX, inZ, field, title, outFileName, plt.close() def _compute_section_x_z(self): - x = _interp_extrap_corner( - self.dsMesh.xIsomipCell[self.sectionCellIndices]) + if 'xIsomipCell' in self.dsMesh.keys(): + x = _interp_extrap_corner( + self.dsMesh.xIsomipCell[self.sectionCellIndices]) + else: + x = _interp_extrap_corner( + self.dsMesh.xCell[self.sectionCellIndices]) nx = len(x) nVertLevels = self.dsMesh.sizes['nVertLevels'] nTime = self.ds.sizes['Time'] @@ -902,8 +914,12 @@ def _compute_cell_patches(dsMesh, mask): patches = [] nVerticesOnCell = dsMesh.nEdgesOnCell.values verticesOnCell = dsMesh.verticesOnCell.values - 1 - xVertex = dsMesh.xIsomipVertex.values - yVertex = dsMesh.yIsomipVertex.values + if 'xIsomipVertex' in dsMesh.keys(): + xVertex = dsMesh.xIsomipVertex.values + yVertex = dsMesh.yIsomipVertex.values + else: + xVertex = dsMesh.xVertex.values + yVertex = dsMesh.yVertex.values for iCell in range(dsMesh.sizes['nCells']): if not mask[iCell]: continue @@ -922,8 +938,12 @@ def _compute_cell_patches(dsMesh, mask): def _compute_section_cell_indices(y, dsMesh): - xCell = dsMesh.xIsomipCell.values - yCell = dsMesh.yIsomipCell.values + if 'xIsomipCell' in dsMesh.keys(): + xCell = dsMesh.xIsomipCell.values + yCell = dsMesh.yIsomipCell.values + else: + xCell = dsMesh.xCell.values + yCell = dsMesh.yCell.values xMin = numpy.amin(xCell) xMax = numpy.amax(xCell) xs = numpy.linspace(xMin, xMax, 10000)