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

import os
import subprocess

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr

from compass.step import Step


[docs] class Viz(Step): """ A step for visualizing drying slope results, as well as comparison with analytical solution and ROMS results. Attributes ---------- damping_coeffs : list of float Rayleigh damping coefficients used for the MPAS-Ocean and ROMS forward runs times : list of str Time in days since the start of the simulation for which ROMS comparison data is available datatypes : list of str The sources of data for comparison to the MPAS-Ocean run """
[docs] def __init__(self, test_case, damping_coeffs=None): """ Create the step Parameters ---------- test_case : compass.TestCase The test case this step belongs to """ super().__init__(test_case=test_case, name='viz') times = ['0.05', '0.15', '0.25', '0.30', '0.40', '0.50'] datatypes = ['analytical', 'ROMS'] self.damping_coeffs = damping_coeffs self.times = times self.datatypes = datatypes self.add_input_file(filename='init.nc', target='../initial_state/initial_state.nc') if damping_coeffs is None: self.add_input_file(filename='output.nc', target='../forward/output.nc') else: for coeff in damping_coeffs: self.add_input_file(filename=f'output_{coeff}.nc', target=f'../forward_{coeff}/output.nc') for time in times: for datatype in datatypes: filename = f'r{coeff}d{time}-{datatype.lower()}'\ '.csv' self.add_input_file(filename=filename, target=filename, database='drying_slope')
[docs] def run(self): """ Run this step of the test case """ section = self.config['paths'] section = self.config['vertical_grid'] section = self.config['drying_slope_viz'] generate_movie = section.getboolean('generate_movie') self._plot_ssh_validation(times=self.times) self._plot_ssh_time_series() if generate_movie: frames_per_second = section.getint('frames_per_second') movie_format = section.get('movie_format') outFolder = 'movie' if not os.path.exists(outFolder): try: os.makedirs(os.path.join(os.getcwd(), outFolder)) except OSError: pass for tidx, itime in enumerate(np.linspace(0, 0.5, 5 * 12 + 1)): self._plot_ssh_validation(times=[itime], tidx=tidx, outFolder=outFolder) self._images_to_movies(framesPerSecond=frames_per_second, outFolder=outFolder, extension=movie_format)
def _forcing(self, t): ssh = 10. * np.sin(t * np.pi / 12.) - 10. return ssh def _plot_ssh_time_series(self, outFolder='.'): """ Plot ssh forcing on the right x boundary as a function of time against the analytical solution. The agreement should be within machine precision if the namelist options are consistent with the Warner et al. (2013) test case. """ colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'} xSsh = np.linspace(0, 12.0, 100) ySsh = 10.0 * np.sin(xSsh * np.pi / 12.0) - 10.0 figsize = [6.4, 4.8] damping_coeffs = self.damping_coeffs if damping_coeffs is None: naxes = 1 ncFilename = ['output.nc'] else: naxes = len(damping_coeffs) ncFilename = [f'output_{damping_coeff}.nc' for damping_coeff in damping_coeffs] fig, _ = plt.subplots(nrows=naxes, ncols=1, figsize=figsize, dpi=100) for i in range(naxes): ax = plt.subplot(naxes, 1, i + 1) ds = xr.open_dataset(ncFilename[i]) ympas = ds.ssh.where(ds.tidalInputMask).mean('nCells').values xmpas = np.linspace(0, 1.0, len(ds.xtime)) * 12.0 ax.plot(xmpas, ympas, marker='o', label='MPAS-O forward', color=colors['MPAS-O']) ax.plot(xSsh, ySsh, lw=3, label='analytical', color=colors['analytical']) ax.set_ylabel('Tidal amplitude (m)') ax.set_xlabel('Time (hrs)') ax.legend(frameon=False) ax.label_outer() ds.close() fig.suptitle('Tidal amplitude forcing (right side)') fig.savefig(f'{outFolder}/ssh_t.png', bbox_inches='tight', dpi=200) plt.close(fig) def _plot_ssh_validation(self, times, tidx=None, outFolder='.'): """ Plot ssh as a function of along-channel distance for all times for which there is validation data """ colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'} locs = [7.2, 2.2, 0.2, 1.2, 4.2, 9.3] locs = 0.92 - np.divide(locs, 11.) damping_coeffs = self.damping_coeffs datatypes = self.datatypes if damping_coeffs is None: naxes = 1 nhandles = 1 ncFilename = ['output.nc'] else: naxes = len(damping_coeffs) nhandles = len(datatypes) + 1 ncFilename = [f'output_{damping_coeff}.nc' for damping_coeff in damping_coeffs] ds_mesh = xr.open_dataset('init.nc') mesh_ymean = ds_mesh.isel(Time=0).groupby('yCell').mean( dim=xr.ALL_DIMS) bottom_depth = mesh_ymean.bottomDepth.values drying_length = self.config.getfloat('drying_slope', 'ly_analysis') drying_length = drying_length * 1e3 x_offset = np.max(mesh_ymean.yCell.values) - drying_length x = (mesh_ymean.yCell.values - x_offset) / 1000.0 xBed = np.linspace(0, 25, 100) yBed = 10.0 / 25.0 * xBed fig, _ = plt.subplots(nrows=naxes, ncols=1, sharex=True) for i in range(naxes): ax = plt.subplot(naxes, 1, i + 1) ds = xr.open_dataset(ncFilename[i]) ds = ds.drop_vars(np.setdiff1d([j for j in ds.variables], ['yCell', 'ssh'])) ax.plot(xBed, yBed, '-k', lw=3) ax.set_xlim(0, 25) ax.set_ylim(-1, 11) ax.invert_yaxis() ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.label_outer() for atime, ay in zip(times, locs): # Plot MPAS-O data # factor of 1e- needed to account for annoying round-off issue # to get right time slices plottime = int((float(atime) / 0.2 + 1e-16) * 24.0) ymean = ds.isel(Time=plottime).groupby('yCell').mean( dim=xr.ALL_DIMS) y = ymean.ssh.values ax.plot(x, -y, label='MPAS-O', color=colors['MPAS-O']) if damping_coeffs is not None: ax.text(0.5, 5, 'r = ' + str(damping_coeffs[i])) # Plot comparison data if tidx is not None: plt.title(f'{atime:03f} days') for atime, ay in zip(self.times, locs): ax.text(1, ay, f'{atime} days', size=8, transform=ax.transAxes) for datatype in datatypes: datafile = f'./r{damping_coeffs[i]}d{atime}-'\ f'{datatype.lower()}.csv' if os.path.exists(datafile): data = pd.read_csv(datafile, header=None) ax.scatter(data[0], data[1], marker='.', color=colors[datatype], label=datatype) else: ax.text(1, ay, f'{atime} days', size=8, transform=ax.transAxes) for datatype in datatypes: datafile = f'./r{damping_coeffs[i]}d{atime}-'\ f'{datatype.lower()}.csv' if os.path.exists(datafile): data = pd.read_csv(datafile, header=None) ax.scatter(data[0], data[1], marker='.', color=colors[datatype], label=datatype) # Plot bottom depth, but line will not be visible unless bottom # depth is incorrect ax.plot(x, bottom_depth, ':k') ax.legend(frameon=False, loc='lower left') ds.close() h, l0 = ax.get_legend_handles_labels() ax.legend(h[:nhandles], l0[:nhandles], frameon=False, loc='lower left') fig.text(0.04, 0.5, 'Channel depth (m)', va='center', rotation='vertical') fig.text(0.5, 0.02, 'Along channel distance (km)', ha='center') filename = f'{outFolder}/ssh_depth_section' if tidx is not None: filename = f'{filename}_t{tidx:03d}' fig.savefig(f'{filename}.png', dpi=200, format='png') plt.close(fig) def _images_to_movies(self, outFolder='.', framesPerSecond=30, extension='mp4', overwrite=True): """ Convert all the image sequences into movies with ffmpeg """ try: os.makedirs(f'{outFolder}/logs') except OSError: pass framesPerSecond = str(framesPerSecond) prefix = 'ssh_depth_section' outFileName = f'{outFolder}/{prefix}.{extension}' if overwrite or not os.path.exists(outFileName): imageFileTemplate = f'{outFolder}/{prefix}_t%03d.png' logFileName = f'{outFolder}/logs/{prefix}.log' with open(logFileName, 'w') as logFile: args = ['ffmpeg', '-y', '-r', framesPerSecond, '-i', imageFileTemplate, '-b:v', '32000k', '-r', framesPerSecond, '-pix_fmt', 'yuv420p', outFileName] print_args = ' '.join(args) print(f'running {print_args}') subprocess.check_call(args, stdout=logFile, stderr=logFile)