# Copyright 2016-2021 The Salish Sea NEMO Project and
# The University of British Columbia
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Functions for common model visualisations
"""
import datetime
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import patches
from salishsea_tools import geo_tools, nc_tools
[docs]
def contour_thalweg(
axes, var, bathy, mesh_mask, clevels=None,
mesh_mask_depth_var='gdept_0', cmap='hsv', land_colour='burlywood',
xcoord_distance=True,
thalweg_file='/home/sallen/MEOPAR/Tools/bathymetry/thalweg_working.txt',
cbar_args=None,
mesh_args=None,
method='contourf'
):
"""Contour the data stored in var along the domain thalweg.
:arg axes: Axes instance to plot thalweg contour on.
:type axes: :py:class:`matplotlib.axes.Axes`
:arg var: Salish Sea NEMO model results variable to be contoured
:type var: :py:class:`numpy.ndarray`
:arg bathy: Salish Sea NEMO model bathymetry dataset
:type bathy: :py:class:`netCDF4.Dataset`
:arg mesh_mask: Salish Sea NEMO model mesh_mask dataset
:type mesh_mask: :py:class:`netCDF4.Dataset`
:arg clevels: argument for determining contour levels. Choices are
1. 'salinity' or 'temperature' for pre-determined levels
used in nowcast.
2. an integer N, for N automatically determined levels.
3. a sequence V of contour levels, which must be in
increasing order.
:type clevels: str or int or iterable
:arg str mesh_mask_depth_var: name of depth variable in :kbd:`mesh_mask`
that is appropriate for :kbd:`var`;
defaults to :kbd:`gdept_0` for NEMO-3.6
tracer variables.
:arg str cmap: matplotlib colormap
:arg str land_colour: matplotlib colour for land
:arg xcoord_distance: plot along thalweg distance (True) or index (False)
:type xcoord_distance: boolean
:arg thalweg_file: Path and file name to read the array of
thalweg grid points from.
:type thalweg_file: str
:arg dict cbar_args: Additional arguments to be passed to the cbar
function (fraction, pad, etc.)
:arg dict mesh_args: Additional arguments to be passed to the contourf
or pcolormesh function
: arg string method: method to use for data display: defaults to
'contourf' but 'pcolormesh' is also accepted
:returns: matplotlib colorbar object
"""
thalweg_pts = np.loadtxt(thalweg_file, delimiter=' ', dtype=int)
depth = mesh_mask.variables[mesh_mask_depth_var][:]
dep_thal, distance, var_thal = load_thalweg(
depth[0, ...], var, bathy['nav_lon'][:], bathy['nav_lat'][:],
thalweg_pts)
if xcoord_distance:
xx_thal = distance
axes.set_xlabel('Distance along thalweg [km]')
else:
xx_thal, _ = np.meshgrid(np.arange(var_thal.shape[-1]), dep_thal[:, 0])
axes.set_xlabel('Thalweg index')
# Determine contour levels
clevels_default = {
'salinity': [
26, 27, 28, 29, 30, 30.2, 30.4, 30.6, 30.8, 31, 32, 33, 34
],
'temperature': [
6.9, 7, 7.5, 8, 8.5, 9, 9.8, 9.9, 10.3, 10.5, 11, 11.5, 12,
13, 14, 15, 16, 17, 18, 19
]
}
if isinstance(clevels, str):
try:
clevels = clevels_default[clevels]
except KeyError:
raise KeyError('no default clevels defined for {}'.format(clevels))
# Prepare for plotting by filling in grid points just above bathymetry
var_plot = _fill_in_bathy(var_thal, mesh_mask, thalweg_pts)
if method == 'pcolormesh':
if mesh_args is None:
mesh = axes.pcolormesh(xx_thal, dep_thal, var_plot, cmap=cmap)
else:
mesh = axes.pcolormesh(xx_thal, dep_thal, var_plot, cmap=cmap, **mesh_args)
axes.set_xlim((np.min(xx_thal),np.max(xx_thal)))
else:
if mesh_args is None:
mesh = axes.contourf(xx_thal, dep_thal, var_plot, clevels, cmap=cmap,
extend='both')
else:
mesh = axes.contourf(xx_thal, dep_thal, var_plot, clevels, cmap=cmap,
extend='both', **mesh_args)
_add_bathy_patch(
xx_thal, bathy['Bathymetry'][:], thalweg_pts, axes, color=land_colour)
if cbar_args is None:
cbar = plt.colorbar(mesh, ax=axes)
else:
cbar = plt.colorbar(mesh, ax=axes, **cbar_args)
axes.invert_yaxis()
axes.set_ylabel('Depth [m]')
return cbar
def _add_bathy_patch(xcoord, bathy, thalweg_pts, ax, color, zmin=450):
"""Add a polygon shaped as the land in the thalweg section
:arg xcoord: x grid along thalweg
:type xcoord: 2D numpy array
:arg bathy: Salish Sea NEMO model bathymetry data
:type bathy: :py:class:`numpy.ndarray`
:arg thalweg_pts: Salish Sea NEMO model grid indices along thalweg
:type thalweg_pts: 2D numpy array
:arg ax: Axes instance to plot thalweg contour on.
:type ax: :py:class:`matplotlib.axes.Axes`
:arg str color: color of bathymetry patch
:arg zmin: minimum depth for plot in meters
:type zmin: float
"""
# Look up bottom bathymetry along thalweg
thalweg_bottom = bathy[thalweg_pts[:, 0], thalweg_pts[:, 1]]
# Construct bathy polygon
poly = np.zeros((thalweg_bottom.shape[0]+2, 2))
poly[0, :] = 0, zmin
poly[1:-1, 0] = xcoord[0, :]
poly[1:-1:, 1] = thalweg_bottom
poly[-1, :] = xcoord[0, -1], zmin
ax.add_patch(patches.Polygon(poly, facecolor=color, edgecolor=color))
[docs]
def load_thalweg(depths, var, lons, lats, thalweg_pts):
"""Returns depths, cumulative distance and variable along thalweg.
:arg depths: depth array for variable. Can be 1D or 3D.
:type depths: :py:class:`numpy.ndarray`
:arg var: 3D Salish Sea NEMO model results variable
:type var: :py:class:`numpy.ndarray`
:arg lons: Salish Sea NEMO model longitude grid data
:type lons: :py:class:`numpy.ndarray`
:arg lats: Salish Sea NEMO model latitude grid data
:type lats: :py:class:`numpy.ndarray`
:arg thalweg_pts: Salish Sea NEMO model grid indices along thalweg
:type thalweg_pts: 2D numpy array
:returns: dep_thal, xx_thal, var_thal, all the same shape
(depth, thalweg length)
"""
lons_thal = lons[thalweg_pts[:, 0], thalweg_pts[:, 1]]
lats_thal = lats[thalweg_pts[:, 0], thalweg_pts[:, 1]]
var_thal = var[:, thalweg_pts[:, 0], thalweg_pts[:, 1]]
xx_thal = geo_tools.distance_along_curve(lons_thal, lats_thal)
xx_thal = xx_thal + np.zeros(var_thal.shape)
if depths.ndim > 1:
dep_thal = depths[:, thalweg_pts[:, 0], thalweg_pts[:, 1]]
else:
_, dep_thal = np.meshgrid(xx_thal[0, :], depths)
return dep_thal, xx_thal, var_thal
def _fill_in_bathy(variable, mesh_mask, thalweg_pts):
"""For each horizontal point in variable, fill in first vertically masked
point with the value just above.
Use mbathy in mesh_mask file to determine level of vertical masking
:arg variable: the variable to be filled
:type variable: 2D numpy array
:arg mesh_mask: Salish Sea NEMO model mesh_mask data
:type mesh_mask: :py:class:`netCDF4.Dataset`
:arg thalweg_pts: Salish Sea NEMO model grid indices along thalweg
:type thalweg_pts: 2D numpy array
:returns: newvar, the filled numpy array
"""
mbathy = mesh_mask.variables['mbathy'][0, :, :]
newvar = np.copy(variable)
mbathy = mbathy[thalweg_pts[:, 0], thalweg_pts[:, 1]]
for i, level in enumerate(mbathy):
newvar[level, i] = variable[level-1, i]
return newvar
[docs]
def contour_layer_grid(axes,data,mask,clevels=10,lat=None,lon=None,cmap=None,var_name=None,
land_colour='burlywood',is_depth_avg=False,is_pcolmesh=False,title='',cbar_args=None,
):
"""Contour 2d data at an arbitrary klevel on the model grid
:arg axes: Axes instance to plot thalweg contour on.
:type axes: :py:class:`matplotlib.axes.Axes`
:arg data: 2D array to be contoured at level k
:type data: :py:class:`numpy.ndarray`
:arg klev: Index of k-level along which to contour
:type klev: int
:arg mask: Mask array with same dimensions as data
:type mask: :py:class:`numpy.ndarray`
:arg clevels: argument for determining contour levels. Choices are
1. an integer N, for N automatically determined levels.
2. a sequence V of contour levels, which must be in
increasing order.
:type clevels: str or int or iterable
:arg lon: Array of longitudes with same length as x dimension of data.
:type lon: :py:class:`numpy.ndarray`
:arg lat: Array of longitudes with same length as x dimension of data.
:type lat: :py:class:`numpy.ndarray`
:arg str cmap: matplotlib colormap
:arg str var_name: Name of variable to plot. Necesssary if cmap=None.
:arg str land_colour: matplotlib colour for land
:arg is_depth_avg: True if data is a depth averaged field (default is False).
:type is_depth_avg: boolean
:arg is_pcolmesh: plot a pcolormesh (True) instead of a contourf (default).
:type is_pcolmesh: boolean
:arg str title: Title string
:arg dict cbar_args: Additional arguments to be passed to the cbar
function (fraction, pad, etc.)
:returns: matplotlib colorbar object
"""
mdata = np.ma.masked_where(mask==0,data)
viz_tools.set_aspect(axes)
if cmap == None:
cbMIN, cbMAX, cmap = visualisations.retrieve_cmap(var_name,is_depth_avg)
cmap = plt.get_cmap(cmocean.cm.algae)
if is_pcolmesh:
mesh = axes.pcolormesh(mdata, cmap=cmap)
else:
mesh= axes.contourf(mdata,clevels,cmap=cmap)
axes.set_xlabel('X index')
axes.set_ylabel('Y index')
axes.set_title(title)
axes.set_axis_bgcolor(land_colour)
if cbar_args is None:
cbar = plt.colorbar(mesh, ax=axes)
else:
cbar = plt.colorbar(mesh, ax=axes, **cbar_args)
return cbar
[docs]
def plot_drifters(ax, DATA, DRIFT_OBJS=None, color='red', cutoff=24, zorder=15):
"""Plot a drifter track from ODL Drifter observations.
:arg time_ind: Time index (current drifter position, track will be visible
up until this point, ex. 'YYYY-mmm-dd HH:MM:SS', format is flexible)
:type time_ind: str or :py:class:`datetime.datetime`
:arg ax: Axis object
:type ax: :py:class:`matplotlib.pyplot.axes`
:arg DATA: Drifter track dataset
:type DATA: :py:class:`xarray.Dataset`
:arg color: Drifter track color
:type color: str
:arg cutoff: Time threshold for color plotting (hours)
:type cutoff: integer
:arg zorder: Plotting layer specifier
:type zorder: integer
:returns: Dictionary of line objects
:rtype: dict > :py:class:`matplotlib.lines.Line2D`
"""
if DATA.time.shape[0] > 0:
# Convert time boundaries to datetime.datetime to allow operations/slicing
starttime = nc_tools.xarraytime_to_datetime(DATA.time[ 0])
endtime = nc_tools.xarraytime_to_datetime(DATA.time[-1])
# Color plot cutoff
time_cutoff = endtime - datetime.timedelta(hours=cutoff)
if DRIFT_OBJS is not None: # --- Update line objects only
# Plot drifter track (gray)
DRIFT_OBJS['L_old'][0].set_data(
DATA.lon.sel(time=slice(starttime, time_cutoff)),
DATA.lat.sel(time=slice(starttime, time_cutoff)))
# Plot drifter track (color)
DRIFT_OBJS['L_new'][0].set_data(
DATA.lon.sel(time=slice(time_cutoff, endtime)),
DATA.lat.sel(time=slice(time_cutoff, endtime)))
# Plot drifter position
DRIFT_OBJS['P'][0].set_data(
DATA.lon.sel(time=endtime, method='nearest'),
DATA.lat.sel(time=endtime, method='nearest'))
else: # ------------------------ Plot new line objects instances
# Define drifter objects dict
DRIFT_OBJS = {}
# Plot drifter track (gray)
DRIFT_OBJS['L_old'] = ax.plot(
DATA.lon.sel(time=slice(starttime, time_cutoff)),
DATA.lat.sel(time=slice(starttime, time_cutoff)),
'-', linewidth=2, color='gray', zorder=zorder)
# Plot drifter track (color)
DRIFT_OBJS['L_new'] = ax.plot(
DATA.lon.sel(time=slice(time_cutoff, endtime)),
DATA.lat.sel(time=slice(time_cutoff, endtime)),
'-', linewidth=2, color=color, zorder=zorder+1)
# Plot drifter position
DRIFT_OBJS['P'] = ax.plot(
DATA.lon.sel(time=endtime, method='nearest'),
DATA.lat.sel(time=endtime, method='nearest'),
'o', color=color, zorder=zorder+2)
else:
if DRIFT_OBJS is not None: # --- Update line objects only
# Update drifter tracks
DRIFT_OBJS['L_old'][0].set_data([], []) # gray
DRIFT_OBJS['L_new'][0].set_data([], []) # color
DRIFT_OBJS['P' ][0].set_data([], []) # position
else:
DRIFT_OBJS = {}
DRIFT_OBJS['L_old'] = ax.plot([], [], '-',
linewidth=2, color='gray', zorder=zorder)
# Plot drifter track (color)
DRIFT_OBJS['L_new'] = ax.plot([], [], '-',
linewidth=2, color=color, zorder=zorder+1)
# Plot drifter position
DRIFT_OBJS['P'] = ax.plot([], [], 'o', color=color, zorder=zorder+2)
return DRIFT_OBJS
[docs]
def plot_tracers(
ax, qty, DATA, C=None, coords='map', clim=[0, 35, 1], cmap='jet', zorder=0
):
"""Plot a horizontal slice of NEMO tracers as filled contours.
.. note::
This function is deprecated.
Plot NEMO results directly using `matplotlib.pyplot.contourf` or
equivalent instead.
"""
raise DeprecationWarning(
'plot_tracers has been deprecated. Plot NEMO results directly using '
'matplotlib.pyplot.contourf or equivalent instead.')
[docs]
def plot_velocity(
ax, model, DATA, Q=None, coords='map', processed=False, spacing=5,
mask=True, color='black', scale=20, headwidth=3, linewidth=0, zorder=5
):
"""Plot a horizontal slice of NEMO or GEM velocities as quiver objects.
Accepts subsampled u and v fields via the **processed** keyword
argument.
.. note::
This function is deprecated.
Plot NEMO results directly using `matplotlib.pyplot.quiver` or
equivalent instead.
"""
raise DeprecationWarning(
'plot_velocity has been deprecated. Plot NEMO results directly using '
'matplotlib.pyplot.quiver or equivalent instead.')
[docs]
def retrieve_cmap(varname,deep_bool):
"""takes 2 args:
string varname - name of a variable from nowcast-green output
boolean deep_bool - indicates whether the variable is depth-integrated or not
returns 2 ints(min and max value of range), and string identifying cmap"""
var_namemap ={'Fraser_tracer': {'varname':'Fraser_tracer'},
'ammonium': {'varname':'NH4'},
'NH4': {'varname':'NH4'},
'biogenic_silicon': {'varname':'bSi'},
'bSi': {'varname':'bSi'},
'ciliates': {'varname':'MYRI'},
'MYRI': {'varname':'MYRI'},
'diatoms': {'varname':'PHY2'},
'PHY2': {'varname':'PHY2'},
'dissolved_organic_nitrogen': {'varname':'dissolved_organic_nitrogen'},
'flagellates': {'varname':'PHY'},
'PHY': {'varname':'PHY'},
'mesozooplankton': {'varname':'MESZ'},
'MESZ': {'varname':'MESZ'},
'microzooplankton': {'varname':'MICZ'},
'MICZ': {'varname':'MICZ'},
'nitrate': {'varname':'NO3'},
'NO3': {'varname':'NO3'},
'particulate_organic_nitrogen': {'varname':'PON'},
'POC': {'varname':'PON'},
'PON': {'varname':'PON'},
'dissolved_organic_nitrogen': {'varname':'DON'},
'DOC': {'varname':'DON'},
'DON': {'varname':'DON'},
'silicon': {'varname':'Si'},
'Si': {'varname':'Si'}}
#dictionary of colour ranges
var_colour_ranges = {
'Fraser_tracer':{'colorBarMinimum': 0.0, 'colorBarMaximum': 140.0,'cmap': 'turbid'},
'MESZ': {'colorBarMinimum': 0.0, 'colorBarMaximum': 3.0,'cmap': 'algae'},
'MICZ': {'colorBarMinimum': 0.0, 'colorBarMaximum': 4.0,'cmap': 'algae'},
'MYRI': {'colorBarMinimum': 0.0, 'colorBarMaximum': 5.0,'cmap': 'algae'},
'NH4': {'colorBarMinimum': 0.0, 'colorBarMaximum': 10.0,'cmap': 'matter'},
'NO3': {'colorBarMinimum': 0.0, 'colorBarMaximum': 40.0,'cmap': 'tempo'},
'PON': {'colorBarMinimum': 0.0, 'colorBarMaximum': 2.0,'cmap': 'amp'},
'DON': {'colorBarMinimum': 0.0, 'colorBarMaximum': 20.0,'cmap': 'amp'},
'O2': {'colorBarMinimum': 0.0, 'colorBarMaximum': 140.0,'cmap': 'turbid'},
'PHY': {'colorBarMinimum': 0.0, 'colorBarMaximum': 6.0, 'cmap': 'algae'},
'PHY2': {'colorBarMinimum': 0.0, 'colorBarMaximum': 15.0,'cmap': 'algae'},
'Si': {'colorBarMinimum': 0.0, 'colorBarMaximum': 70.0,'cmap': 'turbid'},
'bSi': {'colorBarMinimum': 0.0, 'colorBarMaximum': 70.0,'cmap': 'turbid'},
'Fraser_tracer_int':{'colorBarMinimum': 0.0, 'colorBarMaximum': 6500,'cmap': 'turbid'},
'MESZ_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 140,'cmap': 'algae'},
'MICZ_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 350,'cmap': 'algae'},
'MYRI_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 75,'cmap': 'algae'},
'NH4_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 1500,'cmap': 'matter'},
'NO3_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 24000,'cmap': 'tempo'},
'PON_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 600,'cmap': 'amp'},
'DON_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 2500,'cmap': 'amp'},
'O2_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 1000,'cmap': 'turbid'},
'PHY_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 100,'cmap': 'algae'},
'PHY2_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 350,'cmap': 'algae'},
'Si_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 40000,'cmap': 'turbid'},
'bSi_int': {'colorBarMinimum': 0.0, 'colorBarMaximum': 40000,'cmap': 'turbid'},
}
dp = var_namemap[varname]
vn = dp['varname']
if (deep_bool == True):
vn = vn + '_int'
dict_pull = var_colour_ranges[vn]
cbMIN = dict_pull['colorBarMinimum']
print()
cbMAX = dict_pull['colorBarMaximum']
cmap_name = dict_pull['cmap']
return cbMIN, cbMAX, cmap_name