Matplotlib

pytao.plotting.mpl

Classes

pytao.plotting.mpl.MatplotlibGraphManager

MatplotlibGraphManager(tao)

Bases: GraphManager

Matplotlib backend graph manager.

Source code in pytao/plotting/plot.py

def __init__(self, tao: Tao) -> None:
    self.tao = tao
    self.regions = {}
    self._to_place = {}

Functions

pytao.plotting.mpl.MatplotlibGraphManager.plot

plot(template, *, region_name=None, include_layout=True, tight_layout=True, width=None, height=None, layout_height=None, figsize=None, share_x=True, xlim=None, ylim=None, save=None, settings=None, curves=None, axes=None)

Plot a graph with Matplotlib.

Parameters:

Name	Type	Description	Default
template	str	Graph template name.	required
region_name	str	Graph region name.	None
include_layout	bool	Include a layout plot at the bottom, if not already placed and if appropriate (i.e., another plot uses longitudinal coordinates on the x-axis).	True
tight_layout	bool	Apply a tight layout with matplotlib.	True
figsize	(float, float)	Figure size. Alternative to specifying width and height separately. This takes precedence over width and height. Defaults to Matplotlib's `rcParams["figure.figsize"]``.	None
width	float	Width of the whole plot.	None
height	float	Height of the whole plot.	None
layout_height	float	Normalized height of the layout plot - assuming regular plots are of height 1. Default is 0.5 which is configurable with set_defaults.	None
share_x	bool	Share x-axes for all plots.	True
xlim	(float, float)	X axis limits.	None
ylim	(float, float)	Y axis limits.	None
save	Path or str	Save the plot to the given filename.	None
curves	Dict[int, TaoCurveSettings]	Dictionary of curve index to curve settings. These settings will be applied to the placed graph prior to plotting.	None
settings	TaoGraphSettings	Graph customization settings.	None

Returns:

Type	Description
list of graphs	List of plotted graphs.
Figure	To gain access to the resulting plot objects, use the backend's plot method directly.
List[Axes]

Source code in pytao/plotting/mpl.py

def plot(
    self,
    template: str,
    *,
    region_name: Optional[str] = None,
    include_layout: bool = True,
    tight_layout: bool = True,
    width: Optional[float] = None,
    height: Optional[float] = None,
    layout_height: Optional[float] = None,
    figsize: Optional[Tuple[float, float]] = None,
    share_x: bool = True,
    xlim: Optional[Limit] = None,
    ylim: Optional[Limit] = None,
    save: Union[bool, str, pathlib.Path, None] = None,
    settings: Optional[TaoGraphSettings] = None,
    curves: Optional[Dict[int, TaoCurveSettings]] = None,
    axes: Optional[List[matplotlib.axes.Axes]] = None,
):
    """
    Plot a graph with Matplotlib.

    Parameters
    ----------
    template : str
        Graph template name.
    region_name : str, optional
        Graph region name.
    include_layout : bool, optional
        Include a layout plot at the bottom, if not already placed and if
        appropriate (i.e., another plot uses longitudinal coordinates on
        the x-axis).
    tight_layout : bool, default=True
        Apply a tight layout with matplotlib.
    figsize : (float, float), optional
        Figure size. Alternative to specifying `width` and `height`
        separately.  This takes precedence over `width` and `height`.
        Defaults to Matplotlib's `rcParams["figure.figsize"]``.
    width : float, optional
        Width of the whole plot.
    height : float, optional
        Height of the whole plot.
    layout_height : float, optional
        Normalized height of the layout plot - assuming regular plots are
        of height 1.  Default is 0.5 which is configurable with `set_defaults`.
    share_x : bool, default=True
        Share x-axes for all plots.
    xlim : (float, float), optional
        X axis limits.
    ylim : (float, float), optional
        Y axis limits.
    save : pathlib.Path or str, optional
        Save the plot to the given filename.
    curves : Dict[int, TaoCurveSettings], optional
        Dictionary of curve index to curve settings. These settings will be
        applied to the placed graph prior to plotting.
    settings : TaoGraphSettings, optional
        Graph customization settings.

    Returns
    -------
    list of graphs
        List of plotted graphs.
    matplotlib.Figure
        To gain access to the resulting plot objects, use the backend's
        `plot` method directly.
    List[matplotlib.axes.Axes]
    """
    graphs = self.prepare_graphs_by_name(
        template_name=template,
        region_name=region_name,
        curves=curves,
        settings=settings,
        xlim=xlim,
        ylim=ylim,
    )
    if not graphs:
        raise UnsupportedGraphError(f"No supported plots from this template: {template}")

    figsize = get_figsize(figsize, width, height)

    if (
        include_layout
        and not any(isinstance(graph, LatticeLayoutGraph) for graph in graphs)
        and any(graph.is_s_plot for graph in graphs)
    ):
        layout_graph = self.lattice_layout_graph
        graphs.append(layout_graph)
    else:
        include_layout = False

    if axes is not None:
        if len(axes) != len(graphs):
            raise ValueError(
                f"Not enough axes provided. Expected {len(graphs)}, got {len(axes)}"
            )
        fig = axes[0].figure
    else:
        if include_layout:
            layout_height = layout_height or _Defaults.layout_height
            fig, gs = plt.subplots(
                nrows=len(graphs),
                ncols=1,
                sharex=share_x,
                height_ratios=[1] * (len(graphs) - 1) + [layout_height],
                figsize=figsize,
                squeeze=False,
            )
        else:
            fig, gs = plt.subplots(
                nrows=len(graphs),
                ncols=1,
                sharex=share_x,
                figsize=figsize,
                squeeze=False,
            )
        axes = list(gs[:, 0])
        assert axes is not None

    if include_layout:
        layout_graph = self.lattice_layout_graph

    for ax, graph in zip(axes, graphs):
        try:
            plot(graph, ax)
        except UnsupportedGraphError:
            continue

        if isinstance(graph, LatticeLayoutGraph) and len(graphs) > 1:
            # Do not set ylimits if the user specifically requested a layout graph
            yl = None
        else:
            yl = ylim

        setup_matplotlib_ticks(graph, ax, user_xlim=xlim, user_ylim=yl)

    if fig is not None:
        if tight_layout:
            fig.tight_layout()

        if save:
            title = graphs[0].title or f"plot-{time.time()}"
            if save is True:
                save = f"{title}.png"
            logger.info(f"Saving plot to {save!r}")
            fig.savefig(save)

    return graphs, fig, axes

pytao.plotting.mpl.MatplotlibGraphManager.plot_field

plot_field(ele_id, *, colormap=None, radius=0.015, num_points=100, figsize=None, width=4, height=4, x_scale=1000.0, ax=None, save=None)

Plot field information for a given element.

Parameters:

Name	Type	Description	Default
ele_id	str	Element ID.	required
colormap	str	Colormap for the plot. Matplotlib defaults to "PRGn_r", and bokeh defaults to "".	None
radius	float	Radius.	0.015
num_points	int	Number of data points.	100
ax	Axes	The axes to place the plot in.	None
save	Path or str	Save the plot to the given filename.	None

Source code in pytao/plotting/mpl.py

def plot_field(
    self,
    ele_id: str,
    *,
    colormap: Optional[str] = None,
    radius: float = 0.015,
    num_points: int = 100,
    figsize: Optional[Tuple[float, float]] = None,
    width: int = 4,
    height: int = 4,
    x_scale: float = 1e3,
    ax: Optional[matplotlib.axes.Axes] = None,
    save: Union[bool, str, pathlib.Path, None] = None,
):
    """
    Plot field information for a given element.

    Parameters
    ----------
    ele_id : str
        Element ID.
    colormap : str, optional
        Colormap for the plot.
        Matplotlib defaults to "PRGn_r", and bokeh defaults to "".
    radius : float, default=0.015
        Radius.
    num_points : int, default=100
        Number of data points.
    ax : matplotlib.axes.Axes, optional
        The axes to place the plot in.
    save : pathlib.Path or str, optional
        Save the plot to the given filename.
    """
    user_specified_axis = ax is not None

    figsize = get_figsize(figsize, width, height)

    if ax is None:
        _, ax = plt.subplots(figsize=figsize)
    assert ax is not None

    colormap = colormap or _Defaults.colormap

    field = ElementField.from_tao(self.tao, ele_id, num_points=num_points, radius=radius)
    mesh = ax.pcolormesh(
        np.asarray(field.s),
        np.asarray(field.x) * x_scale,
        np.asarray(field.by),
        # vmin=min_field,
        # vmax=max_field,
        cmap=colormap,
    )
    fig = ax.figure
    if fig is not None:
        if not user_specified_axis:
            fig.colorbar(mesh)

        if save:
            if save is True:
                save = f"{ele_id}_field.png"
            if not pathlib.Path(save).suffix:
                save = f"{save}.png"
            logger.info(f"Saving plot to {save!r}")
            fig.savefig(save)

    return field, fig, ax

pytao.plotting.mpl.MatplotlibGraphManager.plot_grid

plot_grid(templates, grid, *, include_layout=False, figsize=None, tight_layout=True, share_x='col', layout_height=None, width=None, height=None, xlim=None, ylim=None, curves=None, settings=None, save=None, axes=None)

Plot graphs on a grid with Matplotlib.

Parameters:

Name	Type	Description	Default
templates	list of str	Graph template names.	required
grid	(nrows, ncols)	Grid the provided graphs into this many rows and columns.	required
include_layout	bool	Include a layout plot at the bottom of each column.	False
tight_layout	bool	Apply a tight layout with matplotlib.	True
figsize	(float, float)	Figure size. Alternative to specifying width and height separately. This takes precedence over width and height. Defaults to Matplotlib's `rcParams["figure.figsize"]``.	None
width	float	Width of the whole plot.	None
height	float	Height of the whole plot.	None
layout_height	int	Normalized height of the layout plot - assuming regular plots are of height 1. Default is 0.5 which is configurable with set_defaults.	None
share_x	(bool, 'row', 'col', 'all')	Share all x-axes (True or "all"), share x-axes in rows ("row") or in columns ("col").	"col"
xlim	list of (float, float)	X axis limits for each graph.	None
ylim	list of (float, float)	Y axis limits for each graph.	None
curves	list of Dict[int, TaoCurveSettings]	One dictionary per graph, with each dictionary mapping the curve index to curve settings. These settings will be applied to the placed graphs prior to plotting.	None
settings	list of TaoGraphSettings	Graph customization settings, per graph.	None
save	Path or str	Save the plot to the given filename.	None

Returns:

Type	Description
list of graphs	List of plotted graphs.
Figure	To gain access to the resulting plot objects, use the backend's plot method directly.
List[List[Axes]]	Gridded axes, accessible with grid[row][col].

Source code in pytao/plotting/mpl.py

def plot_grid(
    self,
    templates: List[str],
    grid: Tuple[int, int],
    *,
    include_layout: bool = False,
    figsize: Optional[Tuple[float, float]] = None,
    tight_layout: bool = True,
    share_x: Union[bool, Literal["row", "col", "all"]] = "col",
    layout_height: Optional[float] = None,
    width: Optional[float] = None,
    height: Optional[float] = None,
    xlim: Union[OptionalLimit, Sequence[OptionalLimit]] = None,
    ylim: Union[OptionalLimit, Sequence[OptionalLimit]] = None,
    curves: Optional[List[Dict[int, TaoCurveSettings]]] = None,
    settings: Optional[List[TaoGraphSettings]] = None,
    save: Union[bool, str, pathlib.Path, None] = None,
    axes: Optional[List[List[matplotlib.axes.Axes]]] = None,
):
    """
    Plot graphs on a grid with Matplotlib.

    Parameters
    ----------
    templates : list of str
        Graph template names.
    grid : (nrows, ncols), optional
        Grid the provided graphs into this many rows and columns.
    include_layout : bool, default=False
        Include a layout plot at the bottom of each column.
    tight_layout : bool, default=True
        Apply a tight layout with matplotlib.
    figsize : (float, float), optional
        Figure size. Alternative to specifying `width` and `height`
        separately.  This takes precedence over `width` and `height`.
        Defaults to Matplotlib's `rcParams["figure.figsize"]``.
    width : float, optional
        Width of the whole plot.
    height : float, optional
        Height of the whole plot.
    layout_height : int, optional
        Normalized height of the layout plot - assuming regular plots are
        of height 1.  Default is 0.5 which is configurable with `set_defaults`.
    share_x : bool, "row", "col", "all", default="col"
        Share all x-axes (`True` or "all"), share x-axes in rows ("row") or
        in columns ("col").
    xlim : list of (float, float), optional
        X axis limits for each graph.
    ylim : list of (float, float), optional
        Y axis limits for each graph.
    curves : list of Dict[int, TaoCurveSettings], optional
        One dictionary per graph, with each dictionary mapping the curve
        index to curve settings. These settings will be applied to the
        placed graphs prior to plotting.
    settings : list of TaoGraphSettings, optional
        Graph customization settings, per graph.
    save : pathlib.Path or str, optional
        Save the plot to the given filename.

    Returns
    -------
    list of graphs
        List of plotted graphs.
    matplotlib.Figure
        To gain access to the resulting plot objects, use the backend's
        `plot` method directly.
    List[List[matplotlib.axes.Axes]]
        Gridded axes, accessible with `grid[row][col]`.
    """

    graphs = self.prepare_grid_by_names(
        template_names=templates,
        curves=curves,
        settings=settings,
        xlim=xlim,
        ylim=ylim,
    )
    nrows, ncols = grid
    height_ratios = None

    figsize = get_figsize(figsize, width, height)

    if include_layout:
        layout_height = layout_height or _Defaults.layout_height
        empty_graph_count = nrows * ncols - len(templates)
        if empty_graph_count < ncols:
            # Add a row for the layout
            nrows += 1
        height_ratios = [1] * (nrows - 1) + [layout_height]

    if axes is not None:
        tight_layout = False
        fig = None
    else:
        fig, gs = plt.subplots(
            nrows=nrows,
            ncols=ncols,
            sharex=share_x,
            figsize=figsize,
            squeeze=False,
            height_ratios=height_ratios,
        )
        axes = [list(gs[row, :]) for row in range(nrows)]
        for row in axes:
            for ax in row:
                ax.set_axis_off()

    all_xlim = fix_grid_limits(xlim, num_graphs=len(graphs))
    all_ylim = fix_grid_limits(ylim, num_graphs=len(graphs))

    rows_cols = [(row, col) for row in range(nrows) for col in range(ncols)]

    for graph, xl, yl, (row, col) in zip(graphs, all_xlim, all_ylim, rows_cols):
        ax = axes[row][col]
        try:
            plot(graph, ax)
        except UnsupportedGraphError:
            continue

        ax.set_axis_on()
        setup_matplotlib_ticks(graph, ax, user_xlim=xl, user_ylim=yl)

    if include_layout:
        layout_graph = self.lattice_layout_graph
        for col in range(ncols):
            ax = axes[-1][col]
            plot(layout_graph, ax)
            ax.set_axis_on()

            xl = None
            if share_x in {"all", "col", True} and nrows > 1:
                try:
                    xl = axes[0][col].get_xlim()
                except IndexError:
                    pass

            setup_matplotlib_ticks(layout_graph, ax, user_xlim=xl, user_ylim=None)

    if tight_layout and fig is not None:
        fig.tight_layout()

    if save and fig is not None:
        title = graphs[0].title or f"plot-{time.time()}"
        if save is True:
            save = f"{title}.png"
        logger.info(f"Saving plot to {save!r}")
        fig.savefig(save)

    return graphs, fig, axes

Functions

pytao.plotting.mpl.set_defaults

set_defaults(layout_height=None, colormap=None, line_width_scale=None, floor_line_width_scale=None, figsize=None, width=None, height=None, dpi=None)

Set default values for Matplotlib plot settings.

Parameters:

Name	Type	Description	Default
layout_height	float	Height of the layout. Default is 0.5.	None
colormap	str	Colormap to use for plotting. Default is "PRGn_r".	None
line_width_scale	float	Scale factor for line widths, excluding floor plan lines. Default is 0.5.	None
floor_line_width_scale	float	Scale factor for floor plan line widths. Default is 0.5.	None
figsize	tuple of float	Size of the figure (width, height). Default is as-configured in matplotlib rcParams.	None
width	float	Width of the figure in inches. Default is as-configured in matplotlib rcParams.	None
height	float	Height of the figure in inches. Default is as-configured in matplotlib rcParams.	None
dpi	int	Dots per inch for the figure. Default is as-configured in matplotlib rcParams.	None

Source code in pytao/plotting/mpl.py

def set_defaults(
    layout_height: Optional[float] = None,
    colormap: Optional[str] = None,
    line_width_scale: Optional[float] = None,
    floor_line_width_scale: Optional[float] = None,
    figsize: Optional[Tuple[float, float]] = None,
    width: Optional[float] = None,
    height: Optional[float] = None,
    dpi: Optional[int] = None,
):
    """
    Set default values for Matplotlib plot settings.

    Parameters
    ----------
    layout_height : float, optional
        Height of the layout. Default is 0.5.
    colormap : str, optional
        Colormap to use for plotting. Default is "PRGn_r".
    line_width_scale : float, optional
        Scale factor for line widths, excluding floor plan lines. Default is 0.5.
    floor_line_width_scale : float, optional
        Scale factor for floor plan line widths. Default is 0.5.
    figsize : tuple of float, optional
        Size of the figure (width, height). Default is as-configured in matplotlib rcParams.
    width : float, optional
        Width of the figure in inches. Default is as-configured in matplotlib rcParams.
    height : float, optional
        Height of the figure in inches. Default is as-configured in matplotlib rcParams.
    dpi : int, optional
        Dots per inch for the figure. Default is as-configured in matplotlib rcParams.
    """

    if layout_height is not None:
        _Defaults.layout_height = layout_height
    if colormap is not None:
        _Defaults.colormap = colormap
    if line_width_scale is not None:
        _Defaults.line_width_scale = line_width_scale
    if floor_line_width_scale is not None:
        _Defaults.floor_line_width_scale = floor_line_width_scale
    if figsize is not None:
        matplotlib.rcParams["figure.figsize"] = figsize
    if width and height:
        matplotlib.rcParams["figure.figsize"] = (width, height)
    if dpi is not None:
        matplotlib.rcParams["figure.dpi"] = dpi

    info = {key: value for key, value in vars(_Defaults).items() if not key.startswith("_")}
    info["figsize"] = matplotlib.rcParams["figure.figsize"]
    info["dpi"] = matplotlib.rcParams["figure.dpi"]
    return info

pytao.plotting.mpl.setup_matplotlib_axis

setup_matplotlib_axis(graph, ax)

Configure limits, title, and basic info for the given axes.

Source code in pytao/plotting/mpl.py

def setup_matplotlib_axis(graph: AnyGraph, ax: matplotlib.axes.Axes):
    """Configure limits, title, and basic info for the given axes."""
    if not graph.show_axes:
        ax.set_axis_off()

    ax.set_title(pgplot.mpl_string(graph.title))
    ax.set_xlabel(pgplot.mpl_string(graph.xlabel))
    ax.set_ylabel(pgplot.mpl_string(graph.ylabel))
    ax.set_axisbelow(True)

    if graph.draw_grid:
        ax.grid(graph.draw_grid, which="major", axis="both")