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

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 for weather change windows
        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()