Issue
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 label
argument 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
Solution
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()
Answered By - JohanC
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.