From 724ee3010e3c5f738883ef9058a70d7c7970d8fb Mon Sep 17 00:00:00 2001
From: Casper Guo <casperg@umich.edu>
Date: Thu, 11 May 2023 21:34:40 -0400
Subject: [PATCH] Strategy example, using session.drivers to plot in correct
 order

---
 examples/plot_strategy.py | 85 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 examples/plot_strategy.py

diff --git a/examples/plot_strategy.py b/examples/plot_strategy.py
new file mode 100644
index 000000000..0ef73bbf8
--- /dev/null
+++ b/examples/plot_strategy.py
@@ -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()