Skip to content

Commit

Permalink
Strategy example, using session.drivers to plot in correct order
Browse files Browse the repository at this point in the history
  • Loading branch information
Casper-Guo committed May 18, 2023
1 parent fd0f9fb commit 724ee30
Showing 1 changed file with 85 additions and 0 deletions.
85 changes: 85 additions & 0 deletions examples/plot_strategy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
"""
Tyre strategies during a race
=============================
Plot all drivers' tyre strategies during a race.
"""

import fastf1
from matplotlib import pyplot as plt

###############################################################################
# Load the race session

session = fastf1.get_session(2022, "Hungary", 'R')
session.load()
laps = session.laps

###############################################################################
# Get the list of driver numbers
drivers = session.drivers
print(drivers)

###############################################################################
# Convert the driver numbers to three letter abbreviations
drivers = [session.get_driver(driver)["Abbreviation"] for driver in drivers]
print(drivers)

###############################################################################
# We need to find the stint length and compound used
# for every stint by every driver.
# We do this by first grouping the laps by the driver,
# the stint number, and the compound.
# And then counting the number of laps in each group.
stints = laps[["Driver", "Stint", "Compound", "LapNumber"]]
stints = stints.groupby(["Driver", "Stint", "Compound"])
stints = stints.count().reset_index()

###############################################################################
# The number in the LapNumber column now stands for the number of observations
# in that group aka the stint length.
stints = stints.rename(columns={"LapNumber": "StintLength"})
print(stints)

###############################################################################
# Now we can plot the strategies for each driver
fig, ax = plt.subplots(figsize=(5, 10))

for driver in drivers:
driver_stints = stints.loc[stints["Driver"] == driver]

previous_stint_end = 0
for idx, row in driver_stints.iterrows():
# each row contains the compound name and stint length
# we can use these information to draw horizontal bars
plt.barh(
y=driver,
width=row["StintLength"],
left=previous_stint_end,
color=fastf1.plotting.COMPOUND_COLORS[row["Compound"]],
edgecolor="black",
fill=True
)

previous_stint_end += row["StintLength"]

# sphinx_gallery_defer_figures

###############################################################################
# Make the plot more readable and intuitive
plt.title("2022 Hungarian Grand Prix Strategies")
plt.xlabel("Lap Number")
plt.grid(False)
# invert the y-axis so drivers that finish higher are closer to the top
ax.invert_yaxis()

# sphinx_gallery_defer_figures

###############################################################################
# Plot aesthetics
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)

plt.tight_layout()
plt.show()

0 comments on commit 724ee30

Please sign in to comment.