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

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

from polaris import Step
from polaris.viz import use_mplstyle


[docs]class Viz(Step): """ A step for plotting the results of a single-column test """
[docs] def __init__(self, component, indir, ideal_age=False): """ Create the step Parameters ---------- component : polaris.Component The component the step belongs to indir : str The subdirectory that the task belongs to, that this step will go into a subdirectory of ideal_age : bool, optional Whether the initial condition should include the ideal age tracer """ super().__init__(component=component, name='viz', indir=indir) self.ideal_age = ideal_age self.add_input_file( filename='initial_state.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 """ use_mplstyle() ideal_age = self.ideal_age ds = xr.load_dataset('output.nc') t_index = ds.sizes['Time'] - 1 t = ds.daysSinceStartOfSim[t_index] t_days = t.values.astype('timedelta64[D]') title = f'final time = {t_days / np.timedelta64(1, "D")} days' fields = {'temperature': 'degC', 'salinity': 'PSU'} if ideal_age: fields['iAge'] = 'seconds' z_mid = ds['zMid'].mean(dim='nCells') z_mid_init = z_mid.isel(Time=0) z_mid_final = z_mid.isel(Time=t_index) for field_name, field_units in fields.items(): if field_name not in ds.keys(): raise ValueError(f'{field_name} not present in output.nc') var = ds[field_name].mean(dim='nCells') var_init = var.isel(Time=0) var_final = var.isel(Time=t_index) plt.figure(figsize=(3, 5)) ax = plt.subplot(111) ax.plot(var_init, z_mid_init, '--k', label='initial') ax.plot(var_final, z_mid_final, '-k', label='final') ax.set_xlabel(f'{field_name} ({field_units})') ax.set_ylabel('z (m)') ax.legend() plt.title(title) plt.tight_layout(pad=0.5) plt.savefig(f'{field_name}.png') plt.close()