# 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
# https://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