def plot_state_components(self, burn: int = None, time: np.array = None, same_scale: bool = True, ylim: tuple = None, fig=None, components=None, **kwargs): """ Plot the contribution of each state model. Args: burn: TODO: Finish this. Use fig.add_gridspec instead of plt.subplots, so that each figure can subgrid if needed. See https://matplotlib.org/3.1.1/tutorials/intermediate/gridspec.html """ if fig is None: fig = plt.figure(constrained_layout=True) if components is None: state_models = self._state_models else: state_models = self._state_models[components] if burn is None: burn = self.suggest_burn() elif burn < 0: burn = 0 if time is None: time = self.original_series.index if same_scale is True and ylim is None: ylim = _find_ylim(state_models, burn) nr, nc = R.plot_grid_shape(len(state_models)) outer_grid = fig.add_gridspec(nr, nc) plot_index = 0 for i in range(nr): for j in range(nc): if plot_index < len(state_models): state_model = state_models[plot_index] state_model.plot_state_contribution( fig=fig, gridspec=outer_grid[i, j], time=time, burn=burn, ylim=ylim, **kwargs) plot_index += 1
def plot_coefficients(self, subset=None, fig=None, ylim=None, same_scale: bool = True, burn: int = None, **kwargs): """ Plot some or all of the model coefficients on the same figure. Args: subset: Identifies the subset of coefficients to be plotted. 'None' means plot all coefficients. Otherwise 'subset' must be an object that can be used to index a numpy array, such as a vector of integers or booleans. fig: The plt.Figure object on which to draw the plots. If None a new Figure will be created and drawn upon function exit. ylim: A pair of numbers giving the lower and upper limits of the Y axes. If specified then all plots will be drawn with the same values. same_scale: This is ignored if ylim is specified. Otherwise, True indicates that the Y axes should all be drawn on the same scale. False means each coefficient is plotted on its own scale. kw_args: Additional keyword arguments are ignored. burn: The number of MCMC iteration to discard as burn-in. Returns: The Figure object on which the plots are drawn. """ if subset is None: coef = self._beta_draws else: coef = self._beta_draws[:, subset, :] if burn is None: burn = self.suggest_burn() if burn > 0: coef = coef[burn:, :, :] call_show = fig is None if fig is None: fig = plt.figure() nr, nc = R.plot_grid_shape(coef.shape[1]) sharey = same_scale or (ylim is not None) ax = fig.subplots(nr, nc, sharex=True, sharey=sharey) index = 0 if same_scale and (ylim is None): ylim = R.data_range(self._beta_draws) for col in range(nc): for row in range(nr): if index < coef.shape[1]: self.plot_single_coefficient(coef[:, index, :], ax=ax[row][col], ylim=ylim) index += 1 if call_show: fig.set_tight_layout(True) fig.show() return fig