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

import cmocean  # noqa: F401
import xarray as xr

from polaris import Step
from polaris.viz import plot_global_mpas_field


[docs]class Viz(Step): """ A step for plotting fields from the cosine bell output Attributes ---------- mesh_name : str The name of the mesh """
[docs] def __init__(self, component, name, subdir, base_mesh, init, forward, mesh_name): """ Create the step Parameters ---------- component : polaris.Component The component the step belongs to name : str The name of the step subdir : str The subdirectory in the test case's work directory for the step base_mesh : polaris.Step The base mesh step init : polaris.Step The init step forward : polaris.Step The init step mesh_name : str The name of the mesh """ super().__init__(component=component, name=name, subdir=subdir) self.add_input_file( filename='mesh.nc', work_dir_target=f'{base_mesh.path}/base_mesh.nc') self.add_input_file( filename='initial_state.nc', work_dir_target=f'{init.path}/initial_state.nc') self.add_input_file( filename='output.nc', work_dir_target=f'{forward.path}/output.nc') self.mesh_name = mesh_name
[docs] def run(self): """ Run this step of the test case """ config = self.config mesh_name = self.mesh_name run_duration = \ config.getfloat('convergence_forward', 'run_duration') / 24. colormap_sections = dict( h='geostrophic_viz_h', u='geostrophic_viz_vel', v='geostrophic_viz_vel') ds_init = xr.open_dataset('initial_state.nc') bottom_depth = ds_init.bottomDepth ds_init = self._process_ds(ds_init, bottom_depth, time_index=0) ds_init.to_netcdf('remapped_init.nc') for var, colormap_section in colormap_sections.items(): plot_global_mpas_field( mesh_filename='mesh.nc', da=ds_init[var], out_filename=f'init_{var}.png', config=config, colormap_section=colormap_section, colorbar_label=config.get(colormap_section, 'label'), title=f'{mesh_name} {var} at init', plot_land=False) ds_out = xr.open_dataset('output.nc') ds_out = self._process_ds(ds_out, bottom_depth, time_index=-1) ds_out.to_netcdf('remapped_final.nc') for var, colormap_section in colormap_sections.items(): plot_global_mpas_field( mesh_filename='mesh.nc', da=ds_out[var], out_filename=f'final_{var}.png', config=config, colormap_section=colormap_section, colorbar_label=config.get(colormap_section, 'label'), title=f'{mesh_name} {var} after {run_duration:g} days', plot_land=False) colormap_sections = dict( h='geostrophic_viz_diff_h', u='geostrophic_viz_diff_vel', v='geostrophic_viz_diff_vel') for var, colormap_section in colormap_sections.items(): diff = ds_out[var] - ds_init[var] plot_global_mpas_field( mesh_filename='mesh.nc', da=diff, out_filename=f'diff_{var}.png', config=config, colormap_section=colormap_section, colorbar_label=config.get(colormap_section, 'label'), title=f'{mesh_name} {var} change after {run_duration:g} days', plot_land=False)
@staticmethod def _process_ds(ds, bottom_depth, time_index): ds_out = ds.isel(Time=time_index, nVertLevels=0) ds_out['h'] = ds_out.ssh + bottom_depth ds_out = ds_out.rename(dict(velocityZonal='u', velocityMeridional='v',)) return ds_out