plot.py

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors
import random
from collections import Counter

def calculate_sizes(data):
    """Calculate sizes based on frequency of values."""
    flat_list = [item for sublist in data for item in sublist]
    count = Counter(flat_list)
    return [count[val] * 50 for val in flat_list]  # Multiply by a factor to scale sizes

# Original Data
times = [[1, 6, 9, 3], [6, 2, 23, 7], [1, 3, 3, 6], [6, 3, 2, 1], [9, 7, 34, 12], [31, 32, 4, 6], [12, 65, 3, 1], [5, 5, 3, 2]]

# New Data with random, higher scale values
new_times = [[random.randint(50, 300) for _ in range(4)] for _ in range(8)]

# Flatten the lists
flat_list = [item for sublist in times for item in sublist]
flat_list_new = [item for sublist in new_times for item in sublist]

# Calculate sizes for scatter points based on frequency
sizes_times = calculate_sizes(times)
sizes_new_times = calculate_sizes(new_times)

# Calculate the overall average for each time point
average_over_time = [np.mean(sublist) for sublist in times]
new_average_over_time = [np.mean(sublist) for sublist in new_times]


# Set up the figure and subplots
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 12))
for ax in [ax1, ax2]:
    ax.set_facecolor('#CECBD6')  # Set background to grey
    ax.grid(True, which='both', color='white', linestyle='-', linewidth=0.5)
    ax.grid(True, which='minor', color='white', linestyle='-', linewidth=0.2)  # White minor gridlines
    ax.minorticks_on()


# Gradient for original data
norm = mcolors.Normalize(vmin=min(flat_list), vmax=max(flat_list))
cmap = mcolors.LinearSegmentedColormap.from_list("val_color_new", ["#4158D0","#EF7D7D", "#FFEC51"])

# Plot original data with gradient and varying sizes
for i, time_set in enumerate(times):
    colors = [cmap(norm(val)) for val in time_set]
    sizes = sizes_times[i * len(time_set):(i + 1) * len(time_set)]
    ax1.scatter([i + 1] * len(time_set), time_set, color=colors, s=sizes,  marker='s')

ax1.plot(range(1, len(average_over_time) + 1), average_over_time, color='black', label='Average Over Time')


# Gradient for new data
norm_new = mcolors.Normalize(vmin=min(flat_list_new), vmax=max(flat_list_new))
cmap_new = mcolors.LinearSegmentedColormap.from_list("val_color_new", ["#4158D0","#EF7D7D", "#FFEC51"])

# Plot new data with gradient and varying sizes
for i, time_set in enumerate(new_times):
    colors = [cmap_new(norm_new(val)) for val in time_set]
    sizes = sizes_new_times[i * len(time_set):(i + 1) * len(time_set)]
    ax2.scatter([i + 1] * len(time_set), time_set, color=colors, marker='s', s=sizes)

ax2.plot(range(1, len(new_average_over_time) + 1), new_average_over_time, color='black', label='Average Over Time')

# Adding titles, labels, and legends
ax1.set_title('Scatter Plot - Dataset 1')
ax1.set_ylabel('Values (Dataset 1)')

ax2.set_title('Scatter Plot - Dataset 2')
ax2.set_xlabel('Time Set')
ax2.set_ylabel('Values (Dataset 2)')

# Show the plot

plt.show()