import matplotlib # noqa: ICN001
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.artist import Artist
from matplotlib.dates import date2num
from matplotlib.lines import Line2D
from numpy import ma
from oceans.ocfis import cart2pol
[docs]
def stick_plot(time, u, v, **kw):
"""Parameters
----------
time: list/arrays of datetime objects
u, v: list/arrays of 2D vector components.
Returns
-------
q: matplotlib's quiver handle for quiverkey.
Examples
--------
>>> from pandas import date_range
>>> time = date_range(start="1990-11-01 00:00", end="1991-2-1 00:00")
>>> u = np.sin(0.1 * time.to_julian_date().values) ** 2 - 0.5
>>> v = np.cos(0.1 * time.to_julian_date().values)
>>> fig, (ax0, ax1, ax2) = plt.subplots(nrows=3, figsize=(10, 6), sharex=True)
>>> q = stick_plot(time, u, v, ax=ax0)
>>> qk = ax0.quiverkey(
... q, 0.2, 0.65, 1, "1 m s$^{-1}$", labelpos="N", coordinates="axes"
... )
>>> l = ax1.plot(time.to_pydatetime(), np.sqrt(u**2 + v**2), label="speed")
>>> l0 = ax2.plot(time.to_pydatetime(), u, label="u")
>>> l1 = ax2.plot(time.to_pydatetime(), v, label="v")
Based on Stephane Raynaud's example from:
https://www.mail-archive.com/matplotlib-users@lists.sourceforge.net/msg18051.html
"""
from pandas import DatetimeIndex
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
width = kw.pop("width", 0.002)
headwidth = kw.pop("headwidth", 0)
headlength = kw.pop("headlength", 0)
headaxislength = kw.pop("headaxislength", 0)
angles = kw.pop("angles", "uv")
ax = kw.pop("ax", None)
msg = (
"Stickplot angles must be `uv` so that if *U*==*V* the angle of the "
"arrow on the plot is 45 degrees CCW from the *x*-axis."
)
if angles != "uv":
raise AssertionError(msg)
if isinstance(time, DatetimeIndex):
time = time.to_pydatetime()
time, u, v = list(map(np.asanyarray, (time, u, v)))
if not ax:
_, ax = plt.subplots()
q = ax.quiver(
date2num(time),
[[0] * len(time)],
u,
v,
angles="uv",
width=width,
headwidth=headwidth,
headlength=headlength,
headaxislength=headaxislength,
**kw,
)
ax.axes.get_yaxis().set_visible(False)
ax.xaxis_date()
return q
[docs]
def landmask(M, color="0.8"):
"""Plot land mask.
Based on trondkristiansen mpl_util.py.
"""
# Make a constant colormap, default = grey
constmap = np.matplotlib.colors.ListedColormap([color])
jmax, imax = M.shape
# X and Y give the grid cell boundaries,
# one more than number of grid cells + 1
# half integers (grid cell centers are integers)
X = -0.5 + np.arange(imax + 1)
Y = -0.5 + np.arange(jmax + 1)
# Draw the mask by pcolor.
M = ma.masked_where(M > 0, M)
plt.pcolor(X, Y, M, shading="flat", cmap=constmap)
[docs]
def level_colormap(levels, cmap=None):
"""Make a colormap based on an increasing sequence of levels.
Based on trondkristiansen.com mpl_util.py.
"""
# Start with an existing colormap.
if not cmap:
cmap = plt.get_cmap()
# Spread the colors maximally.
nlev = len(levels)
S = np.arange(nlev, dtype="float") / (nlev - 1)
A = cmap(S)
# Normalize the levels to interval [0, 1].
levels = np.array(levels, dtype="float")
L = (levels - levels[0]) / (levels[-1] - levels[0])
# Make the color dictionary.
R = [(L[i], A[i, 0], A[i, 0]) for i in range(nlev)]
G = [(L[i], A[i, 1], A[i, 1]) for i in range(nlev)]
B = [(L[i], A[i, 2], A[i, 2]) for i in range(nlev)]
cdict = {"red": tuple(R), "green": tuple(G), "blue": tuple(B)}
return matplotlib.colors.LinearSegmentedColormap(
f"{cmap.name}_levels",
cdict,
256,
)
[docs]
def get_pointsxy(points):
"""Return x, y of the given point object."""
return points.get_xdata(), points.get_ydata()
[docs]
def compass(u, v, **arrowprops):
"""Compass draws a graph that displays the vectors with
components `u` and `v` as arrows from the origin.
Examples
--------
>>> import numpy as np
>>> u = [+0, -0.5, -0.50, +0.90]
>>> v = [+1, +0.5, -0.45, -0.85]
>>> fig, ax = compass(u, v)
"""
# Create plot.
fig, ax = plt.subplots(subplot_kw={"polar": True})
angles, radii = cart2pol(u, v)
# Arrows or sticks?
kw = {"arrowstyle": "->"}
kw.update(arrowprops)
[
ax.annotate("", xy=(angle, radius), xytext=(0, 0), arrowprops=kw)
for angle, radius in zip(angles, radii, strict=True)
]
ax.set_ylim(0, np.max(radii))
return fig, ax
[docs]
def plot_spectrum(data, fs):
"""Plots a Single-Sided Amplitude Spectrum of y(t)."""
n = len(data) # Length of the signal.
k = np.arange(n)
T = n / fs
frq = k / T # Two sides frequency range.
N = list(range(n // 2))
frq = frq[N] # One side frequency range
# FFT computing and normalization.
Y = np.fft.fft(data) / n
Y = Y[N]
# Plotting the spectrum.
plt.semilogx(frq, np.abs(Y), "r")
plt.xlabel("Freq (Hz)")
plt.ylabel("|Y(freq)|")
plt.show()
[docs]
class EditPoints:
"""Edit points on a graph with the mouse. Handles only one set of points.
Key-bindings:
't' toggle on and off. (When on, you can move, delete, or add points.)
'd' delete the point.
'i' insert a point.
Examples
--------
>>> import matplotlib.pyplot as plt
>>> fig, ax = plt.subplots(figsize=(6, 6))
>>> theta = np.arange(0, 2 * np.pi, 0.1)
>>> r = 1.5
>>> xs = r * np.cos(theta)
>>> ys = r * np.sin(theta)
>>> points = ax.plot(xs, ys, "ko")
>>> p = EditPoints(fig, ax, points[0], verbose=True)
>>> _ = ax.set_title("Click and drag a point to move it")
>>> _ = ax.axis([-2, 2, -2, 2])
Based on https://matplotlib.org/examples/event_handling/poly_editor.html
"""
epsilon = 5 # Maximum pixel distance to count as a point hit.
showpoint = True
def __init__(self, fig, ax, points, *, verbose=False):
matplotlib.interactive(b=True)
if points is None:
msg = "First add points to a figure or canvas."
raise RuntimeError(msg)
canvas = fig.canvas
self.ax = ax
self.dragged = None
self.points = points
self.verbose = verbose
x, y = get_pointsxy(points)
self.line = Line2D(
x,
y,
marker="o",
markerfacecolor="r",
linestyle="none",
animated=True,
)
self.ax.add_line(self.line)
self._ind = None # The active point.
canvas.mpl_connect("draw_event", self.draw_callback)
canvas.mpl_connect("button_press_event", self.button_press_callback)
canvas.mpl_connect("key_press_event", self.key_press_callback)
canvas.mpl_connect(
"button_release_event",
self.button_release_callback,
)
canvas.mpl_connect("motion_notify_event", self.motion_notify_callback)
self.canvas = canvas
[docs]
def draw_callback(self, event): # noqa: ARG002
self.background = self.canvas.copy_from_bbox(self.ax.bbox)
self.ax.draw_artist(self.line)
self.canvas.blit(self.ax.bbox)
if self.verbose:
print("\nDrawing...") # noqa: T201
[docs]
def points_changed(self, points):
"""This method is called whenever the points object is called."""
# Only copy the artist props to the line (except visibility).
vis = self.line.get_visible()
Artist.update_from(self.line, points)
# Don't use the points visibility state.
self.line.set_visible(vis)
if self.verbose:
print("\nPoints modified.") # noqa: T201
[docs]
def get_ind_under_point(self, event):
"""Get the index of the point under mouse if within epsilon tolerance."""
# Display coordinates.
arr = self.ax.transData.transform(self.points.get_xydata())
x, y = arr[:, 0], arr[:, 1]
d = np.sqrt((x - event.x) ** 2 + (y - event.y) ** 2)
indseq = np.nonzero(np.equal(d, np.amin(d)))[0]
ind = indseq[0]
if self.verbose:
print(f"d[ind] {d[ind]} epsilon {self.epsilon}") # noqa: T201
if d[ind] >= self.epsilon:
ind = None
if self.verbose:
print(f"\nClicked at ({event.xdata}, {event.ydata})") # noqa: T201
return ind
[docs]
def key_press_callback(self, event): # noqA: C901
"""Whenever a key is pressed."""
if not event.inaxes:
return None
if event.key == "t":
self.showpoint = not self.showpoint
self.line.set_visible(self.showpoint)
if not self.showpoint:
self._ind = None
if self.verbose:
print(f"\nToggle {self.showpoint:d}") # noqa: T201
return get_pointsxy(self.points)
if event.key == "d":
x, y = get_pointsxy(self.points)
ind = self.get_ind_under_point(event)
if ind is not None:
if self.verbose:
print( # noqa: T201
f"\nDeleted ({x[ind]}, {y[ind]}) ind: {ind}",
)
x = np.delete(x, ind)
y = np.delete(y, ind)
self.points.set_xdata(x)
self.points.set_ydata(y)
self.line.set_data(self.points.get_data())
elif event.key == "i":
if self.verbose:
print("Insert point") # noqa: T201
xs = self.points.get_xdata()
ex, ey = event.xdata, event.ydata
for _i in range(len(xs) - 1):
self.points.set_xdata(np.r_[self.points.get_xdata(), ex])
self.points.set_ydata(np.r_[self.points.get_ydata(), ey])
self.line.set_data(self.points.get_data())
if self.verbose:
print(f"\nInserting: ({ex}, {ey})") # noqa: T201
break
self.canvas.draw()
return None
[docs]
def motion_notify_callback(self, event):
"""On mouse movement."""
if not self.showpoint:
return
if not self._ind:
return
if not event.inaxes:
return
if not event.button:
return
x, y = get_pointsxy(self.points)
dx = event.xdata - self.pick_pos[0]
dy = event.ydata - self.pick_pos[1]
x[self._ind] = self.pick_pos[0] + dx
y[self._ind] = self.pick_pos[1] + dy
if self.verbose:
print(f"\nevent.xdata {event.xdata}") # noqa: T201
print(f"\nevent.ydata {event.ydata}") # noqa: T201
self.points.set_xdata(x)
self.points.set_ydata(y)
self.line.set_data(list(zip(self.points.get_data())))
self.canvas.restore_region(self.background)
self.ax.draw_artist(self.line)
self.canvas.blit(self.ax.bbox)
if self.verbose:
print("\nMoving") # noqa: T201
