mst_clustering_3D.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a966ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Minimum Spanning Tree Clustering. This code is derived from the astroML package:\n",
    "\n",
    "Copyright (c) 2012-2013, Jacob Vanderplas All rights reserved.\n",
    "\n",
    "Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:\n",
    "\n",
    "=======\n",
    "Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.\n",
    "Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in \n",
    "the documentation and/or other materials provided with the distribution.\n",
    "\n",
    "THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, \n",
    "BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. \n",
    "IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR \n",
    "CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; \n",
    "OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT \n",
    "(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n",
    "\n",
    "\n",
    "=======\n",
    "\"\"\"\n",
    "import numpy as np\n",
    "\n",
    "from scipy import sparse\n",
    "from sklearn.neighbors import kneighbors_graph\n",
    "from sklearn.mixture import GMM\n",
    "\n",
    "try:\n",
    "    from scipy.sparse.csgraph import \\\n",
    "        minimum_spanning_tree, connected_components\n",
    "except:\n",
    "    raise ValueError(\"scipy v0.11 or greater required \"\n",
    "                     \"for minimum spanning tree\")\n",
    "\n",
    "\n",
    "class HierarchicalClustering(object):\n",
    "    \"\"\"Hierarchical Clustering via Approximate Euclidean Minimum Spanning Tree\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_neighbors : int\n",
    "        number of neighbors of each point used for approximate Euclidean\n",
    "        minimum spanning tree (MST) algorithm.  See Notes below.\n",
    "    edge_cutoff : float\n",
    "        specify a fraction of edges to keep when selecting clusters.\n",
    "        edge_cutoff should be between 0 and 1.\n",
    "    min_cluster_size : int, optional\n",
    "        specify a minimum number of points per cluster.  If not specified,\n",
    "        all clusters will be kept.\n",
    "    Attributes\n",
    "    ----------\n",
    "    X_train_ : ndarray\n",
    "        the training data\n",
    "    full_tree_ : sparse graph\n",
    "        the full approximate Euclidean MST spanning the data\n",
    "    cluster_graph_ : sparse graph\n",
    "        the final (truncated) graph showing clusters\n",
    "    n_components_ : int\n",
    "        the number of clusters found.\n",
    "    labels_ : int\n",
    "        the cluster labels for each training point.  Labels range from -1\n",
    "        to n_components_ - 1: points labeled -1 are in the background (i.e.\n",
    "        their clusters were smaller than min_cluster_size)\n",
    "    Notes\n",
    "    -----\n",
    "    This routine uses an approximate Euclidean minimum spanning tree (MST)\n",
    "    to perform hierarchical clustering.  A true Euclidean minimum spanning\n",
    "    tree naively costs O[N^3].  Graph traversal algorithms only help so much,\n",
    "    because all N^2 edges must be used as candidates.  In this approximate\n",
    "    algorithm, we use k < N edges from each point, so that the cost is only\n",
    "    O[Nk log(Nk)]. For k = N, the approximation is exact; in practice for\n",
    "    well-behaved data sets, the result is exact for k << N.\n",
    "\"\"\"\n",
    "    def __init__(self, n_neighbors=20,\n",
    "                 edge_cutoff=0.9,\n",
    "                 min_cluster_size=1):\n",
    "        self.n_neighbors = n_neighbors\n",
    "        self.edge_cutoff = edge_cutoff\n",
    "        self.min_cluster_size = min_cluster_size\n",
    "\n",
    "    def fit(self, X):\n",
    "        \"\"\"Fit the clustering model\n",
    "        Parameters\n",
    "        ----------\n",
    "        X : array_like\n",
    "            the data to be clustered: shape = [n_samples, n_features]\n",
    "        \"\"\"\n",
    "        X = np.asarray(X, dtype=float)\n",
    "\n",
    "        self.X_train_ = X\n",
    "\n",
    "        # generate a sparse graph using the k nearest neighbors of each point\n",
    "        G = kneighbors_graph(X, n_neighbors=self.n_neighbors, mode='distance')\n",
    "\n",
    "        # Compute the minimum spanning tree of this graph\n",
    "        self.full_tree_ = minimum_spanning_tree(G, overwrite=True)\n",
    "\n",
    "        # Find the cluster labels\n",
    "        self.n_components_, self.labels_, self.cluster_graph_ =\\\n",
    "            self.compute_clusters()\n",
    "\n",
    "        return self\n",
    "\n",
    "    def compute_clusters(self, edge_cutoff=None, min_cluster_size=None):\n",
    "        \"\"\"Compute the clusters given a trained tree\n",
    "        After fit() is called, this method may be called to obtain a\n",
    "        clustering result with a new edge_cutoff and min_cluster_size.\n",
    "        Parameters\n",
    "        ----------\n",
    "        edge_cutoff : float, optional\n",
    "            specify a fraction of edges to keep when selecting clusters.\n",
    "            edge_cutoff should be between 0 and 1.  If not specified,\n",
    "            self.edge_cutoff will be used.\n",
    "        min_cluster_size : int, optional\n",
    "            specify a minimum number of points per cluster.  If not specified,\n",
    "            self.min_cluster_size will be used.\n",
    "        Returns\n",
    "        -------\n",
    "        n_components : int\n",
    "            the number of clusters found\n",
    "        labels : ndarray\n",
    "            the labels of each point.  Labels range from -1 to\n",
    "            n_components_ - 1: points labeled -1 are in the background\n",
    "            (i.e. their clusters were smaller than min_cluster_size)\n",
    "        T_trunc : sparse matrix\n",
    "            the truncated minimum spanning tree\n",
    "        \"\"\"\n",
    "        if edge_cutoff is None:\n",
    "            edge_cutoff = self.edge_cutoff\n",
    "\n",
    "        if min_cluster_size is None:\n",
    "            min_cluster_size = self.min_cluster_size\n",
    "\n",
    "        if not hasattr(self, 'full_tree_'):\n",
    "            raise ValueError(\"must call fit() before calling \"\n",
    "                             \"compute_clusters()\")\n",
    "\n",
    "        T_trunc = self.full_tree_.copy()\n",
    "\n",
    "        # cut-off edges at the percentile given by edge_cutoff\n",
    "        cutoff = np.percentile(T_trunc.data, 100 * edge_cutoff)\n",
    "        T_trunc.data[T_trunc.data > cutoff] = 0\n",
    "        T_trunc.eliminate_zeros()\n",
    "\n",
    "        # find connected components\n",
    "        n_components, labels = connected_components(T_trunc, directed=False)\n",
    "        counts = np.bincount(labels)\n",
    "\n",
    "        # for all components with less than min_cluster_size points, set\n",
    "        # to background, and re-label the clusters\n",
    "        i_bg = np.where(counts < min_cluster_size)[0]\n",
    "\n",
    "        for i in i_bg:\n",
    "            labels[labels == i] = -1\n",
    "\n",
    "        if len(i_bg) > 0:\n",
    "            _, labels = np.unique(labels, return_inverse=True)\n",
    "            labels -= 1\n",
    "            n_components = labels.max() + 1\n",
    "\n",
    "        # eliminate links in T_trunc which are not clusters\n",
    "        I = sparse.eye(len(labels), len(labels))\n",
    "        I.data[0, labels < 0] = 0\n",
    "        T_trunc = I * T_trunc * I\n",
    "\n",
    "        return n_components, labels, T_trunc\n",
    "\n",
    "\n",
    "def get_graph_segments(X, G):\n",
    "    \"\"\"Get graph segments for plotting a 2D graph\n",
    "    Parameters\n",
    "    ----------\n",
    "    X : array_like\n",
    "        the data, of shape [n_samples, 2]\n",
    "    G : array_like or sparse graph\n",
    "        the [n_samples, n_samples] matrix encoding the graph of connectinons\n",
    "        on X\n",
    "    Returns\n",
    "    -------\n",
    "    x_coords, y_coords : ndarrays\n",
    "        the x and y coordinates for plotting the graph.  They are of size\n",
    "        [2, n_links], and can be visualized using\n",
    "        ``plt.plot(x_coords, y_coords, '-k')``\n",
    "    \"\"\"\n",
    "    X = np.asarray(X)\n",
    "    # if (X.ndim != 3) or (X.shape[1] != 3):\n",
    "    #     raise ValueError('shape of X should be (n_samples, 2)')\n",
    "\n",
    "    n_samples = X.shape[0]\n",
    "\n",
    "    G = sparse.coo_matrix(G)\n",
    "    A = X[G.row].T\n",
    "    B = X[G.col].T\n",
    "\n",
    "    x_coords = np.vstack([A[0], B[0]])\n",
    "    y_coords = np.vstack([A[1], B[1]])\n",
    "    z_coords = np.vstack([A[2], B[2]])\n",
    "\n",
    "    return x_coords, y_coords, z_coords\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}