Heatmap updates

docs/tutorials/plotting/heatmaps.ipynb

@@ -25,18 +25,22 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Plotting graphs using heatmap\n",
-    "\n",
-    "### Simulate graphs using weighted stochastic block models\n",
-    "The 2-block model is defined as below:\n",
+    "## Plotting Simple Graphs using heatmap"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A 2-block Stochastic Block Model is defined as below:\n",
     "\n",
     "\\begin{align*}\n",
     "P = \\begin{bmatrix}0.8 & 0.2 \\\\\n",
     "0.2 & 0.8 \n",
     "\\end{bmatrix}\n",
     "\\end{align*}\n",
     "\n",
-    "We generate two weight SBMs where the weights are distributed from a Poisson(3) and Normal(5, 1)."
+    "In simple cases, the model is unweighted. Below, we plot an unweighted SBM."
    ]
   },
   {
@@ -46,25 +50,28 @@
    "outputs": [],
    "source": [
     "from graspologic.simulations import sbm\n",
+    "from graspologic.plot import heatmap\n",
     "\n",
     "n_communities = [50, 50]\n",
     "p = [[0.8, 0.2], \n",
     "     [0.2, 0.8]]\n",
     "\n",
-    "wt = np.random.poisson\n",
-    "wtargs = dict(lam=3)\n",
-    "A_poisson= sbm(n_communities, p, wt=wt, wtargs=wtargs)\n",
-    "\n",
-    "wt = np.random.normal\n",
-    "wtargs = dict(loc=5, scale=1)\n",
-    "A_normal = sbm(n_communities, p, wt=wt, wtargs=wtargs)"
+    "A, labels = sbm(n_communities, p, return_labels=True)\n",
+    "heatmap(A, title=\"Basic Heatmap function\");"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plotting Simple Graphs using binary_heatmap"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Plot the simulated weighted SBMs"
+    "In the unweighted case, we might want to plot a heatmap with binary labels. The `binary_heatmap` utility function is a wrapper on `heatmap` made for this common simple use case."
    ]
   },
   {
@@ -73,10 +80,81 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from graspologic.plot import heatmap\n",
+    "from graspologic.plot import binary_heatmap\n",
     "\n",
-    "title = 'Weighted Stochastic Block Model with Poisson(3)'\n",
+    "binary_heatmap(A, colors=[\"white\", \"black\"], \n",
+    "               colorbar_ticklabels=[\"Edge\", \"No Edge\"],\n",
+    "               title=\"Plotting with Binary Heatmap\");"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plotting with Hierarchy Labels"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If we have labels, we can use them to show communities on a Heatmap."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "heatmap(A, inner_hier_labels=labels)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can plot outer hierarchy labels in addition to inner hierarchy labels."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outer_labels = [\"Outer Labels\"] * 100\n",
+    "heatmap(A, inner_hier_labels=labels,\n",
+    "        outer_hier_labels=outer_labels)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Weighted SBMs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can also use heatmap when our graph is weighted. Here, we generate two weighted SBMs where the weights are distributed from a Poisson(3) and Normal(5, 1)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Draw weights from a Poisson(3) distribution\n",
+    "wt = np.random.poisson\n",
+    "wtargs = dict(lam=3)\n",
+    "A_poisson= sbm(n_communities, p, wt=wt, wtargs=wtargs)\n",
     "\n",
+    "# Plot\n",
+    "title = 'Weighted Stochastic Block Model with \\n weights drawn from a Poisson(3) distribution'\n",
     "fig= heatmap(A_poisson, title=title)"
    ]
   },
@@ -86,18 +164,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "title = 'Weighted Stochastic Block Model with Normal(5, 1)'\n",
+    "# Draw weights from a Normal(5, 1) distribution\n",
+    "wt = np.random.normal\n",
+    "wtargs = dict(loc=5, scale=1)\n",
+    "A_normal = sbm(n_communities, p, wt=wt, wtargs=wtargs)\n",
     "\n",
-    "fig= heatmap(A_normal, title=title)"
+    "# Plot\n",
+    "title = 'Weighted Stochastic Block Model with \\n weights drawn from a Normal(5, 1) distribution'\n",
+    "fig = heatmap(A_normal, title=title)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### You can also change color maps\n",
+    "### Colormaps\n",
     "\n",
-    "See [here](https://matplotlib.org/tutorials/colors/colormaps.html) for a list of colormaps"
+    "You can change colormaps. See [here](https://matplotlib.org/tutorials/colors/colormaps.html) for a list of colormaps."
    ]
   },
   {
@@ -107,8 +190,7 @@
    "outputs": [],
    "source": [
     "title = 'Weighted Stochastic Block Model with Poisson(3)'\n",
-    "\n",
-    "fig= heatmap(A_poisson, title=title, transform=None, cmap=\"binary\", center=None)"
+    "fig = heatmap(A_poisson, title=title, transform=None, cmap=\"binary\", center=None)"
    ]
   },
   {
@@ -203,7 +285,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.0"
+   "version": "3.8.5"
   }
  },
  "nbformat": 4,

graspologic/plot/__init__.py

@@ -6,6 +6,7 @@
 
 from .plot import (
     heatmap,
+    binary_heatmap,
     gridplot,
     pairplot,
     pairplot_with_gmm,
@@ -17,6 +18,7 @@
 
 __all__ = [
     "heatmap",
+    "binary_heatmap",
     "gridplot",
     "pairplot",
     "pairplot_with_gmm",

graspologic/plot/plot.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation and contributors.
 # Licensed under the MIT License.
 
+import warnings
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
@@ -262,18 +263,22 @@ def heatmap(
         if len(xticklabels) != X.shape[1]:
             msg = "xticklabels must have same length {}.".format(X.shape[1])
             raise ValueError(msg)
-    elif not isinstance(xticklabels, bool):
-        msg = "xticklabels must be a bool or a list, not {}".format(type(xticklabels))
+
+    elif not isinstance(xticklabels, (bool, int)):
+        msg = "xticklabels must be a bool, int, or a list, not {}".format(
+            type(xticklabels)
+        )
         raise TypeError(msg)
 
     if isinstance(yticklabels, list):
         if len(yticklabels) != X.shape[0]:
             msg = "yticklabels must have same length {}.".format(X.shape[0])
             raise ValueError(msg)
-    elif not isinstance(yticklabels, bool):
-        msg = "yticklabels must be a bool or a list, not {}".format(type(yticklabels))
+    elif not isinstance(yticklabels, (bool, int)):
+        msg = "yticklabels must be a bool, int, or a list, not {}".format(
+            type(yticklabels)
+        )
         raise TypeError(msg)
-
     # Handle cmap
     if not isinstance(cmap, (str, list, Colormap)):
         msg = "cmap must be a string, list of colors, or matplotlib.colors.Colormap,"
@@ -291,6 +296,11 @@ def heatmap(
         msg = "cbar must be a bool, not {}.".format(type(center))
         raise TypeError(msg)
 
+    # Warning on labels
+    if (inner_hier_labels is None) and (outer_hier_labels is not None):
+        msg = "outer_hier_labels requires inner_hier_labels to be used."
+        warnings.warn(msg)
+
     arr = import_graph(X)
 
     arr = _process_graphs(
@@ -341,6 +351,53 @@ def heatmap(
     return plot
 
 
+def binary_heatmap(
+    X, colors=["white", "black"], colorbar_ticklabels=["No Edge", "Edge"], **kwargs
+):
+    """
+    Plots an unweighted graph as a black-and-white matrix
+    with a binary colorbar.
+
+    Takes the same keyword arguments as ``plot.heatmap``.
+
+    Parameters
+    ----------
+    X : nx.Graph or np.ndarray object
+        Unweighted graph or numpy matrix to plot.
+
+    colors : list-like or np.ndarray
+        A list of exactly two colors to use for the heatmap.
+
+    colorbar_ticklabels : list-like
+        Binary labels to use in the colorbar.
+
+    **kwargs : dict, optional
+        All keyword arguments in ``plot.heatmap``.
+
+    """
+    if len(colors) != 2:
+        raise ValueError("Colors must be length 2")
+    if "center" in kwargs:
+        raise ValueError("Center is not allowed for binary heatmaps.")
+    if "cmap" in kwargs:
+        raise ValueError(
+            "cmap is not allowed in a binary heatmap. To change colors, use the `colors` parameter."
+        )
+    if not (
+        isinstance(colorbar_ticklabels, (list, tuple)) and len(colorbar_ticklabels) == 2):
+        raise ValueError("colorbar_ticklabels must be list-like and length 2.")
+
+    cmap = mpl.colors.ListedColormap(colors)
+    ax = heatmap(X, center=None, cmap=cmap, **kwargs)
+    colorbar = ax.collections[0].colorbar
+    cbar = kwargs.setdefault("cbar", True)
+    if cbar:
+        colorbar.set_ticks([0.25, 0.75])
+        colorbar.set_ticklabels(colorbar_ticklabels)
+        colorbar.ax.set_frame_on(True)
+    return ax
+
+
 def gridplot(
     X,
     labels=None,
@@ -1267,6 +1324,8 @@ def _unique_like(vals):
 # assume that the graph has already been plotted in sorted form
 def _plot_groups(ax, graph, inner_labels, outer_labels=None, fontsize=30):
     inner_labels = np.array(inner_labels)
+    if outer_labels is not None:
+        outer_labels = np.array(outer_labels)
     plot_outer = True
     if outer_labels is None:
         outer_labels = np.ones_like(inner_labels)