Notebook

Develop `tracer_thalweg_and_surface` Figure Module¶

Development of functions for nowcast.figures.research.tracer_thalweg_and_surface web site figure module.

This is an example of developing the functions for a web site figure module in a notebook. It follows the function organization patterns described in Creating a Figure Module docs.

Notebooks like this should be developed in a Nowcast Figures Development Environment so that all of the necessary dependency packages are installed. The development has to be done on a workstation that has the Salish Sea Nowcast system /results/ parition mounted.

If you choose to develop your figure in a notebook, the next step is to create a module like nowcast.figures.research.tracer_thalweg_and_surface and a notebook like TestTracerThalwegAndSurfaceModule that tests it in the nowcast system context.

In [1]:

import datetime
import os
from types import SimpleNamespace

import cmocean
import matplotlib.pyplot as plt
from matplotlib import gridspec
import netCDF4 as nc
import numpy as np
import scipy.io as sio

import salishsea_tools
from salishsea_tools import visualisations as vis
from salishsea_tools import tidetools, viz_tools

import nowcast
import nowcast.figures.website_theme

In [2]:

%matplotlib inline

We know for sure that we need to develop:

a _prep_plot_data() function
a _prep_fig_axes() function
a make_figure() function that the make_plots worker will call

Since the figure will have 2 axes (a vertical thalweg slice, and a horizontal surface slice) we also know that we need to develop 2 axes plotting functions:

_plot_tracer_thalweg()
_plot_tracer_surface()

`_prep_plot_data()` Function¶

In [3]:

def _prep_plot_data(tracer_var, mesh_mask, depth_integrated):
    hr = 19
    sj, ej = 200, 800
    si, ei = 20, 395
    
    tracer_hr = tracer_var[hr]
    masked_tracer_hr = np.ma.masked_where(
        mesh_mask["tmask"][0, ...] == 0, tracer_hr)
    surface_hr = masked_tracer_hr[0, sj:ej, si:ei]
    
    if depth_integrated:
        grid_heights = mesh_mask.variables['e3t_1d'][:][0].reshape(
            tracer_hr.shape[0], 1, 1)
        height_weighted = masked_tracer_hr[:, sj:ej, si:ei] * grid_heights
        surface_hr = height_weighted.sum(axis=0)
        
    return SimpleNamespace(
        tracer_var=tracer_var,
        tracer_hr=tracer_hr,
        surface_hr=surface_hr,
        surface_j_limits=(sj, ej),
        surface_i_limits=(si, ei),
        thalweg_depth_limits=(0, 450),
        thalweg_length_limits=(0, 632),
    )

In [4]:

bathy = nc.Dataset('/results/nowcast-sys/NEMO-forcing/grid/bathy_downonegrid2.nc')

mesh_mask = nc.Dataset('/results/nowcast-sys/NEMO-forcing/grid/mesh_mask_downbyone2.nc')

ptrc_T = nc.Dataset('/results/SalishSea/nowcast-green/25apr17/SalishSea_1h_20170425_20170425_ptrc_T.nc')

tracer_var = ptrc_T.variables['PHY2']

Confirm that _prep_plot_data() is doing what we want. Lots of tests are possible, but we'll just confirm that the shapes of the data arrays to be plotted are what we expect.

In [5]:

plot_data = _prep_plot_data(tracer_var, mesh_mask, depth_integrated=True)

print(plot_data.tracer_hr.shape)
print(plot_data.surface_hr.shape)

(40, 898, 398)
(570, 350)

`_prep_fig_axes() Function`¶

In [6]:

def _prep_fig_axes(figsize, theme):
    fig = plt.figure(
        figsize=figsize, facecolor=theme.COLOURS['figure']['facecolor'])

    gs = gridspec.GridSpec(1, 2, width_ratios=[1.618, 1])

    ax_thalweg = fig.add_subplot(gs[0])
    ax_thalweg.set_axis_bgcolor(theme.COLOURS['axes']['background'])

    ax_surface = fig.add_subplot(gs[1])
    ax_surface.set_axis_bgcolor(theme.COLOURS['axes']['background'])

    return fig, (ax_thalweg, ax_surface)

`_plot_tracer_thalweg()` Function¶

Because our figure is going to show 2 views of the tracer (a vertical slice along the thalweg, and a horizontal slice on the ocean surface) it would be nice to use the same contour colour levels if appropriate. But it won't be appropriate if the values contoured on the surface are depth averaged. It's also possible that the tracer concentration ranges will be different enough in the 2 slices that common contour levels don't make sense.

So, we come up with a function that calculates contour levels for the two axes and decides whether whether the levels are similar enough that one colour bar is sufficient for the figure, or if each axes requires one.

In [7]:

def _calc_clevels(plot_data):
    """Calculates contour levels for the two axes and decides whether whether
    the levels are similar enough that one colour bar is sufficient for the 
    figure, or if each axes requires one.
    """
    percent_98_surf = np.percentile(plot_data.surface_hr.compressed(), 98)
    percent_2_surf = np.percentile(plot_data.surface_hr.compressed(), 2)

    percent_98_grid = np.percentile(
        np.ma.masked_values(plot_data.tracer_hr, 0).compressed(), 98)
    percent_2_grid = np.percentile(
        np.ma.masked_values(plot_data.tracer_hr, 0).compressed(), 2)

    overlap = (
        max(0, min(percent_98_surf, percent_98_grid)
            - max(percent_2_surf, percent_2_grid)))
    magnitude = (
        (percent_98_surf - percent_2_surf) + (percent_98_grid - percent_2_grid))
    if 2 * overlap / magnitude > 0.5:
        max_clevel = max(percent_98_surf, percent_98_grid)
        min_clevel = min(percent_2_surf, percent_2_grid)
        clevels_thalweg = np.arange(
            min_clevel, max_clevel, (max_clevel - min_clevel) / 20.0)
        clevels_surface = clevels_thalweg
        show_thalweg_cbar = False
    else:
        clevels_thalweg = np.arange(
            percent_2_grid, percent_98_grid,
            (percent_98_grid - percent_2_grid) / 20.0)
        clevels_surface = np.arange(
            percent_2_surf, percent_98_surf,
            (percent_98_surf - percent_2_surf) / 20.0)
        show_thalweg_cbar = True
    return clevels_thalweg, clevels_surface, show_thalweg_cbar

In [8]:

def _plot_tracer_thalweg(ax, plot_data, bathy, mesh_mask, cmap, clevels):
    cbar = vis.contour_thalweg(
        ax, plot_data.tracer_hr, bathy, mesh_mask, clevels=clevels, cmap=cmap,
        thalweg_file='/results/nowcast-sys/tools/bathymetry/thalweg_working.txt',
        cbar_args={'fraction': 0.030, 'pad': 0.04, 'aspect': 45}
    )
    return cbar

`_thalweg_axes_labels()` Function¶

We also put the code for labeling and prettifying the axes in a separate function. Along the way, we pull the code that we use to label either colour bar out into a separate function that both _thalweg_axes_labels() and _surface_axes_labels() can make use of.

In [9]:

def _cbar_labels(cbar, contour_intervals, theme, label):
    cbar.set_ticks(contour_intervals)
    cbar.ax.axes.tick_params(labelcolor=theme.COLOURS['cbar']['tick labels'])
    cbar.set_label(
        label,
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis'])

In [10]:

def _thalweg_axes_labels(
    ax, plot_data, show_thalweg_cbar, clevels, cbar, theme
):
    ax.set_xlim(plot_data.thalweg_length_limits)
    ax.set_ylim(
        plot_data.thalweg_depth_limits[1], plot_data.thalweg_depth_limits[0])
    if show_thalweg_cbar:
        label = (
            f'{plot_data.tracer_var.long_name} [{plot_data.tracer_var.units}]')
        _cbar_labels(cbar, clevels[::2], theme, label)
    else:
        cbar.remove()
    ax.set_xlabel(
        'Distance along thalweg [km]', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_ylabel(
        'Depth [m]', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    theme.set_axis_colors(ax)

`_plot_tracer_surface()` Function¶

In [11]:

def _plot_tracer_surface(ax, plot_data, cmap, clevels):
    x, y = np.meshgrid(
        np.arange(*plot_data.surface_i_limits, dtype=int),
        np.arange(*plot_data.surface_j_limits, dtype=int))
    mesh = ax.contourf(
        x, y, plot_data.surface_hr, levels=clevels, cmap=cmap, extend='both')
    cbar = plt.colorbar(mesh, ax=ax, fraction=0.034, pad=0.04, aspect=45)
    return cbar

`_surface_axes_labels()` Function¶

In [12]:

def _surface_axes_labels(
    ax, tracer_var, depth_integrated, clevels, cbar, theme
):
    cbar_units = (
        f'{tracer_var.units}*m' if depth_integrated
        else f'{tracer_var.units}')
    cbar_label = f'{tracer_var.long_name} [{cbar_units}]'
    _cbar_labels(cbar, clevels[::2], theme, cbar_label)
    ax.set_xlabel(
        'Grid x', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_ylabel(
        'Grid y', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_axis_bgcolor('burlywood')
    viz_tools.set_aspect(ax)
    theme.set_axis_colors(ax)

`make_figure()` Function¶

This is is the function that will be called by the nowcast.workers.make_plots worker to return a matplotlib.figure.Figure object.

In [13]:

def make_figure(
    tracer_var, bathy, mesh_mask, cmap,  depth_integrated,
    figsize=(16, 9), theme=nowcast.figures.website_theme
):
    plot_data = _prep_plot_data(tracer_var, mesh_mask, depth_integrated)
    fig, (ax_thalweg, ax_surface) = _prep_fig_axes(figsize, theme)

    clevels_thalweg, clevels_surface, show_thalweg_cbar = _calc_clevels(
        plot_data)

    cbar_thalweg = _plot_tracer_thalweg(
        ax_thalweg, plot_data, bathy, mesh_mask, cmap, clevels_thalweg)
    _thalweg_axes_labels(
        ax_thalweg, plot_data, show_thalweg_cbar, clevels_thalweg,
        cbar_thalweg, theme)

    cbar_surface = _plot_tracer_surface(
        ax_surface, plot_data, cmap, clevels_surface)
    _surface_axes_labels(
        ax_surface, tracer_var, depth_integrated, clevels_surface, cbar_surface,
        theme)
    return fig

Render the Figure¶

The %%timeit cell magic lets us keep an eye on how log the figure takes to process. Setting -n1 -r1 prevents it from processing the figure more than once as it might try to do to generate better statistics.

In [14]:

%%timeit -n1 -r1

fig = make_figure(
    tracer_var, bathy, mesh_mask, cmap=cmocean.cm.algae, depth_integrated=True)

1.93 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [15]:

tracers = {
    'DOC': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'MICZ': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'MYRI': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'NH4': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'NO3': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'O2': {'cmap': cmocean.cm.oxy, 'depth_integrated': False},
    'PHY': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'PHY2': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'POC': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'Si': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'bSi': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
}
for tracer in tracers:
    tracer_var = ptrc_T.variables[tracer]
    fig = make_figure(
        tracer_var, bathy, mesh_mask, 
        cmap=tracers[tracer]['cmap'], 
        depth_integrated=tracers[tracer]['depth_integrated']
    )

In [ ]:

Develop tracer_thalweg_and_surface Figure Module¶

_prep_plot_data() Function¶

_prep_fig_axes() Function¶

_plot_tracer_thalweg() Function¶

_thalweg_axes_labels() Function¶

_plot_tracer_surface() Function¶

_surface_axes_labels() Function¶

make_figure() Function¶