Vis: code cleanup (#100)

Author: dweindl

petab/visualize/plotter.py

@@ -45,6 +45,26 @@ class MPLPlotter(Plotter):
     def __init__(self, figure: Figure, data_provider: DataProvider):
         super().__init__(figure, data_provider)
 
+    @staticmethod
+    def _error_column_for_plot_type_data(plot_type_data: str) -> Optional[str]:
+        """Translate PEtab plotTypeData value to column name of internal
+        data representation
+
+        Parameters
+        ----------
+            plot_type_data: PEtab plotTypeData value
+        Returns
+        -------
+            Name of corresponding column
+        """
+        if plot_type_data == MEAN_AND_SD:
+            return 'sd'
+        if plot_type_data == MEAN_AND_SEM:
+            return 'sem'
+        if plot_type_data == PROVIDED:
+            return 'noise_model'
+        return None
+
     def generate_lineplot(self, ax: 'matplotlib.pyplot.Axes',
                           dataplot: DataPlot,
                           plotTypeData: str) -> None:
@@ -67,14 +87,7 @@ def generate_lineplot(self, ax: 'matplotlib.pyplot.Axes',
         measurements_to_plot, simulations_to_plot = \
             self.data_provider.get_data_to_plot(dataplot,
                                                 plotTypeData == PROVIDED)
-        noise_col = None
-        # set type of noise
-        if plotTypeData == MEAN_AND_SD:
-            noise_col = 'sd'
-        elif plotTypeData == MEAN_AND_SEM:
-            noise_col = 'sem'
-        elif plotTypeData == PROVIDED:
-            noise_col = 'noise_model'
+        noise_col = self._error_column_for_plot_type_data(plotTypeData)
 
         label_base = dataplot.legendEntry
 
@@ -155,14 +168,7 @@ def generate_barplot(self, ax: 'matplotlib.pyplot.Axes',
             Specifies how replicates should be handled.
         """
         # TODO: plotTypeData == REPLICATE?
-        # set type of noise
-        noise_col = None
-        if plotTypeData == MEAN_AND_SD:
-            noise_col = 'sd'
-        elif plotTypeData == MEAN_AND_SEM:
-            noise_col = 'sem'
-        elif plotTypeData == PROVIDED:
-            noise_col = 'noise_model'
+        noise_col = self._error_column_for_plot_type_data(plotTypeData)
 
         simu_colors = None
         measurements_to_plot, simulations_to_plot = \
@@ -222,7 +228,7 @@ def generate_scatterplot(self, ax: 'matplotlib.pyplot.Axes',
         ax.scatter(measurements_to_plot.data_to_plot['mean'],
                    simulations_to_plot.data_to_plot['mean'],
                    label=getattr(dataplot, LEGEND_ENTRY))
-        ax = self._square_plot_equal_ranges(ax)
+        self._square_plot_equal_ranges(ax)
 
     def generate_subplot(self,
                          ax,
@@ -278,11 +284,10 @@ def generate_subplot(self,
             elif subplot.xScale == 'order':
                 ax.set_xscale("linear")
                 # check if conditions are monotone decreasing or increasing
-                if np.all(
-                        np.diff(subplot.conditions) < 0):  # monot. decreasing
-                    xlabel = subplot.conditions[::-1]  # reversing
-                    conditions = range(len(subplot.conditions))[
-                                 ::-1]  # reversing
+                if np.all(np.diff(subplot.conditions) < 0):
+                    # monot. decreasing -> reverse
+                    xlabel = subplot.conditions[::-1]
+                    conditions = range(len(subplot.conditions))[::-1]
                     ax.set_xticks(range(len(conditions)), xlabel)
                 elif np.all(np.diff(subplot.conditions) > 0):
                     xlabel = subplot.conditions
@@ -305,7 +310,7 @@ def ticks(y, _):
         if subplot.yScale == LOG:
             ax.yaxis.set_major_formatter(mtick.FuncFormatter(ticks))
 
-        if not subplot.plotTypeSimulation == BAR_PLOT:
+        if subplot.plotTypeSimulation != BAR_PLOT:
             ax.legend()
         ax.set_title(subplot.plotName)
         if subplot.xlim:
@@ -315,10 +320,8 @@ def ticks(y, _):
         ax.autoscale_view()
 
         # Beautify plots
-        ax.set_xlabel(
-            subplot.xLabel)
-        ax.set_ylabel(
-            subplot.yLabel)
+        ax.set_xlabel(subplot.xLabel)
+        ax.set_ylabel(subplot.yLabel)
 
     def generate_figure(
             self,

petab/visualize/plotting.py

@@ -72,10 +72,8 @@ def add_x_offset(self, offset) -> None:
             self.conditions += offset
 
     def add_y_offset(self, offset):
-        self.data_to_plot['mean'] = \
-            self.data_to_plot['mean'] + offset
-        self.data_to_plot['repl'] = \
-            self.data_to_plot['repl'] + offset
+        self.data_to_plot['mean'] += offset
+        self.data_to_plot['repl'] += offset
 
     def add_offsets(self, x_offset=0, y_offset=0) -> None:
         """
@@ -348,9 +346,9 @@ def save_to_tsv(self, output_file_path: str = 'visuSpec.tsv') -> None:
                                  dataplot.__dict__ if key in
                                  VISUALIZATION_DF_SINGLE_PLOT_LEVEL_COLS}
                 row = {**subplot_level, **dataset_level}
-                for key in row:
+                for key, value in row.items():
                     if key in visu_dict:
-                        visu_dict[key].append(row[key])
+                        visu_dict[key].append(value)
                     else:
                         visu_dict[key] = [row[key]]
         visu_df = pd.DataFrame.from_dict(visu_dict)
@@ -596,7 +594,6 @@ def get_data_to_plot(self, dataplot: DataPlot, provided_noise: bool
                                                        SIMULATION,
                                                        dataplot,
                                                        provided_noise)
-
         return measurements_to_plot, simulations_to_plot
 
 
@@ -874,11 +871,12 @@ def _get_vis_spec_dependent_columns_dict(
         plot_id_column = ['plot%s' % str(ind + 1) for ind, inner_list in
                           enumerate(dataset_id_list) for _ in inner_list]
 
-        columns_dict = {PLOT_ID: plot_id_column,
-                        DATASET_ID: dataset_id_column,
-                        LEGEND_ENTRY: dataset_label_column,
-                        Y_VALUES: yvalues_column}
-        return columns_dict
+        return {
+            PLOT_ID: plot_id_column,
+            DATASET_ID: dataset_id_column,
+            LEGEND_ENTRY: dataset_label_column,
+            Y_VALUES: yvalues_column
+        }
 
     def _create_legend(self, dataset_id: str) -> str:
         """