diff --git a/20news.ipynb b/20news.ipynb
index 5f4c9f1..9b04925 100644
--- a/20news.ipynb
+++ b/20news.ipynb
@@ -223,37 +223,13 @@
    },
    "outputs": [],
    "source": [
-    "def coarsen(A, levels):\n",
-    "    graphs, parents = coarsening.metis(A, levels)\n",
-    "    perms = coarsening.compute_perm(parents)\n",
-    "\n",
-    "    laplacians = []\n",
-    "    for i,A in enumerate(graphs):\n",
-    "        M, M = A.shape\n",
-    "\n",
-    "        # No self-connections.\n",
-    "        if True:\n",
-    "            A = A.tocoo()\n",
-    "            A.setdiag(0)\n",
-    "\n",
-    "        if i < levels:\n",
-    "            A = coarsening.perm_adjacency(A, perms[i])\n",
-    "\n",
-    "        A = A.tocsr()\n",
-    "        A.eliminate_zeros()\n",
-    "        Mnew, Mnew = A.shape\n",
-    "        print('Layer {0}: M_{0} = |V| = {1} nodes ({2} added), |E| = {3} edges'.format(i, Mnew, Mnew-M, A.nnz//2))\n",
-    "\n",
-    "        L = graph.laplacian(A, normalized=FLAGS.normalized_laplacian)\n",
-    "        laplacians.append(L)\n",
-    "    return laplacians, perms[0] if len(perms) > 0 else None\n",
-    "\n",
     "t_start = time.process_time()\n",
     "dist, idx = graph.distance_sklearn_metrics(graph_data, k=FLAGS.number_edges, metric=FLAGS.metric)\n",
     "A = graph.adjacency(dist, idx)\n",
     "print(\"{} > {} edges\".format(A.nnz//2, FLAGS.number_edges*graph_data.shape[0]//2))\n",
     "A = graph.replace_random_edges(A, 0)\n",
-    "L, perm = coarsen(A, FLAGS.coarsening_levels)\n",
+    "graphs, perm = coarsening.coarsen(A, levels=FLAGS.coarsening_levels, self_connections=False)\n",
+    "L = [graph.laplacian(A, normalized=True) for A in graphs]\n",
     "print('Execution time: {:.2f}s'.format(time.process_time() - t_start))\n",
     "#graph.plot_spectrum(L)\n",
     "#del graph_data, A, dist, idx"
diff --git a/coarsening.py b/coarsening.py
index 8a67f95..fc33855 100644
--- a/coarsening.py
+++ b/coarsening.py
@@ -2,6 +2,35 @@
 import scipy.sparse
 
 
+def coarsen(A, levels, self_connections=False):
+    """
+    Coarsen a graph, represented by its adjacency matrix A, at multiple
+    levels.
+    """
+    graphs, parents = metis(A, levels)
+    perms = compute_perm(parents)
+
+    for i, A in enumerate(graphs):
+        M, M = A.shape
+
+        if not self_connections:
+            A = A.tocoo()
+            A.setdiag(0)
+
+        if i < levels:
+            A = perm_adjacency(A, perms[i])
+
+        A = A.tocsr()
+        A.eliminate_zeros()
+        graphs[i] = A
+
+        Mnew, Mnew = A.shape
+        print('Layer {0}: M_{0} = |V| = {1} nodes ({2} added),'
+              '|E| = {3} edges'.format(i, Mnew, Mnew-M, A.nnz//2))
+
+    return graphs, perms[0] if levels > 0 else None
+
+
 def metis(W, levels, rid=None):
     """
     Coarsen a graph multiple times using the METIS algorithm.
diff --git a/mnist.ipynb b/mnist.ipynb
index 1cf29d1..c560008 100644
--- a/mnist.ipynb
+++ b/mnist.ipynb
@@ -78,35 +78,11 @@
     "    print(\"{} > {} edges\".format(A.nnz//2, FLAGS.number_edges*m**2//2))\n",
     "    return A\n",
     "\n",
-    "def coarsen(A, levels):\n",
-    "    graphs, parents = coarsening.metis(A, levels)\n",
-    "    perms = coarsening.compute_perm(parents)\n",
-    "\n",
-    "    laplacians = []\n",
-    "    for i,A in enumerate(graphs):\n",
-    "        M, M = A.shape\n",
-    "\n",
-    "        # No self-connections.\n",
-    "        if True:\n",
-    "            A = A.tocoo()\n",
-    "            A.setdiag(0)\n",
-    "\n",
-    "        if i < levels:\n",
-    "            A = coarsening.perm_adjacency(A, perms[i])\n",
-    "\n",
-    "        A = A.tocsr()\n",
-    "        A.eliminate_zeros()\n",
-    "        Mnew, Mnew = A.shape\n",
-    "        print('Layer {0}: M_{0} = |V| = {1} nodes ({2} added), |E| = {3} edges'.format(i, Mnew, Mnew-M, A.nnz//2))\n",
-    "\n",
-    "        L = graph.laplacian(A, normalized=FLAGS.normalized_laplacian)\n",
-    "        laplacians.append(L)\n",
-    "    return laplacians, perms[0] if len(perms) > 0 else None\n",
-    "\n",
     "t_start = time.process_time()\n",
     "A = grid_graph(28, corners=False)\n",
     "A = graph.replace_random_edges(A, 0)\n",
-    "L, perm = coarsen(A, FLAGS.coarsening_levels)\n",
+    "graphs, perm = coarsening.coarsen(A, levels=FLAGS.coarsening_levels, self_connections=False)\n",
+    "L = [graph.laplacian(A, normalized=True) for A in graphs]\n",
     "print('Execution time: {:.2f}s'.format(time.process_time() - t_start))\n",
     "graph.plot_spectrum(L)\n",
     "del A"
diff --git a/rcv1.ipynb b/rcv1.ipynb
index 6747880..be25ead 100644
--- a/rcv1.ipynb
+++ b/rcv1.ipynb
@@ -227,37 +227,13 @@
    },
    "outputs": [],
    "source": [
-    "def coarsen(A, levels):\n",
-    "    graphs, parents = coarsening.metis(A, levels)\n",
-    "    perms = coarsening.compute_perm(parents)\n",
-    "\n",
-    "    laplacians = []\n",
-    "    for i,A in enumerate(graphs):\n",
-    "        M, M = A.shape\n",
-    "\n",
-    "        # No self-connections.\n",
-    "        if True:\n",
-    "            A = A.tocoo()\n",
-    "            A.setdiag(0)\n",
-    "\n",
-    "        if i < levels:\n",
-    "            A = coarsening.perm_adjacency(A, perms[i])\n",
-    "\n",
-    "        A = A.tocsr()\n",
-    "        A.eliminate_zeros()\n",
-    "        Mnew, Mnew = A.shape\n",
-    "        print('Layer {0}: M_{0} = |V| = {1} nodes ({2} added), |E| = {3} edges'.format(i, Mnew, Mnew-M, A.nnz//2))\n",
-    "\n",
-    "        L = graph.laplacian(A, normalized=FLAGS.normalized_laplacian)\n",
-    "        laplacians.append(L)\n",
-    "    return laplacians, perms[0] if len(perms) > 0 else None\n",
-    "\n",
     "t_start = time.process_time()\n",
     "dist, idx = graph.distance_lshforest(graph_data.astype(np.float64), k=FLAGS.number_edges, metric=FLAGS.metric)\n",
     "A = graph.adjacency(dist.astype(np.float32), idx)\n",
     "print(\"{} > {} edges\".format(A.nnz//2, FLAGS.number_edges*graph_data.shape[0]//2))\n",
     "A = graph.replace_random_edges(A, 0)\n",
-    "L, perm = coarsen(A, FLAGS.coarsening_levels)\n",
+    "graphs, perm = coarsening.coarsen(A, levels=FLAGS.coarsening_levels, self_connections=False)\n",
+    "L = [graph.laplacian(A, normalized=True) for A in graphs]\n",
     "print('Execution time: {:.2f}s'.format(time.process_time() - t_start))\n",
     "#graph.plot_spectrum(L)\n",
     "#del graph_data, A, dist, idx"