DataPlotter.py

import random

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib import cm
import matplotlib as mpl
import numpy as np
from matplotlib.colors import ListedColormap

from DataHandler import DataHandler, WeatherChangeWindow

# TODO: Adjust input parameters for multiple weather change windows
class DataPlotter(DataHandler):
    def plotOvertakes(self, overtakesPerLap: list[int], raceName: str, weatherChangeWindows: list[WeatherChangeWindow] = None):

        # Define variables
        laps: int = len(overtakesPerLap)
        axis = plt.gca()

        # Set x-range of plot
        plt.xlim(0, laps + 1)

        # Define data range
        x_values = np.arange(1, laps + 1)
        plt.plot(x_values, overtakesPerLap)

        # Paint background between first and last tire change
        for weatherChangeWindow in weatherChangeWindows:
            axis.axvspan(weatherChangeWindow.firstLap, weatherChangeWindow.lastLap, color='blue', alpha=0.3)

        # Label stuff
        axis.set_xlabel('Lap')
        axis.set_ylabel('Position changes in lap')
        axis.set_title(raceName)

        # Set ticks on axis
        major_ticks_laps = np.arange(0, laps + 1, 5)
        major_ticks_overtakes = np.arange(0, max(overtakesPerLap) + 1, 5)
        minor_ticks_laps = np.arange(0, laps + 1, 1)
        minor_ticks_overtakes = np.arange(0, max(overtakesPerLap) + 1, 1)
        axis.set_xticks(major_ticks_laps)
        axis.set_xticks(minor_ticks_laps, minor=True)
        axis.set_yticks(major_ticks_overtakes)
        axis.set_yticks(minor_ticks_overtakes, minor=True)

        # Add grid
        axis.grid(which='both')
        axis.grid(which='minor', alpha=0.2)
        axis.grid(which='major', alpha=0.5)

        plt.tight_layout()
        plt.show()


    def plotBackgroundPaintTest(self):

        laps: int = 50
        overtakesPerLap: list[int] = [random.randint(0,10) for _ in range(laps)]
        firstTireChange: int = random.randint(10, 20)
        lastTireChange: int = random.randint(30, 40)

        isWeatherChanging: list[bool] = [False] * laps
        for lap in range(laps):
            if firstTireChange <= lap <= lastTireChange:
                isWeatherChanging[lap] = True

        colormap = mpl.colors.ListedColormap(['white', 'blue'])
        dataFrame = pd.DataFrame(overtakesPerLap)
        dataFrame[1] = isWeatherChanging
        axis = dataFrame[0].plot()
        axis.pcolorfast(axis.get_xlim(), axis.get_ylim(),
                        dataFrame[1].values[np.newaxis],
                        cmap=colormap, alpha=0.3)

        plt.show()