Plotting multiple horizontal bars within the same y tick-CodePudding

I'm working one a personal project where I'd like to plot the best times of each Formula 1 driver during the Qualifying Session (which takes place in 3 parts : Q1, Q2 and Q3)

The plot would be with the drivers' name (order by their final results : first on the top and last in the bottom) on the y-axis and their time on the x-axis as shown below with an Excel representation for 4 drivers :

With Python, I'm able to collect their best time for each session (Q1, Q2 and Q3) that I store as a list (t_Q1,t_Q2,t_Q3) in a list, and their final position. Here is the beginning of the code :

import matplotlib.pyplot as plt
import numpy as np


Drivers_name=['VER','HAM','VET','MSC'] #for the label of the Y-axis
Drivers_position=[1,2,12,20]
Drivers_timings=[[94.352,93.464,92.91],[94.579,93.797,93.119],[95.281,95.5],[96.119]]


fig,ax = plt.subplots()

y=np.arange(len(Drivers_name))

for i in range(len(Drivers_name)):
    if len(Drivers_timings[i])==3:
        eps=0.2
        ax.barh(y-eps,Drivers_timings[i][0],height=eps,color='g', label='Q1')
        ax.barh(y,Drivers_timings[i][1],height=eps,color='r', label='Q2')
        ax.barh(y eps,Drivers_timings[i][1],height=eps,color='b', label='Q3')
    elif len(Drivers_timings[i])==2:
        eps=0.2
        ax.barh(y-eps/2,Drivers_timings[i][0],height=eps,color='g', label='Q1')
        ax.barh(y eps/2,Drivers_timings[i][1],height=eps,color='r', label='Q2')
    elif len(Drivers_timings[i])==1:
        ax.barh(y,Drivers_timings[i][0],height=0.2,color='g', label='Q1')

plt.yticks(y,Drivers_name)
plt.legend()
plt.show()

But this is what I get when I execute it:

We can notice that it does not really represent what I was looking for and that the legend shows the same elements many time (which is normal as the labelargument is in the for loop but I can't thing of any other ideas)

Review : To sum it up, here are the purposes of my post

Obtain the expected plot (the one I get with Excel)
Fixing the legend box

CodePudding user response：

By using ax.barh(y-eps, ...) where y=np.arange(len(Drivers_name)), in each step of the loop bars are created for all y-values. E.g. for i=0 three bars will be created around y=0, but also 3 around y=1, 3 more around y=2 and 3 around y=3. Replacing y by i creates the desired effect.

The legend contains a line for each time label= is used. In the full plot, ax.barh will be called 9 times, each time with label=. A solution is to create the legend from dummy rectangles.

import matplotlib.pyplot as plt
import numpy as np

Drivers_name = ['VER', 'HAM', 'VET', 'MSC']  # for the label of the Y-axis
Drivers_position = [1, 2, 12, 20]
Drivers_timings = [[94.352, 93.464, 92.91], [94.579, 93.797, 93.119], [95.281, 95.5], [96.119]]

colors = ['g', 'r', 'b']
labels = ['Q1', 'Q2', 'Q3']
fig, ax = plt.subplots()
for i in range(len(Drivers_name)):
    eps = 0.2
    if len(Drivers_timings[i]) == 3:
        ax.barh(i - eps, Drivers_timings[i][0], height=eps, color=colors[0])
        ax.barh(i, Drivers_timings[i][1], height=eps, color=colors[1])
        ax.barh(i   eps, Drivers_timings[i][1], height=eps, color=colors[2])
    elif len(Drivers_timings[i]) == 2:
        ax.barh(i - eps / 2, Drivers_timings[i][0], height=eps, color=colors[0])
        ax.barh(i   eps / 2, Drivers_timings[i][1], height=eps, color=colors[1])
    elif len(Drivers_timings[i]) == 1:
        ax.barh(i, Drivers_timings[i][0], height=eps, color=colors[0])

plt.yticks(np.arange(len(Drivers_name)), Drivers_name)
handles = [plt.Rectangle((0, 0), 0, 0, color=color) for color in colors]
plt.legend(handles=handles, labels=labels, bbox_to_anchor=(1.01, 0.99), loc='upper left')
plt.tight_layout()
plt.show()

The code could further be shortened, and at the same time made more general by also drawing the individual bars in an own loop:

colors = ['g', 'r', 'b']
labels = ['Q1', 'Q2', 'Q3']
fig, ax = plt.subplots()
eps = 0.2
for i, driver_timings in enumerate(Drivers_timings):
    num_bars = len(driver_timings)
    first_bar = i - eps * (num_bars - 1) / 2
    for j, (driver_timings_j, color) in enumerate(zip(driver_timings, colors)):
        ax.barh(first_bar   j * eps, driver_timings_j, height=eps, color=color)

plt.yticks(np.arange(len(Drivers_name)), Drivers_name)
handles = [plt.Rectangle((0, 0), 0, 0, color=color) for color in colors]
plt.legend(handles=handles, labels=labels, bbox_to_anchor=(1.01, 0.99), loc='upper left')
plt.tight_layout()
plt.show()