forked from jefferyUstc/python-dotplot
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcore.py
More file actions
225 lines (204 loc) · 10.5 KB
/
Copy pathcore.py
File metadata and controls
225 lines (204 loc) · 10.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
import math
from os import PathLike
from typing import Union, Sequence, Callable
import matplotlib as mpl
import numpy as np
import pandas as pd
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
DEFAULT_ITEM_WIDTH = 0.3
DEFAULT_LEGENDS_WIDTH = .45
MIN_FIGURE_HEIGHT = 3
DEFAULT_BAND_ITEM_LENGTH = DEFAULT_ITEM_HEIGHT
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,
):
"""
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', 'resized_size_data',
'color_data', 'height_item', 'width_item',
'circle_data', 'resized_circle_data', 'annotation_data'
]
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')
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
self.resized_size_data: Union[pd.DataFrame, None] = None
self.resized_circle_data: Union[pd.DataFrame, None] = None
def __get_figure(self):
_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
)
if self.annotation_data is not None:
pass
figure_height = max([self.MIN_FIGURE_HEIGHT, mainplot_height])
figure_width = mainplot_width + self.DEFAULT_LEGENDS_WIDTH
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])
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: 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
):
"""
class method for conveniently constructing DotPlot from tidy data
:param data_frame:
:param item_key:
:param group_key:
:param sizes_key:
:param color_key:
:param selected_item: default None, if specified, this should be subsets of `item_key` in `data_frame`
alternatively, this param can be used as self-defined item order definition.
:param selected_group: Same as `selected_item`, for group order and subset groups
:param sizes_func:
:param color_func:
:param circle_key:
:return:
"""
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)
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)
sizes_df, color_df, circle_df = None, None, None
sizes_df = data_frame.loc[:, data_frame.columns.str.startswith(sizes_key)]
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)
if self.circle_data is not None:
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):
X, Y = self.__get_coordinates(size_factor)
if self.color_data is None:
sct = ax.scatter(X, Y, c='r', cmap=cmap, s=self.resized_size_data.values.flatten(),
edgecolors='none', linewidths=0, vmin=vmin, vmax=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])
ax.set_xticks(range(1, width + 1))
ax.set_yticks(range(1, height + 1))
ax.set_xticklabels(self.size_data.columns.tolist(), rotation='vertical')
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, 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]
_ = 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()
_ = ax_cbar2.set_yticks([0, 1000])
if vmax is None:
vmax = math.ceil(sct.get_array().max())
if vmin is None:
vmin = math.floor(sct.get_array().min())
_ = ax_cbar2.set_yticklabels([vmin, vmax])
_ = ax_cbar2.set_ylabel('-log10(pvalue)')
@staticmethod
def __draw_legend(ax, sct: mpl.collections.PathCollection, size_factor, title, circle=False, color=None):
handles, labels = sct.legend_elements(prop="sizes", alpha=1,
func=lambda x: x / size_factor,
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]]
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)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
def plot(self, size_factor: float = 15,
vmin: float = 0, vmax: float = None,
path: Union[PathLike, None] = None,
cmap: Union[str, mpl.colors.Colormap] = 'Reds'):
"""
:param size_factor: `size factor` * `value` for the actually representation of scatter size in the final figure
:param vmin: `vmin` in `matplotlib.pyplot.scatter`
:param vmax: `vmax` in `matplotlib.pyplot.scatter`
:param path: path to save the figure
:param cmap: color map supported by matplotlib
:return:
"""
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')
else:
ax_circles.axis('off')
self.__draw_color_bar(ax_cbar, scatter, cmap, vmin, vmax)
if path:
fig.savefig(path, dpi=300, bbox_inches='tight') #
return scatter
def __str__(self):
return 'DotPlot object with data point in shape %s' % str(self.size_data.shape)
__repr__ = __str__