Add track-incidents command

f1/api/stats.py

@@ -519,6 +519,36 @@ def compare_lap_telemetry_delta(ref_lap: Telemetry, comp_lap: Telemetry) -> np.n
     return comp_interp - ref_times
 
 
+def get_dnf_results(session: Session):
+    """Filter the results to only drivers who retired and include their final lap."""
+
+    driver_nums = [d for d in session.drivers if session.get_driver(d).dnf]
+    dnfs = session.results.loc[session.results["DriverNumber"].isin(driver_nums)].reset_index(drop=True)
+    # Get the retirement lap number
+    dnfs["LapNumber"] = [session.laps.pick_drivers(d)["LapNumber"].astype(int).max() for d in driver_nums]
+
+    return dnfs
+
+
+def get_track_events(session: Session):
+    """Return a DataFrame with lap number and event description, e.g. safety cars."""
+
+    incidents = (
+        Laps(session.laps.loc[:, ["LapNumber", "TrackStatus"]].dropna())
+        .pick_track_status("1456", how="any")
+        .groupby("LapNumber").min()
+        .reset_index()
+    )
+    # Map the status codes to names
+    incidents["Event"] = incidents["TrackStatus"].apply(utils.map_track_status)
+
+    # Mark the first occurance of consecutive status events by comparing against neighbouring row
+    # Allows filtering to keep only the first lap where the status occured until the next change
+    incidents["Change"] = (incidents["Event"] != incidents["Event"].shift(1)).astype(int)
+
+    return incidents[incidents["Change"] == 1]
+
+
 def results_table(results: pd.DataFrame, name: str) -> tuple[Figure, Axes]:
     """Return a formatted matplotlib table from a session results dataframe.
 
@@ -671,3 +701,32 @@ def sectors_table(df: pd.DataFrame) -> tuple[Figure, Axes]:
     del df
 
     return table.figure, table.ax
+
+
+def incidents_table(df: pd.DataFrame) -> tuple[Figure, Axes]:
+    """Return table listing track retirements and status events."""
+    df = df.rename(columns={"LapNumber": "Lap"})
+    col_defs = [
+        ColDef("Lap", width=0.15, textprops={"weight": "bold"}, border="r"),
+        ColDef("Event", width=0.5, textprops={"ha": "right"})
+    ]
+    # Dynamic size
+    size = (4, (df["Lap"].size / 3.333) + 1)
+    table = plot_table(df, col_defs, "Lap", size)
+
+    # Styling
+    for cell in table.columns["Event"].cells:
+        if cell.content in ("Safety Car", "Virtual Safety Car"):
+            cell.text.set_color("#ffb300")
+            cell.text.set_alpha(0.84)
+        elif cell.content == "Red Flag":
+            cell.text.set_color("#e53935")
+            cell.text.set_alpha(0.84)
+        elif cell.content == "Green Flag":
+            cell.text.set_color("#43a047")
+            cell.text.set_alpha(0.84)
+        else:
+            cell.text.set_color((1, 1, 1, 0.5))
+
+    del df
+    return table.figure, table.ax

f1/cogs/race.py

@@ -10,6 +10,7 @@
 from f1 import options, utils
 from f1.api import ergast, stats
 from f1.config import Config
+from f1.errors import MissingDataError
 from f1.target import MessageTarget
 
 logger = logging.getLogger('f1-bot')
@@ -246,6 +247,46 @@ async def career(self, ctx: ApplicationContext, driver: options.DriverOptionRequ
 
         await target.send(embed=embed)
 
+    @commands.slash_command(
+        name="track-incidents",
+        description="Summary of race events including Safety Cars and retirements.")
+    async def track_incidents(self, ctx: ApplicationContext,
+                              year: options.SeasonOption, round: options.RoundOption):
+        """Outputs a table showing the lap number and event, such as Safety Car or Red Flag."""
+        await utils.check_season(ctx, year)
+        ev = await stats.to_event(year, round)
+        s = await stats.load_session(ev, 'R', laps=True)
+
+        if not s.f1_api_support:
+            raise MissingDataError("Track status data unavailable.")
+
+        # Get the driver DNFs
+        dnfs = stats.get_dnf_results(s)
+        # Combine driver code and dnf reason into single column for merging
+        dnfs["Event"] = dnfs.apply(lambda row: f"Retired: {row['Abbreviation']} ({row['Status']})", axis=1)
+        dnfs = dnfs.loc[:, ["LapNumber", "Event"]]
+
+        # Get track status events grouped by lap number
+        track_events = stats.get_track_events(s)
+
+        if dnfs["Event"].size == 0 and track_events["Event"].size == 0:
+            raise MissingDataError("No track events in this race.")
+
+        # Combine the driver retirements and track status events
+        incidents = pd.concat([
+            dnfs.loc[:, ["LapNumber", "Event"]],
+            track_events.loc[:, ["LapNumber", "Event"]]
+        ], ignore_index=True).sort_values(by="LapNumber").reset_index(drop=True)
+        incidents["LapNumber"] = incidents["LapNumber"].astype(int)
+
+        table, ax = stats.incidents_table(incidents)
+        ax.set_title(
+            f"{ev['EventDate'].year} {ev['EventName']}\nTrack Incidents"
+        ).set_fontsize(12)
+
+        f = utils.plot_to_file(table, f"incidents_{ev['EventDate'].year}_{ev['RoundNumber']}")
+        await MessageTarget(ctx).send(file=f)
+
 
 def setup(bot: discord.Bot):
     bot.add_cog(Race(bot))

f1/utils.py

@@ -342,3 +342,12 @@ def get_driver_or_team_color(id: str, session: Session, team_only=False, api_onl
             pass
 
     return f"#{session.get_driver(id)['TeamColor']}"
+
+
+def map_track_status(status: str):
+    """Return the description of a track status code."""
+    names = {"4": "Safety Car", "5": "Red Flag", "6": "Virtual Safety Car", "1": "Green Flag"}
+    for code in status:
+        if code in names:
+            return names[code]
+    return None