Source code for polaris.ocean.tasks.internal_wave.viz

import cmocean  # noqa: F401
import numpy as np
import xarray as xr

from polaris import Step
from polaris.ocean.viz import compute_transect, plot_transect


[docs]class Viz(Step): """ A step for visualizing a cross-section through the internal wave """
[docs] def __init__(self, component, indir): """ Create the step Parameters ---------- test_case : compass.TestCase The test case this step belongs to """ super().__init__(component=component, name='viz', indir=indir) self.add_input_file( filename='mesh.nc', target='../../../init/culled_mesh.nc') self.add_input_file( filename='init.nc', target='../../../init/initial_state.nc') self.add_input_file( filename='output.nc', target='../forward/output.nc')
[docs] def run(self): """ Run this step of the test case """ ds_mesh = xr.load_dataset('init.nc') ds_init = ds_mesh ds = xr.load_dataset('output.nc') x_mid = ds_mesh.xCell.median() y_min = ds_mesh.yCell.min() y_max = ds_mesh.yCell.max() x = xr.DataArray(data=[x_mid, x_mid], dims=('nPoints',)) y = xr.DataArray(data=[y_min, y_max], dims=('nPoints',)) vmin_temp = np.min(ds.temperature.values) vmax_temp = np.max(ds.temperature.values) vmax_v = np.max(np.abs(ds.vertVelocityTop.values)) tidx = 0 # Plot the initial time ds_transect = compute_transect( x=x, y=y, ds_horiz_mesh=ds_mesh, layer_thickness=ds_init.layerThickness.isel(Time=tidx), bottom_depth=ds_mesh.bottomDepth, min_level_cell=ds_mesh.minLevelCell - 1, max_level_cell=ds_mesh.maxLevelCell - 1, spherical=False) plot_transect(ds_transect, mpas_field=ds_init.temperature.isel(Time=tidx), out_filename='temperature_section_init.png', title='temperature', interface_color='grey', vmin=vmin_temp, vmax=vmax_temp, colorbar_label=r'$^{\circ}$C', cmap='cmo.thermal') tidx = -1 # Plot the final time ds_transect = compute_transect( x=x, y=y, ds_horiz_mesh=ds_mesh, layer_thickness=ds.layerThickness.isel(Time=tidx), bottom_depth=ds_mesh.bottomDepth, min_level_cell=ds_mesh.minLevelCell - 1, max_level_cell=ds_mesh.maxLevelCell - 1, spherical=False) plot_transect(ds_transect, mpas_field=ds.temperature.isel(Time=tidx), out_filename='temperature_section_final.png', title='temperature', interface_color='grey', vmin=vmin_temp, vmax=vmax_temp, colorbar_label=r'$^{\circ}$C', cmap='cmo.thermal') w_values = ds.vertVelocityTop.isel(Time=tidx).values[:, :-1] w = w_values * xr.ones_like(ds.temperature.isel(Time=tidx)) plot_transect( ds_transect, mpas_field=w, out_filename='vertical_velocity_section_final.png', title='vertical velocity', vmin=-vmax_v, vmax=vmax_v, colorbar_label='m/s', cmap='cmo.balance')