Implement the add_gap functionality for the dashboard

app.py

@@ -8,16 +8,9 @@
 from dash import Dash, Input, Output, State, callback
 from plotly import graph_objects as go
 
+import f1_visualization.plotly_dash.graphs as pg
 from f1_visualization._consts import CURRENT_SEASON, SPRINT_FORMATS
-from f1_visualization.plotly_dash.graphs import (
-    stats_distplot,
-    stats_lineplot,
-    stats_scatterplot,
-    strategy_barplot,
-)
-from f1_visualization.plotly_dash.layout import (
-    app_layout,
-)
+from f1_visualization.plotly_dash.layout import app_layout, line_y_options, scatter_y_options
 from f1_visualization.visualization import get_session_info, load_laps
 
 Session_info: TypeAlias = tuple[int, str, list[str]]
@@ -33,17 +26,36 @@ def df_convert_timedelta(df: pd.DataFrame) -> pd.DataFrame:
     The pd.Timedelta type is not JSON serializable.
     Columns with this data type need to be dropped or converted.
     """
-    # The Time column is dropped directly since its information is retained by LapTime
-    df = df.drop(columns=["Time"])
-    # PitOUtTime and PitInTime contains information that we might need later
-    df[["PitInTime", "PitOutTime"]] = df[["PitInTime", "PitOutTime"]].fillna(
-        pd.Timedelta(0, unit="ms")
-    )
-    df["PitInTime"] = df["PitInTime"].dt.total_seconds()
-    df["PitOutTime"] = df["PitOutTime"].dt.total_seconds()
+    timedelta_columns = ["Time", "PitInTime", "PitOutTime"]
+    # usually the Time column has no NaT values
+    # it is included here for consistency
+    df[timedelta_columns] = df[timedelta_columns].ffill()
+
+    for column in timedelta_columns:
+        df[column] = df[column].dt.total_seconds()
     return df
 
 
+def add_gap(driver: str, df_laps: pd.DataFrame) -> pd.DataFrame:
+    """
+    Calculate the gap to a certain driver.
+
+    Compared to the implementation in visualization.py. Here we assume
+    that the Time column has been converted to float type and that df_laps
+    contain laps from one round only.
+
+    The second assumption is checked during merging.
+    """
+    df_driver = df_laps[df_laps["Driver"] == driver][["LapNumber", "Time"]]
+    timing_column_name = f"{driver}Time"
+    df_driver = df_driver.rename(columns={"Time": timing_column_name})
+
+    df_laps = df_laps.merge(df_driver, on="LapNumber", validate="many_to_one")
+    df_laps[f"GapTo{driver}"] = df_laps["Time"] - df_laps[timing_column_name]
+
+    return df_laps.drop(columns=timing_column_name)
+
+
 def configure_lap_numbers_slider(data: dict) -> tuple[int, list[int], dict[int, str]]:
     """Configure range slider based on the number of laps in a session."""
     if not data:
@@ -65,18 +77,6 @@ def configure_lap_numbers_slider(data: dict) -> tuple[int, list[int], dict[int,
 app.layout = app_layout
 
 
-@callback(
-    Output("plotly-warning", "is_open"),
-    Input("plotly-warning-toggle", "n_clicks"),
-    prevent_initial_call=True,
-)
-def toggle_plotly_warning(n_clicks: int) -> bool:
-    """Toggle the visibility of the Plotly warning."""
-    # start with n_clicks = 0 and open
-    # so odd n_clicks imply the warning should be closed
-    return (n_clicks % 2) != 1
-
-
 @callback(
     Output("event", "options"),
     Output("event", "value"),
@@ -137,18 +137,15 @@ def set_session_options(event: str | None, schedule: dict) -> tuple[list[dict],
     Input("season", "value"),
     Input("event", "value"),
     Input("session", "value"),
+    prevent_initial_call=True,
 )
 def enable_load_session(season: int | None, event: str | None, session: str | None) -> bool:
     """Toggles load session button on when the previous three fields are filled."""
     return not (season is not None and event is not None and session is not None)
 
 
 @callback(
-    Output("drivers", "options"),
-    Output("drivers", "value"),
-    Output("drivers", "disabled"),
     Output("session-info", "data"),
-    Output("laps", "data"),
     Input("load-session", "n_clicks"),
     State("season", "value"),
     State("event", "value"),
@@ -162,29 +159,107 @@ def get_session_metadata(
     event: str,
     session: str,
     teammate_comp: bool,
-) -> tuple[list[str], list, bool, Session_info, dict]:
+) -> Session_info:
     """
-    Store session metadata and populate driver dropdown.
+    Store round number, event name, and the list of drivers into browser cache.
 
     Can assume that season, event, and session are all set (not None).
     """
-    round_number, event_name, drivers = get_session_info(
-        season, event, session, teammate_comp=teammate_comp
-    )
+    return get_session_info(season, event, session, teammate_comp=teammate_comp)
+
+
+@callback(
+    Output("laps", "data"),
+    Input("load-session", "n_clicks"),
+    State("season", "value"),
+    State("event", "value"),
+    State("session", "value"),
+    prevent_initial_call=True,
+)
+def get_session_laps(
+    _: int,  # ignores actual_value of n_clicks
+    season: int,
+    event: str,
+    session: str,
+) -> dict:
+    """
+    Save the laps of the selected session into browser cache.
 
+    Can assume that season, event, and session are all set (not None).
+    """
     included_laps = DF_DICT[season][session]
-    included_laps = included_laps[included_laps["RoundNumber"] == round_number]
+    included_laps = included_laps[included_laps["EventName"] == event]
     included_laps = df_convert_timedelta(included_laps)
 
+    return included_laps.to_dict()
+
+
+@callback(
+    Output("drivers", "options"),
+    Output("drivers", "value"),
+    Output("drivers", "disabled"),
+    Output("gap-drivers", "options"),
+    Output("gap-drivers", "value"),
+    Output("gap-drivers", "disabled"),
+    Input("session-info", "data"),
+    prevent_initial_call=True,
+)
+def set_driver_dropdowns(session_info: Session_info):
+    """Configure driver dropdowns."""
+    drivers = session_info[2]
+    return drivers, drivers, False, drivers, None, False
+
+
+@callback(
+    Output("scatter-y", "options"),
+    Output("line-y", "options"),
+    Output("scatter-y", "value"),
+    Output("line-y", "value"),
+    Input("laps", "data"),
+    prevent_initial_call=True,
+)
+def set_y_axis_dropdowns(
+    data: dict,
+) -> tuple[list[dict[str, str]], list[dict[str, str]], str, str]:
+    """Update y axis options based on the columns in the laps dataframe."""
+
+    def readable_gap_col_name(col: str) -> str:
+        """Convert Pandas GapTox column names to the more readable Gap to x."""
+        return f"Gap to {col[-3:]} (s)"
+
+    gap_cols = filter(lambda x: x.startswith("Gap"), data.keys())
+    gap_col_options = [{"label": readable_gap_col_name(col), "value": col} for col in gap_cols]
     return (
-        drivers,
-        drivers,
-        False,
-        (round_number, event_name, drivers),
-        included_laps.to_dict(),
+        scatter_y_options + gap_col_options,
+        line_y_options + gap_col_options,
+        "LapTime",
+        "LapTime",
     )
 
 
+@callback(
+    Output("laps", "data", allow_duplicate=True),
+    Input("add-gap", "n_clicks"),
+    State("gap-drivers", "value"),
+    State("laps", "data"),
+    running=[
+        (Output("gap-drivers", "disabled"), True, False),
+        (Output("add-gap", "disabled"), True, False),
+        (Output("add-gap", "children"), "Calculating...", "Add Gap"),
+        (Output("add-gap", "color"), "warning", "success"),
+    ],
+    prevent_initial_call=True,
+)
+def add_gap_to_driver(_: int, drivers: list[str], data: dict) -> dict:
+    """Amend the dataframe in cache and add driver gap columns."""
+    laps = pd.DataFrame.from_dict(data)
+    for driver in drivers:
+        if f"GapTo{driver}" not in laps.columns:
+            laps = add_gap(driver, laps)
+
+    return laps.to_dict()
+
+
 @callback(
     Output("lap-numbers-scatter", "max"),
     Output("lap-numbers-scatter", "value"),
@@ -227,7 +302,7 @@ def render_strategy_plot(
     included_laps = included_laps[included_laps["Driver"].isin(drivers)]
 
     event_name = session_info[1]
-    fig = strategy_barplot(included_laps, drivers)
+    fig = pg.strategy_barplot(included_laps, drivers)
     fig.update_layout(title=event_name)
     return fig
 
@@ -262,7 +337,7 @@ def render_scatterplot(
         & (included_laps["LapNumber"].isin(lap_interval))
     ]
 
-    fig = stats_scatterplot(included_laps, drivers, y)
+    fig = pg.stats_scatterplot(included_laps, drivers, y)
     event_name = session_info[1]
     fig.update_layout(title=event_name)
 
@@ -301,7 +376,7 @@ def render_lineplot(
         & (included_laps["LapNumber"].isin(lap_interval))
     ]
 
-    fig = stats_lineplot(included_laps, drivers, y, upper_bound)
+    fig = pg.stats_lineplot(included_laps, drivers, y, upper_bound)
     event_name = session_info[1]
     fig.update_layout(title=event_name)
 
@@ -333,7 +408,7 @@ def render_distplot(
         & (included_laps["PctFromFastest"] < (upper_bound - 100))
     ]
 
-    fig = stats_distplot(included_laps, drivers, boxplot)
+    fig = pg.stats_distplot(included_laps, drivers, boxplot)
     event_name = session_info[1]
     fig.update_layout(title=event_name)
 

f1_visualization/plotly_dash/graphs.py

@@ -194,8 +194,6 @@ def stats_lineplot(
     included_laps: pd.DataFrame, drivers: list[str], y: str, upper_bound: int
 ) -> go.Figure:
     """Make lineplots showing a statistic."""
-    # TODO: incorporate the add_gap functionality
-
     # Identify SC and VSC laps before filtering for upper bound
     sc_laps, vsc_laps = _find_sc_laps(included_laps)
 
@@ -230,7 +228,7 @@ def stats_lineplot(
         )
 
     fig = shade_sc_periods(fig, sc_laps, vsc_laps)
-    if y == "Position":
+    if y == "Position" or y.startswith("Gap"):
         fig.update_yaxes(autorange="reversed")
 
     num_laps = included_laps["LapNumber"].max()