diff --git a/.gitignore b/.gitignore index f384bb0..7a341d7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,4 @@ -*.csv *.csv psPlotKit/examples/figs/Water recovery vs LCOW.svg psPlotKit/examples/figs/Water recovery vs LCOW.csv @@ -9,3 +8,4 @@ psPlotKit/examples/figs/Brine salinity vs LCOW.csv *.egg *.egg-info site/ +*.log diff --git a/src/psPlotKit/data_plotter/ps_break_down_plotter.py b/src/psPlotKit/data_plotter/ps_break_down_plotter.py index 33ae2da..5441bc1 100644 --- a/src/psPlotKit/data_plotter/ps_break_down_plotter.py +++ b/src/psPlotKit/data_plotter/ps_break_down_plotter.py @@ -15,19 +15,22 @@ class BreakDownPlotter: def __init__( self, - PsData, + PsData=None, save_location=None, save_folder=None, save_name=None, show_fig=True, + fig=None, + line_indexes={}, + line_color_list=None, + color_dict=None, ): self.save_location = create_save_location(save_location, save_folder) self.show_fig = show_fig self.select_data_key_list = [] self.PsData = PsData - self.define_plot_styles() - self.line_indexes = {} + self.define_plot_styles(line_color_list=line_color_list) self.line_groups = None self.xunit = None self.yunit = None @@ -36,6 +39,10 @@ def __init__( self.save_name = save_name self.hatch_groups = {} self.area_groups = {} + if fig is not None and not isinstance(fig, FigureGenerator): + raise ValueError("fig must be a FigureGenerator object") + self.fig = fig + self.color_dict = color_dict def _select_data(self, xkeys, ykeys): self.PsData.select_data(xkeys, require_all_in_dir=False) @@ -48,6 +55,7 @@ def define_area_groups(self, groups): self.area_groups = groups def _get_color(self, label): + if isinstance(self.line_colors, dict): return self.line_colors[label] else: @@ -153,7 +161,11 @@ def _get_group_options(self, selected_keys, xdata, ydata): opts = None key = None color = None + if skey in self.color_dict.keys(): + color = self.color_dict[skey] for i, key in enumerate(self.hatch_groups): + # First try to get plotting options based on skey + # provided directly in hatch_groups if key in str(skey): opts = self.hatch_groups[key] if opts.get("label") == None: @@ -169,12 +181,16 @@ def _get_group_options(self, selected_keys, xdata, ydata): if isinstance(akey, dict): akey, item = list(akey.items())[0] if isinstance(item, dict): + # if True, the structure of area_groups is a list of dicts with one key? if "label" in item: _label = item["label"] if "color" in item: color = item["color"] else: _label = item + + if akey in self.color_dict.keys(): + color = self.color_dict[akey] if self.check_key_in_dir(skey, akey): if _label is None: _label = akey @@ -198,12 +214,14 @@ def _get_group_options(self, selected_keys, xdata, ydata): cur_line[key] = val if self.hatch_groups != {}: + # If hatch_groups is not an empty dict for h_key in self.hatch_groups: - # print("ss", h_key, akey, skey) if h_key in str(skey): + # print("ss", h_key, akey, skey) _label = tuple([h_key, akey]) plot_label.replace(h_key, "") + # Then we are looking for a color based on h_key if color is None: color = self._get_color(h_key) cur_line["color"] = color @@ -230,19 +248,24 @@ def plotbreakdown( xdata, ydata, axis_options=None, - generate_figure=True, + generate_plot=False, legend_loc="upper left", legend_cols=2, fig_options={}, + ax_idx=None, + **kwargs ): + if ax_idx is None: + ax_idx = 0 + self.fig_options = fig_options + self.fig_options["ax_idx"] = ax_idx + self._select_data(xdata, ydata) self.selected_data = self.PsData.get_selected_data() - self.selected_data.display() - self.generate_groups_lines = self._get_group_options( + self._get_group_options( self.selected_data.keys(), xdata, ydata ) - self.fig_options = fig_options self.index = 0 if axis_options is None: self.axis_options = {} @@ -251,28 +274,27 @@ def plotbreakdown( if self.axis_options.get("xlabel") == None: self.axis_options["xlabel"] = self._get_axis_label( self.xdata_label, self.xunit - ) # all lines shold share units + ) # all lines should share units if self.axis_options.get("ylabel") == None: self.axis_options["ylabel"] = self._get_axis_label( self.ydata_label, self.yunit - ) # all lines shold share units + ) # all lines should share units - self.plot_imported_data() + self.plot_imported_data(ax_idx=ax_idx) - if ( - generate_figure - ): # TODO: other plotters call this generate_plot, should make this consistent - self.generate_figure(loc=legend_loc, cols=legend_cols) + # TODO: other plotters call this generate_plot, should make this consistent + if generate_plot: + self.generate_plot(ax_idx=ax_idx, loc=legend_loc, cols=legend_cols) - def plot_imported_data(self): - if "fig_object" in self.fig_options: - self.fig = self.fig_options.get("fig_object") - else: - self.fig = FigureGenerator() - self.fig.init_figure(**self.fig_options) + def plot_imported_data(self, ax_idx=0): + # if "fig_object" in self.fig_options: + # self.fig = self.fig_options.get("fig_object") + # else: + # self.fig = FigureGenerator() + # self.fig.init_figure(**self.fig_options) plotted_legend = [] for group, items in self.hatch_groups.items(): - self.fig.plot_area([], [], **items) + self.fig.plot_area([], [], ax_idx=ax_idx, **items) old_data = 0 current_data = None for linelabel in self.plot_order: @@ -287,15 +309,13 @@ def plot_imported_data(self): else: current_data = line["ydata"] + old_data line["ydata"] = current_data - if "ax_idx" in self.fig_options: - line["ax_idx"] = self.fig_options.get("ax_idx") - self.fig.plot_area(**line) + self.fig.plot_area(ax_idx=ax_idx, **line) old_data = line["ydata"] + + self.fig.set_axis(ax_idx=ax_idx, **self.axis_options) - def generate_figure(self, loc="upper left", cols=2): - if "ax_idx" in self.fig_options: - self.axis_options["ax_idx"] = self.fig_options["ax_idx"] - self.fig.set_axis(**self.axis_options) + def generate_plot(self, ax_idx=0, loc="upper left", cols=2): + self.fig.set_axis(ax_idx=ax_idx, **self.axis_options) self.fig.add_legend(loc=loc, ncol=cols) if self.save_name is not None: @@ -306,6 +326,9 @@ def generate_figure(self, loc="upper left", cols=2): self.fig.close() + def set_selected_data(self, data): + self.PsData = data + class breakDownPlotter(BreakDownPlotter): _logger.warning("breakDownPlotter is deprecated, please use BreakDownPlotter")