Updated swarmplot to produce legends if plotting by hue column

Author: sunroofgod

dabest/plot_tools.py

@@ -1305,6 +1305,7 @@ def plot(
             0  # x-coordinate of center of each individual swarm of the swarm plot
         )
         x_tick_tabels = []
+
         for group_i, values_i in self.__data_copy.groupby(self.__x):
             x_new = []
             values_i_y = values_i[self.__y]
@@ -1363,16 +1364,31 @@ def plot(
                     )
                 else:
                     facecolor = "none" if not filled else self.__palette[group_i]
+
                 ax.scatter(
                     values_i["x_new"],
                     values_i[self.__y],
                     s=self.__size,
                     zorder=self.__zorder,
                     facecolor=facecolor,
                     edgecolor=self.__palette[group_i],
+                    label=group_i,
                     **kwargs,
                 )
 
+        # Handling of legends
+        # This is currently a workaround because c and cmap is unable to map the labels when calling scatter()
+        # labels has to be used to designate legend labels and handles in scatter() due to the potential calling of ax.get_legend_handles_labels()
+        if self.__hue is not None:
+            for cmap_group_i in self.__palette:
+                ax.scatter(
+                    [],
+                    [],
+                    c=self.__palette[cmap_group_i],
+                    label=cmap_group_i,
+                )
+            handles, labels = ax.get_legend_handles_labels()
+
         ax.get_xaxis().set_ticks(np.arange(x_position))
         ax.get_xaxis().set_ticklabels(x_tick_tabels)
 

nbs/API/plot_tools.ipynb

@@ -1368,6 +1368,7 @@
     "            0  # x-coordinate of center of each individual swarm of the swarm plot\n",
     "        )\n",
     "        x_tick_tabels = []\n",
+    "\n",
     "        for group_i, values_i in self.__data_copy.groupby(self.__x):\n",
     "            x_new = []\n",
     "            values_i_y = values_i[self.__y]\n",
@@ -1426,16 +1427,31 @@
     "                    )\n",
     "                else:\n",
     "                    facecolor = \"none\" if not filled else self.__palette[group_i]\n",
+    "\n",
     "                ax.scatter(\n",
     "                    values_i[\"x_new\"],\n",
     "                    values_i[self.__y],\n",
     "                    s=self.__size,\n",
     "                    zorder=self.__zorder,\n",
     "                    facecolor=facecolor,\n",
     "                    edgecolor=self.__palette[group_i],\n",
+    "                    label=group_i,\n",
     "                    **kwargs,\n",
     "                )\n",
     "\n",
+    "        # Handling of legends\n",
+    "        # This is currently a workaround because c and cmap is unable to map the labels when calling scatter()\n",
+    "        # labels has to be used to designate legend labels and handles in scatter() due to the potential calling of ax.get_legend_handles_labels()\n",
+    "        if self.__hue is not None:\n",
+    "            for cmap_group_i in self.__palette:\n",
+    "                ax.scatter(\n",
+    "                    [],\n",
+    "                    [],\n",
+    "                    c=self.__palette[cmap_group_i],\n",
+    "                    label=cmap_group_i,\n",
+    "                )\n",
+    "            handles, labels = ax.get_legend_handles_labels()\n",
+    "\n",
     "        ax.get_xaxis().set_ticks(np.arange(x_position))\n",
     "        ax.get_xaxis().set_ticklabels(x_tick_tabels)\n",
     "\n",