fix beeswarm plot colorbar issue (#18)

Author: jlyang1990

fasttreeshap/plots/_beeswarm.py

@@ -361,14 +361,14 @@ def beeswarm(shap_values, max_display=10, order=Explanation.abs.mean(0),
         import matplotlib.cm as cm
         m = cm.ScalarMappable(cmap=color)
         m.set_array([0, 1])
-        cb = pl.colorbar(m, ticks=[0, 1], aspect=1000)
+        cb = pl.colorbar(m, ax=pl.gca(), ticks=[0, 1], aspect=80)
         cb.set_ticklabels([labels['FEATURE_VALUE_LOW'], labels['FEATURE_VALUE_HIGH']])
         cb.set_label(color_bar_label, size=12, labelpad=0)
         cb.ax.tick_params(labelsize=11, length=0)
         cb.set_alpha(1)
         cb.outline.set_visible(False)
-        bbox = cb.ax.get_window_extent().transformed(pl.gcf().dpi_scale_trans.inverted())
-        cb.ax.set_aspect((bbox.height - 0.9) * 20)
+        # bbox = cb.ax.get_window_extent().transformed(pl.gcf().dpi_scale_trans.inverted())
+        # cb.ax.set_aspect((bbox.height - 0.9) * 20)
         # cb.draw_all()
 
     pl.gca().xaxis.set_ticks_position('bottom')
@@ -861,14 +861,14 @@ def summary_legacy(shap_values, features=None, feature_names=None, max_display=N
         import matplotlib.cm as cm
         m = cm.ScalarMappable(cmap=cmap if plot_type != "layered_violin" else pl.get_cmap(color))
         m.set_array([0, 1])
-        cb = pl.colorbar(m, ticks=[0, 1], aspect=1000)
+        cb = pl.colorbar(m, ax=pl.gca(), ticks=[0, 1], aspect=80)
         cb.set_ticklabels([labels['FEATURE_VALUE_LOW'], labels['FEATURE_VALUE_HIGH']])
         cb.set_label(color_bar_label, size=12, labelpad=0)
         cb.ax.tick_params(labelsize=11, length=0)
         cb.set_alpha(1)
         cb.outline.set_visible(False)
-        bbox = cb.ax.get_window_extent().transformed(pl.gcf().dpi_scale_trans.inverted())
-        cb.ax.set_aspect((bbox.height - 0.9) * 20)
+        # bbox = cb.ax.get_window_extent().transformed(pl.gcf().dpi_scale_trans.inverted())
+        # cb.ax.set_aspect((bbox.height - 0.9) * 20)
         # cb.draw_all()
 
     pl.gca().xaxis.set_ticks_position('bottom')

fasttreeshap/plots/_scatter.py

@@ -343,19 +343,19 @@ def scatter(shap_values, color="#1E88E5", hist=True, axis_color="#333333", cmap=
             tick_positions = np.array([cname_map[n] for n in cnames])
             tick_positions *= 1 - 1 / len(cnames)
             tick_positions += 0.5 * (chigh - clow) / (chigh - clow + 1)
-            cb = pl.colorbar(p, ticks=tick_positions, ax=ax)
+            cb = pl.colorbar(p, ticks=tick_positions, ax=ax, aspect=80)
             cb.set_ticklabels(cnames)
         else:
-            cb = pl.colorbar(p, ax=ax)
+            cb = pl.colorbar(p, ax=ax, aspect=80)
 
         cb.set_label(feature_names[interaction_index], size=13)
         cb.ax.tick_params(labelsize=11)
         if categorical_interaction:
             cb.ax.tick_params(length=0)
         cb.set_alpha(1)
         cb.outline.set_visible(False)
-        bbox = cb.ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
-        cb.ax.set_aspect((bbox.height - 0.7) * 20)
+        # bbox = cb.ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+        # cb.ax.set_aspect((bbox.height - 0.7) * 20)
 
     # handles any setting of xmax and xmin
     # note that we handle None,float, or "percentile(float)" formats
@@ -690,19 +690,19 @@ def dependence_legacy(ind, shap_values=None, features=None, feature_names=None,
             if len(tick_positions) == 2:
                 tick_positions[0] -= 0.25
                 tick_positions[1] += 0.25
-            cb = pl.colorbar(p, ticks=tick_positions, ax=ax)
+            cb = pl.colorbar(p, ticks=tick_positions, ax=ax, aspect=80)
             cb.set_ticklabels(cnames)
         else:
-            cb = pl.colorbar(p, ax=ax)
+            cb = pl.colorbar(p, ax=ax, aspect=80)
 
         cb.set_label(feature_names[interaction_index], size=13)
         cb.ax.tick_params(labelsize=11)
         if categorical_interaction:
             cb.ax.tick_params(length=0)
         cb.set_alpha(1)
         cb.outline.set_visible(False)
-        bbox = cb.ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
-        cb.ax.set_aspect((bbox.height - 0.7) * 20)
+        # bbox = cb.ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+        # cb.ax.set_aspect((bbox.height - 0.7) * 20)
 
     # handles any setting of xmax and xmin
     # note that we handle None,float, or "percentile(float)" formats