def plot()

in orbit/diagnostics/backtest.py [0:0]

• 47 lines of code
• 5 McCabe index (conditional complexity)

    def plot(self, fig_width=20, show_index=False, strftime_fmt="%Y-%m-%d"):
        """
        Parameters
        ----------
        fig_width : float
        show_index : bool
        strftime_fmt : str

        Returns
        -------
        matplotlib axes object
        """
        _, ax = plt.subplots(figsize=(fig_width, self.n_splits))
        # visualize the train/test windows for each split
        tr_start = list()
        tr_len = list()
        # technically should be just self.forecast_len
        tt_len = list()
        yticks = list(range(self.n_splits))
        for idx, scheme in self._split_scheme.items():
            # fill in indices with the training/test groups
            tr_start.append(list(scheme[TimeSeriesSplitSchemeKeys.TRAIN_IDX.value])[0])
            tr_len.append(len(list(scheme[TimeSeriesSplitSchemeKeys.TRAIN_IDX.value])))
            tt_len.append(self.forecast_len)

        tr_start = np.array(tr_start)
        tr_len = np.array(tr_len)

        # ax.barh(yticks, tr_start, align='center', height=.5, color='black', alpha=0.5)
        ax.barh(
            yticks,
            tr_len,
            align="center",
            height=0.5,
            left=tr_start,
            color=OrbitPal.BLUE.value,
            label="train",
        )
        ax.barh(
            yticks,
            tt_len,
            align="center",
            height=0.5,
            left=tr_start + tr_len,
            color=OrbitPal.ORANGE.value,
            label="test",
        )

        if not show_index and self.date_col is not None:
            xticks_loc = np.array(ax.get_xticks(), dtype=int)
            new_xticks_loc = np.linspace(
                0, len(self.dt_array) - 1, num=len(xticks_loc)
            ).astype(int)
            dt_xticks = self.dt_array[new_xticks_loc]
            dt_xticks = dt_xticks.strftime(strftime_fmt)
            ax.set_xticks(new_xticks_loc)
            ax.set_xticklabels(dt_xticks)

        # some formatting parameters
        middle = 15
        large = 20

        ax.set_yticks(yticks)
        ax.set_ylabel("Split #", fontsize=large)
        ax.invert_yaxis()
        # ax.grid(which="both", color='grey', alpha=0.5)
        ax.tick_params(axis="x", which="major", labelsize=middle)
        ax.set_title("Train/Test Split Scheme", fontsize=large)
        return ax