Better data vizualization? about cosmic-rai HOT 1 OPEN

import matplotlib.pyplot as plt
from cosmic_rAI.data_prep import *

mat1 = np.load('../data/sim_12360_00.npy').item()  # protons
mat2 = np.load('../data/sim_12362_00.npy').item()  # iron
matrices = [mat1, mat2]

event_df = event_df_from_matrices(matrices)
sensor_df = sensor_df_from_matrices(matrices)

# merge sensor data wth (transposed) charge data
vals = sensor_df.merge(
    event_df['charges'].T,
    how='outer', left_index=True, right_index=True)

# seperate high/ low gain
high_vals = vals[vals['gain']=='High'].drop('gain', axis=1)
low_vals = vals[vals['gain']=='Low'].drop('gain', axis=1)

# choose an event
event = 11

# x,y should be same in high_vals and low_vals
x, y = high_vals['pos_x'], high_vals['pos_y']

# construct plot
fig, ax = plt.subplots()
ax.scatter(x, y, s=high_vals[event], c='r', label='high_gain')
ax.scatter(x, y, s=low_vals[event], c='b', label='high_gain')
ax.set_title('Event #{}'.format(event))
ax.legend()