plotting fun stuff in f1

f1plot.py

"""Overlaying speed traces of two laps
======================================

Compare two fastest laps by overlaying their speed traces.
"""

import numpy as np
import matplotlib.pyplot as plt
import fastf1.plotting
import itertools as it


fastf1.Cache.enable_cache('doc_cache')  # replace with your cache directory

# enable some matplotlib patches for plotting timedelta values and load
# FastF1's default color scheme
fastf1.plotting.setup_mpl()

# load a session and its telemetry data
session = fastf1.get_session(2022, 2, 'FP2')
session.load()

##############################################################################
# First, we select the two laps that we want to compare

# teams = ['FER', 'MER', 'RBR']
# drivers = ['LEC', 'RUS', '1']
teams = ['FER', 'RBR', 'MER']
drivers = ['LEC', '1', 'HAM']

# Select which lap you care about
lap_num = []
for driver in drivers:
    print("-"*50)
    print(driver)
    for idx, lap in session.laps.pick_driver(driver).iterlaps():
        print("LAP {} {}/{}: {}".format(idx, lap["Compound"], lap['TyreLife'], lap['LapTime']))
    s = input("Pick lap num: ")
    lap_num.append(s)

telemetry = []
for driver, lap in zip(drivers, lap_num):
    if lap.isnumeric():
        laps = session.laps.pick_driver(driver)
        fastest_lap = laps.loc[int(lap)]
    else:
        fastest_lap = session.laps.pick_driver(driver).pick_fastest()

    # use padding so that there are values outside of the desired range for accurate interpolation later
    car_data = fastest_lap.get_car_data(pad=1, pad_side='both')
    pos_data = fastest_lap.get_pos_data(pad=1, pad_side='both')

    merged_data = car_data.merge_channels(pos_data)

    # slice again to remove the padding and interpolate the exact first and last value
    merged_data = merged_data.slice_by_lap(fastest_lap, interpolate_edges=True)

    telemetry.append(merged_data.add_distance())

# Now we try to detect where corners are by reading the telemetry of all cars
def find_corners(tel):
    corners = []
    tel.reset_index()
    corner_cnt = 1
    in_corner = False
    start, end = None, None

    for _, row in tel.iterrows():
        if (row['CurrentAction'] == 'Cornering') and not in_corner:
            in_corner = True
            s = int(row['Distance'])
            if not start or start > s:
                start = s

        if (row['CurrentAction'] == 'Full Throttle') and in_corner:
            in_corner = False
            corner_cnt += 1
            corners.append((start, int(row['Distance'])))
            start = None
    return corners


corners_by_car = []
# Decorate telemetry with driver action
for tel in telemetry:
    tel.loc[tel['Brake'] > 0, 'CurrentAction'] = 'Cornering'
    tel.loc[tel['Throttle'] > 97, 'CurrentAction'] = 'Full Throttle'
    tel.loc[(tel['Brake'] == 0) & (tel['Throttle'] < 97), 'CurrentAction'] = 'Cornering'

    # find corners for lap
    corners_by_car.append(find_corners(tel))


def overlap(og, corner):
    (start1, end1) = og
    (start2, end2) = corner
    return end1 >= start2 and end2 >= start1

# different telemetries could have a different set of corners
# we put them all together, remove duplicates and then extend a corner
# if different cars create overlapping corners.
def merge_corners(corners_by_car):
    c = sorted(set(it.chain(*corners_by_car)))
    newc = []
    singles = []

    i = 0
    while i < len(c):
        # somehow in saudi arabia this is needed to remove something bad from the data
        if c[i] == (612, 619):
            i += 1
            continue
        if singles == []:
            singles.append(c[i])
            i += 1
            continue

        if any(overlap(s, c[i]) for s in singles):
            singles.append(c[i])
        else:
            corn = (min(s[0] for s in singles), max(s[1] for s in singles))
            newc.append(corn)
            singles = [c[i]]

        i += 1

    if singles:
        newc.append((min(s[0] for s in singles), max(s[1] for s in singles)))

    return newc

corners = merge_corners(corners_by_car)
print(corners)

# this is a very not precise attempt at calculating the time delta of corners.
# we attempt by interpolating  the position and time curves from the 3 cars.
# while the interpolation works, there's not enough points in the source data
# to make this a precise calculation
def mini_pro(stream):
    # Ensure that all samples are interpolated
    dstream_start = stream[1] - stream[0]
    dstream_end = stream[-1] - stream[-2]
    return np.concatenate([[stream[0] - dstream_start], stream, [stream[-1] + dstream_end]])

all_tels = []
for tel in telemetry:
    all_tels.extend(tel['Distance'].tolist())
distances = np.array(sorted(set(all_tels)))
fully_interpolated_laps = []
for i, ref in enumerate(telemetry):
    ltime = mini_pro(ref['Time'].dt.total_seconds().to_numpy())
    ldistance = mini_pro(ref['Distance'].to_numpy())
    lap_time = np.interp(distances, ldistance, ltime)

    fully_interpolated_laps.append(lap_time)



# Now calculate the speeds by corner as well as the center of the corner
# it also calculates the time, but that isn't displayed because it's broken
speeds_by_corner = []
corner_loc = []
for corner in corners:
    speed_by_driver = []
    loc_by_driver = []
    for tel, interp_lap in zip(telemetry, fully_interpolated_laps):
        corner_speeds = tel['Speed'].loc[
                (tel['Distance'] >= corner[0]) & (tel['Distance'] <= corner[1])
            ]
        #print("{} {}".format(len(corner_speeds), corner_speeds.tolist()))
        corner_time_min = interp_lap[distances.searchsorted(corner[0])]
        corner_time_max = interp_lap[distances.searchsorted(corner[1])]

        speed_by_driver.append((corner_speeds.min(), corner_speeds.mean(), corner_time_max-corner_time_min))
        loc = tel['Distance'].loc[(tel['Speed'] == speed_by_driver[-1][0]) & (tel['Distance'] >= corner[0]) & (tel['Distance'] <= corner[1])]
        loc_by_driver.append(
            loc.mean()
        )

    speeds_by_corner.append(speed_by_driver)
    corner_loc.append(sum(loc_by_driver)/len(loc_by_driver))


##############################################################################
# Finally, we create a plot and plot both speed traces.
# We color the individual lines with the driver's team colors.

colors = []
for team in teams:
    # if team == "ALF":
    #     colors.append('yellow')
    #     continue
    colors.append(fastf1.plotting.team_color(team))

metrics = ['Speed', 'Throttle', 'Brake', 'RPM'] #, 'DRS', 'nGear']

fig, axs = plt.subplots(len(metrics)+1)

for corner in corners:
    print(corner, (corner[0]+corner[1])/2)

for i, ax in enumerate(axs):
    for label in ax.xaxis.get_ticklabels():
        label.set_fontsize(3)

    for label in ax.yaxis.get_ticklabels():
        label.set_fontsize(3)

    if i < 1:
        # first graph is to display speeds, we deal with this later
        continue

    metric = metrics[i-1]
    ax.set_ylabel(metric, size=5)

    if metric == 'Speed':
        for corner in corners:
            ax.axvspan(corner[0], corner[1], color='gray', alpha=0.5, lw=0)

    for j, driver in enumerate(drivers):
        color = colors[j]
        tel = telemetry[j]
        ax.plot(tel['Distance'], tel[metric], linewidth=0.2, color=color, label=driver)
        for xc in corner_loc:
            ax.axvline(x=int(xc), linewidth=0.3, color="white")

# Now deal with the first graph
ax = axs[0]
ax.set_xlim(axs[1].get_xlim())
for xc in corner_loc:
    ax.axvline(x=int(xc), linewidth=0.3, color="white")

next_y_pos = it.cycle([1.05, 0.83, 0.60, 0.37, 0.14])
for corner, speed_by_driver  in zip(corners, speeds_by_corner):
    center = (corner[0]+corner[1])/2
    best_speed = max(s[0] for s in speed_by_driver)
    best_avg = max(s[1] for s in speed_by_driver)
    best_time = min(s[2] for s in speed_by_driver)

    text = []
    for driver, driver_speed in zip(drivers, speed_by_driver):
        text.append("{}: min {:+.0f}/avg {:+.2f}".format( #/T {:+.3f}s".format(
            driver,
            driver_speed[0]-best_speed,
            driver_speed[1]-best_avg))
            #(driver_speed[2]-best_time)))

    final = "\n".join(text)
    props = dict(boxstyle='round', facecolor='white', alpha=1, edgecolor='none')
    ax.text(center, next(next_y_pos), final, fontsize=3, color='black', verticalalignment='top', bbox=props)

# last set of things to tidy up
ax = axs[-1]
ax.xaxis.set_visible(True)
ax.legend(loc=4, fontsize=4)
ax.set_xlabel('Distance in m', size='xx-small')
plt.suptitle(f"Fastest Lap Comparison \n "
             f"{session.weekend.name} {session.weekend.year}", size="xx-small")
plt.subplots_adjust(left=0.08, right=0.96, top=0.93, bottom=0.08)
plt.savefig("test.png", dpi=600)
# plt.show()