Ejemplo n.º 1
0
    def setup_plot(cls,
                   width=16,
                   height=4,
                   ncols=1,
                   nrows=1,
                   interactive=None,
                   link_dataframes=None,
                   cursor_delta=None,
                   **kwargs):
        """
        Common helper for setting up a matplotlib plot

        :param width: Width of the plot (inches)
        :type width: int or float

        :param height: Height of each subplot (inches)
        :type height: int or float

        :param ncols: Number of plots on a single row
        :type ncols: int

        :param nrows: Number of plots in a single column
        :type nrows: int

        :param link_dataframes: Link the provided dataframes to the axes using
            :func:`lisa.notebook.axis_link_dataframes`
        :type link_dataframes: list(pandas.DataFrame) or None

        :param cursor_delta: Add two vertical lines set with left and right
            clicks, and show the time delta between them in a widget.
        :type cursor_delta: bool or None

        :param interactive: If ``True``, use the pyplot API of matplotlib,
            which integrates well with notebooks. However, it can lead to
            memory leaks in scripts generating lots of plots, in which case it
            is better to use the non-interactive API. Defaults to ``True`` when
            running under IPython or Jupyter notebook, `False`` otherwise.
        :type interactive: bool

        :Keywords arguments: Extra arguments to pass to
          :obj:`matplotlib.figure.Figure.subplots`

        :returns: tuple(matplotlib.figure.Figure, matplotlib.axes.Axes (or an
          array of, if ``nrows`` > 1))
        """

        figure, axes = make_figure(
            interactive=interactive,
            width=width,
            height=height,
            ncols=ncols,
            nrows=nrows,
            **kwargs,
        )
        if interactive is None:
            interactive = is_running_ipython()

        use_widgets = interactive

        if link_dataframes:
            if not use_widgets:
                cls.get_logger().error(
                    'Dataframes can only be linked to axes in interactive widget plots'
                )
            else:
                for axis in figure.axes:
                    axis_link_dataframes(axis, link_dataframes)

        if cursor_delta or cursor_delta is None and use_widgets:
            if not use_widgets and cursor_delta is not None:
                cls.get_logger().error(
                    'Cursor delta can only be used in interactive widget plots'
                )
            else:
                for axis in figure.axes:
                    axis_cursor_delta(axis)

        for axis in figure.axes:
            axis.relim(visible_only=True)
            axis.autoscale_view(True)

        # Needed for multirow plots to not overlap with each other
        figure.set_tight_layout(dict(h_pad=3.5))
        return figure, axes
Ejemplo n.º 2
0
    def _plot(self,
              df,
              title,
              plot_func,
              facet_rows,
              facet_cols,
              collapse_cols,
              filename=None,
              interactive=None):
        unit_col = self._unit_col

        group_on = list(facet_rows) + list(facet_cols)
        facet_rows_len = len(facet_rows)

        def split_row_col(group):
            row = tuple(group[:facet_rows_len])
            col = tuple(group[facet_rows_len:])
            return (row, col)

        grouped = df.groupby(group_on, observed=True)
        # DataFrame.groupby() return type is "interesting":
        # When grouping on one column only, the group is not a tuple, but the
        # value itself, leading to equally "interesting" bugs.
        fixup_tuple = lambda x: x if isinstance(x, tuple) else (x, )

        unzip = lambda x: zip(*x)
        rows, cols = map(
            sorted,
            map(
                set,
                unzip(
                    map(split_row_col,
                        map(fixup_tuple, map(itemgetter(0), grouped))))))
        nrows, ncols = len(rows), len(cols)

        # Collapse together all the tag columns that are not already in use
        not_collapse = set(group_on) | {unit_col}
        collapse_cols = [
            col for col in self._restrict_cols(collapse_cols, df)
            if col not in not_collapse
        ]
        if len(collapse_cols) > 1:
            collapsed_col = 'group'
            collapse_group = {collapsed_col: collapse_cols}
        elif collapse_cols:
            collapsed_col = collapse_cols[0]
            collapse_group = {}
        else:
            collapsed_col = None
            collapse_group = {}

        figure, axes = make_figure(
            width=16,
            height=16,
            nrows=nrows,
            ncols=ncols,
            interactive=interactive,
        )
        if nrows == 1 and ncols == 1:
            axes = [[axes]]
        elif nrows == 1:
            axes = [axes]
        elif ncols == 1:
            axes = [[ax] for ax in axes]

        figure.set_tight_layout(dict(h_pad=3.5, ))
        figure.suptitle(title, y=1.01, fontsize=30)

        for group, subdf in grouped:
            group = fixup_tuple(group)

            row, col = split_row_col(group)
            ax = axes[rows.index(row)][cols.index(col)]

            if subdf.empty:
                figure.delaxes(ax)
            else:
                subdf = subdf.drop(columns=group_on)
                subdf = self._collapse_cols(subdf, collapse_group)
                group_dict = dict(zip(group_on, group))
                plot_func(subdf, ax, collapsed_col, group_dict)

        if filename:
            # The suptitle is not taken into account by tight layout by default:
            # https://stackoverflow.com/questions/48917631/matplotlib-how-to-return-figure-suptitle
            suptitle = figure._suptitle
            figure.savefig(filename,
                           bbox_extra_artists=[suptitle],
                           bbox_inches='tight')

        return figure