First Draft: A Script to Display Multicolor Voronoi Diagrams in Matplotlib

Code Sample

'''
This module generates 2D voronoi diagrams from a list of points.
The vornoi cells can be colored by supplying a value associated with 
   each node.
'''

__author__ = 'Ed Tate'
__email__  = 'edtate<at>gmail-dot-com'
__website__ = 'exnumerus.blogspot.com'
__license__ = 'Creative Commons Attribute By - http://creativecommons.org/licenses/by/3.0/us/'''

from matplotlib.pyplot import cm

def voronoi2D(xpt,ypt,cpt=None,cmap=None):
    '''
    This function returns a list of line segments which describe the voronoi
        cells formed by the points in zip(xpt,ypt).
    
    If cpt is provided, it identifies which cells should be returned.
        The boundary of the cell about (xpt[i],ypt[i]) is returned 
            if cpt[i]<=threshold.
            
    This function requires qvoronoi.exe in the working directory. 
    The working directory must have permissions for read and write access.
    This function will leave 2 files in the working directory:
        data.txt
        results.txt
    This function will overwrite these files if they already exist.
    '''
    
    if cpt is None:
        # assign a value to cpt for later use
        cpt = [0 for x in xpt]
        
    if cmap is None:
        cmap = cm.gray
    
    # write the data file
    pts_filename = 'data.txt'
    pts_F = open(pts_filename,'w')
    pts_F.write('2 # this is a 2-D input set\n')
    pts_F.write('%i # number of points\n' % len(xpt))
    for i,(x,y) in enumerate(zip(xpt,ypt)):
        pts_F.write('%f %f # data point %i\n' % (x,y,i))
    pts_F.close()

    # trigger the shell command
    import subprocess
    p = subprocess.Popen('qvoronoi TI data.txt TO results.txt p FN Fv QJ', shell=True)
    p.wait()

    # open the results file and parse results
    results = open('results.txt','r')

    # get 'p' results - the vertices of the voronoi diagram
    data = results.readline()
    voronoi_x_list = []
    voronoi_y_list = []
    data = results.readline()
    for i in range(0,int(data)):
        data = results.readline()
        xx,yy,dummy = data.split(' ')    
        voronoi_x_list.append(float(xx))
        voronoi_y_list.append(float(yy))
        
    # get 'FN' results - the voronoi edges
    data = results.readline()
    voronoi_idx_list = []
    for i in range(0,int(data)):
        data = results.readline()
        this_list = data.split(' ')[:-1]
        for j in range(len(this_list)):
            this_list[j]=int(this_list[j])-1
        voronoi_idx_list.append(this_list[1:])
        
    # get 'FV' results - pairs of points which define a voronoi edge
    # combine these results to build a complete representation of the 
    data = results.readline()
    voronoi_dict = {}
    for i in range(0,int(data)):
        data = results.readline().split(' ')

        pair_idx_1 = int(data[1])
        pair_idx_2 = int(data[2])

        vertex_idx_1 = int(data[3])-1
        vertex_idx_2 = int(data[4])-1

        try:
            voronoi_dict[pair_idx_1].append({ 'edge_vertices':[vertex_idx_1,vertex_idx_2],
                                      'neighbor': pair_idx_2 })
        except KeyError:
            voronoi_dict[pair_idx_1] = [{ 'edge_vertices':[vertex_idx_1,vertex_idx_2],
                                      'neighbor': pair_idx_2 } ]

        try:
            voronoi_dict[pair_idx_2].append({ 'edge_vertices':[vertex_idx_1,vertex_idx_2],
                                      'neighbor': pair_idx_1 })
        except KeyError:
            voronoi_dict[pair_idx_2] = [{ 'edge_vertices':[vertex_idx_1,vertex_idx_2],
                                      'neighbor': pair_idx_1 } ]    

                    
    #################
    # generate a collection of voronoi cells
    result_list = []
    for point_idx in voronoi_dict.keys():
        # determine cell color
        this_color = cmap(cpt[point_idx])
        
        # display this cell, so add the data to the edge list
        e_list = []
        for edge in voronoi_dict[point_idx]:
            p1_idx = edge['edge_vertices'][0]
            p2_idx = edge['edge_vertices'][1]
            e_list.append((p1_idx,p2_idx))
        
        # put the vertices points in order so they
        #   walk around the voronoi cells
        p_list = [p1_idx]
        while True:
            p=p_list[-1]
            for e in e_list:
                if p==e[0]:
                    next_p = e[1]
                    break
                elif p==e[1]:
                    next_p = e[0]
                    break
            p_list.append(next_p)
            e_list.remove(e)
            if p_list[0]==p_list[-1]:
                # the cell is closed
                break
            
        # build point list
        this_x_list = []
        this_y_list = []
        if all([p>=0 for p in p_list]):
            for p in p_list:
                if p>=0:
                    this_x_list.append(voronoi_x_list[p])
                    this_y_list.append(voronoi_y_list[p])     
                    
        result_list.append([this_x_list,this_y_list,this_color])
        
    return result_list

if __name__=='__main__':
    
    import random
    
    xpt = [random.random()-0.5 for i in range(0,100)]
    ypt = [random.random()-0.5 for i in range(0,100)]
    cpt = [int(random.random()*255) for i in range(0,100)]

    import matplotlib.pyplot as pp
    
    result_list = voronoi2D(xpt,ypt,cpt,cmap=pp.cm.copper)
    
    pp.figure()
    for item in result_list:
        x_list = item[0]
        y_list = item[1]
        this_color = item[2]
        pp.fill(x_list,y_list,color=this_color,edgecolor='none')
    
   
    pp.axis([-0.5,0.5,-0.5,0.5])

    pp.show()
    
    
    

This work is licensed under a Creative Commons Attribution By license.