def plot_quiver(
trj: TrajaDataFrame,
bins: Optional[Union[int, tuple]] = None,
quiverplot_kws: dict = {},
**kwargs,
) -> Axes:
"""Plot average flow from each grid cell to neighbor.
Args:
bins (int or tuple): Tuple of x,y bin counts; if `bins` is int, bin count of x,
with y inferred from aspect ratio
quiverplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.quiver`
Returns:
ax (:class:`~matplotlib.axes.Axes`): Axes of quiver plot
"""
after_plot_args, _ = _get_after_plot_args(**kwargs)
X, Y, U, V = coords_to_flow(trj, bins)
Z = np.sqrt(U * U + V * V)
fig, ax = plt.subplots()
ax.quiver(X, Y, U, V, units="width", **quiverplot_kws)
ax = _label_axes(trj, ax)
ax.set_aspect("equal")
_process_after_plot_args(**after_plot_args)
return ax
def plot_contour(
trj: TrajaDataFrame,
bins: Optional[Union[int, tuple]] = None,
filled: bool = True,
quiver: bool = True,
contourplot_kws: dict = {},
contourfplot_kws: dict = {},
quiverplot_kws: dict = {},
ax: Axes = None,
**kwargs,
) -> Axes:
"""Plot average flow from each grid cell to neighbor.
Args:
trj: Traja DataFrame
bins (int or tuple): Tuple of x,y bin counts; if `bins` is int, bin count of x,
with y inferred from aspect ratio
filled (bool): Contours filled
quiver (bool): Quiver plot
contourplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.contour`
contourfplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.contourf`
quiverplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.quiver`
ax (optional): Matplotlib Axes
Returns:
ax (:class:`~matplotlib.axes.Axes`): Axes of quiver plot
"""
after_plot_args, _ = _get_after_plot_args(**kwargs)
X, Y, U, V = coords_to_flow(trj, bins)
Z = np.sqrt(U * U + V * V)
if not ax:
_, ax = plt.subplots()
if filled:
cfp = plt.contourf(X, Y, Z, **contourfplot_kws)
plt.colorbar(cfp, ax=ax)
plt.contour(
X, Y, Z, colors="k", linewidths=1, linestyles="solid", **contourplot_kws
)
if quiver:
ax.quiver(X, Y, U, V, units="width", **quiverplot_kws)
ax = _label_axes(trj, ax)
ax.set_aspect("equal")
_process_after_plot_args(**after_plot_args)
return ax
def plot_stream(
trj: TrajaDataFrame,
bins: Optional[Union[int, tuple]] = None,
cmap: str = "jet",
contourfplot_kws: dict = {},
contourplot_kws: dict = {},
streamplot_kws: dict = {},
**kwargs,
):
"""Plot average flow from each grid cell to neighbor.
Args:
bins (int or tuple): Tuple of x,y bin counts; if `bins` is int, bin count of x,
with y inferred from aspect ratio
contourplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.contour`
contourfplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.contourf`
streamplot_kws: Additional keyword arguments for :meth:`~matplotlib.axes.Axes.streamplot`
Returns:
fig (:class:`~matplotlib.figure.Figure`): Figure of stream plot
"""
after_plot_args, _ = _get_after_plot_args(**kwargs)
X, Y, U, V = coords_to_flow(trj, bins)
Z = np.sqrt(U * U + V * V)
fig, ax = plt.subplots()
cfp = plt.contourf(X, Y, Z, **contourfplot_kws)
cp = plt.contour(X,
Y,
Z,
colors="k",
linewidths=1,
linestyles="solid",
**contourplot_kws)
sp = ax.streamplot(X, Y, U, V, color=Z, cmap=cmap, **streamplot_kws)
ax = _label_axes(trj, ax)
ax.set_aspect("equal")
_process_after_plot_args(**after_plot_args)
return fig
def plot_surface(
trj: TrajaDataFrame,
bins: Optional[Union[int, tuple]] = None,
cmap: str = "jet",
**surfaceplot_kws: dict,
) -> Figure:
"""Plot surface of flow from each grid cell to neighbor in 3D.
Args:
bins (int or tuple): Tuple of x,y bin counts; if `bins` is int, bin count of x,
with y inferred from aspect ratio
cmap (str): color map
surfaceplot_kws: Additional keyword arguments for :meth:`~mpl_toolkits.mplot3D.Axes3D.plot_surface`
Returns:
ax (:class:`~matplotlib.axes.Axes`): Axes of quiver plot
"""
from mpl_toolkits.mplot3d import Axes3D
after_plot_args, surfaceplot_kws = _get_after_plot_args(**surfaceplot_kws)
X, Y, U, V = coords_to_flow(trj, bins)
Z = np.sqrt(U * U + V * V)
fig = plt.figure()
ax = fig.gca(projection="3d")
ax.plot_surface(X,
Y,
Z,
cmap=matplotlib.cm.coolwarm,
linewidth=0,
**surfaceplot_kws)
ax = _label_axes(trj, ax)
try:
ax.set_aspect("equal")
except NotImplementedError:
# 3D
pass
_process_after_plot_args(**after_plot_args)
return ax
def bar_plot(trj: TrajaDataFrame,
bins: Union[int, tuple] = None,
**kwargs) -> Axes:
"""Plot trajectory for single animal over period.
Args:
trj (:class:`traja.TrajaDataFrame`): trajectory
bins (int or tuple): number of bins for x and y
**kwargs: additional keyword arguments to :meth:`mpl_toolkits.mplot3d.Axed3D.plot`
Returns:
ax (:class:`~matplotlib.collections.PathCollection`): Axes of plot
"""
# TODO: Add time component
from mpl_toolkits.mplot3d import Axes3D
bins = traja.trajectory._bins_to_tuple(trj, bins)
X, Y, U, V = coords_to_flow(trj, bins)
hist, _ = trip_grid(trj, bins, hist_only=True)
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.set_aspect("equal")
X = X.flatten("F")
Y = Y.flatten("F")
ax.bar3d(
X,
Y,
np.zeros_like(X),
1,
1,
hist.flatten(),
zsort="average",
shade=True,
**kwargs,
)
ax.set(xlabel="x", ylabel="y", zlabel="Frames")
return ax
