add color band annotation, #4

Author: jefferyUstc

dotplot/annotation_bands.py

@@ -0,0 +1,88 @@
+"""
+module for draw annotation bands.
+PS. This implementation mainly refers to the way of seaborn
+"""
+
+import itertools
+from typing import Union, Sequence
+
+import matplotlib as mpl
+import numpy as np
+import pandas as pd
+from matplotlib.colors import to_rgb
+
+mpl.rcParams['pdf.fonttype'] = 42
+mpl.rcParams["font.sans-serif"] = "Arial"
+
+
+def _process_colors(colors: Union[pd.DataFrame, pd.Series],
+                    index_order: Union[Sequence[Union[str, int]], None] = None):
+    if not isinstance(colors, (pd.DataFrame, pd.Series)):
+        raise TypeError('`colors` should be pandas.DataFrame or pandas.Series')
+    if index_order is not None:
+        colors = colors.reindex(index_order)
+    colors = colors.astype(object).fillna('white')  # TODO, to alpha
+    if isinstance(colors, pd.DataFrame):
+        labels = list(colors.columns)
+        colors = colors.T.values
+    else:
+        labels = [''] if colors.name is None else [colors.name]
+        colors = colors.values
+    try:
+        to_rgb(colors[0])
+        colors = list(map(to_rgb, colors))
+    except ValueError:
+        colors = [list(map(to_rgb, item)) for item in colors]
+    return colors, labels
+
+
+def _color_list_to_matrix_and_cmap(colors, axis=0):
+    if any(issubclass(type(item), list) for item in colors):
+        all_colors = set(itertools.chain(*colors))
+        n = len(colors)  # number of fields
+        m = len(colors[0])  # number of observations
+    else:
+        all_colors = set(colors)
+        n = 1
+        m = len(colors)
+        colors = [colors]
+    color_to_value = dict((col, i) for i, col in enumerate(all_colors))
+
+    matrix = np.array([color_to_value[c]
+                       for color in colors for c in color])
+
+    shape = (n, m)
+    matrix = matrix.reshape(shape)
+    if axis == 0:
+        # row-side:
+        matrix = matrix.T
+
+    cmap = mpl.colors.ListedColormap(all_colors)
+    return matrix, cmap
+
+
+def _determine_ticks(labels: Sequence[str]):
+    num = len(labels)
+    return [item + .5 for item in range(num)]
+
+
+def draw_heatmap(colors: Union[pd.DataFrame, pd.Series],
+                 axes: mpl.axes.Axes, axis=0,
+                 index_order: Union[Sequence[Union[str, int]], None] = None, **kwargs):
+    # index_order reverse
+    colors = colors.copy()
+    if index_order is not None:
+        index_order = list(index_order)[::-1]
+    colors, labels = _process_colors(colors, index_order)
+    matrix, cmap = _color_list_to_matrix_and_cmap(colors, axis=axis)
+    axes.pcolormesh(matrix, cmap=cmap, **kwargs)
+    if axis == 1:
+        axes.set_yticks(_determine_ticks(labels))
+        axes.set_yticklabels(labels)
+        _ = axes.set_xticks([])
+    elif axis == 0:
+        axes.set_xticks(_determine_ticks(labels))
+        axes.set_xticklabels(labels, rotation=90)
+        _ = axes.set_yticks([])
+    for item in ['left', 'right', 'bottom', 'top']:
+        axes.spines[item].set_visible(False)

dotplot/core.py

@@ -17,13 +17,13 @@ class DotPlot(object):
     DEFAULT_ITEM_WIDTH = 0.3
     DEFAULT_LEGENDS_WIDTH = .45
     MIN_FIGURE_HEIGHT = 3
-    DEFAULT_BAND_ITEM_LENGTH = DEFAULT_ITEM_HEIGHT
+    DEFAULT_BAND_ITEM_LENGTH = .2
 
-    # TODO implement annotation band
     def __init__(self, df_size: pd.DataFrame,
                  df_color: Union[pd.DataFrame, None] = None,
                  df_circle: Union[pd.DataFrame, None] = None,
-                 df_annotation: Union[pd.DataFrame, None] = None,
+                 row_colors: Union[pd.DataFrame, None] = None,
+                 col_colors: Union[pd.DataFrame, None] = None,
                  ):
         """
         Construction a `DotPlot` object from `df_size` and `df_color`
@@ -33,46 +33,71 @@ def __init__(self, df_size: pd.DataFrame,
         """
         __slots__ = ['size_data', 'resized_size_data',
                      'color_data', 'height_item', 'width_item',
-                     'circle_data', 'resized_circle_data', 'annotation_data'
+                     'circle_data', 'resized_circle_data', 'row_colors', 'col_colors'
                      ]
         if df_color is not None and df_size.shape != df_color.shape:
             raise ValueError('df_size and df_color should have the same dimension')
         if df_circle is not None and df_size.shape != df_circle.shape:
             raise ValueError('df_size and df_circle should have the same dimension')
-        if df_annotation is not None and df_size.shape != df_annotation.shape:
-            raise ValueError('df_size and df_annotation should have the same row number')
+        if row_colors is not None and df_size.shape[0] != len(row_colors):
+            raise ValueError('row_colors has the wrong shape')
+        if col_colors is not None and df_size.shape[1] != len(col_colors):
+            raise ValueError('col_colors has the wrong shape')
 
         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.annotation_data = df_annotation
+        # TODO code logic need to argument
+        self.row_colors = row_colors
+        self.col_colors = col_colors
         self.resized_size_data: Union[pd.DataFrame, None] = None
         self.resized_circle_data: Union[pd.DataFrame, None] = None
 
     def __get_figure(self):
+        """
+        Figure layout
+        :return:
+        """
         _text_max = math.ceil(self.size_data.index.map(len).max() / 15)
         mainplot_height = self.height_item * self.DEFAULT_ITEM_HEIGHT
         mainplot_width = (
                 (_text_max + self.width_item) * self.DEFAULT_ITEM_WIDTH
         )
         figure_height = max([self.MIN_FIGURE_HEIGHT, mainplot_height])
         figure_width = mainplot_width + self.DEFAULT_LEGENDS_WIDTH
-        if self.annotation_data is not None:
-            # figure_width = figure_width + self.DEFAULT_BAND_ITEM_LENGTH * self.annotation_data.shape[1]
-            ...
+        band_width, band_height = 0., 0.
+        if self.row_colors is not None:
+            band_width = self.DEFAULT_BAND_ITEM_LENGTH * self.row_colors.shape[1]
+        if self.col_colors is not None:
+            band_height = self.DEFAULT_BAND_ITEM_LENGTH * self.col_colors.shape[1]
+        figure_width = figure_width + band_width
+        figure_height = figure_height + band_height
+
         plt.style.use('seaborn-white')
         fig = plt.figure(figsize=(figure_width, figure_height))
-        gs = gridspec.GridSpec(nrows=3, ncols=2, wspace=0.15, hspace=0.15,
-                               width_ratios=[mainplot_width, self.DEFAULT_LEGENDS_WIDTH])
-        ax = fig.add_subplot(gs[:, 0])
-        ax_cbar = fig.add_subplot(gs[2, 1])
-        ax_sizes = fig.add_subplot(gs[0, 1])
-        ax_circles = fig.add_subplot(gs[1, 1])
+        gs = gridspec.GridSpec(nrows=2, ncols=3, wspace=0.05, hspace=0.02,
+                               width_ratios=[mainplot_width, band_width, self.DEFAULT_LEGENDS_WIDTH],
+                               height_ratios=[band_height, mainplot_height]
+                               )
+        ax = fig.add_subplot(gs[1, 0])
+        ax_row_bands = fig.add_subplot(gs[1, 1])
+        ax_col_bands = fig.add_subplot(gs[0, 0])
+        ax_abandon = fig.add_subplot(gs[0, 1])
+        legend_gs = gridspec.GridSpecFromSubplotSpec(3, 1, hspace=.1, subplot_spec=gs[1, 2])
+        ax_sizes = fig.add_subplot(legend_gs[0, 0])
+        ax_circles = fig.add_subplot(legend_gs[1, 0])
+        ax_cbar = fig.add_subplot(legend_gs[2, 0])
+
+        _, _ = ax_sizes.axis('off'), ax_circles.axis('off')
+        if self.col_colors is None:
+            ax_col_bands.axis('off')
+        if self.row_colors is None:
+            ax_row_bands.axis('off')
         if self.color_data is None:
             ax_cbar.axis('off')
-        ax_circles.axis('off')
-        return ax, ax_cbar, ax_sizes, ax_circles, fig
+        ax_abandon.axis('off')
+        return ax, ax_cbar, ax_sizes, ax_circles, ax_row_bands, ax_col_bands, fig
 
     @classmethod
     def parse_from_tidy_data(cls, data_frame: pd.DataFrame, item_key: str, group_key: str, sizes_key: str,
@@ -126,7 +151,7 @@ class method for conveniently constructing DotPlot from tidy data
             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):
+    def __get_coordinates(self):
         X = list(range(1, self.width_item + 1)) * self.height_item
         Y = sorted(list(range(1, self.height_item + 1)) * self.width_item)
         return X, Y
@@ -138,7 +163,7 @@ def __draw_dotplot(self, ax, size_factor, cmap, vmin, vmax, **kws):
         for _value in ['dot_title', 'circle_title', 'colorbar_title', 'dot_color', 'circle_color']:
             _ = kws.pop(_value, None)
 
-        X, Y = self.__get_coordinates(size_factor)
+        X, Y = self.__get_coordinates()
         if self.color_data is None:
             sct = ax.scatter(X, Y, c=dot_color, s=self.resized_size_data.values.flatten(),
                              edgecolors='none', linewidths=0, vmin=vmin, vmax=vmax, cmap=cmap, **kws)
@@ -234,7 +259,9 @@ def plot(self, size_factor: float = 15,
              path: Union[PathLike, None] = None,
              cmap: Union[str, mpl.colors.Colormap] = 'Reds',
              cluster_row: bool = False, cluster_col: bool = False,
-             cluster_kws: Union[Dict, None] = None, **kwargs
+             cluster_kws: Union[Dict, None] = None,
+             color_band_kws: Union[Dict, None] = None,
+             **kwargs
              ):
         """
 
@@ -248,12 +275,13 @@ def plot(self, size_factor: float = 15,
         :param cluster_kws, key args for cluster, including `cluster_method`, `cluster_metric`， 'cluster_n'
         :param kwargs: dot_title, circle_title, colorbar_title, dot_color, circle_color
                     other kwargs are passed to `matplotlib.Axes.scatter`
+        :param color_band_kws: this kwargs was passed to `matplotlib.axes.Axes.pcolormesh`
         :return:
         """
         self.__preprocess_data(size_factor, cluster_row=cluster_row, cluster_col=cluster_col,
                                **cluster_kws if cluster_kws is not None else {}
                                )
-        ax, ax_cbar, ax_sizes, ax_circles, fig = self.__get_figure()
+        ax, ax_cbar, ax_sizes, ax_circles, ax_row_bands, ax_col_bands, fig = self.__get_figure()
         scatter, sct_circle = self.__draw_dotplot(ax, size_factor, cmap, vmin, vmax)
         self.__draw_legend(ax_sizes, scatter, size_factor,
                            color=kwargs.get('dot_color', '#58000C'),  # dot legend color
@@ -266,8 +294,20 @@ def plot(self, size_factor: float = 15,
         if self.color_data is not None:
             self.__draw_color_bar(ax_cbar, scatter, cmap, vmin, vmax,
                                   ylabel=kwargs.get('colorbar_title', '-log10(pvalue)'))
+
+        if self.col_colors is not None:
+            from .annotation_bands import draw_heatmap
+            color_band_kws = {} if color_band_kws is None else color_band_kws
+            draw_heatmap(self.col_colors, axes=ax_col_bands,
+                         index_order=self.size_data.columns.tolist(), axis=1, **color_band_kws)
+        if self.row_colors is not None:
+            color_band_kws = {} if color_band_kws is None else color_band_kws
+            from .annotation_bands import draw_heatmap
+            draw_heatmap(self.row_colors, axes=ax_row_bands,
+                         index_order=self.size_data.index.tolist(), axis=0, **color_band_kws)
+
         if path:
-            fig.savefig(path, dpi=300, bbox_inches='tight')  #
+            fig.savefig(path, dpi=300, bbox_inches='tight')
         return scatter
 
     def __str__(self):