00_MergeTreesAndDecoratedMergeTrees.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Merge Trees and Decorated Merge Trees\n",
    "\n",
    "This notebook illustrates some of the basic functionality of our merge tree and decorated merge tree (DMT) analysis framework. The main paper describing this framework is \n",
    "\n",
    "[1] *Decorated Merge Trees for Persistenct Topology* by Justin Curry, Haibin Hang, Washington Mio, Tom Needham and Osman Okutan. [Arxiv](https://arxiv.org/abs/2103.15804). The paper has been published in [Journal of Applied and Computational Topology](https://link.springer.com/article/10.1007/s41468-022-00089-3).\n",
    "\n",
    "We will refer back to specific sections of this paper throughout the notebook. Other notebooks in the GitHub repo show more involved examples and can be used to recreate the figures and experiments in the paper.\n",
    "\n",
    "This notebook will be focused on creating and visualizing merge trees and DMTs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ripser import ripser\n",
    "import persim\n",
    "import networkx as nx\n",
    "import random\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from scipy.cluster.hierarchy import dendrogram, linkage, cut_tree\n",
    "from sklearn.metrics.pairwise import pairwise_distances\n",
    "\n",
    "from DMT_tools import *\n",
    "\n",
    "import ot\n",
    "\n",
    "import time\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Merge Trees\n",
    "\n",
    "Our code includes functions for creating, visualizing and comparing merge trees. The creation and visualization  functionality is demonstrated below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating Merge Trees\n",
    "\n",
    "A *merge tree* is a rooted tree with heights assigned to its nodes. Every non-root node $v$ must have exactly one neighbor whose height is strictly greater than the height of $v$. Merge trees are represented as pairs `T`, `height` containing the graph structure---as a `networkx` graph---and the height function---as a dictionary with keys equal to the node labels of the tree. This is data is stored together in a `MergeTree` class.\n",
    "\n",
    "The code below creates merge trees by directly specifying their combinatorial structure. Later we will see methods for generating merge trees automatically from point clouds or filtered networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initializing data\n",
    "T1=nx.Graph()\n",
    "T1.add_nodes_from(list(range(6)))\n",
    "T1.add_edges_from([(0,4),(1,3),(2,3),(3,4),(4,5)])\n",
    "height1 = {0:0,1:0,2:2,3:3,4:5,5:7}\n",
    "\n",
    "T2 = nx.Graph()\n",
    "T2.add_nodes_from(list(range(4)))\n",
    "T2.add_edges_from([(0,2),(1,2),(2,3)])\n",
    "height2 = {0:0,1:0,2:5,3:7}\n",
    "\n",
    "# Creating merge tree classes\n",
    "MT1 = MergeTree(tree = T1, height = height1)\n",
    "MT2 = MergeTree(tree = T2, height = height2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The tree structure or node heights can be recovered from the class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(MT1.height)\n",
    "\n",
    "nx.draw_kamada_kawai(MT1.tree)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing Merge Trees\n",
    "\n",
    "We can plot the merge trees we created above according to their height data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1.draw(axes = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT2.draw(axes = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimating Merge Tree Interleaving Distance\n",
    "\n",
    "One can define a metric on the space of merge trees called *interleaving distance* which is inspired by bottleneck distance between persistence diagrams (see *Interleaving distance between merge trees* by Morozov et al., 2013 or Section 4.3 of [1]). Interleaving distance is not tractable to compute exactly, but we introduce a method to estimate it by using Gromov-Wasserstein distance from optimal transport theory (see Section 6.5 of [1]). \n",
    "\n",
    "To improve the accuracty, the user can subdivide the merge trees by specifying a mesh. Generally, a smaller mesh yields a more accurate distance, but a finer mesh incurs a higher computational cost. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = 0.5\n",
    "res = merge_tree_interleaving_distance(MT1,MT2,mesh,verbose = False)\n",
    "print('Estimated Interleaving Distance:',res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distance computation can return several outputs, which can be used to visualize the way that the distance computation works. To get more informative output, set `verbose = True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = 0.5\n",
    "res = merge_tree_interleaving_distance(MT1,MT2,mesh,verbose = True)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The 'coupling' key in the dictionary `res` returns a probabilistic matching between the nodes of the subdivided trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (5,5))\n",
    "plt.imshow(res['coupling'])\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distance computation proceeds by estimating an 'optimal labeling' of the merge trees from the coupling. The 'labeling' approach to merge tree interleaving was introduced in Gasparovic et al., *Intrinsic interleaving distance for merge trees*, 2019. \n",
    "\n",
    "The labeling can be visualized as shown below. To display it, we should tell the interleaving distance to return the subdivided trees as well as a verbose results output.\n",
    "\n",
    "Nodes with the same label across the two merge trees in the figure below should be thought of as 'structurally matched'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = 1\n",
    "MT1_sub, MT2_sub, res = merge_tree_interleaving_distance(MT1,MT2,mesh,verbose = True,return_subdivided = True)\n",
    "\n",
    "MT1_sub.draw_with_labels(res['label1'])\n",
    "plt.show()\n",
    "MT2_sub.draw_with_labels(res['label2'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res['dist']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a Merge Tree from Data\n",
    "\n",
    "Generally, we would be generating merge trees from a dataset. Let's start with a toy dataset consisting of a point cloud in 2D. We'll write some functions to generate random point clouds with topological structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def noisy_circle(n_samples, noise_level, center_x, center_y ,radius):\n",
    "\n",
    "    t = np.linspace(0,2*np.pi,n_samples)\n",
    "    x = center_x + radius*np.cos(t)\n",
    "    y = center_y + radius*np.sin(t)\n",
    "    noise = np.random.rand(n_samples,2)\n",
    "    data = np.array([x,y]).T + noise_level*radius*noise\n",
    "\n",
    "    return data\n",
    "\n",
    "def noisy_disk(n_samples,noise_level,center_x,center_y,radius):\n",
    "\n",
    "    t = np.linspace(0,2*np.pi,n_samples)\n",
    "    r = radius*np.random.rand(n_samples)\n",
    "    x = center_x + np.multiply(r,np.cos(t))\n",
    "    y = center_y + np.multiply(r,np.sin(t))\n",
    "    noise = np.random.rand(n_samples,2)\n",
    "    data = np.array([x,y]).T + noise_level*radius*noise\n",
    "\n",
    "    return data\n",
    "\n",
    "def one_disk_two_circles(radii, separation, n_samples_per_shape = 50, noise_level = 0.5, centersy = None):\n",
    "    \n",
    "    \"\"\"\n",
    "    This example produces one disk and two circles or radii r1,r2,r3. The circles are clustered together and \n",
    "    the disk is separated in the x-direction by `separation`\n",
    "    In: triple of radii = [r1,r2,r3], separation distance\n",
    "    \"\"\"\n",
    "    \n",
    "    data = np.zeros([n_samples_per_shape*3,2])\n",
    "    \n",
    "    shapes = ['disk','circle','circle']\n",
    "    r1 = radii[0]\n",
    "    r2 = radii[1]\n",
    "    r3 = radii[2]\n",
    "    centersx = [0,r1+separation+r2,r1+separation+2*r2+r3]\n",
    "    \n",
    "    if centersy is None:\n",
    "        centersy = [0,0,0]\n",
    "    \n",
    "    num_shapes = len(shapes)\n",
    "    \n",
    "    for j in range(num_shapes):\n",
    "        if shapes[j] == 'disk':\n",
    "            shape = noisy_disk(n_samples_per_shape,noise_level,centersx[j],centersy[j],radii[j])\n",
    "        elif shapes[j] == 'circle':\n",
    "            shape = noisy_circle(n_samples_per_shape, noise_level,centersx[j],centersy[j],radii[j])\n",
    "        data[j*n_samples_per_shape:(j+1)*n_samples_per_shape,:] = shape\n",
    "        \n",
    "    return data\n",
    "\n",
    "def three_circles(radii, separation1, separation2, n_samples_per_shape = 50, noise_level = 0.5, centersy = None):\n",
    "    \n",
    "    \n",
    "    data = np.zeros([n_samples_per_shape*3,2])\n",
    "    \n",
    "    shapes = ['circle','circle','circle']\n",
    "    r1 = radii[0]\n",
    "    r2 = radii[1]\n",
    "    r3 = radii[2]\n",
    "    centersx = [0,r1+separation1+r2,r1+separation1+r2]\n",
    "    \n",
    "    if centersy is None:\n",
    "        centersy = [0,0,r2+separation2+r3]\n",
    "    \n",
    "    num_shapes = len(shapes)\n",
    "    \n",
    "    for j in range(num_shapes):\n",
    "        if shapes[j] == 'disk':\n",
    "            shape = noisy_disk(n_samples_per_shape,noise_level,centersx[j],centersy[j],radii[j])\n",
    "        elif shapes[j] == 'circle':\n",
    "            shape = noisy_circle(n_samples_per_shape, noise_level,centersx[j],centersy[j],radii[j])\n",
    "        data[j*n_samples_per_shape:(j+1)*n_samples_per_shape,:] = shape\n",
    "        \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data1 = one_disk_two_circles([1,1,1],3,centersy = [0,0.5,-0.5])\n",
    "\n",
    "plt.scatter(data1[:,0],data1[:,1])\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "\n",
    "data2 = three_circles([1,1,1],3,3)\n",
    "\n",
    "plt.scatter(data2[:,0],data2[:,1])\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create a merge tree by tracking merging of connected components in the Vietoris-Rips complexes as the parameter increases. The result is also known as a **single linkage hierarchical clustering dendrogram**. Passing a point cloud matrix into the merge tree class produces the merge tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1 = MergeTree(pointCloud = data1)\n",
    "MT2 = MergeTree(pointCloud = data2)\n",
    "\n",
    "MT1.draw(axes = True)\n",
    "plt.show()\n",
    "MT2.draw(axes = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simplifying Merge Trees\n",
    "\n",
    "The merge trees shown above are quite complicated. We can get simplified versions by slicing at a given height and merging all nodes below the threshold height to a single node at height zero (more generally, at the lowest leaf height below the threshold point in each branch). This is akin to throwing away bars in a barcode with sufficiently low persistence. The resulting visualizations are much easier to parse and the simplification improves computation time for computing interleaving distance.\n",
    "\n",
    "The code below performes this simplification at a user-defined threshold height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresh = 0.4\n",
    "\n",
    "MT1.threshold(thresh)\n",
    "MT1.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple Distance Comparison\n",
    "\n",
    "Let's check qualitatively that our estimation of merge tree interleaving distance is able to distinguish between the classes that we've defined above.\n",
    "\n",
    "The hyperparameters in the experiment are `threshold`, which controls the height at which the trees are truncated, and `mesh`, which controls how densely the trees are upsampled. Both have an effect on compute time. We've observed experimentally that simplifying by thresholding doesn't tend to degrade results, and sometimes improves accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Hyperparameters\n",
    "threshold = 1\n",
    "mesh = 1\n",
    "\n",
    "\n",
    "# Create dataset of 2 samples for each of the two classes\n",
    "print('Creating Data...')\n",
    "point_clouds = [one_disk_two_circles([1,1,1],3,centersy = [0,0.5,-0.5]),\n",
    "               one_disk_two_circles([1,1,1.1],3,centersy = [0,0.5,-0.5]),\n",
    "                three_circles([1,1,1],3,3),\n",
    "                three_circles([1,1,1.1],3,3)]\n",
    "\n",
    "MTs = []\n",
    "\n",
    "for point_cloud in point_clouds:\n",
    "    MT = MergeTree(pointCloud = point_cloud)\n",
    "    MT.threshold(threshold)\n",
    "    MTs.append(MT)\n",
    "    \n",
    "# Compute a 4x4 pairwise distance matrix\n",
    "\n",
    "print('Computing Pairwise Distances...')\n",
    "distMat = np.zeros([4,4])\n",
    "\n",
    "for i in range(4):\n",
    "    print('row=',i)\n",
    "    MTi = MTs[i]\n",
    "    for j in range(i+1,4):\n",
    "        MTj = MTs[j]\n",
    "        distMat[i,j] = merge_tree_interleaving_distance(MTi,MTj,mesh,verbose = False)\n",
    "        \n",
    "distMat += distMat.T\n",
    "\n",
    "plt.imshow(distMat)\n",
    "plt.title('Distance Matrix',fontsize = 16)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decorated Merge Trees (DMTs)\n",
    "\n",
    "Decorated Merge Trees fuse persistent connectivity data (the merge tree) with higher dimensional persistent homological data (barcodes)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating DMTs from Data\n",
    "\n",
    "The code below creates a DMT from the first dataset created above. Specifically, it creates a *Barcode Decorated Merge Tree* from Section 3.3 of [1]. This attaches a barcode to each leaf of the merge tree.\n",
    "\n",
    "In this example, we are attaching a degree-1 persistent homology barcode to each leaf of the merge tree. The persistent homology computation is handled by `ripser`.\n",
    "\n",
    "The process of decorating a merge tree is handled within the class by the `fit_barcode` method. The degree of the barcode should be specified and the persistent homology computation is handled within the fitting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(data1[:,0],data1[:,1])\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1 = MergeTree(pointCloud = data1)\n",
    "MT1.fit_barcode(degree = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The barcodes attached to the leaves of the merge tree are accessed by the `leaf_barcode` attribute. The leaf barcode is a dictionary of the form {node_id:barcode}, where the barcode is a list of bars (each bar is a list with two elements)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1.leaf_barcode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simplifying DMTs\n",
    "\n",
    "When we simplify a merge tree, as describe above, we can also simplify the leaf decorations. The leaves in the simplified merge tree keep their associated barcode computed above. This is once again handled by the `threshold` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 0.5\n",
    "\n",
    "MT1_thresh = MT1.copy()\n",
    "MT1_thresh.threshold(threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1_thresh.draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1_thresh.leaf_barcode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two apparent loops in the dataset should be represented by two long bars attached to a leaf (or to several leaves which merge at a relatively low height). This corresponds to the two loops lying in a single connected component."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing DMTs\n",
    "\n",
    "We can visualize our DMT by enriching the merge tree with overlaid barcodes. This corresponds to the 'lift decoration' process described in Section 5.4 of [1]. \n",
    "\n",
    "To get a nice looking picture, we only include bars whose persistence is larger than a given threshold. Now there are two parameters involved in visualizing a DMT: the height at which to truncate the tree and the persistence of bars that we want to display overlaid on the tree.\n",
    "\n",
    "We can predetermine a good threshold for persistence of bars by visualizing the persistence diagram. We use the `persim` package and plot in 'lifetime' mode, so that the persistence of a bar is more apparent. The barcode for the merge tree is stored as an attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "persim.plot_diagrams(MT1.barcode,lifetime = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `draw_decorated` method takes thresholds as inputs and produces the following figure. Note that the rendering process can take a while if the merge tree is very complicated. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tree_thresh = 0.2\n",
    "barcode_thresh = 0.4\n",
    "\n",
    "MT1.draw_decorated(tree_thresh,barcode_thresh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observe that the DMT simultaneously illustrates the connected component and loops structure of the dataset, together with how these structures interact!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a code block to draw a DMT for the second dataset defined above, using the same thresholds as for `data1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT2 = MergeTree(pointCloud = data2)\n",
    "MT2.fit_barcode(degree = 1)\n",
    "MT2.draw_decorated(tree_thresh,barcode_thresh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing DMTs\n",
    "\n",
    "Decorated merge trees can be compared using the decorated bottleneck distance defined in Section 4.5 of [1]. We can estimate this distance using a variant of Gromov-Wasserstein distance called [Fused Gromov-Wasserstein Distance](https://arxiv.org/abs/1811.02834). This involves another parameter $\\alpha$, which balances the weight of degree-0 vs. higher degree homology when finding a merge tree interleaving. \n",
    "\n",
    "Details and more examples will be provided in a followup notebook. The basic syntax for computing the distance is shown below. It is computationally beneficial to simplify the merge tree (as above) before computing the distance. This distance computation is currently usable, but is **not yet optimized**---faster versions of this algorithm will be included in future updates to the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT1_thresh = MT1.copy()\n",
    "MT1_thresh.threshold(0.4)\n",
    "MT2_thresh = MT2.copy()\n",
    "MT2_thresh.threshold(0.4)\n",
    "\n",
    "res = DMT_interleaving_distance(MT1_thresh,MT2_thresh,1,verbose = True)\n",
    "\n",
    "print('The estimated decorated bottleneck distance is:', res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}