def plot(self, x: Optional[Iterable], kind: str = 'line', color: str = 'C0', mean: bool = False, median: bool = False, mode: bool = False, std: bool = False, ax: Optional[Axes] = None, **kwargs) -> Axes: """ Plot the function. :param x: Range of values of x to plot p(x) over. :param kind: Kind of plot e.g. 'bar', 'line'. :param color: Optional color for the series. :param mean: Whether to show marker and label for the mean. :param median: Whether to show marker and label for the median. :param mode: Whether to show marker and label for the mode. :param std: Whether to show marker and label for the standard deviation. :param ax: Optional matplotlib axes to plot on. :param kwargs: Additional arguments for the matplotlib plot function. """ if x is None: if (hasattr(self._parent, 'lower_bound') and hasattr(self._parent, 'upper_bound')): x = linspace(self._parent.lower_bound, self._parent.upper_bound, 1001) else: raise ValueError('Must pass x if distribution has no bounds.') data: Series = self.at(x) axf = AxesFormatter(axes=ax) ax = axf.axes if self._method_name in ('pdf', 'cdf', 'logpdf'): if 'label' not in kwargs.keys(): kwargs['label'] = self._parent.label data.plot(kind=kind, color=color, ax=axf.axes, **kwargs) else: raise ValueError('plot not implemented for {}'.format(self._name)) # stats y_min = axf.get_y_min() y_max = axf.get_y_max() x_mean = self._distribution.mean() if mean: axf.add_v_lines(x=x_mean, y_min=y_min, y_max=y_max, line_styles='--', colors=color) axf.add_text(x=x_mean, y=self._distribution.pdf(x_mean), text=f'mean={x_mean: 0.3f}', color=color, h_align='center', v_align='bottom') if median: x_median = self._distribution.median() axf.add_v_lines(x=x_median, y_min=y_min, y_max=y_max, line_styles='-.', colors=color) axf.add_text(x=x_median, y=self._distribution.pdf(x_median), text=f'median={x_median: 0.3f}', color=color, h_align='center', v_align='bottom') if mode: x_mode = self._parent.mode() axf.add_v_lines(x=x_mode, y_min=y_min, y_max=y_max, line_styles='-.', colors=color) axf.add_text(x=x_mode, y=self._distribution.pdf(x_mode), text=f'mode={x_mode: 0.3f}', color=color, h_align='center', v_align='bottom') if std: x_std = self._distribution.std() axf.add_v_lines(x=[x_mean - x_std, x_mean + x_std], y_min=y_min, y_max=y_max, line_styles=':', colors=color) axf.add_text(x=x_mean - x_std / 2, y=self._distribution.pdf(x_mean - x_std / 2), text=f'std={x_std: 0.3f}', color=color, h_align='center', v_align='bottom') ax.set_xlabel(self._parent.x_label) if self._parent.y_label: ax.set_ylabel(self._parent.y_label) else: if self._method_name == 'pdf': ax.set_ylabel('P(X = x)') elif self._method_name == 'cdf': ax.set_ylabel('P(X ≤ x)') elif self._method_name == 'logpdf': ax.set_ylabel('log P(X = x)') else: ax.set_ylabel(self._name) return ax
def plot(self, k: Optional[Iterable[int]], color: str = 'C0', kind: str = 'bar', mean: bool = False, median: bool = False, std: bool = False, ax: Optional[Axes] = None, **kwargs) -> Axes: """ Plot the function. :param k: Range of values of k to plot p(k) over. :param color: Optional color for the series. :param kind: Kind of plot e.g. 'bar', 'line'. :param mean: Whether to show marker and label for the mean. :param median: Whether to show marker and label for the median. :param std: Whether to show marker and label for the standard deviation. :param ax: Optional matplotlib axes to plot on. :param kwargs: Additional arguments for the matplotlib plot function. """ if k is None: if (hasattr(self._parent, 'lower_bound') and hasattr(self._parent, 'upper_bound')): k = range(self._parent.lower_bound, self._parent.upper_bound + 1) else: raise ValueError('Must pass k if distribution has no bounds.') data: Series = self.at(k) axf = AxesFormatter(axes=ax) ax = axf.axes # special kwargs vlines = None if 'vlines' in kwargs.keys(): vlines = kwargs.pop('vlines') if 'label' not in kwargs.keys(): kwargs['label'] = self._parent.label if self._method_name == 'pmf': data.plot(kind=kind, color=color, ax=axf.axes, **kwargs) elif self._method_name == 'cdf': data.plot(kind='line', color=color, drawstyle='steps-post', ax=axf.axes, **kwargs) else: raise ValueError('plot not implemented for {}'.format(self._name)) if vlines: axf.axes.vlines(x=k, ymin=0, ymax=data.values, color=color) y_min = axf.get_y_min() y_max = axf.get_y_max() x_mean = self._distribution.mean() if mean: axf.add_v_lines(x=x_mean, y_min=y_min, y_max=y_max, line_styles='--', colors=color) axf.add_text(x=x_mean, y=self._distribution.pmf(x_mean), text=f'mean={x_mean: 0.3f}', color=color, h_align='center', v_align='bottom') if median: x_median = self._distribution.median() axf.add_v_lines(x=x_median, y_min=y_min, y_max=y_max, line_styles='-.', colors=color) axf.add_text(x=x_median, y=self._distribution.pmf(x_median), text=f'median={x_median: 0.3f}', color=color, h_align='center', v_align='bottom') if std: x_std = self._distribution.std() axf.add_v_lines(x=[x_mean - x_std, x_mean + x_std], y_min=y_min, y_max=y_max, line_styles=':', colors=color) axf.add_text(x=x_mean - x_std / 2, y=self._distribution.pmf(x_mean - x_std / 2), text=f'std={x_std: 0.3f}', color=color, h_align='center', v_align='bottom') ax.set_xlabel(self._parent.x_label) if self._parent.y_label: ax.set_ylabel(self._parent.y_label) else: if self._method_name == 'pmf': ax.set_ylabel('P(K = k)') elif self._method_name == 'cdf': ax.set_ylabel('P(K ≤ k)') else: ax.set_ylabel(self._name) return ax