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