pythonmatplotlibpolygonvoronoigraph-coloring

How to Select Colors for a PatchCollection of Voronoi Polygons and Create a Discrete Colorbar for it?


I'm trying to color a Voronoi object according to the number of neighbors it has. I have created a list of colors according to this number, which ranges from 4 to 7. I then set the array of the PatchCollection to the set of neighbor numbers. This technically works, however, it's selecting some really ugly colors and the colorbar on the side is continuous while it should be discrete. I would prefer to make it so that <= 4 neighbors is blue, 5 neighbors is green, 6 neighbors is grey, >=7 neighbors is red. Any ideas on how to resolve these issues? Code:

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.collections import LineCollection 
    from scipy.spatial import Voronoi
    import curved_analysis as ca
    from matplotlib import patches 
    from matplotlib.collections import PatchCollection

    def vor_plot(particles):
        vor = Voronoi(particles[0,:,:2])

        trial_ridges = vor.ridge_vertices

        line_info = []
        for first, last in trial_ridges:
            if -1 not in (first, last):
                line_info.append([vor.vertices[first], vor.vertices[last]])



        vor_poly = Voronoi(particles[0,:,:2])
        regions = vor_poly.regions

        real_regions = []
        for inner_list in regions:
            if -1 not in inner_list:
                real_regions.append(inner_list)
        real_regions.remove([])
    
        fig, ax = plt.subplots()

        vor_poly = []
        colors = []
        for gon in real_regions:
            xy = vor.vertices[gon]
            vor_poly.append(patches.Polygon(xy))
            colors.append(xy.shape[0])




        lc = LineCollection(line_info, color='k', lw=0.5)

        ax.add_collection(lc)
        ax.scatter(vor.points[:,0], vor.points[:,1], s = 3)
        ax.set_xlim([vor.points[:,0].min()-5, vor.points[:,0].max()+5])
        ax.set_ylim([vor.points[:,1].min()-5, vor.points[:,1].max()+5])

        colors = np.array(colors)
        p = PatchCollection(vor_poly, alpha=0.3)
        p.set_array(colors)
        fig.colorbar(p, ax=ax)
        ax.add_collection(p)
        plt.show()

    if __name__ == "__main__":
        particles = ca.read_xyz("flat.xyz")
        vor_plot(particles)

Solution

  • You can create a ListedColormap listing the desired colors. To decide which number maps to which color, a norm can be used, fixing 4 for the first color and 7 for the last. Both the colormap and norm need to be assigned to the PatchCollection. To position the tick labels, one can divide the range of 4 colored cells into 9 equally-spaced positions and take the ones at odd indexes.

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.collections import LineCollection, PatchCollection
    from scipy.spatial import Voronoi
    from matplotlib import patches
    from matplotlib.colors import ListedColormap
    
    particles = np.random.rand(1, 20, 2) * 100
    
    vor = Voronoi(particles[0, :, :2])
    
    trial_ridges = vor.ridge_vertices
    
    line_info = []
    for first, last in trial_ridges:
        if -1 not in (first, last):
            line_info.append([vor.vertices[first], vor.vertices[last]])
    
    vor_poly = Voronoi(particles[0, :, :2])
    regions = vor_poly.regions
    
    real_regions = []
    for inner_list in regions:
        if -1 not in inner_list:
            real_regions.append(inner_list)
    real_regions.remove([])
    
    fig, ax = plt.subplots()
    
    vor_poly = []
    colors = []
    for gon in real_regions:
        xy = vor.vertices[gon]
        vor_poly.append(patches.Polygon(xy))
        colors.append(xy.shape[0])
    
    lc = LineCollection(line_info, color='k', lw=0.5)
    
    ax.add_collection(lc)
    ax.scatter(vor.points[:, 0], vor.points[:, 1], s=3)
    ax.set_xlim([vor.points[:, 0].min() - 5, vor.points[:, 0].max() + 5])
    ax.set_ylim([vor.points[:, 1].min() - 5, vor.points[:, 1].max() + 5])
    
    cmap = ListedColormap(['dodgerblue', 'limegreen', 'grey', 'crimson'])
    colors = np.array(colors)
    p = PatchCollection(vor_poly, alpha=0.3, cmap=cmap, norm=plt.Normalize(4, 7))
    p.set_array(colors)
    ax.add_collection(p)
    cbar = fig.colorbar(p, ticks=np.linspace(4, 7, 9)[1::2], ax=ax)
    cbar.ax.set_yticklabels(['≤ 4', '5', '6', '≥ 7'])
    plt.show()
    

    example voronoi diagram