diff options
Diffstat (limited to 'bin/utils/bar_chart.py')
-rw-r--r-- | bin/utils/bar_chart.py | 103 |
1 files changed, 103 insertions, 0 deletions
diff --git a/bin/utils/bar_chart.py b/bin/utils/bar_chart.py new file mode 100644 index 0000000..d7b44e2 --- /dev/null +++ b/bin/utils/bar_chart.py @@ -0,0 +1,103 @@ +# Copyright (c) 2017 Egor Tensin <Egor.Tensin@gmail.com> +# This file is part of the "VK scripts" project. +# For details, see https://github.com/egor-tensin/vk-scripts. +# Distributed under the MIT License. + +import matplotlib.pyplot as plt +import numpy as np + +class BarChartBuilder: + _BAR_HEIGHT = 1. + + def __init__(self): + self._fig, self._ax = plt.subplots() + + def set_title(self, title): + self._ax.set_title(title) + + def _get_bar_axis(self): + return self._ax.get_yaxis() + + def _get_value_axis(self): + return self._ax.get_xaxis() + + def set_bar_axis_limits(self, start=None, end=None): + self._ax.set_ylim(bottom=start, top=end) + + def set_value_axis_limits(self, start=None, end=None): + self._ax.set_xlim(left=start, right=end) + + def set_value_grid(self): + self._get_value_axis().grid() + + def get_bar_labels(self): + return self._get_bar_axis().get_ticklabels() + + def get_value_labels(self): + return self._get_value_axis().get_ticklabels() + + def set_value_label_formatter(self, fn): + from matplotlib.ticker import FuncFormatter + self._get_value_axis().set_major_formatter(FuncFormatter(fn)) + + def set_integer_values_only(self): + from matplotlib.ticker import MaxNLocator + self._get_value_axis().set_major_locator(MaxNLocator(integer=True)) + + @staticmethod + def set_property(*args, **kwargs): + plt.setp(*args, **kwargs) + + def _set_size(self, inches, dim=0): + fig_size = self._fig.get_size_inches() + assert len(fig_size) == 2 + fig_size[dim] = inches + self._fig.set_size_inches(fig_size, forward=True) + + def set_width(self, inches): + self._set_size(inches) + + def set_height(self, inches): + self._set_size(inches, dim=1) + + def plot_bars( + self, bar_labels, bar_lengths, + bars_between_ticks=False, + inches_per_bar=1): + + numof_bars = len(bar_labels) + + if not numof_bars: + self.set_height(1) + self._get_bar_axis().set_tick_params(labelleft=False) + return [] + + self.set_height(inches_per_bar * numof_bars) + + bar_offsets = np.arange(numof_bars) * 2 * self._BAR_HEIGHT + self._BAR_HEIGHT + + if bars_between_ticks: + self._get_bar_axis().set_ticks(bar_offsets - self._BAR_HEIGHT) + else: + self._get_bar_axis().set_ticks(bar_offsets) + + bar_axis_min = 0 + bar_axis_max = 2 * self._BAR_HEIGHT * numof_bars + self.set_bar_axis_limits(bar_axis_min, bar_axis_max) + + self._get_bar_axis().set_ticklabels(bar_labels) + + return self._ax.barh( + bar_offsets, bar_lengths, align='center', height=self._BAR_HEIGHT) + + @staticmethod + def show(): + plt.show() + + def save(self, path): + self._fig.savefig(path, bbox_inches='tight') + +if __name__ == '__main__': + builder = BarChartBuilder() + builder.plot_bars(['Name #1', 'Name #2', 'Name #3'], [4, 11, 7]) + builder.show() |