"""
Stores the class for ContourDisplay.
"""
import numpy as np
from scipy.interpolate import Rbf
# Import Local Libs
from .plot import Display
[docs]class ContourDisplay(Display):
"""
This subclass contains routines that are specific to plotting
contour plots from data. It is inherited from Display and therefore
contains all of Display's attributes and methods.
"""
def __init__(self, ds, subplot_shape=(1,), ds_name=None, **kwargs):
super().__init__(ds, subplot_shape, ds_name, **kwargs)
[docs] def create_contour(
self,
fields=None,
time=None,
function='cubic',
subplot_index=(0,),
contour='contourf',
grid_delta=(0.01, 0.01),
grid_buffer=0.1,
twod_dim_value=None,
**kwargs,
):
"""
Extracts, grids, and creates a contour plot. If subplots have not been
added yet, an axis will be created assuming that there is only going
to be one plot.
Parameters
----------
fields : dict
Dictionary of fields to use for x, y, and z data.
time : datetime
Time in which to slice through the datasets.
function : string
Defaults to cubic function for interpolation.
See scipy.interpolate.Rbf for additional options.
subplot_index : 1 or 2D tuple, list, or array
The index of the subplot to set the x range of.
contour : str
Whether to use contour or contourf as plotting function.
Default is contourf
grid_delta : 1D tuple, list, or array
x and y deltas for creating grid.
grid_buffer : float
Buffer to apply to grid.
twod_dim_value : float
If the field is 2D, which dimension value to pull.
I.e. if dim is depths of [5, 10, 50, 100] specifying 50
would index the data at 50
**kwargs : keyword arguments
The keyword arguments for :func:`plt.contour`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
# Get x, y, and z data by looping through each dictionary
# item and extracting data from appropriate time
x = []
y = []
z = []
for dsname in self._ds:
ds = self._ds[dsname]
if dsname not in fields:
continue
field = fields[dsname]
if ds[field[2]].sel(time=time).values.size > 1:
dim_values = ds[ds[field[2]].dims[1]].values
if twod_dim_value is None:
dim_index = 0
else:
dim_index = np.where(dim_values == twod_dim_value)
if dim_index[0].size == 0:
continue
if np.isnan(ds[field[2]].sel(time=time).values[dim_index]):
continue
z.append(ds[field[2]].sel(time=time).values[dim_index].tolist())
else:
if np.isnan(ds[field[2]].sel(time=time).values):
continue
z.append(ds[field[2]].sel(time=time).values.tolist())
if ds[field[0]].values.size > 1:
x.append(ds[field[0]].sel(time=time).values.tolist())
else:
x.append(ds[field[0]].values.tolist())
if ds[field[1]].values.size > 1:
y.append(ds[field[1]].sel(time=time).values.tolist())
else:
y.append(ds[field[1]].values.tolist())
# Create a meshgrid for gridding onto
xs = np.arange(np.min(x) - grid_buffer, np.max(x) + grid_buffer, grid_delta[0])
ys = np.arange(np.min(y) - grid_buffer, np.max(y) + grid_buffer, grid_delta[1])
xi, yi = np.meshgrid(xs, ys)
# Use scipy radial basis function to interpolate data onto grid
rbf = Rbf(x, y, z, function=function)
zi = rbf(xi, yi)
# Create contour plot
if contour == 'contour':
self.contour(xi, yi, zi, subplot_index=subplot_index, **kwargs)
elif 'contourf':
self.contourf(xi, yi, zi, subplot_index=subplot_index, **kwargs)
else:
raise ValueError(
"Invalid contour plot type. Please choose either 'contourf' or 'contour'"
)
return self.axes[subplot_index]
[docs] def contourf(self, x, y, z, subplot_index=(0,), **kwargs):
"""
Base function for filled contours if user already has data gridded.
If subplots have not been added yet, an axis will be created
assuming that there is only going to be one plot.
Parameters
----------
x : array-like
x coordinates or grid for z.
y : array-like
y coordinates or grid for z.
z : array-like(x,y)
Values over which to contour.
subplot_index : 1 or 2D tuple, list, or array
The index of the subplot to set the x range of.
**kwargs : keyword arguments
The keyword arguments for :func:`plt.contourf`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
self.axes[subplot_index].contourf(x, y, z, **kwargs)
return self.axes[subplot_index]
[docs] def contour(self, x, y, z, subplot_index=(0,), **kwargs):
"""
Base function for contours if user already has data gridded.
If subplots have not been added yet, an axis will be created
assuming that there is only going to be one plot.
Parameters
----------
x : array-like
x coordinates or grid for z.
y : array-like
y coordinates or grid for z.
z : array-like(x, y)
Values over which to contour.
subplot_index : 1 or 2D tuple, list, or array
The index of the subplot to set the x range of.
**kwargs : keyword arguments
The keyword arguments for :func:`plt.contour`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
self.axes[subplot_index].contour(x, y, z, **kwargs)
return self.axes[subplot_index]
[docs] def plot_vectors_from_spd_dir(
self,
fields,
time=None,
subplot_index=(0,),
mesh=False,
function='cubic',
grid_delta=(0.01, 0.01),
grid_buffer=0.1,
**kwargs,
):
"""
Extracts, grids, and creates a contour plot.
If subplots have not been added yet, an axis will be created
assuming that there is only going to be one plot.
Parameters
----------
fields : dict
Dictionary of fields to use for x, y, and z data.
time : datetime
Time in which to slice through datasets.
mesh : boolean
Set to True to interpolate u and v to grid and create wind barbs.
function : string
Defaults to cubic function for interpolation.
See scipy.interpolate.Rbf for additional options.
grid_delta : 1D tuple, list, or array
x and y deltas for creating grid.
grid_buffer : float
Buffer to apply to grid.
**kwargs : keyword arguments
The keyword arguments for :func:`plt.barbs`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
# Get x, y, and z data by looping through each dictionary
# item and extracting data from appropriate time
x = []
y = []
wspd = []
wdir = []
for dsname in self._ds:
ds = self._ds[dsname]
field = fields[dsname]
if ds[field[0]].values.size > 1:
x.append(ds[field[0]].sel(time=time).values.tolist())
else:
x.append(ds[field[0]].values.tolist())
if ds[field[1]].values.size > 1:
y.append(ds[field[1]].sel(time=time).values.tolist())
else:
y.append(ds[field[1]].values.tolist())
wspd.append(ds[field[2]].sel(time=time).values.tolist())
wdir.append(ds[field[3]].sel(time=time).values.tolist())
# Calculate u and v
tempu = -np.sin(np.deg2rad(wdir)) * wspd
tempv = -np.cos(np.deg2rad(wdir)) * wspd
if mesh is True:
# Create a meshgrid for gridding onto
xs = np.arange(min(x) - grid_buffer, max(x) + grid_buffer, grid_delta[0])
ys = np.arange(min(y) - grid_buffer, max(y) + grid_buffer, grid_delta[1])
xi, yi = np.meshgrid(xs, ys)
# Use scipy radial basis function to interpolate data onto grid
rbf = Rbf(x, y, tempu, function=function)
u = rbf(xi, yi)
rbf = Rbf(x, y, tempv, function=function)
v = rbf(xi, yi)
else:
xi = x
yi = y
u = tempu
v = tempv
self.barbs(xi, yi, u, v, **kwargs)
return self.axes[subplot_index]
[docs] def barbs(self, x, y, u, v, subplot_index=(0,), **kwargs):
"""
Base function for wind barbs. If subplots have not been added yet,
an axis will be created assuming that there is only going to be
one plot.
Parameters
----------
x : array-like
x coordinates or grid for z.
y : array-like
y coordinates or grid for z.
u : array-like
U component of vector.
v : array-like
V component of vector.
subplot_index : 1 or 2D tuple, list, or array
The index of the subplot to set the x range of.
**kwargs : keyword arguments
The keyword arguments for :func:`plt.barbs`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
self.axes[subplot_index].barbs(x, y, u, v, **kwargs)
return self.axes[subplot_index]
[docs] def plot_station(self, fields, time=None, subplot_index=(0,), text_color='white', **kwargs):
"""
Extracts, grids, and creates a contour plot. If subplots have not
been added yet, an axis will be created assuming that there is only
going to be one plot.
Parameters
----------
fields : dict
Dictionary of fields to use for x, y, and z data.
time : datetime
Time in which to slice through datasets.
subplot_index : 1 or 2D tuple, list, or array
The index of the subplot to set the x range of.
text_color : string
Color for text.
**kwargs : keyword arguments
The keyword arguments for :func:`plt.plot`
Returns
-------
ax : matplotlib axis handle
The matplotlib axis handle of the plot.
"""
# Get x, y, and data by looping through each dictionary
# item and extracting data from appropriate time
for dsname in self._ds:
ds = self._ds[dsname]
field = fields[dsname]
for i, f in enumerate(field):
if i == 0:
if ds[f].values.size > 1:
x = ds[f].sel(time=time).values.tolist()
else:
x = ds[f].values.tolist()
elif i == 1:
if ds[f].values.size > 1:
y = ds[f].sel(time=time).values.tolist()
else:
y = ds[f].values.tolist()
self.axes[subplot_index].plot(x, y, '*', **kwargs)
else:
data = ds[f].sel(time=time).values.tolist()
offset = 0.02
if i == 2:
x1 = x - 3.0 * offset
y1 = y + offset
if i == 3:
x1 = x + offset
y1 = y + offset
if i == 4:
x1 = x + offset
y1 = y - 2.0 * offset
if i == 5:
x1 = x - 3.0 * offset
y1 = y - 2.0 * offset
if data < 5:
string = str(round(data, 1))
else:
string = str(round(data))
self.axes[subplot_index].text(x1, y1, string, color=text_color)
return self.axes[subplot_index]