FIX: per-lap track status incorrect in some cases + cleanup

theOehrly · Mar 24, 2024 · c59c2dc · c59c2dc
1 parent 3098384
commit c59c2dc
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 26 deletions.
diff --git a/fastf1/core.py b/fastf1/core.py
@@ -1847,48 +1847,52 @@ def _add_track_status_to_laps(self, laps):
         if track_status is None:
             return
 
-        # first set all laps to green flag as a starting point
-        laps['TrackStatus'] = '1'
+        # ensure track status is not set
+        laps['TrackStatus'] = ''
 
-        def applicator(new_status, current_status):
-            if current_status == '1':
-                return new_status
-            elif new_status not in current_status:
+        def _applicator(new_status, current_status):
+            if new_status not in current_status:
                 return current_status + new_status
             else:
                 return current_status
 
+        # -- Track Status Timeline
+        #           --> (status before) --|--- status ---|-- next_status -->
+        #                                 |              |
+        #                                 t              next_t
+        # -- Lap Timeline ---------------------------------------------------
+        # Case A (end criterion):    ----> Lap --|
+        # Case B (start criterion):                |---- Lap --->
+        #    (matches B and C)               |-- Lap --|
+        # Case C (full overlap):     |---------- Lap ----------|
+
         if len(track_status['Time']) > 0:
             t = track_status['Time'][0]
             status = track_status['Status'][0]
             for next_t, next_status in zip(track_status['Time'][1:],
                                            track_status['Status'][1:]):
-                if status != '1':
-                    # status change partially in lap and partially outside
-                    sel = (((next_t >= laps['LapStartTime'])
-                            & (laps['LapStartTime'] >= t))
-                           | ((t <= laps['Time']) & (laps['Time'] <= next_t)))
-
-                    laps.loc[sel, 'TrackStatus'] \
-                        = laps.loc[sel, 'TrackStatus'].apply(
-                            lambda curr: applicator(status, curr)
-                    )
-
-                    # status change two times in one lap (short yellow flag)
-                    sel = ((laps['LapStartTime'] <= t)
-                           & (laps['Time'] >= next_t))
 
-                    laps.loc[sel, 'TrackStatus'] \
-                        = laps.loc[sel, 'TrackStatus'].apply(
-                            lambda curr: applicator(status, curr)
-                    )
+                # Case A: The lap ends during the current status
+                sel = ((t <= laps['Time']) & (laps['Time'] <= next_t))
+                # Case B: The lap starts during the current status
+                sel |= ((t <= laps['LapStartTime'])
+                        & (laps['LapStartTime'] <= next_t))
+                # Case C: The lap fully contains the current status
+                sel |= ((laps['LapStartTime'] <= t) & (next_t <= laps['Time']))
+
+                laps.loc[sel, 'TrackStatus'] \
+                    = laps.loc[sel, 'TrackStatus'].apply(
+                        lambda curr: _applicator(status, curr)
+                )
 
                 t = next_t
                 status = next_status
 
-            sel = laps['LapStartTime'] >= t
+            # process the very last status: any lap that ends after this status
+            # started was fully or partially set under this track status
+            sel = (t <= laps['Time'])
             laps.loc[sel, 'TrackStatus'] = laps.loc[sel, 'TrackStatus'].apply(
-                lambda curr: applicator(status, curr)
+                lambda curr: _applicator(status, curr)
             )
 
     @soft_exceptions("first lap time",

diff --git a/fastf1/tests/test_core.py b/fastf1/tests/test_core.py
@@ -6,6 +6,7 @@
     DriverResult,
     Lap,
     Laps,
+    Session,
     SessionResults
 )
 from fastf1.ergast import Ergast
@@ -86,3 +87,39 @@ def test_session_results_constructor_sliced():
 
     assert isinstance(results.loc[:, 'A'], pd.Series)
     assert not isinstance(results.loc[:, 'A'], DriverResult)
+
+
+def test_add_lap_status_to_laps():
+    # TODO: It should really be possible to mock a session object instead of
+    #   modifying an existing one like here. This is incredibly hack.
+    session = fastf1.get_session(2020, 'Italy', 'R')
+
+    laps = Laps(
+        [[pd.Timedelta(minutes=1), pd.Timedelta(minutes=2)],
+         [pd.Timedelta(minutes=2), pd.Timedelta(minutes=3)],
+         [pd.Timedelta(minutes=3), pd.Timedelta(minutes=4)],
+         [pd.Timedelta(minutes=4), pd.Timedelta(minutes=5)],
+         [pd.Timedelta(minutes=5), pd.Timedelta(minutes=6)],
+         [pd.Timedelta(minutes=6), pd.Timedelta(minutes=7)],
+         [pd.Timedelta(minutes=7), pd.Timedelta(minutes=8)]],
+        force_default_cols=False,
+        columns=('LapStartTime', 'Time')
+    )
+
+    status = pd.DataFrame(
+        [[pd.Timedelta(minutes=0), '1', 'AllClear'],
+         [pd.Timedelta(minutes=2.5), '2', 'Yellow'],
+         [pd.Timedelta(minutes=3.25), '6', 'VSCDeployed'],
+         [pd.Timedelta(minutes=3.75), '7', 'VSCEnding'],
+         [pd.Timedelta(minutes=4.25), '1', 'AllClear'],
+         [pd.Timedelta(minutes=6.5), '2', 'Yellow']],
+        columns=('Time', 'Status', 'Message')
+    )
+
+    # modify and reuse the existing session (very hacky but ok here)
+    session._track_status = status
+    session._add_track_status_to_laps(laps)
+
+    expected_per_lap_status = ['1', '12', '267', '71', '1', '12', '2']
+
+    assert (laps['TrackStatus'] == expected_per_lap_status).all()