Merge 3dd9e69 into 9690d37

localgraphclustering/GraphLocal.py

@@ -8,13 +8,22 @@
 import warnings
 import collections as cole
 from .cpp import *
+import random
 
 import gzip
 import bz2
 import lzma
 
 import multiprocessing as mp
 
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.collections import LineCollection
+from mpl_toolkits.mplot3d.art3d import Line3DCollection
+from matplotlib.colors import to_rgb,to_rgba
+from matplotlib.colors import LinearSegmentedColormap
+from matplotlib.colors import Normalize
+
 def _load_from_shared(sabuf, dtype, shape):
     return np.frombuffer(sabuf, dtype=dtype).reshape(shape)
 
@@ -558,8 +567,6 @@ def local_extrema(self,vals,strict=False,reverse=False):
         Parameters
         ----------
 
-        G: GraphLocal
-
         vals: Sequence[float]
             a feature value per node used to find the ex against each other, i.e. conductance
 
@@ -616,3 +623,179 @@ def local_extrema(self,vals,strict=False,reverse=False):
         minvals = vals[minverts]
 
         return minverts, minvals
+
+    def draw(self,coords,alpha=1.0,nodesize=5,linewidth=1,
+         nodealpha=1.0,edgealpha=0.01,nodecolor='r',
+         edgecolor='k',nodemarker='o',setalpha=1.0,
+         setcolor='y',axs=None,fig=None,nodeset=None,
+         groups=None,values=None,cm="Reds"):
+        """
+        standard drawing function of GraphLocal object
+
+        Parameters
+        ----------
+
+        coords: a n-by-2 or n-by-3 array with coordinates for each node of the graph.
+
+        Optional parameters
+        ------------------
+
+        alpha: float (1.0 by default)
+            the overall alpha scaling of the plot, [0,1]
+
+        nodealpha: float (1.0 by default)
+            the overall node alpha scaling of the plot, [0, 1]
+
+        edgealpha: float (1.0 by default)
+            the overall edge alpha scaling of the plot, [0, 1]
+
+        setalpha: float (1.0 by default)
+            the overall set alpha scaling of the plot, [0, 1]
+
+        nodecolor: string or RGB ('r' by default)
+
+        edgecolor: string or RGB ('k' by default)
+
+        setcolor: string or RGB ('y' by default)
+
+        nodemarker: string ('o' by default)
+
+        nodesize: float (5.0 by default)
+
+        linewidth: float (1.0 by default)
+
+        nodeset: Sequence[int] (None by default)
+            a set of nodes to highlight
+
+        groups: Sequence[Sequence[int]] (None by default)
+            node partitions, different colors will be assigned to different groups
+
+        axs,fig: None,None (default) 
+            by default it will create a new figure, or this will plot in axs if not None.
+
+        values: Sequence[float] (None by default)
+            used to determine node colors in a colormap, should have the same length as coords
+
+        cm: string ("Reds" by default)
+            colormap
+
+        Returns
+        -------
+
+        a dictionary with: 
+        fig, ax, setnodes, groupnodes
+        """
+        if axs is None:
+            fig = plt.figure()
+            if len(coords[0]) == 3:
+                axs = fig.add_subplot(111, projection='3d')
+            else:
+                axs = fig.add_subplot(111)
+        axs.set_axis_off()
+        nodeset = set(nodeset) if nodeset is not None else set()
+        nodelist_in = []
+        nodelist_out = []
+        for i in range(self._num_vertices):
+            if i in nodeset:
+                nodelist_in.append(i)
+            else:
+                nodelist_out.append(i)
+
+        if values is not None:
+            #when values are provided, use values and predefined colormap to determine colors
+            cm = plt.get_cmap(cm)
+            node_color_list = np.reshape(np.array(values),len(coords))
+            vmin = min(values)
+            vmax = max(values)
+        else:
+            #when values are not provided, use customed colormap to determine colors
+            colors = [to_rgba(nodecolor,alpha*nodealpha),to_rgba(setcolor,alpha*setalpha)]
+            if groups is not None:
+                colorset = set([to_rgb(c) for c in colors])
+                number_of_colors = len(groups)
+                for i in range(number_of_colors):
+                    new_color = to_rgb("#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)]))
+                    #make sure color hasn't already existed
+                    while new_color in colorset:
+                        new_color = to_rgb("#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)]))
+                    colorset.add(new_color)
+                    colors.append(to_rgba(new_color,alpha*nodealpha))
+            cm = LinearSegmentedColormap.from_list('my_cmap',colors,N=256,gamma=1.0)
+            node_color_list = np.zeros(self._num_vertices)
+            node_color_list[nodelist_in] = 1.0/(len(colors)-1)
+            if groups is not None:
+                for i,g in enumerate(groups):
+                    node_color_list[g] = (2+i)*1.0/(len(colors)-1)
+            vmin = 0.0
+            vmax = 1.0
+
+        if len(coords[0]) == 3:
+            self.draw_3d(coords,axs,cm,nodemarker=nodemarker,nodesize=nodesize,edgealpha=alpha*edgealpha,
+                    linewidth=linewidth,node_color_list=node_color_list,vmin=vmin,vmax=vmax,nodelist_in=nodelist_in,
+                    nodelist_out=nodelist_out,setalpha=alpha*setalpha,nodealpha=alpha*nodealpha,use_values=(values is not None))
+        else:
+            self.draw_2d(coords,axs,cm,nodemarker=nodemarker,nodesize=nodesize,edgealpha=alpha*edgealpha,
+                    linewidth=linewidth,node_color_list=node_color_list,vmin=vmin,vmax=vmax,nodelist_in=nodelist_in,
+                    nodelist_out=nodelist_out,setalpha=alpha*setalpha,nodealpha=alpha*nodealpha,use_values=(values is not None))
+
+        ret_dict = {"fig":fig,"ax":axs,"setnodes":nodelist_in,"groupnodes":groups}
+        return ret_dict
+
+    @staticmethod
+    def draw_star(center,points,pos,edge_pos):
+        for i,p in enumerate(points):
+            if p >= center:
+                edge_pos.append([pos[center],pos[p]])
+
+    def draw_2d(self,pos,axs,cm,nodemarker='o',nodesize=5,edgealpha=0.01,linewidth=1,
+                node_color_list=None,edgecolor='k',nodecolor='r',node_list=None,nodelist_in=None,
+                nodelist_out=None,setalpha=1.0,nodealpha=1.0,use_values=False,vmin=0.0,vmax=1.0):
+        if use_values:
+            axs.scatter([p[0] for p in pos[nodelist_in]],[p[1] for p in pos[nodelist_in]],c=node_color_list[nodelist_in],
+                s=nodesize,marker=nodemarker,cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),alpha=setalpha,zorder=2)
+            axs.scatter([p[0] for p in pos[nodelist_out]],[p[1] for p in pos[nodelist_out]],c=node_color_list[nodelist_out],
+                s=nodesize,marker=nodemarker,cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),alpha=nodealpha,zorder=2)
+        else:
+            if node_color_list is not None:
+                axs.scatter([p[0] for p in pos],[p[1] for p in pos],c=node_color_list,s=nodesize,marker=nodemarker,
+                    cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),zorder=2)
+            else:
+                axs.scatter([p[0] for p in pos],[p[1] for p in pos],c=nodecolor,s=nodesize,marker=nodemarker,zorder=2)
+        node_list = range(self._num_vertices) if node_list is None else node_list
+        edge_pos = []
+        for i in node_list:
+            self.draw_star(i,self.aj[self.ai[i]:self.ai[i+1]],pos,edge_pos)
+        edge_pos = np.asarray(edge_pos)
+        edge_collection = LineCollection(edge_pos,colors=to_rgba(edgecolor,edgealpha),linewidths=linewidth)
+        #make sure edges are at the bottom
+        edge_collection.set_zorder(1)
+        axs.add_collection(edge_collection)
+        axs.autoscale()
+
+    def draw_3d(self,pos,axs,cm,nodemarker='o',nodesize=5,edgealpha=0.01,linewidth=1,
+                node_color_list=None,angle=30,edgecolor='k',nodecolor='r',node_list=None,
+                nodelist_in=None,nodelist_out=None,setalpha=1.0,nodealpha=1.0,use_values=False,vmin=0.0,vmax=1.0):
+        if use_values:
+            axs.scatter([p[0] for p in pos[nodelist_in]],[p[1] for p in pos[nodelist_in]],[p[2] for p in pos[nodelist_in]],c=node_color_list[nodelist_in],
+                s=nodesize,marker=nodemarker,cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),zorder=2,alpha=setalpha)
+            axs.scatter([p[0] for p in pos[nodelist_out]],[p[1] for p in pos[nodelist_out]],[p[2] for p in pos[nodelist_out]],c=node_color_list[nodelist_out],
+                s=nodesize,marker=nodemarker,cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),zorder=2,alpha=nodealpha)
+        else:
+            if node_color_list is not None:
+                axs.scatter([p[0] for p in pos],[p[1] for p in pos],[p[2] for p in pos],c=node_color_list,
+                    s=nodesize,marker=nodemarker,cmap=cm,norm=Normalize(vmin=vmin,vmax=vmax),zorder=2)
+            else:
+                axs.scatter([p[0] for p in pos],[p[1] for p in pos],[p[2] for p in pos],c=nodecolor,
+                    s=nodesize,marker=nodemarker,zorder=2)
+        node_list = range(self._num_vertices) if node_list is None else node_list
+        edge_pos = []
+        for i in node_list:
+            self.draw_star(i,self.aj[self.ai[i]:self.ai[i+1]],pos,edge_pos)
+        edge_pos = np.asarray(edge_pos)
+        edge_collection = Line3DCollection(edge_pos,colors=to_rgba(edgecolor,edgealpha),linewidths=linewidth)
+        #make sure edges are at the bottom
+        edge_collection.set_zorder(1)
+        axs.add_collection(edge_collection)
+        axs.autoscale()
+        # Set the initial view
+        axs.view_init(30, angle)

localgraphclustering/tests/test_algs.py

@@ -1,10 +1,25 @@
 from localgraphclustering import *
 import time
 import numpy as np
+import networkx as nx
+import random
 
 def load_example_graph(vtype,itype):
     return GraphLocal("localgraphclustering/tests/data/dolphins.edges",separator=" ",vtype=vtype,itype=itype)
 
+def generate_random_3Dgraph(n_nodes, radius, seed=None):
+
+    if seed is not None:
+        random.seed(seed)
+
+    # Generate a dict of positions
+    pos = {i: (random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)) for i in range(n_nodes)}
+
+    # Create random 3D network
+    G = nx.random_geometric_graph(n_nodes, radius, pos=pos)
+
+    return G
+
 def test_GraphLocal_methods():
     g = load_example_graph(np.uint32,np.uint32)
     g.largest_component()
@@ -33,6 +48,34 @@ def test_GraphLocal_methods():
     # Test graph with more than one components
     G = GraphLocal("notebooks/datasets/neuro-fmri-01.edges",file_type = "edgelist", separator = " ", header = True)
 
+    # Test drawing fuinctions
+    g = GraphLocal('notebooks/datasets/JohnsHopkins.graphml','graphml','\t')
+    ld_coord = np.loadtxt('notebooks/datasets/JohnHopkins_coord.xy', dtype = 'Float64')
+    idx = np.argsort(ld_coord[:,0])
+    coords = []
+    for i in range(g._num_vertices):
+        coords.append(ld_coord[idx[i],1:3])
+    coords = np.array(coords)
+    # Call the global spectral partitioning algorithm.
+    eig2 = fiedler(g)[0]
+
+    # Round the eigenvector
+    output_sc = sweep_cut(g,eig2)
+
+    # Extract the partition for g and store it.
+    eig2_rounded = output_sc[0]
+    ret_dict = g.draw(coords,edgealpha=0.01,nodealpha=0.5,nodeset=eig2_rounded,values=eig2)
+    ret_dict = g.draw(coords,edgealpha=0.01,nodealpha=0.5,nodeset=eig2_rounded)
+
+    N = generate_random_3Dgraph(n_nodes=200, radius=0.25, seed=1)
+    pos = np.array(list(nx.get_node_attributes(N,'pos').values()))
+    G = GraphLocal()
+    G = G.from_networkx(N)
+    ret_dict = G.draw(pos,edgealpha=0.01,nodealpha=0.5,nodeset=range(100,150),groups=[range(50),range(50,100)],
+                      values=[random.uniform(0, 1) for i in range(200)])
+    ret_dict = G.draw(pos,edgealpha=0.01,nodealpha=0.5,nodeset=range(100,150),groups=[range(50),range(50,100)])
+
+
 def test_sweepcut_self_loop():
     """ This is a regression test for sweep-cuts with self-loops """
     g = GraphLocal()

notebooks/examples_with_visualization.ipynb

@@ -1410,7 +1410,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.4"
+   "version": "3.6.5"
   }
  },
  "nbformat": 4,