diff --git a/examples/plot_strategy.py b/examples/plot_strategy.py new file mode 100644 index 000000000..1fdeeb845 --- /dev/null +++ b/examples/plot_strategy.py @@ -0,0 +1,77 @@ +""" +Tyre strategies during a race +============================= + +Plot all drivers' tyre strategies during a race. +""" + +import fastf1 as f +import fastf1.plotting as p +from matplotlib import pyplot as plt + +############################################################################### +# Load the race session and create the plot + +session = f.get_session(2022, "Hungary", 'R') +session.load() + +fig, ax = plt.subplots(figsize=(5, 10)) + +laps = session.laps + +############################################################################### +# Get the list of drivers +drivers = laps["Driver"].unique() +print(drivers) + +############################################################################### +# We need to find the stint length and compound used +# for every stint by every driver +# We do this by grouping the laps by the driver, +# the stint number, and the compound +# And 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 stands for the stint length +stints = stints.rename(columns={"LapNumber": "StintLength"}) +print(stints) + +############################################################################### +# Now we can plot the strategies for each driver + +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=p.COMPOUND_COLORS[row["Compound"]], + edgecolor="black", + fill=True + ) + + previous_stint_end += row["StintLength"] + +############################################################################### +# Make the plot more readable +plt.title("2022 Hungarian Grand Prix Strategies") +plt.xlabel("Lap Number") +plt.grid(False) +ax.invert_yaxis() + +############################################################################### +# Plot aesthetics +ax.spines['top'].set_visible(False) +ax.spines['right'].set_visible(False) +ax.spines['left'].set_visible(False) + +plt.savefig("test.png")