Plotting

Plotting module for the grav_sim package.

plot_2d_trajectory(sol_state, colors=None, labels=None, legend=False, xlabel='$x$', ylabel='$y$', title=None, marker='o', markersize=6, save_fig=False, save_fig_path=None)

Plot 2D trajectory of objects

Parameters:

Name	Type	Description	Default
sol_state	ndarray	Solution state of the system	required
colors	Optional[list[str]]	Colors of the trajectories, by default None	None
labels	Optional[list[str]]	Labels of the objects, used for legend, by default None	None
legend	bool	Flag to check whether to show legend, by default False	False
xlabel	str	Label of x-axis, by default "\(x\) (AU)"	'$x$'
ylabel	str	Label of y-axis, by default "\(y\) (AU)"	'$y$'
marker	str	Marker for the last position, by default "o"	'o'
markersize	int	Marker size for the last position, by default 6	6
save_fig	bool	Flag to check whether to save the figure, by default False	False
save_fig_path	Optional[str]	Path to save the figure, by default None	None

Source code in grav_sim/plotting.py

def plot_2d_trajectory(
    sol_state: np.ndarray,
    colors: Optional[list[str]] = None,
    labels: Optional[list[str]] = None,
    legend: bool = False,
    xlabel: str = "$x$",
    ylabel: str = "$y$",
    title: Optional[str] = None,
    marker: str = "o",
    markersize: int = 6,
    save_fig: bool = False,
    save_fig_path: Optional[str | Path] = None,
) -> None:
    """Plot 2D trajectory of objects

    Parameters
    ----------
    sol_state : np.ndarray
        Solution state of the system
    colors : Optional[list[str]], optional
        Colors of the trajectories, by default None
    labels : Optional[list[str]], optional
        Labels of the objects, used for legend, by default None
    legend : bool, optional
        Flag to check whether to show legend, by default False
    xlabel : str, optional
        Label of x-axis, by default "$x$ (AU)"
    ylabel : str, optional
        Label of y-axis, by default "$y$ (AU)"
    marker : str, optional
        Marker for the last position, by default "o"
    markersize : int, optional
        Marker size for the last position, by default 6
    save_fig : bool, optional
        Flag to check whether to save the figure, by default False
    save_fig_path : Optional[str], optional
        Path to save the figure, by default None
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, aspect="equal")

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    for i in range(sol_state.shape[1]):
        if colors is not None:
            colors_i = colors[i]
        else:
            colors_i = None

        if labels is not None:
            labels_i = labels[i]
        else:
            labels_i = None

        traj = ax.plot(
            sol_state[:, i, 1],
            sol_state[:, i, 2],
            color=colors_i,
        )
        # Plot the last position with marker
        ax.plot(
            sol_state[-1, i, 1],
            sol_state[-1, i, 2],
            color=traj[0].get_color(),
            label=labels_i,
            marker=marker,
            markersize=markersize,
        )

    if title is not None:
        ax.set_title(title)

    if legend:
        fig.legend(loc="center right", borderaxespad=0.2)
        fig.tight_layout()

    if save_fig:
        plt.savefig(save_fig_path, dpi=300)
    else:
        plt.show()
    plt.close("all")

plot_3d_trajectory(sol_state, colors=None, labels=None, legend=False, xlabel='$x$', ylabel='$y$', zlabel='$z$', title=None, marker='o', markersize=6, save_fig=False, save_fig_path=None)

Plot 3D trajectory of objects

Parameters:

Name	Type	Description	Default
sol_state	ndarray	Solution state of the system	required
colors	Optional[list[str]]	Colors of the trajectories, by default None	None
labels	Optional[list[str]]	Labels of the objects, used for legend, by default None	None
legend	bool	Flag to check whether to show legend, by default False	False
xlabel	str	Label of x-axis, by default "\(x\) (AU)"	'$x$'
ylabel	str	Label of y-axis, by default "\(y\) (AU)"	'$y$'
zlabel	str	Label of z-axis, by default "\(z\) (AU)"	'$z$'
marker	str	Marker for the last position, by default "o"	'o'
markersize	int	Marker size for the last position, by default 6	6
save_fig	bool	Flag to check whether to save the figure, by default False	False
save_fig_path	Optional[str]	Path to save the figure, by default None	None

Source code in grav_sim/plotting.py

def plot_3d_trajectory(
    sol_state: np.ndarray,
    colors: Optional[list[str]] = None,
    labels: Optional[list[str]] = None,
    legend: bool = False,
    xlabel: str = "$x$",
    ylabel: str = "$y$",
    zlabel: str = "$z$",
    title: Optional[str] = None,
    marker: str = "o",
    markersize: int = 6,
    save_fig: bool = False,
    save_fig_path: Optional[str | Path] = None,
) -> None:
    """Plot 3D trajectory of objects

    Parameters
    ----------
    sol_state : np.ndarray
        Solution state of the system
    colors : Optional[list[str]], optional
        Colors of the trajectories, by default None
    labels : Optional[list[str]], optional
        Labels of the objects, used for legend, by default None
    legend : bool, optional
        Flag to check whether to show legend, by default False
    xlabel : str, optional
        Label of x-axis, by default "$x$ (AU)"
    ylabel : str, optional
        Label of y-axis, by default "$y$ (AU)"
    zlabel : str, optional
        Label of z-axis, by default "$z$ (AU)"
    marker : str, optional
        Marker for the last position, by default "o"
    markersize : int, optional
        Marker size for the last position, by default 6
    save_fig : bool, optional
        Flag to check whether to save the figure, by default False
    save_fig_path : Optional[str], optional
        Path to save the figure, by default None
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    ax.set_box_aspect([1.0, 1.0, 1.0])  # type: ignore
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_zlabel(zlabel)  # type: ignore

    for i in range(sol_state.shape[1] // 6):
        if colors is not None:
            colors_i = colors[i]
        else:
            colors_i = None

        if labels is not None:
            labels_i = labels[i]
        else:
            labels_i = None

        traj = ax.plot(
            sol_state[:, i * 3],
            sol_state[:, i * 3 + 1],
            sol_state[:, i * 3 + 2],
            color=colors_i,
        )
        # Plot the last position with marker
        ax.plot(
            sol_state[-1, i * 3],
            sol_state[-1, i * 3 + 1],
            sol_state[-1, i * 3 + 2],
            color=traj[0].get_color(),
            label=labels_i,
            marker=marker,
            markersize=markersize,
        )

    set_3d_axes_equal(ax)

    if title is not None:
        ax.set_title(title)

    if legend:
        ax.legend(loc="center right", bbox_to_anchor=(1.325, 0.5))
        fig.subplots_adjust(
            right=0.7
        )  # Adjust the right boundary of the plot to make room for the legend

    fig.tight_layout()

    if save_fig:
        plt.savefig(save_fig_path, dpi=300)
    else:
        plt.show()
    plt.close("all")

plot_quantity_against_time(quantity, sol_time, is_log_y=False, title=None, xlabel=None, ylabel=None, colors=None, labels=None, legend=False, save_fig=False, save_fig_path=None)

Plot a quantity against time

Parameters:

Name	Type	Description	Default
quantity	ndarray	Quantity to be plotted	required
sol_time	ndarray	Solution time	required
is_log_y	bool	Flag to check whether y-axis is in log scale, by default False	False
title	Optional[str]	Title of the plot, by default None	None
xlabel	Optional[str]	Label of x-axis, by default None	None
ylabel	Optional[str]	Label of y-axis, by default None	None
colors	Optional[list[str]]	Colors of the trajectories, by default None	None
labels	Optional[list[str]]	Labels of the objects, used for legend, by default None	None
legend	bool	Flag to check whether to show legend, by default False	False
save_fig	bool	Flag to check whether to save the figure, by default False	False
save_fig_path	Optional[str]	Path to save the figure, by default None	None

Source code in grav_sim/plotting.py

def plot_quantity_against_time(
    quantity: np.ndarray,
    sol_time: np.ndarray,
    is_log_y: bool = False,
    title: Optional[str] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    colors: Optional[list[str]] = None,
    labels: Optional[list[str]] = None,
    legend: bool = False,
    save_fig: bool = False,
    save_fig_path: Optional[str | Path] = None,
) -> None:
    """Plot a quantity against time

    Parameters
    ----------
    quantity : np.ndarray
        Quantity to be plotted
    sol_time : np.ndarray
        Solution time
    is_log_y : bool, optional
        Flag to check whether y-axis is in log scale, by default False
    title : Optional[str], optional
        Title of the plot, by default None
    xlabel : Optional[str], optional
        Label of x-axis, by default None
    ylabel : Optional[str], optional
        Label of y-axis, by default None
    colors : Optional[list[str]], optional
        Colors of the trajectories, by default None
    labels : Optional[list[str]], optional
        Labels of the objects, used for legend, by default None
    legend : bool, optional
        Flag to check whether to show legend, by default False
    save_fig : bool, optional
        Flag to check whether to save the figure, by default False
    save_fig_path : Optional[str], optional
        Path to save the figure, by default None
    """
    fig = plt.figure()
    ax = fig.add_subplot(111)

    if len(quantity.shape) == 1:
        quantity = quantity.reshape(-1, 1)

    for i in range(quantity.shape[1]):
        if colors is not None:
            colors_i = colors[i]
        else:
            colors_i = None

        if labels is not None:
            labels_i = labels[i]
        else:
            labels_i = None

        ax.plot(sol_time, quantity[:, i], color=colors_i, label=labels_i)

    if is_log_y:
        ax.set_yscale("log")
    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)

    if legend:
        ax.legend(loc="center right", bbox_to_anchor=(1.325, 0.5))
        fig.subplots_adjust(
            right=0.7
        )  # Adjust the right boundary of the plot to make room for the legend

    if save_fig:
        plt.savefig(save_fig_path, dpi=300)
    else:
        plt.show()
    plt.close("all")

set_3d_axes_equal(ax)

Make axes of 3D plot have equal scale

Parameters:

Name	Type	Description	Default
ax	matplotlib axis	The axis to set equal scale	required

Reference

karlo, https://stackoverflow.com/questions/13685386/how-to-set-the-equal-aspect-ratio-for-all-axes-x-y-z

Source code in grav_sim/plotting.py

def set_3d_axes_equal(ax: plt.Axes) -> None:
    """
    Make axes of 3D plot have equal scale

    Parameters
    ----------
    ax : matplotlib axis
        The axis to set equal scale

    Reference
    ---------
    karlo, https://stackoverflow.com/questions/13685386/how-to-set-the-equal-aspect-ratio-for-all-axes-x-y-z
    """

    x_limits = ax.get_xlim3d()  # type: ignore
    y_limits = ax.get_ylim3d()  # type: ignore
    z_limits = ax.get_zlim3d()  # type: ignore

    x_range = abs(x_limits[1] - x_limits[0])
    x_middle = np.mean(x_limits)
    y_range = abs(y_limits[1] - y_limits[0])
    y_middle = np.mean(y_limits)
    z_range = abs(z_limits[1] - z_limits[0])
    z_middle = np.mean(z_limits)

    # The plot bounding box is a sphere in the sense of the infinity
    # norm, hence I call half the max range the plot radius.
    plot_radius = 0.5 * max([x_range, y_range, z_range])

    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])  # type: ignore
    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])  # type: ignore
    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])  # type: ignore

Comments