Source code for compass.ocean.tests.isomip_plus.viz

import os

import matplotlib.pyplot as plt
import numpy as np
import xarray
from mpas_tools.io import write_netcdf

from compass.ocean.haney import compute_haney_number
from compass.ocean.tests.isomip_plus.viz.plot import (
    MoviePlotter,
    TimeSeriesPlotter,
)
from compass.step import Step


[docs] class Viz(Step): """ A step for visualizing output from an ISOMIP+ simulation Attributes ---------- resolution : float The horizontal resolution (km) of the test case experiment : {'Ocean0', 'Ocean1', 'Ocean2'} The ISOMIP+ experiment """
[docs] def __init__(self, test_case, resolution, experiment, tidal_forcing=False): """ Create the step Parameters ---------- test_case : compass.TestCase The test case this step belongs to resolution : float The horizontal resolution (km) of the test case experiment : {'Ocean0', 'Ocean1', 'Ocean2'} The ISOMIP+ experiment """ super().__init__(test_case=test_case, name='viz') self.resolution = resolution self.experiment = experiment self.tidal_forcing = tidal_forcing
[docs] def run(self): """ Run this step of the test case """ config = self.config section = config['isomip_plus_viz'] plot_streamfunctions = section.getboolean('plot_streamfunctions') plot_haney = section.getboolean('plot_haney') frames_per_second = section.getint('frames_per_second') movie_format = section.get('movie_format') section_y = section.getfloat('section_y') section = config['isomip_plus'] min_column_thickness = section.getfloat('min_column_thickness') # show progress only if we're not writing to a log file show_progress = self.log_filename is None expt = self.experiment if self.tidal_forcing: sim_dir = '../performance' else: sim_dir = '../simulation' if not os.path.exists(f'{sim_dir}/timeSeriesStatsMonthly.0001-01-01.nc'): # noqa: E501 sim_dir = '../performance' streamfunction_dir = '../streamfunction' out_dir = '.' dsMesh = xarray.open_dataset(f'{sim_dir}/init.nc') dsOut = xarray.open_dataset(f'{sim_dir}/output.nc') plotter = MoviePlotter(inFolder=sim_dir, streamfunctionFolder=streamfunction_dir, outFolder=f'{out_dir}/plots', expt=expt, sectionY=section_y, dsMesh=dsMesh, ds=dsOut, showProgress=show_progress) if 'time_varying' in expt: plotter.plot_horiz_series(dsOut.ssh, 'ssh', 'ssh', True, cmap='cmo.curl') delice = dsOut.landIcePressure - dsOut.landIcePressure[0, :] plotter.plot_horiz_series(delice, 'delLandIcePressure', 'delLandIcePressure', True, cmap='cmo.curl') plotter.plot_horiz_series(dsOut.velocityX[:, :, 0], 'u', 'u', True, cmap='cmo.balance', vmin=-5e-1, vmax=5e-1) plotter.plot_horiz_series(dsOut.velocityY[:, :, 0], 'v', 'v', True, cmap='cmo.balance', vmin=-5e-1, vmax=5e-1) plotter.plot_horiz_series(dsOut.ssh + dsMesh.bottomDepth, 'H', 'H', True, vmin=min_column_thickness + 1e-10, vmax=700, cmap_set_under='r', cmap_scale='log') if 'tidal' in expt: delssh = dsOut.ssh - dsOut.ssh[0, :] plotter.plot_horiz_series(delssh, 'delssh', 'delssh', True, cmap='cmo.curl', vmin=-1, vmax=1) wct = dsOut.ssh + dsMesh.bottomDepth idx_thin = np.logical_and(wct[0, :] > 1e-1, wct[0, :] < 1) wct_thin = wct[:, idx_thin] wct_mean = wct_thin.mean(dim='nCells').values time = dsOut.daysSinceStartOfSim.values fig = plt.figure() plt.plot(time, wct_mean, '.') fig.set_xlabel('Time (days)') fig.set_ylabel('Mean thickness of thin film (m)') plt.savefig('wct_thin_t.png') plt.close() plotter.plot_horiz_series(wct, 'H', 'H', True, vmin=min_column_thickness + 1e-10, vmax=700, cmap_set_under='r', cmap_scale='log') if os.path.exists(f'{sim_dir}/timeSeriesStatsMonthly.0001-01-01.nc'): ds = xarray.open_mfdataset( '{}/timeSeriesStatsMonthly*.nc'.format(sim_dir), concat_dim='Time', combine='nested') if plot_haney: _compute_and_write_haney_number(dsMesh, ds, out_dir, showProgress=show_progress) tsPlotter = TimeSeriesPlotter(inFolder=sim_dir, outFolder='{}/plots'.format(out_dir), expt=expt) tsPlotter.plot_melt_time_series() tsPlotter = TimeSeriesPlotter( inFolder=sim_dir, outFolder='{}/timeSeriesBelow300m'.format(out_dir), expt=expt) tsPlotter.plot_melt_time_series(sshMax=-300.) mPlotter = MoviePlotter(inFolder=sim_dir, streamfunctionFolder=streamfunction_dir, outFolder='{}/plots'.format(out_dir), expt=expt, sectionY=section_y, dsMesh=dsMesh, ds=ds, showProgress=show_progress) mPlotter.plot_layer_interfaces() if plot_streamfunctions: mPlotter.plot_barotropic_streamfunction() mPlotter.plot_overturning_streamfunction() if plot_haney: mPlotter.plot_haney_number(haneyFolder=out_dir) mPlotter.plot_melt_rates() mPlotter.plot_ice_shelf_boundary_variables() mPlotter.plot_temperature() mPlotter.plot_salinity() mPlotter.plot_potential_density() mPlotter.images_to_movies(outFolder='{}/movies'.format(out_dir), framesPerSecond=frames_per_second, extension=movie_format)
[docs] def file_complete(ds, fileName): """ Find out if the file already has the same number of time slices as the monthly-mean data set """ complete = False if os.path.exists(fileName): with xarray.open_dataset(fileName) as dsCompare: if ds.sizes['Time'] == dsCompare.sizes['Time']: complete = True return complete
def _compute_and_write_haney_number(dsMesh, ds, folder, showProgress=False): """ compute the Haney number rx1 for each edge, and interpolate it to cells. """ haneyFileName = '{}/haney.nc'.format(folder) if file_complete(ds, haneyFileName): return haneyEdge, haneyCell = compute_haney_number( dsMesh, ds.timeMonthly_avg_layerThickness, ds.timeMonthly_avg_ssh, showProgress) dsHaney = xarray.Dataset() dsHaney['xtime_startMonthly'] = ds.xtime_startMonthly dsHaney['xtime_endMonthly'] = ds.xtime_endMonthly dsHaney['haneyEdge'] = haneyEdge dsHaney.haneyEdge.attrs['units'] = 'unitless' dsHaney.haneyEdge.attrs['description'] = 'Haney number on edges' dsHaney['haneyCell'] = haneyCell dsHaney.haneyCell.attrs['units'] = 'unitless' dsHaney.haneyCell.attrs['description'] = 'Haney number on cells' dsHaney = dsHaney.transpose('Time', 'nCells', 'nEdges', 'nVertLevels') write_netcdf(dsHaney, haneyFileName)