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

# TODO: Adjust input parameters for multiple weather change windows
class DataPlotter(DataHandler):
    def plotOvertakesWithTireChangeWindow(self, overtakesPerLap: list[int], firstTireChange: int, lastTireChange: int, raceName: str):

        # TODO: Account for cases where no tire changes (meaning values of 0 or -1)
        # TODO: Let visualization start at 0, but graph at 1
        # Adjust values
        #overtakesPerLap.insert(0, 0) # Insert 0th lap, which cannot have overtakes
        firstTireChange -= 1
        lastTireChange += 1


        laps: int = len(overtakesPerLap)
        plt.xlim(0, laps + 1) # set x-range of plot


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



        axis = plt.gca()

        # 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)

        # Add legend
        axis.legend(bbox_to_anchor=(1.0, 1.02))

        # Paint background between the first and last tire change
        axis.axvspan(firstTireChange, lastTireChange, color='blue', alpha=0.3)



        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()