cl-plot.py

#!/usr/bin/env python3
import numpy as np
import seaborn as sns
import matplotlib.pylab as plt
import os
import sys

np.set_printoptions(threshold=sys.maxsize, linewidth=np.inf)

exp_data = np.zeros((6,7))
obs_data = np.zeros((6,7))
directory = 'monosbb'
for filename in os.listdir(directory):
    path = os.path.join(directory, filename)
    mhs = int(filename.split('-')[1])
    mzp = int(filename.split('-')[2].split('.')[0])
    with open(path, "r") as file:
        lines = file.readlines()
        exp_data[(mhs-50)//20,mzp//500-1] = lines[1].split()[3]
        obs_data[(mhs-50)//20,mzp//500-1] = lines[1].split()[4]
    # print(f"{mhs} {mzp}: {mzp//500},{(mhs-50)//20}")
# print(exp_data)
# print(obs_data)

def plot_heatmap(data, name):
    ax = sns.heatmap(exp_data, annot=True, linewidth=0.5)

    ax.set_xlabel("M_Z' (GeV)")
    ax.set_ylabel("M_S (GeV)")
    ax.set_title(f'CL {name} Results')
    ax.invert_yaxis()
    ax.set_yticklabels(list(range(50, 170, 20)))
    ax.set_xticklabels(list(range(500, 4000, 500)))
    plt.show()
    fig = plt.figure()

    fig.savefig('plot.png')

plot_heatmap(obs_data, 'Observed')