Consistent colour across Matplotlib Groupby Graph

Question

I am attempting to create a scatter pie plot that groups by 2 columns, Column1 & Column 2 where the colour in the pie (if numbers are the same) is decided by Column 3.

See my example below of where I am:

This graph shows Column 1 (y-axis) and Column 2 (x-axis). The colour is dictated by Column 3.

But with the code I use the colours do not stay consistent across graphs and if the same Column 3 appears with a different Column 2 or Column 3 value it assigns it a different colour.

I have attempted using cmaps and manually assigning colours but I cannot keep it consistent across each column 2.

See my current code below:

import pandas as pd 
import matplotlib.pyplot as plt
import matplotlib.ticker as mticks
from matplotlib.font_manager import FontProperties
import numpy as np

def draw_pie(dist, 
             xpos, 
             ypos, 
             size, 
             color,
             ax=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=(70,60))

    # for incremental pie slices
    cumsum = np.cumsum(dist)
    cumsum = cumsum/ cumsum[-1]
    pie = [0] + cumsum.tolist()

    for r1, r2 in zip(pie[:-1], pie[1:]):
        angles = np.linspace(2 * np.pi * r1, 2 * np.pi * r2)
        x = [0] + np.cos(angles).tolist()
        y = [0] + np.sin(angles).tolist()

        xy = np.column_stack([x, y])

        ax.scatter([xpos], [ypos], marker=xy, s=size,c=color)

    return ax


colors = {'Group A':'red', 'Group B':'green', 'Group C':'blue', 'Group D':'yellow', 'Group E':'yellow', 'Group F':'yellow', 'Group G':'yellow', 'Group H':'yellow'}


fig, ax = plt.subplots(figsize=(94,70))
for (x,y), d in dataset.groupby(['Column 1','Column 2']):
    dist = d['Column 3'].value_counts()
    draw_pie(dist, x, y, 50000, ax=ax,color=dataset['Column 3'].map(colors))


params = {'legend.fontsize': 100}
plt.rcParams.update(params)
#plt.legend(dataset["Column 3"],markerscale=.4,frameon=True,framealpha=1,ncol=3,loc=(0.00, -0.3), bbox_to_anchor=(0.0, 0., 0.5, 1.25),handletextpad=1,markerfirst=True,facecolor='lightgrey',mode='expand',borderaxespad=-16)


ax.yaxis.set_major_locator(mticks.MultipleLocator(1))
full = plt.Rectangle((-0.05, 4.25), 2.10, 2, color='g', alpha=0.15)
partial = plt.Rectangle((-0.05, 2.25), 2.10, 2, color='orange', alpha=0.15)
low = plt.Rectangle((-0.05, 0.25), 2.10, 2, color='r', alpha=0.15)
ax.add_patch(full)
ax.add_patch(partial)
ax.add_patch(low)
plt.xticks(fontsize=120)
plt.yticks(fontsize=100)
plt.ylim([0, 6.75]) 
plt.tight_layout()
plt.show()

Ideally the output graph based on the data (I will copy in below) should be like the below graph (I have placed a number in each pie to define what colour should be there)

Here is the full data used for the graph:

Column 1    3       2   Colour Group Desired
First Line  Group A 6   1
First Line  Group A 6   1
First Line  Group A 6   1
First Line  Group C 6   3
First Line  Group B 6   2
First Line  Group B 6   2
First Line  Group B 6   2
First Line  Group A 6   1
First Line  Group A 6   1
First Line  Group C 6   3
First Line  Group A 6   1
Second Line Group A 6   1
Second Line Group A 6   1
Second Line Group A 6   1
Second Line Group C 6   3
Second Line Group B 6   2
Second Line Group B 6   2
Second Line Group B 6   2
Second Line Group A 4.5 1
Second Line Group A 6   1
Second Line Group C 6   3
Second Line Group A 6   1
Third Line  Group A 1   1
Third Line  Group A 6   1
Third Line  Group A 1   1
Third Line  Group C 6   3
Third Line  Group B 3.5 2
Third Line  Group B 3.5 2
Third Line  Group B 3.5 2
Third Line  Group A 1   1
Third Line  Group A 1   1
Third Line  Group C 4   3
Third Line  Group A 1   1

Additionally I would like to add a label in each section of the pie with the count of distinct(Column 3).

user_14536792 · Answer

Currently I came up with the following solution to fix the issue with the colours:

import pandas as pd 
import matplotlib.pyplot as plt
import matplotlib.ticker as mticks
from matplotlib.font_manager import FontProperties
import numpy as np

def draw_pie(dist, 
             xpos, 
             ypos, 
             size, 
             ax=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=(7,6))
    
    # The colors, corresponding to the values 1, 2 and 3:
    # 1 is tab:blue
    # 2 is tab:orange
    # 3 is tab:green
    # Of course, you can change this
    colors = ['tab:blue', 'tab:orange', 'tab:green']
    
    # for incremental pie slices
    cumsum = np.cumsum(dist)
    cumsum = cumsum/ cumsum[-1]
    pie = [0] + cumsum.tolist()

    for r1, r2, i in zip(pie[:-1], pie[1:], range(0, len(dist))):
    
        # If no counts present, skip this one
        if dist[i] == 0:
            continue
    
        angles = np.linspace(2 * np.pi * r1, 2 * np.pi * r2)
        x = [0] + np.cos(angles).tolist()
        y = [0] + np.sin(angles).tolist()

        xy = np.column_stack([x, y])
        ax.scatter([xpos], [ypos], marker=xy, s=size, color=colors[i])

    return ax
    
#colors = {'Group A':'red', 'Group B':'green', 'Group C':'blue', 'Group D':'yellow', 'Group E':'yellow', 'Group F':'yellow', 'Group G':'yellow', 'Group H':'yellow'}

# Read dataset
dataset = pd.read_csv('dataset.csv')

fig, ax = plt.subplots(figsize=(9,5))

for (x,y), d in dataset.groupby(['Column 1','Column 2']):

    # Only interested in the 'Column 3' column, as this one
    # contains the values 1-2-3
    d = d['Colour Group Desired']
    
    # Count how often each value (1-2-3) occurs and store
    # this in a list (count for value i located at list index 
    # i-1)   
    dist = list()
    for i in [1,2,3]:
        dist.append(d[d==i].count())
        
        
    # Call your draw_pie function
    draw_pie(dist, x, y, 500, ax=ax)
   
    
ax.yaxis.set_major_locator(mticks.MultipleLocator(1))
full = plt.Rectangle((-0.05, 4.25), 2.10, 2, color='g', alpha=0.15)
partial = plt.Rectangle((-0.05, 2.25), 2.10, 2, color='orange', alpha=0.15)
low = plt.Rectangle((-0.05, 0), 2.10, 2.25, color='r', alpha=0.15)
ax.add_patch(full)
ax.add_patch(partial)
ax.add_patch(low)
plt.xticks(fontsize=10)
plt.yticks(fontsize=8)
plt.ylim([0, 6.75]) 
plt.tight_layout()
plt.show()

Small note to start with, I changed all sizes by a factor of 10 such that the plot fits on my screen (e.g. the figsize). You probably want to use your original values again, but this doesn't matter for the question anyway.

The first change I made was to the loop body of the for (x,y), d in dataset.groupby(['Column 1','Column 2']) loop. Instead of using dist = d['Column 3'].value_counts(), I create an empty array. Subsequently, I loop over the values 1, 2 and 3. In each iteration, I check how many rows match the specific value and append the outcome to the list. In this way, I end up with a list of size 3, in which the first element corresponds to the amount of rows that equal 1, the second element corresponds to the amount of rows that equal 2 and the third element corresponds to the amount of rows that equal 3. The advantage is that I can also keep track of values that occur 0 times.

Secondly, I changed the draw_pie function a bit. However, since I do not fully understand the meaning of a group in terms of the colour, I commented out the colors dictionary. It looks as if 1 always corresponds to Group A, 2 always corresponds to Group B and 3 always corresponds to Group C. I made use of this observation and defined another colors variable (in the draw_pie function). Instead of a dictionary, colors is now a list (where the first element corresponds to the value 1, the second element corresponds to the value 2 and the third element corresponds to the value 3). I changed your for loop from for r1, r2 in zip(pie[:-1], pie[1:]) to for r1, r2, i in zip(pie[:-1], pie[1:], range(0, len(dist))). The advantage is that I can now use the iteration variable i to get the proper color from the colors list. In addition, I added a small if statement which checks if there are exactly 0 occurences. If that is the case, I just skip the remainder of the loop and draw nothing (if you don't skip these cases, it will draw a very thin line instead. Try this for yourself by removing it).

If I run the code, I get the following result:

Unfortunately, I was not successful in adding the labels. I tried using the Axes.text method, but I couldn't get the labels to be placed at the proper locations.

Edit

I decided to change the body of the draw_pie function. In this new version, we draw an Axes instance at the desired (xpos, ypos) location. This involves some transformations: first a transformation from data coordinates to display coordinates and subsequently a transformation from display coordinates to figure coordinates. See this tutorial for an explanation. The advantage is that we can now plot a pie chart inside the created axes using the Axes.pie method. This method has some nice options, such as adding labels!

However, there is a catch. Before we start drawing the pie charts, we need to already fix the xlim and the ylim values of the main Axes. If we don't do this (and do it after plotting the pie charts), the pie charts will no longer be located at the proper positions. Therefore, I have moved the code which sets the xlim and ylim values before the first time we call the draw_pie function. I also removed to call to plt.tight_layout(), as this will (unfortunately) also cause the pie charts to no longer be located at the proper locations.

As a small side note, I changed the manner in which the background colors are drawn. Instead of using patches, I now use the Axes.axhspan method. With this method, you can still control the y-locations, but the width will extend infinitely (meaning that the colors remain if you scroll left/right). If you don't want this, you can remove it again :).

See the new version of the code below (I note once more, I changed all sizes such that it fits on my computer screen. You probably want to substitute your original sizes again):

import pandas as pd 
import matplotlib.pyplot as plt
import matplotlib.ticker as mticks
from matplotlib.font_manager import FontProperties
import numpy as np

def draw_pie(dist, 
             xpos, 
             ypos, 
             size, 
             ax=None,
             fig=None):
    if ax is None:
        fig, ax = plt.subplots(figsize=(7,6))
    
    # Transform xpos and ypos to figure coordinates by the
    # following steps:
    # 1. First transform xpos and ypos to data coordinates
    # 2. Transform the data coordinates back to figure coordinates
    xfig, yfig   = ax.transData.transform((xpos, ypos))
    trans_to_fig = fig.transFigure.inverted()
    xfig, yfig   = trans_to_fig.transform((xfig, yfig))
    
    # Calculate figure coordinates from the desired pie chart size
    # given in pixels
    size = trans_to_fig.transform((size, 0))[0]
    
    # Add axes at data coordinates xpos and ypos. On these axes,
    # the pie chart will be plotted
    ax_pie = fig.add_axes([xfig-0.5*size, yfig-0.5*size, size, size])
    
    # Plot the pie chart (with some special options)
    textprops = {'color'      : 'w',
                 'fontweight' : 'bold', 
                 'fontsize'   :  10,
                 'va'         : 'center',
                 'ha'         : 'center'
                }            
    labels = [str(i) if not i == 0 else "" for i in dist]
    labeldistance = 0.5
    
    if sum(not x==0 for x in dist) == 1: # Ensures we plot the label in the center
        labeldistance = 0.0              # if we have only one entry
    
    ax_pie.pie(dist, labels=labels, labeldistance=labeldistance, textprops=textprops)

    return ax_pie
    

# Read dataset
dataset = pd.read_csv('dataset.txt')

fig, ax = plt.subplots(figsize=(9,5))

# Important, limits must be set before calling draw_pie function!
# (Otherwise, the data coordinates will change, which will break
# the transform sequence inside the draw_pie function!)
ax.set_xlim([-0.2, 2.2]) # Tweak these values for the desired output
ax.set_ylim([0, 6.75])

# Make sure the string from 'Column 1' is displayed again
ax.set_xticks([0, 1, 2])
ax.set_xticklabels(['First Line', 'Second Line', 'Third Line'])

# Remainder of formatting
ax.yaxis.set_major_locator(mticks.MultipleLocator(1))  
plt.xticks(fontsize=10)
plt.yticks(fontsize=8)

# Define float values for 'Column 1' (easier for transformation,
# we have already put the text back there using ax.set_xticklabels)
column1_to_float = {'First Line':0, 'Second Line':1, 'Third Line':2}

for (x,y), d in dataset.groupby(['Column 1','Column 2']):

    # Only interested in the 'Column 3' column, as this one
    # contains the values 1-2-3
    d = d['Colour Group Desired']
    
    # Count how often each value (1-2-3) occurs and store
    # this in a list (count for value i located at list index 
    # i-1)   
    dist = list()
    for i in [1,2,3]:
        dist.append(d[d==i].count())
        
    # Call your draw_pie function
    draw_pie(dist, column1_to_float[x], y, 100, ax=ax, fig=fig)

# Plot the colours (note: using axhspan, they extend the full 
# horizontal direction, even while scrolling)
ax.axhspan(0   , 2.25, fc='r'     , ec=None, alpha=0.15)
ax.axhspan(2.25, 4.25, fc='orange', ec=None, alpha=0.15)
ax.axhspan(4.25, 6.75, fc='g'     , ec=None, alpha=0.15)

# Unfortunately, tight_layout can no longer be used. If we do use this,
# the pie charts will no longer be at the proper positions...
# plt.tight_layout()
plt.show()

If I run this code, I get the following output:

Just to mention it, you can decrease the size of the pie charts by adjusting the size argument to the draw_pie function (I just liked the output above :) ). But keep in mind that in this case, you also want to decrease the fontsize specified in the textprops dictionary in the body of the draw_pie function. As an example (size = 63, fontsize=7):

Consistent colour across Matplotlib Groupby Graph

1 Answers

Edit