add circles plot, fixed #3

Author: jefferyUstc

dotplot/__init__.py

@@ -8,6 +8,9 @@
 from matplotlib import gridspec
 from matplotlib import pyplot as plt
 
+mpl.rcParams['pdf.fonttype'] = 42
+mpl.rcParams["font.sans-serif"] = "Arial"
+
 
 class DotPlot(object):
     DEFAULT_ITEM_HEIGHT = 0.3
@@ -17,20 +20,26 @@ class DotPlot(object):
 
     def __init__(self, df_size: pd.DataFrame,
                  df_color: Union[pd.DataFrame, None] = None,
+                 df_circle: Union[pd.DataFrame, None] = None,
                  ):
         """
         Construction a `DotPlot` object from `df_size` and `df_color`
 
         :param df_size: the DataFrame object represents the scatter size in dotplot
         :param df_color: the DataFrame object represents the color in dotplot
         """
-        __slots__ = ['size_data', 'color_data', 'height_item', 'width_item', 'resized_size_data']
+        __slots__ = ['size_data', 'resized_size_data',
+                     'color_data', 'height_item', 'width_item',
+                     'circle_data', 'resized_circle_data'
+                     ]
         if (df_color is not None) & (df_size.shape != df_color.shape):
             raise ValueError('df_size and df_color should have the same dimension')
         self.size_data = df_size
         self.color_data = df_color
+        self.circle_data = df_circle
         self.height_item, self.width_item = df_size.shape
         self.resized_size_data: pd.DataFrame
+        self.resized_circle_data: pd.DataFrame
 
     def __get_figure(self):
         _text_max = math.ceil(self.size_data.index.map(len).max() / 15)
@@ -46,12 +55,14 @@ def __get_figure(self):
                                width_ratios=[mainplot_width, self.DEFAULT_LEGENDS_WIDTH])
         ax = fig.add_subplot(gs[:, 0])
         ax_cbar = fig.add_subplot(gs[2, 1])
-        ax_legend = fig.add_subplot(gs[0:2, 1])
-        return ax, ax_cbar, ax_legend, fig
+        ax_sizes = fig.add_subplot(gs[0, 1])
+        ax_circles = fig.add_subplot(gs[1, 1])
+        return ax, ax_cbar, ax_sizes, ax_circles, fig
 
     @classmethod
     def parse_from_tidy_data(cls, data_frame: pd.DataFrame, item_key: str, group_key: str, sizes_key: str,
-                             color_key: str, selected_item: Union[None, Sequence] = None,
+                             color_key: Union[None, str] = None, circle_key: Union[None, str] = None,
+                             selected_item: Union[None, Sequence] = None,
                              selected_group: Union[None, Sequence] = None, *,
                              sizes_func: Union[None, Callable] = None, color_func: Union[None, Callable] = None
                              ):
@@ -69,34 +80,41 @@ class method for conveniently constructing DotPlot from tidy data
         :param selected_group: Same as `selected_item`, for group order and subset groups
         :param sizes_func:
         :param color_func:
+        :param circle_key:
         :return:
         """
-        data_frame = data_frame[[item_key, group_key, sizes_key, color_key]]
+        keys = [v for v in [item_key, group_key, sizes_key, color_key, circle_key] if v is not None]
+        data_frame = data_frame[keys]
         _original_item_order = data_frame[item_key].tolist()
         _original_item_order = _original_item_order[::-1]
         if sizes_func is not None:
             data_frame[sizes_key] = data_frame[sizes_key].map(sizes_func)
         if color_func is not None:
             data_frame[color_key] = data_frame[color_key].map(color_func)
-        data_frame = data_frame.pivot(index=item_key, columns=group_key, values=[color_key, sizes_key])
+        keys.remove(item_key)
+        keys.remove(group_key)
+        data_frame = data_frame.pivot(index=item_key, columns=group_key, values=keys)
         data_frame = data_frame.loc[_original_item_order, :]
         if selected_item is not None:
             data_frame = data_frame.loc[selected_item, :]
         if selected_group is not None:
             data_frame = data_frame.loc[:, selected_group]
-
         data_frame.columns = data_frame.columns.map(lambda x: '_'.join(x))
         data_frame = data_frame.fillna(0)
-        color_df = data_frame.loc[:, data_frame.columns.str.startswith(color_key)]
+
+        sizes_df, color_df, circle_df = (None, None, None)
         sizes_df = data_frame.loc[:, data_frame.columns.str.startswith(sizes_key)]
-        color_df.columns = color_df.columns.map(lambda x: '_'.join(x.split('_')[1:]))
-        sizes_df.columns = sizes_df.columns.map(lambda x: '_'.join(x.split('_')[1:]))
-        return cls(color_df, sizes_df)
+        if color_key is not None:
+            color_df = data_frame.loc[:, data_frame.columns.str.startswith(color_key)]
+        if circle_key is not None:
+            circle_df = data_frame.loc[:, data_frame.columns.str.startswith(circle_key)]
+        return cls(sizes_df, color_df, circle_df)
 
     def __get_coordinates(self, size_factor):
         X = list(range(1, self.width_item + 1)) * self.height_item
         Y = sorted(list(range(1, self.height_item + 1)) * self.width_item)
         self.resized_size_data = self.size_data.applymap(func=lambda x: x * size_factor)
+        self.resized_circle_data = self.circle_data.applymap(func=lambda x: x * size_factor)
         return X, Y
 
     def __draw_dotplot(self, ax, size_factor, cmap, vmin, vmax):
@@ -107,6 +125,10 @@ def __draw_dotplot(self, ax, size_factor, cmap, vmin, vmax):
         else:
             sct = ax.scatter(X, Y, c=self.color_data.values.flatten(), s=self.resized_size_data.values.flatten(),
                              edgecolors='none', linewidths=0, vmin=vmin, vmax=vmax, cmap=cmap)
+        sct_circle = None
+        if self.circle_data is not None:
+            sct_circle = ax.scatter(X, Y, c='', edgecolors='k', marker='o', linestyle='--',
+                                    s=self.resized_circle_data.values.flatten())
         width, height = self.width_item, self.height_item
         ax.set_xlim([0.5, width + 0.5])
         ax.set_ylim([0.6, height + 0.6])
@@ -116,13 +138,13 @@ def __draw_dotplot(self, ax, size_factor, cmap, vmin, vmax):
         ax.set_yticklabels(self.size_data.index.tolist())
         ax.tick_params(axis='y', length=5, labelsize=15, direction='out')
         ax.tick_params(axis='x', length=5, labelsize=15, direction='out')
-        return sct
+        return sct, sct_circle
 
     @staticmethod
     def __draw_color_bar(ax, sct: mpl.collections.PathCollection, cmap, vmin, vmax):
         gradient = np.linspace(1, 0, 500)
         gradient = gradient[:, np.newaxis]
-        im = ax.imshow(gradient, aspect='auto', cmap=cmap, origin='upper', extent=[.2, 0.3, 0.5, -0.5])
+        _ = ax.imshow(gradient, aspect='auto', cmap=cmap, origin='upper', extent=[.2, 0.3, 0.5, -0.5])
         ax.set_xticks([])
         ax.set_yticks([])
         ax_cbar2 = ax.twinx()
@@ -135,16 +157,25 @@ def __draw_color_bar(ax, sct: mpl.collections.PathCollection, cmap, vmin, vmax):
         _ = ax_cbar2.set_ylabel('-log10(pvalue)')
 
     @staticmethod
-    def __draw_legend(ax, sct: mpl.collections.PathCollection, size_factor):
+    def __draw_legend(ax, sct: mpl.collections.PathCollection, size_factor, title, circle=False, color=None):
+        print(id(sct))
         handles, labels = sct.legend_elements(prop="sizes", alpha=1,
                                               func=lambda x: x / size_factor,
-                                              color='#58000C')
+                                              color=color
+                                              )
         if len(handles) > 3:
             handles = np.asarray(handles)
             labels = np.asarray(labels)
             handles = handles[[0, math.ceil(len(handles) / 2), -1]]
             labels = labels[[0, math.ceil(len(labels) / 2), -1]]
-        _ = ax.legend(handles, labels, title="Sizes", loc='center left')  # bbox_to_anchor=(0.9, 0., 0.4, 0.4)
+        if circle:
+            from matplotlib.lines import Line2D
+            for i, _item in enumerate(handles):
+                xdata, ydata = _item.get_data()
+                marker_size = _item.get_markersize()
+                handles[i] = Line2D(xdata, ydata, color='white', marker='$\u25CC$',
+                                    markeredgecolor=color, markersize=marker_size)
+        _ = ax.legend(handles, labels, title=title, loc='center left')  # bbox_to_anchor=(0.9, 0., 0.4, 0.4)
         ax.set_xticks([])
         ax.set_yticks([])
         ax.spines['top'].set_visible(False)
@@ -165,9 +196,11 @@ def plot(self, size_factor: float = 15,
         :param cmap: color map supported by matplotlib
         :return:
         """
-        ax, ax_cbar, ax_legend, fig = self.__get_figure()
-        scatter = self.__draw_dotplot(ax, size_factor, cmap, vmin, vmax)
-        self.__draw_legend(ax_legend, scatter, size_factor)
+        ax, ax_cbar, ax_sizes, ax_circles, fig = self.__get_figure()
+        scatter, sct_circle = self.__draw_dotplot(ax, size_factor, cmap, vmin, vmax)
+        self.__draw_legend(ax_sizes, scatter, size_factor, title='Sizes', color='#58000C')
+        if sct_circle is not None:
+            self.__draw_legend(ax_circles, sct_circle, size_factor, title='Circles', circle=True, color='k')
         self.__draw_color_bar(ax_cbar, scatter, cmap, vmin, vmax)
         if path:
             fig.savefig(path, dpi=300, bbox_inches='tight')  #